feat(console): add large future lint

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. A new lint has been added for large futures. By default it will show a warning when a task's future is greater than or equal to 2048 bytes. This value was chosen as it is the threshold at which Tokio will box a future in debug mode. The lint has been added to the list of default lints.
tokio-rs · Oct 1, 2024 · c309d80 · c309d80
1 parent ca39033
commit c309d80
Show file tree

Hide file tree

Showing 9 changed files with 124 additions and 7 deletions.
diff --git a/README.md b/README.md
@@ -224,8 +224,12 @@ 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 large-future]
+          [possible values: self-wakes, lost-waker, never-yielded,
+          large-future]
 
   -A, --allow <ALLOW_WARNINGS>...
           Allow lint warnings.
@@ -243,9 +247,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,12 @@ async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
                     .spawn(spawn_blocking(5))
                     .unwrap();
             }
+            "huge" => {
+                tokio::task::Builder::new()
+                    .name("huge")
+                    .spawn(huge::<1000>())
+                    .unwrap();
+            }
             "help" | "-h" => {
                 eprintln!("{}", HELP);
                 return Ok(());
@@ -152,6 +158,21 @@ async fn spawn_blocking(seconds: u64) {
     }
 }
 
+#[tracing::instrument]
+async fn huge<const N: usize>() {
+    let mut numbers = [0_usize; N];
+
+    loop {
+        for idx in 0..N {
+            numbers[idx] = idx;
+            tokio::time::sleep(Duration::from_millis(100)).await;
+            (0..=idx).for_each(|jdx| {
+                assert_eq!(numbers[jdx], jdx);
+            });
+        }
+    }
+}
+
 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,7 @@ warnings = [
     'self-wakes',
     'lost-waker',
     'never-yielded',
+    '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,7 @@ pub(crate) enum KnownWarnings {
     SelfWakes,
     LostWaker,
     NeverYielded,
+    LargeFuture,
 }
 
 impl FromStr for KnownWarnings {
@@ -153,6 +160,7 @@ impl FromStr for KnownWarnings {
             "self-wakes" => Ok(KnownWarnings::SelfWakes),
             "lost-waker" => Ok(KnownWarnings::LostWaker),
             "never-yielded" => Ok(KnownWarnings::NeverYielded),
+            "large-future" => Ok(KnownWarnings::LargeFuture),
             _ => Err(format!("unknown warning: {}", s)),
         }
     }
@@ -164,6 +172,7 @@ 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::LargeFuture => warnings::Linter::new(warnings::LargeFuture::default()),
         }
     }
 }
@@ -174,6 +183,7 @@ impl fmt::Display for KnownWarnings {
             KnownWarnings::SelfWakes => write!(f, "self-wakes"),
             KnownWarnings::LostWaker => write!(f, "lost-waker"),
             KnownWarnings::NeverYielded => write!(f, "never-yielded"),
+            KnownWarnings::LargeFuture => write!(f, "large-future"),
         }
     }
 }
@@ -184,6 +194,7 @@ impl KnownWarnings {
             KnownWarnings::SelfWakes,
             KnownWarnings::LostWaker,
             KnownWarnings::NeverYielded,
+            KnownWarnings::LargeFuture,
         ]
     }
 }

diff --git a/tokio-console/src/state/mod.rs b/tokio-console/src/state/mod.rs
@@ -277,6 +277,7 @@ 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";
 
     /// 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,8 @@ 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>,
 }
 
 #[derive(Debug)]
@@ -184,6 +186,7 @@ impl TasksState {
                 let mut name = None;
                 let mut task_id = None;
                 let mut kind = strings.string(String::new());
+                let mut size_bytes = None;
                 let target_field = Field::new(
                     strings.string_ref("target"),
                     FieldValue::Str(meta.target.to_string()),
@@ -210,6 +213,14 @@ 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)
+                            }
                             _ => Some(field),
                         }
                     })
@@ -245,6 +256,7 @@ impl TasksState {
                     warnings: Vec::new(),
                     location,
                     kind,
+                    size_bytes,
                 };
                 if let TaskLintResult::RequiresRecheck = task.lint(linters) {
                     next_pending_lint.insert(task.id);
@@ -506,6 +518,10 @@ impl Task {
     pub(crate) fn location(&self) -> &str {
         &self.location
     }
+
+    pub(crate) fn size_bytes(&self) -> Option<usize> {
+        self.size_bytes
+    }
 }
 
 enum TaskLintResult {

diff --git a/tokio-console/src/warnings.rs b/tokio-console/src/warnings.rs
@@ -258,3 +258,53 @@ impl Warn<Task> for NeverYielded {
         )
     }
 }
+
+/// Warning for if a task's driving future if very large
+#[derive(Clone, Debug)]
+pub(crate) struct LargeFuture {
+    min_size: usize,
+    description: String,
+}
+impl LargeFuture {
+    pub(crate) const DEFAULT_SIZE_BYTES: usize = 2048;
+    pub(crate) fn new(min_size: usize) -> Self {
+        Self {
+            min_size,
+            description: format!("tasks are very large (threshold {} bytes)", min_size),
+        }
+    }
+}
+
+impl Default for LargeFuture {
+    fn default() -> Self {
+        Self::new(Self::DEFAULT_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,12 @@ 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 large-future]
+          [possible values: self-wakes, lost-waker, never-yielded,
+          large-future]
 
   -A, --allow <ALLOW_WARNINGS>...
           Allow lint warnings.
@@ -72,9 +76,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.
Original file line number	Diff line number	Diff line change
Expand Up		@@ -62,3 +62,4 @@ hyper-util = { version = "0.1.6", features = ["tokio"] }

		[dev-dependencies]
		trycmd = "0.15.4"