basic cluster transport protocol working

Author: teskje

Cargo.lock

@@ -20,7 +20,6 @@ use hyper_util::rt::TokioIo;
 use mz_build_info::{BuildInfo, build_info};
 use mz_cloud_resources::AwsExternalIdPrefix;
 use mz_compute::server::ComputeInstanceContext;
-use mz_compute_client::service::proto_compute_server::ProtoComputeServer;
 use mz_http_util::DynamicFilterTarget;
 use mz_orchestrator_tracing::{StaticTracingConfig, TracingCliArgs};
 use mz_ore::cli::{self, CliConfig};
@@ -32,10 +31,9 @@ use mz_persist_client::cache::PersistClientCache;
 use mz_persist_client::cfg::PersistConfig;
 use mz_persist_client::rpc::{GrpcPubSubClient, PersistPubSubClient, PersistPubSubClientConfig};
 use mz_service::emit_boot_diagnostics;
-use mz_service::grpc::{GrpcServer, GrpcServerMetrics, MAX_GRPC_MESSAGE_SIZE};
 use mz_service::secrets::SecretsReaderCliArgs;
+use mz_service::transport;
 use mz_storage::storage_state::StorageInstanceContext;
-use mz_storage_client::client::proto_storage_server::ProtoStorageServer;
 use mz_storage_types::connections::ConnectionContext;
 use mz_txn_wal::operator::TxnsContext;
 use tokio::runtime::Handle;
@@ -316,9 +314,6 @@ async fn run(args: Args) -> Result<(), anyhow::Error> {
         None,
     );
 
-    let grpc_host = args.grpc_host.and_then(|h| (!h.is_empty()).then_some(h));
-    let grpc_server_metrics = GrpcServerMetrics::register_with(&metrics_registry);
-
     // Start storage server.
     let storage_client_builder = mz_storage::serve(
         &metrics_registry,
@@ -335,13 +330,10 @@ async fn run(args: Args) -> Result<(), anyhow::Error> {
     );
     mz_ore::task::spawn(
         || "storage_server",
-        GrpcServer::serve(
-            &grpc_server_metrics,
+        transport::serve(
             args.storage_controller_listen_addr,
             BUILD_INFO.semver_version(),
-            grpc_host.clone(),
             storage_client_builder,
-            |svc| ProtoStorageServer::new(svc).max_decoding_message_size(MAX_GRPC_MESSAGE_SIZE),
         ),
     );
 
@@ -363,13 +355,10 @@ async fn run(args: Args) -> Result<(), anyhow::Error> {
     );
     mz_ore::task::spawn(
         || "compute_server",
-        GrpcServer::serve(
-            &grpc_server_metrics,
+        transport::serve(
             args.compute_controller_listen_addr,
             BUILD_INFO.semver_version(),
-            grpc_host,
             compute_client_builder,
-            |svc| ProtoComputeServer::new(svc).max_decoding_message_size(MAX_GRPC_MESSAGE_SIZE),
         ),
     );
 

src/clusterd/src/lib.rs

@@ -9,8 +9,9 @@
 
 //! A client for replicas of a compute instance.
 
-use std::sync::Arc;
+use std::str::FromStr;
 use std::sync::atomic::{self, AtomicBool};
+use std::sync::Arc;
 use std::time::{Duration, Instant};
 
 use anyhow::bail;
@@ -19,13 +20,15 @@ use mz_cluster_client::client::{ClusterReplicaLocation, ClusterStartupEpoch, Tim
 use mz_compute_types::dyncfgs::ENABLE_COMPUTE_REPLICA_EXPIRATION;
 use mz_dyncfg::ConfigSet;
 use mz_ore::channel::InstrumentedUnboundedSender;
+use mz_ore::netio::SocketAddr;
 use mz_ore::retry::{Retry, RetryState};
 use mz_ore::task::AbortOnDropHandle;
 use mz_service::client::GenericClient;
 use mz_service::params::GrpcClientParameters;
+use mz_service::transport;
 use tokio::select;
 use tokio::sync::mpsc::error::SendError;
-use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender, unbounded_channel};
+use tokio::sync::mpsc::{unbounded_channel, UnboundedReceiver, UnboundedSender};
 use tracing::{debug, info, trace, warn};
 
 use crate::controller::instance::ReplicaResponse;
@@ -36,7 +39,6 @@ use crate::metrics::IntCounter;
 use crate::metrics::ReplicaMetrics;
 use crate::protocol::command::{ComputeCommand, InitialComputeParameters, InstanceConfig};
 use crate::protocol::response::ComputeResponse;
-use crate::service::{ComputeClient, ComputeGrpcClient};
 
 type Client<T> = SequentialHydration<T>;
 
@@ -69,7 +71,6 @@ pub(super) struct ReplicaClient<T> {
 impl<T> ReplicaClient<T>
 where
     T: ComputeControllerTimestamp,
-    ComputeGrpcClient: ComputeClient<T>,
 {
     pub(super) fn spawn(
         id: ReplicaId,
@@ -160,7 +161,6 @@ struct ReplicaTask<T> {
 impl<T> ReplicaTask<T>
 where
     T: ComputeControllerTimestamp,
-    ComputeGrpcClient: ComputeClient<T>,
 {
     /// Asynchronously forwards commands to and responses from a single replica.
     async fn run(self) {
@@ -181,17 +181,15 @@ where
     async fn connect(&self) -> Client<T> {
         let try_connect = |retry: RetryState| {
             let addrs = &self.config.location.ctl_addrs;
-            let dests = addrs
+            let addrs = addrs
                 .iter()
-                .map(|addr| (addr.clone(), self.metrics.clone()))
+                .map(|addr| SocketAddr::from_str(addr).unwrap())
                 .collect();
             let version = self.build_info.semver_version();
-            let client_params = &self.config.grpc_client;
 
             async move {
                 let connect_start = Instant::now();
-                let connect_result =
-                    ComputeGrpcClient::connect_partitioned(dests, version, client_params).await;
+                let connect_result = transport::Client::connect_partitioned(addrs, version).await;
                 self.metrics.observe_connect_time(connect_start.elapsed());
 
                 connect_result.inspect_err(|error| {
@@ -233,7 +231,6 @@ where
     async fn run_message_loop(mut self, mut client: Client<T>) -> Result<(), anyhow::Error>
     where
         T: ComputeControllerTimestamp,
-        ComputeGrpcClient: ComputeClient<T>,
     {
         let id = self.replica_id;
         let incarnation = self.epoch.replica();

src/compute-client/src/controller/replica.rs

@@ -36,7 +36,7 @@ include!(concat!(
 /// from the compute controller.
 ///
 /// [`ComputeCommand`]: super::command::ComputeCommand
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
 pub enum ComputeResponse<T = mz_repr::Timestamp> {
     /// `Frontiers` announces the advancement of the various frontiers of the specified compute
     /// collection.
@@ -320,7 +320,7 @@ impl Arbitrary for FrontiersResponse {
 ///
 /// Note that each `Peek` expects to generate exactly one `PeekResponse`, i.e.
 /// we expect a 1:1 contract between `Peek` and `PeekResponse`.
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
 pub enum PeekResponse {
     /// Returned rows of a successful peek.
     Rows(RowCollection),

src/compute-client/src/protocol/response.rs

@@ -29,7 +29,7 @@ include!(concat!(env!("OUT_DIR"), "/mz_expr.row.collection.rs"));
 ///
 /// Note: the encoding format we use to represent [`Row`]s in this struct is
 /// not stable, and thus should never be persisted durably.
-#[derive(Default, Debug, Clone, PartialEq)]
+#[derive(Default, Debug, Clone, PartialEq, Serialize, Deserialize)]
 pub struct RowCollection {
     /// Contiguous blob of encoded Rows.
     encoded: Bytes,

src/expr/src/row/collection.rs

@@ -23,7 +23,7 @@ use std::{fmt, io};
 use async_trait::async_trait;
 use tokio::fs;
 use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
-use tokio::net::{self, TcpListener, TcpStream, UnixListener, UnixStream};
+use tokio::net::{self, tcp, unix, TcpListener, TcpStream, UnixListener, UnixStream};
 use tonic::transport::server::{Connected, TcpConnectInfo, UdsConnectInfo};
 use tracing::warn;
 
@@ -440,6 +440,21 @@ impl Stream {
             Stream::Unix(stream) => stream,
         }
     }
+
+    /// Splits a stream into a read half and a write half, which can be used to read and write the
+    /// stream concurrently.
+    pub fn split(self) -> (StreamReadHalf, StreamWriteHalf) {
+        match self {
+            Stream::Tcp(stream) => {
+                let (rx, tx) = stream.into_split();
+                (StreamReadHalf::Tcp(rx), StreamWriteHalf::Tcp(tx))
+            }
+            Stream::Unix(stream) => {
+                let (rx, tx) = stream.into_split();
+                (StreamReadHalf::Unix(rx), StreamWriteHalf::Unix(tx))
+            }
+        }
+    }
 }
 
 impl AsyncRead for Stream {
@@ -489,6 +504,56 @@ impl Connected for Stream {
     }
 }
 
+/// Read half of a [`Stream`], created by [`Stream::split`].
+#[derive(Debug)]
+pub enum StreamReadHalf {
+    Tcp(tcp::OwnedReadHalf),
+    Unix(unix::OwnedReadHalf),
+}
+
+impl AsyncRead for StreamReadHalf {
+    fn poll_read(
+        self: Pin<&mut Self>,
+        cx: &mut Context,
+        buf: &mut ReadBuf,
+    ) -> Poll<io::Result<()>> {
+        match self.get_mut() {
+            Self::Tcp(rx) => Pin::new(rx).poll_read(cx, buf),
+            Self::Unix(rx) => Pin::new(rx).poll_read(cx, buf),
+        }
+    }
+}
+
+/// Write half of a [`Stream`], created by [`Stream::split`].
+#[derive(Debug)]
+pub enum StreamWriteHalf {
+    Tcp(tcp::OwnedWriteHalf),
+    Unix(unix::OwnedWriteHalf),
+}
+
+impl AsyncWrite for StreamWriteHalf {
+    fn poll_write(self: Pin<&mut Self>, cx: &mut Context, buf: &[u8]) -> Poll<io::Result<usize>> {
+        match self.get_mut() {
+            Self::Tcp(tx) => Pin::new(tx).poll_write(cx, buf),
+            Self::Unix(tx) => Pin::new(tx).poll_write(cx, buf),
+        }
+    }
+
+    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context) -> Poll<io::Result<()>> {
+        match self.get_mut() {
+            Self::Tcp(tx) => Pin::new(tx).poll_flush(cx),
+            Self::Unix(tx) => Pin::new(tx).poll_flush(cx),
+        }
+    }
+
+    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context) -> Poll<io::Result<()>> {
+        match self.get_mut() {
+            Self::Tcp(tx) => Pin::new(tx).poll_shutdown(cx),
+            Self::Unix(tx) => Pin::new(tx).poll_shutdown(cx),
+        }
+    }
+}
+
 /// Connection information for a [`Stream`].
 #[derive(Debug, Clone)]
 pub enum ConnectInfo {

src/ore/src/netio/socket.rs

@@ -13,8 +13,10 @@ workspace = true
 anyhow = "1.0.95"
 async-stream = "0.3.3"
 async-trait = "0.1.83"
+bincode = "1.3.3"
 clap = { version = "4.5.23", features = ["env", "derive"] }
 crossbeam-channel = "0.5.15"
+flate2 = "1.1.0"
 futures = "0.3.31"
 http = "1.2.0"
 hyper-util = "0.1.6"

src/service/Cargo.toml

@@ -22,3 +22,4 @@ pub mod params;
 pub mod retry;
 pub mod secrets;
 pub mod tracing;
+pub mod transport;