chore: reduce one Option match for elements in object pool (#159)

* chore: reduce one Option match for elements in object pool

Signed-off-by: Andy Lok <[email protected]>

* add LocalParentGuard back again

Signed-off-by: Andy Lok <[email protected]>

* improve doc

* fix test

Signed-off-by: Andy Lok <[email protected]>

---------

Signed-off-by: Andy Lok <[email protected]>

Author: andylokandy

minitrace/Cargo.toml

@@ -16,7 +16,6 @@ keywords = ["tracing", "span", "datadog", "jaeger", "opentelemetry"]
 enable = []
 
 [dependencies]
-defer = "0.1"
 futures = "0.3"
 minitrace-macro = { version = "0.5.0", path = "../minitrace-macro" }
 minstant = "0.1"

minitrace/src/collector/console_reporter.rs

@@ -3,7 +3,7 @@
 use super::global_collector::Reporter;
 use super::SpanRecord;
 
-/// A console reporter that logs span records to the stderr.
+/// A console reporter that prints span records to the stderr.
 pub struct ConsoleReporter;
 
 impl Reporter for ConsoleReporter {

minitrace/src/collector/global_collector.rs

@@ -438,7 +438,7 @@ fn amend_local_span(
             continue;
         }
 
-        let end_unix_time_ns = if span.end_instant == span.begin_instant {
+        let end_time_unix_ns = if span.end_instant == span.begin_instant {
             local_spans.end_time.as_unix_nanos(anchor)
         } else {
             span.end_instant.as_unix_nanos(anchor)
@@ -448,7 +448,7 @@ fn amend_local_span(
             span_id: span.id,
             parent_id,
             begin_time_unix_ns,
-            duration_ns: end_unix_time_ns.saturating_sub(begin_time_unix_ns),
+            duration_ns: end_time_unix_ns.saturating_sub(begin_time_unix_ns),
             name: span.name,
             properties: span.properties.clone(),
             events: vec![],
@@ -476,13 +476,13 @@ fn amend_span(
         return;
     }
 
-    let end_unix_time_ns = raw_span.end_instant.as_unix_nanos(anchor);
+    let end_time_unix_ns = raw_span.end_instant.as_unix_nanos(anchor);
     spans.push(SpanRecord {
         trace_id,
         span_id: raw_span.id,
         parent_id,
         begin_time_unix_ns,
-        duration_ns: end_unix_time_ns.saturating_sub(begin_time_unix_ns),
+        duration_ns: end_time_unix_ns.saturating_sub(begin_time_unix_ns),
         name: raw_span.name,
         properties: raw_span.properties.clone(),
         events: vec![],

minitrace/src/lib.rs

@@ -113,13 +113,15 @@
 //!
 //! {
 //!     let root_span = Span::root("root", SpanContext::random());
+//!
 //!     {
 //!         let child_span = Span::enter_with_parent("a child span", &root_span);
 //!
 //!         // ...
 //!
 //!         // child_span ends here.
 //!     }
+//!
 //!     // root_span ends here.
 //! }
 //!
@@ -145,6 +147,7 @@
 //!
 //! {
 //!     let root = Span::root("root", SpanContext::random());
+//!
 //!     {
 //!         let _guard = root.set_local_parent();
 //!

minitrace/src/local/local_collector.rs

@@ -12,7 +12,7 @@ use crate::local::local_span_stack::LOCAL_SPAN_STACK;
 use crate::util::CollectToken;
 use crate::util::RawSpans;
 
-/// Collector to collect [`LocalSpan`].
+/// A collector to collect [`LocalSpan`].
 ///
 /// `LocalCollector` allows to collect `LocalSpan` manually without a local parent. The collected `LocalSpan` can later be
 /// attached to a parent.

minitrace/src/local/local_span_line.rs

@@ -193,11 +193,11 @@ span []
         span_line2.add_properties(&span, || [("k1".into(), "v1".into())]);
         span_line1.finish_span(span);
 
-        let raw_spans = span_line1.collect(1).unwrap().0.into_inner().1;
+        let raw_spans = span_line1.collect(1).unwrap().0.into_inner();
         assert_eq!(raw_spans.len(), 1);
         assert_eq!(raw_spans[0].properties.len(), 0);
 
-        let raw_spans = span_line2.collect(2).unwrap().0.into_inner().1;
+        let raw_spans = span_line2.collect(2).unwrap().0.into_inner();
         assert!(raw_spans.is_empty());
     }
 
@@ -228,11 +228,11 @@ span []
         let (spans, collect_token) = span_line1.collect(1).unwrap();
         let collect_token = collect_token.unwrap();
         assert_eq!(collect_token.as_slice(), &[item]);
-        assert_eq!(spans.into_inner().1.len(), 1);
+        assert_eq!(spans.into_inner().len(), 1);
 
         let (spans, collect_token) = span_line2.collect(2).unwrap();
         assert!(collect_token.is_none());
-        assert!(spans.into_inner().1.is_empty());
+        assert!(spans.into_inner().is_empty());
     }
 
     #[test]

minitrace/src/local/mod.rs

@@ -12,3 +12,4 @@ pub(crate) mod span_queue;
 pub use self::local_collector::LocalCollector;
 pub use self::local_collector::LocalSpans;
 pub use self::local_span::LocalSpan;
+pub use crate::span::LocalParentGuard;

minitrace/src/span.rs

@@ -226,10 +226,10 @@ impl Span {
     ///
     /// [`LocalSpan`]: crate::local::LocalSpan
     /// [`LocalSpan::enter_with_local_parent()`]: crate::local::LocalSpan::enter_with_local_parent
-    pub fn set_local_parent(&self) -> Option<impl Drop> {
+    pub fn set_local_parent(&self) -> LocalParentGuard {
         #[cfg(not(feature = "enable"))]
         {
-            None::<Span>
+            LocalParentGuard::noop()
         }
 
         #[cfg(feature = "enable")]
@@ -388,10 +388,11 @@ impl Span {
     pub(crate) fn attach_into_stack(
         &self,
         stack: &Rc<RefCell<LocalSpanStack>>,
-    ) -> Option<impl Drop> {
+    ) -> LocalParentGuard {
         self.inner
             .as_ref()
             .map(move |inner| inner.capture_local_spans(stack.clone()))
+            .unwrap_or_else(LocalParentGuard::noop)
     }
 }
 
@@ -409,19 +410,11 @@ impl SpanInner {
     }
 
     #[inline]
-    fn capture_local_spans(&self, stack: Rc<RefCell<LocalSpanStack>>) -> impl Drop {
+    fn capture_local_spans(&self, stack: Rc<RefCell<LocalSpanStack>>) -> LocalParentGuard {
         let token = self.issue_collect_token().collect();
         let collector = LocalCollector::new(Some(token), stack);
-        let collect = self.collect.clone();
-        defer::defer(move || {
-            let (spans, token) = collector.collect_spans_and_token();
-            debug_assert!(token.is_some());
-            let token = token.unwrap_or_else(|| [].iter().collect());
 
-            if !spans.spans.is_empty() {
-                collect.submit_spans(SpanSet::LocalSpansInner(spans), token);
-            }
-        })
+        LocalParentGuard::new(collector, self.collect.clone())
     }
 
     #[inline]
@@ -473,6 +466,51 @@ impl Drop for Span {
     }
 }
 
+/// A guard created by [`Span::set_local_parent()`].
+#[derive(Default)]
+pub struct LocalParentGuard {
+    #[cfg(feature = "enable")]
+    inner: Option<LocalParentGuardInner>,
+}
+
+struct LocalParentGuardInner {
+    collector: LocalCollector,
+    collect: GlobalCollect,
+}
+
+impl LocalParentGuard {
+    pub(crate) fn noop() -> LocalParentGuard {
+        LocalParentGuard {
+            #[cfg(feature = "enable")]
+            inner: None,
+        }
+    }
+
+    pub(crate) fn new(collector: LocalCollector, collect: GlobalCollect) -> LocalParentGuard {
+        LocalParentGuard {
+            #[cfg(feature = "enable")]
+            inner: Some(LocalParentGuardInner { collector, collect }),
+        }
+    }
+}
+
+impl Drop for LocalParentGuard {
+    fn drop(&mut self) {
+        #[cfg(feature = "enable")]
+        if let Some(inner) = self.inner.take() {
+            let (spans, token) = inner.collector.collect_spans_and_token();
+            debug_assert!(token.is_some());
+            if let Some(token) = token {
+                if !spans.spans.is_empty() {
+                    inner
+                        .collect
+                        .submit_spans(SpanSet::LocalSpansInner(spans), token);
+                }
+            }
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::sync::atomic::AtomicUsize;
@@ -495,7 +533,7 @@ mod tests {
     fn noop_basic() {
         let span = Span::noop();
         let stack = Rc::new(RefCell::new(LocalSpanStack::with_capacity(16)));
-        assert!(span.attach_into_stack(&stack).is_none());
+        assert!(span.attach_into_stack(&stack).inner.is_none());
         assert!(stack.borrow_mut().enter_span("span1").is_none());
     }
 
@@ -809,11 +847,11 @@ parent5 []
             {
                 let parent_ctx = SpanContext::random();
                 let root = Span::root("root", parent_ctx, collect.clone());
-                let _g = root.attach_into_stack(&stack).unwrap();
+                let _g = root.attach_into_stack(&stack);
                 let child =
                     Span::enter_with_stack("child", &mut stack.borrow_mut(), collect.clone());
                 {
-                    let _g = child.attach_into_stack(&stack).unwrap();
+                    let _g = child.attach_into_stack(&stack);
                     let _s = Span::enter_with_stack("grandchild", &mut stack.borrow_mut(), collect);
                 }
                 let _s = LocalSpan::enter_with_stack("local", stack);

minitrace/src/util/object_pool.rs

@@ -1,5 +1,6 @@
 // Copyright 2022 TiKV Project Authors. Licensed under Apache-2.0.
 
+use std::mem::ManuallyDrop;
 use std::ops::Deref;
 use std::ops::DerefMut;
 
@@ -21,16 +22,6 @@ impl<T> Pool<T> {
         }
     }
 
-    #[inline]
-    #[allow(dead_code)]
-    pub fn pull(&self) -> Reusable<T> {
-        self.objects
-            .lock()
-            .pop()
-            .map(|obj| Reusable::new(self, obj))
-            .unwrap_or_else(|| Reusable::new(self, (self.init)()))
-    }
-
     #[inline]
     fn batch_pull<'a>(&'a self, n: usize, buffer: &mut Vec<Reusable<'a, T>>) {
         let mut objects = self.objects.lock();
@@ -78,26 +69,33 @@ impl<'a, T> Puller<'a, T> {
 
 pub struct Reusable<'a, T> {
     pool: &'a Pool<T>,
-    obj: Option<T>,
+    obj: ManuallyDrop<T>,
 }
 
 impl<'a, T> Reusable<'a, T> {
     #[inline]
-    pub fn new(pool: &'a Pool<T>, t: T) -> Self {
-        Self { pool, obj: Some(t) }
+    pub fn new(pool: &'a Pool<T>, obj: T) -> Self {
+        Self {
+            pool,
+            obj: ManuallyDrop::new(obj),
+        }
     }
 
     #[inline]
-    pub fn into_inner(mut self) -> (&'a Pool<T>, T) {
-        (self.pool, self.obj.take().unwrap())
+    pub fn into_inner(mut self) -> T {
+        unsafe {
+            let obj = ManuallyDrop::take(&mut self.obj);
+            std::mem::forget(self);
+            obj
+        }
     }
 }
 
 impl<'a, T> std::fmt::Debug for Reusable<'a, T>
 where T: std::fmt::Debug
 {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        self.obj.as_ref().unwrap().fmt(f)
+        self.obj.fmt(f)
     }
 }
 
@@ -116,22 +114,22 @@ impl<'a, T> Deref for Reusable<'a, T> {
 
     #[inline]
     fn deref(&self) -> &Self::Target {
-        self.obj.as_ref().unwrap()
+        &self.obj
     }
 }
 
 impl<'a, T> DerefMut for Reusable<'a, T> {
     #[inline]
     fn deref_mut(&mut self) -> &mut Self::Target {
-        self.obj.as_mut().unwrap()
+        &mut self.obj
     }
 }
 
 impl<'a, T> Drop for Reusable<'a, T> {
     #[inline]
     fn drop(&mut self) {
-        if let Some(obj) = self.obj.take() {
-            self.pool.recycle(obj);
+        unsafe {
+            self.pool.recycle(ManuallyDrop::take(&mut self.obj));
         }
     }
 }

minitrace/src/util/tree.rs

@@ -55,7 +55,7 @@ impl Tree {
     pub fn from_raw_spans(raw_spans: RawSpans) -> Vec<Tree> {
         let mut children = HashMap::new();
 
-        let spans = raw_spans.into_inner().1;
+        let spans = raw_spans.into_inner();
         children.insert(SpanId::default(), ("", vec![], vec![]));
         for span in &spans {
             children.insert(span.id, (span.name, vec![], span.properties.clone()));