diffs_rs/storage/
bytes.rs

1use std::{
2    collections::{HashMap, VecDeque},
3    fs::File,
4    io::{self, Read, Write},
5    os::unix::fs::{FileExt, PermissionsExt},
6    sync::{
7        Arc, Mutex,
8        atomic::{AtomicU64, Ordering},
9    },
10};
11
12use anyhow::{Context, Result, ensure};
13
14const PAGE: usize = 64 * 1024;
15const SOURCE_CACHE: usize = 32 * 1024 * 1024;
16static NEXT_STORE: AtomicU64 = AtomicU64::new(1);
17
18/// Creates an unnamed temporary file readable and writable only by its owner.
19///
20/// On Linux, `tempfile()` uses `O_TMPFILE`, which applies the process umask to mode 0666
21/// (usually 0644) because the file has no name. Setting the mode keeps private data
22/// owner-only on every platform, including through `/proc/<pid>/fd`.
23pub(crate) fn private_tempfile() -> std::io::Result<File> {
24    let file = tempfile::tempfile()?;
25    file.set_permissions(std::fs::Permissions::from_mode(0o600))?;
26    Ok(file)
27}
28
29#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
30pub(super) struct StoreId(u64);
31
32/// A checked range of published bytes, tied to its backing store.
33#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
34pub struct SourceRef {
35    store: StoreId,
36    offset: u64,
37    len: usize,
38}
39
40#[derive(Debug)]
41struct Page {
42    bytes: Box<[u8]>,
43}
44
45#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
46struct PageKey {
47    store: StoreId,
48    number: u64,
49}
50
51#[derive(Default, Debug)]
52struct Pages {
53    entries: HashMap<PageKey, Arc<Page>>,
54    order: VecDeque<PageKey>,
55    bytes: usize,
56}
57
58#[derive(Debug)]
59pub struct ByteStore {
60    id: StoreId,
61    file: File,
62    len: AtomicU64,
63    pages: Arc<Mutex<Pages>>,
64}
65
66/// The sole writer for a store. Appends never modify published bytes.
67pub struct StoreWriter {
68    store: Arc<ByteStore>,
69}
70
71impl StoreWriter {
72    pub fn new() -> Result<Self> {
73        Self::with_pages(Arc::new(Mutex::new(Pages::default())))
74    }
75
76    /// Create a distinct private store within the existing store's cache budget.
77    pub fn sharing(store: &ByteStore) -> Result<Self> {
78        Self::with_pages(store.pages.clone())
79    }
80
81    fn with_pages(pages: Arc<Mutex<Pages>>) -> Result<Self> {
82        let file = private_tempfile().context("Cannot create the private source store")?;
83        Self::with_file(file, pages)
84    }
85
86    pub(crate) fn named_sharing(store: &ByteStore) -> Result<(Self, tempfile::TempPath)> {
87        let (file, path) = tempfile::NamedTempFile::new()?.into_parts();
88        Ok((Self::with_file(file, store.pages.clone())?, path))
89    }
90
91    fn with_file(file: File, pages: Arc<Mutex<Pages>>) -> Result<Self> {
92        let id = NEXT_STORE
93            .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |id| id.checked_add(1))
94            .map(StoreId)
95            .map_err(|_| anyhow::anyhow!("Source store identity exhausted"))?;
96        Ok(Self {
97            store: Arc::new(ByteStore {
98                id,
99                file,
100                len: AtomicU64::new(0),
101                pages,
102            }),
103        })
104    }
105
106    pub fn store(&self) -> Arc<ByteStore> {
107        self.store.clone()
108    }
109
110    pub fn append(&mut self, bytes: &[u8]) -> Result<u64> {
111        let base = self.store.len();
112        let len = u64::try_from(bytes.len()).context("Source length overflow")?;
113        let end = base
114            .checked_add(len)
115            .context("Source store offset overflow")?;
116        self.store
117            .file
118            .write_all_at(bytes, base)
119            .context("Cannot write the source store")?;
120        self.store.len.store(end, Ordering::Release);
121        Ok(base)
122    }
123
124    pub fn finish(self) -> Arc<ByteStore> {
125        self.store
126    }
127
128    /// Reclaim a source job's tail before its snapshot descriptors are published.
129    pub(crate) fn rollback(&mut self, len: u64) -> Result<()> {
130        ensure!(len <= self.store.len(), "Invalid source rollback boundary");
131        let mut pages = self
132            .store
133            .pages
134            .lock()
135            .map_err(|_| anyhow::anyhow!("Source cache lock failed"))?;
136        pages
137            .entries
138            .retain(|key, _| key.store != self.store.id || key.number < len / PAGE as u64);
139        pages
140            .order
141            .retain(|key| key.store != self.store.id || key.number < len / PAGE as u64);
142        pages.bytes = pages.entries.values().map(|page| page.bytes.len()).sum();
143        self.store.file.set_len(len)?;
144        self.store.len.store(len, Ordering::Release);
145        Ok(())
146    }
147}
148
149/// A bounded sequential reader that does not populate the visible-page cache.
150pub struct StoreReader {
151    store: Arc<ByteStore>,
152    position: u64,
153    end: u64,
154}
155
156impl StoreReader {
157    pub(crate) fn remaining(&self) -> u64 {
158        self.end - self.position
159    }
160}
161
162impl Read for StoreReader {
163    fn read(&mut self, output: &mut [u8]) -> io::Result<usize> {
164        let count = (self.end - self.position).min(output.len() as u64) as usize;
165        if count == 0 {
166            return Ok(0);
167        }
168        let count = self
169            .store
170            .file
171            .read_at(&mut output[..count], self.position)?;
172        if count == 0 {
173            return Err(io::ErrorKind::UnexpectedEof.into());
174        }
175        self.position += count as u64;
176        Ok(count)
177    }
178}
179
180impl Write for StoreWriter {
181    fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
182        self.append(bytes).map_err(io::Error::other)?;
183        Ok(bytes.len())
184    }
185
186    fn flush(&mut self) -> io::Result<()> {
187        Ok(())
188    }
189}
190
191impl ByteStore {
192    #[cfg(feature = "owner-census")]
193    pub(crate) fn cache_census(&self) -> Result<serde_json::Value> {
194        let pages = self
195            .pages
196            .lock()
197            .map_err(|_| anyhow::anyhow!("Source cache lock failed"))?;
198        Ok(serde_json::json!({
199            "entries": pages.entries.len(), "map_capacity": pages.entries.capacity(),
200            "payload_counter_bytes": pages.bytes,
201            "recomputed_payload_bytes": pages.entries.values().map(|page| page.bytes.len()).sum::<usize>(),
202            "payload_limit_bytes": SOURCE_CACHE,
203            "page_value_bytes": pages.entries.len() * size_of::<Page>(),
204            "occupied_map_element_bytes": pages.entries.len() * size_of::<(PageKey, Arc<Page>)>(),
205            "fifo_length": pages.order.len(), "fifo_capacity": pages.order.capacity(),
206            "fifo_element_size_bytes": size_of::<PageKey>(),
207        }))
208    }
209
210    pub(crate) fn identity(&self) -> u64 {
211        self.id.0
212    }
213    pub fn reader(self: &Arc<Self>, offset: u64, len: u64) -> Result<StoreReader> {
214        let end = offset.checked_add(len).context("Source range overflow")?;
215        ensure!(
216            end <= self.len(),
217            "Source range is outside the published bytes"
218        );
219        Ok(StoreReader {
220            store: self.clone(),
221            position: offset,
222            end,
223        })
224    }
225
226    pub fn source_reader(self: &Arc<Self>, source: SourceRef) -> Result<StoreReader> {
227        ensure!(
228            source.store == self.id,
229            "Source reference belongs to another store"
230        );
231        self.reader(source.offset, source.len as u64)
232    }
233    pub(super) fn id(&self) -> StoreId {
234        self.id
235    }
236
237    pub fn len(&self) -> u64 {
238        self.len.load(Ordering::Acquire)
239    }
240
241    pub fn is_empty(&self) -> bool {
242        self.len() == 0
243    }
244
245    pub fn source_ref(&self, offset: u64, len: usize) -> Result<SourceRef> {
246        let length = u64::try_from(len).context("Source length overflow")?;
247        let end = offset
248            .checked_add(length)
249            .context("Source range overflow")?;
250        ensure!(
251            end <= self.len(),
252            "Source range is outside the published bytes"
253        );
254        Ok(SourceRef {
255            store: self.id,
256            offset,
257            len,
258        })
259    }
260
261    pub fn read_source(&self, source: SourceRef) -> Result<Vec<u8>> {
262        ensure!(
263            source.store == self.id,
264            "Source reference belongs to another store"
265        );
266        let end = source.offset + source.len as u64;
267        let mut bytes = Vec::with_capacity(source.len);
268        let mut position = source.offset;
269        while position < end {
270            let within = (position % PAGE as u64) as usize;
271            let take = (end - position).min((PAGE - within) as u64) as usize;
272            let page = self.page(position / PAGE as u64, within + take)?;
273            bytes.extend_from_slice(&page.bytes[within..within + take]);
274            position += take as u64;
275        }
276        Ok(bytes)
277    }
278
279    pub(crate) fn read_into(&self, offset: u64, output: &mut [u8]) -> Result<()> {
280        self.source_ref(offset, output.len())?;
281        let mut position = offset;
282        let mut remaining = output;
283        while !remaining.is_empty() {
284            let within = (position % PAGE as u64) as usize;
285            let take = remaining.len().min(PAGE - within);
286            let page = self.page(position / PAGE as u64, within + take)?;
287            remaining[..take].copy_from_slice(&page.bytes[within..within + take]);
288            remaining = &mut remaining[take..];
289            position += take as u64;
290        }
291        Ok(())
292    }
293
294    fn page(&self, number: u64, required: usize) -> Result<Arc<Page>> {
295        let key = PageKey {
296            store: self.id,
297            number,
298        };
299        let mut pages = self
300            .pages
301            .lock()
302            .map_err(|_| anyhow::anyhow!("Source cache lock failed"))?;
303        if let Some(page) = pages
304            .entries
305            .get(&key)
306            .filter(|page| page.bytes.len() >= required)
307        {
308            return Ok(page.clone());
309        }
310        let offset = number
311            .checked_mul(PAGE as u64)
312            .context("Source page offset overflow")?;
313        let len = self
314            .len()
315            .checked_sub(offset)
316            .context("Source page is outside the store")?
317            .min(PAGE as u64) as usize;
318        ensure!(
319            required <= len,
320            "Source range is outside the published bytes"
321        );
322        let mut bytes = vec![0; len].into_boxed_slice();
323        self.file
324            .read_exact_at(&mut bytes, offset)
325            .context("Cannot read the source store")?;
326        let page = Arc::new(Page { bytes });
327        if let Some(previous) = pages.entries.remove(&key) {
328            pages.bytes -= previous.bytes.len();
329            pages.order.retain(|entry| *entry != key);
330        }
331        while pages.bytes + len > SOURCE_CACHE {
332            if let Some(old) = pages.order.pop_front()
333                && let Some(page) = pages.entries.remove(&old)
334            {
335                pages.bytes -= page.bytes.len();
336            }
337        }
338        pages.bytes += len;
339        pages.order.push_back(key);
340        pages.entries.insert(key, page.clone());
341        Ok(page)
342    }
343
344    pub fn read(&self, offset: u64, len: usize) -> Result<Vec<u8>> {
345        self.read_source(self.source_ref(offset, len)?)
346    }
347
348    /// Cached payload across this store and every store created with `StoreWriter::sharing`.
349    pub fn cached_bytes(&self) -> usize {
350        self.pages
351            .lock()
352            .map(|pages| pages.bytes)
353            .unwrap_or_default()
354    }
355}
356
357impl Drop for ByteStore {
358    fn drop(&mut self) {
359        if let Ok(mut pages) = self.pages.lock() {
360            let mut removed = 0;
361            pages.entries.retain(|key, page| {
362                if key.store == self.id {
363                    removed += page.bytes.len();
364                    false
365                } else {
366                    true
367                }
368            });
369            pages.bytes -= removed;
370            pages.order.retain(|key| key.store != self.id);
371        }
372    }
373}
374
375#[cfg(test)]
376mod tests {
377    use super::*;
378    use std::{
379        io::{Read, Seek, SeekFrom},
380        sync::mpsc,
381        thread,
382    };
383
384    #[test]
385    fn source_pages_preserve_bytes_and_readers_through_refresh_and_eviction() {
386        let mut writer = StoreWriter::new().unwrap();
387        let spool = writer.store();
388        let prefix = b"partial page\t\xff\r\n";
389        assert_eq!(writer.append(prefix).unwrap(), 0);
390        let short = spool.page(0, prefix.len()).unwrap();
391        let short_weak = Arc::downgrade(&short);
392        assert_eq!(short.bytes.as_ref(), prefix);
393        assert!(Arc::ptr_eq(&short, &spool.page(0, 1).unwrap()));
394        assert!(spool.read(prefix.len() as u64, 1).is_err());
395
396        let expected: Vec<u8> = (0..PAGE * 2 + 19).map(|n| (n % 251) as u8).collect();
397        writer.append(&expected[prefix.len()..]).unwrap();
398        let full = spool.page(0, PAGE).unwrap();
399        let full_weak = Arc::downgrade(&full);
400        assert!(!Arc::ptr_eq(&short, &full));
401        assert_eq!(short.bytes.as_ref(), prefix);
402        assert_eq!(&full.bytes[..prefix.len()], prefix);
403        assert_eq!(&full.bytes[prefix.len()..], &expected[prefix.len()..PAGE]);
404        assert_eq!(
405            spool.read((PAGE - 7) as u64, 23).unwrap(),
406            expected[PAGE - 7..PAGE + 16]
407        );
408        assert_eq!(
409            spool.read((PAGE * 2) as u64, 19).unwrap(),
410            expected[PAGE * 2..]
411        );
412        assert!(spool.read(spool.len(), 0).unwrap().is_empty());
413        assert!(spool.read(u64::MAX, 2).is_err());
414        drop(short);
415        assert!(short_weak.upgrade().is_none());
416
417        let page_bytes: Vec<u8> = (0..PAGE).map(|n| (n % 239) as u8).collect();
418        let padding = PAGE - (spool.len() as usize % PAGE);
419        writer.append(&page_bytes[..padding]).unwrap();
420        for _ in 0..=SOURCE_CACHE / PAGE {
421            let offset = writer.append(&page_bytes).unwrap();
422            assert_eq!(spool.read(offset, PAGE).unwrap(), page_bytes);
423        }
424        {
425            let pages = spool.pages.lock().unwrap();
426            assert_eq!(pages.bytes, SOURCE_CACHE);
427            assert_eq!(pages.entries.len(), SOURCE_CACHE / PAGE);
428            assert!(!pages.entries.contains_key(&PageKey {
429                store: spool.id,
430                number: 0
431            }));
432        }
433        assert_eq!(&full.bytes[..prefix.len()], prefix);
434        assert_eq!(spool.read(0, prefix.len()).unwrap(), prefix);
435        assert!(!Arc::ptr_eq(&full, &spool.page(0, PAGE).unwrap()));
436        drop(full);
437        assert!(full_weak.upgrade().is_none());
438        let last = spool.page(0, PAGE).unwrap();
439        let last_weak = Arc::downgrade(&last);
440        drop(writer);
441        drop(spool);
442        assert_eq!(&last.bytes[..prefix.len()], prefix);
443        drop(last);
444        assert!(last_weak.upgrade().is_none());
445    }
446
447    #[test]
448    fn patch_and_full_source_stores_share_budget_without_aliasing_or_retained_dead_pages() {
449        let mut files = Vec::new();
450        let patch = crate::storage::read(
451            include_bytes!("../../fixtures/review.patch").as_slice(),
452            |file| {
453                files.push(file);
454                Ok(())
455            },
456        )
457        .unwrap();
458        let first = &files[0];
459        let reference = first.source_ref(first.source_base).unwrap();
460        let original = patch.read_source(reference).unwrap();
461        let patch_page = patch.page(0, 1).unwrap();
462        let mut writers = [
463            StoreWriter::sharing(&patch).unwrap(),
464            StoreWriter::sharing(&patch).unwrap(),
465        ];
466        let mut sources = Vec::new();
467        let mut expected = Vec::new();
468        for (side, writer) in writers.iter_mut().enumerate() {
469            let bytes: Vec<u8> = (0..PAGE).map(|n| ((n + side) % 251) as u8).collect();
470            let mut input = tempfile::tempfile().unwrap();
471            for _ in 0..=SOURCE_CACHE / PAGE / 2 {
472                input.write_all(&bytes).unwrap();
473            }
474            input.seek(SeekFrom::Start(0)).unwrap();
475            assert_eq!(
476                io::copy(&mut input, writer).unwrap(),
477                input.metadata().unwrap().len()
478            );
479            sources.push(writer.store());
480            expected.push(bytes);
481        }
482        for number in 0..=SOURCE_CACHE / PAGE / 2 {
483            for (store, bytes) in sources.iter().zip(&expected) {
484                assert_eq!(store.read((number * PAGE) as u64, PAGE).unwrap(), *bytes);
485                assert_eq!(store.cached_bytes(), patch.cached_bytes());
486                assert!(store.cached_bytes() <= SOURCE_CACHE);
487            }
488        }
489        assert_eq!(patch.cached_bytes(), SOURCE_CACHE);
490        #[cfg(feature = "owner-census")]
491        {
492            let pointers = || {
493                let pages = patch.pages.lock().unwrap();
494                let mut values: Vec<_> = pages
495                    .entries
496                    .iter()
497                    .map(|(key, page)| {
498                        (
499                            key.store.0,
500                            key.number,
501                            Arc::as_ptr(page),
502                            Arc::strong_count(page),
503                        )
504                    })
505                    .collect();
506                values.sort_unstable();
507                (values, pages.order.clone(), Arc::strong_count(&patch.pages))
508            };
509            let before = pointers();
510            let census = patch.cache_census().unwrap();
511            assert_eq!(census["payload_counter_bytes"], SOURCE_CACHE);
512            assert_eq!(census["recomputed_payload_bytes"], SOURCE_CACHE);
513            for store in &sources {
514                assert_eq!(store.cache_census().unwrap(), census);
515            }
516            assert_eq!(patch.cache_census().unwrap(), census);
517            assert_eq!(pointers(), before);
518        }
519        let old = sources[0].source_ref(0, PAGE).unwrap();
520        let new = sources[1].source_ref(0, PAGE).unwrap();
521        assert_ne!(old, new);
522        assert!(sources[0].read_source(new).is_err());
523        assert!(patch.read_source(old).is_err());
524        assert_eq!(patch.read_source(reference).unwrap(), original);
525        assert_eq!(first.source_ref(first.source_base).unwrap(), reference);
526        assert_eq!(
527            patch_page.bytes[..original.len()],
528            patch.read(0, original.len()).unwrap()
529        );
530
531        drop(writers);
532        let old_store = sources.remove(0);
533        let old_id = old_store.id;
534        let retained = old_store.page(0, PAGE).unwrap();
535        let weak = Arc::downgrade(&old_store);
536        drop(old_store);
537        assert!(weak.upgrade().is_none());
538        {
539            let pages = patch.pages.lock().unwrap();
540            assert!(pages.entries.keys().all(|key| key.store != old_id));
541            assert!(pages.order.iter().all(|key| key.store != old_id));
542            assert_eq!(
543                pages.bytes,
544                pages
545                    .entries
546                    .values()
547                    .map(|page| page.bytes.len())
548                    .sum::<usize>()
549            );
550        }
551        assert_eq!(retained.bytes.as_ref(), expected[0]);
552        assert_eq!(sources[0].read_source(new).unwrap(), expected[1]);
553        let replacement = StoreWriter::sharing(&patch).unwrap();
554        assert_ne!(replacement.store.id, old_id);
555        assert!(replacement.store.read_source(old).is_err());
556        let cache = Arc::downgrade(&patch.pages);
557        drop(replacement);
558        drop(sources);
559        drop(files);
560        drop(patch);
561        assert!(cache.upgrade().is_none());
562        assert_eq!(retained.bytes.as_ref(), expected[0]);
563    }
564
565    #[test]
566    fn published_references_survive_concurrent_append_and_writer_completion() {
567        let mut writer = StoreWriter::new().unwrap();
568        let store = writer.store();
569        assert!(store.is_empty());
570        let empty = store.source_ref(0, 0).unwrap();
571        assert!(store.source_ref(1, 0).is_err());
572        assert!(store.source_ref(u64::MAX, 2).is_err());
573        let payload = b"\tconst VALUE = \"caf\xc3\xa9 \xe6\x97\xa5\";\xff\r\n";
574        let input: Vec<u8> = payload
575            .iter()
576            .copied()
577            .cycle()
578            .take(PAGE * 3 + 11)
579            .collect();
580        let mut input_file = tempfile::tempfile().unwrap();
581        input_file.write_all(&input).unwrap();
582        input_file.seek(SeekFrom::Start(0)).unwrap();
583        let (send, receive) = mpsc::sync_channel(1);
584        thread::scope(|scope| {
585            let writer_task = scope.spawn(move || {
586                let mut bytes = [0; 997];
587                loop {
588                    let n = input_file.read(&mut bytes).unwrap();
589                    if n == 0 {
590                        break;
591                    }
592                    let offset = writer.append(&bytes[..n]).unwrap();
593                    send.send(writer.store.source_ref(offset, n).unwrap())
594                        .unwrap();
595                }
596                writer.finish()
597            });
598            let mut references = Vec::new();
599            for reference in receive {
600                references.push(reference);
601                let offset = reference.offset as usize;
602                assert_eq!(
603                    store.read_source(reference).unwrap(),
604                    input[offset..offset + reference.len]
605                );
606                if let Some(&first) = references.first() {
607                    assert_eq!(store.read_source(first).unwrap(), input[..first.len]);
608                }
609            }
610            let completed = writer_task.join().unwrap();
611            assert!(Arc::ptr_eq(&store, &completed));
612            assert_eq!(store.read(0, input.len()).unwrap(), input);
613            assert!(store.read_source(empty).unwrap().is_empty());
614            for reference in references {
615                let offset = reference.offset as usize;
616                assert_eq!(
617                    completed.read_source(reference).unwrap(),
618                    input[offset..offset + reference.len]
619                );
620            }
621        });
622    }
623}
diffs_rs/storage/bytes.rs

diffs_rs/storage/
bytes.rs
