From 734e3653570d933130a0d464f347b02f447a5b67 Mon Sep 17 00:00:00 2001
From: andylokandy <andylokandy@hotmail.com>
Date: Fri, 21 Jun 2024 22:52:36 +0800
Subject: [PATCH 1/2] fix: don't attempt to reuse in foreground thread

Signed-off-by: andylokandy <andylokandy@hotmail.com>
---
 minitrace/benches/object_pool.rs            |  1 +
 minitrace/src/collector/global_collector.rs |  3 +++
 minitrace/src/util/object_pool.rs           | 22 +++++++++++++++++++--
 3 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/minitrace/benches/object_pool.rs b/minitrace/benches/object_pool.rs
index aa694bd1..c21341ec 100644
--- a/minitrace/benches/object_pool.rs
+++ b/minitrace/benches/object_pool.rs
@@ -12,6 +12,7 @@ fn bench_alloc_vec(c: &mut Criterion) {
     for cap in &[1, 10, 100, 1000, 10000, 100000] {
         let vec_pool: Pool<Vec<usize>> = Pool::new(Vec::new, Vec::clear);
         let mut puller = vec_pool.puller(512);
+        minitrace::util::object_pool::enable_reuse_in_current_thread();
         bgroup.bench_function(format!("object-pool/{}", cap), |b| {
             b.iter_batched(
                 || (),
diff --git a/minitrace/src/collector/global_collector.rs b/minitrace/src/collector/global_collector.rs
index f2491ab1..4583ff36 100644
--- a/minitrace/src/collector/global_collector.rs
+++ b/minitrace/src/collector/global_collector.rs
@@ -27,6 +27,7 @@ use crate::collector::SpanSet;
 use crate::collector::TraceId;
 use crate::local::local_collector::LocalSpansInner;
 use crate::local::raw_span::RawSpan;
+use crate::util::object_pool;
 use crate::util::spsc::Receiver;
 use crate::util::spsc::Sender;
 use crate::util::spsc::{self};
@@ -252,6 +253,8 @@ impl GlobalCollector {
     }
 
     fn handle_commands(&mut self, flush: bool) {
+        object_pool::enable_reuse_in_current_thread();
+
         debug_assert!(self.start_collects.is_empty());
         debug_assert!(self.drop_collects.is_empty());
         debug_assert!(self.commit_collects.is_empty());
diff --git a/minitrace/src/util/object_pool.rs b/minitrace/src/util/object_pool.rs
index 8cf58750..b52c3319 100644
--- a/minitrace/src/util/object_pool.rs
+++ b/minitrace/src/util/object_pool.rs
@@ -1,12 +1,28 @@
 // Copyright 2022 TiKV Project Authors. Licensed under Apache-2.0.
 
+use std::cell::Cell;
 use std::mem::ManuallyDrop;
 use std::ops::Deref;
 use std::ops::DerefMut;
 
 use parking_lot::Mutex;
 
+thread_local! {
+    static REUSABLE: Cell<bool> = Cell::new(false);
+}
+
+pub fn enable_reuse_in_current_thread() {
+    REUSABLE.with(|r| r.set(true));
+}
+
+fn is_reusable() -> bool {
+    REUSABLE.with(|r| r.get())
+}
+
 pub struct Pool<T> {
+    // The objects in the pool ready to be reused.
+    // The mutex should only be visited in the global collector, which is guaranteed by `is_reusable`,
+    // so it should not have synchronization overhead.
     objects: Mutex<Vec<T>>,
     init: fn() -> T,
     reset: fn(&mut T),
@@ -46,8 +62,10 @@ impl<T> Pool<T> {
 
     #[inline]
     pub fn recycle(&self, mut obj: T) {
-        (self.reset)(&mut obj);
-        self.objects.lock().push(obj)
+        if is_reusable() {
+            (self.reset)(&mut obj);
+            self.objects.lock().push(obj)
+        }
     }
 }
 

From a488419cb03299c97dbfaf44850c8ecd21129ea0 Mon Sep 17 00:00:00 2001
From: andylokandy <andylokandy@hotmail.com>
Date: Fri, 21 Jun 2024 23:01:22 +0800
Subject: [PATCH 2/2] fix

Signed-off-by: andylokandy <andylokandy@hotmail.com>
---
 minitrace/src/util/object_pool.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/minitrace/src/util/object_pool.rs b/minitrace/src/util/object_pool.rs
index b52c3319..06d54d14 100644
--- a/minitrace/src/util/object_pool.rs
+++ b/minitrace/src/util/object_pool.rs
@@ -8,7 +8,7 @@ use std::ops::DerefMut;
 use parking_lot::Mutex;
 
 thread_local! {
-    static REUSABLE: Cell<bool> = Cell::new(false);
+    static REUSABLE: Cell<bool> = const { Cell::new(false) };
 }
 
 pub fn enable_reuse_in_current_thread() {