Skip to main content

turbopack_trace_server/
store_container.rs

1use std::{
2    ops::{Deref, DerefMut},
3    sync::{
4        RwLock, RwLockReadGuard, RwLockWriteGuard,
5        atomic::{AtomicBool, Ordering},
6    },
7};
8
9use crate::store::Store;
10
11pub struct StoreContainer {
12    store: RwLock<StoreWithGeneration>,
13    want_to_read: AtomicBool,
14}
15
16struct StoreWithGeneration {
17    store: Store,
18    generation: usize,
19}
20
21impl Default for StoreContainer {
22    fn default() -> Self {
23        Self::new()
24    }
25}
26
27impl StoreContainer {
28    pub fn new() -> Self {
29        Self {
30            store: RwLock::new(StoreWithGeneration {
31                store: Store::new(),
32                generation: 0,
33            }),
34            want_to_read: AtomicBool::new(false),
35        }
36    }
37
38    pub fn read(&self) -> StoreReadGuard<'_> {
39        if let Ok(guard) = self.store.try_read() {
40            return StoreReadGuard { guard };
41        }
42        self.want_to_read.store(true, Ordering::Relaxed);
43        let guard = StoreReadGuard {
44            guard: self.store.read().unwrap(),
45        };
46        self.want_to_read.store(false, Ordering::Relaxed);
47        guard
48    }
49
50    pub fn write(&self) -> StoreWriteGuard<'_> {
51        let mut guard = self.store.write().unwrap();
52        guard.generation += 1;
53        StoreWriteGuard { guard }
54    }
55
56    pub fn want_to_read(&self) -> bool {
57        self.want_to_read.load(Ordering::Relaxed)
58    }
59}
60
61pub struct StoreReadGuard<'a> {
62    guard: RwLockReadGuard<'a, StoreWithGeneration>,
63}
64
65impl StoreReadGuard<'_> {
66    pub fn generation(&self) -> usize {
67        self.guard.generation
68    }
69}
70
71impl Deref for StoreReadGuard<'_> {
72    type Target = Store;
73
74    fn deref(&self) -> &Self::Target {
75        &self.guard.store
76    }
77}
78
79pub struct StoreWriteGuard<'a> {
80    guard: RwLockWriteGuard<'a, StoreWithGeneration>,
81}
82
83impl Deref for StoreWriteGuard<'_> {
84    type Target = Store;
85
86    fn deref(&self) -> &Self::Target {
87        &self.guard.store
88    }
89}
90
91impl DerefMut for StoreWriteGuard<'_> {
92    fn deref_mut(&mut self) -> &mut Self::Target {
93        &mut self.guard.store
94    }
95}