feat(meta): iterative streaming scheduler (part 2)

proto/stream_plan.proto

@@ -346,6 +346,9 @@ enum ChainType {
 //   1. MergeNode (as a placeholder) for streaming read.
 //   2. BatchPlanNode for snapshot read.
 message ChainNode {
+  reserved 5;
+  reserved "same_worker_node";
+
   uint32 table_id = 1;
   // The schema of input stream, which will be used to build a MergeNode
   repeated plan_common.Field upstream_fields = 2;
@@ -356,8 +359,6 @@ message ChainNode {
   // large. However, in some cases, e.g., shared state, the barrier cannot be rearranged in ChainNode.
   // ChainType is used to decide which implementation for the ChainNode.
   ChainType chain_type = 4;
-  // Whether to place this chain on the same worker node as upstream actors.
-  bool same_worker_node = 5;
   // Whether the upstream materialize is and this chain should be a singleton.
   // FIXME: This is a workaround for fragmenter since the distribution info will be lost if there's only one
   // fragment in the downstream mview. Remove this when we refactor the fragmenter.
@@ -562,13 +563,11 @@ message Dispatcher {
   repeated uint32 downstream_actor_id = 5;
 }
 
-// Used to place an actor together with another actor in the same worker node.
-message ColocatedActorId {
-  uint32 id = 1;
-}
-
 // A StreamActor is a running fragment of the overall stream graph,
 message StreamActor {
+  reserved 7;
+  reserved "colocated_upstream_actor_id";
+
   uint32 actor_id = 1;
   uint32 fragment_id = 2;
   StreamNode nodes = 3;
@@ -578,8 +577,6 @@ message StreamActor {
   // It is painstaking to traverse through the node tree and get upstream actor id from the root StreamNode.
   // We duplicate the information here to ease the parsing logic in stream manager.
   repeated uint32 upstream_actor_id = 6;
-  // Placement rule for actor, need to stay on the same node as a specified upstream actor.
-  ColocatedActorId colocated_upstream_actor_id = 7;
   // Vnodes that the executors in this actor own.
   // If the fragment is a singleton, this field will not be set and leave a `None`.
   common.Buffer vnode_bitmap = 8;
@@ -619,10 +616,11 @@ message StreamFragmentGraph {
   }
 
   message StreamFragmentEdge {
+    reserved 2;
+    reserved "same_worker_node";
+
     // Dispatch strategy for the fragment.
     DispatchStrategy dispatch_strategy = 1;
-    // Whether the two linked nodes should be placed on the same worker node
-    bool same_worker_node = 2;
     // A unique identifier of this edge. Generally it should be exchange node's operator id. When
     // rewriting fragments into delta joins or when inserting 1-to-1 exchange, there will be
     // virtual links generated.

src/common/src/hash/consistent_hash/mapping.rs

@@ -117,6 +117,15 @@ impl<T: VnodeMappingItem> VnodeMapping<T> {
         self[vnode]
     }
 
+    /// Get the item matched by the virtual nodes indicated by high bits in the given `bitmap`.
+    /// Returns `None` if the no virtual node is set in the bitmap.
+    pub fn get_matched(&self, bitmap: &Bitmap) -> Option<T::Item> {
+        bitmap
+            .iter_ones()
+            .next() // only need to check the first one
+            .map(|i| self.get(VirtualNode::from_index(i)))
+    }
+
     /// Iterate over all items in this mapping, in the order of vnodes.
     pub fn iter(&self) -> impl Iterator<Item = T::Item> + '_ {
         self.data
@@ -144,6 +153,11 @@ impl<T: VnodeMappingItem> VnodeMapping<T> {
         self.data.iter().copied().sorted().dedup()
     }
 
+    /// Returns the item if it's the only item in this mapping, otherwise returns `None`.
+    pub fn to_single(&self) -> Option<T::Item> {
+        self.data.iter().copied().dedup().exactly_one().ok()
+    }
+
     /// Convert this vnode mapping to a mapping from items to bitmaps, where each bitmap represents
     /// the vnodes mapped to the item.
     pub fn to_bitmaps(&self) -> HashMap<T::Item, Bitmap> {

src/frontend/src/optimizer/plan_node/stream_index_scan.rs

@@ -177,7 +177,6 @@ impl StreamIndexScan {
             ],
             node_body: Some(ProstStreamNode::Chain(ChainNode {
                 table_id: self.logical.table_desc().table_id.table_id,
-                same_worker_node: true,
                 chain_type: self.chain_type as i32,
                 // The fields from upstream
                 upstream_fields: self

src/frontend/src/optimizer/plan_node/stream_table_scan.rs

@@ -200,7 +200,6 @@ impl StreamTableScan {
             ],
             node_body: Some(ProstStreamNode::Chain(ChainNode {
                 table_id: self.logical.table_desc().table_id.table_id,
-                same_worker_node: false,
                 chain_type: self.chain_type as i32,
                 // The fields from upstream
                 upstream_fields: self

src/frontend/src/stream_fragmenter/graph/fragment_graph.rs

@@ -53,9 +53,6 @@ pub struct StreamFragmentEdge {
     /// Dispatch strategy for the fragment.
     pub dispatch_strategy: DispatchStrategy,
 
-    /// Whether the two linked nodes should be placed on the same worker node
-    pub same_worker_node: bool,
-
     /// A unique identifier of this edge. Generally it should be exchange node's operator id. When
     /// rewriting fragments into delta joins or when inserting 1-to-1 exchange, there will be
     /// virtual links generated.
@@ -139,7 +136,6 @@ impl StreamFragmentGraph {
             upstream_id,
             downstream_id,
             dispatch_strategy: Some(edge.dispatch_strategy),
-            same_worker_node: edge.same_worker_node,
             link_id: edge.link_id,
         };
 

src/frontend/src/stream_fragmenter/mod.rs

@@ -295,7 +295,6 @@ fn build_fragment(
                         current_fragment.fragment_id,
                         StreamFragmentEdge {
                             dispatch_strategy: exchange_node_strategy,
-                            same_worker_node: false,
                             link_id: child_node.operator_id,
                         },
                     );

src/frontend/src/stream_fragmenter/rewrite/delta_join.rs

@@ -205,7 +205,6 @@ fn build_delta_join_inner(
         lookup_0_frag.fragment_id,
         StreamFragmentEdge {
             dispatch_strategy: dispatch_no_shuffle(),
-            same_worker_node: true,
             link_id: exchange_a0l0.operator_id,
         },
     );
@@ -223,8 +222,6 @@ fn build_delta_join_inner(
                     .map(|x| *x as u32)
                     .collect_vec(),
             ),
-            // stream input doesn't need to be on the same worker node as lookup
-            same_worker_node: false,
             link_id: exchange_a0l1.operator_id,
         },
     );
@@ -242,8 +239,6 @@ fn build_delta_join_inner(
                     .map(|x| *x as u32)
                     .collect_vec(),
             ),
-            // stream input doesn't need to be on the same worker node as lookup
-            same_worker_node: false,
             link_id: exchange_a1l0.operator_id,
         },
     );
@@ -255,7 +250,6 @@ fn build_delta_join_inner(
         lookup_1_frag.fragment_id,
         StreamFragmentEdge {
             dispatch_strategy: dispatch_no_shuffle(),
-            same_worker_node: true,
             link_id: exchange_a1l1.operator_id,
         },
     );
@@ -282,7 +276,6 @@ fn build_delta_join_inner(
         current_fragment.fragment_id,
         StreamFragmentEdge {
             dispatch_strategy: dispatch_consistent_hash_shuffle(node.stream_key.clone()),
-            same_worker_node: false,
             link_id: exchange_l0m.operator_id,
         },
     );
@@ -292,7 +285,6 @@ fn build_delta_join_inner(
         current_fragment.fragment_id,
         StreamFragmentEdge {
             dispatch_strategy: dispatch_consistent_hash_shuffle(node.stream_key.clone()),
-            same_worker_node: false,
             link_id: exchange_l1m.operator_id,
         },
     );

src/meta/src/hummock/manager/compaction_group_manager.rs

@@ -734,7 +734,6 @@ mod tests {
     use risingwave_common::constants::hummock::PROPERTIES_RETENTION_SECOND_KEY;
     use risingwave_hummock_sdk::compaction_group::StaticCompactionGroupId;
     use risingwave_pb::meta::table_fragments::Fragment;
-    use risingwave_pb::stream_plan::StreamEnvironment;
 
     use crate::hummock::manager::compaction_group_manager::CompactionGroupManagerInner;
     use crate::hummock::manager::versioning::Versioning;
@@ -864,7 +863,7 @@ mod tests {
     #[tokio::test]
     async fn test_manager() {
         let (_, compaction_group_manager, ..) = setup_compute_env(8080).await;
-        let table_fragment_1 = TableFragments::new(
+        let table_fragment_1 = TableFragments::for_test(
             TableId::new(10),
             BTreeMap::from([(
                 1,
@@ -874,9 +873,8 @@ mod tests {
                     ..Default::default()
                 },
             )]),
-            StreamEnvironment::default(),
         );
-        let table_fragment_2 = TableFragments::new(
+        let table_fragment_2 = TableFragments::for_test(
             TableId::new(20),
             BTreeMap::from([(
                 2,
@@ -886,7 +884,6 @@ mod tests {
                     ..Default::default()
                 },
             )]),
-            StreamEnvironment::default(),
         );
 
         // Test register_table_fragments