From a6cf14b370275367dcecf1191e60f0bd260250d8 Mon Sep 17 00:00:00 2001
From: Franfran <51274081+iFrostizz@users.noreply.github.com>
Date: Fri, 5 Jul 2024 23:32:50 +0200
Subject: [PATCH] feat(subscriber): add `TOKIO_CONSOLE_BUFFER_CAPACITY`  env
 variable (#568)

When the recorded process spawns a lot of tasks and execute a lot of
operations, a red message can be seen on the top of the console saying
`dropped: 14319181 async_ops, 1601630 tasks, 634586 resources`. That
means that the buffer is not sufficiently big to record these and
someone in the tokio Discord advised to use the
`TOKIO_CONSOLE_BUFFER_CAPACITY` variable to tune that buffer. Turns out
it was not a thing, so this pr adds it.
---
 console-subscriber/src/builder.rs | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/console-subscriber/src/builder.rs b/console-subscriber/src/builder.rs
index 85406c8ec..25eadcf21 100644
--- a/console-subscriber/src/builder.rs
+++ b/console-subscriber/src/builder.rs
@@ -333,6 +333,10 @@ impl Builder {
             self.recording_path = Some(path.into());
         }
 
+        if let Some(capacity) = usize_from_env("TOKIO_CONSOLE_BUFFER_CAPACITY") {
+            self.event_buffer_capacity = capacity;
+        }
+
         self
     }
 
@@ -765,3 +769,14 @@ fn duration_from_env(var_name: &str) -> Option<Duration> {
         ),
     }
 }
+
+fn usize_from_env(var_name: &str) -> Option<usize> {
+    let var = std::env::var(var_name).ok()?;
+    match var.parse::<usize>() {
+        Ok(num) => Some(num),
+        Err(e) => panic!(
+            "failed to parse a usize from `{}={:?}`: {}",
+            var_name, var, e
+        ),
+    }
+}