[Core] Let core crate own network fabric messages

crates/admin/src/cluster_controller/scheduler.rs

@@ -582,6 +582,7 @@ impl Drop for TargetPartitionPlacementState<'_> {
 #[cfg(test)]
 mod tests {
     use std::collections::BTreeMap;
+    use std::convert::Infallible;
     use std::iter;
     use std::num::NonZero;
     use std::time::Duration;
@@ -727,10 +728,10 @@ mod tests {
                 if message.body().target() == TargetName::ControlProcessors {
                     let message = message
                         .try_map(|mut m| {
-                            ControlProcessors::decode(
+                            Ok::<_, Infallible>(ControlProcessors::decode(
                                 &mut m.payload,
                                 restate_types::net::CURRENT_PROTOCOL_VERSION,
-                            )
+                            ))
                         })
                         .unwrap();
                     Some((node_id, message))

crates/bifrost/benches/replicated_loglet_serde.rs

@@ -21,17 +21,18 @@ use rand::distr::Alphanumeric;
 use rand::{Rng, random};
 
 use restate_bifrost::InputRecord;
+use restate_core::network::protobuf::network::Message;
+use restate_core::network::protobuf::network::message::{BinaryMessage, Body};
 use restate_storage_api::deduplication_table::{DedupInformation, EpochSequenceNumber, ProducerId};
 use restate_types::GenerationalNodeId;
 use restate_types::identifiers::{InvocationId, LeaderEpoch, PartitionProcessorRpcRequestId};
 use restate_types::invocation::{
     InvocationTarget, ServiceInvocation, ServiceInvocationSpanContext,
 };
 use restate_types::logs::{LogId, LogletId, LogletOffset, Record, SequenceNumber};
-use restate_types::net::codec::{MessageBodyExt, WireDecode, WireEncode};
+use restate_types::net::codec::{Targeted, WireDecode, WireEncode};
 use restate_types::net::log_server::{LogServerRequestHeader, Store, StoreFlags};
 use restate_types::net::replicated_loglet::{Append, CommonRequestHeader};
-use restate_types::protobuf::node::Message;
 use restate_types::time::MillisSinceEpoch;
 use restate_wal_protocol::{Command, Destination, Envelope};
 
@@ -133,10 +134,13 @@ fn serialize_store_message(payloads: Arc<[Record]>) -> anyhow::Result<Message> {
         known_archived: LogletOffset::INVALID,
     };
 
-    let body = store_message.encode(restate_types::net::ProtocolVersion::V1)?;
+    let body = Body::Encoded(BinaryMessage {
+        target: Store::TARGET.into(),
+        payload: store_message.encode_to_bytes(restate_types::net::ProtocolVersion::V1),
+    });
 
     let message = Message {
-        header: Some(restate_types::protobuf::node::Header {
+        header: Some(restate_core::network::protobuf::network::Header {
             my_nodes_config_version: Some(5),
             my_logs_version: None,
             my_schema_version: None,
@@ -160,10 +164,13 @@ fn serialize_append_message(payloads: Arc<[Record]>) -> anyhow::Result<Message>
         payloads,
     };
 
-    let body = append_message.encode(restate_types::net::ProtocolVersion::V1)?;
+    let body = Body::Encoded(BinaryMessage {
+        target: Append::TARGET.into(),
+        payload: append_message.encode_to_bytes(restate_types::net::ProtocolVersion::V1),
+    });
 
     let message = Message {
-        header: Some(restate_types::protobuf::node::Header {
+        header: Some(restate_core::network::protobuf::network::Header {
             my_nodes_config_version: Some(5),
             my_logs_version: None,
             my_schema_version: None,
@@ -183,7 +190,7 @@ fn deserialize_append_message(serialized_message: Bytes) -> anyhow::Result<Appen
     let body = msg.body.unwrap();
     // we ignore non-deserializable messages (serde errors, or control signals in drain)
     let msg_body = body.try_as_binary_body(restate_types::net::ProtocolVersion::V1)?;
-    Ok(Append::decode(msg_body.payload, protocol_version)?)
+    Ok(Append::decode(msg_body.payload, protocol_version))
 }
 
 fn replicated_loglet_append_serde(c: &mut Criterion) {

crates/core/Cargo.toml

@@ -15,9 +15,10 @@ options_schema = ["dep:schemars"]
 [dependencies]
 workspace-hack = { version = "0.1", path = "../../workspace-hack" }
 
-restate-types = { workspace = true }
 restate-core-derive = { workspace = true, optional = true }
+restate-futures-util = { workspace = true }
 restate-object-store-util = { workspace = true }
+restate-types = { workspace = true }
 
 ahash = { workspace = true }
 anyhow = { workspace = true }
@@ -45,6 +46,7 @@ hyper = { workspace = true }
 hyper-util = { workspace = true, features = ["server-graceful", "server"] }
 metrics = { workspace = true }
 opentelemetry = { workspace = true }
+opentelemetry_sdk = { workspace = true }
 object_store = { workspace = true, features = ["aws"] }
 parking_lot = { workspace = true }
 pin-project-lite = { workspace = true }

crates/core/build.rs

@@ -40,12 +40,23 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     tonic_build::configure()
         .bytes(["."])
+        .enum_attribute("Datagram", "#[derive(::derive_more::From)]")
+        .enum_attribute("Datagram.datagram", "#[derive(::derive_more::From)]")
+        .enum_attribute(
+            "Message.body",
+            "#[derive(::derive_more::IsVariant, ::derive_more::From)]",
+        )
+        .enum_attribute("Body", "#[derive(::derive_more::From)]")
         .file_descriptor_set_path(out_dir.join("core_node_svc_descriptor.bin"))
         // allow older protobuf compiler to be used
         .protoc_arg("--experimental_allow_proto3_optional")
+        .extern_path(".restate.common", "::restate_types::protobuf::common")
         .extern_path(".restate.node", "::restate_types::protobuf::node")
         .compile_protos(
-            &["./protobuf/core_node_svc.proto"],
+            &[
+                "./protobuf/restate/network.proto",
+                "./protobuf/core_node_svc.proto",
+            ],
             &["protobuf", "../types/protobuf"],
         )?;
 

crates/core/protobuf/core_node_svc.proto

@@ -9,13 +9,13 @@
 
 syntax = "proto3";
 
-import "restate/node.proto";
+import "restate/network.proto";
 
 package restate.core_node_svc;
 
 // Service which is only accessible on nodes that are alive.
 service CoreNodeSvc {
   // Create a bidirectional node-to-node stream
-  rpc CreateConnection(stream restate.node.Message)
-      returns (stream restate.node.Message);
-}
+  rpc CreateConnection(stream restate.network.Message)
+      returns (stream restate.network.Message);
+}

crates/types/protobuf/restate/node.proto → crates/core/protobuf/restate/network.proto

@@ -11,7 +11,7 @@ syntax = "proto3";
 
 import "restate/common.proto";
 
-package restate.node;
+package restate.network;
 
 //
 // # Wire Protocol Of Streaming Connections
@@ -44,7 +44,8 @@ message Hello {
   restate.common.ProtocolVersion max_protocol_version = 2;
   string cluster_name = 3;
   // generational node id of sender (who am I)
-  // this is optional for future-proofing with anonymous clients using this protocol
+  // this is optional for future-proofing with anonymous clients using this
+  // protocol
   optional restate.common.GenerationalNodeId my_node_id = 4;
 }
 
@@ -63,6 +64,10 @@ message Message {
     // Connection will be dropped
     DRAIN_CONNECTION = 2;
     CODEC_ERROR = 3;
+    // no more requests will be sent over this stream
+    REQUEST_STREAM_DRAINED = 4;
+    // no more responses will be sent over this stream
+    RESPONSE_STREAM_DRAINED = 5;
   }
   message ConnectionControl {
     Signal signal = 1;