fix(s3select): enforce query and resource limits (#5028)

* fix(s3select): enforce query and resource limits

* fix(s3select): close query resource limit gaps

* fix(s3select): preserve timeout and stream invariants

* fix(s3select): enforce staged query limits

* fix(s3select): preserve policy error compatibility

* fix(s3select): bound error source traversal

crates/s3select-api/src/query/execution.rs

@@ -30,7 +30,7 @@ use crate::{QueryError, QueryResult};
 
 use super::Query;
 use super::logical_planner::Plan;
-use super::session::SessionCtx;
+use super::session::{QueryExecutionTracker, SessionCtx};
 
 pub struct PhaseTimer {
     phase_name: &'static str,
@@ -172,6 +172,7 @@ pub struct QueryStateMachine {
     pub session: SessionCtx,
     pub query: Query,
 
+    query_tracker: Option<QueryExecutionTracker>,
     state: RwLock<QueryState>,
     start: Instant,
 }
@@ -195,11 +196,31 @@ impl QueryStateMachine {
         Self {
             session,
             query,
+            query_tracker: None,
             state: RwLock::new(QueryState::ACCEPTING),
             start: Instant::now(),
         }
     }
 
+    pub fn begin_tracked(query: Query, session: SessionCtx, query_tracker: QueryExecutionTracker) -> QueryResult<Self> {
+        if !session.is_bound_to(&query_tracker) {
+            return Err(QueryError::Cancel);
+        }
+        let mut state_machine = Self::begin(query, session);
+        state_machine.query_tracker = Some(query_tracker);
+        Ok(state_machine)
+    }
+
+    pub fn query_tracker(&self) -> Option<&QueryExecutionTracker> {
+        self.query_tracker.as_ref()
+    }
+
+    pub fn tracker_matches_session(&self) -> bool {
+        self.query_tracker
+            .as_ref()
+            .is_some_and(|query_tracker| self.session.is_bound_to(query_tracker))
+    }
+
     pub fn begin_analyze(&self) {
         self.record_phase_timestamp("analyze", "start");
         self.translate_to(QueryState::RUNNING(RUNNING::ANALYZING));

crates/s3select-api/src/query/session.rs

@@ -13,7 +13,7 @@
 // limitations under the License.
 
 use crate::query::Context;
-use crate::{QueryError, QueryResult, object_store::EcObjectStore};
+use crate::{QueryError, QueryResult, SelectStore, object_store::EcObjectStore};
 use datafusion::{
     arrow::{
         array::{Int32Array, StringArray},
@@ -25,19 +25,237 @@ use datafusion::{
     parquet::arrow::ArrowWriter,
     prelude::SessionContext,
 };
-use std::sync::Arc;
+use parking_lot::Mutex;
+use std::sync::{
+    Arc, Weak,
+    atomic::{AtomicU8, Ordering},
+};
+use tokio::{
+    sync::OwnedSemaphorePermit,
+    task::AbortHandle,
+    time::{Instant, sleep_until},
+};
 use tracing::error;
 
+pub type QueryExecutionGuard = Arc<OwnedSemaphorePermit>;
+
+#[derive(Clone, Default)]
+pub struct QueryExecutionOwner {
+    identity: Arc<()>,
+}
+
+impl QueryExecutionOwner {
+    pub fn new() -> Self {
+        Self { identity: Arc::new(()) }
+    }
+}
+
+#[derive(Clone)]
+pub struct QueryExecutionTracker {
+    inner: Arc<QueryExecutionTrackerInner>,
+}
+
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub enum QueryExecutionStatus {
+    Active,
+    Finished,
+    TimedOut,
+}
+
+const EXECUTION_SETTING_UP: u8 = 0;
+const EXECUTION_ADMITTED: u8 = 1;
+const EXECUTION_PLANNING: u8 = 2;
+const EXECUTION_PLANNED: u8 = 3;
+const EXECUTION_STARTING: u8 = 4;
+const EXECUTION_RUNNING: u8 = 5;
+const EXECUTION_FINISHED: u8 = 6;
+const EXECUTION_TIMED_OUT: u8 = 7;
+
+struct QueryExecutionTrackerInner {
+    owner_identity: Arc<()>,
+    query_guard: Mutex<Option<QueryExecutionGuard>>,
+    deadline: Instant,
+    timeout_seconds: u64,
+    state: AtomicU8,
+    deadline_task: Mutex<Option<AbortHandle>>,
+}
+
+impl QueryExecutionTracker {
+    pub fn new(owner: &QueryExecutionOwner, query_guard: QueryExecutionGuard, deadline: Instant, timeout_seconds: u64) -> Self {
+        let inner = Arc::new(QueryExecutionTrackerInner {
+            owner_identity: Arc::clone(&owner.identity),
+            query_guard: Mutex::new(Some(query_guard)),
+            deadline,
+            timeout_seconds,
+            state: AtomicU8::new(EXECUTION_SETTING_UP),
+            deadline_task: Mutex::new(None),
+        });
+        let deadline_inner = Arc::downgrade(&inner);
+        let deadline_task = tokio::spawn(async move {
+            sleep_until(deadline).await;
+            if let Some(inner) = Weak::upgrade(&deadline_inner) {
+                inner.expire_at_deadline();
+            }
+        });
+        *inner.deadline_task.lock() = Some(deadline_task.abort_handle());
+
+        Self { inner }
+    }
+
+    pub fn deadline(&self) -> Instant {
+        self.inner.deadline
+    }
+
+    pub fn timeout_seconds(&self) -> u64 {
+        self.inner.timeout_seconds
+    }
+
+    pub fn status(&self) -> QueryExecutionStatus {
+        match self.inner.state.load(Ordering::Acquire) {
+            EXECUTION_TIMED_OUT => QueryExecutionStatus::TimedOut,
+            state if state < EXECUTION_FINISHED => QueryExecutionStatus::Active,
+            _ => QueryExecutionStatus::Finished,
+        }
+    }
+
+    pub fn is_owned_by(&self, owner: &QueryExecutionOwner) -> bool {
+        Arc::ptr_eq(&self.inner.owner_identity, &owner.identity)
+    }
+
+    pub(crate) fn is_same_execution(&self, other: &Self) -> bool {
+        Arc::ptr_eq(&self.inner, &other.inner)
+    }
+
+    pub fn mark_admitted(&self, owner: &QueryExecutionOwner) -> bool {
+        self.transition(owner, EXECUTION_SETTING_UP, EXECUTION_ADMITTED)
+    }
+
+    pub fn claim_planning(&self, owner: &QueryExecutionOwner) -> bool {
+        self.transition(owner, EXECUTION_ADMITTED, EXECUTION_PLANNING)
+    }
+
+    pub fn mark_planned(&self, owner: &QueryExecutionOwner) -> bool {
+        self.transition(owner, EXECUTION_PLANNING, EXECUTION_PLANNED)
+    }
+
+    pub fn claim_execution(&self, owner: &QueryExecutionOwner) -> bool {
+        self.transition(owner, EXECUTION_PLANNED, EXECUTION_STARTING)
+    }
+
+    pub fn mark_running(&self, owner: &QueryExecutionOwner) -> bool {
+        self.transition(owner, EXECUTION_STARTING, EXECUTION_RUNNING)
+    }
+
+    pub fn handoff_deadline(&self, owner: &QueryExecutionOwner) {
+        if !self.is_owned_by(owner) || self.inner.state.load(Ordering::Acquire) != EXECUTION_RUNNING {
+            return;
+        }
+        if let Some(deadline_task) = self.inner.deadline_task.lock().take() {
+            deadline_task.abort();
+        }
+    }
+
+    pub fn finish(&self, owner: &QueryExecutionOwner) {
+        if self.is_owned_by(owner) {
+            self.inner.finish();
+        }
+    }
+
+    pub fn expire(&self, owner: &QueryExecutionOwner) {
+        if self.is_owned_by(owner) {
+            self.inner.expire();
+        }
+    }
+
+    pub(crate) fn query_guard(&self) -> Option<QueryExecutionGuard> {
+        let query_guard = self.inner.query_guard.lock();
+        if Instant::now() >= self.inner.deadline || self.status() != QueryExecutionStatus::Active {
+            return None;
+        }
+        query_guard.clone()
+    }
+
+    fn transition(&self, owner: &QueryExecutionOwner, from: u8, to: u8) -> bool {
+        self.is_owned_by(owner)
+            && self
+                .inner
+                .state
+                .compare_exchange(from, to, Ordering::AcqRel, Ordering::Acquire)
+                .is_ok()
+    }
+}
+
+impl QueryExecutionTrackerInner {
+    fn finish(&self) {
+        self.state.fetch_max(EXECUTION_FINISHED, Ordering::AcqRel);
+        self.release();
+    }
+
+    fn expire(&self) {
+        self.mark_timed_out();
+        self.release();
+    }
+
+    fn expire_at_deadline(&self) {
+        if self
+            .mark_timed_out()
+            .is_some_and(|state| matches!(state, EXECUTION_ADMITTED | EXECUTION_PLANNED))
+        {
+            self.release();
+        }
+    }
+
+    fn mark_timed_out(&self) -> Option<u8> {
+        self.state
+            .fetch_update(Ordering::AcqRel, Ordering::Acquire, |state| {
+                (state < EXECUTION_FINISHED).then_some(EXECUTION_TIMED_OUT)
+            })
+            .ok()
+    }
+
+    fn release(&self) {
+        self.query_guard.lock().take();
+        if let Some(deadline_task) = self.deadline_task.lock().take() {
+            deadline_task.abort();
+        }
+    }
+}
+
+impl Drop for QueryExecutionTrackerInner {
+    fn drop(&mut self) {
+        if let Some(deadline_task) = self.deadline_task.get_mut().take() {
+            deadline_task.abort();
+        }
+    }
+}
+
+impl std::fmt::Debug for QueryExecutionTracker {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("QueryExecutionTracker")
+            .field("deadline", &self.deadline())
+            .field("timeout_seconds", &self.timeout_seconds())
+            .field("status", &self.status())
+            .finish_non_exhaustive()
+    }
+}
+
 #[derive(Clone)]
 pub struct SessionCtx {
     _desc: Arc<SessionCtxDesc>,
     inner: SessionState,
+    query_tracker: Option<QueryExecutionTracker>,
 }
 
 impl SessionCtx {
     pub fn inner(&self) -> &SessionState {
         &self.inner
     }
+
+    pub(crate) fn is_bound_to(&self, query_tracker: &QueryExecutionTracker) -> bool {
+        self.query_tracker
+            .as_ref()
+            .is_some_and(|bound_tracker| bound_tracker.is_same_execution(query_tracker))
+    }
 }
 
 #[derive(Clone)]
@@ -45,12 +263,19 @@ pub struct SessionCtxDesc {
     // maybe we need some info
 }
 
-#[derive(Default)]
 pub struct SessionCtxFactory {
     pub is_test: bool,
     pub target_partitions: usize,
 }
 
+pub const DEFAULT_S3SELECT_MEMORY_LIMIT_BYTES: usize = 64 * 1024 * 1024;
+
+impl Default for SessionCtxFactory {
+    fn default() -> Self {
+        Self::new(false)
+    }
+}
+
 impl SessionCtxFactory {
     pub fn new(is_test: bool) -> Self {
         Self {
@@ -65,19 +290,66 @@ impl SessionCtxFactory {
     }
 
     pub async fn create_session_ctx(&self, context: &Context) -> QueryResult<SessionCtx> {
-        let df_session_ctx = self.build_df_session_context(context).await?;
+        self.create_session_ctx_inner(context, None, None, DEFAULT_S3SELECT_MEMORY_LIMIT_BYTES)
+            .await
+    }
+
+    pub async fn create_session_ctx_with_tracker_and_memory_limit(
+        &self,
+        context: &Context,
+        query_tracker: QueryExecutionTracker,
+        memory_limit_bytes: usize,
+    ) -> QueryResult<SessionCtx> {
+        self.create_session_ctx_inner(context, Some(query_tracker), None, memory_limit_bytes)
+            .await
+    }
+
+    #[cfg(test)]
+    async fn create_session_ctx_with_tracker_and_store(
+        &self,
+        context: &Context,
+        query_tracker: QueryExecutionTracker,
+        store: Arc<SelectStore>,
+    ) -> QueryResult<SessionCtx> {
+        self.create_session_ctx_inner(context, Some(query_tracker), Some(store), DEFAULT_S3SELECT_MEMORY_LIMIT_BYTES)
+            .await
+    }
+
+    async fn create_session_ctx_inner(
+        &self,
+        context: &Context,
+        query_tracker: Option<QueryExecutionTracker>,
+        store: Option<Arc<SelectStore>>,
+        memory_limit_bytes: usize,
+    ) -> QueryResult<SessionCtx> {
+        let df_session_ctx = self
+            .build_df_session_context(context, query_tracker.clone(), store, memory_limit_bytes)
+            .await?;
 
         Ok(SessionCtx {
             _desc: Arc::new(SessionCtxDesc {}),
             inner: df_session_ctx.state(),
+            query_tracker,
         })
     }
 
-    async fn build_df_session_context(&self, context: &Context) -> QueryResult<SessionContext> {
+    async fn build_df_session_context(
+        &self,
+        context: &Context,
+        query_tracker: Option<QueryExecutionTracker>,
+        store: Option<Arc<SelectStore>>,
+        memory_limit_bytes: usize,
+    ) -> QueryResult<SessionContext> {
         let path = format!("s3://{}", context.input.bucket);
         let store_url = url::Url::parse(&path).unwrap();
-        let rt = RuntimeEnvBuilder::new().build()?;
+        let memory_limit_bytes = if memory_limit_bytes == 0 {
+            DEFAULT_S3SELECT_MEMORY_LIMIT_BYTES
+        } else {
+            memory_limit_bytes
+        };
+        let rt = RuntimeEnvBuilder::new().with_memory_limit(memory_limit_bytes, 1.0).build()?;
         let config = SessionConfig::new().with_target_partitions(self.target_partitions);
+        let memory_pool = Arc::clone(&rt.memory_pool);
         let df_session_state = SessionStateBuilder::new()
             .with_config(config)
             .with_runtime_env(Arc::new(rt))
@@ -127,8 +399,13 @@ impl SessionCtxFactory {
 
             df_session_state.with_object_store(&store_url, store).build()
         } else {
-            let store: EcObjectStore =
-                EcObjectStore::new(context.input.clone()).map_err(|_| QueryError::NotImplemented { err: String::new() })?;
+            let store: EcObjectStore = match query_tracker {
+                Some(query_tracker) => {
+                    EcObjectStore::new_with_query_tracker(context.input.clone(), memory_pool, query_tracker, store)
+                }
+                None => EcObjectStore::new_with_memory_pool(context.input.clone(), memory_pool),
+            }
+            .map_err(|_| QueryError::NotImplemented { err: String::new() })?;
             df_session_state.with_object_store(&store_url, Arc::new(store)).build()
         };
 
@@ -197,6 +474,7 @@ fn test_parquet_batch(
 #[cfg(test)]
 mod tests {
     use super::*;
+    use datafusion::execution::memory_pool::MemoryLimit;
     use s3s::dto::{
         CSVInput, CSVOutput, ExpressionType, InputSerialization, OutputSerialization, SelectObjectContentInput,
         SelectObjectContentRequest,
@@ -229,6 +507,17 @@ mod tests {
         }
     }
 
+    #[test]
+    fn session_factory_fields_remain_source_compatible() {
+        let factory = SessionCtxFactory {
+            is_test: true,
+            target_partitions: 0,
+        };
+
+        assert!(factory.is_test);
+        assert_eq!(factory.target_partitions, 0);
+    }
+
     #[tokio::test]
     async fn session_factory_applies_target_partitions() {
         let factory = SessionCtxFactory::new(true).with_target_partitions(3);
@@ -250,4 +539,78 @@ mod tests {
 
         assert_eq!(session.inner().config().target_partitions(), SessionConfig::new().target_partitions());
     }
+
+    #[tokio::test]
+    async fn session_factory_applies_memory_limit() {
+        let factory = SessionCtxFactory::new(true);
+        let session = factory
+            .create_session_ctx_inner(&test_context(), None, None, 1024)
+            .await
+            .expect("session should be created with a bounded memory pool");
+
+        assert!(matches!(
+            session.inner().runtime_env().memory_pool.memory_limit(),
+            MemoryLimit::Finite(1024)
+        ));
+    }
+
+    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+    #[serial_test::serial]
+    async fn session_factory_propagates_query_guard_to_ec_store() {
+        let temp_root = tempfile::tempdir().expect("create session test temp root");
+        let env = rustfs_test_utils::TestECStoreEnv::builder()
+            .base_dir(temp_root.path())
+            .init_bucket_metadata(false)
+            .build()
+            .await;
+
+        let admission = Arc::new(tokio::sync::Semaphore::new(1));
+        let permit = Arc::clone(&admission)
+            .acquire_owned()
+            .await
+            .expect("query permit should be available");
+        let query_guard = Arc::new(permit);
+        let query_tracker = QueryExecutionTracker::new(
+            &QueryExecutionOwner::new(),
+            Arc::clone(&query_guard),
+            Instant::now() + std::time::Duration::from_secs(300),
+            300,
+        );
+        let session = SessionCtxFactory::new(false)
+            .create_session_ctx_with_tracker_and_store(&test_context(), query_tracker, Arc::clone(&env.ecstore))
+            .await
+            .expect("production session should be created with the query guard");
+
+        assert!(Arc::strong_count(&query_guard) > 1);
+        drop(session);
+        assert_eq!(Arc::strong_count(&query_guard), 1);
+    }
+
+    #[tokio::test]
+    async fn elapsed_deadline_does_not_yield_query_guard_before_timer_poll() {
+        let admission = Arc::new(tokio::sync::Semaphore::new(1));
+        let permit = Arc::clone(&admission)
+            .acquire_owned()
+            .await
+            .expect("query permit should be available");
+        let query_tracker = QueryExecutionTracker::new(&QueryExecutionOwner::new(), Arc::new(permit), Instant::now(), 0);
+
+        assert!(query_tracker.query_guard().is_none());
+    }
+
+    #[tokio::test]
+    async fn session_factory_default_uses_bounded_memory() {
+        let factory = SessionCtxFactory::default();
+        let session = SessionCtxFactory::new(true)
+            .create_session_ctx(&test_context())
+            .await
+            .expect("default session should be created with a bounded memory pool");
+
+        assert!(!factory.is_test);
+        assert_eq!(factory.target_partitions, 0);
+        assert!(matches!(
+            session.inner().runtime_env().memory_pool.memory_limit(),
+            MemoryLimit::Finite(DEFAULT_S3SELECT_MEMORY_LIMIT_BYTES)
+        ));
+    }
 }