PartitionProcessorManager as long-living service

Cargo.toml

@@ -49,6 +49,7 @@ restate-network = { path = "crates/network" }
 restate-node = { path = "crates/node" }
 restate-node-protocol = { path = "crates/node-protocol" }
 restate-node-services = { path = "crates/node-services" }
+restate-partition-store = { path = "crates/partition-store" }
 restate-queue = { path = "crates/queue" }
 restate-rocksdb = { path = "crates/rocksdb" }
 restate-schema = { path = "crates/schema" }
@@ -60,7 +61,6 @@ restate-service-protocol = { path = "crates/service-protocol" }
 restate-storage-api = { path = "crates/storage-api" }
 restate-storage-query-datafusion = { path = "crates/storage-query-datafusion" }
 restate-storage-query-postgres = { path = "crates/storage-query-postgres" }
-restate-storage-rocksdb = { path = "crates/storage-rocksdb" }
 restate-test-util = { path = "crates/test-util" }
 restate-timer = { path = "crates/timer" }
 restate-timer-queue = { path = "crates/timer-queue" }
@@ -120,7 +120,7 @@ prost-build = "0.12.1"
 prost-types = "0.12.1"
 rand = "0.8.5"
 rayon = { version = "1.10" }
-rocksdb = { version = "0.22.0", features = ["multi-threaded-cf"], git = "/~https://github.com/restatedev/rust-rocksdb", branch="next" }
+rocksdb = { version = "0.22.0", features = ["multi-threaded-cf"], git = "/~https://github.com/restatedev/rust-rocksdb", rev="c2181f2b5da6d7bc201dc858433ed9e1c4bba4b7" }
 rustls = "0.21.6"
 schemars = { version = "0.8", features = ["bytes", "enumset"] }
 serde = { version = "1.0", features = ["derive"] }

crates/benchmarks/src/lib.rs

@@ -102,7 +102,9 @@ pub fn spawn_restate(config: Configuration) -> TaskCenter {
         RocksDbManager::init(Constant::new(config.common))
     });
     tc.spawn(TaskKind::TestRunner, "benchmark", None, async move {
-        let node = Node::new(updateable_config).expect("Restate node must build");
+        let node = Node::create(updateable_config)
+            .await
+            .expect("Restate node must build");
         cloned_tc.run_in_scope("startup", None, node.start()).await
     })
     .unwrap();

crates/core/src/lib.rs

@@ -14,6 +14,7 @@ mod metric_definitions;
 pub mod network;
 mod task_center;
 mod task_center_types;
+pub mod worker_api;
 
 pub use metadata::{
     spawn_metadata_manager, Metadata, MetadataKind, MetadataManager, MetadataWriter, SyncError,

crates/core/src/worker_api/mod.rs

@@ -0,0 +1,13 @@
+// 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.
+
+mod partition_processor_manager;
+
+pub use partition_processor_manager::*;

crates/core/src/worker_api/partition_processor_manager.rs

@@ -0,0 +1,37 @@
+// 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 tokio::sync::{mpsc, oneshot};
+
+use crate::ShutdownError;
+
+#[derive(Debug)]
+pub enum ProcessorsManagerCommand {
+    GetLivePartitions(oneshot::Sender<Vec<PartitionId>>),
+}
+
+#[derive(Debug, Clone)]
+pub struct ProcessorsManagerHandle(mpsc::Sender<ProcessorsManagerCommand>);
+
+impl ProcessorsManagerHandle {
+    pub fn new(sender: mpsc::Sender<ProcessorsManagerCommand>) -> Self {
+        Self(sender)
+    }
+
+    pub async fn get_live_partitions(&self) -> Result<Vec<PartitionId>, ShutdownError> {
+        let (tx, rx) = oneshot::channel();
+        self.0
+            .send(ProcessorsManagerCommand::GetLivePartitions(tx))
+            .await
+            .unwrap();
+        rx.await.map_err(|_| ShutdownError)
+    }
+}

crates/node/src/lib.rs

@@ -100,7 +100,7 @@ pub struct Node {
 }
 
 impl Node {
-    pub fn new(updateable_config: UpdateableConfiguration) -> Result<Self, BuildError> {
+    pub async fn create(updateable_config: UpdateableConfiguration) -> Result<Self, BuildError> {
         let config = updateable_config.pinned();
         // ensure we have cluster admin role if bootstrapping.
         if config.common.allow_bootstrap {
@@ -140,7 +140,7 @@ impl Node {
         metadata_manager.register_in_message_router(&mut router_builder);
         let metadata = metadata_manager.metadata();
         let updating_schema_information = metadata.schema_updateable();
-        let bifrost = BifrostService::new(metadata);
+        let bifrost = BifrostService::new(metadata.clone());
 
         let admin_role = if config.has_role(Role::Admin) {
             Some(AdminRole::new(
@@ -153,14 +153,18 @@ impl Node {
         };
 
         let worker_role = if config.has_role(Role::Worker) {
-            Some(WorkerRole::new(
-                updateable_config.clone(),
-                &mut router_builder,
-                networking.clone(),
-                bifrost.handle(),
-                metadata_store_client,
-                updating_schema_information,
-            )?)
+            Some(
+                WorkerRole::create(
+                    metadata,
+                    updateable_config.clone(),
+                    &mut router_builder,
+                    networking.clone(),
+                    bifrost.handle(),
+                    metadata_store_client,
+                    updating_schema_information,
+                )
+                .await?,
+            )
         } else {
             None
         };
@@ -321,7 +325,7 @@ impl Node {
                 TaskKind::SystemBoot,
                 "worker-init",
                 None,
-                worker_role.start(bifrost),
+                worker_role.start(),
             )?;
         }