Skip to main content

triblespace_core/repo/
objectstore.rs

1use std::array::TryFromSliceError;
2use std::convert::Infallible;
3use std::convert::TryInto;
4use std::error::Error;
5use std::fmt;
6use std::sync::Arc;
7
8use anybytes::Bytes;
9use crossbeam_channel::{bounded, Receiver};
10use futures::Stream;
11use futures::StreamExt;
12use tokio::runtime::Runtime;
13
14use object_store::parse_url;
15use object_store::path::Path;
16use object_store::ObjectStore;
17use object_store::PutMode;
18use object_store::UpdateVersion;
19use object_store::{self};
20use url::Url;
21
22use hex::FromHex;
23
24use crate::blob::encodings::UnknownBlob;
25use crate::blob::Blob;
26use crate::blob::BlobEncoding;
27use crate::blob::IntoBlob;
28use crate::blob::TryFromBlob;
29use crate::id::Id;
30use crate::id::RawId;
31use crate::prelude::blobencodings::SimpleArchive;
32use crate::inline::encodings::hash::Handle;
33use crate::inline::RawInline;
34use crate::inline::Inline;
35use crate::inline::InlineEncoding;
36
37use super::BlobStore;
38use super::BlobStoreGet;
39use super::BlobStoreList;
40use super::BlobStorePut;
41use super::PinStore;
42use super::PushResult;
43
44const BRANCH_INFIX: &str = "branches";
45const BLOB_INFIX: &str = "blobs";
46
47/// Repository backed by an [`object_store`] compatible storage backend.
48///
49/// All data is stored in an external service (e.g. S3, local filesystem) via
50/// the `object_store` crate.
51pub struct ObjectStoreRemote {
52    store: Arc<dyn ObjectStore>,
53    prefix: Path,
54    rt: Arc<Runtime>,}
55
56impl fmt::Debug for ObjectStoreRemote {
57    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
58        f.debug_struct("ObjectStoreRemote")
59            .field("prefix", &self.prefix)
60            .finish()
61    }
62}
63
64impl fmt::Debug for ObjectStoreReader {
65    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
66        f.debug_struct("ObjectStoreReader")
67            .field("prefix", &self.prefix)
68            .finish()
69    }
70}
71
72/// Read-only handle into an [`ObjectStoreRemote`] that can be cloned and shared.
73#[derive(Clone)]
74pub struct ObjectStoreReader {
75    store: Arc<dyn ObjectStore>,
76    prefix: Path,
77    rt: Arc<Runtime>,}
78
79/// Iterator that bridges an async [`Stream`] into blocking iteration via a bounded channel.
80pub struct BlockingIter<T> {
81    rx: Receiver<T>,
82}
83
84impl<T> BlockingIter<T> {
85    fn from_stream<S>(handle: tokio::runtime::Handle, stream: S, capacity: usize) -> Self
86    where
87        S: Stream<Item = T> + Send + 'static,
88        T: Send + 'static,
89    {
90        let (tx, rx) = bounded(capacity);
91        let handle_for_spawn = handle.clone();
92        let handle_for_task = handle.clone();
93        handle_for_spawn.spawn(async move {
94            let mut s = Box::pin(stream);
95            let rt = handle_for_task;
96            while let Some(item) = s.next().await {
97                let tx_clone = tx.clone();
98                let bh = rt.clone();
99                // send on blocking pool to avoid blocking a runtime worker
100                match bh.spawn_blocking(move || tx_clone.send(item)).await {
101                    Ok(Ok(())) => {}
102                    _ => break,
103                }
104            }
105            // tx dropped here -> closes channel
106        });
107        BlockingIter { rx }
108    }
109}
110
111impl<T> Iterator for BlockingIter<T> {
112    type Item = T;
113    fn next(&mut self) -> Option<Self::Item> {
114        self.rx.recv().ok()
115    }
116}
117
118impl PartialEq for ObjectStoreReader {
119    fn eq(&self, other: &Self) -> bool {
120        Arc::ptr_eq(&self.store, &other.store) && self.prefix == other.prefix
121    }
122}
123
124impl Eq for ObjectStoreReader {}
125
126impl ObjectStoreRemote {
127    /// Creates a repository pointing at the object store described by `url`.
128    pub fn with_url(url: &Url) -> Result<ObjectStoreRemote, object_store::Error> {
129        let (store, path) = parse_url(url)?;
130        Ok(ObjectStoreRemote {
131            store: Arc::from(store),
132            prefix: path,
133            rt: Arc::new(
134                tokio::runtime::Builder::new_multi_thread()
135                    .enable_all()
136                    .worker_threads(2)
137                    .build()
138                    .expect("build runtime"),
139            ),        })
140    }
141}
142
143impl BlobStorePut for ObjectStoreRemote
144{
145
146    type PutError = object_store::Error;
147
148    fn put<S, T>(&mut self, item: T) -> Result<Inline<Handle<S>>, Self::PutError>
149    where
150        S: BlobEncoding + 'static,
151        T: IntoBlob<S>,
152        Handle<S>: InlineEncoding,
153    {
154        let blob = item.to_blob();
155        let handle = blob.get_handle();
156        let path = self.prefix.child(BLOB_INFIX).child(hex::encode(handle.raw));
157        let bytes: bytes::Bytes = blob.bytes.into();
158        let result = self.rt.block_on(async {
159            self.store
160                .put_opts(&path, bytes.into(), PutMode::Create.into())
161                .await
162        });
163        match result {
164            Ok(_) | Err(object_store::Error::AlreadyExists { .. }) => Ok(handle),
165            Err(e) => Err(e),
166        }
167    }
168}
169
170impl BlobStore for ObjectStoreRemote
171{
172
173    type Reader = ObjectStoreReader;
174    type ReaderError = Infallible;
175
176    fn reader(&mut self) -> Result<Self::Reader, Self::ReaderError> {
177        Ok(ObjectStoreReader {
178            store: self.store.clone(),
179            prefix: self.prefix.clone(),
180            rt: self.rt.clone(),        })
181    }
182}
183
184impl PinStore for ObjectStoreRemote
185{
186
187    type PinsError = ListBranchesErr;
188    type HeadError = PullBranchErr;
189    type UpdateError = PushBranchErr;
190
191    type ListIter<'a> = BlockingIter<Result<Id, Self::PinsError>>;
192
193    fn pins<'a>(&'a mut self) -> Result<Self::ListIter<'a>, Self::PinsError> {
194        let prefix = self.prefix.child(BRANCH_INFIX);
195        let stream = self.store.list(Some(&prefix)).filter_map(|r| async move {
196            match r {
197                Ok(meta) if meta.size == 0 => None, // tombstoned branch (0-byte object)
198                Ok(meta) => {
199                    let name = match meta.location.filename() {
200                        Some(name) => name,
201                        None => return Some(Err(ListBranchesErr::NotAFile("no filename"))),
202                    };
203                    let digest = match RawId::from_hex(name) {
204                        Ok(digest) => digest,
205                        Err(e) => return Some(Err(ListBranchesErr::BadNameHex(e))),
206                    };
207                    let Some(id) = Id::new(digest) else {
208                        return Some(Err(ListBranchesErr::BadId));
209                    };
210                    Some(Ok(id))
211                }
212                Err(e) => Some(Err(ListBranchesErr::List(e))),
213            }
214        });
215        Ok(BlockingIter::from_stream(
216            self.rt.handle().clone(),
217            stream,
218            16,
219        ))
220    }
221
222    fn head(&mut self, id: Id) -> Result<Option<Inline<Handle<SimpleArchive>>>, Self::HeadError> {
223        let path = self.prefix.child(BRANCH_INFIX).child(hex::encode(id));
224        let result = self.rt.block_on(async { self.store.get(&path).await });
225        match result {
226            Ok(object) => {
227                let bytes = self.rt.block_on(object.bytes())?;
228                if bytes.is_empty() {
229                    return Ok(None);
230                }
231                let value = (&bytes[..]).try_into()?;
232                Ok(Some(Inline::new(value)))
233            }
234            Err(object_store::Error::NotFound { .. }) => Ok(None),
235            Err(e) => Err(PullBranchErr::StoreErr(e)),
236        }
237    }
238
239    fn update(
240        &mut self,
241        id: Id,
242        old: Option<Inline<Handle<SimpleArchive>>>,
243        new: Option<Inline<Handle<SimpleArchive>>>,
244    ) -> Result<PushResult, Self::UpdateError> {
245        let path = self.prefix.child(BRANCH_INFIX).child(hex::encode(id));
246        // We encode "deleted branch" as an empty object. This lets us preserve
247        // CAS semantics for delete via conditional PUT (PutMode::Update), since
248        // `object_store` does not currently expose conditional delete.
249        //
250        // TODO: Once `object_store` supports conditional delete, migrate away
251        // from 0-byte tombstones and treat empty objects as corruption.
252        let new_bytes = match new {
253            Some(new) => bytes::Bytes::copy_from_slice(&new.raw),
254            None => bytes::Bytes::new(),
255        };
256
257        let parse_branch = |bytes: &bytes::Bytes| -> Result<
258            Option<Inline<Handle<SimpleArchive>>>,
259            TryFromSliceError,
260        > {
261            if bytes.is_empty() {
262                return Ok(None);
263            }
264            let value = (&bytes[..]).try_into()?;
265            Ok(Some(Inline::new(value)))
266        };
267
268        if let Some(old_hash) = old {
269            let mut result = self.rt.block_on(async { self.store.get(&path).await });
270            loop {
271                match result {
272                    Ok(obj) => {
273                        let version = UpdateVersion {
274                            e_tag: obj.meta.e_tag.clone(),
275                            version: obj.meta.version.clone(),
276                        };
277                        let stored_bytes = self.rt.block_on(obj.bytes())?;
278                        let stored_hash = parse_branch(&stored_bytes)?;
279                        if stored_hash != Some(old_hash) {
280                            return Ok(PushResult::Conflict(stored_hash));
281                        }
282                        match self.rt.block_on(async {
283                            self.store
284                                .put_opts(
285                                    &path,
286                                    new_bytes.clone().into(),
287                                    PutMode::Update(version).into(),
288                                )
289                                .await
290                        }) {
291                            Ok(_) => return Ok(PushResult::Success()),
292                            Err(object_store::Error::Precondition { .. }) => {
293                                result = self.rt.block_on(async { self.store.get(&path).await });
294                                continue;
295                            }
296                            Err(e) => return Err(PushBranchErr::StoreErr(e)),
297                        }
298                    }
299                    Err(object_store::Error::NotFound { .. }) => {
300                        return Ok(PushResult::Conflict(None))
301                    }
302                    Err(e) => return Err(PushBranchErr::StoreErr(e)),
303                }
304            }
305        } else {
306            loop {
307                match self.rt.block_on(async {
308                    self.store
309                        .put_opts(&path, new_bytes.clone().into(), PutMode::Create.into())
310                        .await
311                }) {
312                    Ok(_) => return Ok(PushResult::Success()),
313                    Err(object_store::Error::AlreadyExists { .. }) => {
314                        let mut result = self.rt.block_on(async { self.store.get(&path).await });
315                        loop {
316                            match result {
317                                Ok(obj) => {
318                                    let version = UpdateVersion {
319                                        e_tag: obj.meta.e_tag.clone(),
320                                        version: obj.meta.version.clone(),
321                                    };
322                                    let stored_bytes = self.rt.block_on(obj.bytes())?;
323                                    let stored_hash = parse_branch(&stored_bytes)?;
324                                    if stored_hash.is_some() {
325                                        return Ok(PushResult::Conflict(stored_hash));
326                                    }
327                                    match self.rt.block_on(async {
328                                        self.store
329                                            .put_opts(
330                                                &path,
331                                                new_bytes.clone().into(),
332                                                PutMode::Update(version).into(),
333                                            )
334                                            .await
335                                    }) {
336                                        Ok(_) => return Ok(PushResult::Success()),
337                                        Err(object_store::Error::Precondition { .. }) => {
338                                            result = self
339                                                .rt
340                                                .block_on(async { self.store.get(&path).await });
341                                            continue;
342                                        }
343                                        Err(e) => return Err(PushBranchErr::StoreErr(e)),
344                                    }
345                                }
346                                Err(object_store::Error::NotFound { .. }) => break, // raced with delete; retry create
347                                Err(e) => return Err(PushBranchErr::StoreErr(e)),
348                            }
349                        }
350                        continue;
351                    }
352                    Err(e) => return Err(PushBranchErr::StoreErr(e)),
353                }
354            }
355        }
356    }
357}
358
359impl crate::repo::StorageClose for ObjectStoreRemote {
360    type Error = Infallible;
361
362    fn close(self) -> Result<(), Self::Error> {
363        // No explicit close necessary for the remote object store adapter.
364        Ok(())
365    }
366}
367
368impl ObjectStoreReader {
369    fn blob_path(&self, handle_hex: String) -> Path {
370        self.prefix.child(BLOB_INFIX).child(handle_hex)
371    }
372}
373
374impl BlobStoreList for ObjectStoreReader
375{
376
377    type Err = ListBlobsErr;
378    type Iter<'a> = BlockingIter<Result<Inline<Handle<UnknownBlob>>, Self::Err>>;
379
380    fn blobs<'a>(&'a self) -> Self::Iter<'a> {
381        let prefix = self.prefix.child(BLOB_INFIX);
382        let stream = self.store.list(Some(&prefix)).map(|r| match r {
383            Ok(meta) => {
384                let blob_name = meta
385                    .location
386                    .filename()
387                    .ok_or(ListBlobsErr::NotAFile("no filename"))?;
388                let digest = RawInline::from_hex(blob_name).map_err(ListBlobsErr::BadNameHex)?;
389                Ok(Inline::new(digest))
390            }
391            Err(e) => Err(ListBlobsErr::List(e)),
392        });
393        BlockingIter::from_stream(self.rt.handle().clone(), stream, 16)
394    }
395}
396
397/// Error returned when retrieving a blob from the object store.
398#[derive(Debug)]
399pub enum GetBlobErr<E: Error> {
400    /// The underlying object store operation failed.
401    Store(object_store::Error),
402    /// The blob bytes could not be converted to the requested type.
403    Conversion(E),
404}
405
406impl<E: Error> fmt::Display for GetBlobErr<E> {
407    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
408        match self {
409            Self::Store(e) => write!(f, "object store error: {e}"),
410            Self::Conversion(e) => write!(f, "conversion error: {e}"),
411        }
412    }
413}
414
415impl<E: Error> Error for GetBlobErr<E> {
416    fn source(&self) -> Option<&(dyn Error + 'static)> {
417        match self {
418            Self::Store(e) => Some(e),
419            Self::Conversion(_) => None,
420        }
421    }
422}
423
424impl<E: Error> From<object_store::Error> for GetBlobErr<E> {
425    fn from(e: object_store::Error) -> Self {
426        Self::Store(e)
427    }
428}
429
430impl BlobStoreGet for ObjectStoreReader
431{
432
433    type GetError<E: Error + Send + Sync + 'static> = GetBlobErr<E>;
434
435    fn get<T, S>(
436        &self,
437        handle: Inline<Handle<S>>,
438    ) -> Result<T, Self::GetError<<T as TryFromBlob<S>>::Error>>
439    where
440        S: BlobEncoding + 'static,
441        T: TryFromBlob<S>,
442        Handle<S>: InlineEncoding,
443    {
444        let path = self.blob_path(hex::encode(handle.raw));
445        let object = self.rt.block_on(async { self.store.get(&path).await })?;
446        let bytes = self.rt.block_on(object.bytes())?;
447        let bytes: Bytes = bytes.into();
448        let blob: Blob<S> = Blob::new(bytes);
449        blob.try_from_blob().map_err(GetBlobErr::Conversion)
450    }
451}
452
453/// Error returned when listing blobs from the object store.
454#[derive(Debug)]
455pub enum ListBlobsErr {
456    /// The underlying list operation failed.
457    List(object_store::Error),
458    /// A listed object had no filename component.
459    NotAFile(&'static str),
460    /// A listed object's filename was not valid hexadecimal.
461    BadNameHex(<RawInline as FromHex>::Error),
462}
463
464impl fmt::Display for ListBlobsErr {
465    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
466        match self {
467            Self::List(e) => write!(f, "list failed: {e}"),
468            Self::NotAFile(e) => write!(f, "list failed: {e}"),
469            Self::BadNameHex(e) => write!(f, "list failed: {e}"),
470        }
471    }
472}
473impl Error for ListBlobsErr {}
474
475impl super::BlobChildren for ObjectStoreReader {}
476
477/// Error returned when listing branches from the object store.
478#[derive(Debug)]
479pub enum ListBranchesErr {
480    /// The underlying list operation failed.
481    List(object_store::Error),
482    /// A listed object had no filename component.
483    NotAFile(&'static str),
484    /// A listed object's filename was not valid hexadecimal.
485    BadNameHex(<RawId as FromHex>::Error),
486    /// The decoded bytes represent the nil identifier.
487    BadId,
488}
489
490impl fmt::Display for ListBranchesErr {
491    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
492        match self {
493            Self::List(e) => write!(f, "list failed: {e}"),
494            Self::NotAFile(e) => write!(f, "list failed: {e}"),
495            Self::BadNameHex(e) => write!(f, "list failed: {e}"),
496            Self::BadId => write!(f, "list failed: bad id"),
497        }
498    }
499}
500impl Error for ListBranchesErr {}
501
502/// Error returned when reading a branch head from the object store.
503#[derive(Debug)]
504pub enum PullBranchErr {
505    /// The stored bytes could not be parsed as a valid handle.
506    ValidationErr(TryFromSliceError),
507    /// The underlying object store operation failed.
508    StoreErr(object_store::Error),
509}
510
511impl fmt::Display for PullBranchErr {
512    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
513        match self {
514            Self::StoreErr(e) => write!(f, "pull failed: {e}"),
515            Self::ValidationErr(e) => write!(f, "pull failed: {e}"),
516        }
517    }
518}
519
520impl Error for PullBranchErr {}
521
522impl From<object_store::Error> for PullBranchErr {
523    fn from(err: object_store::Error) -> Self {
524        Self::StoreErr(err)
525    }
526}
527
528impl From<TryFromSliceError> for PullBranchErr {
529    fn from(err: TryFromSliceError) -> Self {
530        Self::ValidationErr(err)
531    }
532}
533
534/// Error returned when updating a branch head in the object store.
535#[derive(Debug)]
536pub enum PushBranchErr {
537    /// The stored bytes could not be parsed as a valid handle during a compare-and-swap.
538    ValidationErr(TryFromSliceError),
539    /// The underlying object store operation failed.
540    StoreErr(object_store::Error),
541}
542
543impl fmt::Display for PushBranchErr {
544    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
545        match self {
546            Self::ValidationErr(e) => write!(f, "commit failed: {e}"),
547            Self::StoreErr(e) => write!(f, "commit failed: {e}"),
548        }
549    }
550}
551
552impl Error for PushBranchErr {}
553
554impl From<object_store::Error> for PushBranchErr {
555    fn from(err: object_store::Error) -> Self {
556        Self::StoreErr(err)
557    }
558}
559
560impl From<TryFromSliceError> for PushBranchErr {
561    fn from(err: TryFromSliceError) -> Self {
562        Self::ValidationErr(err)
563    }
564}
565
566impl crate::repo::BlobStoreMeta for ObjectStoreReader
567{
568
569    type MetaError = object_store::Error;
570
571    fn metadata<S>(
572        &self,
573        handle: Inline<Handle<S>>,
574    ) -> Result<Option<crate::repo::BlobMetadata>, Self::MetaError>
575    where
576        S: BlobEncoding + 'static,
577        Handle<S>: InlineEncoding,
578    {
579        let handle_hex = hex::encode(handle.raw);
580        let path = self.prefix.child(BLOB_INFIX).child(handle_hex);
581        match self.rt.block_on(async { self.store.head(&path).await }) {
582            Ok(meta) => {
583                let ts = meta.last_modified.timestamp_millis() as u64;
584                let len = meta.size;
585                Ok(Some(crate::repo::BlobMetadata {
586                    timestamp: ts,
587                    length: len,
588                }))
589            }
590            Err(object_store::Error::NotFound { .. }) => Ok(None),
591            Err(e) => Err(e),
592        }
593    }
594}
595
596impl crate::repo::BlobStoreForget for ObjectStoreRemote
597{
598
599    type ForgetError = object_store::Error;
600
601    fn forget<S>(&mut self, handle: Inline<Handle<S>>) -> Result<(), Self::ForgetError>
602    where
603        S: BlobEncoding + 'static,
604        Handle<S>: InlineEncoding,
605    {
606        let handle_hex = hex::encode(handle.raw);
607        let path = self.prefix.child(BLOB_INFIX).child(handle_hex);
608        match self.rt.block_on(async { self.store.delete(&path).await }) {
609            Ok(_) => Ok(()),
610            Err(object_store::Error::NotFound { .. }) => Ok(()),
611            Err(e) => Err(e),
612        }
613    }
614}