terminus_store/storage/
archive.rs

1// File format:
2// <header>
3//  [<filetype present>]*
4//  [<offsets>]*
5//
6
7use std::{
8    collections::HashMap,
9    io::{self, ErrorKind, SeekFrom},
10    ops::Range,
11    path::PathBuf,
12    pin::Pin,
13    sync::{Arc, RwLock},
14    task::Poll,
15};
16
17#[cfg(not(target_os = "windows"))]
18use std::os::unix::fs::MetadataExt;
19#[cfg(target_os = "windows")]
20use std::os::windows::fs::MetadataExt;
21
22use async_trait::async_trait;
23use bytes::{Buf, BufMut, Bytes, BytesMut};
24use lru::LruCache;
25use tokio::{
26    fs::{self, File},
27    io::{AsyncRead, AsyncReadExt, AsyncSeekExt, AsyncWrite, AsyncWriteExt},
28};
29use tokio_util::either::Either;
30
31use tdb_succinct::{
32    logarray_length_from_control_word, smallbitarray::SmallBitArray, LateLogArrayBufBuilder,
33    MonotonicLogArray,
34};
35
36use super::{
37    consts::{LayerFileEnum, FILENAME_ENUM_MAP},
38    locking::{ExclusiveLockedFile, LockedFile},
39    name_to_string, string_to_name, FileLoad, FileStore, PersistentLayerStore, SyncableFile,
40};
41
42#[async_trait]
43pub trait ArchiveBackend: Clone + Send + Sync {
44    type Read: AsyncRead + Unpin + Send;
45    async fn get_layer_bytes(&self, id: [u32; 5]) -> io::Result<Bytes>;
46    async fn get_layer_structure_bytes(
47        &self,
48        id: [u32; 5],
49        file_type: LayerFileEnum,
50    ) -> io::Result<Option<Bytes>>;
51    async fn store_layer_file(&self, id: [u32; 5], bytes: Bytes) -> io::Result<()>;
52    async fn read_layer_structure_bytes_from(
53        &self,
54        id: [u32; 5],
55        file_type: LayerFileEnum,
56        read_from: usize,
57    ) -> io::Result<Self::Read>;
58}
59
60#[async_trait]
61pub trait ArchiveMetadataBackend: Clone + Send + Sync {
62    async fn get_layer_names(&self) -> io::Result<Vec<[u32; 5]>>;
63    async fn layer_exists(&self, id: [u32; 5]) -> io::Result<bool>;
64    async fn layer_size(&self, id: [u32; 5]) -> io::Result<u64>;
65    async fn layer_file_exists(&self, id: [u32; 5], file_type: LayerFileEnum) -> io::Result<bool>;
66    async fn get_layer_structure_size(
67        &self,
68        id: [u32; 5],
69        file_type: LayerFileEnum,
70    ) -> io::Result<usize>;
71    async fn get_rollup(&self, id: [u32; 5]) -> io::Result<Option<[u32; 5]>>;
72    async fn set_rollup(&self, id: [u32; 5], rollup: [u32; 5]) -> io::Result<()>;
73    async fn get_parent(&self, id: [u32; 5]) -> io::Result<Option<[u32; 5]>>;
74}
75
76pub struct BytesAsyncReader(Bytes);
77
78impl AsyncRead for BytesAsyncReader {
79    fn poll_read(
80        self: Pin<&mut Self>,
81        _cx: &mut std::task::Context<'_>,
82        buf: &mut tokio::io::ReadBuf<'_>,
83    ) -> Poll<io::Result<()>> {
84        let bytes = &mut self.get_mut().0;
85        let consumed = if buf.remaining() > bytes.len() {
86            bytes.split_to(bytes.len())
87        } else {
88            bytes.split_to(buf.remaining())
89        };
90
91        buf.put_slice(consumed.as_ref());
92
93        Poll::Ready(Ok(()))
94    }
95}
96
97#[derive(Clone)]
98pub struct DirectoryArchiveBackend {
99    path: PathBuf,
100}
101
102impl DirectoryArchiveBackend {
103    pub fn new(path: PathBuf) -> Self {
104        Self { path }
105    }
106    fn path_for_layer(&self, name: [u32; 5]) -> PathBuf {
107        let mut p = self.path.clone();
108        let name_str = name_to_string(name);
109        p.push(&name_str[0..PREFIX_DIR_SIZE]);
110        p.push(&format!("{}.larch", name_str));
111
112        p
113    }
114
115    fn path_for_rollup(&self, name: [u32; 5]) -> PathBuf {
116        let mut p = self.path.clone();
117        let name_str = name_to_string(name);
118        p.push(&name_str[0..PREFIX_DIR_SIZE]);
119        p.push(&format!("{}.rollup.hex", name_str));
120
121        p
122    }
123}
124
125#[async_trait]
126impl ArchiveBackend for DirectoryArchiveBackend {
127    type Read = ArchiveSliceReader;
128    async fn get_layer_bytes(&self, id: [u32; 5]) -> io::Result<Bytes> {
129        let path = self.path_for_layer(id);
130        let mut options = fs::OpenOptions::new();
131        options.read(true);
132        options.create(false);
133        let mut result = options.open(path).await?;
134        let metadata = result.metadata().await?;
135        #[cfg(target_os = "windows")]
136        let size = metadata.file_size();
137        #[cfg(not(target_os = "windows"))]
138        let size = metadata.size();
139        let mut buf = Vec::with_capacity(size as usize);
140        result.read_to_end(&mut buf).await?;
141        buf.shrink_to_fit();
142
143        Ok(buf.into())
144    }
145
146    async fn get_layer_structure_bytes(
147        &self,
148        id: [u32; 5],
149        file_type: LayerFileEnum,
150    ) -> io::Result<Option<Bytes>> {
151        let path = self.path_for_layer(id);
152        let mut options = tokio::fs::OpenOptions::new();
153        options.read(true);
154        let mut file = options.open(path).await?;
155        let header = ArchiveHeader::parse_from_reader(&mut file).await?;
156        if let Some(range) = header.range_for(file_type) {
157            let mut data = vec![0; range.len()];
158            file.seek(SeekFrom::Current((range.start) as i64)).await?;
159            file.read_exact(&mut data).await?;
160
161            Ok(Some(Bytes::from(data)))
162        } else {
163            Ok(None)
164        }
165    }
166
167    async fn store_layer_file(&self, id: [u32; 5], mut bytes: Bytes) -> io::Result<()> {
168        let path = self.path_for_layer(id);
169        let mut directory_path = path.clone();
170        directory_path.pop();
171        fs::create_dir_all(&directory_path).await?;
172
173        let mut options = tokio::fs::OpenOptions::new();
174        options.create(true);
175        options.write(true);
176        let mut file = options.open(path).await?;
177        while bytes.remaining() > 0 {
178            let chunk = bytes.chunk();
179            let written = file.write(chunk).await?;
180            bytes.advance(written);
181        }
182
183        file.flush().await?;
184        file.sync_all().await?;
185
186        if cfg!(unix) {
187            // ensure the underlying directory record is properly synchronized
188            let mut options = tokio::fs::OpenOptions::new();
189            options.create(false);
190            options.read(true);
191            options.write(false);
192            let dir_fd = options.open(directory_path).await?;
193            dir_fd.sync_all().await?;
194        }
195
196        Ok(())
197    }
198
199    async fn read_layer_structure_bytes_from(
200        &self,
201        id: [u32; 5],
202        file_type: LayerFileEnum,
203        read_from: usize,
204    ) -> io::Result<Self::Read> {
205        let path = self.path_for_layer(id);
206        let mut options = tokio::fs::OpenOptions::new();
207        options.read(true);
208        let mut file = options.open(path).await?;
209        let header = ArchiveHeader::parse_from_reader(&mut file).await?;
210
211        let range = header
212            .range_for(file_type)
213            .ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "slice not found in archive"))?;
214
215        let remaining = range.len() - read_from;
216        file.seek(SeekFrom::Current((range.start + read_from) as i64))
217            .await?;
218
219        Ok(ArchiveSliceReader { file, remaining })
220    }
221}
222
223#[async_trait]
224impl ArchiveMetadataBackend for DirectoryArchiveBackend {
225    async fn get_layer_names(&self) -> io::Result<Vec<[u32; 5]>> {
226        let mut stream = fs::read_dir(&self.path).await?;
227        let mut result = Vec::new();
228        while let Some(direntry) = stream.next_entry().await? {
229            let os_name = direntry.file_name();
230            let name = os_name.to_str().ok_or_else(|| {
231                io::Error::new(
232                    io::ErrorKind::InvalidData,
233                    "unexpected non-utf8 directory name",
234                )
235            })?;
236            if name.ends_with(".larch") && direntry.file_type().await?.is_file() {
237                let name_component = &name[..name.len() - 6];
238                result.push(string_to_name(name_component)?);
239            }
240        }
241
242        Ok(result)
243    }
244
245    async fn layer_exists(&self, id: [u32; 5]) -> io::Result<bool> {
246        let path = self.path_for_layer(id);
247        let metadata = tokio::fs::metadata(path).await;
248        if metadata.is_err() && metadata.as_ref().err().unwrap().kind() == io::ErrorKind::NotFound {
249            // layer itself not found
250            return Ok(false);
251        }
252        // propagate error if it was anything but NotFound
253        metadata?;
254
255        // if we got here it means the layer exists
256        Ok(true)
257    }
258
259    async fn layer_size(&self, id: [u32; 5]) -> io::Result<u64> {
260        let path = self.path_for_layer(id);
261        let metadata = tokio::fs::metadata(path).await?;
262        Ok(metadata.len())
263    }
264
265    async fn layer_file_exists(&self, id: [u32; 5], file_type: LayerFileEnum) -> io::Result<bool> {
266        let path = self.path_for_layer(id);
267        let metadata = tokio::fs::metadata(&path).await;
268        if metadata.is_err() && metadata.as_ref().err().unwrap().kind() == io::ErrorKind::NotFound {
269            // layer itself not found
270            return Ok(false);
271        }
272        // propagate error if it was anything but NotFound
273        metadata?;
274
275        // read header!
276        let mut options = tokio::fs::OpenOptions::new();
277        options.read(true);
278        let mut file = options.open(path).await?;
279        let header = ArchiveFilePresenceHeader::new(file.read_u64().await?);
280
281        Ok(header.is_present(file_type))
282    }
283
284    async fn get_layer_structure_size(
285        &self,
286        id: [u32; 5],
287        file_type: LayerFileEnum,
288    ) -> io::Result<usize> {
289        let path = self.path_for_layer(id);
290        // read header!
291        let mut options = tokio::fs::OpenOptions::new();
292        options.read(true);
293        let mut file = options.open(path).await?;
294        let header = ArchiveHeader::parse_from_reader(&mut file).await?;
295
296        header
297            .size_of(file_type)
298            .ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "slice not found in archive"))
299    }
300
301    async fn get_rollup(&self, id: [u32; 5]) -> io::Result<Option<[u32; 5]>> {
302        // acquire a shared lock on the layer. This ensures nobody will write a rollup file while we're retrieving it.
303        let layer_path = self.path_for_layer(id);
304        let layer_lock = LockedFile::open(layer_path).await;
305        if layer_lock.is_err() && layer_lock.as_ref().err().unwrap().kind() == ErrorKind::NotFound {
306            // no such layer - therefore no such rollup
307            return Ok(None);
308        }
309        let _layer_lock = layer_lock.unwrap();
310
311        let path = self.path_for_rollup(id);
312        let result = fs::read_to_string(path).await;
313
314        if result.is_err() && result.as_ref().err().unwrap().kind() == ErrorKind::NotFound {
315            return Ok(None);
316        }
317        let data = result?;
318        let name = data.lines().skip(1).next().expect(
319            "Expected rollup file to have two lines but was unable to skip to the second line",
320        );
321        Ok(Some(string_to_name(&name)?))
322    }
323
324    async fn set_rollup(&self, id: [u32; 5], rollup: [u32; 5]) -> io::Result<()> {
325        // acquire an exclusive lock on the layer. This ensures nobody tries to lookup the rollup while we're writing it.
326        let layer_path = self.path_for_layer(id);
327        let _layer_lock = ExclusiveLockedFile::open(layer_path).await?;
328
329        let path = self.path_for_rollup(id);
330        let mut data = Vec::with_capacity(43);
331        data.extend_from_slice(b"1\n");
332        data.extend_from_slice(name_to_string(rollup).as_bytes());
333        data.extend_from_slice(b"\n");
334        let mut file = fs::File::create(path).await?;
335        file.write_all(&data).await?;
336        file.flush().await?;
337        file.sync_all().await?;
338
339        Ok(())
340    }
341
342    async fn get_parent(&self, id: [u32; 5]) -> io::Result<Option<[u32; 5]>> {
343        if let Some(parent_bytes) = self
344            .get_layer_structure_bytes(id, LayerFileEnum::Parent)
345            .await?
346        {
347            let parent_string = std::str::from_utf8(&parent_bytes[..40]).unwrap();
348            Ok(Some(string_to_name(parent_string).unwrap()))
349        } else {
350            Ok(None)
351        }
352    }
353}
354
355#[derive(Clone)]
356pub struct LruArchiveBackend<M, D> {
357    cache: Arc<tokio::sync::Mutex<LruCache<[u32; 5], CacheEntry>>>,
358    limit: usize,
359    current: usize,
360    metadata_origin: M,
361    data_origin: D,
362}
363
364#[derive(Clone)]
365enum CacheEntry {
366    Resolving(Arc<tokio::sync::RwLock<Option<Result<Bytes, io::ErrorKind>>>>),
367    Resolved(Bytes),
368}
369
370impl CacheEntry {
371    fn is_resolving(&self) -> bool {
372        if let Self::Resolving(_) = self {
373            true
374        } else {
375            false
376        }
377    }
378}
379
380impl<M, D> LruArchiveBackend<M, D> {
381    pub fn new(metadata_origin: M, data_origin: D, limit: usize) -> Self {
382        let cache = Arc::new(tokio::sync::Mutex::new(LruCache::unbounded()));
383
384        Self {
385            cache,
386            limit,
387            current: 0,
388            metadata_origin,
389            data_origin,
390        }
391    }
392
393    fn limit_bytes(&self) -> usize {
394        self.limit * 1024 * 1024
395    }
396}
397
398impl<M: ArchiveMetadataBackend, D: ArchiveBackend> LruArchiveBackend<M, D> {
399    async fn layer_fits_in_cache(&self, id: [u32; 5]) -> io::Result<bool> {
400        let limit = self.limit_bytes();
401        Ok(limit != 0 && self.layer_size(id).await? as usize <= limit)
402    }
403}
404
405fn ensure_additional_cache_space(cache: &mut LruCache<[u32; 5], CacheEntry>, mut required: usize) {
406    if required == 0 {
407        return;
408    }
409
410    loop {
411        let peek = cache
412            .peek_lru()
413            .expect("cache is empty but stored entries were expected");
414        if peek.1.is_resolving() {
415            // this is a resolving entry, we don't want to pop it.
416            let id = peek.0.clone();
417            cache.promote(&id);
418            continue;
419        }
420        // at this point the lru item is not resolving
421        let entry = cache
422            .pop_lru()
423            .expect("cache is empty but stored entries were expected")
424            .1;
425        if let CacheEntry::Resolved(entry) = entry {
426            if entry.len() >= required {
427                // done!
428                return;
429            }
430
431            // more needs to be popped
432            required -= entry.len();
433        } else {
434            panic!("expected resolved entry but got a resolving");
435        }
436    }
437}
438
439fn ensure_enough_cache_space(
440    cache: &mut LruCache<[u32; 5], CacheEntry>,
441    limit: usize,
442    current: usize,
443    required: usize,
444) -> bool {
445    if required > limit {
446        // this entry is too big for the cache
447        return false;
448    }
449
450    let remaining = limit - current;
451    if remaining < required {
452        // we need to clean up some cache spacew to fit this entry
453        ensure_additional_cache_space(cache, required - remaining);
454    }
455
456    true
457}
458
459fn drop_from_cache(cache: &mut LruCache<[u32; 5], CacheEntry>, id: [u32; 5]) {
460    assert!(cache.contains(&id));
461    cache.demote(&id);
462    cache.pop_lru();
463}
464
465#[async_trait]
466impl<M: ArchiveMetadataBackend, D: ArchiveBackend> ArchiveBackend for LruArchiveBackend<M, D> {
467    type Read = Either<BytesAsyncReader, D::Read>;
468    async fn get_layer_bytes(&self, id: [u32; 5]) -> io::Result<Bytes> {
469        let mut cache = self.cache.lock().await;
470        let cached = cache.get(&id).cloned();
471
472        match cached {
473            Some(CacheEntry::Resolved(bytes)) => Ok(bytes),
474            Some(CacheEntry::Resolving(barrier)) => {
475                // someone is already looking up this layer. we'll wait for them to be done.
476                std::mem::drop(cache);
477                let guard = barrier.read().await;
478                match guard.as_ref().unwrap() {
479                    Ok(bytes) => Ok(bytes.clone()),
480                    Err(kind) => Err(io::Error::new(*kind, "layer resolve failed")),
481                }
482            }
483            None => {
484                // nobody is looking this up yet, it is up to us.
485                let barrier = Arc::new(tokio::sync::RwLock::new(None));
486                let mut result = barrier.write().await;
487                cache.get_or_insert(id, || CacheEntry::Resolving(barrier.clone()));
488
489                // drop the cache while doing the lookup
490                std::mem::drop(cache);
491                let lookup = self.data_origin.get_layer_bytes(id).await;
492
493                *result = Some(lookup.as_ref().map_err(|e| e.kind()).cloned());
494
495                // reacquire cache
496                let mut cache = self.cache.lock().await;
497                match lookup {
498                    Ok(bytes) => {
499                        if ensure_enough_cache_space(
500                            &mut *cache,
501                            self.limit_bytes(),
502                            self.current,
503                            bytes.len(),
504                        ) {
505                            let cached = cache
506                                .get_mut(&id)
507                                .expect("layer resolving entry not found in cache");
508                            *cached = CacheEntry::Resolved(bytes.clone());
509                        } else {
510                            // this entry is uncachable. Just remove the resolving entry
511                            drop_from_cache(&mut *cache, id);
512                        }
513                        Ok(bytes)
514                    }
515                    Err(e) => {
516                        drop_from_cache(&mut *cache, id);
517
518                        Err(e)
519                    }
520                }
521            }
522        }
523    }
524    async fn get_layer_structure_bytes(
525        &self,
526        id: [u32; 5],
527        file_type: LayerFileEnum,
528    ) -> io::Result<Option<Bytes>> {
529        if self.layer_fits_in_cache(id).await? {
530            let bytes = self.get_layer_bytes(id).await?;
531            let archive = Archive::parse(bytes);
532            Ok(archive.slice_for(file_type))
533        } else {
534            self.data_origin
535                .get_layer_structure_bytes(id, file_type)
536                .await
537        }
538    }
539    async fn store_layer_file(&self, id: [u32; 5], bytes: Bytes) -> io::Result<()> {
540        self.data_origin.store_layer_file(id, bytes.clone()).await?;
541
542        let mut cache = self.cache.lock().await;
543        cache.get_or_insert(id, move || CacheEntry::Resolved(bytes));
544
545        Ok(())
546    }
547    async fn read_layer_structure_bytes_from(
548        &self,
549        id: [u32; 5],
550        file_type: LayerFileEnum,
551        read_from: usize,
552    ) -> io::Result<Self::Read> {
553        if self.layer_fits_in_cache(id).await? {
554            let mut bytes = self
555                .get_layer_structure_bytes(id, file_type)
556                .await?
557                .ok_or_else(|| {
558                    io::Error::new(io::ErrorKind::NotFound, "slice not found in archive")
559                })?;
560            bytes.advance(read_from);
561
562            Ok(Either::Left(BytesAsyncReader(bytes)))
563        } else {
564            Ok(Either::Right(
565                self.data_origin
566                    .read_layer_structure_bytes_from(id, file_type, read_from)
567                    .await?,
568            ))
569        }
570    }
571}
572
573#[async_trait]
574impl<M: ArchiveMetadataBackend, D: ArchiveBackend> ArchiveMetadataBackend
575    for LruArchiveBackend<M, D>
576{
577    async fn get_layer_names(&self) -> io::Result<Vec<[u32; 5]>> {
578        self.metadata_origin.get_layer_names().await
579    }
580    async fn layer_exists(&self, id: [u32; 5]) -> io::Result<bool> {
581        if let Some(CacheEntry::Resolved(_)) = self.cache.lock().await.peek(&id) {
582            Ok(true)
583        } else {
584            self.metadata_origin.layer_exists(id).await
585        }
586    }
587    async fn layer_size(&self, id: [u32; 5]) -> io::Result<u64> {
588        if let Some(CacheEntry::Resolved(bytes)) = self.cache.lock().await.peek(&id) {
589            Ok(bytes.len() as u64)
590        } else {
591            self.metadata_origin.layer_size(id).await
592        }
593    }
594    async fn layer_file_exists(&self, id: [u32; 5], file_type: LayerFileEnum) -> io::Result<bool> {
595        if let Some(CacheEntry::Resolved(bytes)) = self.cache.lock().await.peek(&id) {
596            let header = ArchiveFilePresenceHeader::new(bytes.clone().get_u64());
597            Ok(header.is_present(file_type))
598        } else {
599            self.metadata_origin.layer_file_exists(id, file_type).await
600        }
601    }
602    async fn get_layer_structure_size(
603        &self,
604        id: [u32; 5],
605        file_type: LayerFileEnum,
606    ) -> io::Result<usize> {
607        if let Some(CacheEntry::Resolved(bytes)) = self.cache.lock().await.peek(&id) {
608            let (header, _) = ArchiveHeader::parse(bytes.clone());
609
610            if let Some(size) = header.size_of(file_type) {
611                Ok(size)
612            } else {
613                Err(io::Error::new(
614                    io::ErrorKind::NotFound,
615                    format!(
616                        "structure {file_type:?} not found in layer {}",
617                        name_to_string(id)
618                    ),
619                ))
620            }
621        } else {
622            self.metadata_origin
623                .get_layer_structure_size(id, file_type)
624                .await
625        }
626    }
627    async fn get_rollup(&self, id: [u32; 5]) -> io::Result<Option<[u32; 5]>> {
628        self.metadata_origin.get_rollup(id).await
629    }
630    async fn set_rollup(&self, id: [u32; 5], rollup: [u32; 5]) -> io::Result<()> {
631        self.metadata_origin.set_rollup(id, rollup).await
632    }
633
634    async fn get_parent(&self, id: [u32; 5]) -> io::Result<Option<[u32; 5]>> {
635        if let Some(parent_bytes) = self
636            .get_layer_structure_bytes(id, LayerFileEnum::Parent)
637            .await?
638        {
639            let parent_string = std::str::from_utf8(&parent_bytes[..40]).unwrap();
640            Ok(Some(string_to_name(parent_string).unwrap()))
641        } else {
642            Ok(None)
643        }
644    }
645}
646
647pub enum ConstructionFileState {
648    UnderConstruction(BytesMut),
649    Finalizing,
650    Finalized(Bytes),
651}
652
653#[derive(Clone)]
654pub struct ConstructionFile(Arc<RwLock<ConstructionFileState>>);
655
656impl ConstructionFile {
657    fn new() -> Self {
658        Self(Arc::new(RwLock::new(
659            ConstructionFileState::UnderConstruction(BytesMut::new()),
660        )))
661    }
662
663    fn new_finalized(bytes: Bytes) -> Self {
664        Self(Arc::new(RwLock::new(ConstructionFileState::Finalized(
665            bytes,
666        ))))
667    }
668
669    fn is_finalized(&self) -> bool {
670        let guard = self.0.read().unwrap();
671        if let ConstructionFileState::Finalized(_) = &*guard {
672            true
673        } else {
674            false
675        }
676    }
677
678    fn finalized_buf(self) -> Bytes {
679        let guard = self.0.read().unwrap();
680        if let ConstructionFileState::Finalized(bytes) = &*guard {
681            bytes.clone()
682        } else {
683            panic!("tried to get the finalized buf from an unfinalized ConstructionFile");
684        }
685    }
686}
687
688#[async_trait]
689impl FileStore for ConstructionFile {
690    type Write = Self;
691    async fn open_write(&self) -> io::Result<Self::Write> {
692        Ok(self.clone())
693    }
694}
695
696impl AsyncWrite for ConstructionFile {
697    fn poll_write(
698        self: std::pin::Pin<&mut Self>,
699        _cx: &mut std::task::Context<'_>,
700        buf: &[u8],
701    ) -> Poll<Result<usize, io::Error>> {
702        let mut guard = self.0.write().unwrap();
703        match &mut *guard {
704            ConstructionFileState::UnderConstruction(x) => {
705                x.put_slice(buf);
706
707                Poll::Ready(Ok(buf.len()))
708            }
709            _ => Poll::Ready(Err(io::Error::new(
710                io::ErrorKind::Other,
711                "file already written",
712            ))),
713        }
714    }
715
716    fn poll_flush(
717        self: std::pin::Pin<&mut Self>,
718        _cx: &mut std::task::Context<'_>,
719    ) -> Poll<Result<(), io::Error>> {
720        // noop
721        Poll::Ready(Ok(()))
722    }
723
724    fn poll_shutdown(
725        self: std::pin::Pin<&mut Self>,
726        _cx: &mut std::task::Context<'_>,
727    ) -> Poll<Result<(), io::Error>> {
728        // noop
729        Poll::Ready(Ok(()))
730    }
731}
732
733#[async_trait]
734impl SyncableFile for ConstructionFile {
735    async fn sync_all(self) -> io::Result<()> {
736        let mut guard = self.0.write().unwrap();
737        let mut state = ConstructionFileState::Finalizing;
738        std::mem::swap(&mut state, &mut *guard);
739
740        match state {
741            ConstructionFileState::UnderConstruction(x) => {
742                let buf = x.freeze();
743                *guard = ConstructionFileState::Finalized(buf);
744
745                Ok(())
746            }
747            _ => {
748                *guard = state;
749                Err(io::Error::new(io::ErrorKind::Other, "file already written"))
750            }
751        }
752    }
753}
754
755impl AsyncRead for ConstructionFile {
756    fn poll_read(
757        self: Pin<&mut Self>,
758        _cx: &mut std::task::Context<'_>,
759        buf: &mut tokio::io::ReadBuf<'_>,
760    ) -> Poll<io::Result<()>> {
761        let mut guard = self.0.write().unwrap();
762        match &mut *guard {
763            ConstructionFileState::Finalized(x) => {
764                let slice = if buf.remaining() > x.len() {
765                    x.split_to(x.len())
766                } else {
767                    x.split_to(buf.remaining())
768                };
769                buf.put_slice(slice.as_ref());
770
771                Poll::Ready(Ok(()))
772            }
773            _ => Poll::Ready(Err(io::Error::new(
774                io::ErrorKind::Other,
775                "file not yet written",
776            ))),
777        }
778    }
779}
780
781#[async_trait]
782impl FileLoad for ConstructionFile {
783    type Read = Self;
784
785    async fn exists(&self) -> io::Result<bool> {
786        let guard = self.0.read().unwrap();
787        Ok(matches!(&*guard, ConstructionFileState::Finalized(_)))
788    }
789    async fn size(&self) -> io::Result<usize> {
790        let guard = self.0.read().unwrap();
791        match &*guard {
792            ConstructionFileState::Finalized(x) => Ok(x.len()),
793            _ => Err(io::Error::new(
794                io::ErrorKind::NotFound,
795                "file not finalized",
796            )),
797        }
798    }
799
800    async fn open_read_from(&self, offset: usize) -> io::Result<Self::Read> {
801        let guard = self.0.read().unwrap();
802        match &*guard {
803            ConstructionFileState::Finalized(data) => {
804                let mut data = data.clone();
805                if data.len() < offset {
806                    Err(io::Error::new(
807                        io::ErrorKind::UnexpectedEof,
808                        "offset is beyond end of file",
809                    ))
810                } else {
811                    data.advance(offset);
812                    // this is suspicious, why would we need a lock here? Maybe we should have a different reader type from the file type
813                    Ok(ConstructionFile(Arc::new(RwLock::new(
814                        ConstructionFileState::Finalized(data),
815                    ))))
816                }
817            }
818            _ => Err(io::Error::new(
819                io::ErrorKind::NotFound,
820                "file not finalized",
821            )),
822        }
823    }
824
825    async fn map(&self) -> io::Result<Bytes> {
826        let guard = self.0.read().unwrap();
827        match &*guard {
828            ConstructionFileState::Finalized(x) => Ok(x.clone()),
829            _ => Err(io::Error::new(
830                io::ErrorKind::NotFound,
831                "file not finalized",
832            )),
833        }
834    }
835}
836
837#[derive(Debug, Clone)]
838pub struct ArchiveFilePresenceHeader {
839    present_files: SmallBitArray,
840}
841
842impl ArchiveFilePresenceHeader {
843    pub fn new(val: u64) -> Self {
844        Self {
845            present_files: SmallBitArray::new(val),
846        }
847    }
848
849    pub fn from_present<I: Iterator<Item = LayerFileEnum>>(present_files: I) -> Self {
850        let mut val = 0;
851
852        for file in present_files {
853            val |= 1 << (u64::BITS - file as u32 - 1);
854        }
855
856        Self::new(val)
857    }
858
859    pub fn is_present(&self, file: LayerFileEnum) -> bool {
860        self.present_files.get(file as usize)
861    }
862
863    pub fn inner(&self) -> u64 {
864        self.present_files.inner()
865    }
866
867    pub fn file_index(&self, file: LayerFileEnum) -> Option<usize> {
868        if !self.is_present(file) {
869            return None;
870        }
871
872        Some(self.present_files.rank1(file as usize) - 1)
873    }
874}
875
876#[derive(Debug, Clone)]
877pub struct ArchiveHeader {
878    file_presence: ArchiveFilePresenceHeader,
879    file_offsets: MonotonicLogArray,
880}
881
882impl ArchiveHeader {
883    pub fn parse(mut bytes: Bytes) -> (Self, Bytes) {
884        let file_presence = ArchiveFilePresenceHeader::new(bytes.get_u64());
885        let (file_offsets, remainder) = MonotonicLogArray::parse_header_first(bytes)
886            .expect("unable to parse structure offsets");
887
888        (
889            Self {
890                file_presence,
891                file_offsets,
892            },
893            remainder,
894        )
895    }
896
897    pub async fn parse_from_reader<R: AsyncRead + Unpin>(reader: &mut R) -> io::Result<Self> {
898        let file_presence = ArchiveFilePresenceHeader::new(reader.read_u64().await?);
899        let mut logarray_bytes = BytesMut::new();
900        logarray_bytes.resize(8, 0);
901        reader.read_exact(&mut logarray_bytes[0..8]).await?;
902        let len = logarray_length_from_control_word(&logarray_bytes[0..8]);
903        logarray_bytes.reserve(len);
904        unsafe {
905            logarray_bytes.set_len(8 + len);
906        }
907        reader.read_exact(&mut logarray_bytes[8..]).await?;
908
909        let (file_offsets, _) =
910            MonotonicLogArray::parse_header_first(logarray_bytes.freeze()).expect("what the heck");
911
912        Ok(Self {
913            file_presence,
914            file_offsets,
915        })
916    }
917
918    pub fn range_for(&self, file: LayerFileEnum) -> Option<Range<usize>> {
919        if let Some(file_index) = self.file_presence.file_index(file) {
920            let start: usize = if file_index == 0 {
921                0
922            } else {
923                self.file_offsets.entry(file_index - 1) as usize
924            };
925
926            let end: usize = self.file_offsets.entry(file_index) as usize;
927
928            Some(start..end)
929        } else {
930            None
931        }
932    }
933
934    pub fn size_of(&self, file: LayerFileEnum) -> Option<usize> {
935        self.range_for(file).map(|range| range.end - range.start)
936    }
937}
938
939pub struct Archive {
940    pub header: ArchiveHeader,
941    pub data: Bytes,
942}
943
944impl Archive {
945    pub fn parse(bytes: Bytes) -> Self {
946        let (header, data) = ArchiveHeader::parse(bytes);
947
948        Self { header, data }
949    }
950
951    pub async fn parse_from_reader<R: AsyncRead + Unpin>(reader: &mut R) -> io::Result<Self> {
952        let header = ArchiveHeader::parse_from_reader(reader).await?;
953        let data_len = header.file_offsets.entry(header.file_offsets.len() - 1) as usize;
954        let mut data = BytesMut::with_capacity(data_len);
955        data.reserve(data_len);
956        unsafe { data.set_len(data_len) };
957        reader.read_exact(&mut data[..]).await?;
958
959        Ok(Self {
960            header,
961            data: data.freeze(),
962        })
963    }
964
965    pub fn slice_for(&self, file: LayerFileEnum) -> Option<Bytes> {
966        self.header
967            .range_for(file)
968            .map(|range| self.data.slice(range))
969    }
970
971    pub fn size_of(&self, file: LayerFileEnum) -> Option<usize> {
972        self.header.size_of(file)
973    }
974}
975
976#[derive(Clone)]
977pub struct PersistentFileSlice<M, D> {
978    metadata_backend: M,
979    data_backend: D,
980    layer_id: [u32; 5],
981    file_type: LayerFileEnum,
982}
983
984impl<M, D> PersistentFileSlice<M, D> {
985    fn new(
986        metadata_backend: M,
987        data_backend: D,
988        layer_id: [u32; 5],
989        file_type: LayerFileEnum,
990    ) -> Self {
991        Self {
992            metadata_backend,
993            data_backend,
994            layer_id,
995            file_type,
996        }
997    }
998}
999
1000pub struct ArchiveSliceReader {
1001    file: File,
1002    remaining: usize,
1003}
1004
1005impl ArchiveSliceReader {
1006    pub fn new(file: File, remaining: usize) -> Self {
1007        Self { file, remaining }
1008    }
1009
1010    pub fn end_early(&mut self, end: usize) {
1011        self.remaining -= end;
1012    }
1013}
1014
1015impl AsyncRead for ArchiveSliceReader {
1016    fn poll_read(
1017        mut self: Pin<&mut Self>,
1018        cx: &mut std::task::Context<'_>,
1019        buf: &mut tokio::io::ReadBuf<'_>,
1020    ) -> Poll<io::Result<()>> {
1021        if self.remaining == 0 {
1022            return Poll::Ready(Ok(()));
1023        }
1024
1025        let before_len = buf.filled().len();
1026        let read = AsyncRead::poll_read(Pin::new(&mut self.file), cx, buf);
1027        if let Poll::Pending = read {
1028            return Poll::Pending;
1029        }
1030
1031        let after_len = buf.filled().len();
1032        let read_len = after_len - before_len;
1033
1034        if read_len > self.remaining {
1035            buf.set_filled(before_len + self.remaining);
1036            self.remaining = 0;
1037        } else {
1038            self.remaining -= read_len;
1039        }
1040
1041        Poll::Ready(Ok(()))
1042    }
1043}
1044
1045#[async_trait]
1046impl<M: ArchiveMetadataBackend, D: ArchiveBackend> FileLoad for PersistentFileSlice<M, D> {
1047    type Read = D::Read;
1048
1049    async fn exists(&self) -> io::Result<bool> {
1050        self.metadata_backend
1051            .layer_file_exists(self.layer_id, self.file_type)
1052            .await
1053    }
1054
1055    async fn size(&self) -> io::Result<usize> {
1056        self.metadata_backend
1057            .get_layer_structure_size(self.layer_id, self.file_type)
1058            .await
1059    }
1060
1061    async fn open_read_from(&self, offset: usize) -> io::Result<Self::Read> {
1062        self.data_backend
1063            .read_layer_structure_bytes_from(self.layer_id, self.file_type, offset)
1064            .await
1065    }
1066
1067    async fn map(&self) -> io::Result<Bytes> {
1068        self.data_backend
1069            .get_layer_structure_bytes(self.layer_id, self.file_type)
1070            .await?
1071            .ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "slice not found in archive"))
1072    }
1073
1074    async fn map_if_exists(&self) -> io::Result<Option<Bytes>> {
1075        self.data_backend
1076            .get_layer_structure_bytes(self.layer_id, self.file_type)
1077            .await
1078    }
1079}
1080
1081// This is some pretty ridiculous contrived logic but it saves having to refactor some other places which should just take a rollup id in the first place.
1082#[derive(Clone)]
1083pub struct ArchiveRollupFile<M> {
1084    layer_id: [u32; 5],
1085    metadata_backend: M,
1086}
1087
1088#[async_trait]
1089impl<M: ArchiveMetadataBackend> FileLoad for ArchiveRollupFile<M> {
1090    type Read = BytesAsyncReader;
1091
1092    async fn exists(&self) -> io::Result<bool> {
1093        Ok(self
1094            .metadata_backend
1095            .get_rollup(self.layer_id)
1096            .await?
1097            .is_some())
1098    }
1099
1100    async fn size(&self) -> io::Result<usize> {
1101        if self
1102            .metadata_backend
1103            .get_rollup(self.layer_id)
1104            .await?
1105            .is_some()
1106        {
1107            Ok(std::mem::size_of::<[u32; 5]>() + 2)
1108        } else {
1109            Err(io::Error::new(
1110                io::ErrorKind::NotFound,
1111                "layer has no rollup",
1112            ))
1113        }
1114    }
1115
1116    async fn open_read_from(&self, offset: usize) -> io::Result<Self::Read> {
1117        let mut bytes = self.map().await?;
1118        bytes.advance(offset);
1119        Ok(BytesAsyncReader(bytes))
1120    }
1121
1122    async fn map(&self) -> io::Result<Bytes> {
1123        let id = self.metadata_backend.get_rollup(self.layer_id).await?;
1124        if let Some(id) = id {
1125            let mut bytes = Vec::with_capacity(42);
1126            bytes.extend_from_slice(b"1\n");
1127            bytes.extend_from_slice(name_to_string(id).as_bytes());
1128            Ok(bytes.into())
1129        } else {
1130            Err(io::Error::new(
1131                io::ErrorKind::NotFound,
1132                "layer has no rollup",
1133            ))
1134        }
1135    }
1136}
1137
1138#[async_trait]
1139impl<M: ArchiveMetadataBackend + Unpin> FileStore for ArchiveRollupFile<M> {
1140    type Write = ArchiveRollupFileWriter<M>;
1141    async fn open_write(&self) -> io::Result<Self::Write> {
1142        Ok(ArchiveRollupFileWriter {
1143            layer_id: self.layer_id,
1144            data: BytesMut::new(),
1145            metadata_backend: self.metadata_backend.clone(),
1146        })
1147    }
1148}
1149
1150pub struct ArchiveRollupFileWriter<M> {
1151    layer_id: [u32; 5],
1152    data: BytesMut,
1153    metadata_backend: M,
1154}
1155
1156impl<M: ArchiveMetadataBackend + Unpin> AsyncWrite for ArchiveRollupFileWriter<M> {
1157    fn poll_write(
1158        self: std::pin::Pin<&mut Self>,
1159        _cx: &mut std::task::Context<'_>,
1160        buf: &[u8],
1161    ) -> Poll<Result<usize, io::Error>> {
1162        self.get_mut().data.extend_from_slice(buf);
1163
1164        Poll::Ready(Ok(buf.len()))
1165    }
1166
1167    fn poll_flush(
1168        self: Pin<&mut Self>,
1169        _cx: &mut std::task::Context<'_>,
1170    ) -> Poll<Result<(), io::Error>> {
1171        Poll::Ready(Ok(()))
1172    }
1173
1174    fn poll_shutdown(
1175        self: Pin<&mut Self>,
1176        _cx: &mut std::task::Context<'_>,
1177    ) -> Poll<Result<(), io::Error>> {
1178        Poll::Ready(Ok(()))
1179    }
1180}
1181
1182#[async_trait]
1183impl<M: ArchiveMetadataBackend + Unpin> SyncableFile for ArchiveRollupFileWriter<M> {
1184    async fn sync_all(self) -> io::Result<()> {
1185        let rollup_string =
1186            String::from_utf8(self.data.to_vec()).expect("rollup id was not a string");
1187        // first line of this string is going to be a version number. it should be discarded.
1188        let line = rollup_string.lines().skip(1).next().unwrap();
1189        let rollup_id = string_to_name(&line)?;
1190
1191        self.metadata_backend
1192            .set_rollup(self.layer_id, rollup_id)
1193            .await
1194    }
1195}
1196
1197#[derive(Clone)]
1198pub enum ArchiveLayerHandle<M, D> {
1199    Construction(ConstructionFile),
1200    Persistent(PersistentFileSlice<M, D>),
1201    Rollup(ArchiveRollupFile<M>),
1202}
1203
1204#[async_trait]
1205impl<M: ArchiveMetadataBackend + Unpin, D: ArchiveBackend> FileStore for ArchiveLayerHandle<M, D> {
1206    type Write = ArchiveLayerHandleWriter<M>;
1207    async fn open_write(&self) -> io::Result<Self::Write> {
1208        Ok(match self {
1209            Self::Construction(c) => ArchiveLayerHandleWriter::Construction(c.open_write().await?),
1210            Self::Rollup(r) => ArchiveLayerHandleWriter::Rollup(r.open_write().await?),
1211            _ => panic!("cannot write to a persistent file slice"),
1212        })
1213    }
1214}
1215
1216#[async_trait]
1217impl<M: ArchiveMetadataBackend, D: ArchiveBackend> FileLoad for ArchiveLayerHandle<M, D> {
1218    type Read = ArchiveLayerHandleReader<D::Read, BytesAsyncReader>;
1219
1220    async fn exists(&self) -> io::Result<bool> {
1221        match self {
1222            Self::Construction(c) => c.exists().await,
1223            Self::Persistent(p) => p.exists().await,
1224            Self::Rollup(r) => r.exists().await,
1225        }
1226    }
1227    async fn size(&self) -> io::Result<usize> {
1228        match self {
1229            Self::Construction(c) => c.size().await,
1230            Self::Persistent(p) => p.size().await,
1231            Self::Rollup(r) => r.size().await,
1232        }
1233    }
1234
1235    async fn open_read_from(&self, offset: usize) -> io::Result<Self::Read> {
1236        Ok(match self {
1237            Self::Construction(c) => {
1238                ArchiveLayerHandleReader::Construction(c.open_read_from(offset).await?)
1239            }
1240            Self::Persistent(p) => {
1241                ArchiveLayerHandleReader::Persistent(p.open_read_from(offset).await?)
1242            }
1243            Self::Rollup(r) => ArchiveLayerHandleReader::Rollup(r.open_read_from(offset).await?),
1244        })
1245    }
1246
1247    async fn map(&self) -> io::Result<Bytes> {
1248        match self {
1249            Self::Construction(c) => c.map().await,
1250            Self::Persistent(p) => p.map().await,
1251            Self::Rollup(r) => r.map().await,
1252        }
1253    }
1254}
1255
1256pub enum ArchiveLayerHandleReader<P, R> {
1257    Construction(ConstructionFile),
1258    Persistent(P),
1259    Rollup(R),
1260}
1261
1262impl<P: AsyncRead + Unpin, R: AsyncRead + Unpin> AsyncRead for ArchiveLayerHandleReader<P, R> {
1263    fn poll_read(
1264        mut self: Pin<&mut Self>,
1265        cx: &mut std::task::Context<'_>,
1266        buf: &mut tokio::io::ReadBuf<'_>,
1267    ) -> Poll<io::Result<()>> {
1268        match &mut *self {
1269            Self::Construction(c) => AsyncRead::poll_read(Pin::new(c), cx, buf),
1270            Self::Persistent(p) => AsyncRead::poll_read(Pin::new(p), cx, buf),
1271            Self::Rollup(r) => AsyncRead::poll_read(Pin::new(r), cx, buf),
1272        }
1273    }
1274}
1275
1276pub enum ArchiveLayerHandleWriter<M> {
1277    Construction(ConstructionFile),
1278    Rollup(ArchiveRollupFileWriter<M>),
1279}
1280
1281impl<M: ArchiveMetadataBackend + Unpin> AsyncWrite for ArchiveLayerHandleWriter<M> {
1282    fn poll_write(
1283        mut self: std::pin::Pin<&mut Self>,
1284        cx: &mut std::task::Context<'_>,
1285        buf: &[u8],
1286    ) -> Poll<Result<usize, io::Error>> {
1287        match &mut *self {
1288            Self::Construction(c) => AsyncWrite::poll_write(Pin::new(c), cx, buf),
1289            Self::Rollup(r) => AsyncWrite::poll_write(Pin::new(r), cx, buf),
1290        }
1291    }
1292
1293    fn poll_flush(
1294        mut self: std::pin::Pin<&mut Self>,
1295        cx: &mut std::task::Context<'_>,
1296    ) -> Poll<Result<(), io::Error>> {
1297        match &mut *self {
1298            Self::Construction(c) => AsyncWrite::poll_flush(Pin::new(c), cx),
1299            Self::Rollup(r) => AsyncWrite::poll_flush(Pin::new(r), cx),
1300        }
1301    }
1302
1303    fn poll_shutdown(
1304        mut self: std::pin::Pin<&mut Self>,
1305        cx: &mut std::task::Context<'_>,
1306    ) -> Poll<Result<(), io::Error>> {
1307        match &mut *self {
1308            Self::Construction(c) => AsyncWrite::poll_shutdown(Pin::new(c), cx),
1309            Self::Rollup(r) => AsyncWrite::poll_shutdown(Pin::new(r), cx),
1310        }
1311    }
1312}
1313
1314#[async_trait]
1315impl<M: ArchiveMetadataBackend + Unpin> SyncableFile for ArchiveLayerHandleWriter<M> {
1316    async fn sync_all(self) -> io::Result<()> {
1317        match self {
1318            Self::Construction(c) => c.sync_all().await,
1319            Self::Rollup(r) => r.sync_all().await,
1320        }
1321    }
1322}
1323
1324type ArchiveLayerConstructionMap =
1325    Arc<RwLock<HashMap<[u32; 5], HashMap<LayerFileEnum, ConstructionFile>>>>;
1326
1327#[derive(Clone)]
1328pub struct ArchiveLayerStore<M, D> {
1329    metadata_backend: M,
1330    data_backend: D,
1331    construction: ArchiveLayerConstructionMap,
1332}
1333
1334impl<M, D> ArchiveLayerStore<M, D> {
1335    pub fn new(metadata_backend: M, data_backend: D) -> ArchiveLayerStore<M, D> {
1336        ArchiveLayerStore {
1337            metadata_backend,
1338            data_backend,
1339            construction: Default::default(),
1340        }
1341    }
1342
1343    #[doc(hidden)]
1344    pub fn write_bytes(&self, name: [u32; 5], file: LayerFileEnum, bytes: Bytes) {
1345        let mut guard = self.construction.write().unwrap();
1346        if let Some(map) = guard.get_mut(&name) {
1347            if map.contains_key(&file) {
1348                panic!("tried to write bytes to an archive, but file is already open");
1349            }
1350
1351            map.insert(file, ConstructionFile::new_finalized(bytes));
1352        } else {
1353            panic!("tried to write bytes to an archive, but layer is not under construction");
1354        }
1355    }
1356}
1357
1358const PREFIX_DIR_SIZE: usize = 3;
1359
1360#[async_trait]
1361impl<M: ArchiveMetadataBackend + Unpin + 'static, D: ArchiveBackend + 'static> PersistentLayerStore
1362    for ArchiveLayerStore<M, D>
1363{
1364    type File = ArchiveLayerHandle<M, D>;
1365
1366    async fn directories(&self) -> io::Result<Vec<[u32; 5]>> {
1367        let mut result = self.metadata_backend.get_layer_names().await?;
1368
1369        {
1370            let guard = self.construction.read().unwrap();
1371
1372            for name in guard.keys() {
1373                result.push(*name);
1374            }
1375        }
1376
1377        result.sort();
1378        result.dedup();
1379
1380        Ok(result)
1381    }
1382
1383    async fn create_named_directory(&self, name: [u32; 5]) -> io::Result<[u32; 5]> {
1384        if !self.metadata_backend.layer_exists(name).await? {
1385            // layer does not exist yet on disk, good.
1386            let mut guard = self.construction.write().unwrap();
1387            if guard.contains_key(&name) {
1388                // whoops! Looks like layer is already under construction!
1389                panic!("tried to create a new layer which is already under construction");
1390            }
1391
1392            // layer is neither on disk nor in the construction map. Let's create it.
1393            guard.insert(name, HashMap::new());
1394            return Ok(name);
1395        } else {
1396            // still here? That means the file existed, even though it shouldn't!
1397            panic!("tried to create a new layer which already exists");
1398        }
1399    }
1400
1401    async fn directory_exists(&self, name: [u32; 5]) -> io::Result<bool> {
1402        {
1403            let guard = self.construction.read().unwrap();
1404            if guard.contains_key(&name) {
1405                return Ok(true);
1406            }
1407        }
1408
1409        self.metadata_backend.layer_exists(name).await
1410    }
1411
1412    async fn get_file(&self, directory: [u32; 5], name: &str) -> io::Result<Self::File> {
1413        let file_type = FILENAME_ENUM_MAP[name];
1414        if file_type == LayerFileEnum::Rollup {
1415            // special case! This is always coming from disk, in its own file
1416            return Ok(ArchiveLayerHandle::Rollup(ArchiveRollupFile {
1417                layer_id: directory,
1418                metadata_backend: self.metadata_backend.clone(),
1419            }));
1420        }
1421
1422        {
1423            let guard = self.construction.read().unwrap();
1424            if let Some(map) = guard.get(&directory) {
1425                if let Some(file) = map.get(&file_type) {
1426                    return Ok(ArchiveLayerHandle::Construction(file.clone()));
1427                }
1428
1429                // the directory is there but the file is not. We'll have to construct it.
1430                std::mem::drop(guard);
1431                let mut guard = self.construction.write().unwrap();
1432                let map = guard.get_mut(&directory).unwrap();
1433                let file = ConstructionFile::new();
1434                map.insert(file_type, file.clone());
1435
1436                Ok(ArchiveLayerHandle::Construction(file))
1437            } else {
1438                // layer does not appear to be under construction so it has to be in persistent storage
1439                Ok(ArchiveLayerHandle::Persistent(PersistentFileSlice::new(
1440                    self.metadata_backend.clone(),
1441                    self.data_backend.clone(),
1442                    directory,
1443                    file_type,
1444                )))
1445            }
1446        }
1447    }
1448
1449    async fn file_exists(&self, directory: [u32; 5], file: &str) -> io::Result<bool> {
1450        let file_type = FILENAME_ENUM_MAP[file];
1451        if file_type == LayerFileEnum::Rollup {
1452            // special case! This is always coming out of the persistent metadata
1453            return Ok(self.metadata_backend.get_rollup(directory).await?.is_some());
1454        }
1455
1456        {
1457            let guard = self.construction.read().unwrap();
1458            if let Some(map) = guard.get(&directory) {
1459                return Ok(map.contains_key(&file_type));
1460            }
1461        }
1462
1463        self.metadata_backend
1464            .layer_file_exists(directory, file_type)
1465            .await
1466    }
1467
1468    async fn finalize(&self, directory: [u32; 5]) -> io::Result<()> {
1469        let files = {
1470            let mut guard = self.construction.write().unwrap();
1471            guard
1472                .remove(&directory)
1473                .expect("layer to be finalized was not found in construction map")
1474        };
1475
1476        let mut files: Vec<(_, _)> = files
1477            .into_iter()
1478            .filter(|(_file_type, file)| file.is_finalized())
1479            .map(|(file_type, file)| (file_type, file.finalized_buf()))
1480            .collect();
1481        files.sort();
1482        let presence_header =
1483            ArchiveFilePresenceHeader::from_present(files.iter().map(|(t, _)| t).cloned());
1484
1485        let mut offsets = LateLogArrayBufBuilder::new(BytesMut::new());
1486        let mut tally = 0;
1487        for (_file_type, data) in files.iter() {
1488            tally += data.len();
1489            offsets.push(tally as u64);
1490        }
1491
1492        let offsets_buf = offsets.finalize_header_first();
1493
1494        let mut data_buf = BytesMut::with_capacity(tally + 8 + offsets_buf.len());
1495        data_buf.put_u64(presence_header.inner());
1496        data_buf.extend(offsets_buf);
1497        for (_file_type, data) in files {
1498            data_buf.extend(data);
1499        }
1500
1501        self.data_backend
1502            .store_layer_file(directory, data_buf.freeze())
1503            .await
1504    }
1505
1506    async fn layer_parent(&self, name: [u32; 5]) -> io::Result<Option<[u32; 5]>> {
1507        self.metadata_backend.get_parent(name).await
1508    }
1509}
1510
1511#[cfg(test)]
1512mod tests {
1513    use super::*;
1514
1515    #[test]
1516    fn parse_and_use_header() {
1517        let files_present = vec![
1518            LayerFileEnum::PredicateDictionaryBlocks,
1519            LayerFileEnum::NegObjects,
1520            LayerFileEnum::NodeValueIdMapBits,
1521            LayerFileEnum::Parent,
1522            LayerFileEnum::Rollup,
1523            LayerFileEnum::NodeDictionaryBlocks,
1524        ];
1525
1526        let header = ArchiveFilePresenceHeader::from_present(files_present.iter().cloned());
1527
1528        for file in files_present {
1529            assert!(header.is_present(file));
1530        }
1531
1532        assert!(!header.is_present(LayerFileEnum::NodeDictionaryOffsets));
1533    }
1534}
terminus_store/storage/archive.rs

terminus_store/storage/
archive.rs
