ckb_sentry_core/
futures.rs1use std::future::Future;
2use std::pin::Pin;
3use std::sync::Arc;
4use std::task::{Context, Poll};
5
6use crate::Hub;
7
8#[derive(Debug)]
16pub struct SentryFuture<F> {
17 hub: Arc<Hub>,
18 future: F,
19}
20
21impl<F> SentryFuture<F> {
22 pub fn new(hub: Arc<Hub>, future: F) -> Self {
24 Self { hub, future }
25 }
26}
27
28impl<F> Future for SentryFuture<F>
29where
30 F: Future,
31{
32 type Output = F::Output;
33
34 fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
35 let hub = self.hub.clone();
36 let future = unsafe { self.map_unchecked_mut(|s| &mut s.future) };
38 #[cfg(feature = "client")]
39 {
40 Hub::run(hub, || future.poll(cx))
41 }
42 #[cfg(not(feature = "client"))]
43 {
44 let _ = hub;
45 future.poll(cx)
46 }
47 }
48}
49
50pub trait SentryFutureExt: Sized {
52 fn bind_hub<H>(self, hub: H) -> SentryFuture<Self>
56 where
57 H: Into<Arc<Hub>>,
58 {
59 SentryFuture {
60 future: self,
61 hub: hub.into(),
62 }
63 }
64}
65
66impl<F> SentryFutureExt for F where F: Future {}
67
68#[cfg(all(test, feature = "test"))]
69mod tests {
70 use crate::test::with_captured_events;
71 use crate::{capture_message, configure_scope, Hub, Level, SentryFutureExt};
72 use tokio::runtime::Runtime;
73
74 #[test]
75 fn test_futures() {
76 let mut events = with_captured_events(|| {
77 let mut runtime = Runtime::new().unwrap();
78
79 runtime.block_on(async {
81 let task1 = async {
82 configure_scope(|scope| scope.set_transaction(Some("transaction1")));
83 capture_message("oh hai from 1", Level::Info);
84 }
85 .bind_hub(Hub::new_from_top(Hub::current()));
86 let task1 = tokio::task::spawn(task1);
87
88 let task2 = async {
89 configure_scope(|scope| scope.set_transaction(Some("transaction2")));
90 capture_message("oh hai from 2", Level::Info);
91 }
92 .bind_hub(Hub::new_from_top(Hub::current()));
93 let task2 = tokio::task::spawn(task2);
94
95 task1.await.unwrap();
96 task2.await.unwrap();
97 });
98
99 capture_message("oh hai from outside", Level::Info);
100 });
101
102 events.sort_by(|a, b| a.transaction.cmp(&b.transaction));
103 assert_eq!(events.len(), 3);
104 assert_eq!(events[1].transaction, Some("transaction1".into()));
105 assert_eq!(events[2].transaction, Some("transaction2".into()));
106 }
107}