feat(global-concur-control): provide a global controller

Author: georgeh0

docs/docs/core/flow_def.mdx

@@ -370,7 +370,7 @@ These options can be passed in to the following APIs:
 
 *   [`FlowBuilder.add_source()`](#import-from-source): The options above control the processing concurrency of multiple rows from a source. New rows will not be loaded in memory if it'll be over the limit.
 
-    The default value can be specified by [`DefaultExecutionOptions`](/docs/core/settings#defaultexecutionoptions) or corresponding [environment variables](/docs/core/settings#list-of-environment-variables).
+    The default value can be specified by [`GlobalExecutionOptions`](/docs/core/settings#defaultexecutionoptions) or corresponding [environment variables](/docs/core/settings#list-of-environment-variables).
 
 *   [`DataSlice.row()`](#for-each-row):  The options above provides a finer-grained control, to limit the processing concurrency of multiple rows within a table at any level.
 

docs/docs/core/settings.mdx

@@ -60,7 +60,7 @@ You have two ways to launch CocoIndex:
 
 *   `app_namespace` (type: `str`, required): The namespace of the application.
 *   `database` (type: `DatabaseConnectionSpec`, required): The connection to the Postgres database.
-*   `default_execution_options` (type: `DefaultExecutionOptions`, optional): The default execution options for the flow.
+*   `global_execution_options` (type: `GlobalExecutionOptions`, optional): The default execution options for the flow.
 
 ### App Namespace
 
@@ -105,9 +105,9 @@ If you use the Postgres database hosted by [Supabase](https://supabase.com/), pl
 
 :::
 
-### DefaultExecutionOptions
+### GlobalExecutionOptions
 
-`DefaultExecutionOptions` is used to configure the default execution options for the flow. It has the following fields:
+`GlobalExecutionOptions` is used to configure the default execution options for the flow. It has the following fields:
 
 *   `source_max_inflight_rows` (type: `int`, optional): The maximum number of concurrent inflight requests for source operations.
 *   `source_max_inflight_bytes` (type: `int`, optional): The maximum number of concurrent inflight bytes for source operations.
@@ -124,5 +124,5 @@ This is the list of environment variables, each of which has a corresponding fie
 | `COCOINDEX_DATABASE_URL` | `database.url` | Yes |
 | `COCOINDEX_DATABASE_USER` | `database.user` | No |
 | `COCOINDEX_DATABASE_PASSWORD` | `database.password` | No |
-| `COCOINDEX_SOURCE_MAX_INFLIGHT_ROWS` | `default_execution_options.source_max_inflight_rows` | No |
-| `COCOINDEX_SOURCE_MAX_INFLIGHT_BYTES` | `default_execution_options.source_max_inflight_bytes` | No |
+| `COCOINDEX_SOURCE_MAX_INFLIGHT_ROWS` | `global_execution_options.source_max_inflight_rows` | No |
+| `COCOINDEX_SOURCE_MAX_INFLIGHT_BYTES` | `global_execution_options.source_max_inflight_bytes` | No |

python/cocoindex/setting.py

@@ -44,7 +44,7 @@ class DatabaseConnectionSpec:
 
 
 @dataclass
-class DefaultExecutionOptions:
+class GlobalExecutionOptions:
     """Default execution options."""
 
     # The maximum number of concurrent inflight requests.
@@ -81,7 +81,7 @@ class Settings:
 
     database: DatabaseConnectionSpec | None = None
     app_namespace: str = ""
-    default_execution_options: DefaultExecutionOptions | None = None
+    global_execution_options: GlobalExecutionOptions | None = None
 
     @classmethod
     def from_env(cls) -> Self:
@@ -110,14 +110,14 @@ def from_env(cls) -> Self:
             "COCOINDEX_SOURCE_MAX_INFLIGHT_BYTES",
             parse=int,
         )
-        default_execution_options = DefaultExecutionOptions(**exec_kwargs)
+        global_execution_options = GlobalExecutionOptions(**exec_kwargs)
 
         app_namespace = os.getenv("COCOINDEX_APP_NAMESPACE", "")
 
         return cls(
             database=database,
             app_namespace=app_namespace,
-            default_execution_options=default_execution_options,
+            global_execution_options=global_execution_options,
         )
 
 

src/base/spec.rs

@@ -262,6 +262,15 @@ pub struct ExecutionOptions {
     pub max_inflight_bytes: Option<usize>,
 }
 
+impl ExecutionOptions {
+    pub fn get_concur_control_options(&self) -> concur_control::Options {
+        concur_control::Options {
+            max_inflight_rows: self.max_inflight_rows,
+            max_inflight_bytes: self.max_inflight_bytes,
+        }
+    }
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize, Default)]
 pub struct SourceRefreshOptions {
     pub refresh_interval: Option<std::time::Duration>,

src/builder/analyzer.rs

@@ -689,18 +689,11 @@ impl AnalyzerContext {
             .clone();
         let output = op_scope.add_op_output(import_op.name, output_type)?;
 
-        let max_inflight_rows =
-            (import_op.spec.execution_options.max_inflight_rows).or_else(|| {
-                self.lib_ctx
-                    .default_execution_options
-                    .source_max_inflight_rows
-            });
-        let max_inflight_bytes =
-            (import_op.spec.execution_options.max_inflight_bytes).or_else(|| {
-                self.lib_ctx
-                    .default_execution_options
-                    .source_max_inflight_bytes
-            });
+        let concur_control_options = import_op
+            .spec
+            .execution_options
+            .get_concur_control_options();
+        let global_concurrency_controller = self.lib_ctx.global_concurrency_controller.clone();
         let result_fut = async move {
             trace!("Start building executor for source op `{}`", op_name);
             let executor = executor.await?;
@@ -711,9 +704,9 @@ impl AnalyzerContext {
                 primary_key_type,
                 name: op_name,
                 refresh_options: import_op.spec.refresh_options,
-                concurrency_controller: concur_control::ConcurrencyController::new(
-                    max_inflight_rows,
-                    max_inflight_bytes,
+                concurrency_controller: concur_control::CombinedConcurrencyController::new(
+                    &concur_control_options,
+                    global_concurrency_controller,
                 ),
             })
         };
@@ -808,7 +801,8 @@ impl AnalyzerContext {
                 };
                 let op_name = reactive_op.name.clone();
 
-                let exec_options = foreach_op.execution_options.clone();
+                let concur_control_options =
+                    foreach_op.execution_options.get_concur_control_options();
                 async move {
                     Ok(AnalyzedReactiveOp::ForEach(AnalyzedForEachOp {
                         local_field_ref,
@@ -817,8 +811,7 @@ impl AnalyzerContext {
                             .with_context(|| format!("Analyzing foreach op: {op_name}"))?,
                         name: op_name,
                         concurrency_controller: concur_control::ConcurrencyController::new(
-                            exec_options.max_inflight_rows,
-                            exec_options.max_inflight_bytes,
+                            &concur_control_options,
                         ),
                     }))
                 }

src/builder/plan.rs

@@ -57,7 +57,7 @@ pub struct AnalyzedImportOp {
     pub primary_key_type: schema::ValueType,
     pub refresh_options: spec::SourceRefreshOptions,
 
-    pub concurrency_controller: concur_control::ConcurrencyController,
+    pub concurrency_controller: concur_control::CombinedConcurrencyController,
 }
 
 pub struct AnalyzedFunctionExecInfo {

src/execution/source_indexer.rs

@@ -108,7 +108,7 @@ impl SourceIndexingContext {
         key: value::KeyValue,
         source_data: Option<interface::SourceData>,
         update_stats: Arc<stats::UpdateStats>,
-        _concur_permit: concur_control::ConcurrencyControllerPermit,
+        _concur_permit: concur_control::CombinedConcurrencyControllerPermit,
         ack_fn: Option<AckFn>,
         pool: PgPool,
     ) {
@@ -247,36 +247,6 @@ impl SourceIndexingContext {
         }
     }
 
-    // Expected to be called during scan, which has no value.
-    fn process_source_key_if_newer(
-        self: &Arc<Self>,
-        key: value::KeyValue,
-        source_version: SourceVersion,
-        update_stats: &Arc<stats::UpdateStats>,
-        concur_permit: concur_control::ConcurrencyControllerPermit,
-        pool: &PgPool,
-    ) -> Option<impl Future<Output = ()> + Send + 'static> {
-        {
-            let mut state = self.state.lock().unwrap();
-            let scan_generation = state.scan_generation;
-            let row_state = state.rows.entry(key.clone()).or_default();
-            row_state.touched_generation = scan_generation;
-            if row_state
-                .source_version
-                .should_skip(&source_version, Some(update_stats.as_ref()))
-            {
-                return None;
-            }
-        }
-        Some(self.clone().process_source_key(
-            key,
-            None,
-            update_stats.clone(),
-            concur_permit,
-            NO_ACK,
-            pool.clone(),
-        ))
-    }
     pub async fn update(
         self: &Arc<Self>,
         pool: &PgPool,
@@ -336,21 +306,34 @@ impl SourceIndexingContext {
         };
         while let Some(row) = rows_stream.next().await {
             for row in row? {
+                let source_version = SourceVersion::from_current_with_ordinal(
+                    row.ordinal
+                        .ok_or_else(|| anyhow::anyhow!("ordinal is not available"))?,
+                );
+                {
+                    let mut state = self.state.lock().unwrap();
+                    let scan_generation = state.scan_generation;
+                    let row_state = state.rows.entry(row.key.clone()).or_default();
+                    row_state.touched_generation = scan_generation;
+                    if row_state
+                        .source_version
+                        .should_skip(&source_version, Some(update_stats.as_ref()))
+                    {
+                        continue;
+                    }
+                }
                 let concur_permit = import_op
                     .concurrency_controller
                     .acquire(concur_control::BYTES_UNKNOWN_YET)
                     .await?;
-                self.process_source_key_if_newer(
+                join_set.spawn(self.clone().process_source_key(
                     row.key,
-                    SourceVersion::from_current_with_ordinal(
-                        row.ordinal
-                            .ok_or_else(|| anyhow::anyhow!("ordinal is not available"))?,
-                    ),
-                    update_stats,
+                    None,
+                    update_stats.clone(),
                     concur_permit,
-                    pool,
-                )
-                .map(|fut| join_set.spawn(fut));
+                    NO_ACK,
+                    pool.clone(),
+                ));
             }
         }
         while let Some(result) = join_set.join_next().await {

src/lib_context.rs

@@ -196,7 +196,7 @@ pub struct LibContext {
     pub persistence_ctx: Option<PersistenceContext>,
     pub flows: Mutex<BTreeMap<String, Arc<FlowContext>>>,
 
-    pub default_execution_options: settings::DefaultExecutionOptions,
+    pub global_concurrency_controller: Arc<concur_control::ConcurrencyController>,
 }
 
 impl LibContext {
@@ -269,7 +269,12 @@ pub fn create_lib_context(settings: settings::Settings) -> Result<LibContext> {
         db_pools,
         persistence_ctx,
         flows: Mutex::new(BTreeMap::new()),
-        default_execution_options: settings.default_execution_options,
+        global_concurrency_controller: Arc::new(concur_control::ConcurrencyController::new(
+            &concur_control::Options {
+                max_inflight_rows: settings.global_execution_options.source_max_inflight_rows,
+                max_inflight_bytes: settings.global_execution_options.source_max_inflight_bytes,
+            },
+        )),
     })
 }
 

src/settings.rs

@@ -8,7 +8,7 @@ pub struct DatabaseConnectionSpec {
 }
 
 #[derive(Deserialize, Debug, Default)]
-pub struct DefaultExecutionOptions {
+pub struct GlobalExecutionOptions {
     pub source_max_inflight_rows: Option<usize>,
     pub source_max_inflight_bytes: Option<usize>,
 }
@@ -21,7 +21,7 @@ pub struct Settings {
     #[allow(dead_code)] // Used via serialization/deserialization to Python
     pub app_namespace: String,
     #[serde(default)]
-    pub default_execution_options: DefaultExecutionOptions,
+    pub global_execution_options: GlobalExecutionOptions,
 }
 
 #[cfg(test)]

src/utils/concur_control.rs

@@ -48,6 +48,11 @@ impl WeightedSemaphore {
     }
 }
 
+pub struct Options {
+    pub max_inflight_rows: Option<usize>,
+    pub max_inflight_bytes: Option<usize>,
+}
+
 pub struct ConcurrencyControllerPermit {
     _inflight_count_permit: Option<OwnedSemaphorePermit>,
     _inflight_bytes_permit: Option<OwnedSemaphorePermit>,
@@ -61,10 +66,14 @@ pub struct ConcurrencyController {
 pub static BYTES_UNKNOWN_YET: Option<fn() -> usize> = None;
 
 impl ConcurrencyController {
-    pub fn new(max_inflight_count: Option<usize>, max_inflight_bytes: Option<usize>) -> Self {
+    pub fn new(exec_options: &Options) -> Self {
         Self {
-            inflight_count_sem: max_inflight_count.map(|max| Arc::new(Semaphore::new(max))),
-            inflight_bytes_sem: max_inflight_bytes.map(|max| WeightedSemaphore::new(max)),
+            inflight_count_sem: exec_options
+                .max_inflight_rows
+                .map(|max| Arc::new(Semaphore::new(max))),
+            inflight_bytes_sem: exec_options
+                .max_inflight_bytes
+                .map(|max| WeightedSemaphore::new(max)),
         }
     }
 
@@ -82,8 +91,7 @@ impl ConcurrencyController {
         };
         let inflight_bytes_permit = if let Some(sem) = &self.inflight_bytes_sem {
             if let Some(bytes_fn) = bytes_fn {
-                let n = bytes_fn();
-                sem.acquire(n, false).await?
+                sem.acquire(bytes_fn(), false).await?
             } else {
                 Some(sem.acquire_reservation().await?)
             }
@@ -107,3 +115,62 @@ impl ConcurrencyController {
         }
     }
 }
+
+pub struct CombinedConcurrencyControllerPermit {
+    _permit: ConcurrencyControllerPermit,
+    _global_permit: ConcurrencyControllerPermit,
+}
+
+pub struct CombinedConcurrencyController {
+    controller: ConcurrencyController,
+    global_controller: Arc<ConcurrencyController>,
+    needs_num_bytes: bool,
+}
+
+impl CombinedConcurrencyController {
+    pub fn new(exec_options: &Options, global_controller: Arc<ConcurrencyController>) -> Self {
+        Self {
+            controller: ConcurrencyController::new(exec_options),
+            needs_num_bytes: exec_options.max_inflight_bytes.is_some()
+                || global_controller.inflight_bytes_sem.is_some(),
+            global_controller,
+        }
+    }
+
+    pub async fn acquire(
+        &self,
+        bytes_fn: Option<impl FnOnce() -> usize>,
+    ) -> Result<CombinedConcurrencyControllerPermit, AcquireError> {
+        let num_bytes_fn = if let Some(bytes_fn) = bytes_fn
+            && self.needs_num_bytes
+        {
+            let num_bytes = bytes_fn();
+            Some(move || num_bytes)
+        } else {
+            None
+        };
+
+        let permit = self.controller.acquire(num_bytes_fn).await?;
+        let global_permit = self.global_controller.acquire(num_bytes_fn).await?;
+        Ok(CombinedConcurrencyControllerPermit {
+            _permit: permit,
+            _global_permit: global_permit,
+        })
+    }
+
+    pub async fn acquire_bytes_with_reservation<'a>(
+        &'a self,
+        bytes_fn: impl FnOnce() -> usize,
+    ) -> Result<(Option<OwnedSemaphorePermit>, Option<OwnedSemaphorePermit>), AcquireError> {
+        let num_bytes = bytes_fn();
+        let permit = self
+            .controller
+            .acquire_bytes_with_reservation(move || num_bytes)
+            .await?;
+        let global_permit = self
+            .global_controller
+            .acquire_bytes_with_reservation(move || num_bytes)
+            .await?;
+        Ok((permit, global_permit))
+    }
+}