Add a synchronizer and expose keep_in_sync API to Python SDK.

Author: badmonster0

src/base/spec.rs

@@ -164,7 +164,7 @@ pub struct OpSpec {
 
 #[derive(Debug, Clone, Serialize, Deserialize, Default)]
 pub struct SourceRefreshOptions {
-    pub refresh_interval: Option<chrono::Duration>,
+    pub refresh_interval: Option<std::time::Duration>,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]

src/execution/mod.rs

@@ -7,4 +7,7 @@ pub(crate) mod row_indexer;
 pub(crate) mod source_indexer;
 pub(crate) mod stats;
 
+mod synchronizer;
+pub(crate) use synchronizer::*;
+
 mod db_tracking;

src/execution/source_indexer.rs

@@ -1,15 +1,14 @@
-use std::collections::{hash_map, HashMap};
-
 use crate::prelude::*;
 
+use sqlx::PgPool;
+use std::collections::{hash_map, HashMap};
+use tokio::{sync::Semaphore, task::JoinSet};
+
 use super::{
     db_tracking,
     row_indexer::{self, SkippedOr, SourceVersion},
     stats,
 };
-use futures::future::try_join_all;
-use sqlx::PgPool;
-use tokio::{sync::Semaphore, task::JoinSet};
 struct SourceRowIndexingState {
     source_version: SourceVersion,
     processing_sem: Arc<Semaphore>,
@@ -190,7 +189,11 @@ impl SourceIndexingContext {
         );
     }
 
-    async fn update_source(self: &Arc<Self>, pool: &PgPool) -> Result<stats::SourceUpdateInfo> {
+    pub async fn update(
+        self: &Arc<Self>,
+        pool: &PgPool,
+        update_stats: &Arc<stats::UpdateStats>,
+    ) -> Result<()> {
         let plan = self.flow.get_execution_plan().await?;
         let import_op = &plan.import_ops[self.source_idx];
         let mut rows_stream = import_op
@@ -204,13 +207,12 @@ impl SourceIndexingContext {
             state.scan_generation += 1;
             state.scan_generation
         };
-        let update_stats = Arc::new(stats::UpdateStats::default());
         while let Some(row) = rows_stream.next().await {
             for row in row? {
                 self.process_source_key_if_newer(
                     row.key,
                     SourceVersion::from_current(row.ordinal),
-                    &update_stats,
+                    update_stats,
                     pool,
                     &mut join_set,
                 );
@@ -252,25 +254,6 @@ impl SourceIndexingContext {
             }
         }
 
-        Ok(stats::SourceUpdateInfo {
-            source_name: import_op.name.clone(),
-            stats: Arc::unwrap_or_clone(update_stats),
-        })
+        Ok(())
     }
 }
-
-pub async fn update(flow_context: &FlowContext, pool: &PgPool) -> Result<stats::IndexUpdateInfo> {
-    let plan = flow_context.flow.get_execution_plan().await?;
-    let source_update_stats = try_join_all(
-        (0..plan.import_ops.len())
-            .map(|idx| async move {
-                let source_context = flow_context.get_source_indexing_context(idx, pool).await?;
-                source_context.update_source(pool).await
-            })
-            .collect::<Vec<_>>(),
-    )
-    .await?;
-    Ok(stats::IndexUpdateInfo {
-        sources: source_update_stats,
-    })
-}

src/execution/synchronizer.rs

@@ -0,0 +1,106 @@
+use std::time::Instant;
+
+use crate::prelude::*;
+
+use super::stats;
+use sqlx::PgPool;
+use tokio::task::JoinSet;
+
+pub struct FlowSynchronizer {
+    flow_ctx: Arc<FlowContext>,
+    tasks: JoinSet<Result<()>>,
+    sources_update_stats: Vec<Arc<stats::UpdateStats>>,
+}
+
+pub struct FlowSynchronizerOptions {
+    pub keep_refreshed: bool,
+}
+
+async fn sync_source(
+    flow_ctx: Arc<FlowContext>,
+    plan: Arc<plan::ExecutionPlan>,
+    source_update_stats: Arc<stats::UpdateStats>,
+    source_idx: usize,
+    pool: PgPool,
+    keep_refreshed: bool,
+) -> Result<()> {
+    let source_context = flow_ctx
+        .get_source_indexing_context(source_idx, &pool)
+        .await?;
+
+    let mut update_start = Instant::now();
+    source_context.update(&pool, &source_update_stats).await?;
+
+    let import_op = &plan.import_ops[source_idx];
+    if let (true, Some(refresh_interval)) =
+        (keep_refreshed, import_op.refresh_options.refresh_interval)
+    {
+        loop {
+            let elapsed = update_start.elapsed();
+            if elapsed < refresh_interval {
+                tokio::time::sleep(refresh_interval - elapsed).await;
+            }
+            update_start = Instant::now();
+            source_context.update(&pool, &source_update_stats).await?;
+        }
+    }
+    Ok(())
+}
+
+impl FlowSynchronizer {
+    pub async fn start(
+        flow_ctx: Arc<FlowContext>,
+        pool: &PgPool,
+        options: &FlowSynchronizerOptions,
+    ) -> Result<Self> {
+        let plan = flow_ctx.flow.get_execution_plan().await?;
+
+        let mut tasks = JoinSet::new();
+        let sources_update_stats = (0..plan.import_ops.len())
+            .map(|source_idx| {
+                let source_update_stats = Arc::new(stats::UpdateStats::default());
+                tasks.spawn(sync_source(
+                    flow_ctx.clone(),
+                    plan.clone(),
+                    source_update_stats.clone(),
+                    source_idx,
+                    pool.clone(),
+                    options.keep_refreshed,
+                ));
+                source_update_stats
+            })
+            .collect();
+        Ok(Self {
+            flow_ctx,
+            tasks,
+            sources_update_stats,
+        })
+    }
+
+    pub async fn join(&mut self) -> Result<()> {
+        while let Some(result) = self.tasks.join_next().await {
+            if let Err(e) = (|| anyhow::Ok(result??))() {
+                error!("{:?}", e.context("Error in synchronizing a source"));
+            }
+        }
+        Ok(())
+    }
+
+    pub fn abort(&mut self) {
+        self.tasks.abort_all();
+    }
+
+    pub fn index_update_info(&self) -> stats::IndexUpdateInfo {
+        stats::IndexUpdateInfo {
+            sources: std::iter::zip(
+                self.flow_ctx.flow.flow_instance.import_ops.iter(),
+                self.sources_update_stats.iter(),
+            )
+            .map(|(import_op, stats)| stats::SourceUpdateInfo {
+                source_name: import_op.name.clone(),
+                stats: (&**stats).clone(),
+            })
+            .collect(),
+        }
+    }
+}

src/py/mod.rs

@@ -92,6 +92,52 @@ impl IndexUpdateInfo {
 #[pyclass]
 pub struct Flow(pub Arc<FlowContext>);
 
+#[pyclass]
+pub struct FlowSynchronizer(pub async_lock::RwLock<execution::FlowSynchronizer>);
+
+#[pymethods]
+impl FlowSynchronizer {
+    pub fn join(&self, py: Python<'_>) -> PyResult<()> {
+        py.allow_threads(|| {
+            let lib_context = get_lib_context()
+                .ok_or_else(|| PyException::new_err("cocoindex library not initialized"))?;
+            lib_context
+                .runtime
+                .block_on(async {
+                    let mut synchronizer = self.0.write().await;
+                    synchronizer.join().await
+                })
+                .into_py_result()
+        })
+    }
+
+    pub fn abort(&self, py: Python<'_>) -> PyResult<()> {
+        py.allow_threads(|| {
+            let lib_context = get_lib_context()
+                .ok_or_else(|| PyException::new_err("cocoindex library not initialized"))?;
+            lib_context.runtime.block_on(async {
+                let mut synchronizer = self.0.write().await;
+                synchronizer.abort();
+            });
+            Ok(())
+        })
+    }
+
+    pub fn index_update_info(&self, py: Python<'_>) -> PyResult<IndexUpdateInfo> {
+        py.allow_threads(|| {
+            let lib_context = get_lib_context()
+                .ok_or_else(|| PyException::new_err("cocoindex library not initialized"))?;
+            lib_context
+                .runtime
+                .block_on(async {
+                    let synchronizer = self.0.read().await;
+                    anyhow::Ok(IndexUpdateInfo(synchronizer.index_update_info()))
+                })
+                .into_py_result()
+        })
+    }
+}
+
 #[pymethods]
 impl Flow {
     pub fn __str__(&self) -> String {
@@ -113,13 +159,44 @@ impl Flow {
             let update_info = lib_context
                 .runtime
                 .block_on(async {
-                    execution::source_indexer::update(&self.0, &lib_context.pool).await
+                    let mut synchronizer = execution::FlowSynchronizer::start(
+                        self.0.clone(),
+                        &lib_context.pool,
+                        &execution::FlowSynchronizerOptions {
+                            keep_refreshed: false,
+                        },
+                    )
+                    .await?;
+                    synchronizer.join().await?;
+                    anyhow::Ok(synchronizer.index_update_info())
                 })
                 .into_py_result()?;
             Ok(IndexUpdateInfo(update_info))
         })
     }
 
+    pub fn keep_in_sync(&self, py: Python<'_>) -> PyResult<FlowSynchronizer> {
+        py.allow_threads(|| {
+            let lib_context = get_lib_context()
+                .ok_or_else(|| PyException::new_err("cocoindex library not initialized"))?;
+            let synchronizer = lib_context
+                .runtime
+                .block_on(async {
+                    let synchronizer = execution::FlowSynchronizer::start(
+                        self.0.clone(),
+                        &lib_context.pool,
+                        &execution::FlowSynchronizerOptions {
+                            keep_refreshed: false,
+                        },
+                    )
+                    .await?;
+                    anyhow::Ok(synchronizer)
+                })
+                .into_py_result()?;
+            Ok(FlowSynchronizer(async_lock::RwLock::new(synchronizer)))
+        })
+    }
+
     pub fn evaluate_and_dump(
         &self,
         py: Python<'_>,
@@ -308,6 +385,7 @@ fn cocoindex_engine(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<builder::flow_builder::DataSlice>()?;
     m.add_class::<builder::flow_builder::DataScopeRef>()?;
     m.add_class::<Flow>()?;
+    m.add_class::<FlowSynchronizer>()?;
     m.add_class::<TransientFlow>()?;
     m.add_class::<IndexUpdateInfo>()?;
     m.add_class::<SimpleSemanticsQueryHandler>()?;

src/service/flows.rs

@@ -1,11 +1,11 @@
 use crate::prelude::*;
 
+use crate::lib_context::LibContext;
 use crate::{base::schema::DataSchema, ops::interface::SourceExecutorListOptions};
 use crate::{
     execution::memoization,
     execution::{row_indexer, stats},
 };
-use crate::{execution::source_indexer, lib_context::LibContext};
 use axum::{
     extract::{Path, State},
     http::StatusCode,
@@ -171,6 +171,14 @@ pub async fn update(
     State(lib_context): State<Arc<LibContext>>,
 ) -> Result<Json<stats::IndexUpdateInfo>, ApiError> {
     let flow_ctx = lib_context.get_flow_context(&flow_name)?;
-    let update_info = source_indexer::update(&flow_ctx, &lib_context.pool).await?;
-    Ok(Json(update_info))
+    let mut synchronizer = execution::FlowSynchronizer::start(
+        flow_ctx.clone(),
+        &lib_context.pool,
+        &execution::FlowSynchronizerOptions {
+            keep_refreshed: false,
+        },
+    )
+    .await?;
+    synchronizer.join().await?;
+    Ok(Json(synchronizer.index_update_info()))
 }