1mod split_assignment;
16mod worker;
17use std::borrow::BorrowMut;
18use std::cmp::Ordering;
19use std::collections::hash_map::Entry;
20use std::collections::{BTreeMap, BTreeSet, BinaryHeap, HashMap, HashSet};
21use std::sync::Arc;
22use std::time::Duration;
23
24use anyhow::Context;
25use risingwave_common::catalog::DatabaseId;
26use risingwave_common::id::ObjectId;
27use risingwave_common::metrics::{
28 LabelGuardedHistogram, LabelGuardedIntCounter, LabelGuardedIntGauge,
29};
30use risingwave_common::panic_if_debug;
31use risingwave_connector::WithOptionsSecResolved;
32use risingwave_connector::error::ConnectorResult;
33use risingwave_connector::source::{
34 AnySplitEnumerator, ConnectorProperties, SourceEnumeratorContext, SourceEnumeratorInfo,
35 SplitId, SplitImpl, SplitMetaData,
36};
37use risingwave_meta_model::SourceId;
38use risingwave_pb::catalog::Source;
39use risingwave_pb::source::{ConnectorSplit, ConnectorSplits};
40pub use split_assignment::{SplitDiffOptions, SplitState, align_splits, reassign_splits};
41use thiserror_ext::AsReport;
42use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender};
43use tokio::sync::{Mutex, MutexGuard, oneshot};
44use tokio::task::JoinHandle;
45use tokio::time::MissedTickBehavior;
46use tokio::{select, time};
47pub use worker::create_source_worker;
48use worker::{ConnectorSourceWorkerHandle, create_source_worker_async};
49
50use crate::barrier::{BarrierScheduler, Command, ReplaceStreamJobPlan};
51use crate::manager::{MetaSrvEnv, MetadataManager};
52use crate::model::{ActorId, FragmentId, StreamJobFragments};
53use crate::rpc::metrics::MetaMetrics;
54use crate::{MetaError, MetaResult};
55
56pub type SourceManagerRef = Arc<SourceManager>;
57pub type SplitAssignment = HashMap<FragmentId, HashMap<ActorId, Vec<SplitImpl>>>;
59
60pub type SourceSplitAssignment = HashMap<SourceId, DiscoveredSplits>;
66
67#[derive(Debug, Clone)]
72pub enum DiscoveredSplits {
73 Fixed(BTreeMap<Arc<str>, SplitImpl>),
75 Adaptive(SplitImpl),
78}
79
80#[derive(Debug, Clone)]
86pub enum ReplaceJobSplitPlan {
87 Discovered(SourceSplitAssignment),
91 AlignFromPrevious,
97}
98
99pub type ConnectorPropsChange = HashMap<ObjectId, HashMap<String, String>>;
101
102const DEFAULT_SOURCE_TICK_TIMEOUT: Duration = Duration::from_secs(10);
103
104pub struct SourceManager {
107 pub paused: Mutex<()>,
108 barrier_scheduler: BarrierScheduler,
109 core: Mutex<SourceManagerCore>,
110 pub metrics: Arc<MetaMetrics>,
111}
112pub struct SourceManagerCore {
113 metadata_manager: MetadataManager,
114
115 managed_sources: HashMap<SourceId, ConnectorSourceWorkerHandle>,
117 source_fragments: HashMap<SourceId, BTreeSet<FragmentId>>,
119 backfill_fragments: HashMap<SourceId, BTreeSet<(FragmentId, FragmentId)>>,
121
122 env: MetaSrvEnv,
123}
124
125pub struct SourceManagerRunningInfo {
126 pub source_fragments: HashMap<SourceId, BTreeSet<FragmentId>>,
127 pub backfill_fragments: HashMap<SourceId, BTreeSet<(FragmentId, FragmentId)>>,
128}
129
130impl SourceManagerCore {
131 fn new(
132 metadata_manager: MetadataManager,
133 managed_sources: HashMap<SourceId, ConnectorSourceWorkerHandle>,
134 source_fragments: HashMap<SourceId, BTreeSet<FragmentId>>,
135 backfill_fragments: HashMap<SourceId, BTreeSet<(FragmentId, FragmentId)>>,
136 env: MetaSrvEnv,
137 ) -> Self {
138 Self {
139 metadata_manager,
140 managed_sources,
141 source_fragments,
142 backfill_fragments,
143 env,
144 }
145 }
146
147 pub fn apply_source_change(&mut self, source_change: SourceChange) {
149 let mut added_source_fragments = Default::default();
150 let mut added_backfill_fragments = Default::default();
151 let mut finished_backfill_fragments = Default::default();
152 let mut fragment_replacements = Default::default();
153 let mut dropped_source_fragments = Default::default();
154 let mut dropped_source_ids = Default::default();
155 let mut recreate_source_id_map_new_props: Vec<(SourceId, HashMap<String, String>)> =
156 Default::default();
157
158 match source_change {
159 SourceChange::CreateJob {
160 added_source_fragments: added_source_fragments_,
161 added_backfill_fragments: added_backfill_fragments_,
162 } => {
163 added_source_fragments = added_source_fragments_;
164 added_backfill_fragments = added_backfill_fragments_;
165 }
166 SourceChange::CreateJobFinished {
167 finished_backfill_fragments: finished_backfill_fragments_,
168 } => {
169 finished_backfill_fragments = finished_backfill_fragments_;
170 }
171
172 SourceChange::DropMv {
173 dropped_source_fragments: dropped_source_fragments_,
174 } => {
175 dropped_source_fragments = dropped_source_fragments_;
176 }
177 SourceChange::ReplaceJob {
178 dropped_source_fragments: dropped_source_fragments_,
179 added_source_fragments: added_source_fragments_,
180 fragment_replacements: fragment_replacements_,
181 } => {
182 dropped_source_fragments = dropped_source_fragments_;
183 added_source_fragments = added_source_fragments_;
184 fragment_replacements = fragment_replacements_;
185 }
186 SourceChange::DropSource {
187 dropped_source_ids: dropped_source_ids_,
188 } => {
189 dropped_source_ids = dropped_source_ids_;
190 }
191
192 SourceChange::UpdateSourceProps {
193 source_id_map_new_props,
194 } => {
195 for (source_id, new_props) in source_id_map_new_props {
196 recreate_source_id_map_new_props.push((source_id, new_props));
197 }
198 }
199 }
200
201 for source_id in dropped_source_ids {
202 let dropped_fragments = self.source_fragments.remove(&source_id);
203
204 if let Some(handle) = self.managed_sources.remove(&source_id) {
205 handle.terminate(dropped_fragments);
206 }
207 if let Some(_fragments) = self.backfill_fragments.remove(&source_id) {
208 }
215 }
216
217 for (source_id, fragments) in added_source_fragments {
218 self.source_fragments
219 .entry(source_id)
220 .or_default()
221 .extend(fragments);
222 }
223
224 for (source_id, fragments) in added_backfill_fragments {
225 self.backfill_fragments
226 .entry(source_id)
227 .or_default()
228 .extend(fragments);
229 }
230
231 for (source_id, fragments) in finished_backfill_fragments {
232 let handle = self.managed_sources.get(&source_id).unwrap_or_else(|| {
233 panic!(
234 "source {} not found when adding backfill fragments {:?}",
235 source_id, fragments
236 );
237 });
238 handle.finish_backfill(fragments.iter().map(|(id, _up_id)| *id).collect());
239 }
240
241 for (source_id, fragment_ids) in dropped_source_fragments {
242 self.drop_source_fragments(Some(source_id), fragment_ids);
243 }
244
245 for (old_fragment_id, new_fragment_id) in fragment_replacements {
246 self.drop_source_fragments(None, BTreeSet::from([old_fragment_id]));
248
249 for fragment_ids in self.backfill_fragments.values_mut() {
250 let mut new_backfill_fragment_ids = fragment_ids.clone();
251 for (fragment_id, upstream_fragment_id) in fragment_ids.iter() {
252 assert_ne!(
253 fragment_id, upstream_fragment_id,
254 "backfill fragment should not be replaced"
255 );
256 if *upstream_fragment_id == old_fragment_id {
257 new_backfill_fragment_ids.remove(&(*fragment_id, *upstream_fragment_id));
258 new_backfill_fragment_ids.insert((*fragment_id, new_fragment_id));
259 }
260 }
261 *fragment_ids = new_backfill_fragment_ids;
262 }
263 }
264
265 for (source_id, new_props) in recreate_source_id_map_new_props {
266 if let Some(handle) = self.managed_sources.get_mut(&source_id) {
267 let props_wrapper =
270 WithOptionsSecResolved::without_secrets(new_props.into_iter().collect());
271 let props = ConnectorProperties::extract(props_wrapper, false).unwrap(); handle.update_props(props);
273 tracing::info!("update source {source_id} properties in source manager");
274 } else {
275 tracing::info!("job id {source_id} is not registered in source manager");
276 }
277 }
278 }
279
280 fn drop_source_fragments(
281 &mut self,
282 source_id: Option<SourceId>,
283 dropped_fragment_ids: BTreeSet<FragmentId>,
284 ) {
285 if let Some(source_id) = source_id {
286 if let Entry::Occupied(mut entry) = self.source_fragments.entry(source_id) {
287 let mut dropped_ids = vec![];
288 let managed_fragment_ids = entry.get_mut();
289 for fragment_id in &dropped_fragment_ids {
290 managed_fragment_ids.remove(fragment_id);
291 dropped_ids.push(*fragment_id);
292 }
293 if let Some(handle) = self.managed_sources.get(&source_id) {
294 handle.drop_fragments(dropped_ids);
295 } else {
296 panic_if_debug!(
297 "source {source_id} not found when dropping fragment {dropped_ids:?}",
298 );
299 }
300 if managed_fragment_ids.is_empty() {
301 entry.remove();
302 }
303 }
304 } else {
305 for (source_id, fragment_ids) in &mut self.source_fragments {
306 let mut dropped_ids = vec![];
307 for fragment_id in &dropped_fragment_ids {
308 if fragment_ids.remove(fragment_id) {
309 dropped_ids.push(*fragment_id);
310 }
311 }
312 if !dropped_ids.is_empty() {
313 if let Some(handle) = self.managed_sources.get(source_id) {
314 handle.drop_fragments(dropped_ids);
315 } else {
316 panic_if_debug!(
317 "source {source_id} not found when dropping fragment {dropped_ids:?}",
318 );
319 }
320 }
321 }
322 }
323 }
324}
325
326impl SourceManager {
327 const DEFAULT_SOURCE_TICK_INTERVAL: Duration = Duration::from_secs(10);
328
329 pub async fn new(
330 barrier_scheduler: BarrierScheduler,
331 metadata_manager: MetadataManager,
332 metrics: Arc<MetaMetrics>,
333 env: MetaSrvEnv,
334 ) -> MetaResult<Self> {
335 let mut managed_sources = HashMap::new();
336 {
337 let sources = metadata_manager.list_sources().await?;
338 for source in sources {
339 create_source_worker_async(
340 source,
341 &mut managed_sources,
342 metrics.clone(),
343 env.await_tree_reg().clone(),
344 )?
345 }
346 }
347
348 let source_fragments = metadata_manager
349 .catalog_controller
350 .load_source_fragment_ids()
351 .await?
352 .into_iter()
353 .map(|(source_id, fragment_ids)| {
354 (
355 source_id as SourceId,
356 fragment_ids.into_iter().map(|id| id as _).collect(),
357 )
358 })
359 .collect();
360 let backfill_fragments = metadata_manager
361 .catalog_controller
362 .load_backfill_fragment_ids()
363 .await?;
364
365 let core = Mutex::new(SourceManagerCore::new(
366 metadata_manager,
367 managed_sources,
368 source_fragments,
369 backfill_fragments,
370 env,
371 ));
372
373 Ok(Self {
374 barrier_scheduler,
375 core,
376 paused: Mutex::new(()),
377 metrics,
378 })
379 }
380
381 pub async fn validate_source_once(
382 &self,
383 source_id: SourceId,
384 new_source_props: WithOptionsSecResolved,
385 ) -> MetaResult<()> {
386 let props = ConnectorProperties::extract(new_source_props, false).unwrap();
387
388 {
389 let mut enumerator = props
390 .create_split_enumerator(Arc::new(SourceEnumeratorContext {
391 metrics: self.metrics.source_enumerator_metrics.clone(),
392 info: SourceEnumeratorInfo { source_id },
393 }))
394 .await
395 .context("failed to create SplitEnumerator")?;
396
397 validate_enumerator_once(&mut *enumerator).await?;
398 }
399 Ok(())
400 }
401
402 #[await_tree::instrument]
404 pub async fn handle_replace_job(
405 &self,
406 dropped_job_fragments: &StreamJobFragments,
407 added_source_fragments: HashMap<SourceId, BTreeSet<FragmentId>>,
408 replace_plan: &ReplaceStreamJobPlan,
409 ) {
410 let dropped_source_fragments = dropped_job_fragments.stream_source_fragments();
412
413 self.apply_source_change(SourceChange::ReplaceJob {
414 dropped_source_fragments,
415 added_source_fragments,
416 fragment_replacements: replace_plan.fragment_replacements(),
417 })
418 .await;
419 }
420
421 #[await_tree::instrument("apply_source_change({source_change})")]
424 pub async fn apply_source_change(&self, source_change: SourceChange) {
425 let need_force_tick = matches!(source_change, SourceChange::UpdateSourceProps { .. });
426 let updated_source_ids = if let SourceChange::UpdateSourceProps {
427 ref source_id_map_new_props,
428 } = source_change
429 {
430 source_id_map_new_props.keys().cloned().collect::<Vec<_>>()
431 } else {
432 Vec::new()
433 };
434
435 {
436 let mut core = self.core.lock().await;
437 core.apply_source_change(source_change);
438 }
439
440 if need_force_tick {
442 self.force_tick_updated_sources(updated_source_ids).await;
443 }
444 }
445
446 #[await_tree::instrument("register_source({})", source.name)]
448 pub async fn register_source(&self, source: &Source) -> MetaResult<()> {
449 tracing::debug!("register_source: {}", source.get_id());
450 let mut core = self.core.lock().await;
451 let source_id = source.get_id();
452 if core.managed_sources.contains_key(&source_id) {
453 tracing::warn!("source {} already registered", source_id);
454 return Ok(());
455 }
456
457 let handle = create_source_worker(
458 source,
459 self.metrics.clone(),
460 core.env.await_tree_reg().clone(),
461 )
462 .await
463 .context("failed to create source worker")?;
464
465 core.managed_sources.insert(source_id, handle);
466
467 Ok(())
468 }
469
470 pub async fn register_source_with_handle(
472 &self,
473 source_id: SourceId,
474 handle: ConnectorSourceWorkerHandle,
475 ) {
476 let mut core = self.core.lock().await;
477 if core.managed_sources.contains_key(&source_id) {
478 tracing::warn!("source {} already registered", source_id);
479 return;
480 }
481
482 core.managed_sources.insert(source_id, handle);
483 }
484
485 pub async fn get_running_info(&self) -> SourceManagerRunningInfo {
486 let core = self.core.lock().await;
487
488 SourceManagerRunningInfo {
489 source_fragments: core.source_fragments.clone(),
490 backfill_fragments: core.backfill_fragments.clone(),
491 }
492 }
493
494 async fn tick(&self) -> MetaResult<()> {
503 let split_states = {
504 let core_guard = self.core.lock().await;
505 core_guard.reassign_splits().await?
506 };
507
508 for (database_id, split_state) in split_states {
509 if !split_state.split_assignment.is_empty() {
510 let command = Command::SourceChangeSplit(split_state);
511 tracing::info!(command = ?command, "pushing down split assignment command");
512 self.barrier_scheduler
513 .run_command(database_id, command)
514 .await?;
515 }
516 }
517
518 Ok(())
519 }
520
521 pub async fn run(&self) -> MetaResult<()> {
522 let mut ticker = time::interval(Self::DEFAULT_SOURCE_TICK_INTERVAL);
523 ticker.set_missed_tick_behavior(MissedTickBehavior::Skip);
524 loop {
525 ticker.tick().await;
526 let _pause_guard = self.paused.lock().await;
527 if let Err(e) = self.tick().await {
528 tracing::error!(
529 error = %e.as_report(),
530 "error happened while running source manager tick",
531 );
532 }
533 }
534 }
535
536 pub async fn pause_tick(&self) -> MutexGuard<'_, ()> {
538 tracing::debug!("pausing tick lock in source manager");
539 self.paused.lock().await
540 }
541
542 async fn force_tick_updated_sources(&self, updated_source_ids: Vec<SourceId>) {
544 let core = self.core.lock().await;
545 for source_id in updated_source_ids {
546 if let Some(handle) = core.managed_sources.get(&source_id) {
547 tracing::info!("forcing tick for updated source {}", source_id);
548 if let Err(e) = handle.force_tick().await {
549 tracing::warn!(
550 error = %e.as_report(),
551 "failed to force tick for source {} after properties update",
552 source_id
553 );
554 }
555 } else {
556 tracing::warn!(
557 "source {} not found when trying to force tick after update",
558 source_id
559 );
560 }
561 }
562 }
563
564 pub async fn reset_source_splits(&self, source_id: SourceId) -> MetaResult<()> {
568 tracing::warn!(
569 %source_id,
570 "UNSAFE: Resetting source splits - clearing cached state and triggering re-discovery"
571 );
572
573 let core = self.core.lock().await;
574 if let Some(handle) = core.managed_sources.get(&source_id) {
575 {
577 let mut splits_guard = handle.splits.lock().await;
578 tracing::info!(
579 %source_id,
580 prev_splits = ?splits_guard.splits.as_ref().map(|s| s.len()),
581 "Clearing cached splits"
582 );
583 splits_guard.splits = None;
584 }
585
586 tracing::info!(
588 %source_id,
589 "Triggering split re-discovery via force_tick"
590 );
591 handle.force_tick().await.with_context(|| {
592 format!(
593 "failed to force tick for source {} after split reset",
594 source_id
595 )
596 })?;
597
598 tracing::info!(
599 %source_id,
600 "Split reset completed - new splits will be assigned on next tick"
601 );
602 Ok(())
603 } else {
604 Err(anyhow::anyhow!("source {} not found in source manager", source_id).into())
605 }
606 }
607
608 pub async fn validate_inject_source_offsets(
615 &self,
616 source_id: SourceId,
617 split_offsets: &HashMap<String, String>,
618 ) -> MetaResult<Vec<String>> {
619 let (fragment_ids, env) = {
620 let core = self.core.lock().await;
621
622 let _ = core.managed_sources.get(&source_id).ok_or_else(|| {
624 MetaError::invalid_parameter(format!(
625 "source {} not found in source manager",
626 source_id
627 ))
628 })?;
629
630 let mut ids = Vec::new();
631 if let Some(src_frags) = core.source_fragments.get(&source_id) {
632 ids.extend(src_frags.iter().copied());
633 }
634 if let Some(backfill_frags) = core.backfill_fragments.get(&source_id) {
635 ids.extend(
636 backfill_frags
637 .iter()
638 .flat_map(|(id, upstream)| [*id, *upstream]),
639 );
640 }
641 (ids, core.env.clone())
642 };
643
644 if fragment_ids.is_empty() {
645 return Err(MetaError::invalid_parameter(format!(
646 "source {} has no running fragments",
647 source_id
648 )));
649 }
650
651 let guard = env.shared_actor_infos().read_guard();
652 let mut assigned_split_ids = HashSet::new();
653 for fragment_id in fragment_ids {
654 if let Some(fragment) = guard.get_fragment(fragment_id) {
655 for actor in fragment.actors.values() {
656 for split in &actor.splits {
657 assigned_split_ids.insert(split.id().to_string());
658 }
659 }
660 }
661 }
662
663 let mut invalid_splits = Vec::new();
665 for split_id in split_offsets.keys() {
666 if !assigned_split_ids.contains(split_id) {
667 invalid_splits.push(split_id.clone());
668 }
669 }
670
671 if !invalid_splits.is_empty() {
672 return Err(MetaError::invalid_parameter(format!(
673 "invalid split IDs for source {}: {:?}. Valid splits are: {:?}",
674 source_id,
675 invalid_splits,
676 assigned_split_ids.iter().collect::<Vec<_>>()
677 )));
678 }
679
680 tracing::info!(
681 source_id = %source_id,
682 num_splits = split_offsets.len(),
683 "Validated inject source offsets request"
684 );
685
686 Ok(split_offsets.keys().cloned().collect())
687 }
688}
689
690async fn validate_enumerator_once(enumerator: &mut dyn AnySplitEnumerator) -> MetaResult<()> {
691 let _ = tokio::time::timeout(DEFAULT_SOURCE_TICK_TIMEOUT, enumerator.list_splits())
692 .await
693 .context("failed to list splits")??;
694 Ok(())
695}
696
697#[derive(strum::Display, Debug)]
698pub enum SourceChange {
699 CreateJob {
702 added_source_fragments: HashMap<SourceId, BTreeSet<FragmentId>>,
703 added_backfill_fragments: HashMap<SourceId, BTreeSet<(FragmentId, FragmentId)>>,
705 },
706 UpdateSourceProps {
707 source_id_map_new_props: HashMap<SourceId, HashMap<String, String>>,
710 },
711 CreateJobFinished {
715 finished_backfill_fragments: HashMap<SourceId, BTreeSet<(FragmentId, FragmentId)>>,
717 },
718 DropSource { dropped_source_ids: Vec<SourceId> },
720 DropMv {
721 dropped_source_fragments: HashMap<SourceId, BTreeSet<FragmentId>>,
723 },
724 ReplaceJob {
725 dropped_source_fragments: HashMap<SourceId, BTreeSet<FragmentId>>,
726 added_source_fragments: HashMap<SourceId, BTreeSet<FragmentId>>,
727 fragment_replacements: HashMap<FragmentId, FragmentId>,
728 },
729}
730
731pub fn build_actor_connector_splits(
732 splits: &HashMap<ActorId, Vec<SplitImpl>>,
733) -> HashMap<ActorId, ConnectorSplits> {
734 splits
735 .iter()
736 .map(|(&actor_id, splits)| {
737 (
738 actor_id,
739 ConnectorSplits {
740 splits: splits.iter().map(ConnectorSplit::from).collect(),
741 },
742 )
743 })
744 .collect()
745}
746
747pub fn build_actor_split_impls(
748 actor_splits: &HashMap<ActorId, ConnectorSplits>,
749) -> HashMap<ActorId, Vec<SplitImpl>> {
750 actor_splits
751 .iter()
752 .map(|(actor_id, ConnectorSplits { splits })| {
753 (
754 *actor_id,
755 splits
756 .iter()
757 .map(|split| SplitImpl::try_from(split).unwrap())
758 .collect(),
759 )
760 })
761 .collect()
762}