grit_lib/refs/store/
memory.rs

1//! In-memory [`super::RefStore`] for embedders and tests.
2
3use std::ops::ControlFlow;
4use std::sync::{Arc, Mutex, MutexGuard};
5
6use crate::error::Result;
7
8type StoreResult<T> = std::result::Result<T, RefStoreError>;
9use crate::reflog::ReflogEntry;
10
11use super::backend::{apply_transaction, check_transaction, resolve_symbolic, RefBatchState};
12use super::error::RefStoreError;
13use super::transaction::RefTransaction;
14use super::{PreparedRefTransaction, RawRef, RefEntry, RefStorageFormat, RefStore};
15
16#[derive(Debug, Clone)]
17struct MemoryRefStoreInner {
18    /// Shared so iterators can walk a snapshot after releasing the lock; writers go through
19    /// [`Arc::make_mut`], which copies only while such a walk still holds the old state.
20    state: Arc<RefBatchState>,
21    prepared: Option<PreparedBatch>,
22}
23
24impl MemoryRefStoreInner {
25    fn state_mut(&mut self) -> &mut RefBatchState {
26        Arc::make_mut(&mut self.state)
27    }
28}
29
30#[derive(Debug, Clone)]
31struct PreparedBatch {
32    lock_name: String,
33    txn: RefTransaction,
34}
35
36/// Fully in-memory ref database with transactional updates.
37///
38/// Built only on the public [`super::backend`] helpers, so it doubles as a reference for
39/// custom backends.
40#[derive(Debug)]
41pub struct MemoryRefStore {
42    inner: Mutex<MemoryRefStoreInner>,
43}
44
45impl Default for MemoryRefStore {
46    fn default() -> Self {
47        Self::new()
48    }
49}
50
51impl MemoryRefStore {
52    /// Create an empty store.
53    #[must_use]
54    pub fn new() -> Self {
55        Self {
56            inner: Mutex::new(MemoryRefStoreInner {
57                state: Arc::default(),
58                prepared: None,
59            }),
60        }
61    }
62
63    /// Seed a ref without going through a transaction (tests and setup).
64    ///
65    /// # Errors
66    ///
67    /// Returns [`RefStoreError::LockHeld`] when a transaction is prepared.
68    pub fn set_ref(&self, name: impl Into<String>, value: RawRef) -> StoreResult<()> {
69        let mut inner = self.lock();
70        if let Some(active) = &inner.prepared {
71            return Err(RefStoreError::LockHeld {
72                name: active.lock_name.clone(),
73            });
74        }
75        inner.state_mut().refs.insert(name.into(), value);
76        Ok(())
77    }
78
79    fn lock(&self) -> MutexGuard<'_, MemoryRefStoreInner> {
80        self.inner.lock().unwrap_or_else(|e| e.into_inner())
81    }
82
83    /// The current state, detached from the lock so iteration callbacks can re-enter the store.
84    fn snapshot(&self) -> Arc<RefBatchState> {
85        Arc::clone(&self.lock().state)
86    }
87}
88
89impl RefStore for MemoryRefStore {
90    fn format(&self) -> RefStorageFormat {
91        RefStorageFormat::Memory
92    }
93
94    fn read_raw(&self, name: &str) -> Result<Option<RawRef>> {
95        Ok(self.lock().state.refs.get(name).cloned())
96    }
97
98    fn for_each_ref(
99        &self,
100        prefix: &str,
101        f: &mut dyn FnMut(&RefEntry) -> ControlFlow<()>,
102    ) -> Result<()> {
103        let state = self.snapshot();
104        for (name, value) in state.refs.range(prefix.to_owned()..) {
105            if !name.starts_with(prefix) {
106                break;
107            }
108            let peeled = match &value {
109                RawRef::Direct(_) => None,
110                RawRef::Symbolic(_) => resolve_symbolic(&state.refs, name).ok(),
111            };
112            let entry = RefEntry {
113                name: name.clone(),
114                value: value.clone(),
115                peeled,
116            };
117            if f(&entry).is_break() {
118                break;
119            }
120        }
121        Ok(())
122    }
123
124    fn prepare(
125        &self,
126        txn: RefTransaction,
127    ) -> StoreResult<Box<dyn PreparedRefTransaction + Send + '_>> {
128        txn.validate()?;
129        let mut inner = self.lock();
130        if let Some(active) = &inner.prepared {
131            return Err(RefStoreError::LockHeld {
132                name: active.lock_name.clone(),
133            });
134        }
135
136        check_transaction(&inner.state, &txn)?;
137
138        let lock_name = txn
139            .updates()
140            .first()
141            .map(|u| u.name.clone())
142            .or_else(|| txn.reflog_renames().first().map(|r| r.from.clone()))
143            .or_else(|| txn.case_only_renames().first().map(|r| r.from.clone()))
144            .unwrap_or_else(|| "refs".to_owned());
145        inner.prepared = Some(PreparedBatch { lock_name, txn });
146
147        Ok(Box::new(MemoryPrepared {
148            store: self,
149            committed: false,
150        }))
151    }
152
153    fn reflog_exists(&self, name: &str) -> Result<bool> {
154        Ok(self.lock().state.reflogs.contains_key(name))
155    }
156
157    fn for_each_reflog_entry(
158        &self,
159        name: &str,
160        reverse: bool,
161        f: &mut dyn FnMut(&ReflogEntry) -> ControlFlow<()>,
162    ) -> Result<()> {
163        let state = self.snapshot();
164        let Some(entries) = state.reflogs.get(name) else {
165            return Ok(());
166        };
167        if reverse {
168            for entry in entries.iter().rev() {
169                if f(entry).is_break() {
170                    break;
171                }
172            }
173        } else {
174            for entry in entries {
175                if f(entry).is_break() {
176                    break;
177                }
178            }
179        }
180        Ok(())
181    }
182
183    fn create_reflog(&self, name: &str) -> Result<()> {
184        let mut inner = self.lock();
185        if !inner.state.reflogs.contains_key(name) {
186            inner
187                .state_mut()
188                .reflogs
189                .insert(name.to_owned(), Vec::new());
190        }
191        Ok(())
192    }
193
194    fn delete_reflog(&self, name: &str) -> Result<()> {
195        let mut inner = self.lock();
196        if inner.state.reflogs.contains_key(name) {
197            inner.state_mut().reflogs.remove(name);
198        }
199        Ok(())
200    }
201
202    fn replace_reflog(&self, name: &str, entries: Vec<ReflogEntry>) -> Result<()> {
203        self.lock()
204            .state_mut()
205            .reflogs
206            .insert(name.to_owned(), entries);
207        Ok(())
208    }
209
210    fn for_each_reflog_ref(&self, f: &mut dyn FnMut(&str) -> ControlFlow<()>) -> Result<()> {
211        let state = self.snapshot();
212        for name in state.reflogs.keys() {
213            if f(name).is_break() {
214                break;
215            }
216        }
217        Ok(())
218    }
219}
220
221struct MemoryPrepared<'a> {
222    store: &'a MemoryRefStore,
223    committed: bool,
224}
225
226impl PreparedRefTransaction for MemoryPrepared<'_> {
227    fn commit(mut self: Box<Self>) -> StoreResult<()> {
228        let mut inner = self.store.lock();
229        let Some(batch) = inner.prepared.take() else {
230            return Ok(());
231        };
232        let mut trial = inner.state.snapshot();
233        apply_transaction(&mut trial, &batch.txn)?;
234        inner.state = Arc::new(trial);
235        self.committed = true;
236        Ok(())
237    }
238
239    fn abort(mut self: Box<Self>) -> StoreResult<()> {
240        let mut inner = self.store.lock();
241        inner.prepared = None;
242        self.committed = true;
243        Ok(())
244    }
245}
246
247impl Drop for MemoryPrepared<'_> {
248    fn drop(&mut self) {
249        if self.committed {
250            return;
251        }
252        let mut inner = self.store.inner.lock().unwrap_or_else(|e| e.into_inner());
253        inner.prepared = None;
254    }
255}
256
257#[cfg(test)]
258mod tests {
259    use super::*;
260    use crate::objects::ObjectId;
261
262    fn oid(byte: u8) -> ObjectId {
263        let mut bytes = [0u8; 20];
264        bytes[19] = byte;
265        ObjectId::from_bytes(&bytes).expect("valid oid")
266    }
267
268    #[test]
269    fn for_each_ref_peels_symbolic_without_deadlock() {
270        let store = MemoryRefStore::new();
271        store
272            .set_ref("refs/heads/main", RawRef::Direct(oid(1)))
273            .expect("main");
274        store
275            .set_ref("HEAD", RawRef::Symbolic("refs/heads/main".to_owned()))
276            .expect("head");
277
278        let mut peeled = None;
279        store
280            .for_each_ref("HEAD", &mut |entry| {
281                peeled = entry.peeled;
282                ControlFlow::Break(())
283            })
284            .expect("iter");
285        assert_eq!(peeled, Some(oid(1)));
286    }
287
288    /// Run `body` on its own thread and fail if it does not finish, so a callback that
289    /// deadlocks on the store's lock fails the test instead of hanging the suite.
290    fn finishes_promptly(body: impl FnOnce() + Send + 'static) {
291        let (done_tx, done_rx) = std::sync::mpsc::channel();
292        let handle = std::thread::spawn(move || {
293            body();
294            let _ = done_tx.send(());
295        });
296        match done_rx.recv_timeout(std::time::Duration::from_secs(10)) {
297            Ok(()) => handle.join().expect("test thread"),
298            Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => {
299                if let Err(panic) = handle.join() {
300                    std::panic::resume_unwind(panic);
301                }
302            }
303            Err(std::sync::mpsc::RecvTimeoutError::Timeout) => {
304                panic!("callback re-entering the store deadlocked")
305            }
306        }
307    }
308
309    fn sample_reflog(byte: u8) -> Vec<ReflogEntry> {
310        vec![ReflogEntry {
311            old_oid: oid(0),
312            new_oid: oid(byte),
313            identity: "A <a@example.com> 1 +0000".to_owned(),
314            message: format!("m{byte}"),
315        }]
316    }
317
318    #[test]
319    fn for_each_ref_callback_can_read_the_store() {
320        finishes_promptly(|| {
321            let store = MemoryRefStore::new();
322            store
323                .set_ref("refs/heads/main", RawRef::Direct(oid(1)))
324                .expect("main");
325            store
326                .set_ref(
327                    "refs/heads/alias",
328                    RawRef::Symbolic("refs/heads/main".into()),
329                )
330                .expect("alias");
331
332            let mut seen = Vec::new();
333            store
334                .for_each_ref("refs/", &mut |entry| {
335                    let raw = store.read_raw(&entry.name).expect("read_raw");
336                    let resolved = store.resolve(&entry.name).expect("resolve");
337                    let mut nested = 0usize;
338                    store
339                        .for_each_ref("refs/heads/", &mut |_| {
340                            nested += 1;
341                            ControlFlow::Continue(())
342                        })
343                        .expect("nested for_each_ref");
344                    seen.push((entry.name.clone(), raw, resolved, nested));
345                    ControlFlow::Continue(())
346                })
347                .expect("for_each_ref");
348            assert_eq!(
349                seen,
350                vec![
351                    (
352                        "refs/heads/alias".to_owned(),
353                        Some(RawRef::Symbolic("refs/heads/main".into())),
354                        oid(1),
355                        2,
356                    ),
357                    (
358                        "refs/heads/main".to_owned(),
359                        Some(RawRef::Direct(oid(1))),
360                        oid(1),
361                        2,
362                    ),
363                ]
364            );
365        });
366    }
367
368    #[test]
369    fn for_each_reflog_ref_callback_can_read_each_reflog() {
370        finishes_promptly(|| {
371            let store = MemoryRefStore::new();
372            store
373                .replace_reflog("HEAD", sample_reflog(1))
374                .expect("HEAD log");
375            store
376                .replace_reflog("refs/heads/main", sample_reflog(2))
377                .expect("main log");
378
379            let mut dump = Vec::new();
380            store
381                .for_each_reflog_ref(&mut |name| {
382                    assert!(store.reflog_exists(name).expect("reflog_exists"));
383                    store
384                        .for_each_reflog_entry(name, false, &mut |entry| {
385                            assert!(store.read_raw(name).expect("read_raw").is_none());
386                            dump.push((name.to_owned(), entry.new_oid));
387                            ControlFlow::Continue(())
388                        })
389                        .expect("for_each_reflog_entry");
390                    ControlFlow::Continue(())
391                })
392                .expect("for_each_reflog_ref");
393            assert_eq!(
394                dump,
395                vec![
396                    ("HEAD".to_owned(), oid(1)),
397                    ("refs/heads/main".to_owned(), oid(2)),
398                ]
399            );
400        });
401    }
402
403    #[test]
404    fn writes_from_callbacks_succeed_without_changing_the_walk() {
405        finishes_promptly(|| {
406            let store = MemoryRefStore::new();
407            store
408                .set_ref("refs/heads/a", RawRef::Direct(oid(1)))
409                .expect("a");
410            store
411                .set_ref("refs/heads/b", RawRef::Direct(oid(2)))
412                .expect("b");
413            store
414                .replace_reflog("refs/heads/a", sample_reflog(1))
415                .expect("log");
416
417            let mut names = Vec::new();
418            store
419                .for_each_ref("refs/heads/", &mut |entry| {
420                    store
421                        .set_ref("refs/heads/added", RawRef::Direct(oid(3)))
422                        .expect("set_ref from callback");
423                    names.push(entry.name.clone());
424                    ControlFlow::Continue(())
425                })
426                .expect("for_each_ref");
427            assert_eq!(names, vec!["refs/heads/a", "refs/heads/b"]);
428            assert_eq!(
429                store.read_raw("refs/heads/added").expect("read"),
430                Some(RawRef::Direct(oid(3)))
431            );
432
433            let mut entries = 0usize;
434            store
435                .for_each_reflog_entry("refs/heads/a", true, &mut |_| {
436                    store
437                        .replace_reflog("refs/heads/a", Vec::new())
438                        .expect("replace_reflog from callback");
439                    entries += 1;
440                    ControlFlow::Continue(())
441                })
442                .expect("for_each_reflog_entry");
443            assert_eq!(entries, 1);
444
445            let mut logged = Vec::new();
446            store
447                .for_each_reflog_ref(&mut |name| {
448                    store
449                        .delete_reflog(name)
450                        .expect("delete_reflog from callback");
451                    store.create_reflog("refs/heads/zz").expect("create_reflog");
452                    logged.push(name.to_owned());
453                    ControlFlow::Continue(())
454                })
455                .expect("for_each_reflog_ref");
456            assert_eq!(logged, vec!["refs/heads/a"]);
457            assert!(!store.reflog_exists("refs/heads/a").expect("exists"));
458            assert!(store.reflog_exists("refs/heads/zz").expect("exists"));
459        });
460    }
461}
grit_lib/refs/store/memory.rs

grit_lib/refs/store/
memory.rs
