reifydb_core/value/column/pool/
thread_local.rs

1// Copyright (c) reifydb.com 2025
2// This file is licensed under the AGPL-3.0-or-later, see license.md file
3
4//! Thread-local pool storage for avoiding explicit pool passing
5//!
6//! This module provides thread-local storage for pools, allowing ColumnData
7//! operations to access pools without needing to pass them through the entire
8//! call stack.
9
10use std::cell::RefCell;
11
12use super::Pools;
13
14thread_local! {
15    /// Thread-local storage for pools instance
16    static THREAD_POOLS: RefCell<Option<Pools>> = RefCell::new(None);
17}
18
19/// Set the thread-local pools instance for the current thread
20pub fn set_thread_pools(pools: Pools) {
21	THREAD_POOLS.with(|p| {
22		*p.borrow_mut() = Some(pools);
23	});
24}
25
26/// Get a clone of the thread-local pools instance
27/// Returns None if not initialized
28pub fn get_thread_pools() -> Option<Pools> {
29	THREAD_POOLS.with(|p| p.borrow().clone())
30}
31
32/// Get the thread-local pools instance or panic with helpful message
33pub fn thread_pools() -> Pools {
34	get_thread_pools()
35		.expect("Thread-local pools not initialized. Call init_thread_pools() or set_thread_pools() first.")
36}
37
38/// Execute a function with access to thread-local pools
39pub fn with_thread_pools<F, R>(f: F) -> R
40where
41	F: FnOnce(&Pools) -> R,
42{
43	THREAD_POOLS.with(|p| {
44		let pools = p.borrow();
45		let pools = pools.as_ref().expect(
46			"Thread-local pools not initialized. Call init_thread_pools() or set_thread_pools() first.",
47		);
48		f(pools)
49	})
50}
51
52/// Clear thread-local pools for the current thread
53pub fn clear_thread_pools() {
54	THREAD_POOLS.with(|p| {
55		*p.borrow_mut() = None;
56	});
57}
58
59/// Check if thread-local pools are initialized
60pub fn has_thread_pools() -> bool {
61	THREAD_POOLS.with(|p| p.borrow().is_some())
62}
63
64/// Execute a closure with temporary thread-local pools
65/// The previous pools state is restored after the closure completes
66pub fn with_temporary_pools<F, R>(pools: Pools, f: F) -> R
67where
68	F: FnOnce() -> R,
69{
70	let previous = get_thread_pools();
71	set_thread_pools(pools);
72
73	// Use defer pattern to ensure cleanup even on panic
74	struct Cleanup(Option<Pools>);
75	impl Drop for Cleanup {
76		fn drop(&mut self) {
77			match self.0.take() {
78				Some(p) => set_thread_pools(p),
79				None => clear_thread_pools(),
80			}
81		}
82	}
83	let _cleanup = Cleanup(previous);
84
85	f()
86}
87
88#[cfg(test)]
89mod tests {
90	use super::*;
91	use crate::value::column::pool::allocator::PoolAllocator;
92
93	#[test]
94	fn test_thread_local_basic() {
95		// Initially not set
96		assert!(!has_thread_pools());
97		assert!(get_thread_pools().is_none());
98
99		// Set pools
100		let pools = Pools::default();
101		set_thread_pools(pools.clone());
102
103		// Now should be available
104		assert!(has_thread_pools());
105		assert!(get_thread_pools().is_some());
106
107		// Can access via thread_pools()
108		let retrieved = thread_pools();
109		// Can't compare Pools directly, but we got one
110
111		// Clear pools
112		clear_thread_pools();
113		assert!(!has_thread_pools());
114	}
115
116	#[test]
117	fn test_with_thread_pools() {
118		let pools = Pools::default();
119		set_thread_pools(pools);
120
121		let result = with_thread_pools(|p| {
122			// Access pools inside closure
123			p.bool_pool().stats().available
124		});
125
126		// Should get some stats
127		assert_eq!(result, 0); // Empty pool initially
128
129		clear_thread_pools();
130	}
131
132	#[test]
133	fn test_temporary_pools() {
134		// Set initial pools
135		let pools1 = Pools::new(16);
136		set_thread_pools(pools1.clone());
137		assert!(has_thread_pools());
138
139		// Use temporary pools
140		let pools2 = Pools::new(32);
141		let result = with_temporary_pools(pools2, || {
142			// Should have temporary pools here
143			thread_pools().bool_pool().stats().available
144		});
145
146		// Original pools should be restored
147		assert!(has_thread_pools());
148		// Can't easily verify it's the same pools instance, but it's
149		// restored
150
151		clear_thread_pools();
152	}
153
154	#[test]
155	fn test_temporary_pools_with_none() {
156		// Start with no pools
157		assert!(!has_thread_pools());
158
159		// Use temporary pools
160		let pools = Pools::default();
161		with_temporary_pools(pools, || {
162			assert!(has_thread_pools());
163		});
164
165		// Should be back to no pools
166		assert!(!has_thread_pools());
167	}
168
169	#[test]
170	#[should_panic(expected = "Thread-local pools not initialized")]
171	fn test_thread_pools_panics_when_not_set() {
172		clear_thread_pools();
173		thread_pools(); // Should panic
174	}
175}
reifydb_core/value/column/pool/thread_local.rs

reifydb_core/value/column/pool/
thread_local.rs
