cloudlysis/runner/src/gateway/mod.rs

pub const TENANT_ID_METADATA_KEY: &str = shared::HEADER_X_TENANT_ID;
pub const CORRELATION_ID_METADATA_KEY: &str = shared::HEADER_X_CORRELATION_ID;
pub const TRACEPARENT_METADATA_KEY: &str = shared::HEADER_TRACEPARENT;

pub mod proto {
    tonic::include_proto!("aggregate.gateway.v1");
}

#[derive(Clone)]
pub struct GatewayClient {
    inner: proto::command_service_client::CommandServiceClient<tonic::transport::Channel>,
}

impl std::fmt::Debug for GatewayClient {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("GatewayClient").finish_non_exhaustive()
    }
}

impl GatewayClient {
    pub async fn connect(endpoint: &str) -> Result<Self, crate::types::RunnerError> {
        let channel = tonic::transport::Endpoint::from_shared(endpoint.to_string())
            .map_err(|e| crate::types::RunnerError::RuntimeError(e.to_string()))?
            .connect()
            .await
            .map_err(|e| crate::types::RunnerError::RuntimeError(e.to_string()))?;
        let inner = proto::command_service_client::CommandServiceClient::new(channel);
        Ok(Self { inner })
    }

    pub async fn submit_command(
        &mut self,
        request: proto::SubmitCommandRequest,
    ) -> Result<proto::SubmitCommandResponse, tonic::Status> {
        let mut grpc_request = tonic::Request::new(request);

        let tenant_id = grpc_request.get_ref().tenant_id.as_str();
        if !tenant_id.is_empty() {
            let value = tonic::metadata::MetadataValue::try_from(tenant_id).map_err(|e| {
                tonic::Status::invalid_argument(format!("invalid tenant_id metadata: {}", e))
            })?;
            grpc_request
                .metadata_mut()
                .insert(TENANT_ID_METADATA_KEY, value);
        }

        let correlation_id = grpc_request
            .get_ref()
            .metadata
            .get(shared::HEADER_X_CORRELATION_ID)
            .or_else(|| grpc_request.get_ref().metadata.get("correlation_id"))
            .map(|s| s.trim())
            .filter(|s| !s.is_empty())
            .map(|s| s.to_string());
        if let Some(correlation_id) = correlation_id {
            let value =
                tonic::metadata::MetadataValue::try_from(correlation_id.as_str()).map_err(|e| {
                    tonic::Status::invalid_argument(format!(
                        "invalid correlation_id metadata: {}",
                        e
                    ))
                })?;
            grpc_request
                .metadata_mut()
                .insert(CORRELATION_ID_METADATA_KEY, value);
        }

        let traceparent = grpc_request
            .get_ref()
            .metadata
            .get(shared::HEADER_TRACEPARENT)
            .map(|s| s.trim())
            .filter(|s| !s.is_empty())
            .map(|s| s.to_string())
            .or_else(|| {
                grpc_request
                    .get_ref()
                    .metadata
                    .get("trace_id")
                    .map(|s| s.trim())
                    .and_then(|trace_id| {
                        shared::traceparent_from_trace_id(&shared::TraceId::new(trace_id))
                    })
            });
        if let Some(traceparent) = traceparent {
            let value =
                tonic::metadata::MetadataValue::try_from(traceparent.as_str()).map_err(|e| {
                    tonic::Status::invalid_argument(format!("invalid traceparent metadata: {}", e))
                })?;
            grpc_request
                .metadata_mut()
                .insert(TRACEPARENT_METADATA_KEY, value);
        }

        let resp = self.inner.submit_command(grpc_request).await?;
        Ok(resp.into_inner())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn traceparent_is_derived_from_trace_id_when_present() {
        let req = proto::SubmitCommandRequest {
            tenant_id: "t1".to_string(),
            command_id: "c1".to_string(),
            aggregate_id: "a1".to_string(),
            aggregate_type: "User".to_string(),
            payload_json: "{}".to_string(),
            metadata: std::collections::HashMap::from([(
                "trace_id".to_string(),
                "0123456789abcdef0123456789abcdef".to_string(),
            )]),
        };
        let trace_id = req.metadata.get("trace_id").unwrap().as_str();
        let span_id = uuid::Uuid::new_v4().simple().to_string()[..16].to_string();
        let traceparent = format!("00-{trace_id}-{span_id}-01");
        assert!(traceparent.starts_with("00-0123456789abcdef0123456789abcdef-"));
        assert!(traceparent.ends_with("-01"));
    }

    #[tokio::test]
    async fn submit_command_propagates_correlation_and_traceparent_metadata_when_present() {
        use proto::command_service_server::CommandService;

        #[derive(Default)]
        struct Upstream;

        #[tonic::async_trait]
        impl CommandService for Upstream {
            async fn submit_command(
                &self,
                request: tonic::Request<proto::SubmitCommandRequest>,
            ) -> Result<tonic::Response<proto::SubmitCommandResponse>, tonic::Status> {
                let correlation = request
                    .metadata()
                    .get("x-correlation-id")
                    .and_then(|v| v.to_str().ok())
                    .unwrap_or("");
                if correlation != "corr-1" {
                    return Err(tonic::Status::failed_precondition("missing correlation"));
                }

                let traceparent = request
                    .metadata()
                    .get("traceparent")
                    .and_then(|v| v.to_str().ok())
                    .unwrap_or("");
                if traceparent != "00-0123456789abcdef0123456789abcdef-1111111111111111-01" {
                    return Err(tonic::Status::failed_precondition("missing traceparent"));
                }

                Ok(tonic::Response::new(proto::SubmitCommandResponse {
                    events: vec![],
                }))
            }
        }

        let upstream_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let upstream_addr = upstream_listener.local_addr().unwrap();
        drop(upstream_listener);
        tokio::spawn(async move {
            tonic::transport::Server::builder()
                .add_service(proto::command_service_server::CommandServiceServer::new(
                    Upstream,
                ))
                .serve(upstream_addr)
                .await
                .unwrap();
        });
        tokio::time::sleep(std::time::Duration::from_millis(50)).await;

        let mut client = GatewayClient::connect(&format!("http://{}", upstream_addr))
            .await
            .unwrap();

        let req = proto::SubmitCommandRequest {
            tenant_id: "t1".to_string(),
            command_id: "c1".to_string(),
            aggregate_id: "a1".to_string(),
            aggregate_type: "User".to_string(),
            payload_json: "{}".to_string(),
            metadata: std::collections::HashMap::from([
                ("correlation_id".to_string(), "corr-1".to_string()),
                (
                    "traceparent".to_string(),
                    "00-0123456789abcdef0123456789abcdef-1111111111111111-01".to_string(),
                ),
            ]),
        };

        client.submit_command(req).await.unwrap();
    }
}