reifydb_core/value/column/pool/
scoped.rs1use super::{
10 Pools,
11 config::PoolConfig,
12 thread_local::{clear_thread_pools, get_thread_pools, set_thread_pools},
13};
14
15pub struct ScopedPools {
17 previous: Option<Pools>,
18}
19
20impl ScopedPools {
21 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 pub fn default() -> Self {
33 Self::new(Pools::default())
34 }
35
36 pub fn with_config(config: PoolConfig) -> Self {
38 Self::new(Pools::new(config.max_pool_size))
39 }
40
41 pub fn test() -> Self {
43 Self::with_config(PoolConfig::test())
44 }
45
46 pub fn production() -> Self {
48 Self::with_config(PoolConfig::production())
49 }
50
51 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
66pub 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
76pub fn with_default_pools<F, R>(f: F) -> R
78where
79 F: FnOnce() -> R,
80{
81 let _guard = ScopedPools::default();
82 f()
83}
84
85pub 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 clear_thread_pools();
103 assert!(!has_thread_pools());
104
105 {
106 let _guard = ScopedPools::default();
107 assert!(has_thread_pools());
109 }
110
111 assert!(!has_thread_pools());
113 }
114
115 #[test]
116 fn test_scoped_pools_restore_previous() {
117 let initial_pools = Pools::new(8);
119 set_thread_pools(initial_pools);
120 assert!(has_thread_pools());
121
122 {
123 let _guard = ScopedPools::new(Pools::new(16));
125 assert!(has_thread_pools());
126 }
128
129 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 assert!(has_thread_pools());
184
185 drop(_outer);
187 assert!(!has_thread_pools());
188 }
189}