reifydb_core/value/column/pool/
scoped.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//! RAII guards for scoped pool management
5//!
6//! This module provides RAII guards that temporarily set thread-local pools
7//! and restore the previous state when dropped.
8
9use super::{
10	Pools,
11	config::PoolConfig,
12	thread_local::{clear_thread_pools, get_thread_pools, set_thread_pools},
13};
14
15/// RAII guard that sets thread-local pools and restores previous state on drop
16pub struct ScopedPools {
17	previous: Option<Pools>,
18}
19
20impl ScopedPools {
21	/// Create a new scoped pools guard
22	/// Sets the given pools as thread-local and saves the previous state
23	pub fn new(pools: Pools) -> Self {
24		let previous = get_thread_pools();
25		set_thread_pools(pools);
26		Self {
27			previous,
28		}
29	}
30
31	/// Create a scoped guard with new default pools
32	pub fn default() -> Self {
33		Self::new(Pools::default())
34	}
35
36	/// Create a scoped guard with custom configuration
37	pub fn with_config(config: PoolConfig) -> Self {
38		Self::new(Pools::new(config.max_pool_size))
39	}
40
41	/// Create a scoped guard with test configuration
42	pub fn test() -> Self {
43		Self::with_config(PoolConfig::test())
44	}
45
46	/// Create a scoped guard with production configuration
47	pub fn production() -> Self {
48		Self::with_config(PoolConfig::production())
49	}
50
51	/// Create a scoped guard with development configuration
52	pub fn development() -> Self {
53		Self::with_config(PoolConfig::development())
54	}
55}
56
57impl Drop for ScopedPools {
58	fn drop(&mut self) {
59		match self.previous.take() {
60			Some(p) => set_thread_pools(p),
61			None => clear_thread_pools(),
62		}
63	}
64}
65
66/// Execute a function with scoped pools
67/// This is a convenience function that creates a ScopedPools guard
68pub fn with_scoped_pools<F, R>(pools: Pools, f: F) -> R
69where
70	F: FnOnce() -> R,
71{
72	let _guard = ScopedPools::new(pools);
73	f()
74}
75
76/// Execute a function with default scoped pools
77pub fn with_default_pools<F, R>(f: F) -> R
78where
79	F: FnOnce() -> R,
80{
81	let _guard = ScopedPools::default();
82	f()
83}
84
85/// Execute a function with test scoped pools
86pub fn with_test_pools<F, R>(f: F) -> R
87where
88	F: FnOnce() -> R,
89{
90	let _guard = ScopedPools::test();
91	f()
92}
93
94#[cfg(test)]
95mod tests {
96	use super::*;
97	use crate::value::column::pool::thread_local::has_thread_pools;
98
99	#[test]
100	fn test_scoped_pools_basic() {
101		// Start with no pools
102		clear_thread_pools();
103		assert!(!has_thread_pools());
104
105		{
106			let _guard = ScopedPools::default();
107			// Pools should be set within scope
108			assert!(has_thread_pools());
109		}
110
111		// Pools should be cleared after scope
112		assert!(!has_thread_pools());
113	}
114
115	#[test]
116	fn test_scoped_pools_restore_previous() {
117		// Set initial pools
118		let initial_pools = Pools::new(8);
119		set_thread_pools(initial_pools);
120		assert!(has_thread_pools());
121
122		{
123			// Use different pools in scope
124			let _guard = ScopedPools::new(Pools::new(16));
125			assert!(has_thread_pools());
126			// Can't easily verify it's different pools, but it is
127		}
128
129		// Should restore initial pools
130		assert!(has_thread_pools());
131
132		clear_thread_pools();
133	}
134
135	#[test]
136	fn test_with_scoped_pools() {
137		clear_thread_pools();
138
139		let result = with_scoped_pools(Pools::default(), || {
140			assert!(has_thread_pools());
141			42
142		});
143
144		assert_eq!(result, 42);
145		assert!(!has_thread_pools());
146	}
147
148	#[test]
149	fn test_with_default_pools() {
150		clear_thread_pools();
151
152		with_default_pools(|| {
153			assert!(has_thread_pools());
154		});
155
156		assert!(!has_thread_pools());
157	}
158
159	#[test]
160	fn test_with_test_pools() {
161		clear_thread_pools();
162
163		with_test_pools(|| {
164			assert!(has_thread_pools());
165		});
166
167		assert!(!has_thread_pools());
168	}
169
170	#[test]
171	fn test_nested_scoped_pools() {
172		clear_thread_pools();
173
174		let _outer = ScopedPools::default();
175		assert!(has_thread_pools());
176
177		{
178			let _inner = ScopedPools::test();
179			assert!(has_thread_pools());
180		}
181
182		// Should still have outer pools
183		assert!(has_thread_pools());
184
185		// Manually drop to test
186		drop(_outer);
187		assert!(!has_thread_pools());
188	}
189}