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

grit_lib/refs/store/
memory.rs
