subscriber: Add with_buf_capacity_limit to fmt Layer

Shrink the thread-local formatting buffer after each event if its capacity
exceeds the configured limit.

Prevents outlier large log lines from permanently inflating TLS buffers
across worker threads.

Signed-off-by: Jon Doron <jond@wiz.io>

Author: arilou

tracing-subscriber/src/fmt/fmt_layer.rs

@@ -73,6 +73,7 @@ pub struct Layer<
     fmt_span: format::FmtSpanConfig,
     is_ansi: bool,
     log_internal_errors: bool,
+    max_buf_capacity: Option<usize>,
     _inner: PhantomData<fn(S)>,
 }
 
@@ -123,6 +124,7 @@ where
             make_writer: self.make_writer,
             is_ansi: self.is_ansi,
             log_internal_errors: self.log_internal_errors,
+            max_buf_capacity: self.max_buf_capacity,
             _inner: self._inner,
         }
     }
@@ -153,6 +155,7 @@ where
             make_writer: self.make_writer,
             is_ansi: self.is_ansi,
             log_internal_errors: self.log_internal_errors,
+            max_buf_capacity: self.max_buf_capacity,
             _inner: self._inner,
         }
     }
@@ -186,6 +189,7 @@ impl<S, N, E, W> Layer<S, N, E, W> {
             fmt_span: self.fmt_span,
             is_ansi: self.is_ansi,
             log_internal_errors: self.log_internal_errors,
+            max_buf_capacity: self.max_buf_capacity,
             make_writer,
             _inner: self._inner,
         }
@@ -291,6 +295,7 @@ impl<S, N, E, W> Layer<S, N, E, W> {
             fmt_span: self.fmt_span,
             is_ansi: self.is_ansi,
             log_internal_errors: self.log_internal_errors,
+            max_buf_capacity: self.max_buf_capacity,
             make_writer: TestWriter::default(),
             _inner: self._inner,
         }
@@ -358,6 +363,17 @@ impl<S, N, E, W> Layer<S, N, E, W> {
         }
     }
 
+    /// Sets the maximum capacity (in bytes) retained by the thread-local
+    /// formatting buffer between events. After formatting, if the buffer
+    /// capacity exceeds `max_capacity` it is shrunk back. By default there
+    /// is no limit and the buffer retains whatever capacity it grew to.
+    pub fn with_buf_capacity_limit(self, max_capacity: usize) -> Self {
+        Self {
+            max_buf_capacity: Some(max_capacity),
+            ..self
+        }
+    }
+
     /// Updates the [`MakeWriter`] by applying a function to the existing [`MakeWriter`].
     ///
     /// This sets the [`MakeWriter`] that the layer being built will use to write events.
@@ -387,6 +403,7 @@ impl<S, N, E, W> Layer<S, N, E, W> {
             fmt_span: self.fmt_span,
             is_ansi: self.is_ansi,
             log_internal_errors: self.log_internal_errors,
+            max_buf_capacity: self.max_buf_capacity,
             make_writer: f(self.make_writer),
             _inner: self._inner,
         }
@@ -419,6 +436,7 @@ where
             make_writer: self.make_writer,
             is_ansi: self.is_ansi,
             log_internal_errors: self.log_internal_errors,
+            max_buf_capacity: self.max_buf_capacity,
             _inner: self._inner,
         }
     }
@@ -432,6 +450,7 @@ where
             make_writer: self.make_writer,
             is_ansi: self.is_ansi,
             log_internal_errors: self.log_internal_errors,
+            max_buf_capacity: self.max_buf_capacity,
             _inner: self._inner,
         }
     }
@@ -561,6 +580,7 @@ where
             make_writer: self.make_writer,
             is_ansi: self.is_ansi,
             log_internal_errors: self.log_internal_errors,
+            max_buf_capacity: self.max_buf_capacity,
             _inner: self._inner,
         }
     }
@@ -576,6 +596,7 @@ where
             make_writer: self.make_writer,
             is_ansi: self.is_ansi,
             log_internal_errors: self.log_internal_errors,
+            max_buf_capacity: self.max_buf_capacity,
             _inner: self._inner,
         }
     }
@@ -607,6 +628,7 @@ where
             // always disable ANSI escapes in JSON mode!
             is_ansi: false,
             log_internal_errors: self.log_internal_errors,
+            max_buf_capacity: self.max_buf_capacity,
             _inner: self._inner,
         }
     }
@@ -674,6 +696,7 @@ impl<S, N, E, W> Layer<S, N, E, W> {
             make_writer: self.make_writer,
             is_ansi: self.is_ansi,
             log_internal_errors: self.log_internal_errors,
+            max_buf_capacity: self.max_buf_capacity,
             _inner: self._inner,
         }
     }
@@ -705,6 +728,7 @@ impl<S, N, E, W> Layer<S, N, E, W> {
             make_writer: self.make_writer,
             is_ansi: self.is_ansi,
             log_internal_errors: self.log_internal_errors,
+            max_buf_capacity: self.max_buf_capacity,
             _inner: self._inner,
         }
     }
@@ -723,6 +747,7 @@ impl<S> Default for Layer<S> {
             make_writer: io::stdout,
             is_ansi: ansi,
             log_internal_errors: false,
+            max_buf_capacity: None,
             _inner: PhantomData,
         }
     }
@@ -1018,6 +1043,9 @@ where
             }
 
             buf.clear();
+            if let Some(max_cap) = self.max_buf_capacity {
+                buf.shrink_to(max_cap);
+            }
         });
     }
 
@@ -1666,4 +1694,40 @@ mod test {
             actual.as_str()
         );
     }
+
+    #[test]
+    fn buf_capacity_limit_does_not_break_formatting() {
+        // Verify that setting a buf capacity limit still produces correct output
+        // for both small and large events.
+        let make_writer = MockMakeWriter::default();
+        let subscriber = crate::fmt::Subscriber::builder()
+            .with_writer(make_writer.clone())
+            .with_level(false)
+            .with_ansi(false)
+            .with_timer(MockTime)
+            .with_buf_capacity_limit(256)
+            .finish();
+
+        with_default(subscriber, || {
+            // A large event that exceeds the 256-byte limit
+            let big = "x".repeat(1024);
+            tracing::info!(big_field = big.as_str(), "large event");
+
+            // A small event afterwards -- this verifies that the buffer was
+            // usable after being shrunk.
+            tracing::info!("small event");
+        });
+
+        let actual = make_writer.get_string();
+        assert!(
+            actual.contains("large event"),
+            "large event should have been written, got: {}",
+            actual
+        );
+        assert!(
+            actual.contains("small event"),
+            "small event should have been written after buf shrink, got: {}",
+            actual
+        );
+    }
 }

tracing-subscriber/src/fmt/mod.rs

@@ -658,6 +658,20 @@ where
         }
     }
 
+    /// Sets the maximum capacity (in bytes) retained by the thread-local
+    /// formatting buffer between events.
+    ///
+    /// See [`Layer::with_buf_capacity_limit`] for details.
+    pub fn with_buf_capacity_limit(
+        self,
+        max_capacity: usize,
+    ) -> SubscriberBuilder<N, format::Format<L, T>, F, W> {
+        SubscriberBuilder {
+            inner: self.inner.with_buf_capacity_limit(max_capacity),
+            ..self
+        }
+    }
+
     /// Sets whether or not an event's target is displayed.
     pub fn with_target(
         self,