feat(console): add large future lints

In Tokio, the futures for tasks are stored on the stack unless they are explicitly boxed. Having very large futures can be problematic as it can cause a stack overflow. This change makes use of new instrumentation in Tokio (tokio-rs/tokio#6881) which includes the size of the future which drives a task. The size isn't given its own column (we're running out of horizontal space) and appears in the additional fields column. In the case that a future was auto-boxed by Tokio, the original size of the task will also be provided. Two new lints have been added for large futures. The first will detect auto-boxed futures and warn about them. The second will detect futures which are large (over 1024 bytes by default), but have not been auto-boxed by the runtime. Since the new lints depend on the new instrumentation in Tokio, they will only trigger if the instrumented application is running using `tokio` 1.41.0 or later. The version is as yet, unreleased. Both lints have been added to the default list.
tokio-rs · Oct 8, 2024 · 92894ad · 92894ad
1 parent ca39033
commit 92894ad
Show file tree

Hide file tree

Showing 9 changed files with 219 additions and 7 deletions.
diff --git a/README.md b/README.md
@@ -224,8 +224,13 @@ Options:
 
           * `never-yielded` -- Warns when a task has never yielded.
 
-          [default: self-wakes lost-waker never-yielded]
-          [possible values: self-wakes, lost-waker, never-yielded]
+          * `large-future` -- Warnings when the future driving a task
+          occupies a large amount of stack space.
+
+          [default: self-wakes lost-waker never-yielded
+          auto-boxed-future large-future]
+          [possible values: self-wakes, lost-waker, never-yielded,
+          auto-boxed-future, large-future]
 
   -A, --allow <ALLOW_WARNINGS>...
           Allow lint warnings.
@@ -243,9 +248,13 @@ Options:
 
           * `never-yielded` -- Warns when a task has never yielded.
 
+          * `large-future` -- Warnings when the future driving a task
+          occupies a large amount of stack space.
+
           If this is set to `all`, all warnings are allowed.
 
-          [possible values: all, self-wakes, lost-waker, never-yielded]
+          [possible values: all, self-wakes, lost-waker, never-yielded,
+          large-future]
 
       --log-dir <LOG_DIRECTORY>
           Path to a directory to write the console's internal logs to.

diff --git a/console-subscriber/examples/app.rs b/console-subscriber/examples/app.rs
@@ -51,6 +51,23 @@ async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
                     .spawn(spawn_blocking(5))
                     .unwrap();
             }
+            "large" => {
+                tokio::task::Builder::new()
+                    .name("pretty-big")
+                    // Below debug mode auto-boxing limit
+                    .spawn(large_future::<1024>())
+                    .unwrap();
+                tokio::task::Builder::new()
+                    .name("huge")
+                    // Larger than the release mode auto-boxing limit
+                    .spawn(large_future::<20_000>())
+                    .unwrap();
+                tokio::task::Builder::new()
+                    .name("huge-blocking-wait")
+                    // Larger than the release mode auto-boxing limit
+                    .spawn(large_blocking::<20_000>())
+                    .unwrap();
+            }
             "help" | "-h" => {
                 eprintln!("{}", HELP);
                 return Ok(());
@@ -152,6 +169,41 @@ async fn spawn_blocking(seconds: u64) {
     }
 }
 
+#[tracing::instrument]
+async fn large_future<const N: usize>() {
+    let mut numbers = [0_u8; N];
+
+    loop {
+        for idx in 0..N {
+            numbers[idx] = (idx % 256) as u8;
+            tokio::time::sleep(Duration::from_millis(100)).await;
+            (0..=idx).for_each(|jdx| {
+                assert_eq!(numbers[jdx], (jdx % 256) as u8);
+            });
+        }
+    }
+}
+
+async fn large_blocking<const N: usize>() {
+    let numbers = [0_usize; N];
+
+    tokio::task::Builder::new()
+        .name("huge-blocking")
+        .spawn_blocking(move || {
+            let mut numbers = numbers;
+            loop {
+                for idx in 0..N {
+                    numbers[idx] = idx;
+                    std::thread::sleep(Duration::from_millis(100));
+                    (0..=idx).for_each(|jdx| {
+                        assert_eq!(numbers[jdx], jdx);
+                    });
+                }
+            }
+        })
+        .unwrap();
+}
+
 fn self_wake() -> impl Future<Output = ()> {
     struct SelfWake {
         yielded: bool,

diff --git a/tokio-console/Cargo.toml b/tokio-console/Cargo.toml
@@ -62,3 +62,4 @@ hyper-util = { version = "0.1.6", features = ["tokio"] }
 
 [dev-dependencies]
 trycmd = "0.15.4"
+
diff --git a/tokio-console/console.example.toml b/tokio-console/console.example.toml
@@ -4,6 +4,8 @@ warnings = [
     'self-wakes',
     'lost-waker',
     'never-yielded',
+    'auto-boxed-future',
+    'large-future',
 ]
 log_directory = '/tmp/tokio-console/logs'
 retention = '6s'

diff --git a/tokio-console/src/config.rs b/tokio-console/src/config.rs
@@ -63,6 +63,9 @@ pub struct Config {
     /// * `lost-waker` -- Warns when a task is dropped without being woken.
     ///
     /// * `never-yielded` -- Warns when a task has never yielded.
+    ///
+    /// * `large-future` -- Warnings when the future driving a task occupies a large amount of
+    ///                     stack space.
     #[clap(long = "warn", short = 'W', value_delimiter = ',', num_args = 1..)]
     #[clap(default_values_t = KnownWarnings::default_enabled_warnings())]
     pub(crate) warnings: Vec<KnownWarnings>,
@@ -80,9 +83,12 @@ pub struct Config {
     ///
     /// * `never-yielded` -- Warns when a task has never yielded.
     ///
+    /// * `large-future` -- Warnings when the future driving a task occupies a large amount of
+    ///                     stack space.
+    ///
     /// If this is set to `all`, all warnings are allowed.
     ///
-    /// [possible values: all, self-wakes, lost-waker, never-yielded]
+    /// [possible values: all, self-wakes, lost-waker, never-yielded, large-future]
     #[clap(long = "allow", short = 'A', num_args = 1..)]
     pub(crate) allow_warnings: Option<AllowedWarnings>,
 
@@ -143,6 +149,8 @@ pub(crate) enum KnownWarnings {
     SelfWakes,
     LostWaker,
     NeverYielded,
+    AutoBoxedFuture,
+    LargeFuture,
 }
 
 impl FromStr for KnownWarnings {
@@ -153,6 +161,8 @@ impl FromStr for KnownWarnings {
             "self-wakes" => Ok(KnownWarnings::SelfWakes),
             "lost-waker" => Ok(KnownWarnings::LostWaker),
             "never-yielded" => Ok(KnownWarnings::NeverYielded),
+            "auto-boxed-future" => Ok(KnownWarnings::AutoBoxedFuture),
+            "large-future" => Ok(KnownWarnings::LargeFuture),
             _ => Err(format!("unknown warning: {}", s)),
         }
     }
@@ -164,6 +174,8 @@ impl From<&KnownWarnings> for warnings::Linter<Task> {
             KnownWarnings::SelfWakes => warnings::Linter::new(warnings::SelfWakePercent::default()),
             KnownWarnings::LostWaker => warnings::Linter::new(warnings::LostWaker),
             KnownWarnings::NeverYielded => warnings::Linter::new(warnings::NeverYielded::default()),
+            KnownWarnings::AutoBoxedFuture => warnings::Linter::new(warnings::AutoBoxedFuture),
+            KnownWarnings::LargeFuture => warnings::Linter::new(warnings::LargeFuture::default()),
         }
     }
 }
@@ -174,6 +186,8 @@ impl fmt::Display for KnownWarnings {
             KnownWarnings::SelfWakes => write!(f, "self-wakes"),
             KnownWarnings::LostWaker => write!(f, "lost-waker"),
             KnownWarnings::NeverYielded => write!(f, "never-yielded"),
+            KnownWarnings::AutoBoxedFuture => write!(f, "auto-boxed-future"),
+            KnownWarnings::LargeFuture => write!(f, "large-future"),
         }
     }
 }
@@ -184,6 +198,8 @@ impl KnownWarnings {
             KnownWarnings::SelfWakes,
             KnownWarnings::LostWaker,
             KnownWarnings::NeverYielded,
+            KnownWarnings::AutoBoxedFuture,
+            KnownWarnings::LargeFuture,
         ]
     }
 }

diff --git a/tokio-console/src/state/mod.rs b/tokio-console/src/state/mod.rs
@@ -277,6 +277,8 @@ impl Field {
     const KIND: &'static str = "kind";
     const NAME: &'static str = "task.name";
     const TASK_ID: &'static str = "task.id";
+    const SIZE_BYTES: &'static str = "size.bytes";
+    const ORIGINAL_SIZE_BYTES: &'static str = "original_size.bytes";
 
     /// Creates a new Field with a pre-interned `name` and a `FieldValue`.
     fn new(name: InternedStr, value: FieldValue) -> Self {

diff --git a/tokio-console/src/state/tasks.rs b/tokio-console/src/state/tasks.rs
@@ -104,6 +104,10 @@ pub(crate) struct Task {
     location: String,
     /// The kind of task, currently one of task, blocking, block_on, local
     kind: InternedStr,
+    /// The size of the future driving the task
+    size_bytes: Option<usize>,
+    /// The original size of the future (before runtime auto-boxing)
+    original_size_bytes: Option<usize>,
 }
 
 #[derive(Debug)]
@@ -184,6 +188,8 @@ impl TasksState {
                 let mut name = None;
                 let mut task_id = None;
                 let mut kind = strings.string(String::new());
+                let mut size_bytes = None;
+                let mut original_size_bytes = None;
                 let target_field = Field::new(
                     strings.string_ref("target"),
                     FieldValue::Str(meta.target.to_string()),
@@ -210,6 +216,24 @@ impl TasksState {
                                 kind = strings.string(field.value.to_string());
                                 None
                             }
+                            Field::SIZE_BYTES => {
+                                size_bytes = match field.value {
+                                    FieldValue::U64(size_bytes) => Some(size_bytes as usize),
+                                    _ => None,
+                                };
+                                // Include size in pre-formatted fields
+                                Some(field)
+                            }
+                            Field::ORIGINAL_SIZE_BYTES => {
+                                original_size_bytes = match field.value {
+                                    FieldValue::U64(original_size_bytes) => {
+                                        Some(original_size_bytes as usize)
+                                    }
+                                    _ => None,
+                                };
+                                // Include size in pre-formatted fields
+                                Some(field)
+                            }
                             _ => Some(field),
                         }
                     })
@@ -245,6 +269,8 @@ impl TasksState {
                     warnings: Vec::new(),
                     location,
                     kind,
+                    size_bytes,
+                    original_size_bytes,
                 };
                 if let TaskLintResult::RequiresRecheck = task.lint(linters) {
                     next_pending_lint.insert(task.id);
@@ -506,6 +532,14 @@ impl Task {
     pub(crate) fn location(&self) -> &str {
         &self.location
     }
+
+    pub(crate) fn size_bytes(&self) -> Option<usize> {
+        self.size_bytes
+    }
+
+    pub(crate) fn original_size_bytes(&self) -> Option<usize> {
+        self.original_size_bytes
+    }
 }
 
 enum TaskLintResult {

diff --git a/tokio-console/src/warnings.rs b/tokio-console/src/warnings.rs
@@ -258,3 +258,90 @@ impl Warn<Task> for NeverYielded {
         )
     }
 }
+
+/// Warning for if a task's driving future was auto-boxed by the runtime
+#[derive(Clone, Debug, Default)]
+pub(crate) struct AutoBoxedFuture;
+
+impl Warn<Task> for AutoBoxedFuture {
+    fn summary(&self) -> &str {
+        "tasks have been boxed by the runtime due to their size"
+    }
+
+    fn check(&self, task: &Task) -> Warning {
+        let (Some(size_bytes), Some(original_size_bytes)) =
+            (task.size_bytes(), task.original_size_bytes())
+        else {
+            return Warning::Ok;
+        };
+
+        if original_size_bytes != size_bytes {
+            Warning::Warn
+        } else {
+            Warning::Ok
+        }
+    }
+
+    fn format(&self, task: &Task) -> String {
+        let original_size = task
+            .original_size_bytes()
+            .expect("warning should not trigger if original size is None");
+        let boxed_size = task
+            .size_bytes()
+            .expect("warning should not trigger if size is None");
+        format!(
+            "This task was auto-boxed by the runtime due to its size (originally \
+            {original_size} bytes, boxed size {boxed_size} bytes)",
+        )
+    }
+}
+
+/// Warning for if a task's driving future if large
+#[derive(Clone, Debug)]
+pub(crate) struct LargeFuture {
+    min_size: usize,
+    description: String,
+}
+impl LargeFuture {
+    pub(crate) const DEFAULT_MIN_SIZE_BYTES: usize = 1024;
+    pub(crate) fn new(min_size: usize) -> Self {
+        Self {
+            min_size,
+            description: format!("tasks are large (threshold {} bytes)", min_size),
+        }
+    }
+}
+
+impl Default for LargeFuture {
+    fn default() -> Self {
+        Self::new(Self::DEFAULT_MIN_SIZE_BYTES)
+    }
+}
+
+impl Warn<Task> for LargeFuture {
+    fn summary(&self) -> &str {
+        self.description.as_str()
+    }
+
+    fn check(&self, task: &Task) -> Warning {
+        // Don't fire warning for tasks that are not async
+        if task.is_blocking() {
+            return Warning::Ok;
+        }
+
+        if let Some(size_bytes) = task.size_bytes() {
+            if size_bytes >= self.min_size {
+                return Warning::Warn;
+            }
+        }
+        Warning::Ok
+    }
+
+    fn format(&self, task: &Task) -> String {
+        format!(
+            "This task occupies a large amount of stack space ({} bytes)",
+            task.size_bytes()
+                .expect("warning should not trigger if size is None"),
+        )
+    }
+}
diff --git a/tokio-console/tests/cli-ui.stdout b/tokio-console/tests/cli-ui.stdout
@@ -53,8 +53,13 @@ Options:
 
           * `never-yielded` -- Warns when a task has never yielded.
 
-          [default: self-wakes lost-waker never-yielded]
-          [possible values: self-wakes, lost-waker, never-yielded]
+          * `large-future` -- Warnings when the future driving a task
+          occupies a large amount of stack space.
+
+          [default: self-wakes lost-waker never-yielded
+          auto-boxed-future large-future]
+          [possible values: self-wakes, lost-waker, never-yielded,
+          auto-boxed-future, large-future]
 
   -A, --allow <ALLOW_WARNINGS>...
           Allow lint warnings.
@@ -72,9 +77,13 @@ Options:
 
           * `never-yielded` -- Warns when a task has never yielded.
 
+          * `large-future` -- Warnings when the future driving a task
+          occupies a large amount of stack space.
+
           If this is set to `all`, all warnings are allowed.
 
-          [possible values: all, self-wakes, lost-waker, never-yielded]
+          [possible values: all, self-wakes, lost-waker, never-yielded,
+          large-future]
 
       --log-dir <LOG_DIRECTORY>
           Path to a directory to write the console's internal logs to.