iota_storage/object_store/
mod.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::{future::Future, sync::Arc};
6
7use anyhow::{Result, anyhow};
8use async_trait::async_trait;
9use bytes::Bytes;
10use futures::stream::BoxStream;
11use iota_config::object_storage_config::TRANSFER_STALL_TIMEOUT;
12use object_store::{DynObjectStore, ObjectMeta, ObjectStore, ObjectStoreExt, path::Path};
13pub mod http;
14pub mod util;
15
16#[async_trait]
17pub trait ObjectStoreGetExt: std::fmt::Display + Send + Sync + 'static {
18    /// Return the bytes at given path in object store
19    async fn get_bytes(&self, src: &Path) -> Result<Bytes>;
20
21    /// Like [`Self::get_bytes`], additionally invoking `on_bytes` with the
22    /// size of each received chunk while the download is in flight, so
23    /// callers can render live progress. The default implementation reports
24    /// the full size in one call once the download completes; stores that
25    /// stream their responses override it to report per chunk.
26    async fn get_bytes_with_progress(
27        &self,
28        src: &Path,
29        on_bytes: &(dyn Fn(u64) + Send + Sync),
30    ) -> Result<Bytes> {
31        let bytes = self.get_bytes(src).await?;
32        on_bytes(bytes.len() as u64);
33        Ok(bytes)
34    }
35
36    /// Return whether an object exists at the given path.
37    async fn exists(&self, src: &Path) -> Result<bool>;
38
39    /// Return the size in bytes of the object at the given path.
40    async fn object_size(&self, src: &Path) -> Result<u64>;
41}
42
43/// Await `fut`, failing if it does not resolve within
44/// [`TRANSFER_STALL_TIMEOUT`].
45async fn with_stall_timeout<F: Future>(task_name: &str, src: &Path, fut: F) -> Result<F::Output> {
46    tokio::time::timeout(TRANSFER_STALL_TIMEOUT, fut)
47        .await
48        .map_err(|_| {
49            anyhow!("{task_name} for file {src} received nothing for {TRANSFER_STALL_TIMEOUT:?}")
50        })
51}
52
53/// Collect a GET result's payload into contiguous bytes, invoking `on_bytes`
54/// with each chunk's size as it is received.
55pub(crate) async fn collect_get_result_with_progress(
56    result: object_store::GetResult,
57    src: &Path,
58    on_bytes: &(dyn Fn(u64) + Send + Sync),
59) -> Result<Bytes> {
60    use futures::StreamExt;
61    let mut buf = Vec::with_capacity(result.meta.size as usize);
62    let mut stream = result.into_stream();
63    // Fails once no chunk has arrived for `TRANSFER_STALL_TIMEOUT`.
64    while let Some(chunk) = with_stall_timeout("GET result stream", src, stream.next()).await? {
65        let chunk = chunk
66            .map_err(|e| anyhow!("Failed to stream GET result for file {src} with error: {e:?}"))?;
67        on_bytes(chunk.len() as u64);
68        buf.extend_from_slice(&chunk);
69    }
70    Ok(buf.into())
71}
72
73macro_rules! as_ref_get_ext_impl {
74    ($type:ty) => {
75        #[async_trait]
76        impl ObjectStoreGetExt for $type {
77            async fn get_bytes(&self, src: &Path) -> Result<Bytes> {
78                self.as_ref().get_bytes(src).await
79            }
80
81            async fn get_bytes_with_progress(
82                &self,
83                src: &Path,
84                on_bytes: &(dyn Fn(u64) + Send + Sync),
85            ) -> Result<Bytes> {
86                self.as_ref().get_bytes_with_progress(src, on_bytes).await
87            }
88
89            async fn exists(&self, src: &Path) -> Result<bool> {
90                self.as_ref().exists(src).await
91            }
92
93            async fn object_size(&self, src: &Path) -> Result<u64> {
94                self.as_ref().object_size(src).await
95            }
96        }
97    };
98}
99
100as_ref_get_ext_impl!(Arc<dyn ObjectStoreGetExt>);
101as_ref_get_ext_impl!(Box<dyn ObjectStoreGetExt>);
102
103macro_rules! as_ref_get_impl {
104    ($type:ty) => {
105        #[async_trait]
106        impl ObjectStoreGetExt for $type {
107            async fn get_bytes(&self, src: &Path) -> Result<Bytes> {
108                // Collected chunk by chunk rather than with `bytes()` so that a
109                // transfer that stalls part way through hits the stall timeout.
110                self.get_bytes_with_progress(src, &|_| {}).await
111            }
112
113            async fn get_bytes_with_progress(
114                &self,
115                src: &Path,
116                on_bytes: &(dyn Fn(u64) + Send + Sync),
117            ) -> Result<Bytes> {
118                let result = with_stall_timeout("GET request", src, self.get(src))
119                    .await?
120                    .map_err(|e| anyhow!("Failed to get file {src} with error: {e:?}"))?;
121                collect_get_result_with_progress(result, src, on_bytes).await
122            }
123
124            async fn exists(&self, src: &Path) -> Result<bool> {
125                match with_stall_timeout("HEAD request", src, self.head(src)).await? {
126                    Ok(_) => Ok(true),
127                    Err(object_store::Error::NotFound { .. }) => Ok(false),
128                    Err(e) => Err(anyhow!(
129                        "Failed to check if file {src} exists with error: {e:?}"
130                    )),
131                }
132            }
133
134            async fn object_size(&self, src: &Path) -> Result<u64> {
135                with_stall_timeout("HEAD request", src, self.head(src))
136                    .await?
137                    .map(|meta| meta.size)
138                    .map_err(|e| anyhow!("Failed to get size of file {src} with error: {e:?}"))
139            }
140        }
141    };
142}
143
144as_ref_get_impl!(Arc<dyn ObjectStore>);
145as_ref_get_impl!(Box<dyn ObjectStore>);
146
147#[async_trait]
148pub trait ObjectStoreListExt: Send + Sync + 'static {
149    /// List the objects at the given path in object store
150    async fn list_objects(
151        &self,
152        src: Option<&Path>,
153    ) -> BoxStream<'_, object_store::Result<ObjectMeta>>;
154}
155
156macro_rules! as_ref_list_ext_impl {
157    ($type:ty) => {
158        #[async_trait]
159        impl ObjectStoreListExt for $type {
160            async fn list_objects(
161                &self,
162                src: Option<&Path>,
163            ) -> BoxStream<'_, object_store::Result<ObjectMeta>> {
164                self.as_ref().list_objects(src).await
165            }
166        }
167    };
168}
169
170as_ref_list_ext_impl!(Arc<dyn ObjectStoreListExt>);
171as_ref_list_ext_impl!(Box<dyn ObjectStoreListExt>);
172
173#[async_trait]
174impl ObjectStoreListExt for Arc<DynObjectStore> {
175    async fn list_objects(
176        &self,
177        src: Option<&Path>,
178    ) -> BoxStream<'_, object_store::Result<ObjectMeta>> {
179        self.list(src)
180    }
181}
182
183#[async_trait]
184pub trait ObjectStorePutExt: Send + Sync + 'static {
185    /// Write the bytes at the given location in object store
186    async fn put_bytes(&self, src: &Path, bytes: Bytes) -> Result<()>;
187}
188
189macro_rules! as_ref_put_ext_impl {
190    ($type:ty) => {
191        #[async_trait]
192        impl ObjectStorePutExt for $type {
193            async fn put_bytes(&self, src: &Path, bytes: Bytes) -> Result<()> {
194                self.as_ref().put_bytes(src, bytes).await
195            }
196        }
197    };
198}
199
200as_ref_put_ext_impl!(Arc<dyn ObjectStorePutExt>);
201as_ref_put_ext_impl!(Box<dyn ObjectStorePutExt>);
202
203#[async_trait]
204impl ObjectStorePutExt for Arc<DynObjectStore> {
205    async fn put_bytes(&self, src: &Path, bytes: Bytes) -> Result<()> {
206        self.put(src, bytes.into()).await?;
207        Ok(())
208    }
209}
210
211#[async_trait]
212pub trait ObjectStoreDeleteExt: Send + Sync + 'static {
213    /// Delete the object at the given location in object store
214    async fn delete_object(&self, src: &Path) -> Result<()>;
215}
216
217macro_rules! as_ref_delete_ext_impl {
218    ($type:ty) => {
219        #[async_trait]
220        impl ObjectStoreDeleteExt for $type {
221            async fn delete_object(&self, src: &Path) -> Result<()> {
222                self.as_ref().delete_object(src).await
223            }
224        }
225    };
226}
227
228as_ref_delete_ext_impl!(Arc<dyn ObjectStoreDeleteExt>);
229as_ref_delete_ext_impl!(Box<dyn ObjectStoreDeleteExt>);
230
231#[async_trait]
232
233impl ObjectStoreDeleteExt for Arc<DynObjectStore> {
234    async fn delete_object(&self, src: &Path) -> Result<()> {
235        self.delete(src).await?;
236        Ok(())
237    }
238}
239
240#[cfg(test)]
241mod tests {
242    use std::sync::{
243        Arc,
244        atomic::{AtomicU64, Ordering},
245    };
246
247    use bytes::Bytes;
248    use futures::StreamExt;
249    use object_store::{
250        Attributes, GetResult, GetResultPayload, ObjectMeta, ObjectStore, memory::InMemory,
251        path::Path,
252    };
253
254    use crate::object_store::{
255        ObjectStoreGetExt, ObjectStorePutExt, collect_get_result_with_progress,
256    };
257
258    #[tokio::test]
259    async fn test_dyn_object_store_get_bytes() -> anyhow::Result<()> {
260        let store: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
261        let path = Path::from("file1");
262        store
263            .put_bytes(&path, Bytes::from_static(b"Lorem ipsum"))
264            .await?;
265
266        assert_eq!(
267            store.get_bytes(&path).await?,
268            Bytes::from_static(b"Lorem ipsum")
269        );
270        assert!(store.get_bytes(&Path::from("missing")).await.is_err());
271        Ok(())
272    }
273
274    /// A transfer that goes quiet part way through fails rather than hanging.
275    /// Runs on a paused clock, so the stall timeout elapses without the test
276    /// waiting for it.
277    #[tokio::test(start_paused = true)]
278    async fn test_collect_get_result_with_progress_fails_on_a_stalled_stream() {
279        let src = Path::from("file1");
280        let payload = futures::stream::once(async {
281            Ok::<_, object_store::Error>(Bytes::from_static(b"Lorem"))
282        })
283        .chain(futures::stream::pending())
284        .boxed();
285        let result = GetResult {
286            range: 0..11,
287            payload: GetResultPayload::Stream(payload),
288            meta: ObjectMeta {
289                location: src.clone(),
290                last_modified: chrono::Utc::now(),
291                size: 11,
292                e_tag: None,
293                version: None,
294            },
295            attributes: Attributes::new(),
296        };
297
298        let received = AtomicU64::new(0);
299        let err = collect_get_result_with_progress(result, &src, &|n| {
300            received.fetch_add(n, Ordering::Relaxed);
301        })
302        .await
303        .unwrap_err();
304
305        assert!(err.to_string().contains("received nothing"), "{err}");
306        assert_eq!(received.load(Ordering::Relaxed), 5);
307    }
308
309    #[tokio::test]
310    async fn test_dyn_object_store_exists() -> anyhow::Result<()> {
311        let store: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
312        let path = Path::from("file1");
313        store
314            .put_bytes(&path, bytes::Bytes::from_static(b"Lorem ipsum"))
315            .await?;
316
317        assert!(store.exists(&path).await?);
318        assert!(!store.exists(&Path::from("missing")).await?);
319        Ok(())
320    }
321
322    #[tokio::test]
323    async fn test_dyn_object_store_object_size() -> anyhow::Result<()> {
324        let store: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
325        let path = Path::from("file1");
326        store
327            .put_bytes(&path, bytes::Bytes::from_static(b"Lorem ipsum"))
328            .await?;
329
330        assert_eq!(store.object_size(&path).await?, 11);
331        assert!(store.object_size(&Path::from("missing")).await.is_err());
332        Ok(())
333    }
334}
iota_storage/object_store/mod.rs

iota_storage/object_store/
mod.rs
