risingwave_meta/stream/stream_graph/
schedule.rs

1// Copyright 2023 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::{BTreeMap, HashMap};
16
17use anyhow::Context;
18use enum_as_inner::EnumAsInner;
19use itertools::Itertools;
20use risingwave_common::bail;
21use risingwave_common::hash::{ActorAlignmentId, VnodeCountCompat};
22use risingwave_common::util::stream_graph_visitor::visit_fragment;
23use risingwave_connector::source::cdc::{CDC_BACKFILL_MAX_PARALLELISM, CdcScanOptions};
24use risingwave_meta_model::WorkerId;
25use risingwave_pb::common::WorkerNode;
26use risingwave_pb::meta::table_fragments::fragment::{
27    FragmentDistributionType, PbFragmentDistributionType,
28};
29use risingwave_pb::stream_plan::DispatcherType::{self, *};
30
31use crate::MetaResult;
32use crate::model::{ActorId, Fragment};
33use crate::stream::stream_graph::fragment::CompleteStreamFragmentGraph;
34use crate::stream::stream_graph::id::GlobalFragmentId as Id;
35
36type HashMappingId = usize;
37
38/// The internal structure for processing scheduling requirements in the scheduler.
39#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
40enum Req {
41    /// The fragment must be singleton and is scheduled to the given worker id.
42    Singleton,
43    /// The fragment must be hash-distributed and is scheduled by the given hash mapping.
44    Hash(HashMappingId),
45    /// The fragment must have the given vnode count, but can be scheduled anywhere.
46    /// When the vnode count is 1, it means the fragment must be singleton.
47    AnyVnodeCount(usize),
48}
49
50impl Req {
51    /// Equivalent to `Req::AnyVnodeCount(1)`.
52    #[expect(non_upper_case_globals)]
53    const AnySingleton: Self = Self::AnyVnodeCount(1);
54
55    /// Merge two requirements. Returns an error if the requirements are incompatible.
56    ///
57    /// The `mapping_len` function is used to get the vnode count of a hash mapping by its id.
58    fn merge(a: Self, b: Self, mapping_len: impl Fn(HashMappingId) -> usize) -> MetaResult<Self> {
59        // Note that a and b are always different, as they come from a set.
60        let merge = |a, b| match (a, b) {
61            (Self::AnySingleton, Self::Singleton) => Some(Self::Singleton),
62            (Self::AnyVnodeCount(count), Self::Hash(id)) if mapping_len(id) == count => {
63                Some(Self::Hash(id))
64            }
65            _ => None,
66        };
67
68        match merge(a, b).or_else(|| merge(b, a)) {
69            Some(req) => Ok(req),
70            None => bail!("incompatible requirements `{a:?}` and `{b:?}`"),
71        }
72    }
73}
74
75/// Facts as the input of the scheduler.
76#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
77enum Fact {
78    /// An edge in the fragment graph.
79    Edge {
80        from: Id,
81        to: Id,
82        dt: DispatcherType,
83    },
84    /// A scheduling requirement for a fragment.
85    Req { id: Id, req: Req },
86}
87
88crepe::crepe! {
89    @input
90    struct Input(Fact);
91
92    struct Edge(Id, Id, DispatcherType);
93    struct ExternalReq(Id, Req);
94
95    @output
96    struct Requirement(Id, Req);
97
98    // Extract facts.
99    Edge(from, to, dt) <- Input(f), let Fact::Edge { from, to, dt } = f;
100    Requirement(id, req) <- Input(f), let Fact::Req { id, req } = f;
101
102    // The downstream fragment of a `Simple` edge must be singleton.
103    Requirement(y, Req::AnySingleton) <- Edge(_, y, Simple);
104    // Requirements propagate through `NoShuffle` edges.
105    Requirement(x, d) <- Edge(x, y, NoShuffle), Requirement(y, d);
106    Requirement(y, d) <- Edge(x, y, NoShuffle), Requirement(x, d);
107}
108
109/// The distribution (scheduling result) of a fragment.
110#[derive(Debug, Clone, EnumAsInner)]
111pub(super) enum Distribution {
112    /// The fragment is singleton and is scheduled to the given worker slot.
113    Singleton,
114
115    /// The fragment is hash-distributed and is scheduled by the given hash mapping.
116    Hash(usize),
117}
118
119impl Distribution {
120    /// Get the vnode count of the distribution.
121    pub fn vnode_count(&self) -> usize {
122        match self {
123            Distribution::Singleton => 1, // only `SINGLETON_VNODE`
124            Distribution::Hash(vnode_count) => *vnode_count,
125        }
126    }
127
128    /// Create a distribution from a persisted protobuf `Fragment`.
129    pub fn from_fragment(fragment: &Fragment) -> Self {
130        match fragment.distribution_type {
131            FragmentDistributionType::Single => Distribution::Singleton,
132            FragmentDistributionType::Hash => Distribution::Hash(fragment.vnode_count()),
133            PbFragmentDistributionType::Unspecified => {
134                unreachable!()
135            }
136        }
137    }
138
139    /// Convert the distribution to [`PbFragmentDistributionType`].
140    pub fn to_distribution_type(&self) -> PbFragmentDistributionType {
141        match self {
142            Distribution::Singleton => PbFragmentDistributionType::Single,
143            Distribution::Hash(_) => PbFragmentDistributionType::Hash,
144        }
145    }
146}
147
148/// [`Scheduler`] schedules the distribution of fragments in a stream graph.
149pub(super) struct Scheduler {
150    /// The default hash mapping for hash-distributed fragments, if there's no requirement derived.
151    default_vnode_count: usize,
152}
153
154impl Scheduler {
155    /// Create a new [`Scheduler`] with the expected vnode count of the streaming job.
156    pub fn new(expected_vnode_count: usize) -> MetaResult<Self> {
157        Ok(Self {
158            default_vnode_count: expected_vnode_count,
159        })
160    }
161
162    /// Schedule the given complete graph and returns the distribution of each **building
163    /// fragment**.
164    pub fn schedule(
165        &self,
166        graph: &CompleteStreamFragmentGraph,
167    ) -> MetaResult<HashMap<Id, Distribution>> {
168        let existing_distribution = graph.existing_distribution();
169
170        // Build an index map for all hash mappings.
171        let all_hash_mappings = existing_distribution
172            .values()
173            .flat_map(|dist| dist.as_hash())
174            .cloned()
175            .unique()
176            .collect_vec();
177        let hash_mapping_id: HashMap<_, _> = all_hash_mappings
178            .iter()
179            .enumerate()
180            .map(|(i, m)| (*m, i))
181            .collect();
182
183        let mut facts = Vec::new();
184
185        // Singletons.
186        for (&id, fragment) in graph.building_fragments() {
187            if fragment.requires_singleton {
188                facts.push(Fact::Req {
189                    id,
190                    req: Req::AnySingleton,
191                });
192            }
193        }
194        let mut force_parallelism_fragment_ids: HashMap<_, _> = HashMap::default();
195        // Vnode count requirements: if a fragment is going to look up an existing table,
196        // it must have the same vnode count as that table.
197        for (&id, fragment) in graph.building_fragments() {
198            visit_fragment(fragment, |node| {
199                use risingwave_pb::stream_plan::stream_node::NodeBody;
200                let vnode_count = match node {
201                    NodeBody::StreamScan(node) => {
202                        if let Some(table) = &node.arrangement_table {
203                            table.vnode_count()
204                        } else if let Some(table) = &node.table_desc {
205                            table.vnode_count()
206                        } else {
207                            return;
208                        }
209                    }
210                    NodeBody::TemporalJoin(node) => {
211                        if node.is_broadcast {
212                            // Every actor owns a full replicated view of the lookup table, so the
213                            // join fragment follows the left input rather than the lookup table's
214                            // vnode count.
215                            return;
216                        }
217                        node.get_table_desc().unwrap().vnode_count()
218                    }
219                    NodeBody::BatchPlan(node) => node.get_table_desc().unwrap().vnode_count(),
220                    NodeBody::Lookup(node) => node
221                        .get_arrangement_table_info()
222                        .unwrap()
223                        .get_table_desc()
224                        .unwrap()
225                        .vnode_count(),
226                    NodeBody::StreamCdcScan(node) => {
227                        let Some(ref options) = node.options else {
228                            return;
229                        };
230                        let options = CdcScanOptions::from_proto(options);
231                        if options.is_parallelized_backfill() {
232                            force_parallelism_fragment_ids
233                                .insert(id, options.backfill_parallelism as usize);
234                            CDC_BACKFILL_MAX_PARALLELISM as usize
235                        } else {
236                            return;
237                        }
238                    }
239                    _ => return,
240                };
241                facts.push(Fact::Req {
242                    id,
243                    req: Req::AnyVnodeCount(vnode_count),
244                });
245            });
246        }
247        // Distributions of existing fragments.
248        for (id, dist) in existing_distribution {
249            let req = match dist {
250                Distribution::Singleton => Req::Singleton,
251                Distribution::Hash(mapping) => Req::Hash(hash_mapping_id[&mapping]),
252            };
253            facts.push(Fact::Req { id, req });
254        }
255        // Edges.
256        for (from, to, edge) in graph.all_edges() {
257            facts.push(Fact::Edge {
258                from,
259                to,
260                dt: edge.dispatch_strategy.r#type(),
261            });
262        }
263
264        // Run the algorithm to propagate requirements.
265        let mut crepe = Crepe::new();
266        crepe.extend(facts.into_iter().map(Input));
267        let (reqs,) = crepe.run();
268        let reqs = reqs
269            .into_iter()
270            .map(|Requirement(id, req)| (id, req))
271            .into_group_map();
272
273        // Derive scheduling result from requirements.
274        let mut distributions = HashMap::new();
275        for &id in graph.building_fragments().keys() {
276            let dist = match reqs.get(&id) {
277                // Merge all requirements.
278                Some(reqs) => {
279                    let req = (reqs.iter().copied())
280                        .try_reduce(|a, b| Req::merge(a, b, |id| all_hash_mappings[id]))
281                        .with_context(|| {
282                            format!("cannot fulfill scheduling requirements for fragment {id:?}")
283                        })?
284                        .unwrap();
285
286                    // Derive distribution from the merged requirement.
287                    match req {
288                        Req::Singleton => Distribution::Singleton,
289                        Req::Hash(mapping) => Distribution::Hash(all_hash_mappings[mapping]),
290                        Req::AnySingleton => Distribution::Singleton,
291                        Req::AnyVnodeCount(vnode_count) => Distribution::Hash(vnode_count),
292                    }
293                }
294                // No requirement, use the default.
295                None => Distribution::Hash(self.default_vnode_count),
296            };
297
298            distributions.insert(id, dist);
299        }
300
301        tracing::debug!(?distributions, "schedule fragments");
302
303        Ok(distributions)
304    }
305}
306
307/// [`Locations`] represents the locations of the actors.
308#[cfg_attr(test, derive(Default))]
309pub struct Locations {
310    /// actor location map.
311    pub actor_locations: BTreeMap<ActorId, ActorAlignmentId>,
312    /// worker location map.
313    pub worker_locations: HashMap<WorkerId, WorkerNode>,
314}
315
316#[cfg(test)]
317mod tests {
318    use super::*;
319
320    #[derive(Debug)]
321    enum Result {
322        DefaultHash,
323        Required(Req),
324    }
325
326    impl Result {
327        #[expect(non_upper_case_globals)]
328        const DefaultSingleton: Self = Self::Required(Req::AnySingleton);
329    }
330
331    fn run_and_merge(
332        facts: impl IntoIterator<Item = Fact>,
333        mapping_len: impl Fn(HashMappingId) -> usize,
334    ) -> MetaResult<HashMap<Id, Req>> {
335        let mut crepe = Crepe::new();
336        crepe.extend(facts.into_iter().map(Input));
337        let (reqs,) = crepe.run();
338
339        let reqs = reqs
340            .into_iter()
341            .map(|Requirement(id, req)| (id, req))
342            .into_group_map();
343
344        let mut merged = HashMap::new();
345        for (id, reqs) in reqs {
346            let req = (reqs.iter().copied())
347                .try_reduce(|a, b| Req::merge(a, b, &mapping_len))
348                .with_context(|| {
349                    format!("cannot fulfill scheduling requirements for fragment {id:?}")
350                })?
351                .unwrap();
352            merged.insert(id, req);
353        }
354
355        Ok(merged)
356    }
357
358    fn test_success(facts: impl IntoIterator<Item = Fact>, expected: HashMap<Id, Result>) {
359        test_success_with_mapping_len(facts, expected, |_| 0);
360    }
361
362    fn test_success_with_mapping_len(
363        facts: impl IntoIterator<Item = Fact>,
364        expected: HashMap<Id, Result>,
365        mapping_len: impl Fn(HashMappingId) -> usize,
366    ) {
367        let reqs = run_and_merge(facts, mapping_len).unwrap();
368
369        for (id, expected) in expected {
370            match (reqs.get(&id), expected) {
371                (None, Result::DefaultHash) => {}
372                (Some(actual), Result::Required(expected)) if *actual == expected => {}
373                (actual, expected) => panic!(
374                    "unexpected result for fragment {id:?}\nactual: {actual:?}\nexpected: {expected:?}"
375                ),
376            }
377        }
378    }
379
380    fn test_failed(facts: impl IntoIterator<Item = Fact>) {
381        run_and_merge(facts, |_| 0).unwrap_err();
382    }
383
384    // 101
385    #[test]
386    fn test_single_fragment_hash() {
387        #[rustfmt::skip]
388        let facts = [];
389
390        let expected = maplit::hashmap! {
391            101.into() => Result::DefaultHash,
392        };
393
394        test_success(facts, expected);
395    }
396
397    // 101
398    #[test]
399    fn test_single_fragment_singleton() {
400        #[rustfmt::skip]
401        let facts = [
402            Fact::Req { id: 101.into(), req: Req::AnySingleton },
403        ];
404
405        let expected = maplit::hashmap! {
406            101.into() => Result::DefaultSingleton,
407        };
408
409        test_success(facts, expected);
410    }
411
412    // 1 -|-> 101 -->
413    //                103 --> 104
414    // 2 -|-> 102 -->
415    #[test]
416    fn test_scheduling_mv_on_mv() {
417        #[rustfmt::skip]
418        let facts = [
419            Fact::Req { id: 1.into(), req: Req::Hash(1) },
420            Fact::Req { id: 2.into(), req: Req::Singleton },
421            Fact::Edge { from: 1.into(), to: 101.into(), dt: NoShuffle },
422            Fact::Edge { from: 2.into(), to: 102.into(), dt: NoShuffle },
423            Fact::Edge { from: 101.into(), to: 103.into(), dt: Hash },
424            Fact::Edge { from: 102.into(), to: 103.into(), dt: Hash },
425            Fact::Edge { from: 103.into(), to: 104.into(), dt: Simple },
426        ];
427
428        let expected = maplit::hashmap! {
429            101.into() => Result::Required(Req::Hash(1)),
430            102.into() => Result::Required(Req::Singleton),
431            103.into() => Result::DefaultHash,
432            104.into() => Result::DefaultSingleton,
433        };
434
435        test_success(facts, expected);
436    }
437
438    // 1 -|-> 101 --> 103 -->
439    //             X          105
440    // 2 -|-> 102 --> 104 -->
441    #[test]
442    fn test_delta_join() {
443        #[rustfmt::skip]
444        let facts = [
445            Fact::Req { id: 1.into(), req: Req::Hash(1) },
446            Fact::Req { id: 2.into(), req: Req::Hash(2) },
447            Fact::Edge { from: 1.into(), to: 101.into(), dt: NoShuffle },
448            Fact::Edge { from: 2.into(), to: 102.into(), dt: NoShuffle },
449            Fact::Edge { from: 101.into(), to: 103.into(), dt: NoShuffle },
450            Fact::Edge { from: 102.into(), to: 104.into(), dt: NoShuffle },
451            Fact::Edge { from: 101.into(), to: 104.into(), dt: Hash },
452            Fact::Edge { from: 102.into(), to: 103.into(), dt: Hash },
453            Fact::Edge { from: 103.into(), to: 105.into(), dt: Hash },
454            Fact::Edge { from: 104.into(), to: 105.into(), dt: Hash },
455        ];
456
457        let expected = maplit::hashmap! {
458            101.into() => Result::Required(Req::Hash(1)),
459            102.into() => Result::Required(Req::Hash(2)),
460            103.into() => Result::Required(Req::Hash(1)),
461            104.into() => Result::Required(Req::Hash(2)),
462            105.into() => Result::DefaultHash,
463        };
464
465        test_success(facts, expected);
466    }
467
468    // 1 -|-> 101 -->
469    //                103
470    //        102 -->
471    #[test]
472    fn test_singleton_leaf() {
473        #[rustfmt::skip]
474        let facts = [
475            Fact::Req { id: 1.into(), req: Req::Hash(1) },
476            Fact::Edge { from: 1.into(), to: 101.into(), dt: NoShuffle },
477            Fact::Req { id: 102.into(), req: Req::AnySingleton }, // like `Now`
478            Fact::Edge { from: 101.into(), to: 103.into(), dt: Hash },
479            Fact::Edge { from: 102.into(), to: 103.into(), dt: Broadcast },
480        ];
481
482        let expected = maplit::hashmap! {
483            101.into() => Result::Required(Req::Hash(1)),
484            102.into() => Result::DefaultSingleton,
485            103.into() => Result::DefaultHash,
486        };
487
488        test_success(facts, expected);
489    }
490
491    // 1 -|->
492    //        101
493    // 2 -|->
494    #[test]
495    fn test_upstream_hash_shard_failed() {
496        #[rustfmt::skip]
497        let facts = [
498            Fact::Req { id: 1.into(), req: Req::Hash(1) },
499            Fact::Req { id: 2.into(), req: Req::Hash(2) },
500            Fact::Edge { from: 1.into(), to: 101.into(), dt: NoShuffle },
501            Fact::Edge { from: 2.into(), to: 101.into(), dt: NoShuffle },
502        ];
503
504        test_failed(facts);
505    }
506
507    // 1 -|~> 101
508    #[test]
509    fn test_arrangement_backfill_vnode_count() {
510        #[rustfmt::skip]
511        let facts = [
512            Fact::Req { id: 1.into(), req: Req::Hash(1) },
513            Fact::Req { id: 101.into(), req: Req::AnyVnodeCount(128) },
514            Fact::Edge { from: 1.into(), to: 101.into(), dt: Hash },
515        ];
516
517        let expected = maplit::hashmap! {
518            101.into() => Result::Required(Req::AnyVnodeCount(128)),
519        };
520
521        test_success(facts, expected);
522    }
523
524    // 1 -|~> 101
525    #[test]
526    fn test_no_shuffle_backfill_vnode_count() {
527        #[rustfmt::skip]
528        let facts = [
529            Fact::Req { id: 1.into(), req: Req::Hash(1) },
530            Fact::Req { id: 101.into(), req: Req::AnyVnodeCount(128) },
531            Fact::Edge { from: 1.into(), to: 101.into(), dt: NoShuffle },
532        ];
533
534        let expected = maplit::hashmap! {
535            101.into() => Result::Required(Req::Hash(1)),
536        };
537
538        test_success_with_mapping_len(facts, expected, |id| {
539            assert_eq!(id, 1);
540            128
541        });
542    }
543
544    // 1 -|~> 101
545    #[test]
546    fn test_no_shuffle_backfill_mismatched_vnode_count() {
547        #[rustfmt::skip]
548        let facts = [
549            Fact::Req { id: 1.into(), req: Req::Hash(1) },
550            Fact::Req { id: 101.into(), req: Req::AnyVnodeCount(128) },
551            Fact::Edge { from: 1.into(), to: 101.into(), dt: NoShuffle },
552        ];
553
554        // Not specifying `mapping_len` should fail.
555        test_failed(facts);
556    }
557
558    // 1 -|~> 101
559    #[test]
560    fn test_backfill_singleton_vnode_count() {
561        #[rustfmt::skip]
562        let facts = [
563            Fact::Req { id: 1.into(), req: Req::Singleton },
564            Fact::Req { id: 101.into(), req: Req::AnySingleton },
565            Fact::Edge { from: 1.into(), to: 101.into(), dt: NoShuffle }, // or `Simple`
566        ];
567
568        let expected = maplit::hashmap! {
569            101.into() => Result::Required(Req::Singleton),
570        };
571
572        test_success(facts, expected);
573    }
574}
risingwave_meta/stream/stream_graph/schedule.rs

risingwave_meta/stream/stream_graph/
schedule.rs
