turbo_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 StoreContainer {
22    pub fn new() -> Self {
23        Self {
24            store: RwLock::new(StoreWithGeneration {
25                store: Store::new(),
26                generation: 0,
27            }),
28            want_to_read: AtomicBool::new(false),
29        }
30    }
31
32    pub fn read(&self) -> StoreReadGuard<'_> {
33        if let Ok(guard) = self.store.try_read() {
34            return StoreReadGuard { guard };
35        }
36        self.want_to_read.store(true, Ordering::Relaxed);
37        let guard = StoreReadGuard {
38            guard: self.store.read().unwrap(),
39        };
40        self.want_to_read.store(false, Ordering::Relaxed);
41        guard
42    }
43
44    pub fn write(&self) -> StoreWriteGuard<'_> {
45        let mut guard = self.store.write().unwrap();
46        guard.generation += 1;
47        StoreWriteGuard { guard }
48    }
49
50    pub fn want_to_read(&self) -> bool {
51        self.want_to_read.load(Ordering::Relaxed)
52    }
53}
54
55pub struct StoreReadGuard<'a> {
56    guard: RwLockReadGuard<'a, StoreWithGeneration>,
57}
58
59impl StoreReadGuard<'_> {
60    pub fn generation(&self) -> usize {
61        self.guard.generation
62    }
63}
64
65impl Deref for StoreReadGuard<'_> {
66    type Target = Store;
67
68    fn deref(&self) -> &Self::Target {
69        &self.guard.store
70    }
71}
72
73pub struct StoreWriteGuard<'a> {
74    guard: RwLockWriteGuard<'a, StoreWithGeneration>,
75}
76
77impl Deref for StoreWriteGuard<'_> {
78    type Target = Store;
79
80    fn deref(&self) -> &Self::Target {
81        &self.guard.store
82    }
83}
84
85impl DerefMut for StoreWriteGuard<'_> {
86    fn deref_mut(&mut self) -> &mut Self::Target {
87        &mut self.guard.store
88    }
89}
turbo_trace_server/store_container.rs

turbo_trace_server/
store_container.rs
