diff --git a/crates/observe/src/future.rs b/crates/observe/src/future.rs index cf8b6152f4..32b6432f97 100644 --- a/crates/observe/src/future.rs +++ b/crates/observe/src/future.rs @@ -57,38 +57,21 @@ impl Future for Instrumented { // Enter the active spans before entering the inner poll { - let state = this.observe_context.state.write().unwrap(); - - // TODO: Make it a method on state - for guest_span_id in state.active_spans.iter() { - if let Some(span_resource) = state.guest_spans.get(*guest_span_id) { - span_resource.enter(); - } else { - tracing::error!("No span to enter") - } - } + this.observe_context.state.write().unwrap().enter_all(); } let ret = this.inner.poll(cx); // Exit the active spans after exiting the inner poll { - let state = this.observe_context.state.write().unwrap(); - - // TODO: Make it a method on state - for span_id in state.active_spans.iter().rev() { - if let Some(span_resource) = state.guest_spans.get(*span_id) { - span_resource.exit(); - } else { - tracing::error!("span already dropped") - } - } + this.observe_context.state.write().unwrap().exit_all(); } ret } } +/// The context necessary for the observe host component to function. pub struct ObserveContext { state: Arc>, } @@ -107,8 +90,6 @@ impl ObserveContext { } fn drop_all(&self) { - let mut state: std::sync::RwLockWriteGuard = self.state.write().unwrap(); - - state.close_from_back_to(0); + self.state.write().unwrap().close_from_back_to(0); } } diff --git a/crates/observe/src/host_component.rs b/crates/observe/src/host_component.rs index ccaa20a7b7..683c456262 100644 --- a/crates/observe/src/host_component.rs +++ b/crates/observe/src/host_component.rs @@ -120,7 +120,7 @@ pub(crate) struct State { impl State { /// Close all active spans from the top of the stack to the given index. Closing entails exiting /// the inner [tracing] span and removing it from the active spans stack. - pub fn close_from_back_to(&mut self, index: usize) { + pub(crate) fn close_from_back_to(&mut self, index: usize) { self.active_spans .split_off(index) .iter() @@ -133,6 +133,28 @@ impl State { } }); } + + /// Enter the inner [tracing] span for all active spans. + pub(crate) fn enter_all(&self) { + for guest_span_id in self.active_spans.iter() { + if let Some(span_resource) = self.guest_spans.get(*guest_span_id) { + span_resource.enter(); + } else { + tracing::debug!("guest span already dropped") + } + } + } + + /// Exit the inner [tracing] span for all active spans. + pub(crate) fn exit_all(&self) { + for guest_span_id in self.active_spans.iter().rev() { + if let Some(span_resource) = self.guest_spans.get(*guest_span_id) { + span_resource.exit(); + } else { + tracing::debug!("guest span already dropped") + } + } + } } /// The WIT resource Span. Effectively wraps a [tracing] span.