fix(select): pin object snapshot for query lifetime (#5835)

crates/s3select-api/Cargo.toml

@@ -77,11 +77,12 @@ parking_lot.workspace = true
 tokio = { workspace = true, features = ["fs", "rt-multi-thread"] }
 tokio-util = { workspace = true, features = ["io", "compat"] }
 tracing.workspace = true
+uuid.workspace = true
 transform-stream.workspace = true
 url.workspace = true
 
 [dev-dependencies]
-rustfs-test-utils.workspace = true
+rustfs-test-utils = { workspace = true, features = ["put-object-commit-barrier"] }
 serial_test.workspace = true
 
 [lib]

crates/s3select-api/src/lib.rs

@@ -15,22 +15,23 @@
 #![recursion_limit = "256"]
 
 use datafusion::{common::DataFusionError, sql::sqlparser::parser::ParserError};
-use std::fmt::Display;
+use std::{error::Error as StdError, fmt::Display};
 use thiserror::Error;
 
 pub mod object_store;
 pub mod query;
 pub mod server;
 mod storage_api;
+pub use storage_api::SelectObjectSnapshot;
 
 #[cfg(test)]
 mod test;
 
 pub type QueryResult<T> = Result<T, QueryError>;
 pub(crate) use storage_api::crate_boundary::{
-    SELECT_DEFAULT_READ_BUFFER_SIZE, SelectGetObjectReader, SelectObjectInfo, SelectObjectOptions, SelectStorageError,
-    SelectStore, resolve_select_object_store_handle, select_is_err_bucket_not_found, select_is_err_object_not_found,
-    select_is_err_version_not_found,
+    PrepareSelectObjectSnapshotError, SELECT_DEFAULT_READ_BUFFER_SIZE, SelectGetObjectReader, SelectObjectOptions,
+    SelectObjectSnapshotReadError, SelectStorageError, SelectStore, SnapshotConsistencyError, resolve_select_object_store_handle,
+    select_is_err_bucket_not_found, select_is_err_object_not_found, select_is_err_version_not_found,
 };
 
 #[derive(Debug, Error)]
@@ -79,24 +80,24 @@ pub enum S3SelectPolicyError {
     QueryTimeout { seconds: u64 },
 }
 
-impl S3SelectPolicyError {
-    fn from_error<'a>(mut err: &'a (dyn std::error::Error + 'static)) -> Option<&'a Self> {
+impl QueryError {
+    fn source_error<T: StdError + 'static>(&self) -> Option<&T> {
+        let mut err: &(dyn StdError + 'static) = self;
         for _ in 0..16 {
-            if let Some(policy_error) = err.downcast_ref::<Self>() {
-                return Some(policy_error);
+            if let Some(source) = err.downcast_ref::<T>() {
+                return Some(source);
             }
             err = err.source()?;
         }
         None
     }
-}
 
-impl QueryError {
+    pub fn is_snapshot_consistency_error(&self) -> bool {
+        self.source_error::<SnapshotConsistencyError>().is_some()
+    }
+
     pub fn s3_select_policy_error(&self) -> Option<&S3SelectPolicyError> {
-        match self {
-            Self::Datafusion { source } => S3SelectPolicyError::from_error(source.as_ref()),
-            _ => None,
-        }
+        self.source_error()
     }
 }
 
@@ -230,6 +231,17 @@ mod tests {
         ));
     }
 
+    #[test]
+    fn snapshot_consistency_error_is_recoverable_without_string_matching() {
+        let err = QueryError::Datafusion {
+            source: Box::new(DataFusionError::External(Box::new(SelectObjectSnapshotReadError::Consistency(
+                SnapshotConsistencyError::LockLost,
+            )))),
+        };
+
+        assert!(err.is_snapshot_consistency_error());
+    }
+
     #[test]
     fn test_query_error_from_parser_error() {
         let parser_error = ParserError::ParserError("syntax error".to_string());

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

@@ -22,6 +22,7 @@ use super::{
     Query,
     execution::{Output, QueryStateMachine},
     logical_planner::Plan,
+    session::QueryAdmission,
 };
 
 #[async_trait]
@@ -32,6 +33,14 @@ pub trait QueryDispatcher: Send + Sync {
 
     async fn execute_query(&self, query: &Query) -> QueryResult<Output>;
 
+    fn try_reserve_query(&self) -> QueryResult<QueryAdmission> {
+        Ok(QueryAdmission::unmanaged())
+    }
+
+    async fn execute_query_admitted(&self, query: &Query, _admission: QueryAdmission) -> QueryResult<Output> {
+        self.execute_query(query).await
+    }
+
     async fn build_logical_plan(&self, query_state_machine: Arc<QueryStateMachine>) -> QueryResult<Option<Plan>>;
 
     async fn execute_logical_plan(&self, logical_plan: Plan, query_state_machine: Arc<QueryStateMachine>) -> QueryResult<Output>;

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

@@ -15,6 +15,8 @@
 use s3s::dto::SelectObjectContentInput;
 use std::sync::Arc;
 
+use crate::SelectObjectSnapshot;
+
 pub mod analyzer;
 pub mod ast;
 pub mod dispatcher;
@@ -37,12 +39,26 @@ pub struct Context {
 pub struct Query {
     context: Context,
     content: String,
+    snapshot: Option<Arc<SelectObjectSnapshot>>,
 }
 
 impl Query {
     #[inline(always)]
     pub fn new(context: Context, content: String) -> Self {
-        Self { context, content }
+        Self {
+            context,
+            content,
+            snapshot: None,
+        }
+    }
+
+    #[inline(always)]
+    pub fn new_with_snapshot(context: Context, content: String, snapshot: Arc<SelectObjectSnapshot>) -> Self {
+        Self {
+            context,
+            content,
+            snapshot: Some(snapshot),
+        }
     }
 
     pub fn context(&self) -> &Context {
@@ -52,4 +68,8 @@ impl Query {
     pub fn content(&self) -> &str {
         self.content.as_str()
     }
+
+    pub fn snapshot(&self) -> Option<&Arc<SelectObjectSnapshot>> {
+        self.snapshot.as_ref()
+    }
 }

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

@@ -12,14 +12,16 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use crate::SelectObjectSnapshot;
 use crate::query::Context;
-use crate::{QueryError, QueryResult, SelectStore, object_store::EcObjectStore};
+use crate::{QueryError, QueryResult, object_store::EcObjectStore};
 use datafusion::{
     arrow::{
         array::{Int32Array, StringArray},
         datatypes::{DataType, Field, Schema},
         record_batch::RecordBatch,
     },
+    common::DataFusionError,
     execution::{SessionStateBuilder, config::SessionConfig, context::SessionState, runtime_env::RuntimeEnvBuilder},
     object_store::{ObjectStore, ObjectStoreExt, memory::InMemory, path::Path},
     parquet::arrow::ArrowWriter,
@@ -39,6 +41,28 @@ use tracing::error;
 
 pub type QueryExecutionGuard = Arc<OwnedSemaphorePermit>;
 
+/// A one-shot query admission reservation handed from the request boundary to
+/// the dispatcher that owns the corresponding concurrency semaphore.
+pub struct QueryAdmission {
+    query_guard: Option<QueryExecutionGuard>,
+}
+
+impl QueryAdmission {
+    pub fn new(query_guard: QueryExecutionGuard) -> Self {
+        Self {
+            query_guard: Some(query_guard),
+        }
+    }
+
+    pub fn into_query_guard(mut self) -> Option<QueryExecutionGuard> {
+        self.query_guard.take()
+    }
+
+    pub(crate) fn unmanaged() -> Self {
+        Self { query_guard: None }
+    }
+}
+
 #[derive(Clone, Default)]
 pub struct QueryExecutionOwner {
     identity: Arc<()>,
@@ -300,30 +324,30 @@ impl SessionCtxFactory {
         query_tracker: QueryExecutionTracker,
         memory_limit_bytes: usize,
     ) -> QueryResult<SessionCtx> {
-        self.create_session_ctx_inner(context, Some(query_tracker), None, memory_limit_bytes)
+        self.create_session_ctx_inner(context, None, Some(query_tracker), memory_limit_bytes)
             .await
     }
 
-    #[cfg(test)]
-    async fn create_session_ctx_with_tracker_and_store(
+    pub async fn create_session_ctx_with_snapshot_and_tracker_and_memory_limit(
         &self,
         context: &Context,
+        snapshot: Arc<SelectObjectSnapshot>,
         query_tracker: QueryExecutionTracker,
-        store: Arc<SelectStore>,
+        memory_limit_bytes: usize,
     ) -> QueryResult<SessionCtx> {
-        self.create_session_ctx_inner(context, Some(query_tracker), Some(store), DEFAULT_S3SELECT_MEMORY_LIMIT_BYTES)
+        self.create_session_ctx_inner(context, Some(snapshot), Some(query_tracker), memory_limit_bytes)
             .await
     }
 
     async fn create_session_ctx_inner(
         &self,
         context: &Context,
+        snapshot: Option<Arc<SelectObjectSnapshot>>,
         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)
+            .build_df_session_context(context, snapshot, query_tracker.clone(), memory_limit_bytes)
             .await?;
 
         Ok(SessionCtx {
@@ -336,8 +360,8 @@ impl SessionCtxFactory {
     async fn build_df_session_context(
         &self,
         context: &Context,
+        snapshot: Option<Arc<SelectObjectSnapshot>>,
         query_tracker: Option<QueryExecutionTracker>,
-        store: Option<Arc<SelectStore>>,
         memory_limit_bytes: usize,
     ) -> QueryResult<SessionContext> {
         let path = format!("s3://{}", context.input.bucket);
@@ -416,11 +440,13 @@ impl SessionCtxFactory {
         } else {
             let store: EcObjectStore = match query_tracker {
                 Some(query_tracker) => {
-                    EcObjectStore::new_with_query_tracker(context.input.clone(), memory_pool, query_tracker, store)
+                    EcObjectStore::new_with_query_tracker(context.input.clone(), memory_pool, query_tracker, snapshot)
                 }
-                None => EcObjectStore::new_with_memory_pool(context.input.clone(), memory_pool),
+                None => EcObjectStore::new_with_memory_pool(context.input.clone(), memory_pool, snapshot),
             }
-            .map_err(|_| QueryError::NotImplemented { err: String::new() })?;
+            .map_err(|err| QueryError::Datafusion {
+                source: Box::new(DataFusionError::External(Box::new(err))),
+            })?;
             df_session_state.with_object_store(&store_url, Arc::new(store)).build()
         };
 
@@ -498,6 +524,7 @@ mod tests {
         },
         execution::memory_pool::MemoryLimit,
     };
+    use http::HeaderMap;
     use s3s::dto::{
         CSVInput, CSVOutput, ExpressionType, InputSerialization, JSONInput, OutputSerialization, ParquetInput, ScanRange,
         SelectObjectContentInput, SelectObjectContentRequest,
@@ -531,6 +558,16 @@ mod tests {
         }
     }
 
+    async fn prepare_test_snapshot(context: &Context) -> Arc<SelectObjectSnapshot> {
+        let env = crate::storage_api::select_test_ecstore_env().await;
+        Arc::new(
+            env.ecstore
+                .prepare_select_object_snapshot(&context.input.bucket, &context.input.key, &HeaderMap::new(), &Default::default())
+                .await
+                .expect("prepare SelectObjectContent snapshot"),
+        )
+    }
+
     #[test]
     fn session_factory_fields_remain_source_compatible() {
         let factory = SessionCtxFactory {
@@ -679,10 +716,57 @@ mod tests {
         ));
     }
 
+    #[tokio::test]
+    #[serial_test::serial]
+    async fn production_session_preserves_lazy_snapshot_entry() {
+        let _env = crate::storage_api::select_test_ecstore_env().await;
+        let session = SessionCtxFactory::new(false)
+            .create_session_ctx(&test_context())
+            .await
+            .expect("legacy production session should install a lazy object store");
+
+        assert_eq!(session.inner().config().target_partitions(), SessionConfig::new().target_partitions());
+    }
+
+    #[tokio::test]
+    #[serial_test::serial]
+    async fn legacy_tracked_production_session_preserves_lazy_snapshot_entry() {
+        let _env = crate::storage_api::select_test_ecstore_env().await;
+        let permit = Arc::new(tokio::sync::Semaphore::new(1))
+            .acquire_owned()
+            .await
+            .expect("query permit should be available");
+        let tracker = QueryExecutionTracker::new(
+            &QueryExecutionOwner::new(),
+            Arc::new(permit),
+            Instant::now() + std::time::Duration::from_secs(300),
+            300,
+        );
+        let session = SessionCtxFactory::new(false)
+            .create_session_ctx_with_tracker_and_memory_limit(
+                &test_context(),
+                tracker.clone(),
+                DEFAULT_S3SELECT_MEMORY_LIMIT_BYTES,
+            )
+            .await
+            .expect("legacy tracked session should install a lazy object store");
+
+        assert!(session.is_bound_to(&tracker));
+    }
+
     #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
     #[serial_test::serial]
     async fn session_factory_propagates_query_guard_to_ec_store() {
         let env = crate::storage_api::select_test_ecstore_env().await;
+        let mut context = test_context();
+        Arc::make_mut(&mut context.input).bucket = "s3select-query-guard-snapshot".to_string();
+        env.make_bucket(&context.input.bucket, false).await;
+        let mut reader = SelectPutObjReader::from_vec(b"id,name\n1,Alice\n".to_vec());
+        env.ecstore
+            .put_object(&context.input.bucket, &context.input.key, &mut reader, &Default::default())
+            .await
+            .expect("put query guard fixture");
+        let snapshot = prepare_test_snapshot(&context).await;
 
         let admission = Arc::new(tokio::sync::Semaphore::new(1));
         let permit = Arc::clone(&admission)
@@ -697,7 +781,12 @@ mod tests {
             300,
         );
         let session = SessionCtxFactory::new(false)
-            .create_session_ctx_with_tracker_and_store(&test_context(), query_tracker, Arc::clone(&env.ecstore))
+            .create_session_ctx_with_snapshot_and_tracker_and_memory_limit(
+                &context,
+                snapshot,
+                query_tracker,
+                DEFAULT_S3SELECT_MEMORY_LIMIT_BYTES,
+            )
             .await
             .expect("production session should be created with the query guard");
 
@@ -706,6 +795,51 @@ mod tests {
         assert_eq!(Arc::strong_count(&query_guard), 1);
     }
 
+    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+    #[serial_test::serial]
+    async fn session_factory_preserves_snapshot_binding_error_source() {
+        let env = crate::storage_api::select_test_ecstore_env().await;
+        let mut source_context = test_context();
+        Arc::make_mut(&mut source_context.input).bucket = "s3select-session-snapshot-identity".to_string();
+        Arc::make_mut(&mut source_context.input).key = "source.csv".to_string();
+        env.make_bucket(&source_context.input.bucket, false).await;
+        let mut reader = SelectPutObjReader::from_vec(b"source-marker\n".to_vec());
+        env.ecstore
+            .put_object(&source_context.input.bucket, &source_context.input.key, &mut reader, &Default::default())
+            .await
+            .expect("put snapshot identity fixture");
+        let snapshot = prepare_test_snapshot(&source_context).await;
+
+        let mut target_context = source_context.clone();
+        Arc::make_mut(&mut target_context.input).key = "different.csv".to_string();
+        let permit = Arc::new(tokio::sync::Semaphore::new(1))
+            .acquire_owned()
+            .await
+            .expect("query permit should be available");
+        let tracker = QueryExecutionTracker::new(
+            &QueryExecutionOwner::new(),
+            Arc::new(permit),
+            Instant::now() + std::time::Duration::from_secs(300),
+            300,
+        );
+
+        let error = match SessionCtxFactory::new(false)
+            .create_session_ctx_with_snapshot_and_tracker_and_memory_limit(
+                &target_context,
+                snapshot,
+                tracker,
+                DEFAULT_S3SELECT_MEMORY_LIMIT_BYTES,
+            )
+            .await
+        {
+            Ok(_) => panic!("a session must reject a snapshot for a different object"),
+            Err(error) => error,
+        };
+
+        assert!(error.is_snapshot_consistency_error());
+        assert!(error.to_string().contains("snapshot consistency failure"));
+    }
+
     #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
     #[serial_test::serial]
     async fn scan_range_is_preserved_across_large_csv_partition_boundary() {
@@ -721,6 +855,7 @@ mod tests {
         assert!(data.len() > 1024 * 1024);
 
         let mut context = test_context();
+        Arc::make_mut(&mut context.input).bucket = "s3select-scan-range-partition-snapshot".to_string();
         let selected_start = i64::try_from(SELECTED_ROW * ROW_WIDTH).expect("selected row offset should fit in i64");
         Arc::make_mut(&mut context.input).request.scan_range = Some(ScanRange {
             start: Some(selected_start),
@@ -732,6 +867,7 @@ mod tests {
             .put_object(&context.input.bucket, &context.input.key, &mut reader, &Default::default())
             .await
             .expect("put large ScanRange CSV fixture");
+        let snapshot = prepare_test_snapshot(&context).await;
 
         let admission = Arc::new(tokio::sync::Semaphore::new(1));
         let permit = Arc::clone(&admission)
@@ -746,7 +882,12 @@ mod tests {
         );
         let session = SessionCtxFactory::new(false)
             .with_target_partitions(2)
-            .create_session_ctx_with_tracker_and_store(&context, query_tracker, Arc::clone(&env.ecstore))
+            .create_session_ctx_with_snapshot_and_tracker_and_memory_limit(
+                &context,
+                snapshot,
+                query_tracker,
+                DEFAULT_S3SELECT_MEMORY_LIMIT_BYTES,
+            )
             .await
             .expect("create production ScanRange session");
         assert!(!session.inner().config().options().optimizer.repartition_file_scans);

crates/s3select-api/src/server/dbms.rs

@@ -20,6 +20,7 @@ use crate::{
         Query,
         execution::{Output, QueryStateMachineRef},
         logical_planner::Plan,
+        session::QueryAdmission,
     },
 };
 
@@ -44,7 +45,16 @@ impl QueryHandle {
 
 #[async_trait]
 pub trait DatabaseManagerSystem {
+    fn try_reserve_query(&self) -> QueryResult<QueryAdmission> {
+        Ok(QueryAdmission::unmanaged())
+    }
+
     async fn execute(&self, query: &Query) -> QueryResult<QueryHandle>;
+
+    async fn execute_admitted(&self, query: &Query, _admission: QueryAdmission) -> QueryResult<QueryHandle> {
+        self.execute(query).await
+    }
+
     async fn build_query_state_machine(&self, query: Query) -> QueryResult<QueryStateMachineRef>;
     async fn build_logical_plan(&self, query_state_machine: QueryStateMachineRef) -> QueryResult<Option<Plan>>;
     async fn execute_logical_plan(

crates/s3select-api/src/storage_api.rs

@@ -22,35 +22,51 @@ use rustfs_ecstore::api::error::{
 };
 #[cfg(test)]
 pub(crate) use rustfs_ecstore::api::object::PutObjReader as SelectPutObjReader;
+pub use rustfs_ecstore::api::object::SelectObjectSnapshot;
+pub(crate) use rustfs_ecstore::api::object::{
+    PrepareSelectObjectSnapshotError, SelectObjectSnapshotReadError, SnapshotConsistencyError,
+};
 use rustfs_ecstore::api::runtime::object_store_handle as resolve_select_object_store_handle_from_backend;
 pub(crate) use rustfs_ecstore::api::set_disk::DEFAULT_READ_BUFFER_SIZE as SELECT_DEFAULT_READ_BUFFER_SIZE;
 pub(crate) use rustfs_ecstore::api::storage::ECStore as SelectStore;
 use rustfs_storage_api as storage_contracts;
 
+#[cfg(test)]
+static SELECT_TEST_OBJECT_STORE: std::sync::OnceLock<Arc<SelectStore>> = std::sync::OnceLock::new();
+
 pub(crate) mod object_store {
-    pub(crate) use super::storage_contracts::{HTTPRangeSpec, ObjectIO, ObjectOperations};
+    pub(crate) use super::storage_contracts::HTTPRangeSpec;
+    #[cfg(test)]
+    pub(crate) use super::storage_contracts::ObjectIO;
 }
 
 pub(crate) mod crate_boundary {
     pub(crate) use super::{
-        SELECT_DEFAULT_READ_BUFFER_SIZE, SelectGetObjectReader, SelectObjectInfo, SelectObjectOptions, SelectStorageError,
-        SelectStore, resolve_select_object_store_handle, select_is_err_bucket_not_found, select_is_err_object_not_found,
+        PrepareSelectObjectSnapshotError, SELECT_DEFAULT_READ_BUFFER_SIZE, SelectGetObjectReader, SelectObjectOptions,
+        SelectObjectSnapshotReadError, SelectStorageError, SelectStore, SnapshotConsistencyError,
+        resolve_select_object_store_handle, select_is_err_bucket_not_found, select_is_err_object_not_found,
         select_is_err_version_not_found,
     };
 }
 
 pub(crate) type SelectGetObjectReader = <SelectStore as storage_contracts::ObjectIO>::GetObjectReader;
-pub(crate) type SelectObjectInfo = <SelectStore as storage_contracts::ObjectOperations>::ObjectInfo;
 pub(crate) type SelectObjectOptions = <SelectStore as storage_contracts::ObjectOperations>::ObjectOptions;
 
 #[cfg(test)]
 pub(crate) async fn select_test_ecstore_env() -> &'static rustfs_test_utils::TestECStoreEnv {
     static ENV: tokio::sync::OnceCell<rustfs_test_utils::TestECStoreEnv> = tokio::sync::OnceCell::const_new();
-    ENV.get_or_init(|| async { rustfs_test_utils::TestECStoreEnv::builder().build().await })
-        .await
+    let env = ENV
+        .get_or_init(|| async { rustfs_test_utils::TestECStoreEnv::builder().build().await })
+        .await;
+    let _ = SELECT_TEST_OBJECT_STORE.set(Arc::clone(&env.ecstore));
+    env
 }
 
 pub(crate) fn resolve_select_object_store_handle() -> Option<Arc<SelectStore>> {
+    #[cfg(test)]
+    if let Some(store) = SELECT_TEST_OBJECT_STORE.get() {
+        return Some(Arc::clone(store));
+    }
     resolve_select_object_store_handle_from_backend()
 }