feat(meta): introduce consistent hash logic

proto/common.proto

@@ -31,6 +31,10 @@ enum WorkerType {
   UNKNOWN = 2;
 }
 
+message ParallelUnit {
+  uint32 id = 1;
+}
+
 message WorkerNode {
   enum State {
     STARTING = 0;
@@ -40,6 +44,8 @@ message WorkerNode {
   WorkerType type = 2;
   HostAddress host = 3;
   State state = 4;
+  // a mapping from logical key to parallel unit, with logical key as the index of array
+  repeated ParallelUnit parallel_units = 5;
 }
 
 // A cluster can be either a set of OLAP compute nodes, or a set of streaming compute nodes.
@@ -48,3 +54,7 @@ message Cluster {
   repeated common.WorkerNode nodes = 2;
   map<string, string> config = 3;
 }
+
+message HashMapping {
+  repeated ParallelUnit parallel_units = 1;
+}

rust/frontend/src/scheduler/plan_fragmenter.rs

@@ -222,7 +222,7 @@ mod tests {
 
     use std::sync::Arc;
 
-    use risingwave_pb::common::{WorkerNode, WorkerType};
+    use risingwave_pb::common::{ParallelUnit, WorkerNode, WorkerType};
     use risingwave_pb::plan::JoinType;
 
     use crate::optimizer::plan_node::{
@@ -308,18 +308,21 @@ mod tests {
             r#type: WorkerType::ComputeNode as i32,
             host: None,
             state: risingwave_pb::common::worker_node::State::Running as i32,
+            parallel_units: vec![ParallelUnit { id: 1 }],
         };
         let worker2 = WorkerNode {
             id: 1,
             r#type: WorkerType::ComputeNode as i32,
             host: None,
             state: risingwave_pb::common::worker_node::State::Running as i32,
+            parallel_units: vec![ParallelUnit { id: 2 }],
         };
         let worker3 = WorkerNode {
             id: 2,
             r#type: WorkerType::ComputeNode as i32,
             host: None,
             state: risingwave_pb::common::worker_node::State::Running as i32,
+            parallel_units: vec![ParallelUnit { id: 3 }],
         };
         let workers = vec![worker1.clone(), worker2.clone(), worker3.clone()];
         let worker_node_manager = Arc::new(WorkerNodeManager::mock(workers));

rust/meta/src/cluster/mod.rs

@@ -8,28 +8,34 @@ use risingwave_common::error::ErrorCode::InternalError;
 use risingwave_common::error::{Result, RwError};
 use risingwave_common::try_match_expand;
 use risingwave_pb::common::worker_node::State;
-use risingwave_pb::common::{HostAddress, WorkerNode, WorkerType};
+use risingwave_pb::common::{HostAddress, ParallelUnit, WorkerNode, WorkerType};
 use risingwave_pb::meta::subscribe_response::{Info, Operation};
 
 use crate::hummock::HummockManager;
-use crate::manager::{IdCategory, IdGeneratorManagerRef, MetaSrvEnv, NotificationManagerRef};
+use crate::manager::{
+    HashDispatchManager, HashDispatchManagerRef, IdCategory, IdGeneratorManagerRef, MetaSrvEnv,
+    NotificationManagerRef,
+};
 use crate::model::{MetadataModel, Worker};
 use crate::storage::MetaStore;
 
 pub type NodeId = u32;
+pub type ParallelUnitId = u32;
 pub type NodeLocations = HashMap<NodeId, WorkerNode>;
+pub type StoredClusterManagerRef<S> = Arc<StoredClusterManager<S>>;
+
+const DEFAULT_WORKNODE_PARALLEL_DEGREE: usize = 8;
 
 /// [`StoredClusterManager`] manager cluster/worker meta data in [`MetaStore`].
 pub struct StoredClusterManager<S> {
     meta_store_ref: Arc<S>,
     id_gen_manager_ref: IdGeneratorManagerRef<S>,
     hummock_manager_ref: Option<Arc<HummockManager<S>>>,
+    dispatch_manager_ref: HashDispatchManagerRef<S>,
     workers: DashMap<WorkerKey, Worker>,
     nm: NotificationManagerRef,
 }
 
-pub type StoredClusterManagerRef<S> = Arc<StoredClusterManager<S>>;
-
 pub struct WorkerKey(pub HostAddress);
 
 impl PartialEq<Self> for WorkerKey {
@@ -62,15 +68,23 @@ where
             "Worker::list fail"
         )?;
         let worker_map = DashMap::new();
+        let mut compute_nodes = Vec::new();
 
         workers.iter().for_each(|w| {
             worker_map.insert(WorkerKey(w.key().unwrap()), w.clone());
+            if w.worker_type() == WorkerType::ComputeNode {
+                compute_nodes.push(w.worker_node());
+            }
         });
 
+        let dispatch_manager_ref =
+            Arc::new(HashDispatchManager::new(&compute_nodes, meta_store_ref.clone()).await?);
+
         Ok(Self {
             meta_store_ref,
             id_gen_manager_ref: env.id_gen_manager_ref(),
             hummock_manager_ref,
+            dispatch_manager_ref,
             workers: worker_map,
             nm,
         })
@@ -84,19 +98,42 @@ where
         match self.workers.entry(WorkerKey(host_address.clone())) {
             Entry::Occupied(o) => Ok((o.get().to_protobuf(), false)),
             Entry::Vacant(v) => {
-                let id = self
+                let worker_id = self
                     .id_gen_manager_ref
                     .generate::<{ IdCategory::Worker }>()
                     .await?;
+
+                let mut parallel_units = Vec::with_capacity(DEFAULT_WORKNODE_PARALLEL_DEGREE);
+                let start_id = self
+                    .id_gen_manager_ref
+                    .generate_interval::<{ IdCategory::ParallelUnit }>(
+                        DEFAULT_WORKNODE_PARALLEL_DEGREE as i32,
+                    )
+                    .await? as usize;
+                for parallel_unit_id in start_id..start_id + DEFAULT_WORKNODE_PARALLEL_DEGREE {
+                    parallel_units.push(ParallelUnit {
+                        id: parallel_unit_id as u32,
+                    });
+                }
+
                 let worker_node = WorkerNode {
-                    id: id as u32,
+                    id: worker_id as u32,
                     r#type: r#type as i32,
                     host: Some(host_address.clone()),
                     state: State::Starting as i32,
+                    parallel_units,
                 };
+
                 let worker = Worker::from_protobuf(worker_node.clone());
                 worker.insert(&*self.meta_store_ref).await?;
 
+                // Alter consistent hash mapping
+                if r#type == WorkerType::ComputeNode {
+                    self.dispatch_manager_ref
+                        .add_worker_mapping(&worker.worker_node())
+                        .await?;
+                }
+
                 Ok((v.insert(worker).to_protobuf(), true))
             }
         }
@@ -161,6 +198,12 @@ where
                         .release_context_resource(entry.1.to_protobuf().id)
                         .await?;
                 }
+                let worker = entry.1;
+                if worker.worker_type() == WorkerType::ComputeNode {
+                    self.dispatch_manager_ref
+                        .delete_worker_mapping(&worker.worker_node())
+                        .await?;
+                }
                 Ok(())
             }
         }

rust/meta/src/lib.rs

@@ -11,6 +11,7 @@
 #![feature(binary_heap_drain_sorted)]
 #![feature(const_fn_trait_bound)]
 #![feature(option_result_contains)]
+#![feature(map_first_last)]
 
 mod barrier;
 pub mod cluster;