1 #![warn(rust_2018_idioms)]
2 #![cfg(feature = "full")]
3 
4 use std::mem;
5 use std::ops::Drop;
6 use std::sync::atomic::{AtomicU32, Ordering};
7 use std::time::Duration;
8 use tokio::runtime;
9 use tokio::sync::{OnceCell, SetError};
10 use tokio::time;
11 
func1() -> u3212 async fn func1() -> u32 {
13     5
14 }
15 
func2() -> u3216 async fn func2() -> u32 {
17     time::sleep(Duration::from_millis(1)).await;
18     10
19 }
20 
func_err() -> Result<u32, ()>21 async fn func_err() -> Result<u32, ()> {
22     Err(())
23 }
24 
func_ok() -> Result<u32, ()>25 async fn func_ok() -> Result<u32, ()> {
26     Ok(10)
27 }
28 
func_panic() -> u3229 async fn func_panic() -> u32 {
30     time::sleep(Duration::from_millis(1)).await;
31     panic!();
32 }
33 
sleep_and_set() -> u3234 async fn sleep_and_set() -> u32 {
35     // Simulate sleep by pausing time and waiting for another thread to
36     // resume clock when calling `set`, then finding the cell being initialized
37     // by this call
38     time::sleep(Duration::from_millis(2)).await;
39     5
40 }
41 
advance_time_and_set(cell: &'static OnceCell<u32>, v: u32) -> Result<(), SetError<u32>>42 async fn advance_time_and_set(cell: &'static OnceCell<u32>, v: u32) -> Result<(), SetError<u32>> {
43     time::advance(Duration::from_millis(1)).await;
44     cell.set(v)
45 }
46 
47 #[test]
get_or_init()48 fn get_or_init() {
49     let rt = runtime::Builder::new_current_thread()
50         .enable_time()
51         .start_paused(true)
52         .build()
53         .unwrap();
54 
55     static ONCE: OnceCell<u32> = OnceCell::const_new();
56 
57     rt.block_on(async {
58         let handle1 = rt.spawn(async { ONCE.get_or_init(func1).await });
59         let handle2 = rt.spawn(async { ONCE.get_or_init(func2).await });
60 
61         time::advance(Duration::from_millis(1)).await;
62         time::resume();
63 
64         let result1 = handle1.await.unwrap();
65         let result2 = handle2.await.unwrap();
66 
67         assert_eq!(*result1, 5);
68         assert_eq!(*result2, 5);
69     });
70 }
71 
72 #[test]
get_or_init_panic()73 fn get_or_init_panic() {
74     let rt = runtime::Builder::new_current_thread()
75         .enable_time()
76         .build()
77         .unwrap();
78 
79     static ONCE: OnceCell<u32> = OnceCell::const_new();
80 
81     rt.block_on(async {
82         time::pause();
83 
84         let handle1 = rt.spawn(async { ONCE.get_or_init(func1).await });
85         let handle2 = rt.spawn(async { ONCE.get_or_init(func_panic).await });
86 
87         time::advance(Duration::from_millis(1)).await;
88 
89         let result1 = handle1.await.unwrap();
90         let result2 = handle2.await.unwrap();
91 
92         assert_eq!(*result1, 5);
93         assert_eq!(*result2, 5);
94     });
95 }
96 
97 #[test]
set_and_get()98 fn set_and_get() {
99     let rt = runtime::Builder::new_current_thread()
100         .enable_time()
101         .build()
102         .unwrap();
103 
104     static ONCE: OnceCell<u32> = OnceCell::const_new();
105 
106     rt.block_on(async {
107         let _ = rt.spawn(async { ONCE.set(5) }).await;
108         let value = ONCE.get().unwrap();
109         assert_eq!(*value, 5);
110     });
111 }
112 
113 #[test]
get_uninit()114 fn get_uninit() {
115     static ONCE: OnceCell<u32> = OnceCell::const_new();
116     let uninit = ONCE.get();
117     assert!(uninit.is_none());
118 }
119 
120 #[test]
set_twice()121 fn set_twice() {
122     static ONCE: OnceCell<u32> = OnceCell::const_new();
123 
124     let first = ONCE.set(5);
125     assert_eq!(first, Ok(()));
126     let second = ONCE.set(6);
127     assert!(second.err().unwrap().is_already_init_err());
128 }
129 
130 #[test]
set_while_initializing()131 fn set_while_initializing() {
132     let rt = runtime::Builder::new_current_thread()
133         .enable_time()
134         .build()
135         .unwrap();
136 
137     static ONCE: OnceCell<u32> = OnceCell::const_new();
138 
139     rt.block_on(async {
140         time::pause();
141 
142         let handle1 = rt.spawn(async { ONCE.get_or_init(sleep_and_set).await });
143         let handle2 = rt.spawn(async { advance_time_and_set(&ONCE, 10).await });
144 
145         time::advance(Duration::from_millis(2)).await;
146 
147         let result1 = handle1.await.unwrap();
148         let result2 = handle2.await.unwrap();
149 
150         assert_eq!(*result1, 5);
151         assert!(result2.err().unwrap().is_initializing_err());
152     });
153 }
154 
155 #[test]
get_or_try_init()156 fn get_or_try_init() {
157     let rt = runtime::Builder::new_current_thread()
158         .enable_time()
159         .start_paused(true)
160         .build()
161         .unwrap();
162 
163     static ONCE: OnceCell<u32> = OnceCell::const_new();
164 
165     rt.block_on(async {
166         let handle1 = rt.spawn(async { ONCE.get_or_try_init(func_err).await });
167         let handle2 = rt.spawn(async { ONCE.get_or_try_init(func_ok).await });
168 
169         time::advance(Duration::from_millis(1)).await;
170         time::resume();
171 
172         let result1 = handle1.await.unwrap();
173         assert!(result1.is_err());
174 
175         let result2 = handle2.await.unwrap();
176         assert_eq!(*result2.unwrap(), 10);
177     });
178 }
179 
180 #[test]
drop_cell()181 fn drop_cell() {
182     static NUM_DROPS: AtomicU32 = AtomicU32::new(0);
183 
184     struct Foo {}
185 
186     let fooer = Foo {};
187 
188     impl Drop for Foo {
189         fn drop(&mut self) {
190             NUM_DROPS.fetch_add(1, Ordering::Release);
191         }
192     }
193 
194     {
195         let once_cell = OnceCell::new();
196         let prev = once_cell.set(fooer);
197         assert!(prev.is_ok())
198     }
199     assert!(NUM_DROPS.load(Ordering::Acquire) == 1);
200 }
201 
202 #[test]
drop_cell_new_with()203 fn drop_cell_new_with() {
204     static NUM_DROPS: AtomicU32 = AtomicU32::new(0);
205 
206     struct Foo {}
207 
208     let fooer = Foo {};
209 
210     impl Drop for Foo {
211         fn drop(&mut self) {
212             NUM_DROPS.fetch_add(1, Ordering::Release);
213         }
214     }
215 
216     {
217         let once_cell = OnceCell::new_with(Some(fooer));
218         assert!(once_cell.initialized());
219     }
220     assert!(NUM_DROPS.load(Ordering::Acquire) == 1);
221 }
222 
223 #[test]
drop_into_inner()224 fn drop_into_inner() {
225     static NUM_DROPS: AtomicU32 = AtomicU32::new(0);
226 
227     struct Foo {}
228 
229     let fooer = Foo {};
230 
231     impl Drop for Foo {
232         fn drop(&mut self) {
233             NUM_DROPS.fetch_add(1, Ordering::Release);
234         }
235     }
236 
237     let once_cell = OnceCell::new();
238     assert!(once_cell.set(fooer).is_ok());
239     let fooer = once_cell.into_inner();
240     let count = NUM_DROPS.load(Ordering::Acquire);
241     assert!(count == 0);
242     drop(fooer);
243     let count = NUM_DROPS.load(Ordering::Acquire);
244     assert!(count == 1);
245 }
246 
247 #[test]
drop_into_inner_new_with()248 fn drop_into_inner_new_with() {
249     static NUM_DROPS: AtomicU32 = AtomicU32::new(0);
250 
251     struct Foo {}
252 
253     let fooer = Foo {};
254 
255     impl Drop for Foo {
256         fn drop(&mut self) {
257             NUM_DROPS.fetch_add(1, Ordering::Release);
258         }
259     }
260 
261     let once_cell = OnceCell::new_with(Some(fooer));
262     let fooer = once_cell.into_inner();
263     let count = NUM_DROPS.load(Ordering::Acquire);
264     assert!(count == 0);
265     mem::drop(fooer);
266     let count = NUM_DROPS.load(Ordering::Acquire);
267     assert!(count == 1);
268 }
269 
270 #[test]
from()271 fn from() {
272     let cell = OnceCell::from(2);
273     assert_eq!(*cell.get().unwrap(), 2);
274 }
275