1use 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#[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 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
122pub(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 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
466async 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
677fn 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 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
816pub(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), ¤t_source),
1166 (FragmentId::new(2), ¤t_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), ¤t_source),
1238 (FragmentId::new(2), ¤t_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 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 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 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 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 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 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 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 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 should_trigger = true;
1541 }
1542 }
1543
1544 _ => {}
1545 }
1546 }
1547 }
1548 }
1549 }
1550
1551 async fn apply_post_backfill_parallelism(&self, job_id: JobId) -> MetaResult<()> {
1553 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 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 match backfill_parallelism {
1577 Some(backfill_parallelism)
1578 if backfill_parallelism == target
1579 && backfill_adaptive_parallelism_strategy
1580 == target_adaptive_parallelism_strategy =>
1581 {
1582 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 }
1595 None => {
1596 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 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}