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