iota_storage/object_store/
util.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::{
6    collections::BTreeMap, num::NonZeroUsize, ops::Range, path::PathBuf, sync::Arc, time::Duration,
7};
8
9use anyhow::{Context, Result, anyhow};
10use backoff::future::retry;
11use bytes::Bytes;
12use futures::StreamExt;
13use itertools::Itertools;
14use object_store::{DynObjectStore, Error, ObjectStore, ObjectStoreExt, path::Path};
15use serde::{Deserialize, Serialize};
16use tracing::{error, warn};
17use url::Url;
18
19use crate::object_store::{
20    ObjectStoreDeleteExt, ObjectStoreGetExt, ObjectStoreListExt, ObjectStorePutExt,
21};
22
23pub const MANIFEST_FILENAME: &str = "MANIFEST";
24pub const EPOCH_METADATA_FILENAME: &str = "_epoch_metadata.json";
25/// Marker file written to an epoch directory in the store once all files for
26/// that epoch have been written.
27pub const SUCCESS_MARKER: &str = "_SUCCESS";
28
29#[derive(Serialize, Deserialize)]
30pub struct RootManifest {
31    /// Epoch number paired with its end timestamp in ms (or 0 when unknown).
32    pub available_epochs: Vec<(u64, u64)>,
33}
34
35impl RootManifest {
36    pub fn new(available_epochs: Vec<(u64, u64)>) -> Self {
37        RootManifest { available_epochs }
38    }
39
40    pub fn epoch_exists(&self, epoch: u64) -> bool {
41        self.available_epochs.iter().any(|(e, _)| *e == epoch)
42    }
43
44    /// Parse a root MANIFEST from its JSON bytes.
45    pub fn from_bytes(bytes: &[u8]) -> Result<Self> {
46        Ok(serde_json::from_slice(bytes)?)
47    }
48
49    /// Serialize this root MANIFEST into its JSON bytes.
50    pub fn to_bytes(&self) -> Result<Vec<u8>> {
51        Ok(serde_json::to_vec(self)?)
52    }
53}
54
55#[derive(Serialize, Deserialize)]
56pub struct EpochMetadata {
57    pub epoch_end_timestamp_ms: u64,
58}
59
60impl EpochMetadata {
61    pub fn to_bytes(&self) -> Result<Bytes> {
62        Ok(Bytes::from(serde_json::to_vec(self)?))
63    }
64}
65
66pub async fn get<S: ObjectStoreGetExt>(store: &S, src: &Path) -> Result<Bytes> {
67    let bytes = retry(backoff::ExponentialBackoff::default(), || async {
68        store.get_bytes(src).await.map_err(|e| {
69            error!("Failed to read file from object store with error: {:?}", &e);
70            backoff::Error::transient(e)
71        })
72    })
73    .await?;
74    Ok(bytes)
75}
76
77/// Writes bytes in the store with specified path.
78pub async fn put<S: ObjectStorePutExt>(store: &S, src: &Path, bytes: Bytes) -> Result<()> {
79    retry(backoff::ExponentialBackoff::default(), || async {
80        if !bytes.is_empty() {
81            store.put_bytes(src, bytes.clone()).await.map_err(|e| {
82                error!("Failed to write file to object store with error: {:?}", &e);
83                backoff::Error::transient(e)
84            })
85        } else {
86            warn!("Not copying empty file: {:?}", src);
87            Ok(())
88        }
89    })
90    .await?;
91    Ok(())
92}
93
94pub async fn copy_file<S: ObjectStoreGetExt, D: ObjectStorePutExt>(
95    src: &Path,
96    dest: &Path,
97    src_store: &S,
98    dest_store: &D,
99) -> Result<()> {
100    let bytes = get(src_store, src).await?;
101    if !bytes.is_empty() {
102        put(dest_store, dest, bytes).await
103    } else {
104        warn!("Not copying empty file: {:?}", src);
105        Ok(())
106    }
107}
108
109pub async fn delete_files<S: ObjectStoreDeleteExt>(
110    files: &[Path],
111    store: &S,
112    concurrency: NonZeroUsize,
113) -> Result<Vec<()>> {
114    let results: Vec<Result<()>> = futures::stream::iter(files)
115        .map(|f| {
116            retry(backoff::ExponentialBackoff::default(), || async {
117                store.delete_object(f).await.map_err(|e| {
118                    error!("Failed to delete file on object store with error: {:?}", &e);
119                    backoff::Error::transient(e)
120                })
121            })
122        })
123        .boxed()
124        .buffer_unordered(concurrency.into())
125        .collect()
126        .await;
127    results.into_iter().collect()
128}
129
130/// Deletes every object under `path`. Every [`SUCCESS_MARKER`] under `path` is
131/// deleted before the other objects, so a delete that stops partway does not leave the
132/// directory listed as complete without its files. If deleting a marker
133/// fails, nothing else is deleted.
134pub async fn delete_recursively<S: ObjectStoreDeleteExt + ObjectStoreListExt>(
135    path: &Path,
136    store: &S,
137    concurrency: NonZeroUsize,
138) -> Result<Vec<()>> {
139    let mut paths_to_delete = vec![];
140    let mut paths = store.list_objects(Some(path)).await;
141    while let Some(res) = paths.next().await {
142        if let Ok(object_metadata) = res {
143            paths_to_delete.push(object_metadata.location);
144        } else {
145            return Err(res.err().unwrap().into());
146        }
147    }
148    let (markers, files): (Vec<_>, Vec<_>) = paths_to_delete
149        .into_iter()
150        .partition(|location| location.filename() == Some(SUCCESS_MARKER));
151    let mut deleted = delete_files(&markers, store, concurrency).await?;
152    deleted.extend(delete_files(&files, store, concurrency).await?);
153    Ok(deleted)
154}
155
156pub fn path_to_filesystem(local_dir_path: PathBuf, location: &Path) -> anyhow::Result<PathBuf> {
157    // Convert an `object_store::path::Path` to `std::path::PathBuf`
158    let path = std::fs::canonicalize(local_dir_path)?;
159    let mut url = Url::from_file_path(&path)
160        .map_err(|_| anyhow!("Failed to parse input path: {}", path.display()))?;
161    url.path_segments_mut()
162        .map_err(|_| anyhow!("Failed to get path segments: {}", path.display()))?
163        .pop_if_empty()
164        .extend(location.parts());
165    let new_path = url
166        .to_file_path()
167        .map_err(|_| anyhow!("Failed to convert url to path: {}", url.as_str()))?;
168    Ok(new_path)
169}
170
171/// This function will find all child directories in the input store which are
172/// of the form "epoch_num" and return a map of epoch number to the directory
173/// path
174pub async fn find_all_dirs_with_epoch_prefix(
175    store: &Arc<DynObjectStore>,
176    prefix: Option<&Path>,
177) -> anyhow::Result<BTreeMap<u64, Path>> {
178    let mut dirs = BTreeMap::new();
179    let entries = store.list_with_delimiter(prefix).await?;
180    for entry in entries.common_prefixes {
181        if let Some(filename) = entry.filename() {
182            if !filename.starts_with("epoch_") || filename.ends_with(".tmp") {
183                continue;
184            }
185            let epoch = filename
186                .split_once('_')
187                .context("Failed to split dir name")
188                .map(|(_, epoch)| epoch.parse::<u64>())??;
189            dirs.insert(epoch, entry);
190        }
191    }
192    Ok(dirs)
193}
194
195/// Finds all epochs in the store and returns them as a sorted list, paired
196/// with each epoch's end timestamp in ms when its metadata file is present, or
197/// 0 otherwise.
198pub async fn list_all_epochs(object_store: Arc<DynObjectStore>) -> Result<Vec<(u64, u64)>> {
199    let remote_epoch_dirs = find_all_dirs_with_epoch_prefix(&object_store, None).await?;
200    let mut out = vec![];
201    let mut success_marker_found = false;
202    for (epoch, path) in remote_epoch_dirs.iter().sorted() {
203        let success_marker = path.child(SUCCESS_MARKER);
204        let get_result = object_store.get(&success_marker).await;
205        match get_result {
206            Err(_) => {
207                if !success_marker_found {
208                    error!("No success marker found for epoch: {epoch}");
209                }
210            }
211            Ok(_) => {
212                let metadata_path = path.child(EPOCH_METADATA_FILENAME);
213                let epoch_end_timestamp_ms = match object_store.get_bytes(&metadata_path).await {
214                    Ok(bytes) => match serde_json::from_slice::<EpochMetadata>(&bytes) {
215                        Ok(metadata) => metadata.epoch_end_timestamp_ms,
216                        Err(err) => {
217                            warn!("Failed to parse epoch metadata for epoch {epoch}: {err}");
218                            0
219                        }
220                    },
221                    Err(_) => 0,
222                };
223                out.push((*epoch, epoch_end_timestamp_ms));
224                success_marker_found = true;
225            }
226        }
227    }
228    Ok(out)
229}
230
231/// Writes the epochs existed in the store to the root MANIFEST (contains only a
232/// list of epochs in the store) every 300 seconds.
233// TODO: Is 300 seconds too frequent? Or should this be triggered by other
234// events?
235pub async fn run_manifest_update_loop(
236    store: Arc<DynObjectStore>,
237    mut recv: tokio::sync::broadcast::Receiver<()>,
238) -> Result<()> {
239    let mut update_interval = tokio::time::interval(Duration::from_secs(300));
240    loop {
241        tokio::select! {
242            _now = update_interval.tick() => {
243                if let Ok(available_epochs) = list_all_epochs(store.clone()).await {
244                    let manifest_path = Path::from(MANIFEST_FILENAME);
245                    let manifest = RootManifest { available_epochs };
246                    put(&store, &manifest_path, Bytes::from(manifest.to_bytes()?)).await?;
247                }
248            },
249             _ = recv.recv() => break,
250        }
251    }
252    Ok(())
253}
254
255/// This function will find all child directories in the input store which are
256/// of the form "epoch_num" and return a map of epoch number to the directory
257/// path
258pub async fn find_all_files_with_epoch_prefix(
259    store: &Arc<DynObjectStore>,
260    prefix: Option<&Path>,
261) -> anyhow::Result<Vec<Range<u64>>> {
262    let mut ranges = Vec::new();
263    let entries = store.list_with_delimiter(prefix).await?;
264    for entry in entries.objects {
265        let checkpoint_seq_range = entry
266            .location
267            .filename()
268            .ok_or(anyhow!("Illegal file name"))?
269            .split_once('.')
270            .context("Failed to split dir name")?
271            .0
272            .split_once('_')
273            .context("Failed to split dir name")
274            .map(|(start, end)| Range {
275                start: start.parse::<u64>().unwrap(),
276                end: end.parse::<u64>().unwrap(),
277            })?;
278
279        ranges.push(checkpoint_seq_range);
280    }
281    Ok(ranges)
282}
283
284/// This function will find missing epoch directories in the input store and
285/// return a list of such epoch numbers. If the highest epoch directory in the
286/// store is `epoch_N` then it is expected that the store will have all epoch
287/// directories from `epoch_0` to `epoch_N`. Additionally, any epoch directory
288/// should have the passed in marker file present or else that epoch number is
289/// already considered as missing.
290/// The returned list will contain epoch_N+1.
291pub async fn find_missing_epochs_dirs(
292    store: &Arc<DynObjectStore>,
293    success_marker: &str,
294) -> anyhow::Result<Vec<u64>> {
295    let remote_checkpoints_by_epoch = find_all_dirs_with_epoch_prefix(store, None).await?;
296    let mut dirs: Vec<_> = remote_checkpoints_by_epoch.iter().collect();
297    dirs.sort_by_key(|(epoch_num, _path)| *epoch_num);
298    let mut candidate_epoch: u64 = 0;
299    let mut missing_epochs = Vec::new();
300    for (epoch_num, path) in dirs {
301        while candidate_epoch < *epoch_num {
302            // The whole epoch directory is missing
303            missing_epochs.push(candidate_epoch);
304            candidate_epoch += 1;
305            continue;
306        }
307        let success_marker = path.child(success_marker);
308        let get_result = store.get(&success_marker).await;
309        match get_result {
310            Err(Error::NotFound { .. }) => {
311                error!("No success marker found in remote store for epoch: {epoch_num}");
312                missing_epochs.push(*epoch_num);
313            }
314            Err(_) => {
315                // Probably a transient error
316                warn!(
317                    "Failed while trying to read success marker in remote store for epoch: {epoch_num}"
318                );
319            }
320            Ok(_) => {
321                // Nothing to do
322            }
323        }
324        candidate_epoch += 1
325    }
326    missing_epochs.push(candidate_epoch);
327    Ok(missing_epochs)
328}
329
330pub fn get_path(prefix: &str) -> Path {
331    Path::from(prefix)
332}
333
334#[cfg(test)]
335mod tests {
336    use std::{fs, num::NonZeroUsize, sync::Arc, time::Duration};
337
338    use async_trait::async_trait;
339    use bytes::Bytes;
340    use futures::stream::BoxStream;
341    use iota_config::object_storage_config::{ObjectStoreConfig, ObjectStoreType};
342    use object_store::{DynObjectStore, ObjectMeta, ObjectStoreExt, memory::InMemory, path::Path};
343    use tempfile::TempDir;
344
345    use crate::object_store::{
346        ObjectStoreDeleteExt, ObjectStoreListExt,
347        util::{MANIFEST_FILENAME, SUCCESS_MARKER, delete_recursively, put},
348    };
349
350    /// Deletes either only a success marker or everything but it; every other
351    /// delete never finishes, as when the process stops partway through a
352    /// delete.
353    struct StopsDeleting {
354        store: Arc<DynObjectStore>,
355        deletes_marker: bool,
356    }
357
358    #[async_trait]
359    impl ObjectStoreDeleteExt for StopsDeleting {
360        async fn delete_object(&self, src: &Path) -> anyhow::Result<()> {
361            if (src.filename() == Some(SUCCESS_MARKER)) == self.deletes_marker {
362                self.store.delete_object(src).await
363            } else {
364                std::future::pending().await
365            }
366        }
367    }
368
369    #[async_trait]
370    impl ObjectStoreListExt for StopsDeleting {
371        async fn list_objects(
372            &self,
373            src: Option<&Path>,
374        ) -> BoxStream<'_, object_store::Result<ObjectMeta>> {
375            self.store.list_objects(src).await
376        }
377    }
378
379    /// Puts a complete epoch directory into an in-memory store and starts a
380    /// delete of it that stops partway, deleting only the success marker or
381    /// only the other files. Returns the store and the directory.
382    async fn interrupted_delete(deletes_marker: bool) -> (Arc<DynObjectStore>, Path) {
383        let store: Arc<DynObjectStore> = Arc::new(InMemory::new());
384        let epoch_dir = Path::from("epoch_7");
385        for file in [
386            "1_1.obj",
387            "1_1.ref",
388            "EPOCH_INFO",
389            MANIFEST_FILENAME,
390            SUCCESS_MARKER,
391        ] {
392            put(&store, &epoch_dir.child(file), Bytes::from_static(b"data"))
393                .await
394                .unwrap();
395        }
396
397        let stopping_store = StopsDeleting {
398            store: store.clone(),
399            deletes_marker,
400        };
401        let delete =
402            delete_recursively(&epoch_dir, &stopping_store, NonZeroUsize::new(20).unwrap());
403        assert!(
404            tokio::time::timeout(Duration::from_secs(1), delete)
405                .await
406                .is_err(),
407            "delete_recursively must not finish while one of its deletes hangs"
408        );
409        (store, epoch_dir)
410    }
411
412    #[tokio::test(start_paused = true)]
413    async fn an_interrupted_delete_has_removed_the_success_marker() {
414        let (store, epoch_dir) = interrupted_delete(true).await;
415
416        assert!(matches!(
417            store.head(&epoch_dir.child(SUCCESS_MARKER)).await,
418            Err(object_store::Error::NotFound { .. })
419        ));
420        store
421            .head(&epoch_dir.child(MANIFEST_FILENAME))
422            .await
423            .unwrap();
424    }
425
426    #[tokio::test(start_paused = true)]
427    async fn nothing_else_is_deleted_before_the_success_marker() {
428        let (store, epoch_dir) = interrupted_delete(false).await;
429
430        for file in ["1_1.obj", MANIFEST_FILENAME, SUCCESS_MARKER] {
431            store.head(&epoch_dir.child(file)).await.unwrap();
432        }
433    }
434
435    #[tokio::test]
436    pub async fn test_delete_recursively() -> anyhow::Result<()> {
437        let input = TempDir::new()?;
438        let input_path = input.path();
439        let child = input_path.join("child");
440        fs::create_dir(&child)?;
441        let file1 = child.join("file1");
442        fs::write(file1, b"Lorem ipsum")?;
443        let grandchild = child.join("grand_child");
444        fs::create_dir(&grandchild)?;
445        let file2 = grandchild.join("file2");
446        fs::write(file2, b"Lorem ipsum")?;
447
448        let input_store = ObjectStoreConfig {
449            object_store: Some(ObjectStoreType::File),
450            directory: Some(input_path.to_path_buf()),
451            ..Default::default()
452        }
453        .make()?;
454
455        delete_recursively(
456            &Path::from("child"),
457            &input_store,
458            NonZeroUsize::new(1).unwrap(),
459        )
460        .await?;
461
462        assert!(!input_path.join("child").join("file1").exists());
463        assert!(
464            !input_path
465                .join("child")
466                .join("grand_child")
467                .join("file2")
468                .exists()
469        );
470        Ok(())
471    }
472}
iota_storage/object_store/util.rs

iota_storage/object_store/
util.rs
