fix(compute): preserve driver integrations after rebase

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

Author: drew

crates/openshell-core/src/lib.rs

@@ -45,6 +45,7 @@ pub mod provider_credentials;
 pub mod sandbox_env;
 pub mod secrets;
 pub mod settings;
+pub mod spiffe;
 pub mod telemetry;
 pub mod time;
 pub mod transport_errors;

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

@@ -5,6 +5,8 @@
 
 #![allow(clippy::result_large_err)]
 
+pub mod otel_tracing;
+
 use bollard::Docker;
 use bollard::errors::Error as BollardError;
 use bollard::models::{
@@ -54,10 +56,12 @@ use openshell_core::proto_struct::{
 };
 use openshell_core::{Error, Result as CoreResult};
 use std::collections::{HashMap, HashSet};
+use std::future::Future;
 use std::net::{IpAddr, Ipv4Addr, SocketAddr};
 use std::path::{Path, PathBuf};
 use std::pin::Pin;
 use std::sync::Arc;
+use std::task::{Context, Poll};
 use std::time::Duration;
 use tokio::sync::{Mutex, broadcast, mpsc};
 use tokio::task::JoinHandle;
@@ -385,6 +389,119 @@ fn default_true() -> bool {
 type WatchStream =
     Pin<Box<dyn Stream<Item = Result<WatchSandboxesEvent, Status>> + Send + 'static>>;
 
+struct TracedWatchStream {
+    inner: WatchStream,
+    span: tracing::Span,
+    finished: bool,
+}
+
+impl Stream for TracedWatchStream {
+    type Item = Result<WatchSandboxesEvent, Status>;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let span = self.span.clone();
+        let _entered = span.enter();
+        let result = self.inner.as_mut().poll_next(cx);
+        if !self.finished {
+            match &result {
+                Poll::Ready(Some(Err(status))) => {
+                    openshell_otel::mark_error(&self.span);
+                    self.span
+                        .record("rpc.grpc.status_code", status.code() as i32);
+                    self.finished = true;
+                }
+                Poll::Ready(None) => {
+                    self.span
+                        .record("rpc.grpc.status_code", tonic::Code::Ok as i32);
+                    self.finished = true;
+                }
+                Poll::Pending | Poll::Ready(Some(Ok(_))) => {}
+            }
+        }
+        result
+    }
+}
+
+impl Drop for TracedWatchStream {
+    fn drop(&mut self) {
+        if !self.finished {
+            openshell_otel::mark_error(&self.span);
+            self.span
+                .record("rpc.grpc.status_code", tonic::Code::Cancelled as i32);
+        }
+    }
+}
+
+/// Compute-driver service wrapper that preserves the standalone RPC trace
+/// boundary while Docker runs in the gateway process.
+#[derive(Clone)]
+pub struct ComputeDriverService {
+    driver: DockerComputeDriver,
+    trace_in_process_rpc: bool,
+}
+
+impl ComputeDriverService {
+    #[must_use]
+    pub fn new(driver: DockerComputeDriver) -> Self {
+        Self {
+            driver,
+            trace_in_process_rpc: false,
+        }
+    }
+
+    #[must_use]
+    pub fn new_in_process(driver: DockerComputeDriver) -> Self {
+        Self {
+            driver,
+            trace_in_process_rpc: true,
+        }
+    }
+
+    fn in_process_rpc_span(
+        &self,
+        operation: &'static str,
+        method: &'static str,
+    ) -> Option<tracing::Span> {
+        self.trace_in_process_rpc.then(|| {
+            tracing::info_span!(
+                target: "openshell_driver_docker::otel_tracing",
+                "driver_rpc",
+                otel.name = operation,
+                otel.kind = "server",
+                otel.status_code = tracing::field::Empty,
+                rpc.system = "grpc",
+                rpc.service = "openshell.compute.v1.ComputeDriver",
+                rpc.method = method,
+                rpc.grpc.status_code = tracing::field::Empty,
+            )
+        })
+    }
+
+    async fn trace_rpc<T>(
+        &self,
+        operation: &'static str,
+        method: &'static str,
+        future: impl Future<Output = Result<T, Status>>,
+    ) -> Result<T, Status> {
+        use tracing::Instrument as _;
+
+        let Some(span) = self.in_process_rpc_span(operation, method) else {
+            return future.await;
+        };
+        let result = future.instrument(span.clone()).await;
+        match &result {
+            Ok(_) => {
+                span.record("rpc.grpc.status_code", tonic::Code::Ok as i32);
+            }
+            Err(status) => {
+                openshell_otel::mark_error(&span);
+                span.record("rpc.grpc.status_code", status.code() as i32);
+            }
+        }
+        result
+    }
+}
+
 /// Return the first responsive local Docker API socket.
 #[must_use]
 pub fn detect_socket() -> Option<PathBuf> {
@@ -1536,12 +1653,188 @@ impl DockerComputeDriver {
     }
 }
 
+#[tonic::async_trait]
+impl ComputeDriver for ComputeDriverService {
+    type WatchSandboxesStream = WatchStream;
+
+    async fn authenticate_sandbox(
+        &self,
+        request: Request<openshell_core::proto::compute::v1::AuthenticateSandboxRequest>,
+    ) -> Result<Response<openshell_core::proto::compute::v1::AuthenticateSandboxResponse>, Status>
+    {
+        self.trace_rpc(
+            "driver.authenticate_sandbox",
+            "authenticate_sandbox",
+            ComputeDriver::authenticate_sandbox(&self.driver, request),
+        )
+        .await
+    }
+
+    async fn get_capabilities(
+        &self,
+        request: Request<GetCapabilitiesRequest>,
+    ) -> Result<Response<GetCapabilitiesResponse>, Status> {
+        self.trace_rpc(
+            "driver.get_capabilities",
+            "get_capabilities",
+            ComputeDriver::get_capabilities(&self.driver, request),
+        )
+        .await
+    }
+
+    async fn get_gateway_listener_requirements(
+        &self,
+        request: Request<GetGatewayListenerRequirementsRequest>,
+    ) -> Result<Response<GetGatewayListenerRequirementsResponse>, Status> {
+        self.trace_rpc(
+            "driver.get_gateway_listener_requirements",
+            "get_gateway_listener_requirements",
+            ComputeDriver::get_gateway_listener_requirements(&self.driver, request),
+        )
+        .await
+    }
+
+    async fn validate_sandbox_create(
+        &self,
+        request: Request<ValidateSandboxCreateRequest>,
+    ) -> Result<Response<ValidateSandboxCreateResponse>, Status> {
+        self.trace_rpc(
+            "driver.validate_sandbox_create",
+            "validate_sandbox_create",
+            ComputeDriver::validate_sandbox_create(&self.driver, request),
+        )
+        .await
+    }
+
+    async fn get_sandbox(
+        &self,
+        request: Request<GetSandboxRequest>,
+    ) -> Result<Response<GetSandboxResponse>, Status> {
+        self.trace_rpc(
+            "driver.get_sandbox",
+            "get_sandbox",
+            ComputeDriver::get_sandbox(&self.driver, request),
+        )
+        .await
+    }
+
+    async fn list_sandboxes(
+        &self,
+        request: Request<ListSandboxesRequest>,
+    ) -> Result<Response<ListSandboxesResponse>, Status> {
+        self.trace_rpc(
+            "driver.list_sandboxes",
+            "list_sandboxes",
+            ComputeDriver::list_sandboxes(&self.driver, request),
+        )
+        .await
+    }
+
+    async fn create_sandbox(
+        &self,
+        request: Request<CreateSandboxRequest>,
+    ) -> Result<Response<CreateSandboxResponse>, Status> {
+        self.trace_rpc(
+            "driver.create_sandbox",
+            "create_sandbox",
+            ComputeDriver::create_sandbox(&self.driver, request),
+        )
+        .await
+    }
+
+    async fn stop_sandbox(
+        &self,
+        request: Request<StopSandboxRequest>,
+    ) -> Result<Response<StopSandboxResponse>, Status> {
+        self.trace_rpc(
+            "driver.stop_sandbox",
+            "stop_sandbox",
+            ComputeDriver::stop_sandbox(&self.driver, request),
+        )
+        .await
+    }
+
+    async fn start_sandbox(
+        &self,
+        request: Request<StartSandboxRequest>,
+    ) -> Result<Response<StartSandboxResponse>, Status> {
+        self.trace_rpc(
+            "driver.start_sandbox",
+            "start_sandbox",
+            ComputeDriver::start_sandbox(&self.driver, request),
+        )
+        .await
+    }
+
+    async fn delete_sandbox(
+        &self,
+        request: Request<DeleteSandboxRequest>,
+    ) -> Result<Response<DeleteSandboxResponse>, Status> {
+        self.trace_rpc(
+            "driver.delete_sandbox",
+            "delete_sandbox",
+            ComputeDriver::delete_sandbox(&self.driver, request),
+        )
+        .await
+    }
+
+    async fn watch_sandboxes(
+        &self,
+        request: Request<WatchSandboxesRequest>,
+    ) -> Result<Response<Self::WatchSandboxesStream>, Status> {
+        use tracing::Instrument as _;
+
+        let create_stream = ComputeDriver::watch_sandboxes(&self.driver, request);
+        let Some(span) = self.in_process_rpc_span("driver.watch_sandboxes", "watch_sandboxes")
+        else {
+            return create_stream.await;
+        };
+        match create_stream.instrument(span.clone()).await {
+            Ok(response) => Ok(Response::new(Box::pin(TracedWatchStream {
+                inner: response.into_inner(),
+                span,
+                finished: false,
+            }))),
+            Err(status) => {
+                openshell_otel::mark_error(&span);
+                span.record("rpc.grpc.status_code", status.code() as i32);
+                Err(status)
+            }
+        }
+    }
+
+    async fn ensure_workspace(
+        &self,
+        request: Request<EnsureWorkspaceRequest>,
+    ) -> Result<Response<EnsureWorkspaceResponse>, Status> {
+        self.trace_rpc(
+            "driver.ensure_workspace",
+            "ensure_workspace",
+            ComputeDriver::ensure_workspace(&self.driver, request),
+        )
+        .await
+    }
+
+    async fn delete_workspace(
+        &self,
+        request: Request<DeleteWorkspaceRequest>,
+    ) -> Result<Response<DeleteWorkspaceResponse>, Status> {
+        self.trace_rpc(
+            "driver.delete_workspace",
+            "delete_workspace",
+            ComputeDriver::delete_workspace(&self.driver, request),
+        )
+        .await
+    }
+}
+
 #[tonic::async_trait]
 impl ComputeDriver for DockerComputeDriver {
     async fn authenticate_sandbox(
         &self,
         _request: Request<openshell_core::proto::compute::v1::AuthenticateSandboxRequest>,
-    ) -> Result<Response<openshell_core::proto::compute::v1::AuthenticateSandboxResponse>, Status> {
+    ) -> Result<Response<openshell_core::proto::compute::v1::AuthenticateSandboxResponse>, Status>
+    {
         Err(Status::unimplemented(
             "docker does not authenticate sandbox credentials",
         ))