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_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#[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 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
144pub(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 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
488async 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
699fn 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 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
838pub(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, ¤t));
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), ¤t_source),
1241 (FragmentId::new(2), ¤t_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), ¤t_source),
1313 (FragmentId::new(2), ¤t_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 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 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 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 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 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 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 _ = 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 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 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 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 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 should_trigger = true;
1602 }
1603 }
1604
1605 _ => {}
1606 }
1607 }
1608 }
1609 }
1610 }
1611
1612 async fn apply_post_backfill_parallelism(&self, job_id: JobId) -> MetaResult<()> {
1614 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 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 match backfill_parallelism {
1638 Some(backfill_parallelism)
1639 if backfill_parallelism == target
1640 && backfill_adaptive_parallelism_strategy
1641 == target_adaptive_parallelism_strategy =>
1642 {
1643 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 }
1656 None => {
1657 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 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}