iroh_blobs/api/
blobs.rs

1//! API to interact with a local blob store
2//!
3//! This API is for local interactions with the blob store, such as importing
4//! and exporting blobs, observing the bitfield of a blob, and deleting blobs.
5//!
6//! The main entry point is the [`Blobs`] struct.
7use std::{
8    collections::BTreeMap,
9    future::{Future, IntoFuture},
10    io,
11    num::NonZeroU64,
12    path::{Path, PathBuf},
13    pin::Pin,
14};
15
16pub use bao_tree::io::mixed::EncodedItem;
17use bao_tree::{
18    io::{
19        fsm::{ResponseDecoder, ResponseDecoderNext},
20        BaoContentItem, Leaf,
21    },
22    BaoTree, ChunkNum, ChunkRanges,
23};
24use bytes::Bytes;
25use genawaiter::sync::Gen;
26use iroh_io::AsyncStreamWriter;
27use irpc::channel::{mpsc, oneshot};
28use n0_error::AnyError;
29use n0_future::{future, stream, Stream, StreamExt};
30use range_collections::{range_set::RangeSetRange, RangeSet2};
31use ref_cast::RefCast;
32use serde::{Deserialize, Serialize};
33use tracing::trace;
34mod reader;
35pub use reader::BlobReader;
36
37// Public reexports from the proto module.
38//
39// Due to the fact that the proto module is hidden from docs by default,
40// these will appear in the docs as if they were declared here.
41pub use super::proto::{
42    AddProgressItem, Bitfield, BlobDeleteRequest as DeleteOptions, BlobStatus,
43    ExportBaoRequest as ExportBaoOptions, ExportMode, ExportPathRequest as ExportOptions,
44    ExportProgressItem, ExportRangesRequest as ExportRangesOptions,
45    ImportBaoRequest as ImportBaoOptions, ImportMode, ObserveRequest as ObserveOptions,
46};
47use super::{
48    proto::{
49        BatchResponse, BlobStatusRequest, ClearProtectedRequest, CreateTempTagRequest,
50        ExportBaoRequest, ExportRangesItem, ImportBaoRequest, ImportByteStreamRequest,
51        ImportBytesRequest, ImportPathRequest, ListRequest, Scope,
52    },
53    remote::HashSeqChunk,
54    tags::TagInfo,
55    ApiClient, RequestResult, Tags,
56};
57use crate::{
58    api::proto::{BatchRequest, ImportByteStreamUpdate},
59    provider::events::ClientResult,
60    store::IROH_BLOCK_SIZE,
61    util::{temp_tag::TempTag, RecvStreamAsyncStreamReader},
62    BlobFormat, Hash, HashAndFormat,
63};
64
65/// Options for adding bytes.
66#[derive(Debug)]
67pub struct AddBytesOptions {
68    pub data: Bytes,
69    pub format: BlobFormat,
70}
71
72impl<T: Into<Bytes>> From<(T, BlobFormat)> for AddBytesOptions {
73    fn from(item: (T, BlobFormat)) -> Self {
74        let (data, format) = item;
75        Self {
76            data: data.into(),
77            format,
78        }
79    }
80}
81
82/// Blobs API
83#[derive(Debug, Clone, ref_cast::RefCast)]
84#[repr(transparent)]
85pub struct Blobs {
86    client: ApiClient,
87}
88
89impl Blobs {
90    pub(crate) fn ref_from_sender(sender: &ApiClient) -> &Self {
91        Self::ref_cast(sender)
92    }
93
94    /// Creates a new batch scope.
95    ///
96    /// A batch holds open a set of temporary tags that protect blobs from
97    /// garbage collection for as long as the batch is alive. When the batch is
98    /// dropped, all its temporary tags are released and the protected blobs
99    /// become eligible for collection again.
100    ///
101    /// Use a batch whenever a long-running operation needs to protect a blob
102    /// in the middle of a write, e.g. when downloading a large blob that takes
103    /// longer than the GC interval. Without a temp tag, GC may collect the
104    /// partially written data file before the operation completes.
105    ///
106    /// See [`Batch::temp_tag`] to pin a specific hash inside the batch scope.
107    pub async fn batch(&self) -> irpc::Result<Batch<'_>> {
108        let msg = BatchRequest;
109        trace!("{msg:?}");
110        let (tx, rx) = self.client.client_streaming(msg, 32).await?;
111        let scope = rx.await?;
112
113        Ok(Batch {
114            scope,
115            blobs: self,
116            _tx: tx,
117        })
118    }
119
120    /// Create a reader for the given hash. The reader implements [`tokio::io::AsyncRead`] and [`tokio::io::AsyncSeek`]
121    /// and therefore can be used to read the blob's content.
122    ///
123    /// Any access to parts of the blob that are not present will result in an error.
124    ///
125    /// Example:
126    /// ```rust
127    /// use iroh_blobs::{store::mem::MemStore, api::blobs::Blobs};
128    /// use tokio::io::AsyncReadExt;
129    ///
130    /// # async fn example() -> n0_error::Result<()> {
131    /// let store = MemStore::new();
132    /// let tag = store.add_slice(b"Hello, world!").await?;
133    /// let mut reader = store.reader(tag.hash);
134    /// let mut buf = String::new();
135    /// reader.read_to_string(&mut buf).await?;
136    /// assert_eq!(buf, "Hello, world!");
137    /// # Ok(())
138    /// }
139    /// ```
140    pub fn reader(&self, hash: impl Into<Hash>) -> BlobReader {
141        self.reader_with_opts(ReaderOptions { hash: hash.into() })
142    }
143
144    /// Create a reader for the given options. The reader implements [`tokio::io::AsyncRead`] and [`tokio::io::AsyncSeek`]
145    /// and therefore can be used to read the blob's content.
146    ///
147    /// Any access to parts of the blob that are not present will result in an error.
148    pub fn reader_with_opts(&self, options: ReaderOptions) -> BlobReader {
149        BlobReader::new(self.clone(), options)
150    }
151
152    /// Delete a blob.
153    ///
154    /// This function is not public, because it does not work as expected when called manually,
155    /// because blobs are protected from deletion. This is only called from the gc task, which
156    /// clears the protections before.
157    ///
158    /// Users should rely only on garbage collection for blob deletion.
159    pub(crate) async fn delete_with_opts(&self, options: DeleteOptions) -> RequestResult<()> {
160        trace!("{options:?}");
161        self.client.rpc(options).await??;
162        Ok(())
163    }
164
165    /// See [`Self::delete_with_opts`].
166    pub(crate) async fn delete(
167        &self,
168        hashes: impl IntoIterator<Item = impl Into<Hash>>,
169    ) -> RequestResult<()> {
170        self.delete_with_opts(DeleteOptions {
171            hashes: hashes.into_iter().map(Into::into).collect(),
172            force: false,
173        })
174        .await
175    }
176
177    pub fn add_slice(&self, data: impl AsRef<[u8]>) -> AddProgress<'_> {
178        let options = ImportBytesRequest {
179            data: Bytes::copy_from_slice(data.as_ref()),
180            format: crate::BlobFormat::Raw,
181            scope: Scope::GLOBAL,
182        };
183        self.add_bytes_impl(options)
184    }
185
186    pub fn add_bytes(&self, data: impl Into<bytes::Bytes>) -> AddProgress<'_> {
187        let options = ImportBytesRequest {
188            data: data.into(),
189            format: crate::BlobFormat::Raw,
190            scope: Scope::GLOBAL,
191        };
192        self.add_bytes_impl(options)
193    }
194
195    pub fn add_bytes_with_opts(&self, options: impl Into<AddBytesOptions>) -> AddProgress<'_> {
196        let options = options.into();
197        let request = ImportBytesRequest {
198            data: options.data,
199            format: options.format,
200            scope: Scope::GLOBAL,
201        };
202        self.add_bytes_impl(request)
203    }
204
205    fn add_bytes_impl(&self, options: ImportBytesRequest) -> AddProgress<'_> {
206        trace!("{options:?}");
207        let this = self.clone();
208        let stream = Gen::new(|co| async move {
209            let mut receiver = match this.client.server_streaming(options, 32).await {
210                Ok(receiver) => receiver,
211                Err(cause) => {
212                    co.yield_(AddProgressItem::Error(cause.into())).await;
213                    return;
214                }
215            };
216            loop {
217                match receiver.recv().await {
218                    Ok(Some(item)) => co.yield_(item).await,
219                    Err(cause) => {
220                        co.yield_(AddProgressItem::Error(cause.into())).await;
221                        break;
222                    }
223                    Ok(None) => break,
224                }
225            }
226        });
227        AddProgress::new(self, stream)
228    }
229
230    pub fn add_path_with_opts(&self, options: impl Into<AddPathOptions>) -> AddProgress<'_> {
231        let options = options.into();
232        self.add_path_with_opts_impl(ImportPathRequest {
233            path: options.path,
234            mode: options.mode,
235            format: options.format,
236            scope: Scope::GLOBAL,
237        })
238    }
239
240    fn add_path_with_opts_impl(&self, options: ImportPathRequest) -> AddProgress<'_> {
241        trace!("{:?}", options);
242        let client = self.client.clone();
243        let stream = Gen::new(|co| async move {
244            let mut receiver = match client.server_streaming(options, 32).await {
245                Ok(receiver) => receiver,
246                Err(cause) => {
247                    co.yield_(AddProgressItem::Error(cause.into())).await;
248                    return;
249                }
250            };
251            loop {
252                match receiver.recv().await {
253                    Ok(Some(item)) => co.yield_(item).await,
254                    Err(cause) => {
255                        co.yield_(AddProgressItem::Error(cause.into())).await;
256                        break;
257                    }
258                    Ok(None) => break,
259                }
260            }
261        });
262        AddProgress::new(self, stream)
263    }
264
265    pub fn add_path(&self, path: impl AsRef<Path>) -> AddProgress<'_> {
266        self.add_path_with_opts(AddPathOptions {
267            path: path.as_ref().to_owned(),
268            mode: ImportMode::Copy,
269            format: BlobFormat::Raw,
270        })
271    }
272
273    pub async fn add_stream(
274        &self,
275        data: impl Stream<Item = io::Result<Bytes>> + Send + Sync + 'static,
276    ) -> AddProgress<'_> {
277        let inner = ImportByteStreamRequest {
278            format: crate::BlobFormat::Raw,
279            scope: Scope::default(),
280        };
281        let client = self.client.clone();
282        let stream = Gen::new(|co| async move {
283            let (sender, mut receiver) = match client.bidi_streaming(inner, 32, 32).await {
284                Ok(x) => x,
285                Err(cause) => {
286                    co.yield_(AddProgressItem::Error(cause.into())).await;
287                    return;
288                }
289            };
290            let recv = async {
291                loop {
292                    match receiver.recv().await {
293                        Ok(Some(item)) => co.yield_(item).await,
294                        Err(cause) => {
295                            co.yield_(AddProgressItem::Error(cause.into())).await;
296                            break;
297                        }
298                        Ok(None) => break,
299                    }
300                }
301            };
302            let send = async {
303                tokio::pin!(data);
304                while let Some(item) = data.next().await {
305                    sender.send(ImportByteStreamUpdate::Bytes(item?)).await?;
306                }
307                sender.send(ImportByteStreamUpdate::Done).await?;
308                n0_error::Ok(())
309            };
310            let _ = tokio::join!(send, recv);
311        });
312        AddProgress::new(self, stream)
313    }
314
315    pub fn export_ranges(
316        &self,
317        hash: impl Into<Hash>,
318        ranges: impl Into<RangeSet2<u64>>,
319    ) -> ExportRangesProgress {
320        self.export_ranges_with_opts(ExportRangesOptions {
321            hash: hash.into(),
322            ranges: ranges.into(),
323        })
324    }
325
326    pub fn export_ranges_with_opts(&self, options: ExportRangesOptions) -> ExportRangesProgress {
327        trace!("{options:?}");
328        ExportRangesProgress::new(
329            options.ranges.clone(),
330            self.client.server_streaming(options, 32),
331        )
332    }
333
334    pub fn export_bao_with_opts(
335        &self,
336        options: ExportBaoOptions,
337        local_update_cap: usize,
338    ) -> ExportBaoProgress {
339        trace!("{options:?}");
340        ExportBaoProgress::new(self.client.server_streaming(options, local_update_cap))
341    }
342
343    pub fn export_bao(
344        &self,
345        hash: impl Into<Hash>,
346        ranges: impl Into<ChunkRanges>,
347    ) -> ExportBaoProgress {
348        self.export_bao_with_opts(
349            ExportBaoRequest {
350                hash: hash.into(),
351                ranges: ranges.into(),
352            },
353            32,
354        )
355    }
356
357    /// Export a single chunk from the given hash, at the given offset.
358    pub async fn export_chunk(
359        &self,
360        hash: impl Into<Hash>,
361        offset: u64,
362    ) -> super::ExportBaoResult<Leaf> {
363        let base = ChunkNum::full_chunks(offset);
364        let ranges = ChunkRanges::from(base..base + 1);
365        let mut stream = self.export_bao(hash, ranges).stream();
366        while let Some(item) = stream.next().await {
367            match item {
368                EncodedItem::Leaf(leaf) => return Ok(leaf),
369                EncodedItem::Parent(_) => {}
370                EncodedItem::Size(_) => {}
371                EncodedItem::Done => break,
372                EncodedItem::Error(cause) => return Err(cause.into()),
373            }
374        }
375        Err(io::Error::other("unexpected end of stream").into())
376    }
377
378    /// Get the entire blob into a Bytes
379    ///
380    /// This will run out of memory when called for very large blobs, so be careful!
381    pub async fn get_bytes(&self, hash: impl Into<Hash>) -> super::ExportBaoResult<Bytes> {
382        self.export_bao(hash.into(), ChunkRanges::all())
383            .data_to_bytes()
384            .await
385    }
386
387    /// Observe the bitfield of the given hash.
388    pub fn observe(&self, hash: impl Into<Hash>) -> ObserveProgress {
389        self.observe_with_opts(ObserveOptions { hash: hash.into() })
390    }
391
392    pub fn observe_with_opts(&self, options: ObserveOptions) -> ObserveProgress {
393        trace!("{:?}", options);
394        if options.hash == Hash::EMPTY {
395            return ObserveProgress::new(async move {
396                let (tx, rx) = mpsc::channel(1);
397                tx.send(Bitfield::complete(0)).await.ok();
398                Ok(rx)
399            });
400        }
401        ObserveProgress::new(self.client.server_streaming(options, 32))
402    }
403
404    pub fn export_with_opts(&self, options: ExportOptions) -> ExportProgress {
405        trace!("{:?}", options);
406        ExportProgress::new(self.client.server_streaming(options, 32))
407    }
408
409    pub fn export(&self, hash: impl Into<Hash>, target: impl AsRef<Path>) -> ExportProgress {
410        let options = ExportOptions {
411            hash: hash.into(),
412            mode: ExportMode::Copy,
413            target: target.as_ref().to_owned(),
414        };
415        self.export_with_opts(options)
416    }
417
418    /// Import BaoContentItems from a stream.
419    ///
420    /// The store assumes that these are already verified and in the correct order.
421    #[cfg_attr(feature = "hide-proto-docs", doc(hidden))]
422    pub async fn import_bao(
423        &self,
424        hash: impl Into<Hash>,
425        size: NonZeroU64,
426        local_update_cap: usize,
427    ) -> irpc::Result<ImportBaoHandle> {
428        let options = ImportBaoRequest {
429            hash: hash.into(),
430            size,
431        };
432        self.import_bao_with_opts(options, local_update_cap).await
433    }
434
435    #[cfg_attr(feature = "hide-proto-docs", doc(hidden))]
436    pub async fn import_bao_with_opts(
437        &self,
438        options: ImportBaoOptions,
439        local_update_cap: usize,
440    ) -> irpc::Result<ImportBaoHandle> {
441        trace!("{:?}", options);
442        ImportBaoHandle::new(self.client.client_streaming(options, local_update_cap)).await
443    }
444
445    #[cfg_attr(feature = "hide-proto-docs", doc(hidden))]
446    pub async fn import_bao_reader<R: crate::util::RecvStream>(
447        &self,
448        hash: Hash,
449        ranges: ChunkRanges,
450        mut reader: R,
451    ) -> RequestResult<R> {
452        let mut size = [0; 8];
453        reader.recv_exact(&mut size).await?;
454        let size = u64::from_le_bytes(size);
455        let Some(size) = NonZeroU64::new(size) else {
456            return if hash == Hash::EMPTY {
457                Ok(reader)
458            } else {
459                Err(io::Error::other("invalid size for hash").into())
460            };
461        };
462        let tree = BaoTree::new(size.get(), IROH_BLOCK_SIZE);
463        let mut decoder = ResponseDecoder::new(
464            hash.into(),
465            ranges,
466            tree,
467            RecvStreamAsyncStreamReader::new(reader),
468        );
469        let options = ImportBaoOptions { hash, size };
470        let handle = self.import_bao_with_opts(options, 32).await?;
471        let driver = async move {
472            let reader = loop {
473                match decoder.next().await {
474                    ResponseDecoderNext::More((rest, item)) => {
475                        handle.tx.send(item?).await?;
476                        decoder = rest;
477                    }
478                    ResponseDecoderNext::Done(reader) => break reader,
479                };
480            };
481            drop(handle.tx);
482            io::Result::Ok(reader)
483        };
484        let fut = async move { handle.rx.await.map_err(io::Error::other)? };
485        let (reader, res) = tokio::join!(driver, fut);
486        res?;
487        Ok(reader?.into_inner())
488    }
489
490    #[cfg_attr(feature = "hide-proto-docs", doc(hidden))]
491    pub async fn import_bao_bytes(
492        &self,
493        hash: Hash,
494        ranges: ChunkRanges,
495        data: impl Into<Bytes>,
496    ) -> RequestResult<()> {
497        self.import_bao_reader(hash, ranges, data.into()).await?;
498        Ok(())
499    }
500
501    pub fn list(&self) -> BlobsListProgress {
502        let msg = ListRequest;
503        let client = self.client.clone();
504        BlobsListProgress::new(client.server_streaming(msg, 32))
505    }
506
507    pub async fn status(&self, hash: impl Into<Hash>) -> irpc::Result<BlobStatus> {
508        let hash = hash.into();
509        // the empty blob is always complete, regardless of backend state
510        if hash == Hash::EMPTY {
511            return Ok(BlobStatus::Complete { size: 0 });
512        }
513        let msg = BlobStatusRequest { hash };
514        self.client.rpc(msg).await
515    }
516
517    pub async fn has(&self, hash: impl Into<Hash>) -> irpc::Result<bool> {
518        match self.status(hash).await? {
519            BlobStatus::Complete { .. } => Ok(true),
520            _ => Ok(false),
521        }
522    }
523
524    #[allow(dead_code)]
525    pub(crate) async fn clear_protected(&self) -> RequestResult<()> {
526        let msg = ClearProtectedRequest;
527        self.client.rpc(msg).await??;
528        Ok(())
529    }
530}
531
532/// A progress handle for a batch scoped add operation.
533pub struct BatchAddProgress<'a>(AddProgress<'a>);
534
535impl<'a> IntoFuture for BatchAddProgress<'a> {
536    type Output = RequestResult<TempTag>;
537
538    type IntoFuture = Pin<Box<dyn Future<Output = Self::Output> + Send + 'a>>;
539
540    fn into_future(self) -> Self::IntoFuture {
541        Box::pin(self.temp_tag())
542    }
543}
544
545impl<'a> BatchAddProgress<'a> {
546    pub async fn with_named_tag(self, name: impl AsRef<[u8]>) -> RequestResult<HashAndFormat> {
547        self.0.with_named_tag(name).await
548    }
549
550    pub async fn with_tag(self) -> RequestResult<TagInfo> {
551        self.0.with_tag().await
552    }
553
554    pub async fn stream(self) -> impl Stream<Item = AddProgressItem> {
555        self.0.stream().await
556    }
557
558    pub async fn temp_tag(self) -> RequestResult<TempTag> {
559        self.0.temp_tag().await
560    }
561}
562
563/// A scoped lifetime that protects blobs from garbage collection.
564///
565/// All temporary tags created through a `Batch` are kept alive until the batch
566/// is dropped, at which point they are released atomically and the protected
567/// blobs become eligible for GC again.
568///
569/// A typical use case is protecting a blob during a long-running write
570/// (or download) where GC could otherwise collect the partially-written data.
571///
572/// # Example
573///
574/// ```rust
575/// use iroh_blobs::{store::mem::MemStore, Hash};
576///
577/// # async fn example() -> n0_error::Result<()> {
578/// let store = MemStore::new();
579///
580/// // Hash is known upfront (e.g. from a download ticket).
581/// let hash = Hash::new(b"example content");
582///
583/// let batch = store.blobs().batch().await?;
584/// let _guard = batch.temp_tag(hash).await?;
585/// // GC will not collect `hash` for as long as `_guard` and `batch` are alive.
586///
587/// // ... perform long-running write or network download here ...
588///
589/// drop(_guard);
590/// // Any un-upgraded blobs from the batch are now eligible for garbage collection
591/// # Ok(())
592/// # }
593/// ```
594pub struct Batch<'a> {
595    scope: Scope,
596    blobs: &'a Blobs,
597    _tx: mpsc::Sender<BatchResponse>,
598}
599
600impl<'a> Batch<'a> {
601    pub fn add_bytes(&self, data: impl Into<Bytes>) -> BatchAddProgress<'_> {
602        let options = ImportBytesRequest {
603            data: data.into(),
604            format: crate::BlobFormat::Raw,
605            scope: self.scope,
606        };
607        BatchAddProgress(self.blobs.add_bytes_impl(options))
608    }
609
610    pub fn add_bytes_with_opts(&self, options: impl Into<AddBytesOptions>) -> BatchAddProgress<'_> {
611        let options = options.into();
612        BatchAddProgress(self.blobs.add_bytes_impl(ImportBytesRequest {
613            data: options.data,
614            format: options.format,
615            scope: self.scope,
616        }))
617    }
618
619    pub fn add_slice(&self, data: impl AsRef<[u8]>) -> BatchAddProgress<'_> {
620        let options = ImportBytesRequest {
621            data: Bytes::copy_from_slice(data.as_ref()),
622            format: crate::BlobFormat::Raw,
623            scope: self.scope,
624        };
625        BatchAddProgress(self.blobs.add_bytes_impl(options))
626    }
627
628    pub fn add_path_with_opts(&self, options: impl Into<AddPathOptions>) -> BatchAddProgress<'_> {
629        let options = options.into();
630        BatchAddProgress(self.blobs.add_path_with_opts_impl(ImportPathRequest {
631            path: options.path,
632            mode: options.mode,
633            format: options.format,
634            scope: self.scope,
635        }))
636    }
637
638    /// Pins a hash inside this batch scope, protecting it from garbage collection.
639    ///
640    /// The returned [`TempTag`] keeps the blob alive until either the tag or
641    /// the batch is dropped. This is the recommended way to protect a blob
642    /// during a long-running write or download where GC might otherwise collect
643    /// the partially-written data.
644    pub async fn temp_tag(&self, value: impl Into<HashAndFormat>) -> irpc::Result<TempTag> {
645        let value = value.into();
646        let msg = CreateTempTagRequest {
647            scope: self.scope,
648            value,
649        };
650        self.blobs.client.rpc(msg).await
651    }
652}
653
654/// Options for adding data from a file system path.
655#[derive(Debug)]
656pub struct AddPathOptions {
657    pub path: PathBuf,
658    pub format: BlobFormat,
659    pub mode: ImportMode,
660}
661
662/// A progress handle for an import operation.
663///
664/// Internally this is a stream of [`AddProgressItem`] items. Working with this
665/// stream directly can be inconvenient, so this struct provides some convenience
666/// methods to work with the result.
667///
668/// It also implements [`IntoFuture`], so you can await it to get the [`TagInfo`] that
669/// contains the hash of the added content and also protects the content.
670///
671/// If you want access to the stream, you can use the [`AddProgress::stream`] method.
672pub struct AddProgress<'a> {
673    blobs: &'a Blobs,
674    inner: stream::Boxed<AddProgressItem>,
675}
676
677impl<'a> IntoFuture for AddProgress<'a> {
678    type Output = RequestResult<TagInfo>;
679
680    type IntoFuture = Pin<Box<dyn Future<Output = Self::Output> + Send + 'a>>;
681
682    fn into_future(self) -> Self::IntoFuture {
683        Box::pin(self.with_tag())
684    }
685}
686
687impl<'a> AddProgress<'a> {
688    fn new(blobs: &'a Blobs, stream: impl Stream<Item = AddProgressItem> + Send + 'static) -> Self {
689        Self {
690            blobs,
691            inner: Box::pin(stream),
692        }
693    }
694
695    pub async fn temp_tag(self) -> RequestResult<TempTag> {
696        let mut stream = self.inner;
697        while let Some(item) = stream.next().await {
698            match item {
699                AddProgressItem::Done(tt) => return Ok(tt),
700                AddProgressItem::Error(e) => return Err(e.into()),
701                _ => {}
702            }
703        }
704        Err(io::Error::other("unexpected end of stream").into())
705    }
706
707    pub async fn with_named_tag(self, name: impl AsRef<[u8]>) -> RequestResult<HashAndFormat> {
708        let blobs = self.blobs.clone();
709        let tt = self.temp_tag().await?;
710        let haf = tt.hash_and_format();
711        let tags = Tags::ref_from_sender(&blobs.client);
712        tags.set(name, haf).await?;
713        drop(tt);
714        Ok(haf)
715    }
716
717    pub async fn with_tag(self) -> RequestResult<TagInfo> {
718        let blobs = self.blobs.clone();
719        let tt = self.temp_tag().await?;
720        let hash = tt.hash();
721        let format = tt.format();
722        let tags = Tags::ref_from_sender(&blobs.client);
723        let name = tags.create(tt.hash_and_format()).await?;
724        drop(tt);
725        Ok(TagInfo { name, hash, format })
726    }
727
728    pub async fn stream(self) -> impl Stream<Item = AddProgressItem> {
729        self.inner
730    }
731}
732
733/// Options for an async reader for blobs that supports AsyncRead and AsyncSeek.
734#[derive(Debug, Clone, Serialize, Deserialize)]
735pub struct ReaderOptions {
736    pub hash: Hash,
737}
738
739/// An observe result. Awaiting this will return the current state.
740///
741/// Calling [`ObserveProgress::stream`] will return a stream of updates, where
742/// the first item is the current state and subsequent items are updates.
743pub struct ObserveProgress {
744    inner: future::Boxed<irpc::Result<mpsc::Receiver<Bitfield>>>,
745}
746
747impl IntoFuture for ObserveProgress {
748    type Output = RequestResult<Bitfield>;
749
750    type IntoFuture = Pin<Box<dyn Future<Output = Self::Output> + Send>>;
751
752    fn into_future(self) -> Self::IntoFuture {
753        Box::pin(async move {
754            let mut rx = self.inner.await?;
755            match rx.recv().await? {
756                Some(bitfield) => Ok(bitfield),
757                None => Err(io::Error::other("unexpected end of stream").into()),
758            }
759        })
760    }
761}
762
763impl ObserveProgress {
764    fn new(
765        fut: impl Future<Output = irpc::Result<mpsc::Receiver<Bitfield>>> + Send + 'static,
766    ) -> Self {
767        Self {
768            inner: Box::pin(fut),
769        }
770    }
771
772    pub async fn await_completion(self) -> RequestResult<Bitfield> {
773        let mut stream = self.stream().await?;
774        while let Some(item) = stream.next().await {
775            if item.is_complete() {
776                return Ok(item);
777            }
778        }
779        Err(io::Error::other("unexpected end of stream").into())
780    }
781
782    /// Returns an infinite stream of bitfields. The first bitfield is the
783    /// current state, and the following bitfields are updates.
784    ///
785    /// Once a blob is complete, there will be no more updates.
786    pub async fn stream(self) -> irpc::Result<impl Stream<Item = Bitfield>> {
787        let mut rx = self.inner.await?;
788        Ok(Gen::new(|co| async move {
789            while let Ok(Some(item)) = rx.recv().await {
790                co.yield_(item).await;
791            }
792        }))
793    }
794}
795
796/// A progress handle for an export operation.
797///
798/// Internally this is a stream of [`ExportProgress`] items. Working with this
799/// stream directly can be inconvenient, so this struct provides some convenience
800/// methods to work with the result.
801///
802/// To get the underlying stream, use the [`ExportProgress::stream`] method.
803///
804/// It also implements [`IntoFuture`], so you can await it to get the size of the
805/// exported blob.
806pub struct ExportProgress {
807    inner: future::Boxed<irpc::Result<mpsc::Receiver<ExportProgressItem>>>,
808}
809
810impl IntoFuture for ExportProgress {
811    type Output = RequestResult<u64>;
812
813    type IntoFuture = Pin<Box<dyn Future<Output = Self::Output> + Send>>;
814
815    fn into_future(self) -> Self::IntoFuture {
816        Box::pin(self.finish())
817    }
818}
819
820impl ExportProgress {
821    fn new(
822        fut: impl Future<Output = irpc::Result<mpsc::Receiver<ExportProgressItem>>> + Send + 'static,
823    ) -> Self {
824        Self {
825            inner: Box::pin(fut),
826        }
827    }
828
829    pub async fn stream(self) -> impl Stream<Item = ExportProgressItem> {
830        Gen::new(|co| async move {
831            let mut rx = match self.inner.await {
832                Ok(rx) => rx,
833                Err(e) => {
834                    co.yield_(ExportProgressItem::Error(e.into())).await;
835                    return;
836                }
837            };
838            while let Ok(Some(item)) = rx.recv().await {
839                co.yield_(item).await;
840            }
841        })
842    }
843
844    pub async fn finish(self) -> RequestResult<u64> {
845        let mut rx = self.inner.await?;
846        let mut size = None;
847        loop {
848            match rx.recv().await? {
849                Some(ExportProgressItem::Done) => break,
850                Some(ExportProgressItem::Size(s)) => size = Some(s),
851                Some(ExportProgressItem::Error(cause)) => return Err(cause.into()),
852                _ => {}
853            }
854        }
855        if let Some(size) = size {
856            Ok(size)
857        } else {
858            Err(io::Error::other("unexpected end of stream").into())
859        }
860    }
861}
862
863/// A handle for an ongoing bao import operation.
864pub struct ImportBaoHandle {
865    pub tx: mpsc::Sender<BaoContentItem>,
866    pub rx: oneshot::Receiver<super::Result<()>>,
867}
868
869impl ImportBaoHandle {
870    pub(crate) async fn new(
871        fut: impl Future<
872                Output = irpc::Result<(
873                    mpsc::Sender<BaoContentItem>,
874                    oneshot::Receiver<super::Result<()>>,
875                )>,
876            > + Send
877            + 'static,
878    ) -> irpc::Result<Self> {
879        let (tx, rx) = fut.await?;
880        Ok(Self { tx, rx })
881    }
882}
883
884/// A progress handle for a blobs list operation.
885pub struct BlobsListProgress {
886    inner: future::Boxed<irpc::Result<mpsc::Receiver<super::Result<Hash>>>>,
887}
888
889impl BlobsListProgress {
890    fn new(
891        fut: impl Future<Output = irpc::Result<mpsc::Receiver<super::Result<Hash>>>> + Send + 'static,
892    ) -> Self {
893        Self {
894            inner: Box::pin(fut),
895        }
896    }
897
898    pub async fn hashes(self) -> RequestResult<Vec<Hash>> {
899        let mut rx: mpsc::Receiver<Result<Hash, super::Error>> = self.inner.await?;
900        let mut hashes = Vec::new();
901        while let Some(item) = rx.recv().await? {
902            hashes.push(item?);
903        }
904        Ok(hashes)
905    }
906
907    pub async fn stream(self) -> irpc::Result<impl Stream<Item = super::Result<Hash>>> {
908        let mut rx = self.inner.await?;
909        Ok(Gen::new(|co| async move {
910            while let Ok(Some(item)) = rx.recv().await {
911                co.yield_(item).await;
912            }
913        }))
914    }
915}
916
917/// A progress handle for a bao export operation.
918///
919/// Internally, this is a stream of [`EncodedItem`]s. Using this stream directly
920/// is often inconvenient, so there are a number of higher level methods to
921/// process the stream.
922///
923/// You can get access to the underlying stream using the [`ExportBaoProgress::stream`] method.
924pub struct ExportRangesProgress {
925    ranges: RangeSet2<u64>,
926    inner: future::Boxed<irpc::Result<mpsc::Receiver<ExportRangesItem>>>,
927}
928
929impl ExportRangesProgress {
930    fn new(
931        ranges: RangeSet2<u64>,
932        fut: impl Future<Output = irpc::Result<mpsc::Receiver<ExportRangesItem>>> + Send + 'static,
933    ) -> Self {
934        Self {
935            ranges,
936            inner: Box::pin(fut),
937        }
938    }
939}
940
941impl ExportRangesProgress {
942    /// A raw stream of [`ExportRangesItem`]s.
943    ///
944    /// Ranges will be rounded up to chunk boundaries. So if you request a
945    /// range of 0..100, you will get the entire first chunk, 0..1024.
946    ///
947    /// It is up to the caller to clip the ranges to the requested ranges.
948    pub fn stream(self) -> impl Stream<Item = ExportRangesItem> {
949        Gen::new(|co| async move {
950            let mut rx = match self.inner.await {
951                Ok(rx) => rx,
952                Err(e) => {
953                    co.yield_(ExportRangesItem::Error(e.into())).await;
954                    return;
955                }
956            };
957            while let Ok(Some(item)) = rx.recv().await {
958                co.yield_(item).await;
959            }
960        })
961    }
962
963    /// Concatenate all the data into a single `Bytes`.
964    pub async fn concatenate(self) -> RequestResult<Vec<u8>> {
965        let mut rx = self.inner.await?;
966        let mut data = BTreeMap::new();
967        while let Some(item) = rx.recv().await? {
968            match item {
969                ExportRangesItem::Size(_) => {}
970                ExportRangesItem::Data(leaf) => {
971                    data.insert(leaf.offset, leaf.data);
972                }
973                ExportRangesItem::Error(cause) => return Err(cause.into()),
974            }
975        }
976        let mut res = Vec::new();
977        for range in self.ranges.iter() {
978            let (start, end) = match range {
979                RangeSetRange::RangeFrom(range) => (*range.start, u64::MAX),
980                RangeSetRange::Range(range) => (*range.start, *range.end),
981            };
982            for (offset, data) in data.iter() {
983                let cstart = *offset;
984                let cend = *offset + (data.len() as u64);
985                if cstart >= end || cend <= start {
986                    continue;
987                }
988                let start = start.max(cstart);
989                let end = end.min(cend);
990                let data = &data[(start - cstart) as usize..(end - cstart) as usize];
991                res.extend_from_slice(data);
992            }
993        }
994        Ok(res)
995    }
996}
997
998/// A progress handle for a bao export operation.
999///
1000/// Internally, this is a stream of [`EncodedItem`]s. Using this stream directly
1001/// is often inconvenient, so there are a number of higher level methods to
1002/// process the stream.
1003///
1004/// You can get access to the underlying stream using the [`ExportBaoProgress::stream`] method.
1005pub struct ExportBaoProgress {
1006    inner: future::Boxed<irpc::Result<mpsc::Receiver<EncodedItem>>>,
1007}
1008
1009impl ExportBaoProgress {
1010    fn new(
1011        fut: impl Future<Output = irpc::Result<mpsc::Receiver<EncodedItem>>> + Send + 'static,
1012    ) -> Self {
1013        Self {
1014            inner: Box::pin(fut),
1015        }
1016    }
1017
1018    /// Interprets this blob as a hash sequence and returns a stream of hashes.
1019    ///
1020    /// Errors will be reported, but the iterator will nevertheless continue.
1021    /// If you get an error despite having asked for ranges that should be present,
1022    /// this means that the data is corrupted. It can still make sense to continue
1023    /// to get all non-corrupted sections.
1024    pub fn hashes_with_index(
1025        self,
1026    ) -> impl Stream<Item = std::result::Result<(u64, Hash), AnyError>> {
1027        let mut stream = self.stream();
1028        Gen::new(|co| async move {
1029            while let Some(item) = stream.next().await {
1030                let leaf = match item {
1031                    EncodedItem::Leaf(leaf) => leaf,
1032                    EncodedItem::Error(e) => {
1033                        co.yield_(Err(AnyError::from_std(e))).await;
1034                        continue;
1035                    }
1036                    _ => continue,
1037                };
1038                let slice = match HashSeqChunk::try_from(leaf) {
1039                    Ok(slice) => slice,
1040                    Err(e) => {
1041                        co.yield_(Err(e)).await;
1042                        continue;
1043                    }
1044                };
1045                let offset = slice.base();
1046                for (o, hash) in slice.into_iter().enumerate() {
1047                    co.yield_(Ok((offset + o as u64, hash))).await;
1048                }
1049            }
1050        })
1051    }
1052
1053    /// Same as [`Self::hashes_with_index`], but without the indexes.
1054    pub fn hashes(self) -> impl Stream<Item = std::result::Result<Hash, AnyError>> {
1055        self.hashes_with_index().map(|x| x.map(|(_, hash)| hash))
1056    }
1057
1058    pub async fn bao_to_vec(self) -> RequestResult<Vec<u8>> {
1059        let mut data = Vec::new();
1060        let mut stream = self.into_byte_stream();
1061        while let Some(item) = stream.next().await {
1062            data.extend_from_slice(&item?);
1063        }
1064        Ok(data)
1065    }
1066
1067    pub async fn data_to_bytes(self) -> super::ExportBaoResult<Bytes> {
1068        let mut rx = self.inner.await?;
1069        let mut data = Vec::new();
1070        while let Some(item) = rx.recv().await? {
1071            match item {
1072                EncodedItem::Leaf(leaf) => {
1073                    data.push(leaf.data);
1074                }
1075                EncodedItem::Parent(_) => {}
1076                EncodedItem::Size(_) => {}
1077                EncodedItem::Done => break,
1078                EncodedItem::Error(cause) => return Err(cause.into()),
1079            }
1080        }
1081        if data.len() == 1 {
1082            Ok(data.pop().unwrap())
1083        } else {
1084            let mut out = Vec::new();
1085            for item in data {
1086                out.extend_from_slice(&item);
1087            }
1088            Ok(out.into())
1089        }
1090    }
1091
1092    pub async fn data_to_vec(self) -> super::ExportBaoResult<Vec<u8>> {
1093        let mut rx = self.inner.await?;
1094        let mut data = Vec::new();
1095        while let Some(item) = rx.recv().await? {
1096            match item {
1097                EncodedItem::Leaf(leaf) => {
1098                    data.extend_from_slice(&leaf.data);
1099                }
1100                EncodedItem::Parent(_) => {}
1101                EncodedItem::Size(_) => {}
1102                EncodedItem::Done => break,
1103                EncodedItem::Error(cause) => return Err(cause.into()),
1104            }
1105        }
1106        Ok(data)
1107    }
1108
1109    pub async fn write<W: AsyncStreamWriter>(self, target: &mut W) -> super::ExportBaoResult<()> {
1110        let mut rx = self.inner.await?;
1111        while let Some(item) = rx.recv().await? {
1112            match item {
1113                EncodedItem::Size(size) => {
1114                    target.write(&size.to_le_bytes()).await?;
1115                }
1116                EncodedItem::Parent(parent) => {
1117                    let mut data = vec![0u8; 64];
1118                    data[..32].copy_from_slice(parent.pair.0.as_bytes());
1119                    data[32..].copy_from_slice(parent.pair.1.as_bytes());
1120                    target.write(&data).await?;
1121                }
1122                EncodedItem::Leaf(leaf) => {
1123                    target.write_bytes(leaf.data).await?;
1124                }
1125                EncodedItem::Done => break,
1126                EncodedItem::Error(cause) => return Err(cause.into()),
1127            }
1128        }
1129        Ok(())
1130    }
1131
1132    /// Write noq variant that also feeds a progress writer.
1133    pub(crate) async fn write_with_progress<W: crate::util::SendStream>(
1134        self,
1135        writer: &mut W,
1136        progress: &mut impl WriteProgress,
1137        hash: &Hash,
1138        index: u64,
1139    ) -> super::ExportBaoResult<()> {
1140        let mut rx = self.inner.await?;
1141        while let Some(item) = rx.recv().await? {
1142            match item {
1143                EncodedItem::Size(size) => {
1144                    progress.send_transfer_started(index, hash, size).await;
1145                    writer.send(&size.to_le_bytes()).await?;
1146                    progress.log_other_write(8);
1147                }
1148                EncodedItem::Parent(parent) => {
1149                    let mut data = [0u8; 64];
1150                    data[..32].copy_from_slice(parent.pair.0.as_bytes());
1151                    data[32..].copy_from_slice(parent.pair.1.as_bytes());
1152                    writer.send(&data).await?;
1153                    progress.log_other_write(64);
1154                }
1155                EncodedItem::Leaf(leaf) => {
1156                    let len = leaf.data.len();
1157                    writer.send_bytes(leaf.data).await?;
1158                    progress
1159                        .notify_payload_write(index, leaf.offset, len)
1160                        .await?;
1161                }
1162                EncodedItem::Done => break,
1163                EncodedItem::Error(cause) => return Err(cause.into()),
1164            }
1165        }
1166        Ok(())
1167    }
1168
1169    pub fn into_byte_stream(self) -> impl Stream<Item = super::Result<Bytes>> {
1170        self.stream().filter_map(|item| match item {
1171            EncodedItem::Size(size) => {
1172                let size = size.to_le_bytes().to_vec().into();
1173                Some(Ok(size))
1174            }
1175            EncodedItem::Parent(parent) => {
1176                let mut data = vec![0u8; 64];
1177                data[..32].copy_from_slice(parent.pair.0.as_bytes());
1178                data[32..].copy_from_slice(parent.pair.1.as_bytes());
1179                Some(Ok(data.into()))
1180            }
1181            EncodedItem::Leaf(leaf) => Some(Ok(leaf.data)),
1182            EncodedItem::Done => None,
1183            EncodedItem::Error(cause) => Some(Err(cause.into())),
1184        })
1185    }
1186
1187    pub fn stream(self) -> impl Stream<Item = EncodedItem> {
1188        Gen::new(|co| async move {
1189            let mut rx = match self.inner.await {
1190                Ok(rx) => rx,
1191                Err(cause) => {
1192                    co.yield_(EncodedItem::Error(io::Error::other(cause).into()))
1193                        .await;
1194                    return;
1195                }
1196            };
1197            while let Ok(Some(item)) = rx.recv().await {
1198                co.yield_(item).await;
1199            }
1200        })
1201    }
1202}
1203
1204pub(crate) trait WriteProgress {
1205    /// Notify the progress writer that a payload write has happened.
1206    async fn notify_payload_write(&mut self, index: u64, offset: u64, len: usize) -> ClientResult;
1207
1208    /// Log a write of some other data.
1209    fn log_other_write(&mut self, len: usize);
1210
1211    /// Notify the progress writer that a transfer has started.
1212    async fn send_transfer_started(&mut self, index: u64, hash: &Hash, size: u64);
1213}
iroh_blobs/api/blobs.rs

iroh_blobs/api/
blobs.rs
