A rough first implementation of getting tests passing with the opentelemetry based backend

Signed-off-by: Caleb Schoepp <caleb.schoepp@fermyon.com>

Author: calebschoepp

crates/factor-observe/src/future.rs

@@ -6,75 +6,75 @@ use std::{
 
 use crate::State;
 
-pin_project! {
-    struct Instrumented<F> {
-        #[pin]
-        inner: F,
-        observe_context: ObserveContext,
-    }
+// pin_project! {
+//     struct Instrumented<F> {
+//         #[pin]
+//         inner: F,
+//         observe_context: ObserveContext,
+//     }
 
-    impl<F> PinnedDrop for Instrumented<F> {
-        fn drop(this: Pin<&mut Self>) {
-            this.project().observe_context.drop_all();
-        }
-    }
-}
+//     impl<F> PinnedDrop for Instrumented<F> {
+//         fn drop(this: Pin<&mut Self>) {
+//             this.project().observe_context.drop_all();
+//         }
+//     }
+// }
 
-pub trait FutureExt: Future + Sized {
-    /// Manage WASI Observe guest spans.
-    fn manage_wasi_observe_spans(
-        self,
-        observe_context: ObserveContext,
-    ) -> impl Future<Output = Self::Output>;
-}
+// pub trait FutureExt: Future + Sized {
+//     /// Manage WASI Observe guest spans.
+//     fn manage_wasi_observe_spans(
+//         self,
+//         observe_context: ObserveContext,
+//     ) -> impl Future<Output = Self::Output>;
+// }
 
-impl<F: Future> FutureExt for F {
-    fn manage_wasi_observe_spans(
-        self,
-        observe_context: ObserveContext,
-    ) -> impl Future<Output = Self::Output> {
-        Instrumented {
-            inner: self,
-            observe_context,
-        }
-    }
-}
+// impl<F: Future> FutureExt for F {
+//     fn manage_wasi_observe_spans(
+//         self,
+//         observe_context: ObserveContext,
+//     ) -> impl Future<Output = Self::Output> {
+//         Instrumented {
+//             inner: self,
+//             observe_context,
+//         }
+//     }
+// }
 
-impl<F: Future> Future for Instrumented<F> {
-    type Output = F::Output;
+// impl<F: Future> Future for Instrumented<F> {
+//     type Output = F::Output;
 
-    /// Maintains the invariant that all active spans are entered before polling the inner future
-    /// and exited otherwise. If we don't do this then the timing (among many other things) of the
-    /// spans becomes wildly incorrect.
-    fn poll(
-        self: std::pin::Pin<&mut Self>,
-        cx: &mut std::task::Context<'_>,
-    ) -> std::task::Poll<Self::Output> {
-        let this = self.project();
+//     /// Maintains the invariant that all active spans are entered before polling the inner future
+//     /// and exited otherwise. If we don't do this then the timing (among many other things) of the
+//     /// spans becomes wildly incorrect.
+//     fn poll(
+//         self: std::pin::Pin<&mut Self>,
+//         cx: &mut std::task::Context<'_>,
+//     ) -> std::task::Poll<Self::Output> {
+//         let this = self.project();
 
-        // Enter the active spans before entering the inner poll
-        {
-            this.observe_context.state.write().unwrap().enter_all();
-        }
+//         // Enter the active spans before entering the inner poll
+//         {
+//             this.observe_context.state.write().unwrap().enter_all();
+//         }
 
-        let ret = this.inner.poll(cx);
+//         let ret = this.inner.poll(cx);
 
-        // Exit the active spans after exiting the inner poll
-        {
-            this.observe_context.state.write().unwrap().exit_all();
-        }
+//         // Exit the active spans after exiting the inner poll
+//         {
+//             this.observe_context.state.write().unwrap().exit_all();
+//         }
 
-        ret
-    }
-}
+//         ret
+//     }
+// }
 
 /// The context necessary for the observe host component to function.
 pub struct ObserveContext {
     pub(crate) state: Arc<RwLock<State>>,
 }
 
-impl ObserveContext {
-    fn drop_all(&self) {
-        self.state.write().unwrap().close_from_back_to(0);
-    }
-}
+// impl ObserveContext {
+//     fn drop_all(&self) {
+//         self.state.write().unwrap().close_from_back_to(0);
+//     }
+// }

crates/factor-observe/src/host.rs

@@ -1,7 +1,11 @@
 use anyhow::Result;
+use opentelemetry::global::ObjectSafeSpan;
 use opentelemetry::trace::TraceContextExt;
+use opentelemetry::trace::Tracer;
+use opentelemetry::Context;
 use spin_core::async_trait;
 use spin_core::wasmtime::component::Resource;
+use spin_telemetry::traces::WASI_OBSERVE_TRACER;
 use spin_world::wasi::clocks0_2_0::wall_clock::Datetime;
 use spin_world::wasi::observe::traces::{self, KeyValue, Span as WitSpan};
 use tracing_opentelemetry::OpenTelemetrySpanExt;
@@ -13,25 +17,38 @@ impl traces::Host for InstanceState {}
 
 #[async_trait]
 impl traces::HostSpan for InstanceState {
+    // TODO(Caleb): Make this implicit logic make more sense (the indexmap seems wrong)
     async fn start(&mut self, name: String) -> Result<Resource<WitSpan>> {
-        // Create the underlying tracing span
-        let tracing_span = tracing::info_span!("WASI Observe guest", "otel.name" = name);
-        let span_id = tracing_span
-            .context()
-            .span()
-            .span_context()
-            .span_id()
-            .to_string();
+        let mut state = self.state.write().unwrap();
 
-        // Wrap it in a GuestSpan for our own bookkeeping purposes and enter it
-        let guest_span = GuestSpan {
-            name: name.clone(),
-            inner: tracing_span,
+        // TODO(Caleb): Make this cleaner
+        let parent_context = match state.active_spans.is_empty() {
+            true => tracing::Span::current().context(),
+            false => Context::new().with_remote_span_context(
+                state
+                    .guest_spans
+                    .get(*state.active_spans.last().unwrap().1)
+                    .unwrap()
+                    .inner
+                    .span_context()
+                    .clone(),
+            ),
         };
-        guest_span.enter();
+
+        // Create the underlying opentelemetry span
+        let otel_span = WASI_OBSERVE_TRACER
+            .lock()
+            .unwrap()
+            .clone()
+            .unwrap()
+            .start_with_context(name, &parent_context);
+
+        let span_id = otel_span.span_context().span_id().to_string();
+
+        // Wrap it in a GuestSpan for our own bookkeeping purposes
+        let guest_span = GuestSpan { inner: otel_span };
 
         // Put the GuestSpan in our resource table and push it to our stack of active spans
-        let mut state = self.state.write().unwrap();
         let resource_id = state.guest_spans.push(guest_span).unwrap();
         state.active_spans.insert(span_id, resource_id);
 
@@ -45,14 +62,12 @@ impl traces::HostSpan for InstanceState {
     ) -> Result<()> {
         if let Some(guest_span) = self
             .state
-            .try_write()
+            .write()
             .unwrap()
             .guest_spans
             .get_mut(resource.rep())
         {
-            guest_span
-                .inner
-                .set_attribute(attribute.key, attribute.value);
+            guest_span.inner.set_attribute(attribute.into())
         } else {
             tracing::debug!("can't find guest span to set attribute on")
         }
@@ -72,9 +87,7 @@ impl traces::HostSpan for InstanceState {
             .get_mut(resource.rep())
         {
             for attribute in attributes {
-                guest_span
-                    .inner
-                    .set_attribute(attribute.key, attribute.value);
+                guest_span.inner.set_attribute(attribute.into());
             }
         } else {
             tracing::debug!("can't find guest span to set attributes on")
@@ -94,12 +107,24 @@ impl traces::HostSpan for InstanceState {
     }
 
     async fn end(&mut self, resource: Resource<WitSpan>) -> Result<()> {
-        self.safely_close(resource.rep(), false);
+        if let Some(guest_span) = self
+            .state
+            .write()
+            .unwrap()
+            .guest_spans
+            .get_mut(resource.rep())
+        {
+            guest_span.inner.end()
+        } else {
+            tracing::debug!("can't find guest span to end")
+        }
         Ok(())
     }
 
-    fn drop(&mut self, resource: Resource<WitSpan>) -> Result<()> {
-        self.safely_close(resource.rep(), true);
+    fn drop(&mut self, _resource: Resource<WitSpan>) -> Result<()> {
+        // TODO(Caleb): What do we want the dropping behavior to be?
         Ok(())
     }
 }
+
+// TODO(Caleb): Improve debug tracing in failure cases