cf_system_sdk_directory/grpc/

client.rs

//! gRPC client implementation of Directory API
//!
//! This client allows remote gears to discover and resolve services via gRPC.

use anyhow::Result;
use async_trait::async_trait;
use tonic::service::interceptor::InterceptedService;
use tonic::transport::Channel;

use crate::api::{
    DirectoryClient, DirectoryInvalidArgument, DirectoryNotFound, RegisterInstanceInfo,
    ServiceEndpoint, ServiceInstanceInfo,
};
use toolkit_transport_grpc::InternalAuthInterceptor;
use toolkit_transport_grpc::client::{GrpcClientConfig, connect_with_retry};

use crate::{
    DeregisterInstanceRequest, DirectoryServiceClient, GetOpenApiSpecRequest, GrpcServiceEndpoint,
    HeartbeatRequest, InstanceInfo, ListAllInstancesRequest, ListInstancesRequest,
    RegisterInstanceRequest, ResolveGrpcServiceRequest, ResolveRestServiceRequest,
};

/// The directory channel wrapped with the platform-plane
/// [`InternalAuthInterceptor`], which attaches the gear's internal token
/// (`x-toolkit-internal-token`) to every outbound system call. A
/// [`disabled`](InternalAuthInterceptor::disabled) interceptor attaches
/// nothing (Profile 1 / no platform-plane credential).
type AuthedChannel = InterceptedService<Channel, InternalAuthInterceptor>;

/// Map a lookup RPC's `tonic::Status` onto the directory's typed sentinels.
///
/// The status code is the only thing distinguishing "this name is not
/// registered" from "the directory is unreachable", and stringifying the status
/// throws it away. Callers downcast the result: `DirectoryEndpointResolver`
/// turns `DirectoryNotFound` into `Ok(None)` (a provider that has not come up
/// yet — routine during startup) and anything else into a real error.
fn lookup_error(resource: &str, status: &tonic::Status) -> anyhow::Error {
    match status.code() {
        tonic::Code::NotFound => DirectoryNotFound::new(resource.to_owned()).into(),
        tonic::Code::InvalidArgument => {
            DirectoryInvalidArgument::new(status.message().to_owned()).into()
        }
        code => anyhow::anyhow!(
            "directory lookup for {resource} failed: gRPC {code:?}: {}",
            status.message()
        ),
    }
}

/// Map a mutating RPC's `tonic::Status`, preserving the code in the message.
///
/// These calls have no not-found semantics worth typing, but a bare
/// `"gRPC call failed"` hides whether the directory was unreachable
/// (`Unavailable`) or rejected the request (`Internal`).
fn call_error(op: &str, status: &tonic::Status) -> anyhow::Error {
    anyhow::anyhow!(
        "directory {op} failed: gRPC {:?}: {}",
        status.code(),
        status.message()
    )
}

/// gRPC client for Directory API
///
/// This client connects to a remote `DirectoryService` via gRPC and provides
/// typed access to service discovery functionality. It includes:
/// - Configurable timeouts and retries via transport stack
/// - Automatic proto ↔ domain type conversions
/// - Distributed tracing and metrics
/// - Platform-plane credential attachment via an [`InternalAuthInterceptor`]
///   (defaults to attaching nothing; supply one via the `*_with_interceptor`
///   constructors for Profile-3 / shared-secret deployments)
pub struct DirectoryGrpcClient {
    inner: DirectoryServiceClient<AuthedChannel>,
}

impl DirectoryGrpcClient {
    /// Connect to a directory service using default configuration with retries.
    ///
    /// Uses exponential backoff retry logic for reliable connection establishment.
    /// This is the recommended method for `OoP` gears connecting to the master host.
    ///
    /// # Errors
    /// It will return an error when it fails
    pub async fn connect(uri: impl Into<String>) -> Result<Self> {
        let cfg = GrpcClientConfig::new("directory");
        Self::connect_with_retry(uri, &cfg).await
    }

    /// Connect with default configuration + retries, attaching `interceptor`'s
    /// platform-plane credential to every outbound call.
    ///
    /// This is the Profile-3 / shared-secret entry point: the interceptor is
    /// typically built from a
    /// [`ServiceAccountTokenReader`](toolkit_transport_grpc::ServiceAccountTokenReader)
    /// (rotating SA token) or
    /// [`InternalAuthInterceptor::from_token`] (static shared secret).
    ///
    /// # Errors
    /// It will return an error when it fails
    pub async fn connect_with_interceptor(
        uri: impl Into<String>,
        interceptor: InternalAuthInterceptor,
    ) -> Result<Self> {
        let cfg = GrpcClientConfig::new("directory");
        let channel: Channel = connect_with_retry(uri, &cfg).await?;
        Ok(Self::from_channel_with_interceptor(channel, interceptor))
    }

    /// Connect to a directory service with custom configuration and retry logic.
    ///
    /// Uses exponential backoff based on `cfg.max_retries`, `cfg.base_backoff`,
    /// and `cfg.max_backoff` settings.
    ///
    /// # Errors
    /// It will return an error when it fails
    pub async fn connect_with_retry(
        uri: impl Into<String>,
        cfg: &GrpcClientConfig,
    ) -> Result<Self> {
        let channel: Channel = connect_with_retry(uri, cfg).await?;
        Ok(Self::from_channel(channel))
    }

    /// Connect to a directory service without retry logic.
    ///
    /// This method attempts a single connection. Use `connect` or `connect_with_retry`
    /// for production scenarios where the directory service may not be immediately available.
    ///
    /// # Errors
    /// It will return an error when it fails
    pub async fn connect_no_retry(uri: impl Into<String>, cfg: &GrpcClientConfig) -> Result<Self> {
        let uri_string = uri.into();

        // Create endpoint with timeouts from config
        let endpoint = tonic::transport::Endpoint::from_shared(uri_string)?
            .connect_timeout(cfg.connect_timeout)
            .timeout(cfg.rpc_timeout);

        // Connect to the service
        let channel = endpoint.connect().await?;

        if cfg.enable_tracing {
            tracing::debug!(
                service_name = cfg.service_name,
                connect_timeout_ms = cfg.connect_timeout.as_millis(),
                rpc_timeout_ms = cfg.rpc_timeout.as_millis(),
                "directory gRPC client connected"
            );
        }

        Ok(Self::from_channel(channel))
    }

    /// Create from an existing channel (useful for testing or custom setup).
    ///
    /// Attaches no platform-plane credential; use
    /// [`from_channel_with_interceptor`](Self::from_channel_with_interceptor)
    /// to attach one.
    #[must_use]
    pub fn from_channel(channel: Channel) -> Self {
        Self::from_channel_with_interceptor(channel, InternalAuthInterceptor::disabled())
    }

    /// Create from an existing channel, attaching `interceptor`'s platform-plane
    /// credential to every outbound call.
    #[must_use]
    pub fn from_channel_with_interceptor(
        channel: Channel,
        interceptor: InternalAuthInterceptor,
    ) -> Self {
        Self {
            inner: DirectoryServiceClient::with_interceptor(channel, interceptor),
        }
    }
}

#[async_trait]
impl DirectoryClient for DirectoryGrpcClient {
    async fn resolve_grpc_service(&self, service_name: &str) -> Result<ServiceEndpoint> {
        let mut client = self.inner.clone();
        let request = tonic::Request::new(ResolveGrpcServiceRequest {
            service_name: service_name.to_owned(),
        });

        let response = client
            .resolve_grpc_service(request)
            .await
            .map_err(|e| lookup_error(&format!("service {service_name}"), &e))?;

        let proto_response = response.into_inner();
        Ok(ServiceEndpoint::new(proto_response.endpoint_uri))
    }

    async fn resolve_rest_service(&self, gear_name: &str) -> Result<ServiceEndpoint> {
        let mut client = self.inner.clone();
        let request = tonic::Request::new(ResolveRestServiceRequest {
            gear_name: gear_name.to_owned(),
        });

        let response = client
            .resolve_rest_service(request)
            .await
            .map_err(|e| lookup_error(&format!("gear {gear_name}"), &e))?;

        let proto_response = response.into_inner();
        Ok(ServiceEndpoint::new(proto_response.endpoint_uri))
    }

    async fn get_openapi_spec(&self, gear_name: &str) -> Result<String> {
        let mut client = self.inner.clone();
        let request = tonic::Request::new(GetOpenApiSpecRequest {
            gear_name: gear_name.to_owned(),
        });

        let response = client
            .get_open_api_spec(request)
            .await
            .map_err(|e| lookup_error(&format!("openapi spec for gear {gear_name}"), &e))?;

        Ok(response.into_inner().openapi_spec)
    }

    async fn list_instances(&self, gear: &str) -> Result<Vec<ServiceInstanceInfo>> {
        let mut client = self.inner.clone();
        let request = tonic::Request::new(ListInstancesRequest {
            gear_name: gear.to_owned(),
        });

        let response = client
            .list_instances(request)
            .await
            .map_err(|e| lookup_error(&format!("instances of gear {gear}"), &e))?;

        let instances = response
            .into_inner()
            .instances
            .into_iter()
            .map(proto_instance_to_domain)
            .collect();

        Ok(instances)
    }

    async fn list_all_instances(&self) -> Result<Vec<ServiceInstanceInfo>> {
        let mut client = self.inner.clone();
        let response = client
            .list_all_instances(tonic::Request::new(ListAllInstancesRequest {}))
            .await
            .map_err(|e| lookup_error("all instances", &e))?;

        let instances = response
            .into_inner()
            .instances
            .into_iter()
            .map(|proto| {
                let mut info = proto_instance_to_domain(proto);
                info.openapi_spec = None;
                info
            })
            .collect();

        Ok(instances)
    }

    async fn register_instance(&self, info: RegisterInstanceInfo) -> Result<()> {
        let mut client = self.inner.clone();

        // Convert gRPC service endpoints
        let grpc_services = info
            .grpc_services
            .into_iter()
            .map(|(name, ep)| GrpcServiceEndpoint {
                service_name: name,
                endpoint_uri: ep.uri,
            })
            .collect();

        let req = RegisterInstanceRequest {
            gear_name: info.gear,
            instance_id: info.instance_id,
            grpc_services,
            version: info.version.unwrap_or_default(),
            rest_endpoint_uri: info.rest_endpoint.map(|ep| ep.uri),
            openapi_spec: info.openapi_spec,
        };

        client
            .register_instance(tonic::Request::new(req))
            .await
            .map_err(|e| call_error("register_instance", &e))?;

        Ok(())
    }

    async fn deregister_instance(&self, gear: &str, instance_id: &str) -> Result<()> {
        let mut client = self.inner.clone();

        let req = DeregisterInstanceRequest {
            gear_name: gear.to_owned(),
            instance_id: instance_id.to_owned(),
        };

        client
            .deregister_instance(tonic::Request::new(req))
            .await
            .map_err(|e| call_error("deregister_instance", &e))?;

        Ok(())
    }

    async fn send_heartbeat(&self, gear: &str, instance_id: &str) -> Result<()> {
        let mut client = self.inner.clone();

        let req = HeartbeatRequest {
            gear_name: gear.to_owned(),
            instance_id: instance_id.to_owned(),
        };

        client
            .heartbeat(tonic::Request::new(req))
            .await
            .map_err(|e| call_error("heartbeat", &e))?;

        Ok(())
    }
}

/// Convert a proto `InstanceInfo` into the domain [`ServiceInstanceInfo`].
fn proto_instance_to_domain(proto: InstanceInfo) -> ServiceInstanceInfo {
    ServiceInstanceInfo {
        gear: proto.gear_name,
        instance_id: proto.instance_id,
        endpoint: ServiceEndpoint::new(proto.endpoint_uri),
        version: if proto.version.is_empty() {
            None
        } else {
            Some(proto.version)
        },
        rest_endpoint: proto.rest_endpoint_uri.map(ServiceEndpoint::new),
        openapi_spec: proto.openapi_spec,
        openapi_spec_hash: proto.openapi_spec_hash,
        // The list-instances proto response does not break the instance down
        // per gRPC service, so nothing to carry back over the OoP directory
        // transport. The in-process `LocalDirectoryClient` populates this from
        // the live `GearInstance`.
        grpc_services: Vec::new(),
    }
}

#[cfg(test)]
#[cfg_attr(coverage_nightly, coverage(off))]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_grpc_client_can_be_constructed() {
        // Smoke test to ensure types compile and connect
        let endpoint = tonic::transport::Endpoint::from_static("http://[::1]:50051");

        // We can't actually connect without a server, but we can construct the client type
        // This ensures the API is correct
        let channel_result = endpoint.connect().await;

        // It's expected to fail since there's no server, but if it does somehow succeed:
        if let Ok(channel) = channel_result {
            let _client = DirectoryGrpcClient::from_channel(channel);
        }
    }

    #[tokio::test]
    async fn from_channel_constructs_without_connecting() {
        // `connect_lazy` yields a Channel without a live server, so both the
        // default (no-credential) and interceptor-bearing constructors can be
        // exercised offline.
        let channel = Channel::from_static("http://[::1]:50051").connect_lazy();
        let _default = DirectoryGrpcClient::from_channel(channel.clone());
        let _authed = DirectoryGrpcClient::from_channel_with_interceptor(
            channel,
            InternalAuthInterceptor::disabled(),
        );
    }

    #[test]
    fn proto_instance_maps_all_fields_to_domain() {
        let proto = InstanceInfo {
            gear_name: "calc".to_owned(),
            instance_id: "calc-1".to_owned(),
            endpoint_uri: "http://calc:8080".to_owned(),
            version: "1.2.3".to_owned(),
            rest_endpoint_uri: Some("http://calc:8080".to_owned()),
            openapi_spec: Some("{\"openapi\":\"3.1.0\"}".to_owned()),
            openapi_spec_hash: None,
        };
        let domain = proto_instance_to_domain(proto);
        assert_eq!(domain.gear, "calc");
        assert_eq!(domain.instance_id, "calc-1");
        assert_eq!(domain.endpoint.uri, "http://calc:8080");
        assert_eq!(domain.version.as_deref(), Some("1.2.3"));
        assert_eq!(
            domain.rest_endpoint.map(|e| e.uri),
            Some("http://calc:8080".to_owned())
        );
        assert!(domain.openapi_spec.is_some());
    }

    #[test]
    fn proto_instance_maps_empty_version_to_none() {
        let proto = InstanceInfo {
            gear_name: "worker".to_owned(),
            instance_id: "worker-1".to_owned(),
            endpoint_uri: "http://worker:7000".to_owned(),
            version: String::new(),
            rest_endpoint_uri: None,
            openapi_spec: None,
            openapi_spec_hash: None,
        };
        let domain = proto_instance_to_domain(proto);
        // An empty proto version string maps to `None` rather than an empty string.
        assert!(domain.version.is_none());
        assert!(domain.rest_endpoint.is_none());
        assert!(domain.openapi_spec.is_none());
    }
}