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 let prev_splits = handle.splits.take();
577 tracing::info!(
578 %source_id,
579 prev_splits = ?prev_splits.as_ref().map(|s| s.len()),
580 "Clearing cached splits"
581 );
582
583 tracing::info!(
585 %source_id,
586 "Triggering split re-discovery via force_tick"
587 );
588 handle.force_tick().await.with_context(|| {
589 format!(
590 "failed to force tick for source {} after split reset",
591 source_id
592 )
593 })?;
594
595 tracing::info!(
596 %source_id,
597 "Split reset completed - new splits will be assigned on next tick"
598 );
599 Ok(())
600 } else {
601 Err(anyhow::anyhow!("source {} not found in source manager", source_id).into())
602 }
603 }
604
605 pub async fn validate_inject_source_offsets(
612 &self,
613 source_id: SourceId,
614 split_offsets: &HashMap<String, String>,
615 ) -> MetaResult<Vec<String>> {
616 let (fragment_ids, env) = {
617 let core = self.core.lock().await;
618
619 let _ = core.managed_sources.get(&source_id).ok_or_else(|| {
621 MetaError::invalid_parameter(format!(
622 "source {} not found in source manager",
623 source_id
624 ))
625 })?;
626
627 let mut ids = Vec::new();
628 if let Some(src_frags) = core.source_fragments.get(&source_id) {
629 ids.extend(src_frags.iter().copied());
630 }
631 if let Some(backfill_frags) = core.backfill_fragments.get(&source_id) {
632 ids.extend(
633 backfill_frags
634 .iter()
635 .flat_map(|(id, upstream)| [*id, *upstream]),
636 );
637 }
638 (ids, core.env.clone())
639 };
640
641 if fragment_ids.is_empty() {
642 return Err(MetaError::invalid_parameter(format!(
643 "source {} has no running fragments",
644 source_id
645 )));
646 }
647
648 let guard = env.shared_actor_infos().read_guard();
649 let mut assigned_split_ids = HashSet::new();
650 for fragment_id in fragment_ids {
651 if let Some(fragment) = guard.get_fragment(fragment_id) {
652 for actor in fragment.actors.values() {
653 for split in &actor.splits {
654 assigned_split_ids.insert(split.id().to_string());
655 }
656 }
657 }
658 }
659
660 let mut invalid_splits = Vec::new();
662 for split_id in split_offsets.keys() {
663 if !assigned_split_ids.contains(split_id) {
664 invalid_splits.push(split_id.clone());
665 }
666 }
667
668 if !invalid_splits.is_empty() {
669 return Err(MetaError::invalid_parameter(format!(
670 "invalid split IDs for source {}: {:?}. Valid splits are: {:?}",
671 source_id,
672 invalid_splits,
673 assigned_split_ids.iter().collect::<Vec<_>>()
674 )));
675 }
676
677 tracing::info!(
678 source_id = %source_id,
679 num_splits = split_offsets.len(),
680 "Validated inject source offsets request"
681 );
682
683 Ok(split_offsets.keys().cloned().collect())
684 }
685}
686
687async fn validate_enumerator_once(enumerator: &mut dyn AnySplitEnumerator) -> MetaResult<()> {
688 let _ = tokio::time::timeout(DEFAULT_SOURCE_TICK_TIMEOUT, enumerator.list_splits())
689 .await
690 .context("failed to list splits")??;
691 Ok(())
692}
693
694#[derive(strum::Display, Debug)]
695pub enum SourceChange {
696 CreateJob {
699 added_source_fragments: HashMap<SourceId, BTreeSet<FragmentId>>,
700 added_backfill_fragments: HashMap<SourceId, BTreeSet<(FragmentId, FragmentId)>>,
702 },
703 UpdateSourceProps {
704 source_id_map_new_props: HashMap<SourceId, HashMap<String, String>>,
707 },
708 CreateJobFinished {
712 finished_backfill_fragments: HashMap<SourceId, BTreeSet<(FragmentId, FragmentId)>>,
714 },
715 DropSource { dropped_source_ids: Vec<SourceId> },
717 DropMv {
718 dropped_source_fragments: HashMap<SourceId, BTreeSet<FragmentId>>,
720 },
721 ReplaceJob {
722 dropped_source_fragments: HashMap<SourceId, BTreeSet<FragmentId>>,
723 added_source_fragments: HashMap<SourceId, BTreeSet<FragmentId>>,
724 fragment_replacements: HashMap<FragmentId, FragmentId>,
725 },
726}
727
728pub fn build_actor_connector_splits(
729 splits: &HashMap<ActorId, Vec<SplitImpl>>,
730) -> HashMap<ActorId, ConnectorSplits> {
731 splits
732 .iter()
733 .map(|(&actor_id, splits)| {
734 (
735 actor_id,
736 ConnectorSplits {
737 splits: splits.iter().map(ConnectorSplit::from).collect(),
738 },
739 )
740 })
741 .collect()
742}
743
744pub fn build_actor_split_impls(
745 actor_splits: &HashMap<ActorId, ConnectorSplits>,
746) -> HashMap<ActorId, Vec<SplitImpl>> {
747 actor_splits
748 .iter()
749 .map(|(actor_id, ConnectorSplits { splits })| {
750 (
751 *actor_id,
752 splits
753 .iter()
754 .map(|split| SplitImpl::try_from(split).unwrap())
755 .collect(),
756 )
757 })
758 .collect()
759}