Catch reparent being used outside of factor

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

crates/factor-observe/Cargo.toml

@@ -26,7 +26,6 @@ thiserror = "1"
 tokio = { version = "1", features = ["rt-multi-thread"] }
 tracing = "0.1.40"
 tracing-opentelemetry = "0.23.0"
-vaultrs = "0.6.2"
 
 [dev-dependencies]
 toml = "0.5"

crates/factor-observe/src/host.rs

@@ -20,6 +20,16 @@ impl traces::HostSpan for InstanceState {
     async fn start(&mut self, name: String) -> Result<Resource<WitSpan>> {
         let mut state = self.state.write().unwrap();
 
+        if state.active_spans.is_empty() {
+            state.original_host_span_id = Some(
+                tracing::Span::current()
+                    .context()
+                    .span()
+                    .span_context()
+                    .span_id(),
+            );
+        }
+
         // TODO(Caleb): Make this cleaner
         let parent_context = match state.active_spans.is_empty() {
             true => tracing::Span::current().context(),

crates/factor-observe/src/lib.rs

@@ -5,7 +5,7 @@ use std::sync::{Arc, RwLock};
 use indexmap::IndexMap;
 use opentelemetry::{
     global::{self, BoxedTracer, ObjectSafeSpan},
-    trace::TraceContextExt,
+    trace::{SpanId, TraceContextExt},
     Context,
 };
 use spin_factors::{Factor, SelfInstanceBuilder};
@@ -46,6 +46,7 @@ impl Factor for ObserveFactor {
             state: Arc::new(RwLock::new(State {
                 guest_spans: table::Table::new(1024),
                 active_spans: Default::default(),
+                original_host_span_id: None,
             })),
             tracer,
         })
@@ -80,14 +81,20 @@ impl InstanceState {
 /// take Arc references to it.
 pub(crate) struct State {
     /// A resource table that holds the guest spans.
-    pub guest_spans: table::Table<GuestSpan>,
+    pub(crate) guest_spans: table::Table<GuestSpan>,
 
     /// A LIFO stack of guest spans that are currently active.
     ///
     /// Only a reference ID to the guest span is held here. The actual guest span must be looked up
     /// in the `guest_spans` table using the reference ID.
     /// TODO: Fix comment
-    pub active_spans: IndexMap<String, u32>,
+    pub(crate) active_spans: IndexMap<String, u32>,
+
+    /// Id of the last span emitted from within the host before entering the guest.
+    ///
+    /// We use this to avoid accidentally reparenting the original host span as a child of a guest
+    /// span.
+    pub(crate) original_host_span_id: Option<SpanId>,
 }
 
 /// The WIT resource Span. Effectively wraps an [opentelemetry::global::BoxedSpan].
@@ -112,6 +119,18 @@ impl ObserveContext {
             return;
         }
 
+        if let Some(original_host_span_id) = state.original_host_span_id {
+            if tracing::Span::current()
+                .context()
+                .span()
+                .span_context()
+                .span_id()
+                .eq(&original_host_span_id)
+            {
+                panic!("TODO This should not happen")
+            }
+        }
+
         let parent_context = Context::new().with_remote_span_context(
             state
                 .guest_spans