fix(compute): preserve docker tracing after decoupling

Signed-off-by: Drew Newberry <385+drew@users.noreply.github.com>

Author: drew

crates/openshell-driver-docker/src/lib.rs

@@ -55,6 +55,7 @@ use openshell_core::proto_struct::{
     deserialize_optional_non_empty_string_list, struct_to_json_value,
 };
 use openshell_core::{Error, Result as CoreResult};
+use opentelemetry::trace::TraceContextExt as _;
 use std::collections::{HashMap, HashSet};
 use std::future::Future;
 use std::net::{IpAddr, Ipv4Addr, SocketAddr};
@@ -67,7 +68,8 @@ use tokio::sync::{Mutex, broadcast, mpsc};
 use tokio::task::JoinHandle;
 use tokio_stream::wrappers::ReceiverStream;
 use tonic::{Request, Response, Status};
-use tracing::{debug, info, warn};
+use tracing::{Instrument as _, debug, info, warn};
+use tracing_opentelemetry::OpenTelemetrySpanExt as _;
 use url::Url;
 
 const WATCH_BUFFER: usize = 128;
@@ -84,6 +86,28 @@ const HOST_OPENSHELL_INTERNAL: &str = "host.openshell.internal";
 const HOST_DOCKER_INTERNAL: &str = "host.docker.internal";
 const DOCKER_NETWORK_DRIVER: &str = "bridge";
 
+fn provisioning_span(
+    parent: &opentelemetry::Context,
+    sandbox: &DriverSandbox,
+    image_ref: &str,
+) -> tracing::Span {
+    let span = tracing::info_span!(
+        parent: None,
+        "docker.provision",
+        otel.name = "docker.provision",
+        otel.status_code = tracing::field::Empty,
+        sandbox.id = %sandbox.id,
+        sandbox.name = %sandbox.name,
+        image.ref = %image_ref,
+    );
+    let parent_span_context = parent.span().span_context().clone();
+    if parent_span_context.is_valid() {
+        let parent = opentelemetry::Context::new().with_remote_span_context(parent_span_context);
+        let _ = span.set_parent(parent);
+    }
+    span
+}
+
 /// Gateway-local configuration for the Docker compute driver.
 #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
 #[serde(default, deny_unknown_fields)]
@@ -877,7 +901,7 @@ impl DockerComputeDriver {
             &sandbox.id,
             "Scheduled",
             format!("Docker sandbox accepted for image \"{image}\""),
-            HashMap::from([("image_ref".to_string(), image)]),
+            HashMap::from([("image_ref".to_string(), image.clone())]),
         );
         self.publish_sandbox_snapshot(pending_sandbox_snapshot(
             sandbox,
@@ -889,9 +913,14 @@ impl DockerComputeDriver {
         let driver = self.clone();
         let sandbox_for_task = sandbox.clone();
         let sandbox_id = sandbox.id.clone();
-        let task = tokio::spawn(async move {
-            driver.provision_sandbox(sandbox_for_task).await;
-        });
+        let parent = tracing::Span::current().context();
+        let provisioning_span = provisioning_span(&parent, sandbox, &image);
+        let task = tokio::spawn(
+            async move {
+                driver.provision_sandbox(sandbox_for_task).await;
+            }
+            .instrument(provisioning_span),
+        );
 
         let mut pending = self.pending.lock().await;
         if let Some(record) = pending.get_mut(&sandbox_id) {
@@ -914,20 +943,42 @@ impl DockerComputeDriver {
         }
     }
 
+    #[tracing::instrument(
+        name = "docker.provision_sandbox",
+        skip(self, sandbox),
+        fields(
+            otel.name = "docker.provision_sandbox",
+            otel.status_code = tracing::field::Empty,
+            sandbox.id = %sandbox.id,
+            sandbox.name = %sandbox.name,
+        )
+    )]
     async fn provision_sandbox_inner(
         &self,
         sandbox: &DriverSandbox,
     ) -> Result<(), DockerProvisioningFailure> {
+        let span_status = openshell_otel::ErrorStatusGuard::current();
         let validated = Self::validated_sandbox(sandbox, &self.config).map_err(|status| {
             DockerProvisioningFailure::new("ContainerCreateFailed", status.message())
         })?;
         let template = validated.template;
-        let image = self
-            .ensure_image_available(&sandbox.id, &template.image)
-            .await
-            .map_err(|status| {
-                DockerProvisioningFailure::new("ImagePullFailed", status.message())
-            })?;
+        let image = async {
+            openshell_otel::record_error_result(
+                self.ensure_image_available(&sandbox.id, &template.image)
+                    .await
+                    .map_err(|status| {
+                        DockerProvisioningFailure::new("ImagePullFailed", status.message())
+                    }),
+            )
+        }
+        .instrument(tracing::info_span!(
+            "docker.prepare_image",
+            otel.name = "docker.prepare_image",
+            otel.status_code = tracing::field::Empty,
+            sandbox.id = %sandbox.id,
+            image.ref = %template.image,
+        ))
+        .await?;
         let token_file_created = write_sandbox_token_file(sandbox, &self.config)
             .await
             .map_err(|status| {
@@ -961,33 +1012,58 @@ impl DockerComputeDriver {
             }
             DockerProvisioningFailure::new("ContainerCreateFailed", status.message())
         })?;
-        self.docker
-            .create_container(
-                Some(
-                    CreateContainerOptionsBuilder::default()
-                        .name(container_name.as_str())
-                        .build(),
-                ),
-                create_body,
+        async {
+            openshell_otel::record_error_result(
+                self.docker
+                    .create_container(
+                        Some(
+                            CreateContainerOptionsBuilder::default()
+                                .name(container_name.as_str())
+                                .build(),
+                        ),
+                        create_body,
+                    )
+                    .await
+                    .map_err(|err| {
+                        if token_file_created {
+                            cleanup_sandbox_token_file(sandbox, &self.config);
+                        }
+                        DockerProvisioningFailure::from_status(
+                            "ContainerCreateFailed",
+                            create_status_from_docker_error("create docker sandbox container", err),
+                        )
+                    }),
             )
-            .await
-            .map_err(|err| {
-                if token_file_created {
-                    cleanup_sandbox_token_file(sandbox, &self.config);
-                }
-                DockerProvisioningFailure::from_status(
-                    "ContainerCreateFailed",
-                    create_status_from_docker_error("create docker sandbox container", err),
-                )
-            })?;
+        }
+        .instrument(tracing::info_span!(
+            "docker.create_container",
+            otel.name = "docker.create_container",
+            otel.status_code = tracing::field::Empty,
+            sandbox.id = %sandbox.id,
+            container.name = %container_name,
+        ))
+        .await?;
         self.publish_docker_progress(
             &sandbox.id,
             "Created",
             format!("Created Docker container \"{container_name}\""),
             HashMap::from([("container_name".to_string(), container_name.clone())]),
         );
 
-        if let Err(err) = self.docker.start_container(&container_name, None).await {
+        let start_result = async {
+            openshell_otel::record_error_result(
+                self.docker.start_container(&container_name, None).await,
+            )
+        }
+        .instrument(tracing::info_span!(
+            "docker.start_container",
+            otel.name = "docker.start_container",
+            otel.status_code = tracing::field::Empty,
+            sandbox.id = %sandbox.id,
+            container.name = %container_name,
+        ))
+        .await;
+        if let Err(err) = start_result {
             let cleanup = self
                 .docker
                 .remove_container(
@@ -1028,7 +1104,7 @@ impl DockerComputeDriver {
             );
         }
 
-        Ok(())
+        span_status.finish(Ok(()))
     }
 
     async fn delete_sandbox_inner(
@@ -1142,17 +1218,29 @@ impl DockerComputeDriver {
     /// Returns `Ok(true)` when a container existed and was started (or was
     /// already running), `Ok(false)` when no managed container is found for
     /// the sandbox, and `Err(...)` for any Docker failure.
+    #[tracing::instrument(
+        name = "docker.start_sandbox",
+        skip(self),
+        fields(
+            otel.name = "docker.start_sandbox",
+            otel.status_code = tracing::field::Empty,
+            sandbox.id = %sandbox_id,
+            sandbox.name = %sandbox_name,
+        )
+    )]
     pub async fn start_sandbox(
         &self,
         sandbox_id: &str,
         sandbox_name: &str,
     ) -> Result<bool, Status> {
+        let span_status = openshell_otel::ErrorStatusGuard::current();
+        require_sandbox_identifier(sandbox_id, sandbox_name)?;
         self.lifecycle_event_fences.begin_start(sandbox_id);
         let result = self
             .start_sandbox_with_lifecycle_fence(sandbox_id, sandbox_name)
             .await;
         self.lifecycle_event_fences.finish_start(sandbox_id);
-        result
+        span_status.finish(result)
     }
 
     async fn start_sandbox_with_lifecycle_fence(
@@ -1924,38 +2012,59 @@ impl ComputeDriver for DockerComputeDriver {
         }))
     }
 
+    #[tracing::instrument(
+        name = "docker.schedule_sandbox",
+        skip(self, request),
+        fields(
+            otel.name = "docker.schedule_sandbox",
+            otel.status_code = tracing::field::Empty,
+            sandbox.id = %request.get_ref().sandbox.as_ref().map_or("", |sandbox| sandbox.id.as_str()),
+            sandbox.name = %request.get_ref().sandbox.as_ref().map_or("", |sandbox| sandbox.name.as_str()),
+        )
+    )]
     async fn create_sandbox(
         &self,
         request: Request<CreateSandboxRequest>,
     ) -> Result<Response<CreateSandboxResponse>, Status> {
+        let span_status = openshell_otel::ErrorStatusGuard::current();
         let sandbox = request
             .into_inner()
             .sandbox
             .ok_or_else(|| Status::invalid_argument("sandbox is required"))?;
         self.create_sandbox_inner(&sandbox).await?;
-        Ok(Response::new(CreateSandboxResponse {}))
+        span_status.finish(Ok(Response::new(CreateSandboxResponse {})))
     }
 
+    #[tracing::instrument(
+        name = "docker.stop_sandbox",
+        skip(self, request),
+        fields(
+            otel.name = "docker.stop_sandbox",
+            otel.status_code = tracing::field::Empty,
+            sandbox.id = %request.get_ref().sandbox_id,
+            sandbox.name = %request.get_ref().sandbox_name,
+        )
+    )]
     async fn stop_sandbox(
         &self,
         request: Request<StopSandboxRequest>,
     ) -> Result<Response<StopSandboxResponse>, Status> {
+        let span_status = openshell_otel::ErrorStatusGuard::current();
         let request = request.into_inner();
         require_sandbox_identifier(&request.sandbox_id, &request.sandbox_name)?;
 
         self.stop_sandbox_inner(&request.sandbox_id, &request.sandbox_name)
             .await?;
         self.publish_container_snapshot(&request.sandbox_id, &request.sandbox_name)
             .await?;
-        Ok(Response::new(StopSandboxResponse {}))
+        span_status.finish(Ok(Response::new(StopSandboxResponse {})))
     }
 
     async fn start_sandbox(
         &self,
         request: Request<StartSandboxRequest>,
     ) -> Result<Response<StartSandboxResponse>, Status> {
         let request = request.into_inner();
-        require_sandbox_identifier(&request.sandbox_id, &request.sandbox_name)?;
         if !Self::start_sandbox(self, &request.sandbox_id, &request.sandbox_name).await? {
             return Err(Status::not_found("sandbox not found"));
         }
@@ -1964,10 +2073,21 @@ impl ComputeDriver for DockerComputeDriver {
         Ok(Response::new(StartSandboxResponse {}))
     }
 
+    #[tracing::instrument(
+        name = "docker.delete_sandbox",
+        skip(self, request),
+        fields(
+            otel.name = "docker.delete_sandbox",
+            otel.status_code = tracing::field::Empty,
+            sandbox.id = %request.get_ref().sandbox_id,
+            sandbox.name = %request.get_ref().sandbox_name,
+        )
+    )]
     async fn delete_sandbox(
         &self,
         request: Request<DeleteSandboxRequest>,
     ) -> Result<Response<DeleteSandboxResponse>, Status> {
+        let span_status = openshell_otel::ErrorStatusGuard::current();
         let request = request.into_inner();
         require_sandbox_identifier(&request.sandbox_id, &request.sandbox_name)?;
 
@@ -1986,7 +2106,7 @@ impl ComputeDriver for DockerComputeDriver {
             });
         }
 
-        Ok(Response::new(DeleteSandboxResponse { deleted }))
+        span_status.finish(Ok(Response::new(DeleteSandboxResponse { deleted })))
     }
 
     async fn watch_sandboxes(

crates/openshell-server/src/lib.rs

@@ -1409,7 +1409,7 @@ async fn build_compute_runtime(
                 inherited_config_keys: registration.inherited_config_keys,
             };
             let instance = registration.factory.build(build_context).await?;
-            let runtime = match instance {
+            match instance {
                 ComputeDriverInstance::InProcess(driver) => ComputeRuntime::from_driver(
                     registration.name,
                     driver,
@@ -1439,8 +1439,7 @@ async fn build_compute_runtime(
                         Error::execution(format!("failed to create compute runtime: {error}"))
                     })?
                 }
-            };
-            runtime
+            }
         }
         ConfiguredComputeDriver::Remote { name } => {
             let remote_config =
@@ -1453,7 +1452,7 @@ async fn build_compute_runtime(
             let endpoint = compute::connect_remote_compute_driver(name, &remote_config.socket_path)
                 .await
                 .map_err(|e| Error::execution(format!("failed to create compute runtime: {e}")))?;
-            let runtime = ComputeRuntime::new_remote_driver(
+            ComputeRuntime::new_remote_driver(
                 endpoint,
                 store,
                 sandbox_index,
@@ -1462,8 +1461,7 @@ async fn build_compute_runtime(
                 supervisor_sessions,
             )
             .await
-            .map_err(|e| Error::execution(format!("failed to create compute runtime: {e}")))?;
-            runtime
+            .map_err(|e| Error::execution(format!("failed to create compute runtime: {e}")))?
         }
     };