Attach to cluster controller through Networking

This PR includes:
- Attachment of PPM to cluster controller now use Networking
- PPM observes partition processor status through a buffered watch mechanism
- PPM can now send control messages (unused at this PR) to processor for future use.
- PPM collects state information from running processors. This will be used in a follow PR to response to controller requests about partitions.


Cluster controller grpc service is kept for external tooling integration (CLI, etc.)

crates/cluster-controller/Cargo.toml

@@ -14,6 +14,8 @@ options_schema = ["dep:schemars"]
 [dependencies]
 restate-core = { workspace = true }
 restate-errors = { workspace = true }
+restate-network = { workspace = true }
+restate-node-protocol = { workspace = true }
 restate-types = { workspace = true }
 
 anyhow = { workspace = true }

crates/cluster-controller/src/service.rs

@@ -9,7 +9,20 @@
 // by the Apache License, Version 2.0.
 
 use codederror::CodedError;
-use restate_core::cancellation_watcher;
+use futures::stream::BoxStream;
+use futures::StreamExt;
+
+use restate_network::Networking;
+use restate_node_protocol::cluster_controller::{
+    Action, AttachRequest, AttachResponse, RunMode, RunPartition,
+};
+use restate_node_protocol::common::{KeyRange, RequestId};
+use restate_types::partition_table::FixedPartitionTable;
+
+use restate_core::network::{MessageRouterBuilder, NetworkSender};
+use restate_core::{cancellation_watcher, task_center, Metadata, ShutdownError, TaskCenter};
+use restate_node_protocol::MessageEnvelope;
+use restate_types::{GenerationalNodeId, Version};
 
 #[derive(Debug, thiserror::Error, CodedError)]
 pub enum Error {
@@ -18,19 +31,97 @@ pub enum Error {
     Error,
 }
 
-#[derive(Debug, Default)]
-pub struct Service {}
+pub struct Service {
+    metadata: Metadata,
+    networking: Networking,
+    incoming_messages: BoxStream<'static, MessageEnvelope<AttachRequest>>,
+}
+
+impl Service {
+    pub fn new(
+        metadata: Metadata,
+        networking: Networking,
+        router_builder: &mut MessageRouterBuilder,
+    ) -> Self {
+        let incoming_messages = router_builder.subscribe_to_stream(10);
+        Service {
+            metadata,
+            networking,
+            incoming_messages,
+        }
+    }
+}
 
-// todo: Replace with proper handle
 pub struct ClusterControllerHandle;
 
 impl Service {
     pub fn handle(&self) -> ClusterControllerHandle {
         ClusterControllerHandle
     }
 
-    pub async fn run(self) -> anyhow::Result<()> {
-        let _ = cancellation_watcher().await;
+    pub async fn run(mut self) -> anyhow::Result<()> {
+        // Make sure we have partition table before starting
+        let _ = self.metadata.wait_for_partition_table(Version::MIN).await?;
+
+        let mut shutdown = std::pin::pin!(cancellation_watcher());
+        let tc = task_center();
+        loop {
+            tokio::select! {
+                Some(message) = self.incoming_messages.next() => {
+                    let (from, message) = message.split();
+                    self.handle_attach_request(&tc, from, message)?;
+                }
+              _ = &mut shutdown => {
+                    return Ok(());
+                }
+            }
+        }
+    }
+
+    fn handle_attach_request(
+        &mut self,
+        tc: &TaskCenter,
+        from: GenerationalNodeId,
+        request: AttachRequest,
+    ) -> Result<(), ShutdownError> {
+        let partition_table = self
+            .metadata
+            .partition_table()
+            .expect("partition table is loaded before run");
+        let networking = self.networking.clone();
+        let response = self.create_attachment_response(&partition_table, from, request.request_id);
+        tc.spawn(
+            restate_core::TaskKind::Disposable,
+            "attachment-response",
+            None,
+            async move { Ok(networking.send(from.into(), &response).await?) },
+        )?;
         Ok(())
     }
+
+    fn create_attachment_response(
+        &self,
+        partition_table: &FixedPartitionTable,
+        _node: GenerationalNodeId,
+        request_id: RequestId,
+    ) -> AttachResponse {
+        // simulating a plan after initial attachement
+        let actions = partition_table
+            .partitioner()
+            .map(|(partition_id, key_range)| {
+                Action::RunPartition(RunPartition {
+                    partition_id,
+                    key_range_inclusive: KeyRange {
+                        from: *key_range.start(),
+                        to: *key_range.end(),
+                    },
+                    mode: RunMode::Leader,
+                })
+            })
+            .collect();
+        AttachResponse {
+            request_id,
+            actions,
+        }
+    }
 }

crates/core/src/worker_api/partition_processor_manager.rs

@@ -8,16 +8,57 @@
 // the Business Source License, use of this software will be governed
 // by the Apache License, Version 2.0.
 
-use restate_types::identifiers::PartitionId;
 use tokio::sync::{mpsc, oneshot};
 
+use restate_node_protocol::cluster_controller::RunMode;
+use restate_types::identifiers::{LeaderEpoch, PartitionId};
+use restate_types::logs::Lsn;
+use restate_types::time::MillisSinceEpoch;
+use restate_types::GenerationalNodeId;
+
 use crate::ShutdownError;
 
 #[derive(Debug)]
 pub enum ProcessorsManagerCommand {
     GetLivePartitions(oneshot::Sender<Vec<PartitionId>>),
 }
 
+#[derive(Debug, Clone, Eq, PartialEq)]
+pub enum ReplayStatus {
+    Starting,
+    Active,
+    CatchingUp { target_tail_lsn: Lsn },
+}
+
+#[derive(Debug, Clone)]
+pub struct PartitionProcessorStatus {
+    pub updated_at: MillisSinceEpoch,
+    pub planned_mode: RunMode,
+    pub effective_mode: Option<RunMode>,
+    pub last_observed_leader_epoch: Option<LeaderEpoch>,
+    pub last_observed_leader_node: Option<GenerationalNodeId>,
+    pub last_applied_log_lsn: Option<Lsn>,
+    pub last_record_applied_at: Option<MillisSinceEpoch>,
+    pub skipped_records: u64,
+    pub replay_status: ReplayStatus,
+}
+
+impl PartitionProcessorStatus {
+    pub fn new(planned_mode: RunMode) -> Self {
+        Self {
+            updated_at: MillisSinceEpoch::now(),
+            planned_mode,
+            effective_mode: None,
+            last_observed_leader_epoch: None,
+            last_observed_leader_node: None,
+            last_applied_log_lsn: None,
+            last_record_applied_at: None,
+            skipped_records: 0,
+            replay_status: ReplayStatus::Starting,
+        }
+    }
+}
+
 #[derive(Debug, Clone)]
 pub struct ProcessorsManagerHandle(mpsc::Sender<ProcessorsManagerCommand>);
 

crates/node-protocol/Cargo.toml

@@ -23,6 +23,7 @@ flexbuffers = { workspace = true }
 prost = { workspace = true }
 prost-types = { workspace = true }
 serde = { workspace = true}
+serde_with = { workspace = true }
 strum = { workspace = true }
 strum_macros = { workspace = true }
 thiserror = { workspace = true }

crates/node-protocol/proto/common.proto

@@ -22,16 +22,17 @@ message NodeId {
 }
 
 // A generic type for versioned metadata
-message Version {
-  uint32 value = 1;
-}
+message Version { uint32 value = 1; }
 
-// The target service for a message
+// The handle name or type tag of the message. For every target there must be
+// exactly one message handler implementation.
 enum TargetName {
   TargetName_UNKNOWN = 0;
   METADATA_MANAGER = 1;
   INGRESS = 2;
   LOCAL_METADATA_STORE = 3;
   LOCAL_METADATA_STORE_CLIENT = 4;
+  ATTACH_REQUEST = 5;
+  ATTACH_RESPONSE = 6;
+  PARTITION_PROCESSOR_MANAGER_REQUESTS = 7;
 }
-

crates/node-protocol/src/cluster_controller.rs

@@ -0,0 +1,51 @@
+// Copyright (c) 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.
+
+use restate_types::identifiers::PartitionId;
+use serde::{Deserialize, Serialize};
+
+use crate::common::{KeyRange, RequestId, TargetName};
+use crate::define_rpc;
+
+define_rpc! {
+    @request = AttachRequest,
+    @response = AttachResponse,
+    @request_target = TargetName::AttachRequest,
+    @response_target = TargetName::AttachResponse,
+}
+
+#[derive(Debug, Default, Clone, Serialize, Deserialize)]
+pub struct AttachRequest {
+    pub request_id: RequestId,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct AttachResponse {
+    pub request_id: RequestId,
+    pub actions: Vec<Action>,
+}
+
+#[derive(Debug, Clone, Copy, Serialize, Deserialize, Eq, PartialEq)]
+pub enum RunMode {
+    Leader,
+    Follower,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub enum Action {
+    RunPartition(RunPartition),
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct RunPartition {
+    pub partition_id: PartitionId,
+    pub key_range_inclusive: KeyRange,
+    pub mode: RunMode,
+}

crates/node-protocol/src/common.rs

@@ -8,8 +8,11 @@
 // the Business Source License, use of this software will be governed
 // by the Apache License, Version 2.0.
 
+use std::ops::RangeInclusive;
 use std::sync::atomic::AtomicUsize;
 
+use restate_types::identifiers::PartitionKey;
+
 include!(concat!(env!("OUT_DIR"), "/dev.restate.common.rs"));
 
 pub static MIN_SUPPORTED_PROTOCOL_VERSION: ProtocolVersion = ProtocolVersion::Flexbuffers;
@@ -105,6 +108,18 @@ impl From<restate_types::GenerationalNodeId> for NodeId {
     }
 }
 
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+pub struct KeyRange {
+    pub from: PartitionKey,
+    pub to: PartitionKey,
+}
+
+impl From<KeyRange> for RangeInclusive<PartitionKey> {
+    fn from(val: KeyRange) -> Self {
+        RangeInclusive::new(val.from, val.to)
+    }
+}
+
 // write tests for RequestId
 #[cfg(test)]
 mod tests {

crates/node-protocol/src/lib.rs

@@ -8,6 +8,7 @@
 // the Business Source License, use of this software will be governed
 // by the Apache License, Version 2.0.
 
+pub mod cluster_controller;
 pub mod codec;
 pub mod common;
 mod error;
@@ -50,6 +51,10 @@ impl<M> MessageEnvelope<M> {
     pub fn split(self) -> (GenerationalNodeId, M) {
         (self.peer, self.body)
     }
+
+    pub fn body(&self) -> &M {
+        &self.body
+    }
 }
 
 impl<M: RpcMessage> MessageEnvelope<M> {

crates/node-services/proto/cluster_ctrl_svc.proto

@@ -14,13 +14,9 @@ import "common.proto";
 package dev.restate.cluster_ctrl;
 
 service ClusterCtrlSvc {
-  // Attach worker at cluster controller
-  rpc AttachNode(AttachmentRequest) returns (AttachmentResponse);
-
+  rpc GetClusterState(ClusterStateRequest) returns (ClusterStateResponse);
 }
 
-message AttachmentRequest {
-  optional dev.restate.common.NodeId node_id = 1;
-}
+message ClusterStateRequest {}
 
-message AttachmentResponse {}
+message ClusterStateResponse {}

crates/node/src/lib.rs

@@ -145,7 +145,10 @@ impl Node {
         let admin_role = if config.has_role(Role::Admin) {
             Some(AdminRole::new(
                 updateable_config.clone(),
+                metadata.clone(),
+                networking.clone(),
                 metadata_manager.writer(),
+                &mut router_builder,
                 metadata_store_client.clone(),
             )?)
         } else {