use chrono::{DateTime, Utc};
use lru::LruCache;
use radixdb_catalog::ObjectId;
use radixdb_core::time_compat::{system_time_now, Instant};
use rustc_hash::FxHashMap;
use std::cell::RefCell;
use std::collections::BinaryHeap;
use std::num::NonZeroUsize;
use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
use std::sync::{Arc, Condvar, LazyLock, Mutex};
use std::time::Duration;
const SCALAR_SUBQUERY_CACHE_SIZE: usize = 128;
const IN_SUBQUERY_CACHE_SIZE: usize = 128;
const SEMI_JOIN_CACHE_SIZE: usize = 256;
use crate::hash_table::{JoinHashState, JoinHashTable, JoinMemoryOwner, JoinMemoryReservation};
use radixdb_core::ParamVec;
use radixdb_core::{CompactArc, StringMap};
use radixdb_core::{Result, Row, Value, ValueMap, ValueSet};
use radixdb_procedural::BudgetOwner;
pub(crate) trait StoredFunctionInvoker: std::fmt::Debug + Send + Sync {
fn invoke(self: Arc<Self>, name: &str, arguments: &[Value]) -> Result<Value>;
fn external_equal(&self, left: &Value, right: &Value) -> Result<bool> {
let _ = (left, right);
Err(radixdb_core::Error::NotSupported(
"external equality is unavailable in this execution context".to_owned(),
))
}
fn external_compare(&self, left: &Value, right: &Value) -> Result<std::cmp::Ordering> {
let _ = (left, right);
Err(radixdb_core::Error::NotSupported(
"external ordering is unavailable in this execution context".to_owned(),
))
}
fn external_input(&self, type_name: &str, input: &Value) -> Result<Value> {
let _ = (type_name, input);
Err(radixdb_core::Error::NotSupported(
"external type input is unavailable in this execution context".to_owned(),
))
}
fn external_output(&self, value: &Value, target_type: radixdb_core::DataType) -> Result<Value> {
let _ = (value, target_type);
Err(radixdb_core::Error::NotSupported(
"external type output is unavailable in this execution context".to_owned(),
))
}
}
pub(crate) const STORED_OPERATOR_CALL_PREFIX: &str = "\u{1f}operator:";
static EMPTY_PARAMS: LazyLock<CompactArc<ParamVec>> =
LazyLock::new(|| CompactArc::new(ParamVec::new()));
static EMPTY_DATABASE: LazyLock<Arc<Option<String>>> = LazyLock::new(|| Arc::new(None));
static EMPTY_SESSION_VARS: LazyLock<Arc<AHashMap<String, Value>>> =
LazyLock::new(|| Arc::new(AHashMap::new()));
pub(crate) fn is_system_context_name(name: &str) -> bool {
matches!(
name.to_ascii_uppercase().as_str(),
"CURRENT_PRINCIPAL"
| "CURRENT_EFFECTIVE_PRINCIPAL"
| "CURRENT_TRANSACTION_ID"
| "CURRENT_STATEMENT_TIMESTAMP"
| "CURRENT_REQUEST_ID"
| "CURRENT_IDEMPOTENCY_KEY"
| "CURRENT_JOB_ID"
| "CURRENT_JOB_ATTEMPT"
| "CURRENT_JOB_SCHEDULED_AT"
)
}
use crate::expression::RowFilter;
use smallvec::SmallVec;
type ScalarSubqueryCacheEntry = (SmallVec<[CompactArc<str>; 2]>, Value);
thread_local! {
static SCALAR_SUBQUERY_CACHE: RefCell<LruCache<String, ScalarSubqueryCacheEntry>> =
RefCell::new(LruCache::new(NonZeroUsize::new(SCALAR_SUBQUERY_CACHE_SIZE).unwrap()));
static EXISTS_PREDICATE_CACHE: RefCell<FxHashMap<String, RowFilter>> =
RefCell::new(FxHashMap::default());
static ACTIVE_QUERY_CANCELLATION: RefCell<Vec<CancellationHandle>> = const {
RefCell::new(Vec::new())
};
}
pub fn clear_exists_predicate_cache() {
EXISTS_PREDICATE_CACHE.with(|cache| cache.borrow_mut().clear());
}
pub fn get_cached_exists_predicate(key: &str) -> Option<RowFilter> {
EXISTS_PREDICATE_CACHE.with(|cache| cache.borrow().get(key).cloned())
}
pub fn cache_exists_predicate(key: String, filter: RowFilter) {
EXISTS_PREDICATE_CACHE.with(|cache| {
cache.borrow_mut().insert(key, filter);
});
}
#[doc(hidden)]
pub struct StatementCancellationScope;
impl Drop for StatementCancellationScope {
fn drop(&mut self) {
ACTIVE_QUERY_CANCELLATION.with(|stack| {
stack.borrow_mut().pop();
});
}
}
#[doc(hidden)]
pub fn current_query_is_cancelled() -> bool {
ACTIVE_QUERY_CANCELLATION.with(|stack| {
stack
.borrow()
.last()
.is_some_and(CancellationHandle::is_cancelled)
})
}
#[inline]
#[doc(hidden)]
pub fn with_current_query_cancellation<T>(
callback: impl FnOnce(Option<&dyn radixdb_functions::FunctionCancellation>) -> T,
) -> T {
ACTIVE_QUERY_CANCELLATION.with(|stack| {
let stack = stack.borrow();
callback(
stack
.last()
.map(|handle| handle as &dyn radixdb_functions::FunctionCancellation),
)
})
}
#[inline]
#[doc(hidden)]
pub fn check_current_query_cancelled() -> Result<()> {
if current_query_is_cancelled() {
Err(radixdb_core::Error::QueryCancelled)
} else {
Ok(())
}
}
pub fn clear_scalar_subquery_cache() {
SCALAR_SUBQUERY_CACHE.with(|cache| {
cache.borrow_mut().clear();
});
}
#[inline]
pub fn invalidate_scalar_subquery_cache_for_table(table_name: &str) {
SCALAR_SUBQUERY_CACHE.with(|cache| {
let mut c = cache.borrow_mut();
if c.is_empty() {
return;
}
let keys_to_remove: Vec<String> = c
.iter()
.filter(|(_, (tables, _))| tables.iter().any(|t| t.eq_ignore_ascii_case(table_name)))
.map(|(k, _)| k.clone())
.collect();
for key in keys_to_remove {
c.pop(&key);
}
});
}
pub fn get_cached_scalar_subquery(key: &str) -> Option<Value> {
SCALAR_SUBQUERY_CACHE.with(|cache| cache.borrow_mut().get(key).map(|(_, v)| v.clone()))
}
pub fn cache_scalar_subquery(key: String, tables: SmallVec<[CompactArc<str>; 2]>, value: Value) {
SCALAR_SUBQUERY_CACHE.with(|cache| {
cache.borrow_mut().put(key, (tables, value));
});
}
type InSubqueryCacheEntry = (SmallVec<[CompactArc<str>; 2]>, Vec<Value>);
thread_local! {
static IN_SUBQUERY_CACHE: RefCell<LruCache<String, InSubqueryCacheEntry>> =
RefCell::new(LruCache::new(NonZeroUsize::new(IN_SUBQUERY_CACHE_SIZE).unwrap()));
}
pub fn clear_in_subquery_cache() {
IN_SUBQUERY_CACHE.with(|cache| {
cache.borrow_mut().clear();
});
}
#[inline]
pub fn invalidate_in_subquery_cache_for_table(table_name: &str) {
IN_SUBQUERY_CACHE.with(|cache| {
let mut c = cache.borrow_mut();
if c.is_empty() {
return;
}
let keys_to_remove: Vec<String> = c
.iter()
.filter(|(_, (tables, _))| tables.iter().any(|t| t.eq_ignore_ascii_case(table_name)))
.map(|(k, _)| k.clone())
.collect();
for key in keys_to_remove {
c.pop(&key);
}
});
}
pub fn get_cached_in_subquery(key: &str) -> Option<Vec<Value>> {
IN_SUBQUERY_CACHE.with(|cache| cache.borrow_mut().get(key).map(|(_, v)| v.clone()))
}
pub fn cache_in_subquery(key: String, tables: SmallVec<[CompactArc<str>; 2]>, values: Vec<Value>) {
IN_SUBQUERY_CACHE.with(|cache| {
cache.borrow_mut().put(key, (tables, values));
});
}
use radixdb_sql::ast::{Expression, SelectStatement};
pub fn extract_table_names_for_cache(stmt: &SelectStatement) -> SmallVec<[CompactArc<str>; 2]> {
let mut tables = SmallVec::new();
if let Some(ref table_expr) = stmt.table_expr {
collect_real_table_names(table_expr, &mut tables);
}
tables
}
fn collect_real_table_names(source: &Expression, tables: &mut SmallVec<[CompactArc<str>; 2]>) {
match source {
Expression::TableSource(ts) => {
tables.push(CompactArc::from(ts.name.value_lower.as_str()));
}
Expression::JoinSource(js) => {
collect_real_table_names(&js.left, tables);
collect_real_table_names(&js.right, tables);
}
Expression::SubquerySource(ss) => {
if let Some(ref table_expr) = ss.subquery.table_expr {
collect_real_table_names(table_expr, tables);
}
}
_ => {}
}
}
use ahash::AHashMap;
use std::hash::{Hash, Hasher};
type SemiJoinCacheEntry = (CompactArc<str>, CompactArc<ValueSet>);
#[inline]
pub fn compute_semi_join_cache_key(table: &str, column: &str, pred_hash: u64) -> u64 {
let mut hasher = rustc_hash::FxHasher::default();
table.hash(&mut hasher);
column.hash(&mut hasher);
pred_hash.hash(&mut hasher);
hasher.finish()
}
thread_local! {
static SEMI_JOIN_CACHE: RefCell<LruCache<u64, SemiJoinCacheEntry>> =
RefCell::new(LruCache::new(NonZeroUsize::new(SEMI_JOIN_CACHE_SIZE).unwrap()));
}
pub fn clear_semi_join_cache() {
SEMI_JOIN_CACHE.with(|cache| {
cache.borrow_mut().clear();
});
}
#[inline]
pub fn invalidate_semi_join_cache_for_table(table_name: &str) {
SEMI_JOIN_CACHE.with(|cache| {
let mut c = cache.borrow_mut();
if c.is_empty() {
return;
}
let keys_to_remove: Vec<u64> = c
.iter()
.filter(|(_, (key_table, _))| key_table.eq_ignore_ascii_case(table_name))
.map(|(k, _)| *k)
.collect();
for key in keys_to_remove {
c.pop(&key);
}
});
}
#[inline]
pub fn get_cached_semi_join(key_hash: u64) -> Option<CompactArc<ValueSet>> {
SEMI_JOIN_CACHE.with(|cache| {
cache
.borrow_mut()
.get(&key_hash)
.map(|(_, v)| CompactArc::clone(v))
})
}
#[inline]
pub fn cache_semi_join_arc(key_hash: u64, table: &str, values: CompactArc<ValueSet>) {
SEMI_JOIN_CACHE.with(|cache| {
cache
.borrow_mut()
.put(key_hash, (CompactArc::from(table), values));
});
}
use radixdb_storage::traits::Index;
thread_local! {
static EXISTS_INDEX_CACHE: RefCell<FxHashMap<String, std::sync::Arc<dyn Index>>> = RefCell::new(FxHashMap::default());
}
pub fn clear_exists_index_cache() {
EXISTS_INDEX_CACHE.with(|cache| {
cache.borrow_mut().clear();
});
}
pub fn get_cached_exists_index(key: &str) -> Option<std::sync::Arc<dyn Index>> {
EXISTS_INDEX_CACHE.with(|cache| cache.borrow().get(key).cloned())
}
pub fn cache_exists_index(key: String, index: std::sync::Arc<dyn Index>) {
EXISTS_INDEX_CACHE.with(|cache| {
cache.borrow_mut().insert(key, index);
});
}
pub type RowFetcher =
Box<dyn Fn(&[i64]) -> radixdb_core::Result<radixdb_core::RowVec> + Send + Sync>;
pub type RowCounter = Box<dyn Fn(&[i64]) -> usize + Send + Sync>;
thread_local! {
static EXISTS_FETCHER_CACHE: RefCell<FxHashMap<String, std::sync::Arc<RowFetcher>>> = RefCell::new(FxHashMap::default());
}
thread_local! {
static COUNT_COUNTER_CACHE: RefCell<FxHashMap<String, std::sync::Arc<RowCounter>>> = RefCell::new(FxHashMap::default());
}
pub fn clear_exists_fetcher_cache() {
EXISTS_FETCHER_CACHE.with(|cache| {
cache.borrow_mut().clear();
});
}
pub fn clear_count_counter_cache() {
COUNT_COUNTER_CACHE.with(|cache| {
cache.borrow_mut().clear();
});
}
pub fn get_cached_exists_fetcher(key: &str) -> Option<std::sync::Arc<RowFetcher>> {
EXISTS_FETCHER_CACHE.with(|cache| cache.borrow().get(key).cloned())
}
pub fn get_cached_count_counter(key: &str) -> Option<std::sync::Arc<RowCounter>> {
COUNT_COUNTER_CACHE.with(|cache| cache.borrow().get(key).cloned())
}
pub fn cache_exists_fetcher(key: String, fetcher: RowFetcher) {
EXISTS_FETCHER_CACHE.with(|cache| {
cache.borrow_mut().insert(key, std::sync::Arc::new(fetcher));
});
}
pub fn cache_count_counter(key: String, counter: RowCounter) {
COUNT_COUNTER_CACHE.with(|cache| {
cache.borrow_mut().insert(key, std::sync::Arc::new(counter));
});
}
thread_local! {
static EXISTS_SCHEMA_CACHE: RefCell<FxHashMap<String, CompactArc<Vec<String>>>> = RefCell::new(FxHashMap::default());
}
pub fn clear_exists_schema_cache() {
EXISTS_SCHEMA_CACHE.with(|cache| {
cache.borrow_mut().clear();
});
}
pub fn get_cached_exists_schema(key: &str) -> Option<CompactArc<Vec<String>>> {
EXISTS_SCHEMA_CACHE.with(|cache| cache.borrow().get(key).cloned())
}
pub fn cache_exists_schema(key: String, columns: CompactArc<Vec<String>>) {
EXISTS_SCHEMA_CACHE.with(|cache| {
cache.borrow_mut().insert(key, columns);
});
}
thread_local! {
static EXISTS_PRED_KEY_CACHE: RefCell<FxHashMap<usize, String>> = RefCell::new(FxHashMap::default());
}
pub fn clear_exists_pred_key_cache() {
EXISTS_PRED_KEY_CACHE.with(|cache| {
cache.borrow_mut().clear();
});
}
#[inline]
pub fn get_cached_exists_pred_key(subquery_ptr: usize) -> Option<String> {
EXISTS_PRED_KEY_CACHE.with(|cache| cache.borrow().get(&subquery_ptr).cloned())
}
#[inline]
pub fn cache_exists_pred_key(subquery_ptr: usize, pred_key: String) {
EXISTS_PRED_KEY_CACHE.with(|cache| {
cache.borrow_mut().insert(subquery_ptr, pred_key);
});
}
thread_local! {
static BATCH_AGGREGATE_CACHE: RefCell<FxHashMap<String, CompactArc<ValueMap<Value>>>> = RefCell::new(FxHashMap::default());
}
pub fn clear_batch_aggregate_cache() {
BATCH_AGGREGATE_CACHE.with(|cache| {
let mut c = cache.borrow_mut();
c.clear();
c.shrink_to_fit();
});
}
pub fn get_cached_batch_aggregate(key: &str) -> Option<CompactArc<ValueMap<Value>>> {
BATCH_AGGREGATE_CACHE.with(|cache| cache.borrow().get(key).cloned())
}
pub fn cache_batch_aggregate(key: String, values: ValueMap<Value>) {
BATCH_AGGREGATE_CACHE.with(|cache| {
cache.borrow_mut().insert(key, CompactArc::new(values));
});
}
#[derive(Clone)]
pub struct BatchAggregateLookupInfo {
pub cache_key: String,
pub outer_column_lower: String,
pub outer_qualified_lower: Option<String>,
pub is_count: bool,
}
thread_local! {
static BATCH_AGGREGATE_INFO_CACHE: RefCell<FxHashMap<usize, Option<Arc<BatchAggregateLookupInfo>>>> = RefCell::new(FxHashMap::default());
}
pub fn clear_batch_aggregate_info_cache() {
BATCH_AGGREGATE_INFO_CACHE.with(|cache| {
let mut c = cache.borrow_mut();
c.clear();
c.shrink_to_fit();
});
}
#[inline]
pub fn get_cached_batch_aggregate_info(
subquery_ptr: usize,
) -> Option<Option<Arc<BatchAggregateLookupInfo>>> {
BATCH_AGGREGATE_INFO_CACHE.with(|cache| cache.borrow().get(&subquery_ptr).cloned())
}
#[inline]
pub fn cache_batch_aggregate_info(
subquery_ptr: usize,
info: Option<BatchAggregateLookupInfo>,
) -> Option<Arc<BatchAggregateLookupInfo>> {
let arc_info = info.map(Arc::new);
let result = arc_info.clone();
BATCH_AGGREGATE_INFO_CACHE.with(|cache| {
cache.borrow_mut().insert(subquery_ptr, arc_info);
});
result
}
#[derive(Clone)]
pub struct ExistsCorrelationInfo {
pub outer_column: String,
pub outer_table: Option<String>,
pub inner_column: String,
pub inner_table: String,
pub outer_column_lower: String,
pub outer_qualified_lower: Option<String>,
pub additional_predicate: Option<Expression>,
pub index_cache_key: String,
}
thread_local! {
static EXISTS_CORRELATION_CACHE: RefCell<FxHashMap<usize, Option<Arc<ExistsCorrelationInfo>>>> = RefCell::new(FxHashMap::default());
}
pub fn clear_exists_correlation_cache() {
EXISTS_CORRELATION_CACHE.with(|cache| {
cache.borrow_mut().clear();
});
}
pub fn clear_executor_thread_local_caches() {
SCALAR_SUBQUERY_CACHE.with(|cache| {
cache.borrow_mut().clear();
});
IN_SUBQUERY_CACHE.with(|cache| {
cache.borrow_mut().clear();
});
SEMI_JOIN_CACHE.with(|cache| {
cache.borrow_mut().clear();
});
EXISTS_INDEX_CACHE.with(|cache| {
let mut c = cache.borrow_mut();
c.clear();
c.shrink_to_fit();
});
EXISTS_FETCHER_CACHE.with(|cache| {
let mut c = cache.borrow_mut();
c.clear();
c.shrink_to_fit();
});
COUNT_COUNTER_CACHE.with(|cache| {
let mut c = cache.borrow_mut();
c.clear();
c.shrink_to_fit();
});
EXISTS_SCHEMA_CACHE.with(|cache| {
let mut c = cache.borrow_mut();
c.clear();
c.shrink_to_fit();
});
EXISTS_PRED_KEY_CACHE.with(|cache| {
let mut c = cache.borrow_mut();
c.clear();
c.shrink_to_fit();
});
BATCH_AGGREGATE_CACHE.with(|cache| {
let mut c = cache.borrow_mut();
c.clear();
c.shrink_to_fit();
});
BATCH_AGGREGATE_INFO_CACHE.with(|cache| {
let mut c = cache.borrow_mut();
c.clear();
c.shrink_to_fit();
});
EXISTS_CORRELATION_CACHE.with(|cache| {
let mut c = cache.borrow_mut();
c.clear();
c.shrink_to_fit();
});
radixdb_storage::expression::clear_regex_cache();
radixdb_storage::expression::clear_like_regex_cache();
radixdb_core::row_vec::clear_row_vec_pool();
radixdb_core::row_vec::clear_row_id_vec_pool();
radixdb_storage::mvcc::clear_version_map_pools();
}
pub fn clear_all_thread_local_caches() {
clear_executor_thread_local_caches();
clear_exists_predicate_cache();
crate::expression::clear_program_cache();
crate::query::clear_join_dependency_projection_cache();
crate::utils::clear_join_projection_lookup_cache();
crate::query_classification::clear_classification_cache();
}
#[inline]
pub fn get_cached_exists_correlation(
subquery_ptr: usize,
) -> Option<Option<Arc<ExistsCorrelationInfo>>> {
EXISTS_CORRELATION_CACHE.with(|cache| cache.borrow().get(&subquery_ptr).cloned())
}
#[inline]
pub fn cache_exists_correlation(
subquery_ptr: usize,
info: Option<ExistsCorrelationInfo>,
) -> Option<Arc<ExistsCorrelationInfo>> {
let arc_info = info.map(Arc::new);
let result = arc_info.clone();
EXISTS_CORRELATION_CACHE.with(|cache| {
cache.borrow_mut().insert(subquery_ptr, arc_info);
});
result
}
#[derive(Debug, Clone)]
pub struct ExecutionContext {
session: SessionState,
query: QueryState,
}
#[derive(Debug, Clone)]
struct SessionState {
auto_commit: bool,
current_database: Arc<Option<String>>,
session_vars: Arc<AHashMap<String, Value>>,
principal_id: ObjectId,
}
#[derive(Debug, Clone)]
struct QueryState {
effective_principal_id: Option<ObjectId>,
params: CompactArc<ParamVec>,
named_params: Arc<FxHashMap<String, Value>>,
cancelled: Arc<AtomicBool>,
timed_out: Arc<AtomicBool>,
active_reference_expands: Arc<AtomicUsize>,
parent_cancelled: Option<Arc<AtomicBool>>,
timeout_ms: u64,
join_hash_state_max_bytes: usize,
join_hash_states: Arc<Mutex<JoinHashStateCache>>,
join_memory_owner: Arc<JoinMemoryOwner>,
view_depth: usize,
query_depth: usize,
outer_row: Option<FxHashMap<CompactArc<str>, Value>>,
outer_columns: Option<CompactArc<Vec<String>>>,
cte_data: Option<Arc<CteDataMap>>,
transaction_id: Option<u64>,
stored_function_invoker: Option<Arc<dyn StoredFunctionInvoker>>,
procedural_budget: Option<BudgetOwner>,
public_scan_budget: Option<Arc<PublicScanBudget>>,
}
#[derive(Debug)]
struct PublicScanBudget {
remaining: AtomicUsize,
}
impl PublicScanBudget {
fn new(max_rows: usize) -> Self {
Self {
remaining: AtomicUsize::new(max_rows),
}
}
fn claim(&self, rows: usize) -> Result<()> {
self.remaining
.fetch_update(Ordering::AcqRel, Ordering::Acquire, |remaining| {
remaining.checked_sub(rows)
})
.map(|_| ())
.map_err(|_| {
radixdb_core::Error::invalid_argument("public read scanned-row budget exceeded")
})
}
}
#[doc(hidden)]
pub type CteMaterializedRows = Arc<std::sync::OnceLock<CompactArc<Vec<Row>>>>;
#[doc(hidden)]
pub type CteData = (
CompactArc<Vec<String>>,
CompactArc<Vec<(i64, Row)>>,
CteMaterializedRows,
);
#[doc(hidden)]
pub type CteDataMap = StringMap<CteData>;
#[derive(Default)]
struct JoinHashStateCache {
states: Vec<JoinHashState>,
}
impl std::fmt::Debug for JoinHashStateCache {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("JoinHashStateCache")
.field("states", &self.states.len())
.finish()
}
}
impl JoinHashStateCache {
fn get(
&mut self,
build_rows: &CompactArc<Vec<Row>>,
key_indices: &[usize],
) -> Option<JoinHashState> {
let index = self
.states
.iter()
.position(|state| state.matches(build_rows, key_indices))?;
let state = self.states.remove(index);
let result = state.clone();
self.states.push(state);
Some(result)
}
fn insert(&mut self, state: JoinHashState) {
self.states.push(state);
}
fn pop_lru(&mut self) -> Option<JoinHashState> {
(!self.states.is_empty()).then(|| self.states.remove(0))
}
}
impl Default for ExecutionContext {
fn default() -> Self {
Self::new()
}
}
impl ExecutionContext {
pub fn new() -> Self {
let mut system_values = FxHashMap::default();
system_values.insert(
"CURRENT_PRINCIPAL".to_owned(),
Value::uuid(ObjectId::BOOTSTRAP_OWNER.into_bytes()),
);
system_values.insert(
"CURRENT_EFFECTIVE_PRINCIPAL".to_owned(),
Value::uuid(ObjectId::BOOTSTRAP_OWNER.into_bytes()),
);
system_values.insert(
"CURRENT_STATEMENT_TIMESTAMP".to_owned(),
Value::timestamp(system_time_now().into()),
);
Self {
session: SessionState {
auto_commit: true,
current_database: EMPTY_DATABASE.clone(),
session_vars: EMPTY_SESSION_VARS.clone(),
principal_id: ObjectId::BOOTSTRAP_OWNER,
},
query: QueryState {
effective_principal_id: None,
params: EMPTY_PARAMS.clone(),
named_params: Arc::new(system_values),
cancelled: Arc::new(AtomicBool::new(false)),
timed_out: Arc::new(AtomicBool::new(false)),
active_reference_expands: Arc::new(AtomicUsize::new(0)),
parent_cancelled: None,
timeout_ms: 0,
join_hash_state_max_bytes: crate::hash_table::DEFAULT_JOIN_HASH_STATE_MAX_BYTES,
join_hash_states: Arc::new(Mutex::new(JoinHashStateCache::default())),
join_memory_owner: Arc::new(JoinMemoryOwner::default()),
view_depth: 0,
query_depth: 0,
outer_row: None,
outer_columns: None,
cte_data: None,
transaction_id: None,
stored_function_invoker: None,
procedural_budget: None,
public_scan_budget: None,
},
}
}
#[doc(hidden)]
pub fn enter_statement_scope(&self) -> StatementCancellationScope {
ACTIVE_QUERY_CANCELLATION.with(|stack| {
stack.borrow_mut().push(self.cancellation_handle());
});
StatementCancellationScope
}
#[doc(hidden)]
pub fn enter_reference_expand(&self) -> ReferenceExpandExecutionGuard {
self.query
.active_reference_expands
.fetch_add(1, Ordering::AcqRel);
ReferenceExpandExecutionGuard {
active: Arc::clone(&self.query.active_reference_expands),
}
}
#[doc(hidden)]
pub fn active_reference_expands(&self) -> usize {
self.query.active_reference_expands.load(Ordering::Acquire)
}
pub fn with_params(params: ParamVec) -> Self {
let mut context = Self::new();
context.query.params = CompactArc::new(params);
context
}
pub fn with_named_params(named_params: FxHashMap<String, Value>) -> Self {
let mut context = Self::new();
let system_values = context.query.named_params.clone();
let mut admitted = named_params;
admitted.retain(|name, _| !is_system_context_name(name));
admitted.extend(
system_values
.iter()
.map(|(name, value)| (name.clone(), value.clone())),
);
context.query.named_params = Arc::new(admitted);
context
}
pub fn get_param(&self, index: usize) -> Option<&Value> {
if index == 0 || index > self.query.params.len() {
None
} else {
self.query.params.get(index - 1)
}
}
pub fn get_named_param(&self, name: &str) -> Option<&Value> {
self.query.named_params.get(name)
}
#[inline]
#[doc(hidden)]
pub fn join_hash_state_max_bytes(&self) -> usize {
self.query.join_hash_state_max_bytes
}
#[doc(hidden)]
pub fn reserve_join_memory(&self, bytes: usize) -> Option<JoinMemoryReservation> {
loop {
if let Some(reservation) = self
.query
.join_memory_owner
.try_reserve(bytes, self.query.join_hash_state_max_bytes)
{
return Some(reservation);
}
let evicted = self
.query
.join_hash_states
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.pop_lru();
let evicted = evicted?;
drop(evicted);
}
}
#[doc(hidden)]
pub fn retained_join_memory_bytes(&self) -> usize {
self.query.join_memory_owner.retained_bytes()
}
#[doc(hidden)]
pub fn peak_join_memory_bytes(&self) -> usize {
self.query.join_memory_owner.peak_bytes()
}
#[doc(hidden)]
pub fn join_hash_state_for(
&self,
build_rows: CompactArc<Vec<Row>>,
key_indices: &[usize],
retain_for_reuse: bool,
) -> Option<JoinHashState> {
if key_indices.is_empty() {
return None;
}
if retain_for_reuse {
let mut cache = self
.query
.join_hash_states
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if let Some(state) = cache.get(&build_rows, key_indices) {
radixdb_storage::instrumentation::record_join_hash_state_reuse();
return Some(state);
}
}
let state_bytes = JoinHashTable::estimated_retained_bytes(build_rows.len())?;
let reservation = self.reserve_join_memory(state_bytes)?;
let state = JoinHashState::build_reserved(build_rows, key_indices, reservation);
radixdb_storage::instrumentation::record_join_hash_state_build();
if retain_for_reuse {
let mut cache = self
.query
.join_hash_states
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if let Some(existing) = cache.get(state.build_rows(), key_indices) {
radixdb_storage::instrumentation::record_join_hash_state_reuse();
return Some(existing);
}
cache.insert(state.clone());
}
Some(state)
}
#[doc(hidden)]
pub fn join_hash_state_with_bloom_for(
&self,
build_rows: CompactArc<Vec<Row>>,
key_indices: &[usize],
bloom_builder: &mut impl crate::hash_table::JoinHashObserver,
) -> Option<JoinHashState> {
if key_indices.is_empty() {
return None;
}
let state_bytes = JoinHashTable::estimated_retained_bytes(build_rows.len())?
.checked_add(bloom_builder.retained_bytes())?;
let reservation = self.reserve_join_memory(state_bytes)?;
let state = JoinHashState::build_with_bloom_reserved(
build_rows,
key_indices,
bloom_builder,
reservation,
);
radixdb_storage::instrumentation::record_join_hash_state_build();
Some(state)
}
pub fn params(&self) -> &[Value] {
&self.query.params
}
pub fn params_arc(&self) -> &CompactArc<ParamVec> {
&self.query.params
}
pub fn named_params(&self) -> &FxHashMap<String, Value> {
&self.query.named_params
}
pub fn named_params_arc(&self) -> &Arc<FxHashMap<String, Value>> {
&self.query.named_params
}
#[doc(hidden)]
pub(crate) fn stored_function_invoker(&self) -> Option<&Arc<dyn StoredFunctionInvoker>> {
self.query.stored_function_invoker.as_ref()
}
#[doc(hidden)]
pub(crate) fn with_stored_function_invoker(
mut self,
invoker: Arc<dyn StoredFunctionInvoker>,
) -> Self {
self.query.stored_function_invoker = Some(invoker);
self
}
#[doc(hidden)]
pub(crate) fn procedural_budget(&self) -> Option<&BudgetOwner> {
self.query.procedural_budget.as_ref()
}
#[doc(hidden)]
pub(crate) fn with_procedural_budget(mut self, budget: BudgetOwner) -> Self {
self.query.procedural_budget = Some(budget);
self
}
pub fn param_count(&self) -> usize {
self.query.params.len()
}
pub fn set_params(&mut self, params: ParamVec) {
self.query.params = CompactArc::new(params);
}
pub fn add_param(&mut self, value: Value) {
CompactArc::make_mut(&mut self.query.params).push(value);
}
pub fn set_named_param(&mut self, name: impl Into<String>, value: Value) {
let name = name.into();
if !is_system_context_name(&name) {
Arc::make_mut(&mut self.query.named_params).insert(name, value);
}
}
pub fn auto_commit(&self) -> bool {
self.session.auto_commit
}
pub fn set_auto_commit(&mut self, auto_commit: bool) {
self.session.auto_commit = auto_commit;
}
pub fn is_cancelled(&self) -> bool {
self.query.cancelled.load(Ordering::Relaxed)
|| self
.query
.parent_cancelled
.as_ref()
.is_some_and(|cancelled| cancelled.load(Ordering::Relaxed))
}
#[doc(hidden)]
pub fn did_time_out(&self) -> bool {
self.query.timed_out.load(Ordering::Acquire)
}
pub fn cancel(&self) {
self.query.cancelled.store(true, Ordering::Relaxed);
}
pub fn cancellation_handle(&self) -> CancellationHandle {
CancellationHandle {
cancelled: self.query.cancelled.clone(),
parent_cancelled: self.query.parent_cancelled.clone(),
}
}
#[doc(hidden)]
pub fn bind_parent_cancellation(&mut self, handle: &CancellationHandle) {
self.query.parent_cancelled = Some(Arc::clone(&handle.cancelled));
}
pub fn current_database(&self) -> Option<&str> {
self.session.current_database.as_ref().as_deref()
}
pub fn set_current_database(&mut self, database: impl Into<String>) {
self.session.current_database = Arc::new(Some(database.into()));
}
pub const fn principal_id(&self) -> ObjectId {
self.session.principal_id
}
pub fn with_principal_id(&self, principal_id: ObjectId) -> Self {
let mut nested = self.clone();
nested.session.principal_id = principal_id;
nested.query.effective_principal_id = None;
nested
.set_system_context_value("CURRENT_PRINCIPAL", Value::uuid(principal_id.into_bytes()));
nested.set_system_context_value(
"CURRENT_EFFECTIVE_PRINCIPAL",
Value::uuid(principal_id.into_bytes()),
);
nested
}
pub const fn effective_principal_id(&self) -> ObjectId {
match self.query.effective_principal_id {
Some(principal) => principal,
None => self.session.principal_id,
}
}
pub fn with_effective_principal_id(&self, principal_id: ObjectId) -> Self {
let mut nested = self.clone();
nested.query.effective_principal_id = Some(principal_id);
nested.set_system_context_value(
"CURRENT_EFFECTIVE_PRINCIPAL",
Value::uuid(principal_id.into_bytes()),
);
nested
}
#[doc(hidden)]
pub fn with_public_scan_limit(&self, max_rows: usize) -> Self {
let mut nested = self.clone();
nested.query.public_scan_budget = Some(Arc::new(PublicScanBudget::new(max_rows)));
nested
}
#[doc(hidden)]
pub fn claim_public_scan_rows(&self, rows: usize) -> Result<()> {
self.query
.public_scan_budget
.as_ref()
.map_or(Ok(()), |budget| budget.claim(rows))
}
#[doc(hidden)]
pub fn has_public_scan_budget(&self) -> bool {
self.query.public_scan_budget.is_some()
}
pub fn get_session_var(&self, name: &str) -> Option<&Value> {
self.session.session_vars.get(name)
}
pub fn set_session_var(&mut self, name: impl Into<String>, value: Value) {
Arc::make_mut(&mut self.session.session_vars).insert(name.into(), value);
}
pub fn timeout_ms(&self) -> u64 {
self.query.timeout_ms
}
pub fn set_timeout_ms(&mut self, timeout_ms: u64) {
self.query.timeout_ms = timeout_ms;
}
pub fn has_timeout(&self) -> bool {
self.query.timeout_ms > 0
}
pub fn view_depth(&self) -> usize {
self.query.view_depth
}
pub fn with_incremented_view_depth(&self) -> Self {
let mut nested = self.clone();
nested.query.view_depth += 1;
nested.query.query_depth += 1;
nested
}
pub fn with_incremented_query_depth(&self) -> Self {
let mut nested = self.clone();
nested.query.query_depth += 1;
nested
}
pub fn outer_row(&self) -> Option<&FxHashMap<CompactArc<str>, Value>> {
self.query.outer_row.as_ref()
}
#[doc(hidden)]
pub fn take_outer_row(&mut self) -> Option<FxHashMap<CompactArc<str>, Value>> {
self.query.outer_row.take()
}
#[doc(hidden)]
pub fn query_depth(&self) -> usize {
self.query.query_depth
}
pub fn outer_columns(&self) -> Option<&[String]> {
self.query.outer_columns.as_ref().map(|v| v.as_slice())
}
pub fn with_outer_row(
&self,
outer_row: FxHashMap<CompactArc<str>, Value>,
outer_columns: CompactArc<Vec<String>>,
) -> Self {
let mut nested = self.with_incremented_query_depth();
nested.query.outer_row = Some(outer_row);
nested.query.outer_columns = Some(outer_columns);
nested
}
pub fn get_cte(&self, name: &str) -> Option<&CteData> {
self.query
.cte_data
.as_ref()
.and_then(|data| data.get(&name.to_lowercase()))
}
#[inline]
pub fn get_cte_by_lower(&self, name_lower: &str) -> Option<&CteData> {
self.query
.cte_data
.as_ref()
.and_then(|data| data.get(name_lower))
}
#[doc(hidden)]
pub fn get_cte_materialized_rows_by_lower(
&self,
name_lower: &str,
) -> Option<CompactArc<Vec<Row>>> {
let (_, rows_with_ids, materialized) = self.get_cte_by_lower(name_lower)?;
Some(
materialized
.get_or_init(|| {
CompactArc::new(rows_with_ids.iter().map(|(_, row)| row.clone()).collect())
})
.clone(),
)
}
pub fn has_cte(&self, name: &str) -> bool {
self.query
.cte_data
.as_ref()
.is_some_and(|data| data.contains_key(&name.to_lowercase()))
}
#[inline]
pub fn has_cte_by_lower(&self, name_lower: &str) -> bool {
self.query
.cte_data
.as_ref()
.is_some_and(|data| data.contains_key(name_lower))
}
pub fn with_cte_data(&self, cte_data: Arc<CteDataMap>) -> Self {
let mut nested = self.clone();
nested.query.cte_data = Some(cte_data);
nested
}
pub fn transaction_id(&self) -> Option<u64> {
self.query.transaction_id
}
pub fn set_transaction_id(&mut self, txn_id: u64) {
self.query.transaction_id = Some(txn_id);
if let Ok(txn_id) = i64::try_from(txn_id) {
self.set_system_context_value("CURRENT_TRANSACTION_ID", Value::Integer(txn_id));
}
}
pub fn with_transaction_id(&self, txn_id: u64) -> Self {
let mut nested = self.clone();
nested.set_transaction_id(txn_id);
nested
}
#[doc(hidden)]
pub fn set_request_id(&mut self, request_id: u64) -> Result<()> {
let request_id = i64::try_from(request_id).map_err(|_| {
radixdb_core::Error::invalid_argument("request ID exceeds the SQL INTEGER domain")
})?;
self.set_system_context_value("CURRENT_REQUEST_ID", Value::Integer(request_id));
Ok(())
}
#[doc(hidden)]
pub fn set_job_context(
&mut self,
idempotency_key: &str,
job_id: ObjectId,
attempt: u32,
scheduled_at: DateTime<Utc>,
) {
self.set_system_context_value("CURRENT_IDEMPOTENCY_KEY", Value::text(idempotency_key));
self.set_system_context_value("CURRENT_JOB_ID", Value::uuid(job_id.into_bytes()));
self.set_system_context_value("CURRENT_JOB_ATTEMPT", Value::Integer(i64::from(attempt)));
self.set_system_context_value("CURRENT_JOB_SCHEDULED_AT", Value::Timestamp(scheduled_at));
}
fn set_system_context_value(&mut self, name: &str, value: Value) {
Arc::make_mut(&mut self.query.named_params).insert(name.to_owned(), value);
}
pub fn check_cancelled(&self) -> Result<()> {
if self.is_cancelled() {
Err(radixdb_core::Error::QueryCancelled)
} else {
Ok(())
}
}
}
#[doc(hidden)]
pub struct ReferenceExpandExecutionGuard {
active: Arc<AtomicUsize>,
}
impl Drop for ReferenceExpandExecutionGuard {
fn drop(&mut self) {
self.active.fetch_sub(1, Ordering::AcqRel);
}
}
#[derive(Debug, Clone)]
pub struct CancellationHandle {
cancelled: Arc<AtomicBool>,
parent_cancelled: Option<Arc<AtomicBool>>,
}
impl CancellationHandle {
pub fn cancel(&self) {
self.cancelled.store(true, Ordering::Relaxed);
}
pub fn is_cancelled(&self) -> bool {
self.cancelled.load(Ordering::Relaxed)
|| self
.parent_cancelled
.as_ref()
.is_some_and(|cancelled| cancelled.load(Ordering::Relaxed))
}
}
impl radixdb_functions::FunctionCancellation for CancellationHandle {
#[inline]
fn is_cancelled(&self) -> bool {
Self::is_cancelled(self)
}
}
struct TimeoutEntry {
deadline: Instant,
id: u64,
cancel_handle: CancellationHandle,
cancelled: Arc<AtomicBool>,
timed_out: Arc<AtomicBool>,
}
impl PartialEq for TimeoutEntry {
fn eq(&self, other: &Self) -> bool {
self.deadline == other.deadline && self.id == other.id
}
}
impl Eq for TimeoutEntry {}
impl PartialOrd for TimeoutEntry {
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
Some(self.cmp(other))
}
}
impl Ord for TimeoutEntry {
fn cmp(&self, other: &Self) -> std::cmp::Ordering {
other.deadline.cmp(&self.deadline)
}
}
struct TimeoutManagerState {
timeouts: BinaryHeap<TimeoutEntry>,
}
struct TimeoutManager {
state: Mutex<TimeoutManagerState>,
condvar: Condvar,
next_id: AtomicU64,
}
impl TimeoutManager {
fn new() -> Arc<Self> {
let manager = Arc::new(Self {
state: Mutex::new(TimeoutManagerState {
timeouts: BinaryHeap::new(),
}),
condvar: Condvar::new(),
next_id: AtomicU64::new(1),
});
let manager_clone = Arc::clone(&manager);
std::thread::Builder::new()
.name("radixdb-timeout-manager".to_string())
.spawn(move || {
manager_clone.run();
})
.expect("Failed to spawn timeout manager thread");
manager
}
fn run(&self) {
loop {
let mut state = self.state.lock().unwrap();
let now = Instant::now();
while let Some(entry) = state.timeouts.peek() {
if entry.deadline <= now {
let entry = state.timeouts.pop().unwrap();
if !entry.cancelled.load(Ordering::Relaxed) {
entry.timed_out.store(true, Ordering::Release);
entry.cancel_handle.cancel();
}
} else {
break;
}
}
let wait_duration = if let Some(entry) = state.timeouts.peek() {
entry.deadline.saturating_duration_since(now)
} else {
Duration::from_secs(3600) };
if wait_duration.is_zero() {
continue; }
let (_state, _timeout_result) =
self.condvar.wait_timeout(state, wait_duration).unwrap();
}
}
fn register(
&self,
timeout_ms: u64,
cancel_handle: CancellationHandle,
cancelled: Arc<AtomicBool>,
timed_out: Arc<AtomicBool>,
) -> u64 {
let id = self.next_id.fetch_add(1, Ordering::Relaxed);
let deadline = Instant::now() + Duration::from_millis(timeout_ms);
let entry = TimeoutEntry {
deadline,
id,
cancel_handle,
cancelled,
timed_out,
};
let mut state = self.state.lock().unwrap();
let was_empty = state.timeouts.is_empty();
let is_earliest = state.timeouts.peek().is_none_or(|e| deadline < e.deadline);
state.timeouts.push(entry);
if was_empty || is_earliest {
self.condvar.notify_one();
}
id
}
fn unregister(&self, id: u64) {
let mut state = self.state.lock().unwrap();
let removed_earliest = state.timeouts.peek().is_some_and(|entry| entry.id == id);
state.timeouts.retain(|entry| entry.id != id);
drop(state);
if removed_earliest {
self.condvar.notify_one();
}
}
}
fn global_timeout_manager() -> &'static Arc<TimeoutManager> {
use std::sync::OnceLock;
static MANAGER: OnceLock<Arc<TimeoutManager>> = OnceLock::new();
MANAGER.get_or_init(TimeoutManager::new)
}
#[doc(hidden)]
pub fn pending_timeout_count_for(ctx: &ExecutionContext) -> usize {
global_timeout_manager()
.state
.lock()
.unwrap()
.timeouts
.iter()
.filter(|entry| Arc::ptr_eq(&entry.timed_out, &ctx.query.timed_out))
.count()
}
pub struct TimeoutGuard {
registration_id: u64,
cancelled: Arc<AtomicBool>,
}
impl TimeoutGuard {
pub fn new(ctx: &ExecutionContext) -> Option<Self> {
let timeout_ms = ctx.timeout_ms();
if timeout_ms == 0 {
return None;
}
let cancel_handle = ctx.cancellation_handle();
let cancelled = Arc::new(AtomicBool::new(false));
let registration_id = global_timeout_manager().register(
timeout_ms,
cancel_handle,
Arc::clone(&cancelled),
Arc::clone(&ctx.query.timed_out),
);
Some(Self {
registration_id,
cancelled,
})
}
}
impl Drop for TimeoutGuard {
fn drop(&mut self) {
self.cancelled.store(true, Ordering::Relaxed);
global_timeout_manager().unregister(self.registration_id);
}
}
pub struct ExecutionContextBuilder {
ctx: ExecutionContext,
}
impl ExecutionContextBuilder {
pub fn new() -> Self {
Self {
ctx: ExecutionContext::new(),
}
}
pub fn params(mut self, params: ParamVec) -> Self {
self.ctx.query.params = CompactArc::new(params);
self
}
pub fn param(mut self, value: Value) -> Self {
let mut v = (*self.ctx.query.params).clone();
v.push(value);
self.ctx.query.params = CompactArc::new(v);
self
}
pub fn named_param(mut self, name: impl Into<String>, value: Value) -> Self {
self.ctx.set_named_param(name, value);
self
}
pub fn auto_commit(mut self, auto_commit: bool) -> Self {
self.ctx.session.auto_commit = auto_commit;
self
}
pub fn database(mut self, database: impl Into<String>) -> Self {
self.ctx.session.current_database = Arc::new(Some(database.into()));
self
}
pub fn principal_id(mut self, principal_id: ObjectId) -> Self {
self.ctx = self.ctx.with_principal_id(principal_id);
self
}
pub fn session_var(mut self, name: impl Into<String>, value: Value) -> Self {
let mut variables = (*self.ctx.session.session_vars).clone();
variables.insert(name.into(), value);
self.ctx.session.session_vars = Arc::new(variables);
self
}
pub fn timeout_ms(mut self, timeout_ms: u64) -> Self {
self.ctx.query.timeout_ms = timeout_ms;
self
}
#[doc(hidden)]
pub fn join_hash_state_max_bytes(mut self, max_bytes: usize) -> Self {
self.ctx.query.join_hash_state_max_bytes = max_bytes;
self
}
pub fn build(self) -> ExecutionContext {
self.ctx
}
}
impl Default for ExecutionContextBuilder {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use rustc_hash::FxHashMap;
struct TestHashObserver {
retained_bytes: usize,
observed: usize,
}
impl TestHashObserver {
fn new(retained_bytes: usize) -> Self {
Self {
retained_bytes,
observed: 0,
}
}
}
impl crate::hash_table::JoinHashObserver for TestHashObserver {
fn insert_raw_hash(&mut self, _hash: u64) {
self.observed += 1;
}
fn retained_bytes(&self) -> usize {
self.retained_bytes
}
}
#[test]
fn test_context_new() {
let ctx = ExecutionContext::new();
assert_eq!(ctx.param_count(), 0);
assert!(ctx.auto_commit());
assert!(!ctx.is_cancelled());
}
#[test]
fn test_context_with_params() {
let ctx = ExecutionContext::with_params(smallvec::smallvec![
Value::Integer(1),
Value::text("hello")
]);
assert_eq!(ctx.param_count(), 2);
assert_eq!(ctx.get_param(1), Some(&Value::Integer(1)));
assert_eq!(ctx.get_param(2), Some(&Value::text("hello")));
assert_eq!(ctx.get_param(0), None); assert_eq!(ctx.get_param(3), None); }
#[test]
fn test_context_named_params() {
let mut params = FxHashMap::default();
params.insert("name".to_string(), Value::text("Alice"));
params.insert("age".to_string(), Value::Integer(30));
let ctx = ExecutionContext::with_named_params(params);
assert_eq!(ctx.get_named_param("name"), Some(&Value::text("Alice")));
assert_eq!(ctx.get_named_param("age"), Some(&Value::Integer(30)));
assert_eq!(ctx.get_named_param("unknown"), None);
}
#[test]
fn test_context_cancellation() {
let ctx = ExecutionContext::new();
assert!(!ctx.is_cancelled());
let handle = ctx.cancellation_handle();
assert!(!handle.is_cancelled());
handle.cancel();
assert!(ctx.is_cancelled());
assert!(handle.is_cancelled());
}
#[test]
fn test_context_check_cancelled() {
let ctx = ExecutionContext::new();
assert!(ctx.check_cancelled().is_ok());
ctx.cancel();
assert!(ctx.check_cancelled().is_err());
}
#[test]
fn test_context_session_vars() {
let mut ctx = ExecutionContext::new();
ctx.set_session_var("timezone", Value::text("UTC"));
assert_eq!(ctx.get_session_var("timezone"), Some(&Value::text("UTC")));
assert_eq!(ctx.get_session_var("unknown"), None);
}
#[test]
fn test_context_builder() {
let ctx = ExecutionContextBuilder::new()
.params(smallvec::smallvec![Value::Integer(1)])
.param(Value::Integer(2))
.named_param("name", Value::text("test"))
.auto_commit(false)
.database("mydb")
.timeout_ms(5000)
.build();
assert_eq!(ctx.param_count(), 2);
assert_eq!(ctx.get_param(1), Some(&Value::Integer(1)));
assert_eq!(ctx.get_param(2), Some(&Value::Integer(2)));
assert_eq!(ctx.get_named_param("name"), Some(&Value::text("test")));
assert!(!ctx.auto_commit());
assert_eq!(ctx.current_database(), Some("mydb"));
assert_eq!(ctx.timeout_ms(), 5000);
}
#[cfg(feature = "test-hooks")]
#[test]
fn request_local_join_hash_cache_evicts_to_its_byte_budget() {
let retained = crate::hash_table::JoinHashTable::estimated_retained_bytes(1)
.expect("one-row hash state size");
let ctx = ExecutionContextBuilder::new()
.join_hash_state_max_bytes(retained)
.build();
let first = CompactArc::new(vec![Row::from_values(vec![Value::Integer(1)])]);
let second = CompactArc::new(vec![Row::from_values(vec![Value::Integer(2)])]);
radixdb_storage::instrumentation::begin_join_execution_probe();
drop(
ctx.join_hash_state_for(CompactArc::clone(&first), &[0], true)
.expect("first state"),
);
drop(
ctx.join_hash_state_for(CompactArc::clone(&second), &[0], true)
.expect("second state"),
);
drop(
ctx.join_hash_state_for(CompactArc::clone(&second), &[0], true)
.expect("second state reuse"),
);
drop(
ctx.join_hash_state_for(CompactArc::clone(&first), &[0], true)
.expect("first state rebuild after eviction"),
);
let probe = radixdb_storage::instrumentation::end_join_execution_probe();
assert_eq!(probe.hash_state_builds, 3);
assert_eq!(probe.hash_state_reuses, 1);
}
#[test]
fn active_and_cached_hash_states_share_one_request_budget() {
let retained = JoinHashTable::estimated_retained_bytes(1).unwrap();
let ctx = ExecutionContextBuilder::new()
.join_hash_state_max_bytes(retained)
.build();
let cached_rows = CompactArc::new(vec![Row::from_values(vec![Value::Integer(1)])]);
let active_cached = ctx
.join_hash_state_for(CompactArc::clone(&cached_rows), &[0], true)
.expect("first state fits common budget");
assert_eq!(ctx.retained_join_memory_bytes(), retained);
let rejected = ctx.join_hash_state_for(
CompactArc::new(vec![Row::from_values(vec![Value::Integer(2)])]),
&[0],
false,
);
assert!(rejected.is_none());
assert_eq!(ctx.retained_join_memory_bytes(), retained);
drop(active_cached);
assert_eq!(ctx.retained_join_memory_bytes(), 0);
let admitted = ctx
.join_hash_state_for(
CompactArc::new(vec![Row::from_values(vec![Value::Integer(2)])]),
&[0],
false,
)
.expect("released request budget admits next transient state");
assert_eq!(ctx.retained_join_memory_bytes(), retained);
drop(admitted);
assert_eq!(ctx.retained_join_memory_bytes(), 0);
}
#[test]
fn context_clones_cannot_each_spend_the_join_budget() {
let retained = JoinHashTable::estimated_retained_bytes(1).unwrap();
let ctx = ExecutionContextBuilder::new()
.join_hash_state_max_bytes(retained)
.build();
let clone = ctx.clone();
let first = ctx
.join_hash_state_for(
CompactArc::new(vec![Row::from_values(vec![Value::Integer(1)])]),
&[0],
false,
)
.unwrap();
assert!(clone
.join_hash_state_for(
CompactArc::new(vec![Row::from_values(vec![Value::Integer(2)])]),
&[0],
false,
)
.is_none());
drop(first);
assert_eq!(clone.retained_join_memory_bytes(), 0);
}
#[test]
fn bloom_and_hash_share_the_same_join_memory_reservation() {
let rows: CompactArc<Vec<Row>> = CompactArc::new(
(0..100)
.map(|value| Row::from_values(vec![Value::Integer(value)]))
.collect(),
);
let mut builder = TestHashObserver::new(256);
let hash_bytes = JoinHashTable::estimated_retained_bytes(rows.len()).unwrap();
let bloom_bytes = crate::hash_table::JoinHashObserver::retained_bytes(&builder);
assert!(bloom_bytes > 0);
let too_small = ExecutionContextBuilder::new()
.join_hash_state_max_bytes(hash_bytes)
.build();
assert!(too_small
.join_hash_state_with_bloom_for(CompactArc::clone(&rows), &[0], &mut builder)
.is_none());
assert_eq!(too_small.retained_join_memory_bytes(), 0);
let exact = ExecutionContextBuilder::new()
.join_hash_state_max_bytes(hash_bytes + bloom_bytes)
.build();
let state = exact
.join_hash_state_with_bloom_for(rows, &[0], &mut builder)
.expect("combined hash and bloom bytes fit exactly");
assert_eq!(builder.observed, 100);
assert_eq!(exact.retained_join_memory_bytes(), hash_bytes + bloom_bytes);
drop(state);
assert_eq!(exact.retained_join_memory_bytes(), 0);
}
#[test]
fn r6_l01_b_completed_timeout_is_unregistered_immediately() {
let ctx = ExecutionContextBuilder::new().timeout_ms(60_000).build();
let guard = TimeoutGuard::new(&ctx).expect("timeout guard");
assert_eq!(pending_timeout_count_for(&ctx), 1);
drop(guard);
assert_eq!(pending_timeout_count_for(&ctx), 0);
assert!(!ctx.is_cancelled());
}
}