fix: defensive reads for all thread locals (#161)

zhongzc · web-flow · commit 771354914195 · 2023-08-20T17:14:22.000+08:00
* fix: defensive reads for all thread locals

Signed-off-by: Zhenchi &lt;zhongzc_arch@outlook.com&gt;

* fix: clippy

Signed-off-by: Zhenchi &lt;zhongzc_arch@outlook.com&gt;

* address comment

Signed-off-by: Zhenchi &lt;zhongzc_arch@outlook.com&gt;

* address comment

Signed-off-by: Zhenchi &lt;zhongzc_arch@outlook.com&gt;

* address comment

Signed-off-by: Zhenchi &lt;zhongzc_arch@outlook.com&gt;

---------

Signed-off-by: Zhenchi &lt;zhongzc_arch@outlook.com&gt;
diff --git a/minitrace/src/collector/global_collector.rs b/minitrace/src/collector/global_collector.rs
@@ -50,11 +50,13 @@ fn register_receiver(rx: Receiver<CollectCommand>) {
 }
 
 fn send_command(cmd: CollectCommand) {
-    COMMAND_SENDER.with(|sender| sender.send(cmd).ok());
+    COMMAND_SENDER.try_with(|sender| sender.send(cmd).ok()).ok();
 }
 
 fn force_send_command(cmd: CollectCommand) {
-    COMMAND_SENDER.with(|sender| sender.force_send(cmd));
+    COMMAND_SENDER
+        .try_with(|sender| sender.force_send(cmd))
+        .ok();
 }
 
 /// Sets the reporter and its configuration for the current application.
@@ -432,7 +434,7 @@ fn amend_local_span(
                 timestamp_unix_ns: begin_time_unix_ns,
                 properties: span.properties.clone(),
             };
-            events.entry(parent_id).or_insert(vec![]).push(event);
+            events.entry(parent_id).or_default().push(event);
             continue;
         }
 
@@ -470,7 +472,7 @@ fn amend_span(
             timestamp_unix_ns: begin_time_unix_ns,
             properties: raw_span.properties.clone(),
         };
-        events.entry(parent_id).or_insert(vec![]).push(event);
+        events.entry(parent_id).or_default().push(event);
         return;
     }
 
diff --git a/minitrace/src/collector/id.rs b/minitrace/src/collector/id.rs
@@ -14,15 +14,17 @@ impl SpanId {
     #[inline]
     /// Create a non-zero `SpanId`
     pub(crate) fn next_id() -> SpanId {
-        LOCAL_ID_GENERATOR.with(|g| {
-            let (prefix, mut suffix) = g.get();
+        LOCAL_ID_GENERATOR
+            .try_with(|g| {
+                let (prefix, mut suffix) = g.get();
 
-            suffix = suffix.wrapping_add(1);
+                suffix = suffix.wrapping_add(1);
 
-            g.set((prefix, suffix));
+                g.set((prefix, suffix));
 
-            SpanId(((prefix as u64) << 32) | (suffix as u64))
-        })
+                SpanId(((prefix as u64) << 32) | (suffix as u64))
+            })
+            .unwrap_or_else(|_| SpanId(rand::random()))
     }
 }
 
diff --git a/minitrace/src/collector/mod.rs b/minitrace/src/collector/mod.rs
@@ -167,7 +167,8 @@ impl SpanContext {
 
         #[cfg(feature = "enable")]
         {
-            let stack = LOCAL_SPAN_STACK.with(Rc::clone);
+            let stack = LOCAL_SPAN_STACK.try_with(Rc::clone).ok()?;
+
             let mut stack = stack.borrow_mut();
             let collect_token = stack.current_collect_token()?[0];
 
diff --git a/minitrace/src/event.rs b/minitrace/src/event.rs
@@ -1,7 +1,6 @@
 // Copyright 2023 TiKV Project Authors. Licensed under Apache-2.0.
 
 use std::borrow::Cow;
-use std::rc::Rc;
 
 use crate::local::local_span_stack::LOCAL_SPAN_STACK;
 use crate::Span;
@@ -55,9 +54,9 @@ impl Event {
     {
         #[cfg(feature = "enable")]
         {
-            let stack = LOCAL_SPAN_STACK.with(Rc::clone);
-            let mut stack = stack.borrow_mut();
-            stack.add_event(name, properties);
+            LOCAL_SPAN_STACK
+                .try_with(|stack| stack.borrow_mut().add_event(name, properties))
+                .ok();
         }
     }
 }
diff --git a/minitrace/src/local/local_collector.rs b/minitrace/src/local/local_collector.rs
@@ -40,6 +40,7 @@ use crate::util::RawSpans;
 /// [`Span`]: crate::Span
 /// [`LocalSpan`]: crate::local::LocalSpan
 #[must_use]
+#[derive(Default)]
 pub struct LocalCollector {
     #[cfg(feature = "enable")]
     inner: Option<LocalCollectorInner>,
@@ -98,13 +99,14 @@ impl LocalCollector {
     pub fn start() -> Self {
         #[cfg(not(feature = "enable"))]
         {
-            LocalCollector {}
+            LocalCollector::default()
         }
 
         #[cfg(feature = "enable")]
         {
-            let stack = LOCAL_SPAN_STACK.with(Rc::clone);
-            Self::new(None, stack)
+            LOCAL_SPAN_STACK
+                .try_with(|stack| Self::new(None, stack.clone()))
+                .unwrap_or_default()
         }
     }
 
diff --git a/minitrace/src/local/local_span.rs b/minitrace/src/local/local_span.rs
@@ -12,6 +12,7 @@ use crate::local::local_span_stack::LOCAL_SPAN_STACK;
 ///
 /// [`Span`]: crate::Span
 #[must_use]
+#[derive(Default)]
 pub struct LocalSpan {
     #[cfg(feature = "enable")]
     inner: Option<LocalSpanInner>,
@@ -42,13 +43,14 @@ impl LocalSpan {
     pub fn enter_with_local_parent(name: &'static str) -> Self {
         #[cfg(not(feature = "enable"))]
         {
-            LocalSpan {}
+            LocalSpan::default()
         }
 
         #[cfg(feature = "enable")]
         {
-            let stack = LOCAL_SPAN_STACK.with(Rc::clone);
-            Self::enter_with_stack(name, stack)
+            LOCAL_SPAN_STACK
+                .try_with(|stack| Self::enter_with_stack(name, stack.clone()))
+                .unwrap_or_default()
         }
     }
 
diff --git a/minitrace/src/local/span_queue.rs b/minitrace/src/local/span_queue.rs
@@ -35,7 +35,7 @@ impl SpanQueue {
 
         let span = RawSpan::begin_with(
             SpanId::next_id(),
-            self.next_parent_id.unwrap_or(SpanId::default()),
+            self.next_parent_id.unwrap_or_default(),
             Instant::now(),
             name,
             false,
@@ -74,7 +74,7 @@ impl SpanQueue {
 
         let mut span = RawSpan::begin_with(
             SpanId::next_id(),
-            self.next_parent_id.unwrap_or(SpanId::default()),
+            self.next_parent_id.unwrap_or_default(),
             Instant::now(),
             name,
             true,
diff --git a/minitrace/src/span.rs b/minitrace/src/span.rs
@@ -23,6 +23,7 @@ use crate::util::CollectToken;
 
 /// A thread-safe span.
 #[must_use]
+#[derive(Default)]
 pub struct Span {
     #[cfg(feature = "enable")]
     pub(crate) inner: Option<SpanInner>,
@@ -196,9 +197,11 @@ impl Span {
         {
             #[cfg(not(test))]
             let collect = GlobalCollect;
-            LOCAL_SPAN_STACK.with(move |stack| {
-                Self::enter_with_stack(name, &mut (*stack).borrow_mut(), collect)
-            })
+            LOCAL_SPAN_STACK
+                .try_with(move |stack| {
+                    Self::enter_with_stack(name, &mut (*stack).borrow_mut(), collect)
+                })
+                .unwrap_or_default()
         }
     }
 
@@ -231,7 +234,9 @@ impl Span {
 
         #[cfg(feature = "enable")]
         {
-            LOCAL_SPAN_STACK.with(|s| self.attach_into_stack(s))
+            LOCAL_SPAN_STACK
+                .try_with(|s| self.attach_into_stack(s))
+                .unwrap_or_default()
         }
     }
 
diff --git a/minitrace/src/util/mod.rs b/minitrace/src/util/mod.rs
@@ -30,12 +30,16 @@ pub type CollectToken = Reusable<'static, Vec<CollectTokenItem>>;
 
 impl Default for RawSpans {
     fn default() -> Self {
-        RAW_SPANS_PULLER.with(|puller| puller.borrow_mut().pull())
+        RAW_SPANS_PULLER
+            .try_with(|puller| puller.borrow_mut().pull())
+            .unwrap_or_else(|_| Reusable::new(&*RAW_SPANS_POOL, vec![]))
     }
 }
 
 fn new_collect_token(items: impl IntoIterator<Item = CollectTokenItem>) -> CollectToken {
-    let mut token = COLLECT_TOKEN_ITEMS_PULLER.with(|puller| puller.borrow_mut().pull());
+    let mut token = COLLECT_TOKEN_ITEMS_PULLER
+        .try_with(|puller| puller.borrow_mut().pull())
+        .unwrap_or_else(|_| Reusable::new(&*COLLECT_TOKEN_ITEMS_POOL, vec![]));
     token.extend(items);
     token
 }

Original file line number	Diff line number	Diff line change
`@@ -50,11 +50,13 @@ fn register_receiver(rx: Receiver<CollectCommand>) {`
`50`	`50`	`}`
`51`	`51`
`52`	`52`	`fn send_command(cmd: CollectCommand) {`
`53`		`- COMMAND_SENDER.with(\|sender\| sender.send(cmd).ok());`
	`53`	`+ COMMAND_SENDER.try_with(\|sender\| sender.send(cmd).ok()).ok();`
`54`	`54`	`}`
`55`	`55`
`56`	`56`	`fn force_send_command(cmd: CollectCommand) {`
`57`		`- COMMAND_SENDER.with(\|sender\| sender.force_send(cmd));`
	`57`	`+ COMMAND_SENDER`
	`58`	`+ .try_with(\|sender\| sender.force_send(cmd))`
	`59`	`+ .ok();`
`58`	`60`	`}`
`59`	`61`
`60`	`62`	`/// Sets the reporter and its configuration for the current application.`
`@@ -432,7 +434,7 @@ fn amend_local_span(`
`432`	`434`	`timestamp_unix_ns: begin_time_unix_ns,`
`433`	`435`	`properties: span.properties.clone(),`
`434`	`436`	`};`
`435`		`- events.entry(parent_id).or_insert(vec![]).push(event);`
	`437`	`+ events.entry(parent_id).or_default().push(event);`
`436`	`438`	`continue;`
`437`	`439`	`}`
`438`	`440`
`@@ -470,7 +472,7 @@ fn amend_span(`
`470`	`472`	`timestamp_unix_ns: begin_time_unix_ns,`
`471`	`473`	`properties: raw_span.properties.clone(),`
`472`	`474`	`};`
`473`		`- events.entry(parent_id).or_insert(vec![]).push(event);`
	`475`	`+ events.entry(parent_id).or_default().push(event);`
`474`	`476`	`return;`
`475`	`477`	`}`
`476`	`478`
Original file line number	Diff line number	Diff line change
`@@ -167,7 +167,8 @@ impl SpanContext {`
`167`	`167`
`168`	`168`	`#[cfg(feature = "enable")]`
`169`	`169`	`{`
`170`		`- let stack = LOCAL_SPAN_STACK.with(Rc::clone);`
	`170`	`+ let stack = LOCAL_SPAN_STACK.try_with(Rc::clone).ok()?;`
	`171`	`+`
`171`	`172`	`let mut stack = stack.borrow_mut();`
`172`	`173`	`let collect_token = stack.current_collect_token()?[0];`
`173`	`174`
Original file line number	Diff line number	Diff line change
`@@ -1,7 +1,6 @@`
`1`	`1`	`// Copyright 2023 TiKV Project Authors. Licensed under Apache-2.0.`
`2`	`2`
`3`	`3`	`use std::borrow::Cow;`
`4`		`-use std::rc::Rc;`
`5`	`4`
`6`	`5`	`use crate::local::local_span_stack::LOCAL_SPAN_STACK;`
`7`	`6`	`use crate::Span;`
`@@ -55,9 +54,9 @@ impl Event {`
`55`	`54`	`{`
`56`	`55`	`#[cfg(feature = "enable")]`
`57`	`56`	`{`
`58`		`- let stack = LOCAL_SPAN_STACK.with(Rc::clone);`
`59`		`- let mut stack = stack.borrow_mut();`
`60`		`- stack.add_event(name, properties);`
	`57`	`+ LOCAL_SPAN_STACK`
	`58`	`+ .try_with(\|stack\| stack.borrow_mut().add_event(name, properties))`
	`59`	`+ .ok();`
`61`	`60`	`}`
`62`	`61`	`}`
`63`	`62`	`}`
Original file line number	Diff line number	Diff line change
`@@ -40,6 +40,7 @@ use crate::util::RawSpans;`
`40`	`40`	/// [`Span`]: crate::Span
`41`	`41`	/// [`LocalSpan`]: crate::local::LocalSpan
`42`	`42`	`#[must_use]`
	`43`	`+#[derive(Default)]`
`43`	`44`	`pub struct LocalCollector {`
`44`	`45`	`#[cfg(feature = "enable")]`
`45`	`46`	`inner: Option<LocalCollectorInner>,`
`@@ -98,13 +99,14 @@ impl LocalCollector {`
`98`	`99`	`pub fn start() -> Self {`
`99`	`100`	`#[cfg(not(feature = "enable"))]`
`100`	`101`	`{`
`101`		`- LocalCollector {}`
	`102`	`+ LocalCollector::default()`
`102`	`103`	`}`
`103`	`104`
`104`	`105`	`#[cfg(feature = "enable")]`
`105`	`106`	`{`
`106`		`- let stack = LOCAL_SPAN_STACK.with(Rc::clone);`
`107`		`- Self::new(None, stack)`
	`107`	`+ LOCAL_SPAN_STACK`
	`108`	`+ .try_with(\|stack\| Self::new(None, stack.clone()))`
	`109`	`+ .unwrap_or_default()`
`108`	`110`	`}`
`109`	`111`	`}`
`110`	`112`
Original file line number	Diff line number	Diff line change
`@@ -12,6 +12,7 @@ use crate::local::local_span_stack::LOCAL_SPAN_STACK;`
`12`	`12`	`///`
`13`	`13`	/// [`Span`]: crate::Span
`14`	`14`	`#[must_use]`
	`15`	`+#[derive(Default)]`
`15`	`16`	`pub struct LocalSpan {`
`16`	`17`	`#[cfg(feature = "enable")]`
`17`	`18`	`inner: Option<LocalSpanInner>,`
`@@ -42,13 +43,14 @@ impl LocalSpan {`
`42`	`43`	`pub fn enter_with_local_parent(name: &'static str) -> Self {`
`43`	`44`	`#[cfg(not(feature = "enable"))]`
`44`	`45`	`{`
`45`		`- LocalSpan {}`
	`46`	`+ LocalSpan::default()`
`46`	`47`	`}`
`47`	`48`
`48`	`49`	`#[cfg(feature = "enable")]`
`49`	`50`	`{`
`50`		`- let stack = LOCAL_SPAN_STACK.with(Rc::clone);`
`51`		`- Self::enter_with_stack(name, stack)`
	`51`	`+ LOCAL_SPAN_STACK`
	`52`	`+ .try_with(\|stack\| Self::enter_with_stack(name, stack.clone()))`
	`53`	`+ .unwrap_or_default()`
`52`	`54`	`}`
`53`	`55`	`}`
`54`	`56`
Original file line number	Diff line number	Diff line change
`@@ -23,6 +23,7 @@ use crate::util::CollectToken;`
`23`	`23`
`24`	`24`	`/// A thread-safe span.`
`25`	`25`	`#[must_use]`
	`26`	`+#[derive(Default)]`
`26`	`27`	`pub struct Span {`
`27`	`28`	`#[cfg(feature = "enable")]`
`28`	`29`	`pub(crate) inner: Option<SpanInner>,`
`@@ -196,9 +197,11 @@ impl Span {`
`196`	`197`	`{`
`197`	`198`	`#[cfg(not(test))]`
`198`	`199`	`let collect = GlobalCollect;`
`199`		`- LOCAL_SPAN_STACK.with(move \|stack\| {`
`200`		`- Self::enter_with_stack(name, &mut (*stack).borrow_mut(), collect)`
`201`		`- })`
	`200`	`+ LOCAL_SPAN_STACK`
	`201`	`+ .try_with(move \|stack\| {`
	`202`	`+ Self::enter_with_stack(name, &mut (*stack).borrow_mut(), collect)`
	`203`	`+ })`
	`204`	`+ .unwrap_or_default()`
`202`	`205`	`}`
`203`	`206`	`}`
`204`	`207`
`@@ -231,7 +234,9 @@ impl Span {`
`231`	`234`
`232`	`235`	`#[cfg(feature = "enable")]`
`233`	`236`	`{`
`234`		`- LOCAL_SPAN_STACK.with(\|s\| self.attach_into_stack(s))`
	`237`	`+ LOCAL_SPAN_STACK`
	`238`	`+ .try_with(\|s\| self.attach_into_stack(s))`
	`239`	`+ .unwrap_or_default()`
`235`	`240`	`}`
`236`	`241`	`}`
`237`	`242`
Original file line number	Diff line number	Diff line change
`@@ -30,12 +30,16 @@ pub type CollectToken = Reusable<'static, Vec<CollectTokenItem>>;`
`30`	`30`
`31`	`31`	`impl Default for RawSpans {`
`32`	`32`	`fn default() -> Self {`
`33`		`- RAW_SPANS_PULLER.with(\|puller\| puller.borrow_mut().pull())`
	`33`	`+ RAW_SPANS_PULLER`
	`34`	`+ .try_with(\|puller\| puller.borrow_mut().pull())`
	`35`	`+ .unwrap_or_else(\|_\| Reusable::new(&*RAW_SPANS_POOL, vec![]))`
`34`	`36`	`}`
`35`	`37`	`}`
`36`	`38`
`37`	`39`	`fn new_collect_token(items: impl IntoIterator<Item = CollectTokenItem>) -> CollectToken {`
`38`		`- let mut token = COLLECT_TOKEN_ITEMS_PULLER.with(\|puller\| puller.borrow_mut().pull());`
	`40`	`+ let mut token = COLLECT_TOKEN_ITEMS_PULLER`
	`41`	`+ .try_with(\|puller\| puller.borrow_mut().pull())`
	`42`	`+ .unwrap_or_else(\|_\| Reusable::new(&*COLLECT_TOKEN_ITEMS_POOL, vec![]));`
`39`	`43`	`token.extend(items);`
`40`	`44`	`token`
`41`	`45`	`}`