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

grit_lib/refs/store/
memory.rs
