vendor/tracing/tests/support/subscriber.rs

   1 #![allow(missing_docs)]
   2 use super::{
   3     event::MockEvent,
   4     field as mock_field,
   5     span::{MockSpan, NewSpan},
   6     Parent,
   7 };
   8 use std::{
   9     collections::{HashMap, VecDeque},
  10     fmt,
  11     sync::{
  12         atomic::{AtomicUsize, Ordering},
  13         Arc, Mutex,
  14     },
  15 };
  16 use tracing::{
  17     level_filters::LevelFilter,
  18     span::{self, Attributes, Id},
  19     Event, Metadata, Subscriber,
  20 };
  21
  22 #[derive(Debug, Eq, PartialEq)]
  23 enum Expect {
  24     Event(MockEvent),
  25     Enter(MockSpan),
  26     Exit(MockSpan),
  27     CloneSpan(MockSpan),
  28     DropSpan(MockSpan),
  29     Visit(MockSpan, mock_field::Expect),
  30     NewSpan(NewSpan),
  31     Nothing,
  32 }
  33
  34 struct SpanState {
  35     name: &'static str,
  36     refs: usize,
  37 }
  38
  39 struct Running<F: Fn(&Metadata<'_>) -> bool> {
  40     spans: Mutex<HashMap<Id, SpanState>>,
  41     expected: Arc<Mutex<VecDeque<Expect>>>,
  42     current: Mutex<Vec<Id>>,
  43     ids: AtomicUsize,
  44     max_level: Option<LevelFilter>,
  45     filter: F,
  46 }
  47
  48 pub struct MockSubscriber<F: Fn(&Metadata<'_>) -> bool> {
  49     expected: VecDeque<Expect>,
  50     max_level: Option<LevelFilter>,
  51     filter: F,
  52 }
  53
  54 pub struct MockHandle(Arc<Mutex<VecDeque<Expect>>>);
  55
  56 pub fn mock() -> MockSubscriber<fn(&Metadata<'_>) -> bool> {
  57     MockSubscriber {
  58         expected: VecDeque::new(),
  59         filter: (|_: &Metadata<'_>| true) as for<'r, 's> fn(&'r Metadata<'s>) -> _,
  60         max_level: None,
  61     }
  62 }
  63
  64 impl<F> MockSubscriber<F>
  65 where
  66     F: Fn(&Metadata<'_>) -> bool + 'static,
  67 {
  68     pub fn enter(mut self, span: MockSpan) -> Self {
  69         self.expected.push_back(Expect::Enter(span));
  70         self
  71     }
  72
  73     pub fn event(mut self, event: MockEvent) -> Self {
  74         self.expected.push_back(Expect::Event(event));
  75         self
  76     }
  77
  78     pub fn exit(mut self, span: MockSpan) -> Self {
  79         self.expected.push_back(Expect::Exit(span));
  80         self
  81     }
  82
  83     pub fn clone_span(mut self, span: MockSpan) -> Self {
  84         self.expected.push_back(Expect::CloneSpan(span));
  85         self
  86     }
  87
  88     #[allow(deprecated)]
  89     pub fn drop_span(mut self, span: MockSpan) -> Self {
  90         self.expected.push_back(Expect::DropSpan(span));
  91         self
  92     }
  93
  94     pub fn done(mut self) -> Self {
  95         self.expected.push_back(Expect::Nothing);
  96         self
  97     }
  98
  99     pub fn record<I>(mut self, span: MockSpan, fields: I) -> Self
 100     where
 101         I: Into<mock_field::Expect>,
 102     {
 103         self.expected.push_back(Expect::Visit(span, fields.into()));
 104         self
 105     }
 106
 107     pub fn new_span<I>(mut self, new_span: I) -> Self
 108     where
 109         I: Into<NewSpan>,
 110     {
 111         self.expected.push_back(Expect::NewSpan(new_span.into()));
 112         self
 113     }
 114
 115     pub fn with_filter<G>(self, filter: G) -> MockSubscriber<G>
 116     where
 117         G: Fn(&Metadata<'_>) -> bool + 'static,
 118     {
 119         MockSubscriber {
 120             expected: self.expected,
 121             filter,
 122             max_level: self.max_level,
 123         }
 124     }
 125
 126     pub fn with_max_level_hint(self, hint: impl Into<LevelFilter>) -> Self {
 127         Self {
 128             max_level: Some(hint.into()),
 129             ..self
 130         }
 131     }
 132
 133     pub fn run(self) -> impl Subscriber {
 134         let (subscriber, _) = self.run_with_handle();
 135         subscriber
 136     }
 137
 138     pub fn run_with_handle(self) -> (impl Subscriber, MockHandle) {
 139         let expected = Arc::new(Mutex::new(self.expected));
 140         let handle = MockHandle(expected.clone());
 141         let subscriber = Running {
 142             spans: Mutex::new(HashMap::new()),
 143             expected,
 144             current: Mutex::new(Vec::new()),
 145             ids: AtomicUsize::new(1),
 146             filter: self.filter,
 147             max_level: self.max_level,
 148         };
 149         (subscriber, handle)
 150     }
 151 }
 152
 153 impl<F> Subscriber for Running<F>
 154 where
 155     F: Fn(&Metadata<'_>) -> bool + 'static,
 156 {
 157     fn enabled(&self, meta: &Metadata<'_>) -> bool {
 158         (self.filter)(meta)
 159     }
 160
 161     fn max_level_hint(&self) -> Option<LevelFilter> {
 162         self.max_level.clone()
 163     }
 164
 165     fn record(&self, id: &Id, values: &span::Record<'_>) {
 166         let spans = self.spans.lock().unwrap();
 167         let mut expected = self.expected.lock().unwrap();
 168         let span = spans
 169             .get(id)
 170             .unwrap_or_else(|| panic!("no span for ID {:?}", id));
 171         println!("record: {}; id={:?}; values={:?};", span.name, id, values);
 172         let was_expected = if let Some(Expect::Visit(_, _)) = expected.front() {
 173             true
 174         } else {
 175             false
 176         };
 177         if was_expected {
 178             if let Expect::Visit(expected_span, mut expected_values) = expected.pop_front().unwrap()
 179             {
 180                 if let Some(name) = expected_span.name() {
 181                     assert_eq!(name, span.name);
 182                 }
 183                 let mut checker = expected_values.checker(format!("span {}: ", span.name));
 184                 values.record(&mut checker);
 185                 checker.finish();
 186             }
 187         }
 188     }
 189
 190     fn event(&self, event: &Event<'_>) {
 191         let name = event.metadata().name();
 192         println!("event: {};", name);
 193         match self.expected.lock().unwrap().pop_front() {
 194             None => {}
 195             Some(Expect::Event(mut expected)) => {
 196                 let spans = self.spans.lock().unwrap();
 197                 expected.check(event);
 198                 match expected.parent {
 199                     Some(Parent::ExplicitRoot) => {
 200                         assert!(
 201                             event.is_root(),
 202                             "expected {:?} to be an explicit root event",
 203                             name
 204                         );
 205                     }
 206                     Some(Parent::Explicit(expected_parent)) => {
 207                         let actual_parent =
 208                             event.parent().and_then(|id| spans.get(id)).map(|s| s.name);
 209                         assert_eq!(
 210                             Some(expected_parent.as_ref()),
 211                             actual_parent,
 212                             "expected {:?} to have explicit parent {:?}",
 213                             name,
 214                             expected_parent,
 215                         );
 216                     }
 217                     Some(Parent::ContextualRoot) => {
 218                         assert!(
 219                             event.is_contextual(),
 220                             "expected {:?} to have a contextual parent",
 221                             name
 222                         );
 223                         assert!(
 224                             self.current.lock().unwrap().last().is_none(),
 225                             "expected {:?} to be a root, but we were inside a span",
 226                             name
 227                         );
 228                     }
 229                     Some(Parent::Contextual(expected_parent)) => {
 230                         assert!(
 231                             event.is_contextual(),
 232                             "expected {:?} to have a contextual parent",
 233                             name
 234                         );
 235                         let stack = self.current.lock().unwrap();
 236                         let actual_parent =
 237                             stack.last().and_then(|id| spans.get(id)).map(|s| s.name);
 238                         assert_eq!(
 239                             Some(expected_parent.as_ref()),
 240                             actual_parent,
 241                             "expected {:?} to have contextual parent {:?}",
 242                             name,
 243                             expected_parent,
 244                         );
 245                     }
 246                     None => {}
 247                 }
 248             }
 249             Some(ex) => ex.bad(format_args!("observed event {:?}", event)),
 250         }
 251     }
 252
 253     fn record_follows_from(&self, _span: &Id, _follows: &Id) {
 254         // TODO: it should be possible to expect spans to follow from other spans
 255     }
 256
 257     fn new_span(&self, span: &Attributes<'_>) -> Id {
 258         let meta = span.metadata();
 259         let id = self.ids.fetch_add(1, Ordering::SeqCst);
 260         let id = Id::from_u64(id as u64);
 261         println!(
 262             "new_span: name={:?}; target={:?}; id={:?};",
 263             meta.name(),
 264             meta.target(),
 265             id
 266         );
 267         let mut expected = self.expected.lock().unwrap();
 268         let was_expected = match expected.front() {
 269             Some(Expect::NewSpan(_)) => true,
 270             _ => false,
 271         };
 272         let mut spans = self.spans.lock().unwrap();
 273         if was_expected {
 274             if let Expect::NewSpan(mut expected) = expected.pop_front().unwrap() {
 275                 let name = meta.name();
 276                 expected
 277                     .span
 278                     .metadata
 279                     .check(meta, format_args!("span `{}`", name));
 280                 let mut checker = expected.fields.checker(name.to_string());
 281                 span.record(&mut checker);
 282                 checker.finish();
 283                 match expected.parent {
 284                     Some(Parent::ExplicitRoot) => {
 285                         assert!(
 286                             span.is_root(),
 287                             "expected {:?} to be an explicit root span",
 288                             name
 289                         );
 290                     }
 291                     Some(Parent::Explicit(expected_parent)) => {
 292                         let actual_parent =
 293                             span.parent().and_then(|id| spans.get(id)).map(|s| s.name);
 294                         assert_eq!(
 295                             Some(expected_parent.as_ref()),
 296                             actual_parent,
 297                             "expected {:?} to have explicit parent {:?}",
 298                             name,
 299                             expected_parent,
 300                         );
 301                     }
 302                     Some(Parent::ContextualRoot) => {
 303                         assert!(
 304                             span.is_contextual(),
 305                             "expected {:?} to have a contextual parent",
 306                             name
 307                         );
 308                         assert!(
 309                             self.current.lock().unwrap().last().is_none(),
 310                             "expected {:?} to be a root, but we were inside a span",
 311                             name
 312                         );
 313                     }
 314                     Some(Parent::Contextual(expected_parent)) => {
 315                         assert!(
 316                             span.is_contextual(),
 317                             "expected {:?} to have a contextual parent",
 318                             name
 319                         );
 320                         let stack = self.current.lock().unwrap();
 321                         let actual_parent =
 322                             stack.last().and_then(|id| spans.get(id)).map(|s| s.name);
 323                         assert_eq!(
 324                             Some(expected_parent.as_ref()),
 325                             actual_parent,
 326                             "expected {:?} to have contextual parent {:?}",
 327                             name,
 328                             expected_parent,
 329                         );
 330                     }
 331                     None => {}
 332                 }
 333             }
 334         }
 335         spans.insert(
 336             id.clone(),
 337             SpanState {
 338                 name: meta.name(),
 339                 refs: 1,
 340             },
 341         );
 342         id
 343     }
 344
 345     fn enter(&self, id: &Id) {
 346         let spans = self.spans.lock().unwrap();
 347         if let Some(span) = spans.get(id) {
 348             println!("enter: {}; id={:?};", span.name, id);
 349             match self.expected.lock().unwrap().pop_front() {
 350                 None => {}
 351                 Some(Expect::Enter(ref expected_span)) => {
 352                     if let Some(name) = expected_span.name() {
 353                         assert_eq!(name, span.name);
 354                     }
 355                 }
 356                 Some(ex) => ex.bad(format_args!("entered span {:?}", span.name)),
 357             }
 358         };
 359         self.current.lock().unwrap().push(id.clone());
 360     }
 361
 362     fn exit(&self, id: &Id) {
 363         if std::thread::panicking() {
 364             // `exit()` can be called in `drop` impls, so we must guard against
 365             // double panics.
 366             println!("exit {:?} while panicking", id);
 367             return;
 368         }
 369         let spans = self.spans.lock().unwrap();
 370         let span = spans
 371             .get(id)
 372             .unwrap_or_else(|| panic!("no span for ID {:?}", id));
 373         println!("exit: {}; id={:?};", span.name, id);
 374         match self.expected.lock().unwrap().pop_front() {
 375             None => {}
 376             Some(Expect::Exit(ref expected_span)) => {
 377                 if let Some(name) = expected_span.name() {
 378                     assert_eq!(name, span.name);
 379                 }
 380                 let curr = self.current.lock().unwrap().pop();
 381                 assert_eq!(
 382                     Some(id),
 383                     curr.as_ref(),
 384                     "exited span {:?}, but the current span was {:?}",
 385                     span.name,
 386                     curr.as_ref().and_then(|id| spans.get(id)).map(|s| s.name)
 387                 );
 388             }
 389             Some(ex) => ex.bad(format_args!("exited span {:?}", span.name)),
 390         };
 391     }
 392
 393     fn clone_span(&self, id: &Id) -> Id {
 394         let name = self.spans.lock().unwrap().get_mut(id).map(|span| {
 395             let name = span.name;
 396             println!("clone_span: {}; id={:?}; refs={:?};", name, id, span.refs);
 397             span.refs += 1;
 398             name
 399         });
 400         if name.is_none() {
 401             println!("clone_span: id={:?};", id);
 402         }
 403         let mut expected = self.expected.lock().unwrap();
 404         let was_expected = if let Some(Expect::CloneSpan(ref span)) = expected.front() {
 405             assert_eq!(name, span.name());
 406             true
 407         } else {
 408             false
 409         };
 410         if was_expected {
 411             expected.pop_front();
 412         }
 413         id.clone()
 414     }
 415
 416     fn drop_span(&self, id: Id) {
 417         let mut is_event = false;
 418         let name = if let Ok(mut spans) = self.spans.try_lock() {
 419             spans.get_mut(&id).map(|span| {
 420                 let name = span.name;
 421                 if name.contains("event") {
 422                     is_event = true;
 423                 }
 424                 println!("drop_span: {}; id={:?}; refs={:?};", name, id, span.refs);
 425                 span.refs -= 1;
 426                 name
 427             })
 428         } else {
 429             None
 430         };
 431         if name.is_none() {
 432             println!("drop_span: id={:?}", id);
 433         }
 434         if let Ok(mut expected) = self.expected.try_lock() {
 435             let was_expected = match expected.front() {
 436                 Some(Expect::DropSpan(ref span)) => {
 437                     // Don't assert if this function was called while panicking,
 438                     // as failing the assertion can cause a double panic.
 439                     if !::std::thread::panicking() {
 440                         assert_eq!(name, span.name());
 441                     }
 442                     true
 443                 }
 444                 Some(Expect::Event(_)) => {
 445                     if !::std::thread::panicking() {
 446                         assert!(is_event);
 447                     }
 448                     true
 449                 }
 450                 _ => false,
 451             };
 452             if was_expected {
 453                 expected.pop_front();
 454             }
 455         }
 456     }
 457 }
 458
 459 impl MockHandle {
 460     pub fn assert_finished(&self) {
 461         if let Ok(ref expected) = self.0.lock() {
 462             assert!(
 463                 !expected.iter().any(|thing| thing != &Expect::Nothing),
 464                 "more notifications expected: {:?}",
 465                 **expected
 466             );
 467         }
 468     }
 469 }
 470
 471 impl Expect {
 472     fn bad<'a>(&self, what: fmt::Arguments<'a>) {
 473         match self {
 474             Expect::Event(e) => panic!("expected event {}, but {} instead", e, what,),
 475             Expect::Enter(e) => panic!("expected to enter {} but {} instead", e, what,),
 476             Expect::Exit(e) => panic!("expected to exit {} but {} instead", e, what,),
 477             Expect::CloneSpan(e) => panic!("expected to clone {} but {} instead", e, what,),
 478             Expect::DropSpan(e) => panic!("expected to drop {} but {} instead", e, what,),
 479             Expect::Visit(e, fields) => {
 480                 panic!("expected {} to record {} but {} instead", e, fields, what,)
 481             }
 482             Expect::NewSpan(e) => panic!("expected {} but {} instead", e, what),
 483             Expect::Nothing => panic!("expected nothing else to happen, but {} instead", what,),
 484         }
 485     }
 486 }