feat(console): add large future lints (#587)

hds · hawkw · web-flow · commit ae1723091fcc · 2024-10-09T14:18:03.000+02:00
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. Co-authored-by: Eliza Weisman <eliza@elizas.website>
diff --git a/README.md b/README.md
@@ -224,8 +224,17 @@ 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]
+          * `auto-boxed-future` -- Warnings when the future driving a
+          task was automatically boxed by the runtime because it was
+          large.
+          
+          * `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 +252,17 @@ Options:
           
           * `never-yielded` -- Warns when a task has never yielded.
           
+          * `auto-boxed-future` -- Warnings when the future driving a
+          task was automatically boxed by the runtime because it was
+          large.
+          
+          * `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, auto-boxed-future]
 
       --log-dir <LOG_DIRECTORY>
           Path to a directory to write the console's internal logs to.
diff --git a/console-subscriber/README.md b/console-subscriber/README.md
@@ -130,6 +130,9 @@ Other instrumentation is added in later Tokio releases:
 
 * [Tokio v1.21.0] or later is required to use newest `task::Builder::spawn*` APIs.
 
+* [Tokio v1.41.0] (as yet unreleased) or later is required for task future sizes and the related
+  tokio-console lints `auto-boxed-future` and `large-future`.
+
 [Tokio v1.0.0]: https://github.com/tokio-rs/tokio/releases/tag/tokio-1.0.0
 [Tokio v1.7.0]: https://github.com/tokio-rs/tokio/releases/tag/tokio-1.7.0
 [Tokio v1.12.0]:https://github.com/tokio-rs/tokio/releases/tag/tokio-1.12.0
@@ -147,6 +150,7 @@ Other instrumentation is added in later Tokio releases:
 [init]: https://docs.rs/console-subscriber/latest/console_subscriber/fn.init.html
 [compile_time_filters]: https://docs.rs/tracing/latest/tracing/level_filters/index.html#compile-time-filters
 [Tokio v1.21.0]: https://github.com/tokio-rs/tokio/releases/tag/tokio-1.21.0
+[Tokio v1.41.0]: https://github.com/tokio-rs/tokio/releases/tag/tokio-1.41.0
 
 ### Adding the Console Subscriber
 
diff --git a/console-subscriber/examples/app.rs b/console-subscriber/examples/app.rs
@@ -12,6 +12,8 @@ OPTIONS:
     burn        Includes a (misbehaving) task that spins CPU with self-wakes
     coma        Includes a (misbehaving) task that forgets to register a waker
     noyield     Includes a (misbehaving) task that spawns tasks that never yield
+    blocking    Includes a blocking task that  (not misbehaving)
+    large       Includes tasks that are driven by futures that are larger than recommended
 "#;
 
 #[tokio::main]
@@ -51,6 +53,19 @@ 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();
+                large_blocking::<20_000>();
+            }
             "help" | "-h" => {
                 eprintln!("{}", HELP);
                 return Ok(());
@@ -152,6 +167,42 @@ 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);
+            });
+        }
+    }
+}
+
+fn large_blocking<const N: usize>() {
+    let numbers = [0_u8; N];
+
+    tokio::task::Builder::new()
+        .name("huge-blocking")
+        .spawn_blocking(move || {
+            let mut numbers = numbers;
+
+            loop {
+                for idx in 0..N {
+                    numbers[idx] = (idx % 256) as u8;
+                    std::thread::sleep(Duration::from_millis(100));
+                    (0..=idx).for_each(|jdx| {
+                        assert_eq!(numbers[jdx], (jdx % 256) as u8);
+                    });
+                }
+            }
+        })
+        .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,12 @@ pub struct Config {
     /// * `lost-waker` -- Warns when a task is dropped without being woken.
     ///
     /// * `never-yielded` -- Warns when a task has never yielded.
+    ///
+    /// * `auto-boxed-future` -- Warnings when the future driving a task was automatically boxed by
+    ///                          the runtime because it was large.
+    ///
+    /// * `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 +86,15 @@ pub struct Config {
     ///
     /// * `never-yielded` -- Warns when a task has never yielded.
     ///
+    /// * `auto-boxed-future` -- Warnings when the future driving a task was automatically boxed by
+    ///                          the runtime because it was large.
+    ///
+    /// * `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, auto-boxed-future]
     #[clap(long = "allow", short = 'A', num_args = 1..)]
     pub(crate) allow_warnings: Option<AllowedWarnings>,
 
@@ -143,6 +155,8 @@ pub(crate) enum KnownWarnings {
     SelfWakes,
     LostWaker,
     NeverYielded,
+    AutoBoxedFuture,
+    LargeFuture,
 }
 
 impl FromStr for KnownWarnings {
@@ -153,6 +167,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 +180,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 +192,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 +204,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's future was auto-boxed by the runtime when spawning, 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 {} bytes or larger", 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

Original file line number	Diff line number	Diff line change
`@@ -62,3 +62,4 @@ hyper-util = { version = "0.1.6", features = ["tokio"] }`
`62`	`62`
`63`	`63`	`[dev-dependencies]`
`64`	`64`	`trycmd = "0.15.4"`
	`65`	`+`
Original file line number	Diff line number	Diff line change
`@@ -4,6 +4,8 @@ warnings = [`
`4`	`4`	`'self-wakes',`
`5`	`5`	`'lost-waker',`
`6`	`6`	`'never-yielded',`
	`7`	`+ 'auto-boxed-future',`
	`8`	`+ 'large-future',`
`7`	`9`	`]`
`8`	`10`	`log_directory = '/tmp/tokio-console/logs'`
`9`	`11`	`retention = '6s'`