1use chrono::{DateTime, Utc};
21use lru::LruCache;
22use radixdb_catalog::ObjectId;
23use radixdb_core::time_compat::{system_time_now, Instant};
24use rustc_hash::FxHashMap;
25use std::cell::RefCell;
26use std::collections::BinaryHeap;
27use std::num::NonZeroUsize;
28use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
29use std::sync::{Arc, Condvar, LazyLock, Mutex};
30use std::time::Duration;
31
32const SCALAR_SUBQUERY_CACHE_SIZE: usize = 128;
35const IN_SUBQUERY_CACHE_SIZE: usize = 128;
36const SEMI_JOIN_CACHE_SIZE: usize = 256;
37
38use crate::hash_table::{JoinHashState, JoinHashTable, JoinMemoryOwner, JoinMemoryReservation};
39use radixdb_core::ParamVec;
40use radixdb_core::{CompactArc, StringMap};
41use radixdb_core::{Result, Row, Value, ValueMap, ValueSet};
42use radixdb_procedural::BudgetOwner;
43
44pub(crate) trait StoredFunctionInvoker: std::fmt::Debug + Send + Sync {
48 fn invoke(self: Arc<Self>, name: &str, arguments: &[Value]) -> Result<Value>;
49
50 fn external_equal(&self, left: &Value, right: &Value) -> Result<bool> {
51 let _ = (left, right);
52 Err(radixdb_core::Error::NotSupported(
53 "external equality is unavailable in this execution context".to_owned(),
54 ))
55 }
56
57 fn external_compare(&self, left: &Value, right: &Value) -> Result<std::cmp::Ordering> {
58 let _ = (left, right);
59 Err(radixdb_core::Error::NotSupported(
60 "external ordering is unavailable in this execution context".to_owned(),
61 ))
62 }
63
64 fn external_input(&self, type_name: &str, input: &Value) -> Result<Value> {
65 let _ = (type_name, input);
66 Err(radixdb_core::Error::NotSupported(
67 "external type input is unavailable in this execution context".to_owned(),
68 ))
69 }
70
71 fn external_output(&self, value: &Value, target_type: radixdb_core::DataType) -> Result<Value> {
72 let _ = (value, target_type);
73 Err(radixdb_core::Error::NotSupported(
74 "external type output is unavailable in this execution context".to_owned(),
75 ))
76 }
77}
78
79pub(crate) const STORED_OPERATOR_CALL_PREFIX: &str = "\u{1f}operator:";
80
81static EMPTY_PARAMS: LazyLock<CompactArc<ParamVec>> =
85 LazyLock::new(|| CompactArc::new(ParamVec::new()));
86static EMPTY_DATABASE: LazyLock<Arc<Option<String>>> = LazyLock::new(|| Arc::new(None));
87static EMPTY_SESSION_VARS: LazyLock<Arc<AHashMap<String, Value>>> =
88 LazyLock::new(|| Arc::new(AHashMap::new()));
89
90pub(crate) fn is_system_context_name(name: &str) -> bool {
91 matches!(
92 name.to_ascii_uppercase().as_str(),
93 "CURRENT_PRINCIPAL"
94 | "CURRENT_EFFECTIVE_PRINCIPAL"
95 | "CURRENT_TRANSACTION_ID"
96 | "CURRENT_STATEMENT_TIMESTAMP"
97 | "CURRENT_REQUEST_ID"
98 | "CURRENT_IDEMPOTENCY_KEY"
99 | "CURRENT_JOB_ID"
100 | "CURRENT_JOB_ATTEMPT"
101 | "CURRENT_JOB_SCHEDULED_AT"
102 )
103}
104
105use crate::expression::RowFilter;
111use smallvec::SmallVec;
112
113type ScalarSubqueryCacheEntry = (SmallVec<[CompactArc<str>; 2]>, Value);
115
116thread_local! {
117 static SCALAR_SUBQUERY_CACHE: RefCell<LruCache<String, ScalarSubqueryCacheEntry>> =
118 RefCell::new(LruCache::new(NonZeroUsize::new(SCALAR_SUBQUERY_CACHE_SIZE).unwrap()));
119 static EXISTS_PREDICATE_CACHE: RefCell<FxHashMap<String, RowFilter>> =
120 RefCell::new(FxHashMap::default());
121 static ACTIVE_QUERY_CANCELLATION: RefCell<Vec<CancellationHandle>> = const {
122 RefCell::new(Vec::new())
123 };
124}
125
126pub fn clear_exists_predicate_cache() {
128 EXISTS_PREDICATE_CACHE.with(|cache| cache.borrow_mut().clear());
129}
130
131pub fn get_cached_exists_predicate(key: &str) -> Option<RowFilter> {
133 EXISTS_PREDICATE_CACHE.with(|cache| cache.borrow().get(key).cloned())
134}
135
136pub fn cache_exists_predicate(key: String, filter: RowFilter) {
138 EXISTS_PREDICATE_CACHE.with(|cache| {
139 cache.borrow_mut().insert(key, filter);
140 });
141}
142
143#[doc(hidden)]
144pub struct StatementCancellationScope;
145
146impl Drop for StatementCancellationScope {
147 fn drop(&mut self) {
148 ACTIVE_QUERY_CANCELLATION.with(|stack| {
149 stack.borrow_mut().pop();
150 });
151 }
152}
153
154#[doc(hidden)]
155pub fn current_query_is_cancelled() -> bool {
156 ACTIVE_QUERY_CANCELLATION.with(|stack| {
157 stack
158 .borrow()
159 .last()
160 .is_some_and(CancellationHandle::is_cancelled)
161 })
162}
163
164#[inline]
167#[doc(hidden)]
168pub fn with_current_query_cancellation<T>(
169 callback: impl FnOnce(Option<&dyn radixdb_functions::FunctionCancellation>) -> T,
170) -> T {
171 ACTIVE_QUERY_CANCELLATION.with(|stack| {
172 let stack = stack.borrow();
173 callback(
174 stack
175 .last()
176 .map(|handle| handle as &dyn radixdb_functions::FunctionCancellation),
177 )
178 })
179}
180
181#[inline]
182#[doc(hidden)]
183pub fn check_current_query_cancelled() -> Result<()> {
184 if current_query_is_cancelled() {
185 Err(radixdb_core::Error::QueryCancelled)
186 } else {
187 Ok(())
188 }
189}
190
191pub fn clear_scalar_subquery_cache() {
195 SCALAR_SUBQUERY_CACHE.with(|cache| {
196 cache.borrow_mut().clear();
197 });
198}
199
200#[inline]
203pub fn invalidate_scalar_subquery_cache_for_table(table_name: &str) {
204 SCALAR_SUBQUERY_CACHE.with(|cache| {
205 let mut c = cache.borrow_mut();
206 if c.is_empty() {
207 return;
208 }
209 let keys_to_remove: Vec<String> = c
211 .iter()
212 .filter(|(_, (tables, _))| tables.iter().any(|t| t.eq_ignore_ascii_case(table_name)))
213 .map(|(k, _)| k.clone())
214 .collect();
215 for key in keys_to_remove {
216 c.pop(&key);
217 }
218 });
219}
220
221pub fn get_cached_scalar_subquery(key: &str) -> Option<Value> {
223 SCALAR_SUBQUERY_CACHE.with(|cache| cache.borrow_mut().get(key).map(|(_, v)| v.clone()))
224}
225
226pub fn cache_scalar_subquery(key: String, tables: SmallVec<[CompactArc<str>; 2]>, value: Value) {
228 SCALAR_SUBQUERY_CACHE.with(|cache| {
229 cache.borrow_mut().put(key, (tables, value));
230 });
231}
232
233type InSubqueryCacheEntry = (SmallVec<[CompactArc<str>; 2]>, Vec<Value>);
241
242thread_local! {
243 static IN_SUBQUERY_CACHE: RefCell<LruCache<String, InSubqueryCacheEntry>> =
244 RefCell::new(LruCache::new(NonZeroUsize::new(IN_SUBQUERY_CACHE_SIZE).unwrap()));
245}
246
247pub fn clear_in_subquery_cache() {
251 IN_SUBQUERY_CACHE.with(|cache| {
252 cache.borrow_mut().clear();
253 });
254}
255
256#[inline]
259pub fn invalidate_in_subquery_cache_for_table(table_name: &str) {
260 IN_SUBQUERY_CACHE.with(|cache| {
261 let mut c = cache.borrow_mut();
262 if c.is_empty() {
263 return;
264 }
265 let keys_to_remove: Vec<String> = c
267 .iter()
268 .filter(|(_, (tables, _))| tables.iter().any(|t| t.eq_ignore_ascii_case(table_name)))
269 .map(|(k, _)| k.clone())
270 .collect();
271 for key in keys_to_remove {
272 c.pop(&key);
273 }
274 });
275}
276
277pub fn get_cached_in_subquery(key: &str) -> Option<Vec<Value>> {
279 IN_SUBQUERY_CACHE.with(|cache| cache.borrow_mut().get(key).map(|(_, v)| v.clone()))
280}
281
282pub fn cache_in_subquery(key: String, tables: SmallVec<[CompactArc<str>; 2]>, values: Vec<Value>) {
284 IN_SUBQUERY_CACHE.with(|cache| {
285 cache.borrow_mut().put(key, (tables, values));
286 });
287}
288
289use radixdb_sql::ast::{Expression, SelectStatement};
290
291pub fn extract_table_names_for_cache(stmt: &SelectStatement) -> SmallVec<[CompactArc<str>; 2]> {
295 let mut tables = SmallVec::new();
296 if let Some(ref table_expr) = stmt.table_expr {
297 collect_real_table_names(table_expr, &mut tables);
298 }
299 tables
300}
301
302fn collect_real_table_names(source: &Expression, tables: &mut SmallVec<[CompactArc<str>; 2]>) {
304 match source {
305 Expression::TableSource(ts) => {
306 tables.push(CompactArc::from(ts.name.value_lower.as_str()));
308 }
309 Expression::JoinSource(js) => {
310 collect_real_table_names(&js.left, tables);
311 collect_real_table_names(&js.right, tables);
312 }
313 Expression::SubquerySource(ss) => {
314 if let Some(ref table_expr) = ss.subquery.table_expr {
316 collect_real_table_names(table_expr, tables);
317 }
318 }
319 _ => {}
320 }
321}
322
323use ahash::AHashMap;
328use std::hash::{Hash, Hasher};
329
330type SemiJoinCacheEntry = (CompactArc<str>, CompactArc<ValueSet>);
332
333#[inline]
335pub fn compute_semi_join_cache_key(table: &str, column: &str, pred_hash: u64) -> u64 {
336 let mut hasher = rustc_hash::FxHasher::default();
337 table.hash(&mut hasher);
338 column.hash(&mut hasher);
339 pred_hash.hash(&mut hasher);
340 hasher.finish()
341}
342
343thread_local! {
344 static SEMI_JOIN_CACHE: RefCell<LruCache<u64, SemiJoinCacheEntry>> =
345 RefCell::new(LruCache::new(NonZeroUsize::new(SEMI_JOIN_CACHE_SIZE).unwrap()));
346}
347
348pub fn clear_semi_join_cache() {
352 SEMI_JOIN_CACHE.with(|cache| {
353 cache.borrow_mut().clear();
354 });
355}
356
357#[inline]
360pub fn invalidate_semi_join_cache_for_table(table_name: &str) {
361 SEMI_JOIN_CACHE.with(|cache| {
362 let mut c = cache.borrow_mut();
363 if c.is_empty() {
364 return;
365 }
366 let keys_to_remove: Vec<u64> = c
368 .iter()
369 .filter(|(_, (key_table, _))| key_table.eq_ignore_ascii_case(table_name))
370 .map(|(k, _)| *k)
371 .collect();
372 for key in keys_to_remove {
373 c.pop(&key);
374 }
375 });
376}
377
378#[inline]
380pub fn get_cached_semi_join(key_hash: u64) -> Option<CompactArc<ValueSet>> {
381 SEMI_JOIN_CACHE.with(|cache| {
382 cache
383 .borrow_mut()
384 .get(&key_hash)
385 .map(|(_, v)| CompactArc::clone(v))
386 })
387}
388
389#[inline]
391pub fn cache_semi_join_arc(key_hash: u64, table: &str, values: CompactArc<ValueSet>) {
392 SEMI_JOIN_CACHE.with(|cache| {
393 cache
394 .borrow_mut()
395 .put(key_hash, (CompactArc::from(table), values));
396 });
397}
398
399use radixdb_storage::traits::Index;
402thread_local! {
403 static EXISTS_INDEX_CACHE: RefCell<FxHashMap<String, std::sync::Arc<dyn Index>>> = RefCell::new(FxHashMap::default());
404}
405
406pub fn clear_exists_index_cache() {
408 EXISTS_INDEX_CACHE.with(|cache| {
409 cache.borrow_mut().clear();
410 });
411}
412
413pub fn get_cached_exists_index(key: &str) -> Option<std::sync::Arc<dyn Index>> {
415 EXISTS_INDEX_CACHE.with(|cache| cache.borrow().get(key).cloned())
416}
417
418pub fn cache_exists_index(key: String, index: std::sync::Arc<dyn Index>) {
420 EXISTS_INDEX_CACHE.with(|cache| {
421 cache.borrow_mut().insert(key, index);
422 });
423}
424
425pub type RowFetcher =
427 Box<dyn Fn(&[i64]) -> radixdb_core::Result<radixdb_core::RowVec> + Send + Sync>;
428
429pub type RowCounter = Box<dyn Fn(&[i64]) -> usize + Send + Sync>;
432
433thread_local! {
436 static EXISTS_FETCHER_CACHE: RefCell<FxHashMap<String, std::sync::Arc<RowFetcher>>> = RefCell::new(FxHashMap::default());
437}
438
439thread_local! {
442 static COUNT_COUNTER_CACHE: RefCell<FxHashMap<String, std::sync::Arc<RowCounter>>> = RefCell::new(FxHashMap::default());
443}
444
445pub fn clear_exists_fetcher_cache() {
447 EXISTS_FETCHER_CACHE.with(|cache| {
448 cache.borrow_mut().clear();
449 });
450}
451
452pub fn clear_count_counter_cache() {
454 COUNT_COUNTER_CACHE.with(|cache| {
455 cache.borrow_mut().clear();
456 });
457}
458
459pub fn get_cached_exists_fetcher(key: &str) -> Option<std::sync::Arc<RowFetcher>> {
461 EXISTS_FETCHER_CACHE.with(|cache| cache.borrow().get(key).cloned())
462}
463
464pub fn get_cached_count_counter(key: &str) -> Option<std::sync::Arc<RowCounter>> {
466 COUNT_COUNTER_CACHE.with(|cache| cache.borrow().get(key).cloned())
467}
468
469pub fn cache_exists_fetcher(key: String, fetcher: RowFetcher) {
471 EXISTS_FETCHER_CACHE.with(|cache| {
472 cache.borrow_mut().insert(key, std::sync::Arc::new(fetcher));
473 });
474}
475
476pub fn cache_count_counter(key: String, counter: RowCounter) {
478 COUNT_COUNTER_CACHE.with(|cache| {
479 cache.borrow_mut().insert(key, std::sync::Arc::new(counter));
480 });
481}
482
483thread_local! {
486 static EXISTS_SCHEMA_CACHE: RefCell<FxHashMap<String, CompactArc<Vec<String>>>> = RefCell::new(FxHashMap::default());
487}
488
489pub fn clear_exists_schema_cache() {
491 EXISTS_SCHEMA_CACHE.with(|cache| {
492 cache.borrow_mut().clear();
493 });
494}
495
496pub fn get_cached_exists_schema(key: &str) -> Option<CompactArc<Vec<String>>> {
498 EXISTS_SCHEMA_CACHE.with(|cache| cache.borrow().get(key).cloned())
499}
500
501pub fn cache_exists_schema(key: String, columns: CompactArc<Vec<String>>) {
503 EXISTS_SCHEMA_CACHE.with(|cache| {
504 cache.borrow_mut().insert(key, columns);
505 });
506}
507
508thread_local! {
511 static EXISTS_PRED_KEY_CACHE: RefCell<FxHashMap<usize, String>> = RefCell::new(FxHashMap::default());
512}
513
514pub fn clear_exists_pred_key_cache() {
516 EXISTS_PRED_KEY_CACHE.with(|cache| {
517 cache.borrow_mut().clear();
518 });
519}
520
521#[inline]
523pub fn get_cached_exists_pred_key(subquery_ptr: usize) -> Option<String> {
524 EXISTS_PRED_KEY_CACHE.with(|cache| cache.borrow().get(&subquery_ptr).cloned())
525}
526
527#[inline]
529pub fn cache_exists_pred_key(subquery_ptr: usize, pred_key: String) {
530 EXISTS_PRED_KEY_CACHE.with(|cache| {
531 cache.borrow_mut().insert(subquery_ptr, pred_key);
532 });
533}
534
535thread_local! {
539 static BATCH_AGGREGATE_CACHE: RefCell<FxHashMap<String, CompactArc<ValueMap<Value>>>> = RefCell::new(FxHashMap::default());
540}
541
542pub fn clear_batch_aggregate_cache() {
544 BATCH_AGGREGATE_CACHE.with(|cache| {
545 let mut c = cache.borrow_mut();
546 c.clear();
547 c.shrink_to_fit();
548 });
549}
550
551pub fn get_cached_batch_aggregate(key: &str) -> Option<CompactArc<ValueMap<Value>>> {
553 BATCH_AGGREGATE_CACHE.with(|cache| cache.borrow().get(key).cloned())
554}
555
556pub fn cache_batch_aggregate(key: String, values: ValueMap<Value>) {
558 BATCH_AGGREGATE_CACHE.with(|cache| {
559 cache.borrow_mut().insert(key, CompactArc::new(values));
560 });
561}
562
563#[derive(Clone)]
565pub struct BatchAggregateLookupInfo {
566 pub cache_key: String,
568 pub outer_column_lower: String,
570 pub outer_qualified_lower: Option<String>,
572 pub is_count: bool,
574}
575
576thread_local! {
580 static BATCH_AGGREGATE_INFO_CACHE: RefCell<FxHashMap<usize, Option<Arc<BatchAggregateLookupInfo>>>> = RefCell::new(FxHashMap::default());
581}
582
583pub fn clear_batch_aggregate_info_cache() {
585 BATCH_AGGREGATE_INFO_CACHE.with(|cache| {
586 let mut c = cache.borrow_mut();
587 c.clear();
588 c.shrink_to_fit();
589 });
590}
591
592#[inline]
595pub fn get_cached_batch_aggregate_info(
596 subquery_ptr: usize,
597) -> Option<Option<Arc<BatchAggregateLookupInfo>>> {
598 BATCH_AGGREGATE_INFO_CACHE.with(|cache| cache.borrow().get(&subquery_ptr).cloned())
599}
600
601#[inline]
604pub fn cache_batch_aggregate_info(
605 subquery_ptr: usize,
606 info: Option<BatchAggregateLookupInfo>,
607) -> Option<Arc<BatchAggregateLookupInfo>> {
608 let arc_info = info.map(Arc::new);
609 let result = arc_info.clone();
610 BATCH_AGGREGATE_INFO_CACHE.with(|cache| {
611 cache.borrow_mut().insert(subquery_ptr, arc_info);
612 });
613 result
614}
615
616#[derive(Clone)]
619pub struct ExistsCorrelationInfo {
620 pub outer_column: String,
622 pub outer_table: Option<String>,
624 pub inner_column: String,
626 pub inner_table: String,
628 pub outer_column_lower: String,
630 pub outer_qualified_lower: Option<String>,
632 pub additional_predicate: Option<Expression>,
634 pub index_cache_key: String,
636}
637
638thread_local! {
642 static EXISTS_CORRELATION_CACHE: RefCell<FxHashMap<usize, Option<Arc<ExistsCorrelationInfo>>>> = RefCell::new(FxHashMap::default());
643}
644
645pub fn clear_exists_correlation_cache() {
647 EXISTS_CORRELATION_CACHE.with(|cache| {
648 cache.borrow_mut().clear();
649 });
650}
651
652pub fn clear_executor_thread_local_caches() {
656 SCALAR_SUBQUERY_CACHE.with(|cache| {
658 cache.borrow_mut().clear();
659 });
660 IN_SUBQUERY_CACHE.with(|cache| {
661 cache.borrow_mut().clear();
662 });
663 SEMI_JOIN_CACHE.with(|cache| {
664 cache.borrow_mut().clear();
665 });
666 EXISTS_INDEX_CACHE.with(|cache| {
668 let mut c = cache.borrow_mut();
669 c.clear();
670 c.shrink_to_fit();
671 });
672 EXISTS_FETCHER_CACHE.with(|cache| {
673 let mut c = cache.borrow_mut();
674 c.clear();
675 c.shrink_to_fit();
676 });
677 COUNT_COUNTER_CACHE.with(|cache| {
678 let mut c = cache.borrow_mut();
679 c.clear();
680 c.shrink_to_fit();
681 });
682 EXISTS_SCHEMA_CACHE.with(|cache| {
683 let mut c = cache.borrow_mut();
684 c.clear();
685 c.shrink_to_fit();
686 });
687 EXISTS_PRED_KEY_CACHE.with(|cache| {
688 let mut c = cache.borrow_mut();
689 c.clear();
690 c.shrink_to_fit();
691 });
692 BATCH_AGGREGATE_CACHE.with(|cache| {
693 let mut c = cache.borrow_mut();
694 c.clear();
695 c.shrink_to_fit();
696 });
697 BATCH_AGGREGATE_INFO_CACHE.with(|cache| {
698 let mut c = cache.borrow_mut();
699 c.clear();
700 c.shrink_to_fit();
701 });
702 EXISTS_CORRELATION_CACHE.with(|cache| {
703 let mut c = cache.borrow_mut();
704 c.clear();
705 c.shrink_to_fit();
706 });
707
708 radixdb_storage::expression::clear_regex_cache();
710 radixdb_storage::expression::clear_like_regex_cache();
711
712 radixdb_core::row_vec::clear_row_vec_pool();
714 radixdb_core::row_vec::clear_row_id_vec_pool();
715
716 radixdb_storage::mvcc::clear_version_map_pools();
718}
719
720pub fn clear_all_thread_local_caches() {
722 clear_executor_thread_local_caches();
723 clear_exists_predicate_cache();
724 crate::expression::clear_program_cache();
725 crate::query::clear_join_dependency_projection_cache();
726 crate::utils::clear_join_projection_lookup_cache();
727 crate::query_classification::clear_classification_cache();
728}
729
730#[inline]
733pub fn get_cached_exists_correlation(
734 subquery_ptr: usize,
735) -> Option<Option<Arc<ExistsCorrelationInfo>>> {
736 EXISTS_CORRELATION_CACHE.with(|cache| cache.borrow().get(&subquery_ptr).cloned())
737}
738
739#[inline]
742pub fn cache_exists_correlation(
743 subquery_ptr: usize,
744 info: Option<ExistsCorrelationInfo>,
745) -> Option<Arc<ExistsCorrelationInfo>> {
746 let arc_info = info.map(Arc::new);
747 let result = arc_info.clone();
748 EXISTS_CORRELATION_CACHE.with(|cache| {
749 cache.borrow_mut().insert(subquery_ptr, arc_info);
750 });
751 result
752}
753
754#[derive(Debug, Clone)]
762pub struct ExecutionContext {
763 session: SessionState,
764 query: QueryState,
765}
766
767#[derive(Debug, Clone)]
771struct SessionState {
772 auto_commit: bool,
774 current_database: Arc<Option<String>>,
776 session_vars: Arc<AHashMap<String, Value>>,
778 principal_id: ObjectId,
781}
782
783#[derive(Debug, Clone)]
787struct QueryState {
788 effective_principal_id: Option<ObjectId>,
792 params: CompactArc<ParamVec>,
794 named_params: Arc<FxHashMap<String, Value>>,
796 cancelled: Arc<AtomicBool>,
798 timed_out: Arc<AtomicBool>,
800 active_reference_expands: Arc<AtomicUsize>,
803 parent_cancelled: Option<Arc<AtomicBool>>,
806 timeout_ms: u64,
808 join_hash_state_max_bytes: usize,
811 join_hash_states: Arc<Mutex<JoinHashStateCache>>,
814 join_memory_owner: Arc<JoinMemoryOwner>,
817 view_depth: usize,
819 query_depth: usize,
822 outer_row: Option<FxHashMap<CompactArc<str>, Value>>,
827 outer_columns: Option<CompactArc<Vec<String>>>,
829 cte_data: Option<Arc<CteDataMap>>,
832 transaction_id: Option<u64>,
834 stored_function_invoker: Option<Arc<dyn StoredFunctionInvoker>>,
836 procedural_budget: Option<BudgetOwner>,
839 public_scan_budget: Option<Arc<PublicScanBudget>>,
843}
844
845#[derive(Debug)]
846struct PublicScanBudget {
847 remaining: AtomicUsize,
848}
849
850impl PublicScanBudget {
851 fn new(max_rows: usize) -> Self {
852 Self {
853 remaining: AtomicUsize::new(max_rows),
854 }
855 }
856
857 fn claim(&self, rows: usize) -> Result<()> {
858 self.remaining
859 .fetch_update(Ordering::AcqRel, Ordering::Acquire, |remaining| {
860 remaining.checked_sub(rows)
861 })
862 .map(|_| ())
863 .map_err(|_| {
864 radixdb_core::Error::invalid_argument("public read scanned-row budget exceeded")
865 })
866 }
867}
868
869#[doc(hidden)]
872pub type CteMaterializedRows = Arc<std::sync::OnceLock<CompactArc<Vec<Row>>>>;
873#[doc(hidden)]
874pub type CteData = (
875 CompactArc<Vec<String>>,
876 CompactArc<Vec<(i64, Row)>>,
877 CteMaterializedRows,
878);
879
880#[doc(hidden)]
884pub type CteDataMap = StringMap<CteData>;
885
886#[derive(Default)]
887struct JoinHashStateCache {
888 states: Vec<JoinHashState>,
889}
890
891impl std::fmt::Debug for JoinHashStateCache {
892 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
893 formatter
894 .debug_struct("JoinHashStateCache")
895 .field("states", &self.states.len())
896 .finish()
897 }
898}
899
900impl JoinHashStateCache {
901 fn get(
902 &mut self,
903 build_rows: &CompactArc<Vec<Row>>,
904 key_indices: &[usize],
905 ) -> Option<JoinHashState> {
906 let index = self
907 .states
908 .iter()
909 .position(|state| state.matches(build_rows, key_indices))?;
910 let state = self.states.remove(index);
911 let result = state.clone();
912 self.states.push(state);
913 Some(result)
914 }
915
916 fn insert(&mut self, state: JoinHashState) {
917 self.states.push(state);
918 }
919
920 fn pop_lru(&mut self) -> Option<JoinHashState> {
921 (!self.states.is_empty()).then(|| self.states.remove(0))
922 }
923}
924
925impl Default for ExecutionContext {
926 fn default() -> Self {
927 Self::new()
928 }
929}
930
931impl ExecutionContext {
932 pub fn new() -> Self {
935 let mut system_values = FxHashMap::default();
936 system_values.insert(
937 "CURRENT_PRINCIPAL".to_owned(),
938 Value::uuid(ObjectId::BOOTSTRAP_OWNER.into_bytes()),
939 );
940 system_values.insert(
941 "CURRENT_EFFECTIVE_PRINCIPAL".to_owned(),
942 Value::uuid(ObjectId::BOOTSTRAP_OWNER.into_bytes()),
943 );
944 system_values.insert(
945 "CURRENT_STATEMENT_TIMESTAMP".to_owned(),
946 Value::timestamp(system_time_now().into()),
947 );
948 Self {
949 session: SessionState {
950 auto_commit: true,
951 current_database: EMPTY_DATABASE.clone(),
952 session_vars: EMPTY_SESSION_VARS.clone(),
953 principal_id: ObjectId::BOOTSTRAP_OWNER,
954 },
955 query: QueryState {
956 effective_principal_id: None,
957 params: EMPTY_PARAMS.clone(),
958 named_params: Arc::new(system_values),
959 cancelled: Arc::new(AtomicBool::new(false)),
960 timed_out: Arc::new(AtomicBool::new(false)),
961 active_reference_expands: Arc::new(AtomicUsize::new(0)),
962 parent_cancelled: None,
963 timeout_ms: 0,
964 join_hash_state_max_bytes: crate::hash_table::DEFAULT_JOIN_HASH_STATE_MAX_BYTES,
965 join_hash_states: Arc::new(Mutex::new(JoinHashStateCache::default())),
966 join_memory_owner: Arc::new(JoinMemoryOwner::default()),
967 view_depth: 0,
968 query_depth: 0,
969 outer_row: None,
970 outer_columns: None,
971 cte_data: None,
972 transaction_id: None,
973 stored_function_invoker: None,
974 procedural_budget: None,
975 public_scan_budget: None,
976 },
977 }
978 }
979
980 #[doc(hidden)]
981 pub fn enter_statement_scope(&self) -> StatementCancellationScope {
982 ACTIVE_QUERY_CANCELLATION.with(|stack| {
983 stack.borrow_mut().push(self.cancellation_handle());
984 });
985 StatementCancellationScope
986 }
987
988 #[doc(hidden)]
989 pub fn enter_reference_expand(&self) -> ReferenceExpandExecutionGuard {
990 self.query
991 .active_reference_expands
992 .fetch_add(1, Ordering::AcqRel);
993 ReferenceExpandExecutionGuard {
994 active: Arc::clone(&self.query.active_reference_expands),
995 }
996 }
997
998 #[doc(hidden)]
999 pub fn active_reference_expands(&self) -> usize {
1000 self.query.active_reference_expands.load(Ordering::Acquire)
1001 }
1002
1003 pub fn with_params(params: ParamVec) -> Self {
1005 let mut context = Self::new();
1006 context.query.params = CompactArc::new(params);
1007 context
1008 }
1009
1010 pub fn with_named_params(named_params: FxHashMap<String, Value>) -> Self {
1012 let mut context = Self::new();
1013 let system_values = context.query.named_params.clone();
1014 let mut admitted = named_params;
1015 admitted.retain(|name, _| !is_system_context_name(name));
1016 admitted.extend(
1017 system_values
1018 .iter()
1019 .map(|(name, value)| (name.clone(), value.clone())),
1020 );
1021 context.query.named_params = Arc::new(admitted);
1022 context
1023 }
1024
1025 pub fn get_param(&self, index: usize) -> Option<&Value> {
1027 if index == 0 || index > self.query.params.len() {
1028 None
1029 } else {
1030 self.query.params.get(index - 1)
1031 }
1032 }
1033
1034 pub fn get_named_param(&self, name: &str) -> Option<&Value> {
1036 self.query.named_params.get(name)
1037 }
1038
1039 #[inline]
1040 #[doc(hidden)]
1041 pub fn join_hash_state_max_bytes(&self) -> usize {
1042 self.query.join_hash_state_max_bytes
1043 }
1044
1045 #[doc(hidden)]
1050 pub fn reserve_join_memory(&self, bytes: usize) -> Option<JoinMemoryReservation> {
1051 loop {
1052 if let Some(reservation) = self
1053 .query
1054 .join_memory_owner
1055 .try_reserve(bytes, self.query.join_hash_state_max_bytes)
1056 {
1057 return Some(reservation);
1058 }
1059
1060 let evicted = self
1061 .query
1062 .join_hash_states
1063 .lock()
1064 .unwrap_or_else(std::sync::PoisonError::into_inner)
1065 .pop_lru();
1066 let evicted = evicted?;
1067 drop(evicted);
1070 }
1071 }
1072
1073 #[doc(hidden)]
1074 pub fn retained_join_memory_bytes(&self) -> usize {
1075 self.query.join_memory_owner.retained_bytes()
1076 }
1077
1078 #[doc(hidden)]
1079 pub fn peak_join_memory_bytes(&self) -> usize {
1080 self.query.join_memory_owner.peak_bytes()
1081 }
1082
1083 #[doc(hidden)]
1087 pub fn join_hash_state_for(
1088 &self,
1089 build_rows: CompactArc<Vec<Row>>,
1090 key_indices: &[usize],
1091 retain_for_reuse: bool,
1092 ) -> Option<JoinHashState> {
1093 if key_indices.is_empty() {
1094 return None;
1095 }
1096 if retain_for_reuse {
1097 let mut cache = self
1098 .query
1099 .join_hash_states
1100 .lock()
1101 .unwrap_or_else(std::sync::PoisonError::into_inner);
1102 if let Some(state) = cache.get(&build_rows, key_indices) {
1103 radixdb_storage::instrumentation::record_join_hash_state_reuse();
1104 return Some(state);
1105 }
1106 }
1107
1108 let state_bytes = JoinHashTable::estimated_retained_bytes(build_rows.len())?;
1109 let reservation = self.reserve_join_memory(state_bytes)?;
1110 let state = JoinHashState::build_reserved(build_rows, key_indices, reservation);
1111 radixdb_storage::instrumentation::record_join_hash_state_build();
1112 if retain_for_reuse {
1113 let mut cache = self
1114 .query
1115 .join_hash_states
1116 .lock()
1117 .unwrap_or_else(std::sync::PoisonError::into_inner);
1118 if let Some(existing) = cache.get(state.build_rows(), key_indices) {
1119 radixdb_storage::instrumentation::record_join_hash_state_reuse();
1120 return Some(existing);
1121 }
1122 cache.insert(state.clone());
1123 }
1124 Some(state)
1125 }
1126
1127 #[doc(hidden)]
1130 pub fn join_hash_state_with_bloom_for(
1131 &self,
1132 build_rows: CompactArc<Vec<Row>>,
1133 key_indices: &[usize],
1134 bloom_builder: &mut impl crate::hash_table::JoinHashObserver,
1135 ) -> Option<JoinHashState> {
1136 if key_indices.is_empty() {
1137 return None;
1138 }
1139 let state_bytes = JoinHashTable::estimated_retained_bytes(build_rows.len())?
1140 .checked_add(bloom_builder.retained_bytes())?;
1141 let reservation = self.reserve_join_memory(state_bytes)?;
1142 let state = JoinHashState::build_with_bloom_reserved(
1143 build_rows,
1144 key_indices,
1145 bloom_builder,
1146 reservation,
1147 );
1148 radixdb_storage::instrumentation::record_join_hash_state_build();
1149 Some(state)
1150 }
1151
1152 pub fn params(&self) -> &[Value] {
1154 &self.query.params
1155 }
1156
1157 pub fn params_arc(&self) -> &CompactArc<ParamVec> {
1160 &self.query.params
1161 }
1162
1163 pub fn named_params(&self) -> &FxHashMap<String, Value> {
1165 &self.query.named_params
1166 }
1167
1168 pub fn named_params_arc(&self) -> &Arc<FxHashMap<String, Value>> {
1171 &self.query.named_params
1172 }
1173
1174 #[doc(hidden)]
1175 pub(crate) fn stored_function_invoker(&self) -> Option<&Arc<dyn StoredFunctionInvoker>> {
1176 self.query.stored_function_invoker.as_ref()
1177 }
1178
1179 #[doc(hidden)]
1180 pub(crate) fn with_stored_function_invoker(
1181 mut self,
1182 invoker: Arc<dyn StoredFunctionInvoker>,
1183 ) -> Self {
1184 self.query.stored_function_invoker = Some(invoker);
1185 self
1186 }
1187
1188 #[doc(hidden)]
1189 pub(crate) fn procedural_budget(&self) -> Option<&BudgetOwner> {
1190 self.query.procedural_budget.as_ref()
1191 }
1192
1193 #[doc(hidden)]
1194 pub(crate) fn with_procedural_budget(mut self, budget: BudgetOwner) -> Self {
1195 self.query.procedural_budget = Some(budget);
1196 self
1197 }
1198
1199 pub fn param_count(&self) -> usize {
1201 self.query.params.len()
1202 }
1203
1204 pub fn set_params(&mut self, params: ParamVec) {
1206 self.query.params = CompactArc::new(params);
1207 }
1208
1209 pub fn add_param(&mut self, value: Value) {
1211 CompactArc::make_mut(&mut self.query.params).push(value);
1212 }
1213
1214 pub fn set_named_param(&mut self, name: impl Into<String>, value: Value) {
1216 let name = name.into();
1217 if !is_system_context_name(&name) {
1218 Arc::make_mut(&mut self.query.named_params).insert(name, value);
1219 }
1220 }
1221
1222 pub fn auto_commit(&self) -> bool {
1224 self.session.auto_commit
1225 }
1226
1227 pub fn set_auto_commit(&mut self, auto_commit: bool) {
1229 self.session.auto_commit = auto_commit;
1230 }
1231
1232 pub fn is_cancelled(&self) -> bool {
1234 self.query.cancelled.load(Ordering::Relaxed)
1235 || self
1236 .query
1237 .parent_cancelled
1238 .as_ref()
1239 .is_some_and(|cancelled| cancelled.load(Ordering::Relaxed))
1240 }
1241
1242 #[doc(hidden)]
1243 pub fn did_time_out(&self) -> bool {
1244 self.query.timed_out.load(Ordering::Acquire)
1245 }
1246
1247 pub fn cancel(&self) {
1249 self.query.cancelled.store(true, Ordering::Relaxed);
1250 }
1251
1252 pub fn cancellation_handle(&self) -> CancellationHandle {
1254 CancellationHandle {
1255 cancelled: self.query.cancelled.clone(),
1256 parent_cancelled: self.query.parent_cancelled.clone(),
1257 }
1258 }
1259
1260 #[doc(hidden)]
1263 pub fn bind_parent_cancellation(&mut self, handle: &CancellationHandle) {
1264 self.query.parent_cancelled = Some(Arc::clone(&handle.cancelled));
1265 }
1266
1267 pub fn current_database(&self) -> Option<&str> {
1269 self.session.current_database.as_ref().as_deref()
1270 }
1271
1272 pub fn set_current_database(&mut self, database: impl Into<String>) {
1274 self.session.current_database = Arc::new(Some(database.into()));
1275 }
1276
1277 pub const fn principal_id(&self) -> ObjectId {
1279 self.session.principal_id
1280 }
1281
1282 pub fn with_principal_id(&self, principal_id: ObjectId) -> Self {
1284 let mut nested = self.clone();
1285 nested.session.principal_id = principal_id;
1286 nested.query.effective_principal_id = None;
1287 nested
1288 .set_system_context_value("CURRENT_PRINCIPAL", Value::uuid(principal_id.into_bytes()));
1289 nested.set_system_context_value(
1290 "CURRENT_EFFECTIVE_PRINCIPAL",
1291 Value::uuid(principal_id.into_bytes()),
1292 );
1293 nested
1294 }
1295
1296 pub const fn effective_principal_id(&self) -> ObjectId {
1299 match self.query.effective_principal_id {
1300 Some(principal) => principal,
1301 None => self.session.principal_id,
1302 }
1303 }
1304
1305 pub fn with_effective_principal_id(&self, principal_id: ObjectId) -> Self {
1308 let mut nested = self.clone();
1309 nested.query.effective_principal_id = Some(principal_id);
1310 nested.set_system_context_value(
1311 "CURRENT_EFFECTIVE_PRINCIPAL",
1312 Value::uuid(principal_id.into_bytes()),
1313 );
1314 nested
1315 }
1316
1317 #[doc(hidden)]
1321 pub fn with_public_scan_limit(&self, max_rows: usize) -> Self {
1322 let mut nested = self.clone();
1323 nested.query.public_scan_budget = Some(Arc::new(PublicScanBudget::new(max_rows)));
1324 nested
1325 }
1326
1327 #[doc(hidden)]
1328 pub fn claim_public_scan_rows(&self, rows: usize) -> Result<()> {
1329 self.query
1330 .public_scan_budget
1331 .as_ref()
1332 .map_or(Ok(()), |budget| budget.claim(rows))
1333 }
1334
1335 #[doc(hidden)]
1336 pub fn has_public_scan_budget(&self) -> bool {
1337 self.query.public_scan_budget.is_some()
1338 }
1339
1340 pub fn get_session_var(&self, name: &str) -> Option<&Value> {
1342 self.session.session_vars.get(name)
1343 }
1344
1345 pub fn set_session_var(&mut self, name: impl Into<String>, value: Value) {
1347 Arc::make_mut(&mut self.session.session_vars).insert(name.into(), value);
1348 }
1349
1350 pub fn timeout_ms(&self) -> u64 {
1352 self.query.timeout_ms
1353 }
1354
1355 pub fn set_timeout_ms(&mut self, timeout_ms: u64) {
1357 self.query.timeout_ms = timeout_ms;
1358 }
1359
1360 pub fn has_timeout(&self) -> bool {
1362 self.query.timeout_ms > 0
1363 }
1364
1365 pub fn view_depth(&self) -> usize {
1367 self.query.view_depth
1368 }
1369
1370 pub fn with_incremented_view_depth(&self) -> Self {
1374 let mut nested = self.clone();
1375 nested.query.view_depth += 1;
1376 nested.query.query_depth += 1;
1377 nested
1378 }
1379
1380 pub fn with_incremented_query_depth(&self) -> Self {
1383 let mut nested = self.clone();
1384 nested.query.query_depth += 1;
1385 nested
1386 }
1387
1388 pub fn outer_row(&self) -> Option<&FxHashMap<CompactArc<str>, Value>> {
1390 self.query.outer_row.as_ref()
1391 }
1392
1393 #[doc(hidden)]
1395 pub fn take_outer_row(&mut self) -> Option<FxHashMap<CompactArc<str>, Value>> {
1396 self.query.outer_row.take()
1397 }
1398
1399 #[doc(hidden)]
1401 pub fn query_depth(&self) -> usize {
1402 self.query.query_depth
1403 }
1404
1405 pub fn outer_columns(&self) -> Option<&[String]> {
1407 self.query.outer_columns.as_ref().map(|v| v.as_slice())
1408 }
1409
1410 pub fn with_outer_row(
1414 &self,
1415 outer_row: FxHashMap<CompactArc<str>, Value>,
1416 outer_columns: CompactArc<Vec<String>>,
1417 ) -> Self {
1418 let mut nested = self.with_incremented_query_depth();
1419 nested.query.outer_row = Some(outer_row);
1420 nested.query.outer_columns = Some(outer_columns);
1421 nested
1422 }
1423
1424 pub fn get_cte(&self, name: &str) -> Option<&CteData> {
1427 self.query
1428 .cte_data
1429 .as_ref()
1430 .and_then(|data| data.get(&name.to_lowercase()))
1431 }
1432
1433 #[inline]
1437 pub fn get_cte_by_lower(&self, name_lower: &str) -> Option<&CteData> {
1438 self.query
1439 .cte_data
1440 .as_ref()
1441 .and_then(|data| data.get(name_lower))
1442 }
1443
1444 #[doc(hidden)]
1447 pub fn get_cte_materialized_rows_by_lower(
1448 &self,
1449 name_lower: &str,
1450 ) -> Option<CompactArc<Vec<Row>>> {
1451 let (_, rows_with_ids, materialized) = self.get_cte_by_lower(name_lower)?;
1452 Some(
1453 materialized
1454 .get_or_init(|| {
1455 CompactArc::new(rows_with_ids.iter().map(|(_, row)| row.clone()).collect())
1456 })
1457 .clone(),
1458 )
1459 }
1460
1461 pub fn has_cte(&self, name: &str) -> bool {
1463 self.query
1464 .cte_data
1465 .as_ref()
1466 .is_some_and(|data| data.contains_key(&name.to_lowercase()))
1467 }
1468
1469 #[inline]
1472 pub fn has_cte_by_lower(&self, name_lower: &str) -> bool {
1473 self.query
1474 .cte_data
1475 .as_ref()
1476 .is_some_and(|data| data.contains_key(name_lower))
1477 }
1478
1479 pub fn with_cte_data(&self, cte_data: Arc<CteDataMap>) -> Self {
1482 let mut nested = self.clone();
1483 nested.query.cte_data = Some(cte_data);
1484 nested
1485 }
1486
1487 pub fn transaction_id(&self) -> Option<u64> {
1489 self.query.transaction_id
1490 }
1491
1492 pub fn set_transaction_id(&mut self, txn_id: u64) {
1494 self.query.transaction_id = Some(txn_id);
1495 if let Ok(txn_id) = i64::try_from(txn_id) {
1496 self.set_system_context_value("CURRENT_TRANSACTION_ID", Value::Integer(txn_id));
1497 }
1498 }
1499
1500 pub fn with_transaction_id(&self, txn_id: u64) -> Self {
1502 let mut nested = self.clone();
1503 nested.set_transaction_id(txn_id);
1504 nested
1505 }
1506
1507 #[doc(hidden)]
1509 pub fn set_request_id(&mut self, request_id: u64) -> Result<()> {
1510 let request_id = i64::try_from(request_id).map_err(|_| {
1511 radixdb_core::Error::invalid_argument("request ID exceeds the SQL INTEGER domain")
1512 })?;
1513 self.set_system_context_value("CURRENT_REQUEST_ID", Value::Integer(request_id));
1514 Ok(())
1515 }
1516
1517 #[doc(hidden)]
1519 pub fn set_job_context(
1520 &mut self,
1521 idempotency_key: &str,
1522 job_id: ObjectId,
1523 attempt: u32,
1524 scheduled_at: DateTime<Utc>,
1525 ) {
1526 self.set_system_context_value("CURRENT_IDEMPOTENCY_KEY", Value::text(idempotency_key));
1527 self.set_system_context_value("CURRENT_JOB_ID", Value::uuid(job_id.into_bytes()));
1528 self.set_system_context_value("CURRENT_JOB_ATTEMPT", Value::Integer(i64::from(attempt)));
1529 self.set_system_context_value("CURRENT_JOB_SCHEDULED_AT", Value::Timestamp(scheduled_at));
1530 }
1531
1532 fn set_system_context_value(&mut self, name: &str, value: Value) {
1533 Arc::make_mut(&mut self.query.named_params).insert(name.to_owned(), value);
1534 }
1535
1536 pub fn check_cancelled(&self) -> Result<()> {
1538 if self.is_cancelled() {
1539 Err(radixdb_core::Error::QueryCancelled)
1540 } else {
1541 Ok(())
1542 }
1543 }
1544}
1545
1546#[doc(hidden)]
1547pub struct ReferenceExpandExecutionGuard {
1548 active: Arc<AtomicUsize>,
1549}
1550
1551impl Drop for ReferenceExpandExecutionGuard {
1552 fn drop(&mut self) {
1553 self.active.fetch_sub(1, Ordering::AcqRel);
1554 }
1555}
1556
1557#[derive(Debug, Clone)]
1559pub struct CancellationHandle {
1560 cancelled: Arc<AtomicBool>,
1561 parent_cancelled: Option<Arc<AtomicBool>>,
1562}
1563
1564impl CancellationHandle {
1565 pub fn cancel(&self) {
1567 self.cancelled.store(true, Ordering::Relaxed);
1568 }
1569
1570 pub fn is_cancelled(&self) -> bool {
1572 self.cancelled.load(Ordering::Relaxed)
1573 || self
1574 .parent_cancelled
1575 .as_ref()
1576 .is_some_and(|cancelled| cancelled.load(Ordering::Relaxed))
1577 }
1578}
1579
1580impl radixdb_functions::FunctionCancellation for CancellationHandle {
1581 #[inline]
1582 fn is_cancelled(&self) -> bool {
1583 Self::is_cancelled(self)
1584 }
1585}
1586
1587struct TimeoutEntry {
1596 deadline: Instant,
1598 id: u64,
1600 cancel_handle: CancellationHandle,
1602 cancelled: Arc<AtomicBool>,
1604 timed_out: Arc<AtomicBool>,
1606}
1607
1608impl PartialEq for TimeoutEntry {
1609 fn eq(&self, other: &Self) -> bool {
1610 self.deadline == other.deadline && self.id == other.id
1611 }
1612}
1613
1614impl Eq for TimeoutEntry {}
1615
1616impl PartialOrd for TimeoutEntry {
1617 fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
1618 Some(self.cmp(other))
1619 }
1620}
1621
1622impl Ord for TimeoutEntry {
1623 fn cmp(&self, other: &Self) -> std::cmp::Ordering {
1624 other.deadline.cmp(&self.deadline)
1626 }
1627}
1628
1629struct TimeoutManagerState {
1631 timeouts: BinaryHeap<TimeoutEntry>,
1633}
1634
1635struct TimeoutManager {
1637 state: Mutex<TimeoutManagerState>,
1639 condvar: Condvar,
1641 next_id: AtomicU64,
1643}
1644
1645impl TimeoutManager {
1646 fn new() -> Arc<Self> {
1648 let manager = Arc::new(Self {
1649 state: Mutex::new(TimeoutManagerState {
1650 timeouts: BinaryHeap::new(),
1651 }),
1652 condvar: Condvar::new(),
1653 next_id: AtomicU64::new(1),
1654 });
1655
1656 let manager_clone = Arc::clone(&manager);
1658 std::thread::Builder::new()
1659 .name("radixdb-timeout-manager".to_string())
1660 .spawn(move || {
1661 manager_clone.run();
1662 })
1663 .expect("Failed to spawn timeout manager thread");
1664
1665 manager
1666 }
1667
1668 fn run(&self) {
1670 loop {
1671 let mut state = self.state.lock().unwrap();
1672
1673 let now = Instant::now();
1675 while let Some(entry) = state.timeouts.peek() {
1676 if entry.deadline <= now {
1677 let entry = state.timeouts.pop().unwrap();
1678 if !entry.cancelled.load(Ordering::Relaxed) {
1680 entry.timed_out.store(true, Ordering::Release);
1681 entry.cancel_handle.cancel();
1682 }
1683 } else {
1684 break;
1685 }
1686 }
1687
1688 let wait_duration = if let Some(entry) = state.timeouts.peek() {
1690 entry.deadline.saturating_duration_since(now)
1691 } else {
1692 Duration::from_secs(3600) };
1695
1696 if wait_duration.is_zero() {
1698 continue; }
1700 let (_state, _timeout_result) =
1701 self.condvar.wait_timeout(state, wait_duration).unwrap();
1702 }
1703 }
1704
1705 fn register(
1707 &self,
1708 timeout_ms: u64,
1709 cancel_handle: CancellationHandle,
1710 cancelled: Arc<AtomicBool>,
1711 timed_out: Arc<AtomicBool>,
1712 ) -> u64 {
1713 let id = self.next_id.fetch_add(1, Ordering::Relaxed);
1714 let deadline = Instant::now() + Duration::from_millis(timeout_ms);
1715
1716 let entry = TimeoutEntry {
1717 deadline,
1718 id,
1719 cancel_handle,
1720 cancelled,
1721 timed_out,
1722 };
1723
1724 let mut state = self.state.lock().unwrap();
1725 let was_empty = state.timeouts.is_empty();
1726 let is_earliest = state.timeouts.peek().is_none_or(|e| deadline < e.deadline);
1727
1728 state.timeouts.push(entry);
1729
1730 if was_empty || is_earliest {
1732 self.condvar.notify_one();
1733 }
1734
1735 id
1736 }
1737
1738 fn unregister(&self, id: u64) {
1739 let mut state = self.state.lock().unwrap();
1740 let removed_earliest = state.timeouts.peek().is_some_and(|entry| entry.id == id);
1741 state.timeouts.retain(|entry| entry.id != id);
1742 drop(state);
1743 if removed_earliest {
1744 self.condvar.notify_one();
1745 }
1746 }
1747}
1748
1749fn global_timeout_manager() -> &'static Arc<TimeoutManager> {
1751 use std::sync::OnceLock;
1752 static MANAGER: OnceLock<Arc<TimeoutManager>> = OnceLock::new();
1753 MANAGER.get_or_init(TimeoutManager::new)
1754}
1755
1756#[doc(hidden)]
1757pub fn pending_timeout_count_for(ctx: &ExecutionContext) -> usize {
1758 global_timeout_manager()
1759 .state
1760 .lock()
1761 .unwrap()
1762 .timeouts
1763 .iter()
1764 .filter(|entry| Arc::ptr_eq(&entry.timed_out, &ctx.query.timed_out))
1765 .count()
1766}
1767
1768pub struct TimeoutGuard {
1771 registration_id: u64,
1772 cancelled: Arc<AtomicBool>,
1774}
1775
1776impl TimeoutGuard {
1777 pub fn new(ctx: &ExecutionContext) -> Option<Self> {
1780 let timeout_ms = ctx.timeout_ms();
1781 if timeout_ms == 0 {
1782 return None;
1783 }
1784
1785 let cancel_handle = ctx.cancellation_handle();
1786 let cancelled = Arc::new(AtomicBool::new(false));
1787
1788 let registration_id = global_timeout_manager().register(
1790 timeout_ms,
1791 cancel_handle,
1792 Arc::clone(&cancelled),
1793 Arc::clone(&ctx.query.timed_out),
1794 );
1795
1796 Some(Self {
1797 registration_id,
1798 cancelled,
1799 })
1800 }
1801}
1802
1803impl Drop for TimeoutGuard {
1804 fn drop(&mut self) {
1805 self.cancelled.store(true, Ordering::Relaxed);
1807 global_timeout_manager().unregister(self.registration_id);
1808 }
1809}
1810
1811pub struct ExecutionContextBuilder {
1813 ctx: ExecutionContext,
1814}
1815
1816impl ExecutionContextBuilder {
1817 pub fn new() -> Self {
1819 Self {
1820 ctx: ExecutionContext::new(),
1821 }
1822 }
1823
1824 pub fn params(mut self, params: ParamVec) -> Self {
1826 self.ctx.query.params = CompactArc::new(params);
1827 self
1828 }
1829
1830 pub fn param(mut self, value: Value) -> Self {
1832 let mut v = (*self.ctx.query.params).clone();
1833 v.push(value);
1834 self.ctx.query.params = CompactArc::new(v);
1835 self
1836 }
1837
1838 pub fn named_param(mut self, name: impl Into<String>, value: Value) -> Self {
1840 self.ctx.set_named_param(name, value);
1841 self
1842 }
1843
1844 pub fn auto_commit(mut self, auto_commit: bool) -> Self {
1846 self.ctx.session.auto_commit = auto_commit;
1847 self
1848 }
1849
1850 pub fn database(mut self, database: impl Into<String>) -> Self {
1852 self.ctx.session.current_database = Arc::new(Some(database.into()));
1853 self
1854 }
1855
1856 pub fn principal_id(mut self, principal_id: ObjectId) -> Self {
1858 self.ctx = self.ctx.with_principal_id(principal_id);
1859 self
1860 }
1861
1862 pub fn session_var(mut self, name: impl Into<String>, value: Value) -> Self {
1864 let mut variables = (*self.ctx.session.session_vars).clone();
1865 variables.insert(name.into(), value);
1866 self.ctx.session.session_vars = Arc::new(variables);
1867 self
1868 }
1869
1870 pub fn timeout_ms(mut self, timeout_ms: u64) -> Self {
1872 self.ctx.query.timeout_ms = timeout_ms;
1873 self
1874 }
1875
1876 #[doc(hidden)]
1879 pub fn join_hash_state_max_bytes(mut self, max_bytes: usize) -> Self {
1880 self.ctx.query.join_hash_state_max_bytes = max_bytes;
1881 self
1882 }
1883
1884 pub fn build(self) -> ExecutionContext {
1886 self.ctx
1887 }
1888}
1889
1890impl Default for ExecutionContextBuilder {
1891 fn default() -> Self {
1892 Self::new()
1893 }
1894}
1895
1896#[cfg(test)]
1897mod tests {
1898 use super::*;
1899 use rustc_hash::FxHashMap;
1900
1901 struct TestHashObserver {
1902 retained_bytes: usize,
1903 observed: usize,
1904 }
1905
1906 impl TestHashObserver {
1907 fn new(retained_bytes: usize) -> Self {
1908 Self {
1909 retained_bytes,
1910 observed: 0,
1911 }
1912 }
1913 }
1914
1915 impl crate::hash_table::JoinHashObserver for TestHashObserver {
1916 fn insert_raw_hash(&mut self, _hash: u64) {
1917 self.observed += 1;
1918 }
1919
1920 fn retained_bytes(&self) -> usize {
1921 self.retained_bytes
1922 }
1923 }
1924
1925 #[test]
1926 fn test_context_new() {
1927 let ctx = ExecutionContext::new();
1928 assert_eq!(ctx.param_count(), 0);
1929 assert!(ctx.auto_commit());
1930 assert!(!ctx.is_cancelled());
1931 }
1932
1933 #[test]
1934 fn test_context_with_params() {
1935 let ctx = ExecutionContext::with_params(smallvec::smallvec![
1936 Value::Integer(1),
1937 Value::text("hello")
1938 ]);
1939 assert_eq!(ctx.param_count(), 2);
1940 assert_eq!(ctx.get_param(1), Some(&Value::Integer(1)));
1941 assert_eq!(ctx.get_param(2), Some(&Value::text("hello")));
1942 assert_eq!(ctx.get_param(0), None); assert_eq!(ctx.get_param(3), None); }
1945
1946 #[test]
1947 fn test_context_named_params() {
1948 let mut params = FxHashMap::default();
1949 params.insert("name".to_string(), Value::text("Alice"));
1950 params.insert("age".to_string(), Value::Integer(30));
1951
1952 let ctx = ExecutionContext::with_named_params(params);
1953 assert_eq!(ctx.get_named_param("name"), Some(&Value::text("Alice")));
1954 assert_eq!(ctx.get_named_param("age"), Some(&Value::Integer(30)));
1955 assert_eq!(ctx.get_named_param("unknown"), None);
1956 }
1957
1958 #[test]
1959 fn test_context_cancellation() {
1960 let ctx = ExecutionContext::new();
1961 assert!(!ctx.is_cancelled());
1962
1963 let handle = ctx.cancellation_handle();
1964 assert!(!handle.is_cancelled());
1965
1966 handle.cancel();
1967 assert!(ctx.is_cancelled());
1968 assert!(handle.is_cancelled());
1969 }
1970
1971 #[test]
1972 fn test_context_check_cancelled() {
1973 let ctx = ExecutionContext::new();
1974 assert!(ctx.check_cancelled().is_ok());
1975
1976 ctx.cancel();
1977 assert!(ctx.check_cancelled().is_err());
1978 }
1979
1980 #[test]
1981 fn test_context_session_vars() {
1982 let mut ctx = ExecutionContext::new();
1983 ctx.set_session_var("timezone", Value::text("UTC"));
1984
1985 assert_eq!(ctx.get_session_var("timezone"), Some(&Value::text("UTC")));
1986 assert_eq!(ctx.get_session_var("unknown"), None);
1987 }
1988
1989 #[test]
1990 fn test_context_builder() {
1991 let ctx = ExecutionContextBuilder::new()
1992 .params(smallvec::smallvec![Value::Integer(1)])
1993 .param(Value::Integer(2))
1994 .named_param("name", Value::text("test"))
1995 .auto_commit(false)
1996 .database("mydb")
1997 .timeout_ms(5000)
1998 .build();
1999
2000 assert_eq!(ctx.param_count(), 2);
2001 assert_eq!(ctx.get_param(1), Some(&Value::Integer(1)));
2002 assert_eq!(ctx.get_param(2), Some(&Value::Integer(2)));
2003 assert_eq!(ctx.get_named_param("name"), Some(&Value::text("test")));
2004 assert!(!ctx.auto_commit());
2005 assert_eq!(ctx.current_database(), Some("mydb"));
2006 assert_eq!(ctx.timeout_ms(), 5000);
2007 }
2008
2009 #[cfg(feature = "test-hooks")]
2010 #[test]
2011 fn request_local_join_hash_cache_evicts_to_its_byte_budget() {
2012 let retained = crate::hash_table::JoinHashTable::estimated_retained_bytes(1)
2013 .expect("one-row hash state size");
2014 let ctx = ExecutionContextBuilder::new()
2015 .join_hash_state_max_bytes(retained)
2016 .build();
2017 let first = CompactArc::new(vec![Row::from_values(vec![Value::Integer(1)])]);
2018 let second = CompactArc::new(vec![Row::from_values(vec![Value::Integer(2)])]);
2019
2020 radixdb_storage::instrumentation::begin_join_execution_probe();
2021 drop(
2022 ctx.join_hash_state_for(CompactArc::clone(&first), &[0], true)
2023 .expect("first state"),
2024 );
2025 drop(
2026 ctx.join_hash_state_for(CompactArc::clone(&second), &[0], true)
2027 .expect("second state"),
2028 );
2029 drop(
2030 ctx.join_hash_state_for(CompactArc::clone(&second), &[0], true)
2031 .expect("second state reuse"),
2032 );
2033 drop(
2034 ctx.join_hash_state_for(CompactArc::clone(&first), &[0], true)
2035 .expect("first state rebuild after eviction"),
2036 );
2037 let probe = radixdb_storage::instrumentation::end_join_execution_probe();
2038
2039 assert_eq!(probe.hash_state_builds, 3);
2040 assert_eq!(probe.hash_state_reuses, 1);
2041 }
2042
2043 #[test]
2044 fn active_and_cached_hash_states_share_one_request_budget() {
2045 let retained = JoinHashTable::estimated_retained_bytes(1).unwrap();
2046 let ctx = ExecutionContextBuilder::new()
2047 .join_hash_state_max_bytes(retained)
2048 .build();
2049 let cached_rows = CompactArc::new(vec![Row::from_values(vec![Value::Integer(1)])]);
2050
2051 let active_cached = ctx
2052 .join_hash_state_for(CompactArc::clone(&cached_rows), &[0], true)
2053 .expect("first state fits common budget");
2054 assert_eq!(ctx.retained_join_memory_bytes(), retained);
2055
2056 let rejected = ctx.join_hash_state_for(
2059 CompactArc::new(vec![Row::from_values(vec![Value::Integer(2)])]),
2060 &[0],
2061 false,
2062 );
2063 assert!(rejected.is_none());
2064 assert_eq!(ctx.retained_join_memory_bytes(), retained);
2065
2066 drop(active_cached);
2067 assert_eq!(ctx.retained_join_memory_bytes(), 0);
2068 let admitted = ctx
2069 .join_hash_state_for(
2070 CompactArc::new(vec![Row::from_values(vec![Value::Integer(2)])]),
2071 &[0],
2072 false,
2073 )
2074 .expect("released request budget admits next transient state");
2075 assert_eq!(ctx.retained_join_memory_bytes(), retained);
2076 drop(admitted);
2077 assert_eq!(ctx.retained_join_memory_bytes(), 0);
2078 }
2079
2080 #[test]
2081 fn context_clones_cannot_each_spend_the_join_budget() {
2082 let retained = JoinHashTable::estimated_retained_bytes(1).unwrap();
2083 let ctx = ExecutionContextBuilder::new()
2084 .join_hash_state_max_bytes(retained)
2085 .build();
2086 let clone = ctx.clone();
2087 let first = ctx
2088 .join_hash_state_for(
2089 CompactArc::new(vec![Row::from_values(vec![Value::Integer(1)])]),
2090 &[0],
2091 false,
2092 )
2093 .unwrap();
2094 assert!(clone
2095 .join_hash_state_for(
2096 CompactArc::new(vec![Row::from_values(vec![Value::Integer(2)])]),
2097 &[0],
2098 false,
2099 )
2100 .is_none());
2101 drop(first);
2102 assert_eq!(clone.retained_join_memory_bytes(), 0);
2103 }
2104
2105 #[test]
2106 fn bloom_and_hash_share_the_same_join_memory_reservation() {
2107 let rows: CompactArc<Vec<Row>> = CompactArc::new(
2108 (0..100)
2109 .map(|value| Row::from_values(vec![Value::Integer(value)]))
2110 .collect(),
2111 );
2112 let mut builder = TestHashObserver::new(256);
2113 let hash_bytes = JoinHashTable::estimated_retained_bytes(rows.len()).unwrap();
2114 let bloom_bytes = crate::hash_table::JoinHashObserver::retained_bytes(&builder);
2115 assert!(bloom_bytes > 0);
2116
2117 let too_small = ExecutionContextBuilder::new()
2118 .join_hash_state_max_bytes(hash_bytes)
2119 .build();
2120 assert!(too_small
2121 .join_hash_state_with_bloom_for(CompactArc::clone(&rows), &[0], &mut builder)
2122 .is_none());
2123 assert_eq!(too_small.retained_join_memory_bytes(), 0);
2124
2125 let exact = ExecutionContextBuilder::new()
2126 .join_hash_state_max_bytes(hash_bytes + bloom_bytes)
2127 .build();
2128 let state = exact
2129 .join_hash_state_with_bloom_for(rows, &[0], &mut builder)
2130 .expect("combined hash and bloom bytes fit exactly");
2131 assert_eq!(builder.observed, 100);
2132 assert_eq!(exact.retained_join_memory_bytes(), hash_bytes + bloom_bytes);
2133 drop(state);
2134 assert_eq!(exact.retained_join_memory_bytes(), 0);
2135 }
2136
2137 #[test]
2138 fn r6_l01_b_completed_timeout_is_unregistered_immediately() {
2139 let ctx = ExecutionContextBuilder::new().timeout_ms(60_000).build();
2140 let guard = TimeoutGuard::new(&ctx).expect("timeout guard");
2141 assert_eq!(pending_timeout_count_for(&ctx), 1);
2142 drop(guard);
2143 assert_eq!(pending_timeout_count_for(&ctx), 0);
2144 assert!(!ctx.is_cancelled());
2145 }
2146}