Clean up the host a bunch

Signed-off-by: Caleb Schoepp <[email protected]>

Author: calebschoepp

crates/factor-otel/src/host.rs

@@ -27,11 +27,11 @@ impl wasi_otel::Host for InstanceState {
             );
         }
 
-        // Put the span in our map and push it on to our stack of active spans TODO: Fix comment
+        // Track the guest spans context in our ordered map
         let span_context: opentelemetry::trace::SpanContext = span_data.span_context.into();
-        let span_id = span_context.span_id();
-        state.guest_span_contexts.insert(span_id, span_context);
-        state.active_spans.insert(span_id);
+        state
+            .guest_span_contexts
+            .insert(span_context.span_id(), span_context);
 
         Ok(())
     }
@@ -42,17 +42,13 @@ impl wasi_otel::Host for InstanceState {
         let span_context: opentelemetry::trace::SpanContext = span_data.span_context.clone().into();
         let span_id: opentelemetry::trace::SpanId = span_context.span_id();
 
-        self.processor.on_end(span_data.into());
+        if state.guest_span_contexts.shift_remove(&span_id).is_none() {
+            Err(anyhow!("Trying to end a span that was not started"))?;
+        }
 
-        if let Some(_guest_span) = state.guest_span_contexts.get_mut(&span_id) {
-            // Remove the span from active_spans
-            state.active_spans.shift_remove(&span_id);
+        self.processor.on_end(span_data.into());
 
-            Ok(())
-        } else {
-            // TODO: This seems to be wrong
-            Err(anyhow!("BUG: cannot find resource in table"))
-        }
+        Ok(())
     }
 
     async fn current_span_context(&mut self) -> Result<wasi_otel::SpanContext> {

crates/factor-otel/src/lib.rs

@@ -1,13 +1,12 @@
 mod host;
 
 use std::{
-    collections::HashMap,
     sync::{Arc, RwLock},
     time::Duration,
 };
 
 use anyhow::bail;
-use indexmap::IndexSet;
+use indexmap::IndexMap;
 use opentelemetry::{
     trace::{SpanContext, SpanId, TraceContextExt},
     Context,
@@ -53,7 +52,6 @@ impl Factor for OtelFactor {
         Ok(InstanceState {
             state: Arc::new(RwLock::new(State {
                 guest_span_contexts: Default::default(),
-                active_spans: Default::default(),
                 original_host_span_id: None,
             })),
             processor: self.processor.clone(),
@@ -63,7 +61,6 @@ impl Factor for OtelFactor {
 
 impl OtelFactor {
     pub fn new() -> anyhow::Result<Self> {
-        // TODO: Configuring the processor should move to init
         // This will configure the exporter based on the OTEL_EXPORTER_* environment variables.
         let exporter = match OtlpProtocol::traces_protocol_from_env() {
             OtlpProtocol::Grpc => opentelemetry_otlp::SpanExporter::builder()
@@ -111,16 +108,13 @@ impl SelfInstanceBuilder for InstanceState {}
 /// This data lives here rather than directly on InstanceState so that we can have multiple things
 /// take Arc references to it.
 pub(crate) struct State {
-    /// A mapping between immutable [SpanId]s and the actual [BoxedSpan] created by our tracer.
-    // TODO: Rename to not include "guest"
-    // TODO: Merge with active_spans
-    pub(crate) guest_span_contexts: HashMap<SpanId, SpanContext>,
-
-    /// A stack of [SpanIds] for all the active spans. The topmost span is the active span.
+    /// An order-preserved mapping between immutable [SpanId]s of guest created spans and their
+    /// corresponding [SpanContext].
     ///
-    /// When a span is ended it is removed from this stack (regardless of whether is the
-    /// active span) and all other spans are shifted back to retain relative order.
-    pub(crate) active_spans: IndexSet<SpanId>,
+    /// The topmost [SpanId] is the last active span. When a span is ended it is removed from this
+    /// map (regardless of whether it is the active span) and all other spans are shifted back to
+    /// retain relative order.
+    pub(crate) guest_span_contexts: IndexMap<SpanId, SpanContext>,
 
     /// Id of the last span emitted from within the host before entering the guest.
     ///
@@ -129,12 +123,6 @@ pub(crate) struct State {
     pub(crate) original_host_span_id: Option<SpanId>,
 }
 
-// /// The WIT resource Span. Effectively wraps an [opentelemetry::global::BoxedSpan].
-// pub struct GuestSpan {
-//     /// The [opentelemetry::global::BoxedSpan] we use to do the actual tracing work.
-//     pub inner: BoxedSpan,
-// }
-
 /// Manages access to the OtelFactor state for the purpose of maintaining proper span
 /// parent/child relationships when WASI Otel spans are being created.
 pub struct OtelContext {
@@ -170,10 +158,10 @@ impl OtelContext {
     ///          | spin_key_value.get |
     /// ```
     ///
-    ///  Setting the guest spans parent as the host is trivially done. However, the more difficult
-    /// task is having the host factor spans be children of the guest span.
-    /// [`OtelContext::reparent_tracing_span`] handles this by reparenting the current span to be
-    /// a child of the last active guest span (which is tracked internally in the otel factor).
+    ///  Setting the guest spans parent as the host is enabled through current_span_context.
+    /// However, the more difficult task is having the host factor spans be children of the guest
+    /// span. [`OtelContext::reparent_tracing_span`] handles this by reparenting the current span to
+    /// be a child of the last active guest span (which is tracked internally in the otel factor).
     ///
     /// Note that if the otel factor is not in your [`RuntimeFactors`] than this is effectively a
     /// no-op.
@@ -191,7 +179,7 @@ impl OtelContext {
         };
 
         // If there are no active guest spans then there is nothing to do
-        let Some(active_span) = state.active_spans.last() else {
+        let Some((_, active_span_context)) = state.guest_span_contexts.last() else {
             return;
         };
 
@@ -209,8 +197,7 @@ impl OtelContext {
         }
 
         // Now reparent the current span to the last active guest span
-        let span_context = state.guest_span_contexts.get(active_span).unwrap().clone();
-        let parent_context = Context::new().with_remote_span_context(span_context);
+        let parent_context = Context::new().with_remote_span_context(active_span_context.clone());
         tracing::Span::current().set_parent(parent_context);
     }
 }