Skip to main content

risingwave_meta/stream/
refresh_manager.rs

1// Copyright 2025 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
15use std::collections::{HashMap, HashSet};
16use std::sync::Arc;
17use std::time::{Duration, Instant};
18
19use anyhow::anyhow;
20use chrono::{DateTime, Duration as ChronoDuration, Utc};
21use parking_lot::Mutex;
22use risingwave_common::catalog::{DatabaseId, FragmentTypeFlag, TableId};
23use risingwave_common::metrics::LabelGuardedIntCounter;
24use risingwave_common::util::epoch::Epoch;
25use risingwave_meta_model::ActorId;
26use risingwave_meta_model::refresh_job::{self, RefreshState};
27use risingwave_pb::catalog::table::OptionalAssociatedSourceId;
28use risingwave_pb::id::SourceId;
29use risingwave_pb::meta::{RefreshRequest, RefreshResponse};
30use thiserror_ext::AsReport;
31use tokio::sync::{Notify, oneshot};
32use tokio::task::JoinHandle;
33
34use crate::barrier::{BarrierScheduler, Command, SharedActorInfos};
35use crate::manager::{MetaSrvEnv, MetadataManager};
36use crate::rpc::metrics::GLOBAL_META_METRICS;
37use crate::{MetaError, MetaResult};
38
39pub type GlobalRefreshManagerRef = Arc<GlobalRefreshManager>;
40
41pub struct GlobalRefreshManager {
42    metadata_manager: MetadataManager,
43    barrier_scheduler: BarrierScheduler,
44    shared_actor_infos: SharedActorInfos,
45    progress_trackers: Mutex<GlobalRefreshTableProgressTracker>,
46    refresh_job_metrics: Mutex<HashMap<TableId, RefreshJobMetrics>>,
47    scheduler_notify: Notify,
48    scheduler_interval: Duration,
49}
50
51impl GlobalRefreshManager {
52    pub async fn start(
53        metadata_manager: MetadataManager,
54        barrier_scheduler: BarrierScheduler,
55        env: &MetaSrvEnv,
56        scheduler_interval: Duration,
57    ) -> MetaResult<(GlobalRefreshManagerRef, JoinHandle<()>, oneshot::Sender<()>)> {
58        let shared_actor_infos = env.shared_actor_infos().clone();
59        let manager = Arc::new(Self {
60            metadata_manager: metadata_manager.clone(),
61            barrier_scheduler,
62            shared_actor_infos,
63            progress_trackers: Mutex::new(GlobalRefreshTableProgressTracker::default()),
64            refresh_job_metrics: Mutex::new(HashMap::new()),
65            scheduler_notify: Notify::new(),
66            scheduler_interval,
67        });
68
69        manager
70            .metadata_manager
71            .reset_all_refresh_jobs_to_idle()
72            .await?;
73        manager.sync_refreshable_jobs().await?;
74
75        let (shutdown_tx, shutdown_rx) = oneshot::channel();
76        let join_handle = Self::spawn_scheduler(manager.clone(), shutdown_rx);
77
78        Ok((manager, join_handle, shutdown_tx))
79    }
80
81    fn spawn_scheduler(
82        manager: GlobalRefreshManagerRef,
83        mut shutdown_rx: oneshot::Receiver<()>,
84    ) -> JoinHandle<()> {
85        let scheduler_interval = manager.scheduler_interval;
86        tokio::spawn(async move {
87            let mut interval = tokio::time::interval(scheduler_interval);
88            loop {
89                tokio::select! {
90                    _ = interval.tick() => {
91                        if let Err(err) = manager.handle_scheduler_tick().await {
92                            tracing::warn!(error = %err.as_report(), "refresh scheduler tick failed");
93                        }
94                    }
95                    _ = manager.scheduler_notify.notified() => {
96                        if let Err(err) = manager.handle_scheduler_tick().await {
97                            tracing::warn!(error = %err.as_report(), "refresh scheduler tick failed");
98                        }
99                    }
100                    _ = &mut shutdown_rx => {
101                        tracing::info!("refresh scheduler shutting down");
102                        break;
103                    }
104                }
105            }
106        })
107    }
108
109    pub async fn trigger_manual_refresh(
110        self: &Arc<Self>,
111        request: RefreshRequest,
112        shared_actor_infos: &SharedActorInfos,
113    ) -> MetaResult<RefreshResponse> {
114        let table_id = request.table_id;
115        let associated_source_id = request.associated_source_id;
116        tracing::info!(%table_id, %associated_source_id, "trigger manual refresh");
117
118        self.ensure_refreshable(table_id, associated_source_id)
119            .await?;
120
121        let result = self
122            .execute_refresh(table_id, associated_source_id, shared_actor_infos)
123            .await;
124
125        match result {
126            Ok(_) => Ok(RefreshResponse { status: None }),
127            Err(err) => Err(err),
128        }
129    }
130
131    pub async fn mark_refresh_complete(&self, table_id: TableId) -> MetaResult<()> {
132        self.metadata_manager
133            .update_refresh_job_status(table_id, RefreshState::Idle, None, true)
134            .await?;
135        self.remove_progress_tracker(table_id, "success");
136        tracing::info!(%table_id, "Table refresh completed, state updated to Idle");
137        Ok(())
138    }
139
140    pub fn mark_list_stage_finished(
141        &self,
142        table_id: TableId,
143        actors: &HashSet<ActorId>,
144    ) -> MetaResult<bool> {
145        let mut guard = self.progress_trackers.lock();
146        let tracker = guard.inner.get_mut(&table_id).ok_or_else(|| {
147            MetaError::from(anyhow!("Table tracker not found for table {}", table_id))
148        })?;
149        tracker.report_list_finished(actors.iter().copied());
150        tracker.is_list_finished()
151    }
152
153    pub fn mark_load_stage_finished(
154        &self,
155        table_id: TableId,
156        actors: &HashSet<ActorId>,
157    ) -> MetaResult<bool> {
158        let mut guard = self.progress_trackers.lock();
159        let tracker = guard.inner.get_mut(&table_id).ok_or_else(|| {
160            MetaError::from(anyhow!("Table tracker not found for table {}", table_id))
161        })?;
162        tracker.report_load_finished(actors.iter().copied());
163        tracker.is_load_finished()
164    }
165
166    pub fn remove_trackers_by_database(&self, database_id: DatabaseId) {
167        let mut guard = self.progress_trackers.lock();
168        guard.remove_tracker_by_database_id(database_id);
169    }
170
171    pub fn notify_scheduler(&self) {
172        self.scheduler_notify.notify_one();
173    }
174
175    async fn handle_scheduler_tick(self: &Arc<Self>) -> MetaResult<()> {
176        let jobs = self.metadata_manager.list_refresh_jobs().await?;
177        let active_table_ids = jobs.iter().map(|job| job.table_id).collect::<HashSet<_>>();
178        self.refresh_job_metrics
179            .lock()
180            .retain(|table_id, _| active_table_ids.contains(table_id));
181        for job in jobs {
182            if let Err(err) = self.try_trigger_scheduled_refresh(&job).await {
183                tracing::warn!(
184                    table_id = %job.table_id,
185                    error = %err.as_report(),
186                    "failed to trigger scheduled refresh"
187                );
188            }
189        }
190        Ok(())
191    }
192
193    async fn sync_refreshable_jobs(&self) -> MetaResult<()> {
194        let table_ids = self.metadata_manager.list_refreshable_table_ids().await?;
195        for table_id in table_ids {
196            self.metadata_manager.ensure_refresh_job(table_id).await?;
197        }
198        Ok(())
199    }
200
201    async fn try_trigger_scheduled_refresh(
202        self: &Arc<Self>,
203        job: &refresh_job::Model,
204    ) -> MetaResult<()> {
205        if job.current_status != RefreshState::Idle {
206            self.refresh_job_metrics
207                .lock()
208                .entry(job.table_id)
209                .or_insert_with(|| RefreshJobMetrics::new(job.table_id))
210                .cron_miss_count
211                .inc();
212            tracing::warn!(table_id = %job.table_id, "skip scheduled refresh: current status is not idle: {:?}", job.current_status);
213            return Ok(());
214        }
215        let Some(interval_secs) = job.trigger_interval_secs else {
216            return Ok(());
217        };
218        if interval_secs <= 0 {
219            return Ok(());
220        }
221
222        let interval = ChronoDuration::seconds(interval_secs);
223        let last_run = if let Some(last_run) = job.last_trigger_time {
224            last_run
225        } else {
226            self.metadata_manager
227                .get_table_catalog_by_ids(&[job.table_id])
228                .await?
229                .first()
230                .map(|t| {
231                    Epoch(t.created_at_epoch())
232                        .as_timestamptz()
233                        .to_datetime_utc()
234                        .timestamp_millis()
235                })
236                .unwrap()
237        };
238        let now = Utc::now().naive_utc();
239        if now.signed_duration_since(
240            DateTime::from_timestamp_millis(last_run)
241                .unwrap()
242                .naive_utc(),
243        ) < interval
244        {
245            return Ok(());
246        }
247
248        let table = self
249            .metadata_manager
250            .catalog_controller
251            .get_table_by_id(job.table_id)
252            .await?;
253        if !table.refreshable {
254            return Ok(());
255        }
256
257        let Some(OptionalAssociatedSourceId::AssociatedSourceId(src_id)) =
258            table.optional_associated_source_id
259        else {
260            tracing::warn!(
261                table_id = %job.table_id,
262                "skip scheduled refresh: missing associated source id"
263            );
264            return Ok(());
265        };
266        let associated_source_id = src_id;
267
268        // Increment cron job trigger counter
269        self.refresh_job_metrics
270            .lock()
271            .entry(job.table_id)
272            .or_insert_with(|| RefreshJobMetrics::new(job.table_id))
273            .cron_trigger_count
274            .inc();
275        tracing::info!(table_id = %job.table_id, "trigger scheduled refresh at interval {:?}", interval);
276
277        self.ensure_refreshable(job.table_id, associated_source_id)
278            .await?;
279        self.execute_refresh(job.table_id, associated_source_id, &self.shared_actor_infos)
280            .await?;
281        Ok(())
282    }
283
284    async fn execute_refresh(
285        self: &Arc<Self>,
286        table_id: TableId,
287        associated_source_id: SourceId,
288        shared_actor_infos: &SharedActorInfos,
289    ) -> MetaResult<()> {
290        let trigger_time = Utc::now().naive_utc();
291        let database_id = self
292            .metadata_manager
293            .catalog_controller
294            .get_object_database_id(table_id)
295            .await?;
296
297        let job_fragments = self
298            .metadata_manager
299            .get_job_fragments_by_id(table_id.as_job_id())
300            .await?;
301
302        let mut tracker = SingleTableRefreshProgressTracker::new();
303        {
304            let fragment_info_guard = shared_actor_infos.read_guard();
305            for (fragment_id, fragment) in &job_fragments.fragments {
306                if fragment
307                    .fragment_type_mask
308                    .contains(FragmentTypeFlag::Source)
309                    && !fragment.fragment_type_mask.contains(FragmentTypeFlag::Dml)
310                {
311                    let fragment_info = fragment_info_guard
312                        .get_fragment(*fragment_id)
313                        .ok_or_else(|| MetaError::fragment_not_found(*fragment_id))?;
314                    tracker.expected_list_actors.extend(
315                        fragment_info
316                            .actors
317                            .keys()
318                            .map(|actor_id| *actor_id as ActorId),
319                    );
320                }
321
322                if fragment
323                    .fragment_type_mask
324                    .contains(FragmentTypeFlag::FsFetch)
325                    && let Some(fragment_info) = fragment_info_guard.get_fragment(*fragment_id)
326                {
327                    tracker.expected_fetch_actors.extend(
328                        fragment_info
329                            .actors
330                            .keys()
331                            .map(|actor_id| *actor_id as ActorId),
332                    );
333                }
334            }
335        }
336
337        self.register_progress_tracker(table_id, database_id, tracker);
338
339        self.metadata_manager
340            .update_refresh_job_status(
341                table_id,
342                RefreshState::Refreshing,
343                Some(trigger_time),
344                false,
345            )
346            .await?;
347
348        let refresh_command = Command::Refresh {
349            table_id,
350            associated_source_id,
351        };
352
353        let result = self
354            .barrier_scheduler
355            .run_command(database_id, refresh_command)
356            .await;
357
358        match result {
359            Ok(_) => {
360                tracing::info!(table_id = %table_id, "refresh command scheduled");
361                Ok(())
362            }
363            Err(err) => {
364                tracing::error!(
365                    error = %err.as_report(),
366                    table_id = %table_id,
367                    "failed to execute refresh command"
368                );
369                self.metadata_manager
370                    .update_refresh_job_status(table_id, RefreshState::Idle, None, false)
371                    .await?;
372                self.remove_progress_tracker(table_id, "failure");
373                Err(anyhow!(err)
374                    .context(format!("Failed to refresh table {}", table_id))
375                    .into())
376            }
377        }
378    }
379
380    async fn ensure_refreshable(
381        &self,
382        table_id: TableId,
383        associated_source_id: SourceId,
384    ) -> MetaResult<()> {
385        let table = self
386            .metadata_manager
387            .catalog_controller
388            .get_table_by_id(table_id)
389            .await?;
390
391        if !table.refreshable {
392            return Err(MetaError::invalid_parameter(format!(
393                "Table '{}' is not refreshable. Only tables created with REFRESHABLE flag support refresh.",
394                table.name
395            )));
396        }
397
398        if table.optional_associated_source_id != Some(associated_source_id.into()) {
399            return Err(MetaError::invalid_parameter(format!(
400                "Table '{}' is not associated with source '{}'. table.optional_associated_source_id: {:?}",
401                table.name, associated_source_id, table.optional_associated_source_id
402            )));
403        }
404
405        let refresh_job_state = self
406            .metadata_manager
407            .catalog_controller
408            .get_refresh_job_state_by_table_id(table_id)
409            .await?;
410        if refresh_job_state != RefreshState::Idle {
411            return Err(MetaError::invalid_parameter(format!(
412                "Table '{}' is not in idle state. Current state: {:?}",
413                table.name, refresh_job_state
414            )));
415        }
416
417        Ok(())
418    }
419
420    fn register_progress_tracker(
421        &self,
422        table_id: TableId,
423        database_id: DatabaseId,
424        tracker: SingleTableRefreshProgressTracker,
425    ) {
426        let mut guard = self.progress_trackers.lock();
427        guard.inner.insert(table_id, tracker);
428        guard
429            .table_id_by_database_id
430            .entry(database_id)
431            .or_default()
432            .insert(table_id);
433    }
434
435    pub fn remove_progress_tracker(&self, table_id: TableId, status: &str) {
436        let elapsed = {
437            let mut guard = self.progress_trackers.lock();
438            let elapsed = guard
439                .inner
440                .remove(&table_id)
441                .map(|entry| entry.start_time.elapsed().as_secs());
442            guard.table_id_by_database_id.values_mut().for_each(|set| {
443                set.remove(&table_id);
444            });
445            elapsed
446        };
447        if let Some(elapsed) = elapsed {
448            let table_id_label = table_id.to_string();
449            GLOBAL_META_METRICS
450                .refresh_job_duration
451                .with_guarded_label_values(&[&table_id_label, status])
452                .set(elapsed);
453            let mut metrics = self.refresh_job_metrics.lock();
454            let metrics = metrics
455                .entry(table_id)
456                .or_insert_with(|| RefreshJobMetrics::new(table_id))
457                .finished
458                .entry(status.to_owned())
459                .or_insert_with(|| RefreshFinishedMetrics::new(table_id, status));
460            metrics.count.inc();
461        }
462    }
463}
464
465struct RefreshJobMetrics {
466    cron_trigger_count: LabelGuardedIntCounter,
467    cron_miss_count: LabelGuardedIntCounter,
468    finished: HashMap<String, RefreshFinishedMetrics>,
469}
470
471impl RefreshJobMetrics {
472    fn new(table_id: TableId) -> Self {
473        let table_id = table_id.to_string();
474        Self {
475            cron_trigger_count: GLOBAL_META_METRICS
476                .refresh_cron_job_trigger_cnt
477                .with_guarded_label_values(&[&table_id]),
478            cron_miss_count: GLOBAL_META_METRICS
479                .refresh_cron_job_miss_cnt
480                .with_guarded_label_values(&[&table_id]),
481            finished: HashMap::new(),
482        }
483    }
484}
485
486struct RefreshFinishedMetrics {
487    count: LabelGuardedIntCounter,
488}
489
490impl RefreshFinishedMetrics {
491    fn new(table_id: TableId, status: &str) -> Self {
492        let table_id = table_id.to_string();
493        Self {
494            count: GLOBAL_META_METRICS
495                .refresh_job_finish_cnt
496                .with_guarded_label_values(&[&table_id, status]),
497        }
498    }
499}
500
501#[derive(Default, Debug)]
502pub struct GlobalRefreshTableProgressTracker {
503    pub inner: HashMap<TableId, SingleTableRefreshProgressTracker>,
504    pub table_id_by_database_id: HashMap<DatabaseId, HashSet<TableId>>,
505}
506
507impl GlobalRefreshTableProgressTracker {
508    pub fn remove_tracker_by_database_id(&mut self, database_id: DatabaseId) {
509        if let Some(table_ids) = self.table_id_by_database_id.remove(&database_id) {
510            for table_id in table_ids {
511                self.inner.remove(&table_id);
512            }
513        }
514    }
515}
516
517#[derive(Debug)]
518pub struct SingleTableRefreshProgressTracker {
519    pub expected_list_actors: HashSet<ActorId>,
520    pub expected_fetch_actors: HashSet<ActorId>,
521    pub list_finished_actors: HashSet<ActorId>,
522    pub fetch_finished_actors: HashSet<ActorId>,
523
524    pub start_time: Instant,
525}
526
527impl SingleTableRefreshProgressTracker {
528    pub fn new() -> Self {
529        Self {
530            expected_list_actors: HashSet::new(),
531            expected_fetch_actors: HashSet::new(),
532            list_finished_actors: HashSet::new(),
533            fetch_finished_actors: HashSet::new(),
534            start_time: Instant::now(),
535        }
536    }
537
538    pub fn report_list_finished(&mut self, actor_ids: impl Iterator<Item = ActorId>) {
539        self.list_finished_actors.extend(actor_ids);
540    }
541
542    pub fn is_list_finished(&self) -> MetaResult<bool> {
543        if self.list_finished_actors.len() >= self.expected_list_actors.len() {
544            if self.expected_list_actors == self.list_finished_actors {
545                Ok(true)
546            } else {
547                Err(MetaError::from(anyhow!(
548                    "list finished actors mismatch: expected: {:?}, actual: {:?}",
549                    self.expected_list_actors,
550                    self.list_finished_actors
551                )))
552            }
553        } else {
554            Ok(false)
555        }
556    }
557
558    pub fn report_load_finished(&mut self, actor_ids: impl Iterator<Item = ActorId>) {
559        self.fetch_finished_actors.extend(actor_ids);
560    }
561
562    pub fn is_load_finished(&self) -> MetaResult<bool> {
563        if self.fetch_finished_actors.len() >= self.expected_fetch_actors.len() {
564            if self.expected_fetch_actors == self.fetch_finished_actors {
565                Ok(true)
566            } else {
567                Err(MetaError::from(anyhow!(
568                    "fetch finished actors mismatch: expected: {:?}, actual: {:?}",
569                    self.expected_fetch_actors,
570                    self.fetch_finished_actors
571                )))
572            }
573        } else {
574            Ok(false)
575        }
576    }
577}