feat: add timeout to prevent long-running functions (Closes #658)

AdwitaSingh1711 · AdwitaSingh1711 · commit edd09dffabc4 · 2025-11-02T16:47:42.000+05:30
diff --git a/python/cocoindex/flow.py b/python/cocoindex/flow.py
@@ -486,6 +486,7 @@ class _SourceRefreshOptions:
 class _ExecutionOptions:
     max_inflight_rows: int | None = None
     max_inflight_bytes: int | None = None
+    timeout: int | None = None
 
 
 class FlowBuilder:
diff --git a/python/cocoindex/op.py b/python/cocoindex/op.py
@@ -154,6 +154,7 @@ class OpArgs:
     - batching: Whether the executor will be batched.
     - behavior_version: The behavior version of the executor. Cache will be invalidated if it
       changes. Must be provided if `cache` is True.
+    - timeout: Timeout in seconds for this function execution. None means use default (300s).
     - arg_relationship: It specifies the relationship between an input argument and the output,
       e.g. `(ArgRelationship.CHUNKS_BASE_TEXT, "content")` means the output is chunks for the
       input argument with name `content`.
@@ -163,6 +164,7 @@ class OpArgs:
     cache: bool = False
     batching: bool = False
     behavior_version: int | None = None
+    timeout: int | None = None
     arg_relationship: tuple[ArgRelationship, str] | None = None
 
 
@@ -201,6 +203,7 @@ def _register_op_factory(
 
     class _WrappedExecutor:
         _executor: Any
+        _spec: Any
         _args_info: list[_ArgInfo]
         _kwargs_info: dict[str, _ArgInfo]
         _result_encoder: Callable[[Any], Any]
@@ -390,6 +393,12 @@ def enable_cache(self) -> bool:
         def behavior_version(self) -> int | None:
             return op_args.behavior_version
 
+        def timeout(self) -> int | None:  # ✅ Add this
+            if op_args.timeout is not None:
+                return op_args.timeout
+
+            return None
+
     if category == OpCategory.FUNCTION:
         _engine.register_function_factory(
             op_kind,
diff --git a/src/base/spec.rs b/src/base/spec.rs
@@ -234,6 +234,9 @@ pub struct ExecutionOptions {
 
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub max_inflight_bytes: Option<usize>,
+
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub timeout: Option<std::time::Duration>,
 }
 
 impl ExecutionOptions {
@@ -289,6 +292,9 @@ impl fmt::Display for ImportOpSpec {
 pub struct TransformOpSpec {
     pub inputs: Vec<OpArgBinding>,
     pub op: OpSpec,
+
+    #[serde(default)]
+    pub execution_options: ExecutionOptions,
 }
 
 impl SpecFormatter for TransformOpSpec {
diff --git a/src/builder/analyzer.rs b/src/builder/analyzer.rs
@@ -13,6 +13,7 @@ use crate::{
 };
 use futures::future::{BoxFuture, try_join3};
 use futures::{FutureExt, future::try_join_all};
+use tokio::time::Duration;
 
 #[derive(Debug)]
 pub(super) enum ValueTypeBuilder {
@@ -727,16 +728,22 @@ impl AnalyzerContext {
                 let output =
                     op_scope.add_op_output(reactive_op.name.clone(), output_enriched_type)?;
                 let op_name = reactive_op.name.clone();
+                let op_kind = op.op.kind.clone();
+                let execution_options_timeout = op.execution_options.timeout;
                 async move {
                             trace!("Start building executor for transform op `{op_name}`");
                             let executor = executor.await.with_context(|| {
                                 format!("Preparing for transform op: {op_name}")
                             })?;
                             let enable_cache = executor.enable_cache();
                             let behavior_version = executor.behavior_version();
+                            let timeout = executor.timeout()
+                                .or(execution_options_timeout)
+                                .or(Some(Duration::from_secs(300)));
                             trace!("Finished building executor for transform op `{op_name}`, enable cache: {enable_cache}, behavior version: {behavior_version:?}");
                             let function_exec_info = AnalyzedFunctionExecInfo {
                                 enable_cache,
+                                timeout,
                                 behavior_version,
                                 fingerprinter: logic_fingerprinter
                                     .with(&behavior_version)?,
@@ -751,6 +758,7 @@ impl AnalyzerContext {
                             }
                             Ok(AnalyzedReactiveOp::Transform(AnalyzedTransformOp {
                                 name: op_name,
+                                op_kind,
                                 inputs: input_value_mappings,
                                 function_exec_info,
                                 executor,
diff --git a/src/builder/flow_builder.rs b/src/builder/flow_builder.rs
@@ -461,6 +461,7 @@ impl FlowBuilder {
                     })
                     .collect(),
                 op: spec,
+                execution_options: Default::default(),
             }),
         };
 
diff --git a/src/builder/plan.rs b/src/builder/plan.rs
@@ -3,6 +3,7 @@ use crate::prelude::*;
 
 use crate::ops::interface::*;
 use crate::utils::fingerprint::{Fingerprint, Fingerprinter};
+use std::time::Duration;
 
 #[derive(Debug, Clone, PartialEq, Eq, Serialize)]
 pub struct AnalyzedLocalFieldReference {
@@ -63,6 +64,7 @@ pub struct AnalyzedImportOp {
 
 pub struct AnalyzedFunctionExecInfo {
     pub enable_cache: bool,
+    pub timeout: Option<Duration>,
     pub behavior_version: Option<u32>,
 
     /// Fingerprinter of the function's behavior.
@@ -73,6 +75,7 @@ pub struct AnalyzedFunctionExecInfo {
 
 pub struct AnalyzedTransformOp {
     pub name: String,
+    pub op_kind: String,
     pub inputs: Vec<AnalyzedValueMapping>,
     pub function_exec_info: AnalyzedFunctionExecInfo,
     pub executor: Box<dyn SimpleFunctionExecutor>,
diff --git a/src/execution/evaluator.rs b/src/execution/evaluator.rs
@@ -2,6 +2,8 @@ use crate::prelude::*;
 
 use anyhow::{Context, Ok};
 use futures::future::try_join_all;
+use log::warn;
+use tokio::time::Duration;
 
 use crate::base::value::EstimatedByteSize;
 use crate::builder::{AnalyzedTransientFlow, plan::*};
@@ -368,10 +370,13 @@ async fn evaluate_op_scope(
     for reactive_op in op_scope.reactive_ops.iter() {
         match reactive_op {
             AnalyzedReactiveOp::Transform(op) => {
+                let transform_key = format!("transform/{}{}", op_scope.scope_qualifier, op.name);
+
+                //     eprintln!("🔍 DEBUG: Transform op '{}' (function: {}) starting, timeout: {:?}",
+                //   op.name, op.op_kind, op.function_exec_info.timeout);
+
                 // Track transform operation start
                 if let Some(ref op_stats) = operation_in_process_stats {
-                    let transform_key =
-                        format!("transform/{}{}", op_scope.scope_qualifier, op.name);
                     op_stats.start_processing(&transform_key, 1);
                 }
 
@@ -380,6 +385,28 @@ async fn evaluate_op_scope(
                     input_values.push(value?);
                 }
 
+                let timeout_duration = op
+                    .function_exec_info
+                    .timeout
+                    .unwrap_or(Duration::from_secs(300));
+                let warn_duration = Duration::from_secs(30);
+
+                let op_name_for_warning = op.name.clone();
+                let op_kind_for_warning = op.op_kind.clone();
+                let warn_handle = tokio::spawn(async move {
+                    tokio::time::sleep(warn_duration).await;
+                    // eprintln!("WARNING: Function '{}' is taking longer than 30s", op_name_for_warning);
+                    // warn!("Function '{}' is taking longer than 30s", op_name_for_warning);
+                    eprintln!(
+                        "⚠️  WARNING: Function '{}' ({}) is taking longer than 30s",
+                        op_kind_for_warning, op_name_for_warning
+                    ); // ✅ Show both
+                    warn!(
+                        "Function '{}' ({}) is taking longer than 30s",
+                        op_kind_for_warning, op_name_for_warning
+                    );
+                });
+                // Execute with timeout
                 let result = if op.function_exec_info.enable_cache {
                     let output_value_cell = memory.get_cache_entry(
                         || {
@@ -393,27 +420,88 @@ async fn evaluate_op_scope(
                         &op.function_exec_info.output_type,
                         /*ttl=*/ None,
                     )?;
-                    evaluate_with_cell(output_value_cell.as_ref(), move || {
+
+                    let eval_future = evaluate_with_cell(output_value_cell.as_ref(), move || {
                         op.executor.evaluate(input_values)
-                    })
-                    .await
-                    .and_then(|v| head_scope.define_field(&op.output, &v))
+                    });
+
+                    // Handle timeout
+                    let timeout_result = tokio::time::timeout(timeout_duration, eval_future).await;
+                    if timeout_result.is_err() {
+                        Err(anyhow!(
+                            // "Function '{}' timed out after {} seconds",
+                            "Function '{}' ({}) timed out after {} seconds",
+                            op.op_kind,
+                            op.name,
+                            timeout_duration.as_secs()
+                        ))
+                    } else {
+                        timeout_result
+                            .unwrap()
+                            .and_then(|v| head_scope.define_field(&op.output, &v))
+                    }
                 } else {
-                    op.executor
-                        .evaluate(input_values)
-                        .await
-                        .and_then(|v| head_scope.define_field(&op.output, &v))
-                }
-                .with_context(|| format!("Evaluating Transform op `{}`", op.name,));
+                    let eval_future = op.executor.evaluate(input_values);
+
+                    // Handle timeout
+                    let timeout_result = tokio::time::timeout(timeout_duration, eval_future).await;
+                    if timeout_result.is_err() {
+                        Err(anyhow!(
+                            // "Function '{}' timed out after {} seconds",
+                            "Function '{}' ({}) timed out after {} seconds",
+                            op.op_kind,
+                            op.name,
+                            timeout_duration.as_secs()
+                        ))
+                    } else {
+                        timeout_result
+                            .unwrap()
+                            .and_then(|v| head_scope.define_field(&op.output, &v))
+                    }
+                };
+
+                warn_handle.abort();
 
                 // Track transform operation completion
                 if let Some(ref op_stats) = operation_in_process_stats {
-                    let transform_key =
-                        format!("transform/{}{}", op_scope.scope_qualifier, op.name);
                     op_stats.finish_processing(&transform_key, 1);
                 }
 
-                result?
+                result.with_context(|| format!("Evaluating Transform op `{}`", op.name))?
+                // let result = if op.function_exec_info.enable_cache {
+                //     let output_value_cell = memory.get_cache_entry(
+                //         || {
+                //             Ok(op
+                //                 .function_exec_info
+                //                 .fingerprinter
+                //                 .clone()
+                //                 .with(&input_values)?
+                //                 .into_fingerprint())
+                //         },
+                //         &op.function_exec_info.output_type,
+                //         /*ttl=*/ None,
+                //     )?;
+                //     evaluate_with_cell(output_value_cell.as_ref(), move || {
+                //         op.executor.evaluate(input_values)
+                //     })
+                //     .await
+                //     .and_then(|v| head_scope.define_field(&op.output, &v))
+                // } else {
+                //     op.executor
+                //         .evaluate(input_values)
+                //         .await
+                //         .and_then(|v| head_scope.define_field(&op.output, &v))
+                // }
+                // .with_context(|| format!("Evaluating Transform op `{}`", op.name,));
+
+                // // Track transform operation completion
+                // if let Some(ref op_stats) = operation_in_process_stats {
+                //     let transform_key =
+                //         format!("transform/{}{}", op_scope.scope_qualifier, op.name);
+                //     op_stats.finish_processing(&transform_key, 1);
+                // }
+
+                // result?
             }
 
             AnalyzedReactiveOp::ForEach(op) => {
diff --git a/src/execution/live_updater.rs b/src/execution/live_updater.rs
@@ -172,6 +172,7 @@ impl SourceUpdateTask {
                 let mut change_stream = change_stream;
                 let retry_options = retryable::RetryOptions {
                     retry_timeout: None,
+                    per_call_timeout: None,
                     initial_backoff: std::time::Duration::from_secs(5),
                     max_backoff: std::time::Duration::from_secs(60),
                 };
diff --git a/src/ops/interface.rs b/src/ops/interface.rs
@@ -185,6 +185,11 @@ pub trait SimpleFunctionExecutor: Send + Sync {
     fn behavior_version(&self) -> Option<u32> {
         None
     }
+
+    /// Returns None to use the default timeout (300s)
+    fn timeout(&self) -> Option<std::time::Duration> {
+        None
+    }
 }
 
 #[async_trait]
diff --git a/src/ops/py_factory.rs b/src/ops/py_factory.rs
@@ -43,6 +43,7 @@ struct PyFunctionExecutor {
 
     enable_cache: bool,
     behavior_version: Option<u32>,
+    timeout: Option<std::time::Duration>,
 }
 
 impl PyFunctionExecutor {
@@ -112,6 +113,10 @@ impl interface::SimpleFunctionExecutor for Arc<PyFunctionExecutor> {
     fn behavior_version(&self) -> Option<u32> {
         self.behavior_version
     }
+
+    fn timeout(&self) -> Option<std::time::Duration> {
+        self.timeout
+    }
 }
 
 struct PyBatchedFunctionExecutor {
@@ -121,6 +126,7 @@ struct PyBatchedFunctionExecutor {
 
     enable_cache: bool,
     behavior_version: Option<u32>,
+    timeout: Option<std::time::Duration>,
 }
 
 #[async_trait]
@@ -237,7 +243,7 @@ impl interface::SimpleFunctionFactory for PyFunctionFactory {
                     .as_ref()
                     .ok_or_else(|| anyhow!("Python execution context is missing"))?
                     .clone();
-                let (prepare_fut, enable_cache, behavior_version) =
+                let (prepare_fut, enable_cache, behavior_version, timeout) =
                     Python::with_gil(|py| -> anyhow::Result<_> {
                         let prepare_coro = executor
                             .call_method(py, "prepare", (), None)
@@ -257,7 +263,12 @@ impl interface::SimpleFunctionFactory for PyFunctionFactory {
                             .call_method(py, "behavior_version", (), None)
                             .to_result_with_py_trace(py)?
                             .extract::<Option<u32>>(py)?;
-                        Ok((prepare_fut, enable_cache, behavior_version))
+                        let timeout = executor
+                            .call_method(py, "timeout", (), None)
+                            .to_result_with_py_trace(py)?
+                            .extract::<Option<u64>>(py)?
+                            .map(std::time::Duration::from_secs);
+                        Ok((prepare_fut, enable_cache, behavior_version, timeout))
                     })?;
                 prepare_fut.await?;
                 let executor: Box<dyn interface::SimpleFunctionExecutor> = if self.batching {
@@ -268,6 +279,7 @@ impl interface::SimpleFunctionFactory for PyFunctionFactory {
                             result_type,
                             enable_cache,
                             behavior_version,
+                            timeout,
                         }
                         .into_fn_executor(),
                     )
@@ -280,6 +292,7 @@ impl interface::SimpleFunctionFactory for PyFunctionFactory {
                         result_type,
                         enable_cache,
                         behavior_version,
+                        timeout,
                     }))
                 };
                 Ok(executor)
diff --git a/src/utils/retryable.rs b/src/utils/retryable.rs

Original file line number	Diff line number	Diff line change
`@@ -185,6 +185,11 @@ pub trait SimpleFunctionExecutor: Send + Sync {`
`185`	`185`	`fn behavior_version(&self) -> Option<u32> {`
`186`	`186`	`None`
`187`	`187`	`}`
	`188`	`+`
	`189`	`+ /// Returns None to use the default timeout (300s)`
	`190`	`+ fn timeout(&self) -> Option<std::time::Duration> {`
	`191`	`+ None`
	`192`	`+ }`
`188`	`193`	`}`
`189`	`194`
`190`	`195`	`#[async_trait]`