• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 use std::{
2     pin::Pin,
3     task::{Context, Poll},
4 };
5 
6 pub mod event;
7 pub mod field;
8 mod metadata;
9 pub mod span;
10 pub mod subscriber;
11 
12 #[derive(Debug, Eq, PartialEq)]
13 pub enum Parent {
14     ContextualRoot,
15     Contextual(String),
16     ExplicitRoot,
17     Explicit(String),
18 }
19 
20 pub struct PollN<T, E> {
21     and_return: Option<Result<T, E>>,
22     finish_at: usize,
23     polls: usize,
24 }
25 
26 impl Parent {
check_parent_name( &self, parent_name: Option<&str>, provided_parent: Option<tracing_core::span::Id>, ctx: impl std::fmt::Display, collector_name: &str, )27     pub fn check_parent_name(
28         &self,
29         parent_name: Option<&str>,
30         provided_parent: Option<tracing_core::span::Id>,
31         ctx: impl std::fmt::Display,
32         collector_name: &str,
33     ) {
34         match self {
35             Parent::ExplicitRoot => {
36                 assert!(
37                     provided_parent.is_none(),
38                     "[{}] expected {} to be an explicit root, but its parent was actually {:?} (name: {:?})",
39                     collector_name,
40                     ctx,
41                     provided_parent,
42                     parent_name,
43                 );
44             }
45             Parent::Explicit(expected_parent) => {
46                 assert_eq!(
47                     Some(expected_parent.as_ref()),
48                     parent_name,
49                     "[{}] expected {} to have explicit parent {}, but its parent was actually {:?} (name: {:?})",
50                     collector_name,
51                     ctx,
52                     expected_parent,
53                     provided_parent,
54                     parent_name,
55                 );
56             }
57             Parent::ContextualRoot => {
58                 assert!(
59                     provided_parent.is_none(),
60                     "[{}] expected {} to have a contextual parent, but its parent was actually {:?} (name: {:?})",
61                     collector_name,
62                     ctx,
63                     provided_parent,
64                     parent_name,
65                 );
66                 assert!(
67                     parent_name.is_none(),
68                     "[{}] expected {} to be contextual a root, but we were inside span {:?}",
69                     collector_name,
70                     ctx,
71                     parent_name,
72                 );
73             }
74             Parent::Contextual(expected_parent) => {
75                 assert!(provided_parent.is_none(),
76                     "[{}] expected {} to have a contextual parent\nbut its parent was actually {:?} (name: {:?})",
77                     collector_name,
78                     ctx,
79                     provided_parent,
80                     parent_name,
81                 );
82                 assert_eq!(
83                     Some(expected_parent.as_ref()),
84                     parent_name,
85                     "[{}] expected {} to have contextual parent {:?}, but got {:?}",
86                     collector_name,
87                     ctx,
88                     expected_parent,
89                     parent_name,
90                 );
91             }
92         }
93     }
94 }
95 
96 impl<T, E> std::future::Future for PollN<T, E>
97 where
98     T: Unpin,
99     E: Unpin,
100 {
101     type Output = Result<T, E>;
poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output>102     fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
103         let this = self.get_mut();
104 
105         this.polls += 1;
106         if this.polls == this.finish_at {
107             let value = this.and_return.take().expect("polled after ready");
108 
109             Poll::Ready(value)
110         } else {
111             cx.waker().wake_by_ref();
112             Poll::Pending
113         }
114     }
115 }
116 
117 impl PollN<(), ()> {
new_ok(finish_at: usize) -> Self118     pub fn new_ok(finish_at: usize) -> Self {
119         Self {
120             and_return: Some(Ok(())),
121             finish_at,
122             polls: 0,
123         }
124     }
125 
new_err(finish_at: usize) -> Self126     pub fn new_err(finish_at: usize) -> Self {
127         Self {
128             and_return: Some(Err(())),
129             finish_at,
130             polls: 0,
131         }
132     }
133 }
134 
135 #[cfg(feature = "tokio-test")]
block_on_future<F>(future: F) -> F::Output where F: std::future::Future,136 pub fn block_on_future<F>(future: F) -> F::Output
137 where
138     F: std::future::Future,
139 {
140     use tokio_test::task;
141 
142     let mut task = task::spawn(future);
143     loop {
144         if let Poll::Ready(v) = task.poll() {
145             break v;
146         }
147     }
148 }
149