Move networking and connection manager into network module

crates/metadata-store/src/lib.rs

@@ -13,6 +13,7 @@ mod grpc_svc;
 pub mod local;
 pub mod raft;
 mod util;
+mod network;
 
 use bytestring::ByteString;
 use restate_core::metadata_store::VersionedValue;

crates/metadata-store/src/raft/connection_manager.rs → crates/metadata-store/src/network/connection_manager.rs

@@ -1,4 +1,4 @@
-// Copyright (c) 2024 - Restate Software, Inc., Restate GmbH.
+// Copyright (c) 2023 - 2024 Restate Software, Inc., Restate GmbH.
 // All rights reserved.
 //
 // Use of this software is governed by the Business Source License
@@ -8,7 +8,6 @@
 // the Business Source License, use of this software will be governed
 // by the Apache License, Version 2.0.
 
-use crate::raft::grpc_svc::RaftMessage;
 use futures::StreamExt;
 use protobuf::Message as ProtobufMessage;
 use raft::prelude::Message;
@@ -20,6 +19,7 @@ use tokio::sync::mpsc::error::TrySendError;
 use tokio_stream::wrappers::ReceiverStream;
 use tonic::codegen::BoxStream;
 use tracing::{debug, instrument};
+use crate::network::grpc_svc::RaftMessage;
 
 #[derive(Debug, thiserror::Error)]
 pub enum ConnectionError {

crates/metadata-store/src/raft/handler.rs → crates/metadata-store/src/network/handler.rs

@@ -1,4 +1,4 @@
-// Copyright (c) 2024 - Restate Software, Inc., Restate GmbH.
+// Copyright (c) 2023 - 2024 Restate Software, Inc., Restate GmbH.
 // All rights reserved.
 //
 // Use of this software is governed by the Business Source License
@@ -8,12 +8,13 @@
 // the Business Source License, use of this software will be governed
 // by the Apache License, Version 2.0.
 
-use crate::raft::connection_manager::{ConnectionError, ConnectionManager};
-use crate::raft::grpc_svc::raft_metadata_store_svc_server::RaftMetadataStoreSvc;
-use crate::raft::grpc_svc::RaftMessage;
 use std::str::FromStr;
 use tonic::codegen::BoxStream;
 use tonic::{Request, Response, Status, Streaming};
+use crate::network::connection_manager::ConnectionError;
+use crate::network::ConnectionManager;
+use crate::network::grpc_svc::raft_metadata_store_svc_server::RaftMetadataStoreSvc;
+use crate::network::grpc_svc::RaftMessage;
 
 pub const RAFT_PEER_METADATA_KEY: &str = "x-restate-raft-peer";
 

crates/metadata-store/src/network/mod.rs

@@ -0,0 +1,20 @@
+// Copyright (c) 2023 - 2024 Restate Software, Inc., Restate GmbH.
+// All rights reserved.
+//
+// Use of this software is governed by the Business Source License
+// included in the LICENSE file.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0.
+
+mod connection_manager;
+mod grpc_svc;
+mod handler;
+mod networking;
+
+pub use connection_manager::ConnectionManager;
+pub use grpc_svc::FILE_DESCRIPTOR_SET;
+pub use grpc_svc::raft_metadata_store_svc_server::RaftMetadataStoreSvcServer;
+pub use handler::RaftMetadataStoreHandler;
+pub use networking::Networking;

crates/metadata-store/src/raft/networking.rs → crates/metadata-store/src/network/networking.rs

@@ -1,4 +1,4 @@
-// Copyright (c) 2024 - Restate Software, Inc., Restate GmbH.
+// Copyright (c) 2023 - 2024 Restate Software, Inc., Restate GmbH.
 // All rights reserved.
 //
 // Use of this software is governed by the Business Source License
@@ -8,9 +8,6 @@
 // the Business Source License, use of this software will be governed
 // by the Apache License, Version 2.0.
 
-use crate::raft::connection_manager::ConnectionManager;
-use crate::raft::grpc_svc::RaftMessage;
-use crate::raft::handler::RAFT_PEER_METADATA_KEY;
 use bytes::{BufMut, BytesMut};
 use futures::FutureExt;
 use protobuf::Message as ProtobufMessage;
@@ -26,6 +23,9 @@ use tokio_stream::wrappers::ReceiverStream;
 use tonic::metadata::MetadataValue;
 use tonic::IntoStreamingRequest;
 use tracing::{debug, trace};
+use crate::network::connection_manager::ConnectionManager;
+use crate::network::grpc_svc::RaftMessage;
+use crate::network::handler::RAFT_PEER_METADATA_KEY;
 
 #[derive(Debug, thiserror::Error)]
 pub enum TrySendError<T> {
@@ -131,7 +131,7 @@ impl Networking {
                 );
 
                 async move {
-                    let mut raft_client = crate::raft::grpc_svc::raft_metadata_store_svc_client::RaftMetadataStoreSvcClient::new(channel);
+                    let mut raft_client = crate::network::grpc_svc::raft_metadata_store_svc_client::RaftMetadataStoreSvcClient::new(channel);
                     let (outgoing_tx, outgoing_rx) = mpsc::channel(128);
 
                     let mut request = ReceiverStream::new(outgoing_rx).into_streaming_request();

crates/metadata-store/src/raft/mod.rs

@@ -8,10 +8,6 @@
 // the Business Source License, use of this software will be governed
 // by the Apache License, Version 2.0.
 
-mod connection_manager;
-pub mod grpc_svc;
-mod handler;
-mod networking;
 pub mod service;
 mod storage;
 mod store;

crates/metadata-store/src/raft/service.rs

@@ -12,12 +12,8 @@ use crate::grpc::handler::MetadataStoreHandler;
 use crate::grpc::server::GrpcServer;
 use crate::grpc::service_builder::GrpcServiceBuilder;
 use crate::grpc_svc::metadata_store_svc_server::MetadataStoreSvcServer;
-use crate::raft::connection_manager::ConnectionManager;
-use crate::raft::grpc_svc::raft_metadata_store_svc_server::RaftMetadataStoreSvcServer;
-use crate::raft::handler::RaftMetadataStoreHandler;
-use crate::raft::networking::Networking;
 use crate::raft::store::RaftMetadataStore;
-use crate::{grpc_svc, Error, MetadataStoreService};
+use crate::{grpc_svc, network, Error, MetadataStoreService};
 use assert2::let_assert;
 use futures::TryFutureExt;
 use restate_core::{TaskCenter, TaskKind};
@@ -26,6 +22,7 @@ use restate_types::health::HealthStatus;
 use restate_types::live::BoxedLiveLoad;
 use restate_types::protobuf::common::MetadataServerStatus;
 use tokio::sync::mpsc;
+use crate::network::{ConnectionManager, Networking, RaftMetadataStoreHandler, RaftMetadataStoreSvcServer};
 
 pub struct RaftMetadataStoreService {
     health_status: HealthStatus<MetadataServerStatus>,
@@ -70,6 +67,8 @@ impl MetadataStoreService for RaftMetadataStoreService {
         builder.add_service(MetadataStoreSvcServer::new(MetadataStoreHandler::new(
             store.request_sender(),
         )));
+
+        builder.register_file_descriptor_set_for_reflection(network::FILE_DESCRIPTOR_SET);
         builder.add_service(RaftMetadataStoreSvcServer::new(
             RaftMetadataStoreHandler::new(connection_manager),
         ));

crates/metadata-store/src/raft/store.rs

@@ -8,7 +8,6 @@
 // the Business Source License, use of this software will be governed
 // by the Apache License, Version 2.0.
 
-use crate::raft::networking::Networking;
 use crate::raft::storage;
 use crate::raft::storage::RocksDbStorage;
 use crate::{
@@ -35,6 +34,7 @@ use tokio::time::MissedTickBehavior;
 use tracing::{debug, info, warn};
 use tracing_slog::TracingSlogDrain;
 use ulid::Ulid;
+use crate::network::Networking;
 
 #[derive(Debug, thiserror::Error)]
 pub enum BuildError {