|
| 1 | +use std::time::Instant; |
| 2 | + |
| 3 | +use crate::prelude::*; |
| 4 | + |
| 5 | +use super::stats; |
| 6 | +use sqlx::PgPool; |
| 7 | +use tokio::task::JoinSet; |
| 8 | + |
| 9 | +pub struct FlowSynchronizer { |
| 10 | + flow_ctx: Arc<FlowContext>, |
| 11 | + tasks: JoinSet<Result<()>>, |
| 12 | + sources_update_stats: Vec<Arc<stats::UpdateStats>>, |
| 13 | +} |
| 14 | + |
| 15 | +pub struct FlowSynchronizerOptions { |
| 16 | + pub keep_refreshed: bool, |
| 17 | +} |
| 18 | + |
| 19 | +async fn sync_source( |
| 20 | + flow_ctx: Arc<FlowContext>, |
| 21 | + plan: Arc<plan::ExecutionPlan>, |
| 22 | + source_update_stats: Arc<stats::UpdateStats>, |
| 23 | + source_idx: usize, |
| 24 | + pool: PgPool, |
| 25 | + keep_refreshed: bool, |
| 26 | +) -> Result<()> { |
| 27 | + let source_context = flow_ctx |
| 28 | + .get_source_indexing_context(source_idx, &pool) |
| 29 | + .await?; |
| 30 | + |
| 31 | + let mut update_start = Instant::now(); |
| 32 | + source_context.update(&pool, &source_update_stats).await?; |
| 33 | + |
| 34 | + let import_op = &plan.import_ops[source_idx]; |
| 35 | + if let (true, Some(refresh_interval)) = |
| 36 | + (keep_refreshed, import_op.refresh_options.refresh_interval) |
| 37 | + { |
| 38 | + loop { |
| 39 | + let elapsed = update_start.elapsed(); |
| 40 | + if elapsed < refresh_interval { |
| 41 | + tokio::time::sleep(refresh_interval - elapsed).await; |
| 42 | + } |
| 43 | + update_start = Instant::now(); |
| 44 | + source_context.update(&pool, &source_update_stats).await?; |
| 45 | + } |
| 46 | + } |
| 47 | + Ok(()) |
| 48 | +} |
| 49 | + |
| 50 | +impl FlowSynchronizer { |
| 51 | + pub async fn start( |
| 52 | + flow_ctx: Arc<FlowContext>, |
| 53 | + pool: &PgPool, |
| 54 | + options: &FlowSynchronizerOptions, |
| 55 | + ) -> Result<Self> { |
| 56 | + let plan = flow_ctx.flow.get_execution_plan().await?; |
| 57 | + |
| 58 | + let mut tasks = JoinSet::new(); |
| 59 | + let sources_update_stats = (0..plan.import_ops.len()) |
| 60 | + .map(|source_idx| { |
| 61 | + let source_update_stats = Arc::new(stats::UpdateStats::default()); |
| 62 | + tasks.spawn(sync_source( |
| 63 | + flow_ctx.clone(), |
| 64 | + plan.clone(), |
| 65 | + source_update_stats.clone(), |
| 66 | + source_idx, |
| 67 | + pool.clone(), |
| 68 | + options.keep_refreshed, |
| 69 | + )); |
| 70 | + source_update_stats |
| 71 | + }) |
| 72 | + .collect(); |
| 73 | + Ok(Self { |
| 74 | + flow_ctx, |
| 75 | + tasks, |
| 76 | + sources_update_stats, |
| 77 | + }) |
| 78 | + } |
| 79 | + |
| 80 | + pub async fn join(&mut self) -> Result<()> { |
| 81 | + while let Some(result) = self.tasks.join_next().await { |
| 82 | + if let Err(e) = (|| anyhow::Ok(result??))() { |
| 83 | + error!("{:?}", e.context("Error in synchronizing a source")); |
| 84 | + } |
| 85 | + } |
| 86 | + Ok(()) |
| 87 | + } |
| 88 | + |
| 89 | + pub fn abort(&mut self) { |
| 90 | + self.tasks.abort_all(); |
| 91 | + } |
| 92 | + |
| 93 | + pub fn index_update_info(&self) -> stats::IndexUpdateInfo { |
| 94 | + stats::IndexUpdateInfo { |
| 95 | + sources: std::iter::zip( |
| 96 | + self.flow_ctx.flow.flow_instance.import_ops.iter(), |
| 97 | + self.sources_update_stats.iter(), |
| 98 | + ) |
| 99 | + .map(|(import_op, stats)| stats::SourceUpdateInfo { |
| 100 | + source_name: import_op.name.clone(), |
| 101 | + stats: (&**stats).clone(), |
| 102 | + }) |
| 103 | + .collect(), |
| 104 | + } |
| 105 | + } |
| 106 | +} |
0 commit comments