Skip to main content

risingwave_meta/stream/
source_manager.rs

1// Copyright 2022 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15mod 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>;
57/// Actor-level split assignment, used inside the barrier worker after actor rendering.
58pub type SplitAssignment = HashMap<FragmentId, HashMap<ActorId, Vec<SplitImpl>>>;
59
60/// Source-level split assignment: maps source ID to the discovered splits enum.
61/// Unlike [`SplitAssignment`], this does not include fragment- or actor-level assignment.
62/// The [`DiscoveredSplits`] enum is preserved all the way through to the barrier
63/// command, and only resolved to concrete per-fragment, per-actor splits when the
64/// barrier worker performs the actual split assignment.
65pub type SourceSplitAssignment = HashMap<SourceId, DiscoveredSplits>;
66
67/// Discovered splits from the source enumerator.
68///
69/// For adaptive splits, the single template split is returned and must be expanded
70/// to per-actor splits after actor rendering inside the barrier worker.
71#[derive(Debug, Clone)]
72pub enum DiscoveredSplits {
73    /// Fixed splits discovered from the external source.
74    Fixed(BTreeMap<Arc<str>, SplitImpl>),
75    /// A single adaptive split template that needs to be expanded per actor.
76    /// The `fill_adaptive_split` function should be called after actor rendering.
77    Adaptive(SplitImpl),
78}
79
80/// Describes how splits should be resolved for a replace stream job.
81///
82/// This is determined during Phase 1 (before the barrier command) and stored in
83/// [`ReplaceStreamJobPlan`](crate::barrier::ReplaceStreamJobPlan).
84/// The barrier worker uses this to decide the split resolution strategy in Phase 2.
85#[derive(Debug, Clone)]
86pub enum ReplaceJobSplitPlan {
87    /// Fresh discovered splits. Used when the replace job has no existing downstream
88    /// consumers (i.e., `replace_upstream` is empty), or the job is not a source.
89    /// Resolved via [`SourceManager::resolve_fragment_to_actor_splits`] in Phase 2.
90    Discovered(SourceSplitAssignment),
91    /// Splits need to be aligned with the previous source fragment in Phase 2.
92    /// This happens when replacing a source that has existing downstream consumers.
93    /// Resolved via [`SourceManager::resolve_replace_source_splits`] in Phase 2.
94    ///
95    /// Contains the no-shuffle mapping needed for split alignment.
96    AlignFromPrevious,
97}
98
99// ALTER CONNECTOR parameters, specifying the new parameters to be set for each job_id (source_id/sink_id)
100pub type ConnectorPropsChange = HashMap<ObjectId, HashMap<String, String>>;
101
102const DEFAULT_SOURCE_TICK_TIMEOUT: Duration = Duration::from_secs(10);
103
104/// `SourceManager` keeps fetching the latest split metadata from the external source services ([`worker::ConnectorSourceWorker::tick`]),
105/// and sends a split assignment command if split changes detected ([`Self::tick`]).
106pub 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 source loops
116    managed_sources: HashMap<SourceId, ConnectorSourceWorkerHandle>,
117    /// Fragments associated with each source
118    source_fragments: HashMap<SourceId, BTreeSet<FragmentId>>,
119    /// `source_id` -> `(fragment_id, upstream_fragment_id)`
120    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    /// Updates states after all kinds of source change.
148    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                // TODO: enable this assertion after we implemented cleanup for backfill fragments
209                // debug_assert!(
210                //     fragments.is_empty(),
211                //     "when dropping source, there should be no backfill fragments, got: {:?}",
212                //     fragments
213                // );
214            }
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            // TODO: add source_id to the fragment_replacements to avoid iterating all sources
247            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                // the update here should not involve fragments change and split change
268                // Or we need to drop and recreate the source worker instead of updating inplace
269                let props_wrapper =
270                    WithOptionsSecResolved::without_secrets(new_props.into_iter().collect());
271                let props = ConnectorProperties::extract(props_wrapper, false).unwrap(); // already checked when sending barrier
272                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    /// For replacing job (alter table/source, create sink into table).
403    #[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        // Extract the fragments that include source operators.
411        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    /// Updates states after all kinds of source change.
422    /// e.g., split change (`post_collect` barrier) or scaling (`post_apply_reschedule`).
423    #[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        // Force tick for updated source workers
441        if need_force_tick {
442            self.force_tick_updated_sources(updated_source_ids).await;
443        }
444    }
445
446    /// create and register connector worker for source.
447    #[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    /// register connector worker for source.
471    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    /// Checks whether the external source metadata has changed, and sends a split assignment command
495    /// if it has.
496    ///
497    /// This is also how a newly created `SourceExecutor` is initialized.
498    /// (force `tick` in `Self::create_source_worker`)
499    ///
500    /// The command will first updates `SourceExecutor`'s splits, and finally calls `Self::apply_source_change`
501    /// to update states in `SourceManager`.
502    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    /// Pause the tick loop in source manager until the returned guard is dropped.
537    pub async fn pause_tick(&self) -> MutexGuard<'_, ()> {
538        tracing::debug!("pausing tick lock in source manager");
539        self.paused.lock().await
540    }
541
542    /// Force tick for specific updated source workers after properties update.
543    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    /// Reset source split assignments by clearing the cached split state
565    /// and triggering re-discovery. This is an UNSAFE operation that may
566    /// cause data duplication or loss depending on the connector.
567    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            // Clear the cached splits to force re-discovery
576            {
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            // Force a tick to re-discover splits
587            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    /// Validate split offsets before injecting them.
609    /// Returns `Ok(applied_split_ids)` if all validations pass, otherwise returns an error.
610    ///
611    /// Validations performed:
612    /// 1. Source exists in source manager
613    /// 2. All requested split IDs exist in the source's current splits (runtime assignment)
614    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            // Check if source exists
623            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        // Validate all requested split IDs exist
664        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    /// `CREATE SOURCE` (shared), or `CREATE MV`.
700    /// This is applied after the job is successfully created (`post_collect` barrier).
701    CreateJob {
702        added_source_fragments: HashMap<SourceId, BTreeSet<FragmentId>>,
703        /// (`source_id`, -> (`source_backfill_fragment_id`, `upstream_source_fragment_id`))
704        added_backfill_fragments: HashMap<SourceId, BTreeSet<(FragmentId, FragmentId)>>,
705    },
706    UpdateSourceProps {
707        // the new properties to be set for each source_id
708        // and the props should not affect split assignment and fragments
709        source_id_map_new_props: HashMap<SourceId, HashMap<String, String>>,
710    },
711    /// `CREATE SOURCE` (shared), or `CREATE MV` is _finished_ (backfill is done).
712    /// This is applied after `wait_streaming_job_finished`.
713    /// XXX: Should we merge `CreateJob` into this?
714    CreateJobFinished {
715        /// (`source_id`, -> (`source_backfill_fragment_id`, `upstream_source_fragment_id`))
716        finished_backfill_fragments: HashMap<SourceId, BTreeSet<(FragmentId, FragmentId)>>,
717    },
718    /// `DROP SOURCE` or `DROP MV`
719    DropSource { dropped_source_ids: Vec<SourceId> },
720    DropMv {
721        // FIXME: we should consider source backfill fragments here for MV on shared source.
722        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}