gsadmin/rustfs
1
0
Fork 0
mirror of https://github.com/rustfs/rustfs.git synced 2026-08-10 15:16:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
* 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
This commit is contained in:
GatewayJ
2026-07-25 18:44:53 +08:00
committed by GitHub
parent 2dc4d0b651
commit 0364523dad
13 changed files with 3228 additions and 222 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/s3select-api/src/lib.rs
+85
View File
@@ -64,6 +64,48 @@ pub enum QueryError {
StoreError { e: String },
}
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum S3SelectPolicyError {
#[error("Unsupported S3 Select SQL structure: {message}")]
UnsupportedSqlStructure { message: String },
#[error("S3 Select query concurrency limit reached")]
QueryConcurrencyLimit,
#[error("S3 Select query exceeded the {seconds}-second execution limit")]
QueryTimeout { seconds: u64 },
}
impl S3SelectPolicyError {
fn from_error<'a>(mut err: &'a (dyn std::error::Error + 'static)) -> Option<&'a Self> {
for _ in 0..16 {
if let Some(policy_error) = err.downcast_ref::<Self>() {
return Some(policy_error);
}
err = err.source()?;
}
None
}
}
impl QueryError {
pub fn s3_select_policy_error(&self) -> Option<&S3SelectPolicyError> {
match self {
Self::Datafusion { source } => S3SelectPolicyError::from_error(source.as_ref()),
_ => None,
}
}
}
impl From<S3SelectPolicyError> for QueryError {
fn from(value: S3SelectPolicyError) -> Self {
Self::Datafusion {
source: Box::new(DataFusionError::External(Box::new(value))),
}
}
}
impl From<DataFusionError> for QueryError {
fn from(value: DataFusionError) -> Self {
match value {
@@ -116,9 +158,23 @@ mod tests {
};
assert_eq!(err.to_string(), "Multi-statement not allow, found num:2, sql:SELECT 1; SELECT 2;");
let err = S3SelectPolicyError::UnsupportedSqlStructure {
message: "JOIN is not supported".to_string(),
};
assert_eq!(err.to_string(), "Unsupported S3 Select SQL structure: JOIN is not supported");
let err = QueryError::Cancel;
assert_eq!(err.to_string(), "The query has been canceled");
assert_eq!(
S3SelectPolicyError::QueryConcurrencyLimit.to_string(),
"S3 Select query concurrency limit reached"
);
assert_eq!(
S3SelectPolicyError::QueryTimeout { seconds: 300 }.to_string(),
"S3 Select query exceeded the 300-second execution limit"
);
let err = QueryError::FunctionNotExists {
name: "my_func".to_string(),
};
@@ -143,6 +199,35 @@ mod tests {
}
}
#[test]
fn query_error_variants_remain_source_compatible() {
fn exhaustive_match(err: QueryError) {
match err {
QueryError::Datafusion { .. }
| QueryError::NotImplemented { .. }
| QueryError::MultiStatement { .. }
| QueryError::BuildQueryDispatcher { .. }
| QueryError::Cancel
| QueryError::Parser { .. }
| QueryError::FunctionNotExists { .. }
| QueryError::FunctionExists { .. }
| QueryError::StoreError { .. } => {}
}
}
exhaustive_match(QueryError::Cancel);
}
#[test]
fn policy_error_is_recoverable_from_query_error() {
let err: QueryError = S3SelectPolicyError::QueryTimeout { seconds: 300 }.into();
assert!(matches!(
err.s3_select_policy_error(),
Some(S3SelectPolicyError::QueryTimeout { seconds: 300 })
));
}
#[test]
fn test_query_error_from_parser_error() {
let parser_error = ParserError::ParserError("syntax error".to_string());
crates/s3select-api/src/object_store.rs
+576 -64
View File
@@ -14,20 +14,27 @@
use crate::{
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,
SelectStore,
query::session::{QueryExecutionGuard, QueryExecutionTracker},
resolve_select_object_store_handle, select_is_err_bucket_not_found, select_is_err_object_not_found,
select_is_err_version_not_found,
};
use async_trait::async_trait;
use bytes::Bytes;
use chrono::Utc;
use datafusion::object_store::{
Attributes, CopyOptions, Error as o_Error, GetOptions, GetRange, GetResult, GetResultPayload, ListResult, MultipartUpload,
ObjectMeta, ObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult, Result, path::Path,
use datafusion::{
common::{DataFusionError, runtime::SpawnedTask},
execution::memory_pool::{MemoryConsumer, MemoryPool, UnboundedMemoryPool},
object_store::{
Attributes, CopyOptions, Error as o_Error, GetOptions, GetRange, GetResult, GetResultPayload, ListResult,
MultipartUpload, ObjectMeta, ObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult, Result, path::Path,
},
};
use futures::pin_mut;
use futures::{Stream, StreamExt, future::ready, stream};
use futures_core::stream::BoxStream;
use http::{HeaderMap, HeaderValue, header::HeaderName};
use parking_lot::Mutex;
use rustfs_common::DEFAULT_DELIMITER;
use s3s::S3Result;
use s3s::dto::SelectObjectContentInput;
@@ -49,6 +56,13 @@ fn select_default_read_buffer_size_u64() -> u64 {
u64::try_from(SELECT_DEFAULT_READ_BUFFER_SIZE).unwrap_or(u64::MAX)
}
fn validated_object_size(size: i64) -> Result<u64> {
u64::try_from(size).map_err(|err| o_Error::Generic {
store: "EcObjectStore",
source: Box::new(err),
})
}
/// Maximum allowed object size for JSON DOCUMENT mode.
///
/// JSON DOCUMENT format requires loading the entire file into memory for DOM
@@ -61,8 +75,13 @@ fn select_default_read_buffer_size_u64() -> u64 {
/// size limit.
///
/// Default: 128 MiB. This matches the AWS S3 Select limit for JSON DOCUMENT
/// inputs.
/// inputs. The query memory pool also applies: RustFS reserves 64 times the
/// input size for parsing and output. With the default 64 MiB query memory
/// limit, JSON DOCUMENT inputs larger than 1 MiB are rejected; raise
/// `RUSTFS_S3SELECT_MEMORY_LIMIT_BYTES` to process larger inputs, up to this
/// hard cap.
pub const MAX_JSON_DOCUMENT_BYTES: u64 = 128 * 1024 * 1024;
const JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER: usize = 64;
pub const INVALID_SCAN_RANGE_MESSAGE: &str =
"The value of a parameter in ScanRange element is invalid. Check the service API documentation and try again.";
@@ -79,6 +98,8 @@ pub struct EcObjectStore {
/// expression. When set, `flatten_json_document_to_ndjson` navigates to
/// this key in the root JSON object before flattening.
json_sub_path: Option<String>,
memory_pool: Arc<dyn MemoryPool>,
query_tracker: Option<QueryExecutionTracker>,
store: Arc<SelectStore>,
}
@@ -108,7 +129,29 @@ pub struct InvalidScanRange;
impl EcObjectStore {
pub fn new(input: Arc<SelectObjectContentInput>) -> S3Result<Self> {
let Some(store) = resolve_select_object_store_handle() else {
Self::build(input, Arc::new(UnboundedMemoryPool::default()), None, None)
}
pub(crate) fn new_with_memory_pool(input: Arc<SelectObjectContentInput>, memory_pool: Arc<dyn MemoryPool>) -> S3Result<Self> {
Self::build(input, memory_pool, None, None)
}
pub(crate) fn new_with_query_tracker(
input: Arc<SelectObjectContentInput>,
memory_pool: Arc<dyn MemoryPool>,
query_tracker: QueryExecutionTracker,
store: Option<Arc<SelectStore>>,
) -> S3Result<Self> {
Self::build(input, memory_pool, Some(query_tracker), store)
}
fn build(
input: Arc<SelectObjectContentInput>,
memory_pool: Arc<dyn MemoryPool>,
query_tracker: Option<QueryExecutionTracker>,
store: Option<Arc<SelectStore>>,
) -> S3Result<Self> {
let Some(store) = store.or_else(resolve_select_object_store_handle) else {
return Err(s3_error!(InternalError, "ec store not inited"));
};
@@ -151,6 +194,8 @@ impl EcObjectStore {
delimiter,
is_json_document,
json_sub_path,
memory_pool,
query_tracker,
store,
})
}
@@ -213,13 +258,30 @@ impl EcObjectStore {
if range.is_empty() {
return Ok(Bytes::new());
}
let reader = self.object_reader(Some(http_range_spec_from_range(range)), opts).await?;
let reader = self
.object_reader(Some(http_range_spec_from_range(range.clone())), opts)
.await?;
let object_size = validated_object_size(reader.object_info.size)?;
let resolved_range = GetRange::Bounded(range)
.as_range(object_size)
.map_err(|err| o_Error::Generic {
store: "EcObjectStore",
source: Box::new(err),
})?;
let expected_size = usize::try_from(resolved_range.end - resolved_range.start).map_err(|err| o_Error::Generic {
store: "EcObjectStore",
source: Box::new(err),
})?;
let mut reader = reader.stream;
let mut bytes = Vec::new();
reader.read_to_end(&mut bytes).await.map_err(|err| o_Error::Generic {
store: "EcObjectStore",
source: Box::new(err),
})?;
if bytes.len() < expected_size {
return Err(incomplete_object_stream_error(expected_size - bytes.len()));
}
bytes.truncate(expected_size);
Ok(Bytes::from(bytes))
}
@@ -421,7 +483,7 @@ impl ObjectStore for EcObjectStore {
let opts = self.object_options(&options);
let needs_scan_context = options.range.is_none() && !options.head && self.input.request.scan_range.is_some();
let source_size = if needs_scan_context {
Some(self.object_info(&opts).await?.size as u64)
Some(validated_object_size(self.object_info(&opts).await?.size)?)
} else {
None
};
@@ -442,7 +504,10 @@ impl ObjectStore for EcObjectStore {
self.object_reader(range, &opts).await?
};
let original_size = source_size.unwrap_or(reader.object_info.size as u64);
let original_size = match source_size {
Some(source_size) => source_size,
None => validated_object_size(reader.object_info.size)?,
};
let etag = reader.object_info.etag;
let version = reader.object_info.version_id.map(|version| version.to_string());
let attributes = Attributes::default();
@@ -473,20 +538,14 @@ impl ObjectStore for EcObjectStore {
// which must load the whole file into memory; rejecting oversized
// files upfront is safer than risking OOM. Users should convert
// their data to JSON LINES (NDJSON) format for large files.
if original_size > MAX_JSON_DOCUMENT_BYTES {
return Err(o_Error::Generic {
store: "EcObjectStore",
source: format!(
"JSON DOCUMENT object is {original_size} bytes, which exceeds the \
maximum allowed size of {MAX_JSON_DOCUMENT_BYTES} bytes \
({} MiB). Convert the input to JSON LINES (NDJSON) to process \
large files.",
MAX_JSON_DOCUMENT_BYTES / (1024 * 1024)
)
.into(),
});
}
let stream = json_document_ndjson_stream(reader.stream, original_size, self.json_sub_path.clone());
validate_json_document_size(original_size)?;
let stream = json_document_ndjson_stream(
reader.stream,
original_size,
self.json_sub_path.clone(),
Arc::clone(&self.memory_pool),
self.query_tracker.clone(),
);
GetResultPayload::Stream(stream)
} else if let Some((_, scan_range)) = scan_context {
let delimiter = self.record_delimiter();
@@ -503,6 +562,7 @@ impl ObjectStore for EcObjectStore {
scan_range,
include_header && header.is_none(),
read_start,
original_size,
)
.boxed();
let stream = if let Some(header) = header {
@@ -664,6 +724,7 @@ struct ScanRangeState<S> {
record: Vec<u8>,
pending: VecDeque<Bytes>,
done: bool,
expected_end: u64,
}
fn scan_range_stream<S>(
@@ -672,6 +733,7 @@ fn scan_range_stream<S>(
range: SelectScanRange,
include_header: bool,
base_offset: u64,
expected_end: u64,
) -> BoxStream<'static, Result<Bytes>>
where
S: Stream<Item = std::io::Result<Bytes>> + Send + Unpin + 'static,
@@ -686,6 +748,7 @@ where
record: Vec::new(),
pending: VecDeque::new(),
done: false,
expected_end,
};
stream::unfold(state, |mut state| async move {
@@ -709,6 +772,10 @@ where
));
}
None => {
if state.offset < state.expected_end {
state.done = true;
return Some((Err(incomplete_object_stream_error(state.expected_end - state.offset)), state));
}
state.finish_pending_record();
state.done = true;
}
@@ -855,11 +922,57 @@ fn json_document_ndjson_stream(
stream: Box<dyn tokio::io::AsyncRead + Unpin + Send + Sync>,
original_size: u64,
json_sub_path: Option<String>,
memory_pool: Arc<dyn MemoryPool>,
query_tracker: Option<QueryExecutionTracker>,
) -> futures_core::stream::BoxStream<'static, Result<Bytes>> {
json_document_ndjson_stream_with_parser(
stream,
original_size,
json_sub_path,
memory_pool,
query_tracker,
|all_bytes, json_sub_path| parse_json_document_to_lines(&all_bytes, json_sub_path.as_deref()),
)
}
fn json_document_ndjson_stream_with_parser<P>(
stream: Box<dyn tokio::io::AsyncRead + Unpin + Send + Sync>,
original_size: u64,
json_sub_path: Option<String>,
memory_pool: Arc<dyn MemoryPool>,
query_tracker: Option<QueryExecutionTracker>,
parser: P,
) -> futures_core::stream::BoxStream<'static, Result<Bytes>>
where
P: FnOnce(Vec<u8>, Option<String>) -> std::io::Result<Vec<Bytes>> + Send + 'static,
{
AsyncTryStream::<Bytes, o_Error, _>::new(|mut y| async move {
// Compact JSON can expand substantially into a serde_json DOM and
// per-record output buffers, so reserve a conservative upper bound
// before the source buffer is allocated.
let buffer_capacity = usize::try_from(original_size).map_err(|_| o_Error::Generic {
store: "EcObjectStore",
source: Box::new(DataFusionError::ResourcesExhausted(format!(
"JSON DOCUMENT input size {original_size} does not fit in memory"
))),
})?;
let reservation_bytes = buffer_capacity
.checked_mul(JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER)
.ok_or_else(|| o_Error::Generic {
store: "EcObjectStore",
source: Box::new(DataFusionError::ResourcesExhausted(format!(
"JSON DOCUMENT memory reservation overflow for {original_size} input bytes"
))),
})?;
let reservation = MemoryConsumer::new("S3 Select JSON document").register(&memory_pool);
reservation.try_resize(reservation_bytes).map_err(|err| o_Error::Generic {
store: "EcObjectStore",
source: Box::new(err),
})?;
pin_mut!(stream);
// ── 1. Read phase (lazy: only runs when the stream is polled) ────
let mut all_bytes = Vec::with_capacity(original_size as usize);
let mut all_bytes = Vec::with_capacity(buffer_capacity);
stream
.take(original_size)
.read_to_end(&mut all_bytes)
@@ -868,18 +981,26 @@ fn json_document_ndjson_stream(
store: "EcObjectStore",
source: Box::new(e),
})?;
if all_bytes.len() != buffer_capacity {
return Err(incomplete_object_stream_error(buffer_capacity - all_bytes.len()));
}
// ── 2. Parse phase (blocking thread pool, non-blocking runtime) ──
let lines = tokio::task::spawn_blocking(move || parse_json_document_to_lines(&all_bytes, json_sub_path.as_deref()))
.await
.map_err(|e| o_Error::Generic {
store: "EcObjectStore",
source: e.to_string().into(),
})?
.map_err(|e| o_Error::Generic {
store: "EcObjectStore",
source: Box::new(e),
})?;
let pending_query_guard = PendingQueryExecutionGuard::new(query_tracker);
let task_query_guard = pending_query_guard.task_state();
let (lines, _reservation, _query_guard) = SpawnedTask::spawn_blocking(move || {
let query_guard = PendingQueryExecutionGuard::start(&task_query_guard)?;
parser(all_bytes, json_sub_path).map(|lines| (lines, reservation, query_guard))
})
.await
.map_err(|e| o_Error::Generic {
store: "EcObjectStore",
source: e.to_string().into(),
})?
.map_err(|e| o_Error::Generic {
store: "EcObjectStore",
source: Box::new(e),
})?;
// ── 3. Yield phase (one Bytes per NDJSON line) ───────────────────
for line in lines {
@@ -890,6 +1011,66 @@ fn json_document_ndjson_stream(
.boxed()
}
struct PendingQueryExecutionGuard {
state: Arc<Mutex<QueryExecutionGuardState>>,
}
enum QueryExecutionGuardState {
Pending(Option<QueryExecutionTracker>),
Started,
Cancelled,
}
impl PendingQueryExecutionGuard {
fn new(query_tracker: Option<QueryExecutionTracker>) -> Self {
Self {
state: Arc::new(Mutex::new(QueryExecutionGuardState::Pending(query_tracker))),
}
}
fn task_state(&self) -> Arc<Mutex<QueryExecutionGuardState>> {
Arc::clone(&self.state)
}
fn start(state: &Mutex<QueryExecutionGuardState>) -> std::io::Result<Option<QueryExecutionGuard>> {
let mut state = state.lock();
match std::mem::replace(&mut *state, QueryExecutionGuardState::Started) {
QueryExecutionGuardState::Pending(None) => Ok(None),
QueryExecutionGuardState::Pending(Some(query_tracker)) => query_tracker.query_guard().map(Some).ok_or_else(|| {
std::io::Error::new(std::io::ErrorKind::Interrupted, "JSON DOCUMENT parse was cancelled before it started")
}),
QueryExecutionGuardState::Cancelled => {
*state = QueryExecutionGuardState::Cancelled;
Err(std::io::Error::new(
std::io::ErrorKind::Interrupted,
"JSON DOCUMENT parse was cancelled before it started",
))
}
QueryExecutionGuardState::Started => {
*state = QueryExecutionGuardState::Started;
Err(std::io::Error::other("JSON DOCUMENT parse started more than once"))
}
}
}
}
impl Drop for PendingQueryExecutionGuard {
fn drop(&mut self) {
let query_guard = {
let mut state = self.state.lock();
match std::mem::replace(&mut *state, QueryExecutionGuardState::Cancelled) {
QueryExecutionGuardState::Pending(query_guard) => query_guard,
QueryExecutionGuardState::Started => {
*state = QueryExecutionGuardState::Started;
None
}
QueryExecutionGuardState::Cancelled => None,
}
};
drop(query_guard);
}
}
/// Parse a JSON DOCUMENT (a single JSON value, possibly multi-line) into a
/// list of NDJSON lines – one [`Bytes`] per record.
///
@@ -907,19 +1088,11 @@ fn parse_json_document_to_lines(bytes: &[u8], json_sub_path: Option<&str>) -> st
// Navigate into the sub-path when the root is an object and a path was
// extracted from the SQL FROM clause (e.g. `FROM s3object.employees`).
let value = if let Some(path) = json_sub_path {
if let serde_json::Value::Object(ref obj) = root {
match obj.get(path) {
Some(sub) => sub.clone(),
// Path not found – fall back to emitting the whole root object.
None => root,
}
} else {
// Root is already an array or scalar; ignore the path hint.
root
let value = match (root, json_sub_path) {
(serde_json::Value::Object(mut object), Some(path)) => {
object.remove(path).unwrap_or_else(|| serde_json::Value::Object(object))
}
} else {
root
(root, _) => root,
};
let mut lines: Vec<Bytes> = Vec::new();
@@ -979,30 +1152,55 @@ where
y.yield_ok(bytes).await;
}
if remaining > 0 {
return Err(o_Error::Generic {
store: "EcObjectStore",
source: Box::new(std::io::Error::new(
std::io::ErrorKind::UnexpectedEof,
format!("object stream ended with {remaining} bytes remaining"),
)),
});
return Err(incomplete_object_stream_error(remaining));
}
Ok(())
})
}
fn validate_json_document_size(original_size: u64) -> Result<()> {
if original_size <= MAX_JSON_DOCUMENT_BYTES {
return Ok(());
}
Err(o_Error::Generic {
store: "EcObjectStore",
source: Box::new(DataFusionError::ResourcesExhausted(format!(
"JSON DOCUMENT object is {original_size} bytes, which exceeds the maximum allowed size of \
{MAX_JSON_DOCUMENT_BYTES} bytes ({} MiB). Convert the input to JSON LINES (NDJSON) to process large files.",
MAX_JSON_DOCUMENT_BYTES / (1024 * 1024)
))),
})
}
fn incomplete_object_stream_error(remaining: impl std::fmt::Display) -> o_Error {
o_Error::Generic {
store: "EcObjectStore",
source: Box::new(std::io::Error::new(
std::io::ErrorKind::UnexpectedEof,
format!("object stream ended with {remaining} bytes remaining"),
)),
}
}
#[cfg(test)]
mod test {
use super::{
SelectScanRange, bytes_stream, convert_field_delimiter_stream, extract_json_sub_path_from_expression, find_delimiter,
flatten_json_document_to_ndjson, http_range_spec_from_get_range, replace_symbol, scan_range_from_bounds,
scan_range_read_start, scan_range_stream, select_read_headers,
EcObjectStore, JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER, SelectScanRange, bytes_stream,
convert_field_delimiter_stream, extract_json_sub_path_from_expression, find_delimiter, flatten_json_document_to_ndjson,
http_range_spec_from_get_range, json_document_ndjson_stream, json_document_ndjson_stream_with_parser, replace_symbol,
scan_range_from_bounds, scan_range_read_start, scan_range_stream, select_read_headers, validate_json_document_size,
validated_object_size,
};
use crate::query::session::{QueryExecutionGuard, QueryExecutionOwner, QueryExecutionTracker};
use crate::storage_api::SelectPutObjReader;
use crate::storage_api::object_store::ObjectIO as _;
use bytes::Bytes;
use datafusion::object_store::{self, GetRange};
use datafusion::object_store::{GetOptions, GetResultPayload, ObjectStore as _, path::Path};
use datafusion::{
common::DataFusionError,
execution::memory_pool::{GreedyMemoryPool, MemoryPool},
object_store::{self, GetOptions, GetRange, GetResultPayload, ObjectStore as _, path::Path},
};
use futures::{StreamExt, TryStreamExt, stream};
use s3s::dto::{
CSVInput, CSVOutput, ExpressionType, InputSerialization, OutputSerialization, SelectObjectContentInput,
@@ -1017,16 +1215,28 @@ mod test {
atomic::{AtomicUsize, Ordering},
};
#[test]
fn ec_object_store_constructor_remains_source_compatible() {
let _constructor: fn(Arc<SelectObjectContentInput>) -> s3s::S3Result<EcObjectStore> = EcObjectStore::new;
}
use tokio::sync::Semaphore;
#[test]
fn test_replace() {
let result = replace_symbol(b"&&", b"dandan&&is&&best");
assert_eq!(result, b"dandan,is,best");
}
#[test]
fn test_validated_object_size_rejects_negative_metadata() {
assert_eq!(validated_object_size(0).expect("zero object size should be valid"), 0);
assert!(validated_object_size(-1).is_err());
}
#[tokio::test]
async fn test_scan_range_stream_keeps_header_and_selected_record() {
let chunks = stream::iter(vec![Ok::<_, std::io::Error>(Bytes::from_static(b"h1,h2\n1,a\n2,b\n3,c\n"))]);
let mut stream = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange::new(10, 11), true, 0);
let mut stream = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange::new(10, 11), true, 0, 18);
let mut output = Vec::new();
while let Some(bytes) = stream.next().await {
output.extend_from_slice(&bytes.unwrap());
@@ -1037,7 +1247,7 @@ mod test {
#[tokio::test]
async fn test_scan_range_stream_skips_record_when_start_is_in_middle() {
let chunks = stream::iter(vec![Ok::<_, std::io::Error>(Bytes::from_static(b"1,a\n2,b\n3,c\n"))]);
let mut stream = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange::new(2, 7), false, 0);
let mut stream = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange::new(2, 7), false, 0, 12);
let mut output = Vec::new();
while let Some(bytes) = stream.next().await {
output.extend_from_slice(&bytes.unwrap());
@@ -1048,7 +1258,7 @@ mod test {
#[tokio::test]
async fn test_scan_range_stream_keeps_record_when_end_is_in_middle() {
let chunks = stream::iter(vec![Ok::<_, std::io::Error>(Bytes::from_static(b"1,a\n2,b\n3,c\n"))]);
let mut stream = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange::new(0, 5), false, 0);
let mut stream = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange::new(0, 5), false, 0, 12);
let mut output = Vec::new();
while let Some(bytes) = stream.next().await {
output.extend_from_slice(&bytes.unwrap());
@@ -1059,7 +1269,7 @@ mod test {
#[tokio::test]
async fn test_scan_range_stream_uses_base_offset_for_range_reader() {
let chunks = stream::iter(vec![Ok::<_, std::io::Error>(Bytes::from_static(b"\n2,b\n3,c\n"))]);
let mut stream = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange::new(4, 7), false, 3);
let mut stream = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange::new(4, 7), false, 3, 12);
let mut output = Vec::new();
while let Some(bytes) = stream.next().await {
output.extend_from_slice(&bytes.unwrap());
@@ -1074,7 +1284,7 @@ mod test {
Ok::<_, std::io::Error>(Bytes::from_static(b"\n1,a\r\n2,b\r")),
Ok::<_, std::io::Error>(Bytes::from_static(b"\n3,c\r\n")),
]);
let mut stream = scan_range_stream(chunks, b"\r\n".to_vec(), SelectScanRange::new(12, 14), true, 0);
let mut stream = scan_range_stream(chunks, b"\r\n".to_vec(), SelectScanRange::new(12, 14), true, 0, 22);
let mut output = Vec::new();
while let Some(bytes) = stream.next().await {
output.extend_from_slice(&bytes.unwrap());
@@ -1082,6 +1292,26 @@ mod test {
assert_eq!(output, b"h1,h2\r\n2,b\r\n");
}
#[tokio::test]
async fn test_scan_range_stream_rejects_early_eof() {
let chunks = stream::iter(vec![Ok::<_, std::io::Error>(Bytes::from_static(b"1,a\n"))]);
let mut output = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange::new(0, 7), false, 0, 8);
assert_eq!(output.next().await.expect("first stream item").expect("first record"), b"1,a\n"[..]);
let err = output
.next()
.await
.expect("early EOF error")
.expect_err("short ScanRange stream must fail");
let object_store::Error::Generic { source, .. } = err else {
panic!("expected generic object store error");
};
let source = source.downcast_ref::<std::io::Error>().expect("I/O error source");
assert_eq!(source.kind(), std::io::ErrorKind::UnexpectedEof);
assert!(source.to_string().contains("4 bytes remaining"));
assert!(output.next().await.is_none());
}
#[test]
fn test_scan_range_read_start_keeps_full_delimiter_boundary() {
let range = SelectScanRange::new(10, 20);
@@ -1157,7 +1387,7 @@ mod test {
#[tokio::test]
async fn test_scan_range_output_can_convert_field_delimiter() {
let chunks = stream::iter(vec![Ok::<_, std::io::Error>(Bytes::from_static(b"a&&1\nb&&2\n"))]);
let stream = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange::new(0, 10), false, 0);
let stream = scan_range_stream(chunks, b"\n".to_vec(), SelectScanRange::new(0, 10), false, 0, 10);
let mut stream = convert_field_delimiter_stream(stream, Some("&&".to_string()));
let mut output = Vec::new();
while let Some(bytes) = stream.next().await {
@@ -1195,6 +1425,7 @@ mod test {
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial_test::serial]
async fn test_get_opts_validates_raw_length_before_delimiter_conversion() {
let temp_root = tempfile::tempdir().expect("create s3select test temp root");
let env = rustfs_test_utils::TestECStoreEnv::builder()
@@ -1242,6 +1473,8 @@ mod test {
delimiter: "&&".to_string(),
is_json_document: false,
json_sub_path: None,
memory_pool: Arc::new(GreedyMemoryPool::new(1024)),
query_tracker: None,
store: env.ecstore,
};
@@ -1255,6 +1488,13 @@ mod test {
let chunks: Vec<Bytes> = stream.try_collect().await.expect("collect converted object stream");
assert_eq!(chunks.concat(), b"a,1\n");
let requested_range = 3..10;
let ranges = store
.get_ranges(&Path::from(object), std::slice::from_ref(&requested_range))
.await
.expect("bounded range past EOF should return the object remainder");
assert_eq!(ranges, vec![Bytes::from_static(b"1\n")]);
}
#[tokio::test]
@@ -1302,6 +1542,278 @@ mod test {
assert!(output.next().await.is_none());
}
#[tokio::test]
async fn test_json_document_stream_respects_query_memory_pool() {
let input = b"{}".to_vec();
let required = input.len() * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER;
let memory_pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(required - 1));
let mut output = json_document_ndjson_stream(
Box::new(std::io::Cursor::new(input.clone())),
input.len() as u64,
None,
memory_pool,
None,
);
let err = output
.next()
.await
.expect("memory error")
.expect_err("reservation should exceed the pool");
let object_store::Error::Generic { source, .. } = err else {
panic!("expected generic object store error");
};
assert!(matches!(
source.downcast_ref::<DataFusionError>(),
Some(DataFusionError::ResourcesExhausted(_))
));
}
#[tokio::test]
async fn test_json_document_stream_releases_memory_reservation() {
let input = b"[1,2]".to_vec();
let required = input.len() * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER;
let memory_pool = Arc::new(GreedyMemoryPool::new(required));
let output: Vec<Bytes> = json_document_ndjson_stream(
Box::new(std::io::Cursor::new(input.clone())),
input.len() as u64,
None,
memory_pool.clone(),
None,
)
.try_collect()
.await
.expect("JSON conversion should fit the pool");
assert_eq!(output, vec![Bytes::from_static(b"1\n"), Bytes::from_static(b"2\n")]);
assert_eq!(memory_pool.reserved(), 0);
}
#[tokio::test]
async fn test_json_document_stream_rejects_early_eof() {
let input = b"{}".to_vec();
let memory_pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(4 * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER));
let mut output = json_document_ndjson_stream(Box::new(std::io::Cursor::new(input)), 4, None, memory_pool, None);
let err = output
.next()
.await
.expect("early EOF error")
.expect_err("short JSON document must fail");
let object_store::Error::Generic { source, .. } = err else {
panic!("expected generic object store error");
};
let source = source.downcast_ref::<std::io::Error>().expect("I/O error source");
assert_eq!(source.kind(), std::io::ErrorKind::UnexpectedEof);
assert!(source.to_string().contains("2 bytes remaining"));
assert!(output.next().await.is_none());
}
#[test]
fn test_json_document_size_error_is_resource_exhausted() {
assert!(validate_json_document_size(super::MAX_JSON_DOCUMENT_BYTES).is_ok());
let err = validate_json_document_size(super::MAX_JSON_DOCUMENT_BYTES + 1).expect_err("oversized JSON document must fail");
let object_store::Error::Generic { source, .. } = err else {
panic!("expected generic object store error");
};
assert!(matches!(
source.downcast_ref::<DataFusionError>(),
Some(DataFusionError::ResourcesExhausted(_))
));
}
#[test]
fn test_json_document_queued_parse_releases_query_guard_when_cancelled() {
let runtime = tokio::runtime::Builder::new_multi_thread()
.worker_threads(2)
.max_blocking_threads(1)
.enable_all()
.build()
.expect("build test runtime");
runtime.block_on(async {
let (blocking_started_tx, blocking_started_rx) = tokio::sync::oneshot::channel();
let (release_blocking_tx, release_blocking_rx) = std::sync::mpsc::channel();
let blocker = tokio::task::spawn_blocking(move || {
let _ = blocking_started_tx.send(());
release_blocking_rx.recv().expect("release blocking worker");
});
blocking_started_rx.await.expect("blocking worker should start");
let admission = Arc::new(Semaphore::new(1));
let permit = Arc::clone(&admission)
.acquire_owned()
.await
.expect("query permit should be available");
let query_guard: QueryExecutionGuard = Arc::new(permit);
let query_tracker = QueryExecutionTracker::new(
&QueryExecutionOwner::new(),
query_guard,
tokio::time::Instant::now() + std::time::Duration::from_secs(30),
30,
);
let input = b"{}".to_vec();
let memory_pool: Arc<dyn MemoryPool> =
Arc::new(GreedyMemoryPool::new(input.len() * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER));
let mut output = json_document_ndjson_stream(
Box::new(std::io::Cursor::new(input.clone())),
input.len() as u64,
None,
memory_pool,
Some(query_tracker),
);
{
let next = output.next();
futures::pin_mut!(next);
assert!(futures::poll!(next.as_mut()).is_pending());
}
drop(output);
let recovered_permit =
tokio::time::timeout(std::time::Duration::from_secs(5), Arc::clone(&admission).acquire_owned())
.await
.expect("queued JSON parse should be cancelled")
.expect("query admission should remain open");
release_blocking_tx.send(()).expect("release blocking worker");
blocker.await.expect("blocking worker should finish");
drop(recovered_permit);
assert_eq!(admission.available_permits(), 1);
});
}
#[test]
fn test_json_document_expired_queued_parse_does_not_start() {
let runtime = tokio::runtime::Builder::new_multi_thread()
.worker_threads(2)
.max_blocking_threads(1)
.enable_all()
.build()
.expect("build test runtime");
runtime.block_on(async {
let (blocking_started_tx, blocking_started_rx) = tokio::sync::oneshot::channel();
let (release_blocking_tx, release_blocking_rx) = std::sync::mpsc::channel();
let blocker = tokio::task::spawn_blocking(move || {
let _ = blocking_started_tx.send(());
release_blocking_rx.recv().expect("release blocking worker");
});
blocking_started_rx.await.expect("blocking worker should start");
let admission = Arc::new(Semaphore::new(1));
let permit = Arc::clone(&admission)
.acquire_owned()
.await
.expect("query permit should be available");
let owner = QueryExecutionOwner::new();
let query_tracker = QueryExecutionTracker::new(
&owner,
Arc::new(permit),
tokio::time::Instant::now() + std::time::Duration::from_secs(30),
30,
);
let input = b"{}".to_vec();
let memory_pool: Arc<dyn MemoryPool> =
Arc::new(GreedyMemoryPool::new(input.len() * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER));
let parser_started = Arc::new(std::sync::atomic::AtomicBool::new(false));
let parser_started_in_task = Arc::clone(&parser_started);
let mut output = json_document_ndjson_stream_with_parser(
Box::new(std::io::Cursor::new(input.clone())),
input.len() as u64,
None,
memory_pool,
Some(query_tracker.clone()),
move |_, _| {
parser_started_in_task.store(true, std::sync::atomic::Ordering::SeqCst);
Ok(vec![Bytes::from_static(b"{}\n")])
},
);
{
let next = output.next();
futures::pin_mut!(next);
assert!(futures::poll!(next.as_mut()).is_pending());
}
query_tracker.expire(&owner);
assert_eq!(admission.available_permits(), 1);
release_blocking_tx.send(()).expect("release blocking worker");
blocker.await.expect("blocking worker should finish");
let err = tokio::time::timeout(std::time::Duration::from_secs(5), output.next())
.await
.expect("queued parser should resume")
.expect("queued parser should return an error")
.expect_err("expired queued parser must not run");
let object_store::Error::Generic { source, .. } = err else {
panic!("expected generic object store error");
};
let source = source.downcast_ref::<std::io::Error>().expect("I/O error source");
assert_eq!(source.kind(), std::io::ErrorKind::Interrupted);
assert!(!parser_started.load(std::sync::atomic::Ordering::SeqCst));
});
}
#[test]
fn test_json_document_started_parse_retains_query_guard_when_cancelled() {
let runtime = tokio::runtime::Builder::new_multi_thread()
.worker_threads(2)
.max_blocking_threads(1)
.enable_all()
.build()
.expect("build test runtime");
runtime.block_on(async {
let admission = Arc::new(Semaphore::new(1));
let permit = Arc::clone(&admission)
.acquire_owned()
.await
.expect("query permit should be available");
let query_guard: QueryExecutionGuard = Arc::new(permit);
let query_tracker = QueryExecutionTracker::new(
&QueryExecutionOwner::new(),
query_guard,
tokio::time::Instant::now() + std::time::Duration::from_secs(30),
30,
);
let input = b"{}".to_vec();
let memory_pool: Arc<dyn MemoryPool> =
Arc::new(GreedyMemoryPool::new(input.len() * JSON_DOCUMENT_MEMORY_RESERVATION_MULTIPLIER));
let (parse_started_tx, parse_started_rx) = tokio::sync::oneshot::channel();
let (release_parse_tx, release_parse_rx) = std::sync::mpsc::channel();
let mut output = json_document_ndjson_stream_with_parser(
Box::new(std::io::Cursor::new(input.clone())),
input.len() as u64,
None,
memory_pool,
Some(query_tracker),
move |_, _| {
let _ = parse_started_tx.send(());
release_parse_rx.recv().expect("release JSON parser");
Ok(vec![Bytes::from_static(b"{}\n")])
},
);
{
let next = output.next();
futures::pin_mut!(next);
assert!(futures::poll!(next.as_mut()).is_pending());
}
parse_started_rx.await.expect("JSON parser should start");
drop(output);
assert!(Arc::clone(&admission).try_acquire_owned().is_err());
release_parse_tx.send(()).expect("release JSON parser");
let recovered_permit =
tokio::time::timeout(std::time::Duration::from_secs(5), Arc::clone(&admission).acquire_owned())
.await
.expect("started JSON parse should release the query guard")
.expect("query admission should remain open");
drop(recovered_permit);
assert_eq!(admission.available_permits(), 1);
});
}
/// A JSON array is split into one NDJSON line per element.
#[test]
fn test_flatten_array_produces_one_line_per_element() {
crates/s3select-api/src/query/execution.rs
+22 -1
View File
@@ -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
+371 -8
View File
@@ -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)
));
}
}