refactor: uniform client wrappers (#1614)

* refactor: better client wrapper for provisioner * refactor: better client wrapper for logger * refactor: better client wrapper for builder * refactor: better client wrapper for runtime * refactor: trimmed unused dependencies * refactor: fix optional deps * refactor: more optional dependency fixes
shuttle-hq · Feb 6, 2024 · 5b5eb69 · 5b5eb69
1 parent 8c6e931
commit 5b5eb69
Show file tree

Hide file tree

Showing 21 changed files with 197 additions and 327 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/cargo-shuttle/src/lib.rs b/cargo-shuttle/src/lib.rs
@@ -19,7 +19,6 @@ use args::{ConfirmationArgs, GenerateCommand};
 use clap_mangen::Man;
 
 use shuttle_common::{
-    claims::{ClaimService, InjectPropagation},
     constants::{
         API_URL_DEFAULT, DEFAULT_IDLE_MINUTES, EXECUTABLE_DIRNAME, SHUTTLE_CLI_DOCS_URL,
         SHUTTLE_GH_ISSUE_URL, SHUTTLE_IDLE_DOCS_URL, SHUTTLE_INSTALL_DOCS_URL, SHUTTLE_LOGIN_URL,
@@ -37,9 +36,8 @@ use shuttle_common::{
     },
     resource, semvers_are_compatible, ApiKey, LogItem, VersionInfo,
 };
-use shuttle_proto::runtime::{
-    runtime_client::RuntimeClient, LoadRequest, StartRequest, StopRequest,
-};
+use shuttle_proto::runtime;
+use shuttle_proto::runtime::{LoadRequest, StartRequest, StopRequest};
 use shuttle_service::runner;
 use shuttle_service::{
     builder::{build_workspace, BuiltService},
@@ -67,7 +65,6 @@ use tokio::io::{AsyncBufReadExt, BufReader};
 use tokio::process::Child;
 use tokio::task::JoinHandle;
 use tokio::time::{sleep, Duration};
-use tonic::transport::Channel;
 use tonic::Status;
 use tracing::{debug, error, trace, warn};
 use uuid::Uuid;
@@ -938,12 +935,7 @@ impl Shuttle {
         provisioner_server: &JoinHandle<Result<(), tonic::transport::Error>>,
         idx: u16,
         provisioner_port: u16,
-    ) -> Result<
-        Option<(
-            Child,
-            RuntimeClient<ClaimService<InjectPropagation<Channel>>>,
-        )>,
-    > {
+    ) -> Result<Option<(Child, runtime::Client)>> {
         let crate_directory = service.crate_directory();
         let secrets_path = if crate_directory.join("Secrets.dev.toml").exists() {
             crate_directory.join("Secrets.dev.toml")
@@ -1160,7 +1152,7 @@ impl Shuttle {
 
     async fn stop_runtime(
         runtime: &mut Child,
-        runtime_client: &mut RuntimeClient<ClaimService<InjectPropagation<Channel>>>,
+        runtime_client: &mut runtime::Client,
     ) -> Result<(), Status> {
         let stop_request = StopRequest {};
         trace!(?stop_request, "stopping service");
@@ -1178,14 +1170,8 @@ impl Shuttle {
     }
 
     async fn add_runtime_info(
-        runtime: Option<(
-            Child,
-            RuntimeClient<ClaimService<InjectPropagation<Channel>>>,
-        )>,
-        existing_runtimes: &mut Vec<(
-            Child,
-            RuntimeClient<ClaimService<InjectPropagation<Channel>>>,
-        )>,
+        runtime: Option<(Child, runtime::Client)>,
+        existing_runtimes: &mut Vec<(Child, runtime::Client)>,
         extra_servers: &[&JoinHandle<Result<(), tonic::transport::Error>>],
     ) -> Result<(), Status> {
         match runtime {
@@ -1269,10 +1255,7 @@ impl Shuttle {
                 .expect("Can not get the SIGINT signal receptor");
 
         // Start all the services.
-        let mut runtimes: Vec<(
-            Child,
-            RuntimeClient<ClaimService<InjectPropagation<Channel>>>,
-        )> = Vec::new();
+        let mut runtimes: Vec<(Child, runtime::Client)> = Vec::new();
 
         Shuttle::find_available_port(&mut run_args, services.len());
 
@@ -1423,10 +1406,7 @@ impl Shuttle {
         let (provisioner_server, provisioner_port) = Shuttle::setup_local_provisioner().await?;
 
         // Start all the services.
-        let mut runtimes: Vec<(
-            Child,
-            RuntimeClient<ClaimService<InjectPropagation<Channel>>>,
-        )> = Vec::new();
+        let mut runtimes: Vec<(Child, runtime::Client)> = Vec::new();
 
         Shuttle::find_available_port(&mut run_args, services.len());
 

diff --git a/common-tests/src/builder.rs b/common-tests/src/builder.rs
@@ -4,21 +4,13 @@ use std::{
 };
 
 use portpicker::pick_unused_port;
-use shuttle_common::claims::{ClaimLayer, InjectPropagationLayer};
 use shuttle_proto::builder::{
-    builder_client::BuilderClient,
+    self,
     builder_server::{Builder, BuilderServer},
 };
-use tonic::transport::{Endpoint, Server};
-use tower::ServiceBuilder;
+use tonic::transport::Server;
 
-pub async fn get_mocked_builder_client(
-    builder: impl Builder,
-) -> BuilderClient<
-    shuttle_common::claims::ClaimService<
-        shuttle_common::claims::InjectPropagation<tonic::transport::Channel>,
-    >,
-> {
+pub async fn get_mocked_builder_client(builder: impl Builder) -> builder::Client {
     let builder_addr = SocketAddr::new(Ipv4Addr::LOCALHOST.into(), pick_unused_port().unwrap());
     let builder_uri = format!("http://{}", builder_addr);
     tokio::spawn(async move {
@@ -31,16 +23,5 @@ pub async fn get_mocked_builder_client(
     // Wait for the builder server to start before creating a client.
     tokio::time::sleep(Duration::from_millis(200)).await;
 
-    let channel = Endpoint::try_from(builder_uri.to_string())
-        .unwrap()
-        .connect()
-        .await
-        .expect("failed to connect to builder");
-
-    let channel = ServiceBuilder::new()
-        .layer(ClaimLayer)
-        .layer(InjectPropagationLayer)
-        .service(channel);
-
-    BuilderClient::new(channel)
+    builder::get_client(builder_uri.parse().unwrap()).await
 }
diff --git a/common-tests/src/logger.rs b/common-tests/src/logger.rs
@@ -4,20 +4,14 @@ use std::{
 };
 
 use portpicker::pick_unused_port;
-use shuttle_common::claims::{ClaimLayer, InjectPropagationLayer};
 use shuttle_proto::logger::{
-    logger_client::LoggerClient,
+    self,
     logger_server::{Logger, LoggerServer},
     LogLine, LogsRequest, LogsResponse, StoreLogsRequest, StoreLogsResponse,
 };
 use tokio::sync::mpsc;
 use tokio_stream::wrappers::ReceiverStream;
-use tonic::{
-    async_trait,
-    transport::{Endpoint, Server},
-    Request, Response, Status,
-};
-use tower::ServiceBuilder;
+use tonic::{async_trait, transport::Server, Request, Response, Status};
 
 pub struct MockedLogger;
 
@@ -47,13 +41,7 @@ impl Logger for MockedLogger {
     }
 }
 
-pub async fn get_mocked_logger_client(
-    logger: impl Logger,
-) -> LoggerClient<
-    shuttle_common::claims::ClaimService<
-        shuttle_common::claims::InjectPropagation<tonic::transport::Channel>,
-    >,
-> {
+pub async fn get_mocked_logger_client(logger: impl Logger) -> logger::Client {
     let logger_addr = SocketAddr::new(Ipv4Addr::LOCALHOST.into(), pick_unused_port().unwrap());
     let logger_uri = format!("http://{}", logger_addr);
     tokio::spawn(async move {
@@ -66,16 +54,5 @@ pub async fn get_mocked_logger_client(
     // Wait for the logger server to start before creating a client.
     tokio::time::sleep(Duration::from_millis(200)).await;
 
-    let channel = Endpoint::try_from(logger_uri.to_string())
-        .unwrap()
-        .connect()
-        .await
-        .expect("failed to connect to logger");
-
-    let channel = ServiceBuilder::new()
-        .layer(ClaimLayer)
-        .layer(InjectPropagationLayer)
-        .service(channel);
-
-    LoggerClient::new(channel)
+    logger::get_client(logger_uri.parse().unwrap()).await
 }
diff --git a/deployer/src/args.rs b/deployer/src/args.rs
@@ -4,7 +4,6 @@ use clap::Parser;
 use fqdn::FQDN;
 use hyper::Uri;
 use shuttle_common::models::project::ProjectName;
-use tonic::transport::Endpoint;
 
 /// Program to handle the deploys for a single project
 /// Handling includes, building, testing, and running each service
@@ -21,7 +20,7 @@ pub struct Args {
 
     /// Address to connect to the logger service
     #[clap(long, default_value = "http://logger:8000")]
-    pub logger_uri: Endpoint,
+    pub logger_uri: Uri,
 
     /// FQDN where the proxy can be reached at
     #[clap(long)]
@@ -61,7 +60,7 @@ pub struct Args {
 
     /// Address to reach the builder service at
     #[clap(long, default_value = "http://builder:8000")]
-    pub builder_uri: Endpoint,
+    pub builder_uri: Uri,
 
     /// Uri to folder to store all artifacts
     #[clap(long, default_value = "/tmp")]

diff --git a/deployer/src/deployment/mod.rs b/deployer/src/deployment/mod.rs
@@ -4,7 +4,7 @@ use std::{
 };
 
 use shuttle_common::log::LogRecorder;
-use shuttle_proto::{builder::builder_client::BuilderClient, logger::logger_client::LoggerClient};
+use shuttle_proto::{builder, logger};
 use tokio::{
     sync::{mpsc, Mutex},
     task::JoinSet,
@@ -31,26 +31,14 @@ const RUN_BUFFER_SIZE: usize = 100;
 
 pub struct DeploymentManagerBuilder<LR, ADG, DU, RM, QC> {
     build_log_recorder: Option<LR>,
-    logs_fetcher: Option<
-        LoggerClient<
-            shuttle_common::claims::ClaimService<
-                shuttle_common::claims::InjectPropagation<tonic::transport::Channel>,
-            >,
-        >,
-    >,
+    logs_fetcher: Option<logger::Client>,
     active_deployment_getter: Option<ADG>,
     artifacts_path: Option<PathBuf>,
     runtime_manager: Option<Arc<Mutex<RuntimeManager>>>,
     deployment_updater: Option<DU>,
     resource_manager: Option<RM>,
     queue_client: Option<QC>,
-    builder_client: Option<
-        BuilderClient<
-            shuttle_common::claims::ClaimService<
-                shuttle_common::claims::InjectPropagation<tonic::transport::Channel>,
-            >,
-        >,
-    >,
+    builder_client: Option<builder::Client>,
 }
 
 impl<LR, ADG, DU, RM, QC> DeploymentManagerBuilder<LR, ADG, DU, RM, QC>
@@ -67,29 +55,13 @@ where
         self
     }
 
-    pub fn log_fetcher(
-        mut self,
-        logs_fetcher: LoggerClient<
-            shuttle_common::claims::ClaimService<
-                shuttle_common::claims::InjectPropagation<tonic::transport::Channel>,
-            >,
-        >,
-    ) -> Self {
+    pub fn log_fetcher(mut self, logs_fetcher: logger::Client) -> Self {
         self.logs_fetcher = Some(logs_fetcher);
 
         self
     }
 
-    pub fn builder_client(
-        mut self,
-        builder_client: Option<
-            BuilderClient<
-                shuttle_common::claims::ClaimService<
-                    shuttle_common::claims::InjectPropagation<tonic::transport::Channel>,
-                >,
-            >,
-        >,
-    ) -> Self {
+    pub fn builder_client(mut self, builder_client: Option<builder::Client>) -> Self {
         self.builder_client = builder_client;
 
         self
@@ -195,11 +167,7 @@ pub struct DeploymentManager {
     queue_send: QueueSender,
     run_send: RunSender,
     runtime_manager: Arc<Mutex<RuntimeManager>>,
-    logs_fetcher: LoggerClient<
-        shuttle_common::claims::ClaimService<
-            shuttle_common::claims::InjectPropagation<tonic::transport::Channel>,
-        >,
-    >,
+    logs_fetcher: logger::Client,
     _join_set: Arc<Mutex<JoinSet<()>>>,
     builds_path: PathBuf,
 }
@@ -259,13 +227,7 @@ impl DeploymentManager {
         self.builds_path.as_path()
     }
 
-    pub fn logs_fetcher(
-        &self,
-    ) -> &LoggerClient<
-        shuttle_common::claims::ClaimService<
-            shuttle_common::claims::InjectPropagation<tonic::transport::Channel>,
-        >,
-    > {
+    pub fn logs_fetcher(&self) -> &logger::Client {
         &self.logs_fetcher
     }
 }

diff --git a/deployer/src/deployment/queue.rs b/deployer/src/deployment/queue.rs
@@ -14,8 +14,7 @@ use shuttle_common::{
     log::LogRecorder,
     LogItem,
 };
-use shuttle_proto::builder::builder_client::BuilderClient;
-use shuttle_proto::builder::BuildRequest;
+use shuttle_proto::builder::{self, BuildRequest};
 use shuttle_service::builder::{build_workspace, BuiltService};
 use tar::Archive;
 use tokio::{
@@ -42,13 +41,7 @@ pub async fn task(
     deployment_updater: impl DeploymentUpdater,
     log_recorder: impl LogRecorder,
     queue_client: impl BuildQueueClient,
-    builder_client: Option<
-        BuilderClient<
-            shuttle_common::claims::ClaimService<
-                shuttle_common::claims::InjectPropagation<tonic::transport::Channel>,
-            >,
-        >,
-    >,
+    builder_client: Option<builder::Client>,
     builds_path: PathBuf,
 ) {
     info!("Queue task started");