From 4b578a53627defc06823d193e0ce53f68b49d07c Mon Sep 17 00:00:00 2001
From: Caleb Schoepp <caleb.schoepp@fermyon.com>
Date: Thu, 2 May 2024 16:40:04 -0600
Subject: [PATCH] Refactoring 2

Signed-off-by: Caleb Schoepp <caleb.schoepp@fermyon.com>
---
 crates/observe/src/future.rs         | 27 ++++-----------------------
 crates/observe/src/host_component.rs | 24 +++++++++++++++++++++++-
 2 files changed, 27 insertions(+), 24 deletions(-)

diff --git a/crates/observe/src/future.rs b/crates/observe/src/future.rs
index cf8b6152f..32b6432f9 100644
--- a/crates/observe/src/future.rs
+++ b/crates/observe/src/future.rs
@@ -57,38 +57,21 @@ impl<F: Future> Future for Instrumented<F> {
 
         // 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<RwLock<State>>,
 }
@@ -107,8 +90,6 @@ impl ObserveContext {
     }
 
     fn drop_all(&self) {
-        let mut state: std::sync::RwLockWriteGuard<State> = 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 ccaa20a7b..683c45626 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.