From 5bf5afaec83d4e68cbfd5220cd760a7940e7dda5 Mon Sep 17 00:00:00 2001
From: Jason Volk <jason@zemos.net>
Date: Sun, 30 Mar 2025 01:54:55 +0000
Subject: [PATCH] instrument tokio before/after poll hooks

Signed-off-by: Jason Volk <jason@zemos.net>
---
 src/main/runtime.rs | 24 ++++++++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/src/main/runtime.rs b/src/main/runtime.rs
index b1657289..920476db 100644
--- a/src/main/runtime.rs
+++ b/src/main/runtime.rs
@@ -61,6 +61,8 @@ pub(super) fn new(args: &Args) -> Result<tokio::runtime::Runtime> {
 	#[cfg(tokio_unstable)]
 	builder
 		.on_task_spawn(task_spawn)
+		.on_before_task_poll(task_enter)
+		.on_after_task_poll(task_leave)
 		.on_task_terminate(task_terminate);
 
 	#[cfg(tokio_unstable)]
@@ -215,3 +217,25 @@ fn task_spawn(meta: &tokio::runtime::TaskMeta<'_>) {}
 	),
 )]
 fn task_terminate(meta: &tokio::runtime::TaskMeta<'_>) {}
+
+#[cfg(tokio_unstable)]
+#[tracing::instrument(
+	name = "enter",
+	level = "trace",
+	skip_all,
+	fields(
+		id = %meta.id()
+	),
+)]
+fn task_enter(meta: &tokio::runtime::TaskMeta<'_>) {}
+
+#[cfg(tokio_unstable)]
+#[tracing::instrument(
+	name = "leave",
+	level = "trace",
+	skip_all,
+	fields(
+		id = %meta.id()
+	),
+)]
+fn task_leave(meta: &tokio::runtime::TaskMeta<'_>) {}