use std::collections::hash_map::Entry;
use std::collections::HashMap;
use sentry_core::HubSwitchGuard;
use tracing_core::span::Id as SpanId;
pub(super) struct SpanGuardStack {
guards: HashMap<SpanId, Vec<HubSwitchGuard>>,
}
impl SpanGuardStack {
pub(super) fn new() -> Self {
Self {
guards: HashMap::new(),
}
}
pub(super) fn push(&mut self, id: SpanId, guard: HubSwitchGuard) {
self.guards.entry(id).or_default().push(guard);
}
#[must_use]
pub(super) fn pop(&mut self, id: SpanId) -> Option<HubSwitchGuard> {
match self.guards.entry(id) {
Entry::Occupied(mut entry) => {
let stack = entry.get_mut();
let guard = stack.pop();
if stack.is_empty() {
entry.remove();
}
guard
}
Entry::Vacant(_) => None,
}
}
}
#[cfg(test)]
mod tests {
use super::SpanGuardStack;
use sentry_core::{Hub, HubSwitchGuard};
use std::sync::Arc;
use tracing_core::span::Id as SpanId;
#[test]
fn pop_is_lifo() {
let initial = Hub::current();
let hub_a = Arc::new(Hub::new_from_top(initial.clone()));
let hub_b = Arc::new(Hub::new_from_top(hub_a.clone()));
let mut stack = SpanGuardStack::new();
let id = SpanId::from_u64(1);
stack.push(id.clone(), HubSwitchGuard::new(hub_a.clone()));
assert!(Arc::ptr_eq(&Hub::current(), &hub_a));
stack.push(id.clone(), HubSwitchGuard::new(hub_b.clone()));
assert!(Arc::ptr_eq(&Hub::current(), &hub_b));
drop(stack.pop(id.clone()).expect("guard for hub_b"));
assert!(Arc::ptr_eq(&Hub::current(), &hub_a));
drop(stack.pop(id.clone()).expect("guard for hub_a"));
assert!(Arc::ptr_eq(&Hub::current(), &initial));
assert!(stack.pop(id).is_none());
}
}