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> = 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..06d54d14 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 = const { 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 { + // 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>, init: fn() -> T, reset: fn(&mut T), @@ -46,8 +62,10 @@ impl Pool { #[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) + } } }