Skip to main content

risingwave_meta/stream/
scale.rs

1// Copyright 2022 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, HashSet};
16use std::sync::Arc;
17use std::time::Duration;
18
19use anyhow::anyhow;
20use futures::future;
21use itertools::Itertools;
22use risingwave_common::bail;
23use risingwave_common::catalog::{DatabaseId, TableId};
24use risingwave_common::hash::{ActorMapping, VnodeBitmapExt};
25use risingwave_connector::source::{SplitId, SplitMetaData};
26use risingwave_meta_model::{
27    StreamingParallelism, WorkerId, fragment, fragment_relation, object, streaming_job,
28};
29use risingwave_pb::common::{WorkerNode, WorkerType};
30use risingwave_pb::stream_plan::{PbDispatchOutputMapping, PbDispatcher};
31use sea_orm::{ConnectionTrait, JoinType, QuerySelect, RelationTrait};
32use thiserror_ext::AsReport;
33use tokio::sync::oneshot::Receiver;
34use tokio::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard, oneshot};
35use tokio::task::JoinHandle;
36use tokio::time::{Instant, MissedTickBehavior};
37
38use crate::barrier::{Command, Reschedule, RescheduleContext, ReschedulePlan};
39use crate::controller::scale::{
40    FragmentRenderMap, LoadedFragmentContext, NoShuffleEnsemble,
41    find_fragment_no_shuffle_dags_detailed, load_fragment_context, load_fragment_context_for_jobs,
42};
43use crate::error::bail_invalid_parameter;
44use crate::manager::{ActiveStreamingWorkerNodes, LocalNotification, MetaSrvEnv, MetadataManager};
45use crate::model::{ActorId, FragmentId, StreamActor, StreamActorWithDispatchers};
46use crate::stream::{GlobalStreamManager, SourceManagerRef};
47use crate::{MetaError, MetaResult};
48
49#[derive(Debug, Clone, Eq, PartialEq)]
50pub struct WorkerReschedule {
51    pub worker_actor_diff: BTreeMap<WorkerId, isize>,
52}
53
54use risingwave_common::id::JobId;
55use risingwave_meta_model::DispatcherType;
56use risingwave_meta_model::fragment::DistributionType;
57use risingwave_meta_model::prelude::{Fragment, FragmentRelation, RefreshJob, StreamingJob};
58use risingwave_meta_model::refresh_job::{self, RefreshState};
59use sea_orm::ActiveValue::Set;
60use sea_orm::{
61    ActiveModelTrait, ColumnTrait, EntityTrait, IntoActiveModel, QueryFilter, TransactionTrait,
62};
63
64use crate::controller::fragment::{InflightActorInfo, InflightFragmentInfo};
65use crate::controller::utils::{
66    StreamingJobExtraInfo, compose_dispatchers, get_streaming_job_extra_info,
67};
68
69pub type ScaleControllerRef = Arc<ScaleController>;
70
71// Field order is intentional: derived `Ord` is used to canonicalize actor
72// layout for no-op detection, so new fields must preserve meaningful
73// comparison order.
74#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
75struct NormalizedActor {
76    worker_id: WorkerId,
77    vnode_bitmap: Option<Vec<usize>>,
78    splits: Vec<SplitId>,
79}
80
81#[derive(Clone, Debug, PartialEq, Eq)]
82struct NormalizedFragmentLayout {
83    distribution_type: DistributionType,
84    actors: Vec<NormalizedActor>,
85}
86
87fn normalize_actor_info(info: &InflightActorInfo) -> NormalizedActor {
88    let vnode_bitmap = info.vnode_bitmap.as_ref().map(|bitmap| {
89        bitmap
90            .iter_vnodes()
91            .map(|vnode| vnode.to_index())
92            .collect_vec()
93    });
94    let mut splits = info.splits.iter().map(SplitMetaData::id).collect_vec();
95    splits.sort_unstable();
96    NormalizedActor {
97        worker_id: info.worker_id,
98        vnode_bitmap,
99        splits,
100    }
101}
102
103fn build_normalized_fragment_layout(
104    fragment_info: &InflightFragmentInfo,
105) -> NormalizedFragmentLayout {
106    // Actor ids are intentionally erased here. For no-op detection we only care
107    // about the semantic layout carried by worker placement, vnode ownership,
108    // and split assignment.
109    let mut actors = fragment_info
110        .actors
111        .values()
112        .map(normalize_actor_info)
113        .collect_vec();
114    actors.sort_unstable();
115
116    NormalizedFragmentLayout {
117        distribution_type: fragment_info.distribution_type,
118        actors,
119    }
120}
121
122/// Compare the rendered (preview) layout against the current in-flight layout
123/// to decide whether a reschedule is a no-op.
124///
125/// Dispatcher layout is intentionally not compared here. On the reschedule path
126/// it is a deterministic function of the fragment actor layout together with the
127/// stable fragment-relation metadata held in `RescheduleContext`, so fragment-level
128/// equality already implies dispatcher equality.
129///
130/// Only fragments present in `render_result` are compared. This is intentional:
131/// this function sits on the reschedule path, which only re-renders fragments
132/// that were loaded into the `RescheduleContext`. Fragment creation or deletion
133/// is handled by separate DDL / recovery paths, not by reschedule, so fragments
134/// outside the render set are irrelevant here. In other words, this is a subset
135/// check over the rendered fragments, not a full bidirectional equality check.
136pub(crate) fn rendered_layout_matches_current(
137    render_result: &FragmentRenderMap,
138    all_prev_fragments: &HashMap<FragmentId, &InflightFragmentInfo>,
139) -> MetaResult<bool> {
140    let all_rendered_fragments: HashMap<_, _> = render_result
141        .values()
142        .flat_map(|jobs| jobs.values())
143        .flatten()
144        .map(|(fragment_id, info)| (*fragment_id, info))
145        .collect();
146
147    for (fragment_id, rendered_fragment) in &all_rendered_fragments {
148        let Some(prev_fragment) = all_prev_fragments.get(fragment_id).copied() else {
149            return Err(MetaError::from(anyhow!(
150                "previous fragment info for {} not found",
151                fragment_id
152            )));
153        };
154
155        let rendered_layout = build_normalized_fragment_layout(rendered_fragment);
156        let current_layout = build_normalized_fragment_layout(prev_fragment);
157
158        if rendered_layout != current_layout {
159            return Ok(false);
160        }
161    }
162
163    Ok(true)
164}
165
166pub struct ScaleController {
167    pub metadata_manager: MetadataManager,
168
169    pub source_manager: SourceManagerRef,
170
171    pub env: MetaSrvEnv,
172
173    /// We will acquire lock during DDL to prevent scaling operations on jobs that are in the creating state.
174    /// e.g., a MV cannot be rescheduled during foreground backfill.
175    pub reschedule_lock: RwLock<()>,
176}
177
178impl ScaleController {
179    pub fn new(
180        metadata_manager: &MetadataManager,
181        source_manager: SourceManagerRef,
182        env: MetaSrvEnv,
183    ) -> Self {
184        Self {
185            metadata_manager: metadata_manager.clone(),
186            source_manager,
187            env,
188            reschedule_lock: RwLock::new(()),
189        }
190    }
191
192    pub async fn resolve_related_no_shuffle_jobs(
193        &self,
194        jobs: &[JobId],
195    ) -> MetaResult<HashSet<JobId>> {
196        let inner = self.metadata_manager.catalog_controller.inner.read().await;
197        let txn = inner.db.begin().await?;
198
199        let fragment_ids: Vec<_> = Fragment::find()
200            .select_only()
201            .column(fragment::Column::FragmentId)
202            .filter(fragment::Column::JobId.is_in(jobs.to_vec()))
203            .into_tuple()
204            .all(&txn)
205            .await?;
206        let ensembles = find_fragment_no_shuffle_dags_detailed(&txn, &fragment_ids).await?;
207        let related_fragments = ensembles
208            .iter()
209            .flat_map(|ensemble| ensemble.fragments())
210            .collect_vec();
211
212        let job_ids: Vec<_> = Fragment::find()
213            .select_only()
214            .column(fragment::Column::JobId)
215            .filter(fragment::Column::FragmentId.is_in(related_fragments))
216            .into_tuple()
217            .all(&txn)
218            .await?;
219
220        let job_ids = job_ids.into_iter().collect();
221
222        Ok(job_ids)
223    }
224
225    pub async fn reschedule_inplace(
226        &self,
227        policy: HashMap<JobId, ReschedulePolicy>,
228    ) -> MetaResult<HashMap<DatabaseId, Command>> {
229        let inner = self.metadata_manager.catalog_controller.inner.write().await;
230        let txn = inner.db.begin().await?;
231        ensure_not_refreshing(&txn, policy.keys().copied()).await?;
232
233        for (table_id, target) in &policy {
234            let streaming_job = StreamingJob::find_by_id(*table_id)
235                .one(&txn)
236                .await?
237                .ok_or_else(|| MetaError::catalog_id_not_found("table", table_id))?;
238
239            let max_parallelism = streaming_job.max_parallelism;
240
241            let mut streaming_job = streaming_job.into_active_model();
242
243            match &target {
244                ReschedulePolicy::Parallelism(p) | ReschedulePolicy::Both(p, _) => {
245                    if let StreamingParallelism::Fixed(n) = p.parallelism
246                        && n > max_parallelism as usize
247                    {
248                        bail!(format!(
249                            "specified parallelism {n} should not exceed max parallelism {max_parallelism}"
250                        ));
251                    }
252
253                    streaming_job.parallelism = Set(p.parallelism.clone());
254                    if matches!(p.parallelism, StreamingParallelism::Fixed(_)) {
255                        streaming_job.adaptive_parallelism_strategy = Set(None);
256                    } else {
257                        streaming_job.adaptive_parallelism_strategy =
258                            Set(p.adaptive_parallelism_strategy.clone());
259                    }
260                }
261                _ => {}
262            }
263
264            match &target {
265                ReschedulePolicy::ResourceGroup(r) | ReschedulePolicy::Both(_, r) => {
266                    streaming_job.specific_resource_group = Set(r.resource_group.clone());
267                }
268                _ => {}
269            }
270
271            StreamingJob::update(streaming_job).exec(&txn).await?;
272        }
273
274        let job_ids: HashSet<JobId> = policy.keys().copied().collect();
275        let commands = build_reschedule_intent_for_jobs(&txn, job_ids).await?;
276
277        txn.commit().await?;
278
279        Ok(commands)
280    }
281
282    pub async fn reschedule_backfill_parallelism_inplace(
283        &self,
284        policy: HashMap<JobId, Option<ParallelismPolicy>>,
285    ) -> MetaResult<HashMap<DatabaseId, Command>> {
286        if policy.is_empty() {
287            return Ok(HashMap::new());
288        }
289
290        let inner = self.metadata_manager.catalog_controller.inner.write().await;
291        let txn = inner.db.begin().await?;
292
293        for (table_id, policy) in &policy {
294            let streaming_job = StreamingJob::find_by_id(*table_id)
295                .one(&txn)
296                .await?
297                .ok_or_else(|| MetaError::catalog_id_not_found("table", table_id))?;
298
299            let max_parallelism = streaming_job.max_parallelism;
300
301            let mut streaming_job = streaming_job.into_active_model();
302
303            if let Some(ParallelismPolicy {
304                parallelism: StreamingParallelism::Fixed(n),
305                ..
306            }) = policy
307                && *n > max_parallelism as usize
308            {
309                bail!(format!(
310                    "specified backfill parallelism {n} should not exceed max parallelism {max_parallelism}"
311                ));
312            }
313
314            if let Some(policy) = policy {
315                streaming_job.backfill_parallelism = Set(Some(policy.parallelism.clone()));
316                if matches!(policy.parallelism, StreamingParallelism::Fixed(_)) {
317                    streaming_job.backfill_adaptive_parallelism_strategy = Set(None);
318                } else {
319                    streaming_job.backfill_adaptive_parallelism_strategy =
320                        Set(policy.adaptive_parallelism_strategy.clone());
321                }
322            } else {
323                streaming_job.backfill_parallelism = Set(None);
324                streaming_job.backfill_adaptive_parallelism_strategy = Set(None);
325            }
326            streaming_job.update(&txn).await?;
327        }
328
329        let jobs = policy.keys().copied().collect();
330
331        let command = build_reschedule_intent_for_jobs(&txn, jobs).await?;
332
333        txn.commit().await?;
334
335        Ok(command)
336    }
337
338    pub async fn reschedule_fragment_inplace(
339        &self,
340        policy: HashMap<risingwave_meta_model::FragmentId, Option<StreamingParallelism>>,
341    ) -> MetaResult<HashMap<DatabaseId, Command>> {
342        if policy.is_empty() {
343            return Ok(HashMap::new());
344        }
345
346        let inner = self.metadata_manager.catalog_controller.inner.write().await;
347        let txn = inner.db.begin().await?;
348
349        let fragment_id_list = policy.keys().copied().collect_vec();
350
351        let existing_fragment_ids: HashSet<_> = Fragment::find()
352            .select_only()
353            .column(fragment::Column::FragmentId)
354            .filter(fragment::Column::FragmentId.is_in(fragment_id_list.clone()))
355            .into_tuple::<risingwave_meta_model::FragmentId>()
356            .all(&txn)
357            .await?
358            .into_iter()
359            .collect();
360
361        if let Some(missing_fragment_id) = fragment_id_list
362            .iter()
363            .find(|fragment_id| !existing_fragment_ids.contains(*fragment_id))
364        {
365            return Err(MetaError::catalog_id_not_found(
366                "fragment",
367                *missing_fragment_id,
368            ));
369        }
370
371        let mut target_ensembles = vec![];
372
373        for ensemble in find_fragment_no_shuffle_dags_detailed(&txn, &fragment_id_list).await? {
374            let entry_fragment_ids = ensemble.entry_fragments().collect_vec();
375
376            let desired_parallelism = match entry_fragment_ids
377                .iter()
378                .filter_map(|fragment_id| policy.get(fragment_id).cloned())
379                .dedup()
380                .collect_vec()
381                .as_slice()
382            {
383                [] => {
384                    bail_invalid_parameter!(
385                        "none of the entry fragments {:?} were included in the reschedule request; \
386                         provide at least one entry fragment id",
387                        entry_fragment_ids
388                    );
389                }
390                [parallelism] => parallelism.clone(),
391                parallelisms => {
392                    bail!(
393                        "conflicting reschedule policies for fragments in the same no-shuffle ensemble: {:?}",
394                        parallelisms
395                    );
396                }
397            };
398
399            let fragments = Fragment::find()
400                .filter(fragment::Column::FragmentId.is_in(entry_fragment_ids))
401                .all(&txn)
402                .await?;
403
404            debug_assert!(
405                fragments
406                    .iter()
407                    .map(|fragment| fragment.parallelism.as_ref())
408                    .all_equal(),
409                "entry fragments in the same ensemble should share the same parallelism"
410            );
411
412            let current_parallelism = fragments
413                .first()
414                .and_then(|fragment| fragment.parallelism.clone());
415
416            if current_parallelism == desired_parallelism {
417                continue;
418            }
419
420            for fragment in fragments {
421                let mut fragment = fragment.into_active_model();
422                fragment.parallelism = Set(desired_parallelism.clone());
423                Fragment::update(fragment).exec(&txn).await?;
424            }
425
426            target_ensembles.push(ensemble);
427        }
428
429        if target_ensembles.is_empty() {
430            txn.commit().await?;
431            return Ok(HashMap::new());
432        }
433
434        let target_fragment_ids: HashSet<FragmentId> = target_ensembles
435            .iter()
436            .flat_map(|ensemble| ensemble.component_fragments())
437            .collect();
438        let job_ids: Vec<JobId> = Fragment::find()
439            .select_only()
440            .column(fragment::Column::JobId)
441            .filter(fragment::Column::FragmentId.is_in(target_fragment_ids.iter().copied()))
442            .into_tuple()
443            .all(&txn)
444            .await?;
445        ensure_not_refreshing(&txn, job_ids).await?;
446        let commands = build_reschedule_intent_for_fragments(&txn, target_fragment_ids).await?;
447
448        txn.commit().await?;
449
450        Ok(commands)
451    }
452
453    async fn rerender(&self, jobs: HashSet<JobId>) -> MetaResult<HashMap<DatabaseId, Command>> {
454        if jobs.is_empty() {
455            return Ok(HashMap::new());
456        }
457
458        let inner = self.metadata_manager.catalog_controller.inner.read().await;
459        let txn = inner.db.begin().await?;
460        let commands = build_reschedule_intent_for_jobs(&txn, jobs).await?;
461        txn.commit().await?;
462        Ok(commands)
463    }
464}
465
466/// Rescheduling a job while its table is being refreshed would change the actor set the cycle
467/// is tracked against.
468async fn ensure_not_refreshing(
469    txn: &impl ConnectionTrait,
470    job_ids: impl IntoIterator<Item = JobId>,
471) -> MetaResult<()> {
472    let table_ids = job_ids
473        .into_iter()
474        .map(|job_id| job_id.as_mv_table_id())
475        .collect_vec();
476    let refreshing: Vec<TableId> = RefreshJob::find()
477        .select_only()
478        .column(refresh_job::Column::TableId)
479        .filter(refresh_job::Column::TableId.is_in(table_ids))
480        .filter(refresh_job::Column::CurrentStatus.ne(RefreshState::Idle))
481        .into_tuple()
482        .all(txn)
483        .await?;
484    if let Some(table_id) = refreshing.first() {
485        bail!(
486            "Cannot reschedule job {} because its table is being refreshed",
487            table_id.as_job_id()
488        );
489    }
490    Ok(())
491}
492
493async fn build_reschedule_intent_for_jobs(
494    txn: &impl ConnectionTrait,
495    job_ids: HashSet<JobId>,
496) -> MetaResult<HashMap<DatabaseId, Command>> {
497    if job_ids.is_empty() {
498        return Ok(HashMap::new());
499    }
500
501    let job_id_list = job_ids.iter().copied().collect_vec();
502    let database_jobs: Vec<(DatabaseId, JobId)> = StreamingJob::find()
503        .select_only()
504        .column(object::Column::DatabaseId)
505        .column(streaming_job::Column::JobId)
506        .join(JoinType::LeftJoin, streaming_job::Relation::Object.def())
507        .filter(streaming_job::Column::JobId.is_in(job_id_list.clone()))
508        .into_tuple()
509        .all(txn)
510        .await?;
511
512    if database_jobs.len() != job_ids.len() {
513        let returned_jobs: HashSet<JobId> =
514            database_jobs.iter().map(|(_, job_id)| *job_id).collect();
515        let missing = job_ids.difference(&returned_jobs).copied().collect_vec();
516        return Err(MetaError::catalog_id_not_found(
517            "streaming job",
518            format!("{missing:?}"),
519        ));
520    }
521
522    let reschedule_context = load_reschedule_context_for_jobs(txn, job_ids).await?;
523    if reschedule_context.is_empty() {
524        return Ok(HashMap::new());
525    }
526
527    let commands = reschedule_context
528        .into_database_contexts()
529        .into_iter()
530        .map(|(database_id, context)| {
531            (
532                database_id,
533                Command::RescheduleIntent {
534                    context,
535                    reschedule_plan: None,
536                },
537            )
538        })
539        .collect();
540
541    Ok(commands)
542}
543
544async fn build_reschedule_intent_for_fragments(
545    txn: &impl ConnectionTrait,
546    fragment_ids: HashSet<FragmentId>,
547) -> MetaResult<HashMap<DatabaseId, Command>> {
548    if fragment_ids.is_empty() {
549        return Ok(HashMap::new());
550    }
551
552    let fragment_id_list = fragment_ids.iter().copied().collect_vec();
553    let fragment_databases: Vec<(FragmentId, DatabaseId)> = Fragment::find()
554        .select_only()
555        .column(fragment::Column::FragmentId)
556        .column(object::Column::DatabaseId)
557        .join(JoinType::LeftJoin, fragment::Relation::Object.def())
558        .filter(fragment::Column::FragmentId.is_in(fragment_id_list.clone()))
559        .into_tuple()
560        .all(txn)
561        .await?;
562
563    if fragment_databases.len() != fragment_ids.len() {
564        let returned: HashSet<FragmentId> = fragment_databases
565            .iter()
566            .map(|(fragment_id, _)| *fragment_id)
567            .collect();
568        let missing = fragment_ids.difference(&returned).copied().collect_vec();
569        return Err(MetaError::catalog_id_not_found(
570            "fragment",
571            format!("{missing:?}"),
572        ));
573    }
574
575    let ensembles = find_fragment_no_shuffle_dags_detailed(txn, &fragment_id_list).await?;
576    let reschedule_context = load_reschedule_context_for_ensembles(txn, ensembles).await?;
577    if reschedule_context.is_empty() {
578        return Ok(HashMap::new());
579    }
580
581    let commands = reschedule_context
582        .into_database_contexts()
583        .into_iter()
584        .map(|(database_id, context)| {
585            (
586                database_id,
587                Command::RescheduleIntent {
588                    context,
589                    reschedule_plan: None,
590                },
591            )
592        })
593        .collect();
594
595    Ok(commands)
596}
597
598async fn load_reschedule_context_for_jobs(
599    txn: &impl ConnectionTrait,
600    job_ids: HashSet<JobId>,
601) -> MetaResult<RescheduleContext> {
602    let loaded = load_fragment_context_for_jobs(txn, job_ids).await?;
603    build_reschedule_context_from_loaded(txn, loaded).await
604}
605
606async fn load_reschedule_context_for_ensembles(
607    txn: &impl ConnectionTrait,
608    ensembles: Vec<NoShuffleEnsemble>,
609) -> MetaResult<RescheduleContext> {
610    let loaded = load_fragment_context(txn, ensembles).await?;
611    build_reschedule_context_from_loaded(txn, loaded).await
612}
613
614async fn build_reschedule_context_from_loaded(
615    txn: &impl ConnectionTrait,
616    loaded: LoadedFragmentContext,
617) -> MetaResult<RescheduleContext> {
618    if loaded.is_empty() {
619        return Ok(RescheduleContext::empty());
620    }
621
622    let job_ids = loaded.job_map.keys().copied().collect_vec();
623    let job_extra_info = get_streaming_job_extra_info(txn, job_ids).await?;
624
625    let fragment_ids = loaded
626        .job_fragments
627        .values()
628        .flat_map(|fragments| fragments.keys().copied())
629        .collect_vec();
630
631    let upstreams: Vec<(FragmentId, FragmentId, DispatcherType)> = FragmentRelation::find()
632        .select_only()
633        .columns([
634            fragment_relation::Column::TargetFragmentId,
635            fragment_relation::Column::SourceFragmentId,
636            fragment_relation::Column::DispatcherType,
637        ])
638        .filter(fragment_relation::Column::TargetFragmentId.is_in(fragment_ids.clone()))
639        .into_tuple()
640        .all(txn)
641        .await?;
642
643    let mut upstream_fragments = HashMap::new();
644    for (fragment, upstream, dispatcher) in upstreams {
645        upstream_fragments
646            .entry(fragment as FragmentId)
647            .or_insert(HashMap::new())
648            .insert(upstream as FragmentId, dispatcher);
649    }
650
651    let downstreams = FragmentRelation::find()
652        .filter(fragment_relation::Column::SourceFragmentId.is_in(fragment_ids.clone()))
653        .all(txn)
654        .await?;
655
656    let mut downstream_fragments = HashMap::new();
657    let mut downstream_relations = HashMap::new();
658    for relation in downstreams {
659        let source_fragment_id = relation.source_fragment_id as FragmentId;
660        let target_fragment_id = relation.target_fragment_id as FragmentId;
661        downstream_fragments
662            .entry(source_fragment_id)
663            .or_insert(HashMap::new())
664            .insert(target_fragment_id, relation.dispatcher_type);
665        downstream_relations.insert((source_fragment_id, target_fragment_id), relation);
666    }
667
668    Ok(RescheduleContext {
669        loaded,
670        job_extra_info,
671        upstream_fragments,
672        downstream_fragments,
673        downstream_relations,
674    })
675}
676
677/// Build a `Reschedule` by diffing the previously materialized fragment state against
678/// the newly rendered actor layout.
679///
680/// This function assumes a full rebuild (no kept actors) and produces:
681/// - actor additions/removals and vnode bitmap updates
682/// - dispatcher updates for upstream/downstream fragments
683/// - updated split assignments for source actors
684///
685/// `upstream_fragments`/`downstream_fragments` describe neighbor fragments and dispatcher types,
686/// while `all_actor_dispatchers` contains the new dispatcher list for each actor. `job_extra_info`
687/// supplies job-level context for building new actors.
688fn diff_fragment(
689    prev_fragment_info: &InflightFragmentInfo,
690    curr_actors: &HashMap<ActorId, InflightActorInfo>,
691    upstream_fragments: HashMap<FragmentId, DispatcherType>,
692    downstream_fragments: HashMap<FragmentId, DispatcherType>,
693    all_actor_dispatchers: HashMap<ActorId, Vec<PbDispatcher>>,
694    job_extra_info: Option<&StreamingJobExtraInfo>,
695) -> MetaResult<Reschedule> {
696    let prev_ids: HashSet<_> = prev_fragment_info.actors.keys().cloned().collect();
697    let curr_ids: HashSet<_> = curr_actors.keys().cloned().collect();
698
699    let removed_actors: HashSet<_> = &prev_ids - &curr_ids;
700    let added_actor_ids: HashSet<_> = &curr_ids - &prev_ids;
701    let kept_ids: HashSet<_> = prev_ids.intersection(&curr_ids).cloned().collect();
702    debug_assert!(
703        kept_ids.is_empty(),
704        "kept actors found in scale; expected full rebuild, prev={prev_ids:?}, curr={curr_ids:?}, kept={kept_ids:?}"
705    );
706
707    let mut added_actors = HashMap::new();
708    for &actor_id in &added_actor_ids {
709        let InflightActorInfo { worker_id, .. } = curr_actors
710            .get(&actor_id)
711            .ok_or_else(|| anyhow!("BUG: Worker not found for new actor {}", actor_id))?;
712
713        added_actors
714            .entry(*worker_id)
715            .or_insert_with(Vec::new)
716            .push(actor_id);
717    }
718
719    let mut vnode_bitmap_updates = HashMap::new();
720    for actor_id in kept_ids {
721        let prev_actor = &prev_fragment_info.actors[&actor_id];
722        let curr_actor = &curr_actors[&actor_id];
723
724        // Check if the vnode distribution has changed.
725        if prev_actor.vnode_bitmap != curr_actor.vnode_bitmap
726            && let Some(bitmap) = curr_actor.vnode_bitmap.clone()
727        {
728            vnode_bitmap_updates.insert(actor_id, bitmap);
729        }
730    }
731
732    let upstream_dispatcher_mapping =
733        if let DistributionType::Hash = prev_fragment_info.distribution_type {
734            let actor_mapping = curr_actors
735                .iter()
736                .map(
737                    |(
738                        actor_id,
739                        InflightActorInfo {
740                            worker_id: _,
741                            vnode_bitmap,
742                            ..
743                        },
744                    )| { (*actor_id, vnode_bitmap.clone().unwrap()) },
745                )
746                .collect();
747            Some(ActorMapping::from_bitmaps(&actor_mapping))
748        } else {
749            None
750        };
751
752    let upstream_fragment_dispatcher_ids = upstream_fragments
753        .iter()
754        .filter(|&(_, dispatcher_type)| *dispatcher_type != DispatcherType::NoShuffle)
755        .map(|(upstream_fragment, _)| (*upstream_fragment, prev_fragment_info.fragment_id))
756        .collect();
757
758    let downstream_fragment_ids = downstream_fragments
759        .iter()
760        .filter(|&(_, dispatcher_type)| *dispatcher_type != DispatcherType::NoShuffle)
761        .map(|(fragment_id, _)| *fragment_id)
762        .collect();
763
764    let extra_info = job_extra_info.cloned().unwrap_or_default();
765    let expr_context = extra_info.stream_context().to_expr_context();
766    let job_definition = extra_info.job_definition;
767    let config_override = extra_info.config_override;
768
769    let newly_created_actors: HashMap<ActorId, (StreamActorWithDispatchers, WorkerId)> =
770        added_actor_ids
771            .iter()
772            .map(|actor_id| {
773                let actor = StreamActor {
774                    actor_id: *actor_id,
775                    fragment_id: prev_fragment_info.fragment_id,
776                    vnode_bitmap: curr_actors[actor_id].vnode_bitmap.clone(),
777                    mview_definition: job_definition.clone(),
778                    expr_context: Some(expr_context.clone()),
779                    config_override: config_override.clone(),
780                };
781                (
782                    *actor_id,
783                    (
784                        (
785                            actor,
786                            all_actor_dispatchers
787                                .get(actor_id)
788                                .cloned()
789                                .unwrap_or_default(),
790                        ),
791                        curr_actors[actor_id].worker_id,
792                    ),
793                )
794            })
795            .collect();
796
797    let actor_splits = curr_actors
798        .iter()
799        .map(|(&actor_id, info)| (actor_id, info.splits.clone()))
800        .collect();
801
802    let reschedule = Reschedule {
803        added_actors,
804        removed_actors,
805        vnode_bitmap_updates,
806        upstream_fragment_dispatcher_ids,
807        upstream_dispatcher_mapping,
808        downstream_fragment_ids,
809        actor_splits,
810        newly_created_actors,
811    };
812
813    Ok(reschedule)
814}
815
816/// Build executable reschedule plans from the rendered fragment layout.
817///
818/// Callers are expected to run the no-op layout check before invoking this
819/// function so command construction stays focused on plan materialization.
820pub(crate) fn build_reschedule_commands(
821    render_result: FragmentRenderMap,
822    context: RescheduleContext,
823    all_prev_fragments: HashMap<FragmentId, &InflightFragmentInfo>,
824) -> MetaResult<HashMap<DatabaseId, ReschedulePlan>> {
825    if render_result.is_empty() {
826        return Ok(HashMap::new());
827    }
828
829    let RescheduleContext {
830        job_extra_info,
831        upstream_fragments: mut all_upstream_fragments,
832        downstream_fragments: mut all_downstream_fragments,
833        mut downstream_relations,
834        ..
835    } = context;
836
837    let fragment_ids = render_result
838        .values()
839        .flat_map(|jobs| jobs.values())
840        .flatten()
841        .map(|(fragment_id, _)| *fragment_id)
842        .collect_vec();
843
844    let all_related_fragment_ids: HashSet<_> = fragment_ids
845        .iter()
846        .copied()
847        .chain(all_upstream_fragments.values().flatten().map(|(id, _)| *id))
848        .chain(
849            all_downstream_fragments
850                .values()
851                .flatten()
852                .map(|(id, _)| *id),
853        )
854        .collect();
855
856    for fragment_id in all_related_fragment_ids {
857        if !all_prev_fragments.contains_key(&fragment_id) {
858            return Err(MetaError::from(anyhow!(
859                "previous fragment info for {fragment_id} not found"
860            )));
861        }
862    }
863
864    let all_rendered_fragments: HashMap<_, _> = render_result
865        .values()
866        .flat_map(|jobs| jobs.values())
867        .flatten()
868        .map(|(fragment_id, info)| (*fragment_id, info))
869        .collect();
870
871    let mut commands = HashMap::new();
872
873    for (database_id, jobs) in &render_result {
874        let mut all_fragment_actors = HashMap::new();
875        let mut reschedules = HashMap::new();
876
877        for (job_id, fragment_id, fragment_info) in jobs.iter().flat_map(|(job_id, fragments)| {
878            fragments
879                .iter()
880                .map(move |(fragment_id, info)| (job_id, fragment_id, info))
881        }) {
882            let InflightFragmentInfo {
883                distribution_type,
884                actors,
885                ..
886            } = fragment_info;
887
888            let upstream_fragments = all_upstream_fragments
889                .remove(&(*fragment_id as FragmentId))
890                .unwrap_or_default();
891            let downstream_fragments = all_downstream_fragments
892                .remove(&(*fragment_id as FragmentId))
893                .unwrap_or_default();
894
895            let fragment_actors: HashMap<_, _> = upstream_fragments
896                .keys()
897                .copied()
898                .chain(downstream_fragments.keys().copied())
899                .map(|fragment_id| {
900                    all_prev_fragments
901                        .get(&fragment_id)
902                        .map(|fragment| {
903                            (
904                                fragment_id,
905                                fragment.actors.keys().copied().collect::<HashSet<_>>(),
906                            )
907                        })
908                        .ok_or_else(|| {
909                            MetaError::from(anyhow!(
910                                "fragment {} not found in previous state",
911                                fragment_id
912                            ))
913                        })
914                })
915                .collect::<MetaResult<_>>()?;
916
917            all_fragment_actors.extend(fragment_actors);
918
919            let source_fragment_actors = actors
920                .iter()
921                .map(|(actor_id, info)| (*actor_id, info.vnode_bitmap.clone()))
922                .collect();
923
924            let mut all_actor_dispatchers: HashMap<_, Vec<_>> = HashMap::new();
925
926            for downstream_fragment_id in downstream_fragments.keys() {
927                let target_fragment_actors =
928                    match all_rendered_fragments.get(downstream_fragment_id) {
929                        None => {
930                            let external_fragment = all_prev_fragments
931                                .get(downstream_fragment_id)
932                                .ok_or_else(|| {
933                                    MetaError::from(anyhow!(
934                                        "fragment {} not found in previous state",
935                                        downstream_fragment_id
936                                    ))
937                                })?;
938
939                            external_fragment
940                                .actors
941                                .iter()
942                                .map(|(actor_id, info)| (*actor_id, info.vnode_bitmap.clone()))
943                                .collect()
944                        }
945                        Some(downstream_rendered) => downstream_rendered
946                            .actors
947                            .iter()
948                            .map(|(actor_id, info)| (*actor_id, info.vnode_bitmap.clone()))
949                            .collect(),
950                    };
951
952                let target_fragment_distribution = *distribution_type;
953
954                let fragment_relation::Model {
955                    source_fragment_id: _,
956                    target_fragment_id: _,
957                    dispatcher_type,
958                    dist_key_indices,
959                    output_indices,
960                    output_type_mapping,
961                } = downstream_relations
962                    .remove(&(
963                        *fragment_id as FragmentId,
964                        *downstream_fragment_id as FragmentId,
965                    ))
966                    .ok_or_else(|| {
967                        MetaError::from(anyhow!(
968                            "downstream relation missing for {} -> {}",
969                            fragment_id,
970                            downstream_fragment_id
971                        ))
972                    })?;
973
974                let pb_mapping = PbDispatchOutputMapping {
975                    indices: output_indices.into_u32_array(),
976                    types: output_type_mapping.unwrap_or_default().to_protobuf(),
977                };
978
979                let (dispatchers, _) = compose_dispatchers(
980                    *distribution_type,
981                    &source_fragment_actors,
982                    *downstream_fragment_id,
983                    target_fragment_distribution,
984                    &target_fragment_actors,
985                    dispatcher_type,
986                    dist_key_indices.into_u32_array(),
987                    pb_mapping,
988                );
989
990                for (actor_id, dispatcher) in dispatchers {
991                    all_actor_dispatchers
992                        .entry(actor_id)
993                        .or_default()
994                        .push(dispatcher);
995                }
996            }
997
998            let prev_fragment = all_prev_fragments.get(&{ *fragment_id }).ok_or_else(|| {
999                MetaError::from(anyhow!(
1000                    "fragment {} not found in previous state",
1001                    fragment_id
1002                ))
1003            })?;
1004
1005            let reschedule = diff_fragment(
1006                prev_fragment,
1007                actors,
1008                upstream_fragments,
1009                downstream_fragments,
1010                all_actor_dispatchers,
1011                job_extra_info.get(job_id),
1012            )?;
1013
1014            reschedules.insert(*fragment_id as FragmentId, reschedule);
1015        }
1016
1017        let command = ReschedulePlan {
1018            reschedules,
1019            fragment_actors: all_fragment_actors,
1020        };
1021
1022        debug_assert!(
1023            command
1024                .reschedules
1025                .values()
1026                .all(|reschedule| reschedule.vnode_bitmap_updates.is_empty()),
1027            "reschedule plan carries vnode_bitmap_updates, expected full rebuild"
1028        );
1029
1030        commands.insert(*database_id, command);
1031    }
1032
1033    Ok(commands)
1034}
1035
1036#[cfg(test)]
1037mod tests {
1038    use std::collections::{HashMap, HashSet};
1039
1040    use risingwave_common::bitmap::Bitmap;
1041    use risingwave_common::catalog::FragmentTypeMask;
1042    use risingwave_connector::source::SplitImpl;
1043    use risingwave_connector::source::test_source::TestSourceSplit;
1044
1045    use super::*;
1046
1047    fn actor(worker_id: impl Into<WorkerId>, vnode_bitmap: Option<Bitmap>) -> InflightActorInfo {
1048        InflightActorInfo {
1049            worker_id: worker_id.into(),
1050            vnode_bitmap,
1051            splits: vec![],
1052        }
1053    }
1054
1055    fn actor_with_splits(
1056        worker_id: impl Into<WorkerId>,
1057        vnode_bitmap: Option<Bitmap>,
1058        splits: Vec<SplitImpl>,
1059    ) -> InflightActorInfo {
1060        InflightActorInfo {
1061            worker_id: worker_id.into(),
1062            vnode_bitmap,
1063            splits,
1064        }
1065    }
1066
1067    fn test_split(id: &str, offset: &str) -> SplitImpl {
1068        SplitImpl::Test(TestSourceSplit {
1069            id: id.into(),
1070            properties: HashMap::new(),
1071            offset: offset.to_owned(),
1072        })
1073    }
1074
1075    fn fragment(
1076        fragment_id: FragmentId,
1077        actors: impl IntoIterator<Item = (ActorId, InflightActorInfo)>,
1078    ) -> InflightFragmentInfo {
1079        InflightFragmentInfo {
1080            fragment_id,
1081            distribution_type: DistributionType::Hash,
1082            fragment_type_mask: FragmentTypeMask::empty(),
1083            vnode_count: 8,
1084            nodes: Default::default(),
1085            actors: actors.into_iter().collect(),
1086            state_table_ids: HashSet::new(),
1087        }
1088    }
1089
1090    fn render_result(fragments: Vec<(FragmentId, InflightFragmentInfo)>) -> FragmentRenderMap {
1091        HashMap::from([(
1092            DatabaseId::new(1),
1093            HashMap::from([(
1094                JobId::new(1),
1095                fragments.into_iter().collect::<HashMap<_, _>>(),
1096            )]),
1097        )])
1098    }
1099
1100    #[test]
1101    fn rendered_layout_matches_current_ignores_actor_ids_across_fragments() {
1102        let current_source = fragment(
1103            FragmentId::new(1),
1104            [
1105                (
1106                    ActorId::new(101),
1107                    actor(1, Some(Bitmap::from_indices(8, [0, 1, 2, 3]))),
1108                ),
1109                (
1110                    ActorId::new(102),
1111                    actor(2, Some(Bitmap::from_indices(8, [4, 5, 6, 7]))),
1112                ),
1113            ],
1114        );
1115        let current_target = fragment(
1116            FragmentId::new(2),
1117            [
1118                (
1119                    ActorId::new(201),
1120                    actor(3, Some(Bitmap::from_indices(8, [0, 1, 2, 3]))),
1121                ),
1122                (
1123                    ActorId::new(202),
1124                    actor(4, Some(Bitmap::from_indices(8, [4, 5, 6, 7]))),
1125                ),
1126            ],
1127        );
1128
1129        let rendered = render_result(vec![
1130            (
1131                FragmentId::new(1),
1132                fragment(
1133                    FragmentId::new(1),
1134                    [
1135                        (
1136                            ActorId::new(1001),
1137                            actor(1, Some(Bitmap::from_indices(8, [0, 1, 2, 3]))),
1138                        ),
1139                        (
1140                            ActorId::new(1002),
1141                            actor(2, Some(Bitmap::from_indices(8, [4, 5, 6, 7]))),
1142                        ),
1143                    ],
1144                ),
1145            ),
1146            (
1147                FragmentId::new(2),
1148                fragment(
1149                    FragmentId::new(2),
1150                    [
1151                        (
1152                            ActorId::new(2001),
1153                            actor(3, Some(Bitmap::from_indices(8, [0, 1, 2, 3]))),
1154                        ),
1155                        (
1156                            ActorId::new(2002),
1157                            actor(4, Some(Bitmap::from_indices(8, [4, 5, 6, 7]))),
1158                        ),
1159                    ],
1160                ),
1161            ),
1162        ]);
1163
1164        let prev = HashMap::from([
1165            (FragmentId::new(1), &current_source),
1166            (FragmentId::new(2), &current_target),
1167        ]);
1168
1169        assert!(rendered_layout_matches_current(&rendered, &prev,).unwrap());
1170    }
1171
1172    #[test]
1173    fn rendered_layout_matches_current_detects_target_fragment_layout_changes() {
1174        let current_source = fragment(
1175            FragmentId::new(1),
1176            [
1177                (
1178                    ActorId::new(101),
1179                    actor(1, Some(Bitmap::from_indices(8, [0, 1, 2, 3]))),
1180                ),
1181                (
1182                    ActorId::new(102),
1183                    actor(2, Some(Bitmap::from_indices(8, [4, 5, 6, 7]))),
1184                ),
1185            ],
1186        );
1187        let current_target = fragment(
1188            FragmentId::new(2),
1189            [
1190                (
1191                    ActorId::new(201),
1192                    actor(3, Some(Bitmap::from_indices(8, [0, 1, 2, 3]))),
1193                ),
1194                (
1195                    ActorId::new(202),
1196                    actor(4, Some(Bitmap::from_indices(8, [4, 5, 6, 7]))),
1197                ),
1198            ],
1199        );
1200
1201        let rendered = render_result(vec![
1202            (
1203                FragmentId::new(1),
1204                fragment(
1205                    FragmentId::new(1),
1206                    [
1207                        (
1208                            ActorId::new(1001),
1209                            actor(1, Some(Bitmap::from_indices(8, [0, 1, 2, 3]))),
1210                        ),
1211                        (
1212                            ActorId::new(1002),
1213                            actor(2, Some(Bitmap::from_indices(8, [4, 5, 6, 7]))),
1214                        ),
1215                    ],
1216                ),
1217            ),
1218            (
1219                FragmentId::new(2),
1220                fragment(
1221                    FragmentId::new(2),
1222                    [
1223                        (
1224                            ActorId::new(2001),
1225                            actor(3, Some(Bitmap::from_indices(8, [0, 1, 2, 3, 4, 5]))),
1226                        ),
1227                        (
1228                            ActorId::new(2002),
1229                            actor(4, Some(Bitmap::from_indices(8, [6, 7]))),
1230                        ),
1231                    ],
1232                ),
1233            ),
1234        ]);
1235
1236        let prev = HashMap::from([
1237            (FragmentId::new(1), &current_source),
1238            (FragmentId::new(2), &current_target),
1239        ]);
1240
1241        assert!(!rendered_layout_matches_current(&rendered, &prev,).unwrap());
1242    }
1243
1244    #[test]
1245    fn rendered_layout_matches_current_ignores_split_offsets() {
1246        let fragment_id = FragmentId::new(1);
1247        let prev_fragment = fragment(
1248            fragment_id,
1249            [(
1250                ActorId::new(101),
1251                actor_with_splits(
1252                    1,
1253                    Some(Bitmap::from_indices(8, [0, 1, 2, 3])),
1254                    vec![test_split("split-0", "100")],
1255                ),
1256            )],
1257        );
1258        let rendered = render_result(vec![(
1259            fragment_id,
1260            fragment(
1261                fragment_id,
1262                [(
1263                    ActorId::new(1001),
1264                    actor_with_splits(
1265                        1,
1266                        Some(Bitmap::from_indices(8, [0, 1, 2, 3])),
1267                        vec![test_split("split-0", "200")],
1268                    ),
1269                )],
1270            ),
1271        )]);
1272        let prev = HashMap::from([(fragment_id, &prev_fragment)]);
1273
1274        assert!(rendered_layout_matches_current(&rendered, &prev,).unwrap());
1275    }
1276}
1277
1278#[derive(Clone, Debug, Eq, PartialEq)]
1279pub struct ParallelismPolicy {
1280    pub parallelism: StreamingParallelism,
1281    pub adaptive_parallelism_strategy: Option<String>,
1282}
1283
1284#[derive(Clone, Debug)]
1285pub struct ResourceGroupPolicy {
1286    pub resource_group: Option<String>,
1287}
1288
1289#[derive(Clone, Debug)]
1290pub enum ReschedulePolicy {
1291    Parallelism(ParallelismPolicy),
1292    ResourceGroup(ResourceGroupPolicy),
1293    Both(ParallelismPolicy, ResourceGroupPolicy),
1294}
1295
1296impl GlobalStreamManager {
1297    #[await_tree::instrument("acquire_reschedule_read_guard")]
1298    pub async fn reschedule_lock_read_guard(&self) -> RwLockReadGuard<'_, ()> {
1299        self.scale_controller.reschedule_lock.read().await
1300    }
1301
1302    #[await_tree::instrument("acquire_reschedule_write_guard")]
1303    pub async fn reschedule_lock_write_guard(&self) -> RwLockWriteGuard<'_, ()> {
1304        self.scale_controller.reschedule_lock.write().await
1305    }
1306
1307    /// When new worker nodes joined, or the parallelism of existing worker nodes changed,
1308    /// examines if there are any jobs can be scaled, and scales them if found.
1309    ///
1310    /// This method will iterate over all `CREATED` jobs, and can be repeatedly called.
1311    ///
1312    /// Returns
1313    /// - `Ok(false)` if no jobs can be scaled;
1314    /// - `Ok(true)` if some jobs are scaled, and it is possible that there are more jobs can be scaled.
1315    async fn trigger_parallelism_control(&self) -> MetaResult<bool> {
1316        tracing::info!("trigger parallelism control");
1317
1318        let _reschedule_job_lock = self.reschedule_lock_write_guard().await;
1319
1320        let creating_streaming_jobs = self.metadata_manager.list_creating_jobs().await?;
1321
1322        let mut blocked_jobs = self
1323            .metadata_manager
1324            .collect_reschedule_blocked_jobs_for_creating_jobs(&creating_streaming_jobs, true)
1325            .await?;
1326        // A refreshing table is skipped like a job blocked by a creating one and picked up by a
1327        // later pass once its cycle ends.
1328        blocked_jobs.extend(
1329            self.metadata_manager
1330                .list_refresh_jobs()
1331                .await?
1332                .into_iter()
1333                .filter(|job| job.current_status != RefreshState::Idle)
1334                .map(|job| job.table_id.as_job_id()),
1335        );
1336        let has_blocked_jobs = !blocked_jobs.is_empty();
1337
1338        let database_objects: HashMap<risingwave_meta_model::DatabaseId, Vec<JobId>> = self
1339            .metadata_manager
1340            .catalog_controller
1341            .list_streaming_job_with_database()
1342            .await?;
1343
1344        let job_ids = database_objects
1345            .iter()
1346            .flat_map(|(database_id, job_ids)| {
1347                job_ids
1348                    .iter()
1349                    .enumerate()
1350                    .map(move |(idx, job_id)| (idx, database_id, job_id))
1351            })
1352            .sorted_by(|(idx_a, database_a, _), (idx_b, database_b, _)| {
1353                idx_a.cmp(idx_b).then(database_a.cmp(database_b))
1354            })
1355            .map(|(_, database_id, job_id)| (*database_id, *job_id))
1356            .filter(|(_, job_id)| !blocked_jobs.contains(job_id))
1357            .collect_vec();
1358
1359        if job_ids.is_empty() {
1360            tracing::info!("no streaming jobs for scaling, maybe an empty cluster");
1361            // Retry periodically while some jobs are temporarily blocked by creating
1362            // unreschedulable backfill jobs. This allows us to scale them
1363            // automatically once the creating jobs finish.
1364            return Ok(has_blocked_jobs);
1365        }
1366
1367        let active_workers =
1368            ActiveStreamingWorkerNodes::new_snapshot(self.metadata_manager.clone()).await?;
1369
1370        tracing::info!(
1371            "trigger parallelism control for jobs: {:#?}, workers {:#?}",
1372            job_ids,
1373            active_workers.current()
1374        );
1375
1376        let batch_size = match self.env.opts.parallelism_control_batch_size {
1377            0 => job_ids.len(),
1378            n => n,
1379        };
1380
1381        tracing::info!(
1382            "total {} streaming jobs, batch size {}, schedulable worker ids: {:?}",
1383            job_ids.len(),
1384            batch_size,
1385            active_workers.current()
1386        );
1387
1388        let batches: Vec<_> = job_ids
1389            .into_iter()
1390            .chunks(batch_size)
1391            .into_iter()
1392            .map(|chunk| chunk.collect_vec())
1393            .collect();
1394
1395        for batch in batches {
1396            let jobs = batch.iter().map(|(_, job_id)| *job_id).collect();
1397
1398            let commands = self.scale_controller.rerender(jobs).await?;
1399
1400            let futures = commands.into_iter().map(|(database_id, command)| {
1401                let barrier_scheduler = self.barrier_scheduler.clone();
1402                async move { barrier_scheduler.run_command(database_id, command).await }
1403            });
1404
1405            let _results = future::try_join_all(futures).await?;
1406        }
1407
1408        Ok(has_blocked_jobs)
1409    }
1410
1411    /// Handles notification of worker node activation and deletion, and triggers parallelism control.
1412    async fn run(&self, mut shutdown_rx: Receiver<()>) {
1413        tracing::info!("starting automatic parallelism control monitor");
1414
1415        let check_period =
1416            Duration::from_secs(self.env.opts.parallelism_control_trigger_period_sec);
1417
1418        let mut ticker = tokio::time::interval_at(
1419            Instant::now()
1420                + Duration::from_secs(self.env.opts.parallelism_control_trigger_first_delay_sec),
1421            check_period,
1422        );
1423        ticker.set_missed_tick_behavior(MissedTickBehavior::Skip);
1424
1425        let (local_notification_tx, mut local_notification_rx) =
1426            tokio::sync::mpsc::unbounded_channel();
1427
1428        self.env
1429            .notification_manager()
1430            .insert_local_sender(local_notification_tx);
1431
1432        // waiting for the first tick
1433        ticker.tick().await;
1434
1435        let worker_nodes = self
1436            .metadata_manager
1437            .list_active_streaming_compute_nodes()
1438            .await
1439            .expect("list active streaming compute nodes");
1440
1441        let mut worker_cache: BTreeMap<_, _> = worker_nodes
1442            .into_iter()
1443            .map(|worker| (worker.id, worker))
1444            .collect();
1445
1446        let mut should_trigger = false;
1447
1448        loop {
1449            tokio::select! {
1450                biased;
1451
1452                _ = &mut shutdown_rx => {
1453                    tracing::info!("Stream manager is stopped");
1454                    break;
1455                }
1456
1457                _ = ticker.tick(), if should_trigger => {
1458                    let include_workers = worker_cache.keys().copied().collect_vec();
1459
1460                    if include_workers.is_empty() {
1461                        tracing::debug!("no available worker nodes");
1462                        should_trigger = false;
1463                        continue;
1464                    }
1465
1466                    match self.trigger_parallelism_control().await {
1467                        Ok(cont) => {
1468                            should_trigger = cont;
1469                        }
1470                        Err(e) => {
1471                            tracing::warn!(error = %e.as_report(), "Failed to trigger scale out, waiting for next tick to retry after {}s", ticker.period().as_secs());
1472                            ticker.reset();
1473                        }
1474                    }
1475                }
1476
1477                notification = local_notification_rx.recv() => {
1478                    let notification = notification.expect("local notification channel closed in loop of stream manager");
1479
1480                    // Only maintain the cache for streaming compute nodes.
1481                    let worker_is_streaming_compute = |worker: &WorkerNode| {
1482                        worker.get_type() == Ok(WorkerType::ComputeNode)
1483                            && worker.property.as_ref().unwrap().is_streaming
1484                    };
1485
1486                    match notification {
1487                        // TODO(#25478): add a local notification path for relevant session
1488                        // parameter updates and trigger a control round here when the cluster
1489                        // defaults for stream/backfill parallelism change.
1490                        LocalNotification::SystemParamsChange(_) => {}
1491                        LocalNotification::WorkerNodeActivated(worker) => {
1492                            if !worker_is_streaming_compute(&worker) {
1493                                continue;
1494                            }
1495
1496                            tracing::info!(worker = %worker.id, "worker activated notification received");
1497
1498                            let prev_worker = worker_cache.insert(worker.id, worker.clone());
1499
1500                            match prev_worker {
1501                                Some(prev_worker) if prev_worker.compute_node_parallelism() != worker.compute_node_parallelism()  => {
1502                                    tracing::info!(worker = %worker.id, "worker parallelism changed");
1503                                    should_trigger = true;
1504                                }
1505                                Some(prev_worker) if prev_worker.resource_group() != worker.resource_group()  => {
1506                                    tracing::info!(worker = %worker.id, "worker label changed");
1507                                    should_trigger = true;
1508                                }
1509                                None => {
1510                                    tracing::info!(worker = %worker.id, "new worker joined");
1511                                    should_trigger = true;
1512                                }
1513                                _ => {}
1514                            }
1515                        }
1516
1517                        // Since our logic for handling passive scale-in is within the barrier manager,
1518                        // there’s not much we can do here. All we can do is proactively remove the entries from our cache.
1519                        LocalNotification::WorkerNodeDeleted(worker) => {
1520                            if !worker_is_streaming_compute(&worker) {
1521                                continue;
1522                            }
1523
1524                            match worker_cache.remove(&worker.id) {
1525                                Some(prev_worker) => {
1526                                    tracing::info!(worker = %prev_worker.id, "worker removed from stream manager cache");
1527                                }
1528                                None => {
1529                                    tracing::warn!(worker = %worker.id, "worker not found in stream manager cache, but it was removed");
1530                                }
1531                            }
1532                        }
1533
1534                        LocalNotification::StreamingJobBackfillFinished(job_id) => {
1535                            tracing::debug!(job_id = %job_id, "received backfill finished notification");
1536                            if let Err(e) = self.apply_post_backfill_parallelism(job_id).await {
1537                                tracing::warn!(job_id = %job_id, error = %e.as_report(), "failed to restore parallelism after backfill");
1538                                // Retry in the next periodic tick. This avoids triggering
1539                                // unnecessary full control passes when restore succeeds.
1540                                should_trigger = true;
1541                            }
1542                        }
1543
1544                        _ => {}
1545                    }
1546                }
1547            }
1548        }
1549    }
1550
1551    /// Restores a streaming job's parallelism to its target value after backfill completes.
1552    async fn apply_post_backfill_parallelism(&self, job_id: JobId) -> MetaResult<()> {
1553        // Fetch both the target parallelism (final desired state) and the backfill parallelism
1554        // (temporary parallelism used during backfill phase) from the catalog.
1555        let Some((
1556            target,
1557            target_adaptive_parallelism_strategy,
1558            backfill_parallelism,
1559            backfill_adaptive_parallelism_strategy,
1560        )) = self
1561            .metadata_manager
1562            .catalog_controller
1563            .get_job_parallelisms(job_id)
1564            .await?
1565        else {
1566            // The job may have been dropped before this notification is processed.
1567            // Treat it as a benign no-op so we don't trigger unnecessary retries.
1568            tracing::debug!(
1569                job_id = %job_id,
1570                "streaming job not found when applying post-backfill parallelism, skip"
1571            );
1572            return Ok(());
1573        };
1574
1575        // Determine if we need to reschedule based on the backfill configuration.
1576        match backfill_parallelism {
1577            Some(backfill_parallelism)
1578                if backfill_parallelism == target
1579                    && backfill_adaptive_parallelism_strategy
1580                        == target_adaptive_parallelism_strategy =>
1581            {
1582                // Backfill parallelism matches target - no reschedule needed since the job
1583                // is already running at the desired parallelism.
1584                tracing::debug!(
1585                    job_id = %job_id,
1586                    ?backfill_parallelism,
1587                    ?target,
1588                    "backfill parallelism equals job parallelism, skip reschedule"
1589                );
1590                return Ok(());
1591            }
1592            Some(_) => {
1593                // Backfill parallelism differs from target - proceed to restore target parallelism.
1594            }
1595            None => {
1596                // No backfill parallelism was configured, meaning the job was created without
1597                // a special backfill override. No reschedule is necessary.
1598                tracing::debug!(
1599                    job_id = %job_id,
1600                    ?target,
1601                    "no backfill parallelism configured, skip post-backfill reschedule"
1602                );
1603                return Ok(());
1604            }
1605        }
1606
1607        // Reschedule the job to restore its target parallelism.
1608        tracing::info!(
1609            job_id = %job_id,
1610            ?target,
1611            ?backfill_parallelism,
1612            "restoring parallelism after backfill via reschedule"
1613        );
1614        let policy = ReschedulePolicy::Parallelism(ParallelismPolicy {
1615            parallelism: target,
1616            adaptive_parallelism_strategy: target_adaptive_parallelism_strategy,
1617        });
1618        if let Err(e) = self.reschedule_streaming_job(job_id, policy, false).await {
1619            tracing::warn!(job_id = %job_id, error = %e.as_report(), "reschedule after backfill failed");
1620            return Err(e);
1621        }
1622
1623        tracing::info!(job_id = %job_id, "parallelism reschedule after backfill submitted");
1624        Ok(())
1625    }
1626
1627    pub fn start_auto_parallelism_monitor(
1628        self: Arc<Self>,
1629    ) -> (JoinHandle<()>, oneshot::Sender<()>) {
1630        tracing::info!("Automatic parallelism scale-out is enabled for streaming jobs");
1631        let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();
1632        let join_handle = tokio::spawn(async move {
1633            self.run(shutdown_rx).await;
1634        });
1635
1636        (join_handle, shutdown_tx)
1637    }
1638}