surrealdb-core 3.3.1

//! Timeout operator for query time limits.
//!
//! Wraps an input operator stream and enforces a maximum execution
//! duration. If the timeout expires before the input completes, a
//! `QueryTimedout` error is returned.

use std::sync::Arc;

use common::future::stream::{self, Yielder};
use futures::StreamExt;

use crate::err::EngineError;
use crate::exec::{
	AccessMode, CardinalityHint, ContextLevel, ExecOperator, ExecutionContext, FlowResult,
	OperatorMetrics, PhysicalExpr, ValueBatchStream, buffer_stream, monitor_stream,
};
use crate::expr::{ControlFlow, ControlFlowExt};
use crate::val::Duration;

/// Applies a timeout to the execution of its input operator.
///
/// If the timeout expires before the input stream completes, an error is returned.
/// This is typically applied as the outermost operator for a query.
#[derive(Debug, Clone)]
pub struct Timeout {
	pub(crate) input: Arc<dyn ExecOperator>,
	/// The timeout duration. If None, no timeout is applied.
	pub(crate) timeout: Option<Arc<dyn PhysicalExpr>>,
	pub(crate) metrics: Arc<OperatorMetrics>,
}

impl Timeout {
	pub(crate) fn new(
		input: Arc<dyn ExecOperator>,
		timeout: Option<Arc<dyn PhysicalExpr>>,
	) -> Self {
		Self {
			input,
			timeout,
			metrics: Arc::new(OperatorMetrics::new()),
		}
	}
}
impl ExecOperator for Timeout {
	fn name(&self) -> &'static str {
		"Timeout"
	}

	fn attrs(&self) -> Vec<(String, String)> {
		if let Some(timeout) = &self.timeout {
			vec![("duration".to_string(), timeout.to_sql())]
		} else {
			vec![]
		}
	}

	fn required_context(&self) -> ContextLevel {
		// Combine timeout expression context with child operator context
		let timeout_ctx =
			self.timeout.as_ref().map(|e| e.required_context()).unwrap_or(ContextLevel::Root);
		timeout_ctx.max(self.input.required_context())
	}

	fn access_mode(&self) -> AccessMode {
		// Timeout is transparent to access mode
		self.input.access_mode()
	}

	fn cardinality_hint(&self) -> CardinalityHint {
		self.input.cardinality_hint()
	}

	fn children(&self) -> Vec<&Arc<dyn ExecOperator>> {
		vec![&self.input]
	}

	fn metrics(&self) -> Option<&OperatorMetrics> {
		Some(&self.metrics)
	}

	fn expressions(&self) -> Vec<(&str, &Arc<dyn PhysicalExpr>)> {
		if let Some(timeout) = &self.timeout {
			vec![("timeout", timeout)]
		} else {
			vec![]
		}
	}

	fn output_ordering(&self) -> crate::exec::OutputOrdering {
		self.input.output_ordering()
	}

	fn execute(&self, ctx: &ExecutionContext) -> FlowResult<ValueBatchStream> {
		let mut input_stream = buffer_stream(
			self.input.execute(ctx)?,
			self.input.access_mode(),
			self.input.cardinality_hint(),
			ctx.root().ctx.config.exec.operator_buffer_size,
		);

		// If no timeout is specified, just pass through the input stream
		let Some(timeout_expr) = &self.timeout else {
			return Ok(monitor_stream(input_stream, "Timeout", &self.metrics));
		};

		// Evaluate the timeout expression to get the duration
		let timeout_expr = Arc::clone(timeout_expr);
		let ctx = ctx.clone();

		let stream = stream::try_async_stream(async move |mut yielder: Yielder<_>| {
			use crate::exec::EvalContext;

			// Evaluate the timeout expression (no row context needed)
			let eval_ctx = EvalContext::from_exec_ctx(&ctx);
			let timeout_value = timeout_expr.evaluate(eval_ctx).await?;

			// Convert to duration
			let duration: Duration =
				timeout_value.coerce_to::<Duration>().context("Invalid timeout value")?;

			// Convert our Duration to std::time::Duration for tokio
			let std_duration: std::time::Duration = duration.0;

			// Create a timeout future
			let timeout_instant = common::time::Instant::now() + std_duration;

			loop {
				// Check if we've exceeded the timeout
				let remaining =
					timeout_instant.saturating_duration_since(common::time::Instant::now());
				if remaining.is_zero() {
					Err(ControlFlow::Err(anyhow::anyhow!(EngineError::QueryTimedout(duration.0))))?;
				}

				// Wait for next batch with timeout
				let batch_result = common::time::timeout(remaining, input_stream.next()).await;

				match batch_result {
					Ok(Some(batch)) => {
						yielder.emit(batch?).await;
					}
					Ok(None) => {
						// Stream completed normally
						break;
					}
					Err(_) => {
						// Timeout expired
						Err(ControlFlow::Err(anyhow::anyhow!(EngineError::QueryTimedout(
							duration.0
						))))?;
					}
				}
			}
			Ok(())
		});

		Ok(monitor_stream(Box::pin(stream), "Timeout", &self.metrics))
	}
}