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, TryStreamExt};
13use indicatif::ProgressBar;
14use itertools::Itertools;
15use object_store::{DynObjectStore, Error, ObjectStore, ObjectStoreExt, path::Path};
16use serde::{Deserialize, Serialize};
17use tokio::time::Instant;
18use tracing::{error, warn};
19use url::Url;
20
21use crate::object_store::{
22    ObjectStoreDeleteExt, ObjectStoreGetExt, ObjectStoreListExt, ObjectStorePutExt,
23};
24
25pub const MANIFEST_FILENAME: &str = "MANIFEST";
26pub const EPOCH_METADATA_FILENAME: &str = "_epoch_metadata.json";
27
28#[derive(Serialize, Deserialize)]
29pub struct RootManifest {
30    /// Epoch number paired with its end timestamp in ms (or 0 when unknown).
31    pub available_epochs: Vec<(u64, u64)>,
32}
33
34impl RootManifest {
35    pub fn new(available_epochs: Vec<(u64, u64)>) -> Self {
36        RootManifest { available_epochs }
37    }
38
39    pub fn epoch_exists(&self, epoch: u64) -> bool {
40        self.available_epochs.iter().any(|(e, _)| *e == epoch)
41    }
42
43    /// Parse a root MANIFEST from its JSON bytes.
44    pub fn from_bytes(bytes: &[u8]) -> Result<Self> {
45        Ok(serde_json::from_slice(bytes)?)
46    }
47
48    /// Serialize this root MANIFEST into its JSON bytes.
49    pub fn to_bytes(&self) -> Result<Vec<u8>> {
50        Ok(serde_json::to_vec(self)?)
51    }
52}
53
54#[derive(Serialize, Deserialize)]
55pub struct EpochMetadata {
56    pub epoch_end_timestamp_ms: u64,
57}
58
59impl EpochMetadata {
60    pub fn to_bytes(&self) -> Result<Bytes> {
61        Ok(Bytes::from(serde_json::to_vec(self)?))
62    }
63}
64
65#[derive(Debug, Clone)]
66pub struct PerEpochManifest {
67    pub lines: Vec<String>,
68}
69
70impl PerEpochManifest {
71    pub fn new(lines: Vec<String>) -> Self {
72        PerEpochManifest { lines }
73    }
74
75    pub fn serialize_as_newline_delimited(&self) -> String {
76        self.lines.join("\n")
77    }
78
79    pub fn deserialize_from_newline_delimited(s: &str) -> PerEpochManifest {
80        PerEpochManifest {
81            lines: s.lines().map(String::from).collect(),
82        }
83    }
84
85    // Method to filter lines by a given prefix
86    pub fn filter_by_prefix(&self, prefix: &str) -> PerEpochManifest {
87        let filtered_lines = self
88            .lines
89            .iter()
90            .filter(|line| line.starts_with(prefix))
91            .cloned()
92            .collect();
93
94        PerEpochManifest {
95            lines: filtered_lines,
96        }
97    }
98}
99
100pub async fn get<S: ObjectStoreGetExt>(store: &S, src: &Path) -> Result<Bytes> {
101    let bytes = retry(backoff::ExponentialBackoff::default(), || async {
102        store.get_bytes(src).await.map_err(|e| {
103            error!("Failed to read file from object store with error: {:?}", &e);
104            backoff::Error::transient(e)
105        })
106    })
107    .await?;
108    Ok(bytes)
109}
110
111/// Writes bytes in the store with specified path.
112pub async fn put<S: ObjectStorePutExt>(store: &S, src: &Path, bytes: Bytes) -> Result<()> {
113    retry(backoff::ExponentialBackoff::default(), || async {
114        if !bytes.is_empty() {
115            store.put_bytes(src, bytes.clone()).await.map_err(|e| {
116                error!("Failed to write file to object store with error: {:?}", &e);
117                backoff::Error::transient(e)
118            })
119        } else {
120            warn!("Not copying empty file: {:?}", src);
121            Ok(())
122        }
123    })
124    .await?;
125    Ok(())
126}
127
128pub async fn copy_file<S: ObjectStoreGetExt, D: ObjectStorePutExt>(
129    src: &Path,
130    dest: &Path,
131    src_store: &S,
132    dest_store: &D,
133) -> Result<()> {
134    let bytes = get(src_store, src).await?;
135    if !bytes.is_empty() {
136        put(dest_store, dest, bytes).await
137    } else {
138        warn!("Not copying empty file: {:?}", src);
139        Ok(())
140    }
141}
142
143pub async fn copy_files<S: ObjectStoreGetExt, D: ObjectStorePutExt>(
144    src: &[Path],
145    dest: &[Path],
146    src_store: &S,
147    dest_store: &D,
148    concurrency: NonZeroUsize,
149    progress_bar: Option<ProgressBar>,
150) -> Result<Vec<()>> {
151    let mut instant = Instant::now();
152    let progress_bar_clone = progress_bar.clone();
153    // Copies files from dest to src in parallel, and updates the progress bar if
154    // it's provided
155    let results = futures::stream::iter(src.iter().zip(dest.iter()))
156        .map(|(path_in, path_out)| async move {
157            let ret = copy_file(path_in, path_out, src_store, dest_store).await;
158            Ok((path_out.clone(), ret))
159        })
160        .boxed()
161        .buffer_unordered(concurrency.into())
162        .try_for_each(|(path, ret)| {
163            if let Some(progress_bar_clone) = &progress_bar_clone {
164                progress_bar_clone.inc(1);
165                progress_bar_clone.set_message(format!("file: {path}"));
166                instant = Instant::now();
167            }
168            futures::future::ready(ret)
169        })
170        .await;
171    Ok(results.into_iter().collect())
172}
173
174/// Copies all files in the directory from the source store to the destination
175/// store.
176pub async fn copy_recursively<S: ObjectStoreGetExt + ObjectStoreListExt, D: ObjectStorePutExt>(
177    dir: &Path,
178    src_store: &S,
179    dest_store: &D,
180    concurrency: NonZeroUsize,
181) -> Result<Vec<()>> {
182    let mut input_paths = vec![];
183    let mut output_paths = vec![];
184    let mut paths = src_store.list_objects(Some(dir)).await;
185    while let Some(res) = paths.next().await {
186        if let Ok(object_metadata) = res {
187            input_paths.push(object_metadata.location.clone());
188            output_paths.push(object_metadata.location);
189        } else {
190            return Err(res.err().unwrap().into());
191        }
192    }
193    copy_files(
194        &input_paths,
195        &output_paths,
196        src_store,
197        dest_store,
198        concurrency,
199        None,
200    )
201    .await
202}
203
204pub async fn delete_files<S: ObjectStoreDeleteExt>(
205    files: &[Path],
206    store: &S,
207    concurrency: NonZeroUsize,
208) -> Result<Vec<()>> {
209    let results: Vec<Result<()>> = futures::stream::iter(files)
210        .map(|f| {
211            retry(backoff::ExponentialBackoff::default(), || async {
212                store.delete_object(f).await.map_err(|e| {
213                    error!("Failed to delete file on object store with error: {:?}", &e);
214                    backoff::Error::transient(e)
215                })
216            })
217        })
218        .boxed()
219        .buffer_unordered(concurrency.into())
220        .collect()
221        .await;
222    results.into_iter().collect()
223}
224
225pub async fn delete_recursively<S: ObjectStoreDeleteExt + ObjectStoreListExt>(
226    path: &Path,
227    store: &S,
228    concurrency: NonZeroUsize,
229) -> Result<Vec<()>> {
230    let mut paths_to_delete = vec![];
231    let mut paths = store.list_objects(Some(path)).await;
232    while let Some(res) = paths.next().await {
233        if let Ok(object_metadata) = res {
234            paths_to_delete.push(object_metadata.location);
235        } else {
236            return Err(res.err().unwrap().into());
237        }
238    }
239    delete_files(&paths_to_delete, store, concurrency).await
240}
241
242pub fn path_to_filesystem(local_dir_path: PathBuf, location: &Path) -> anyhow::Result<PathBuf> {
243    // Convert an `object_store::path::Path` to `std::path::PathBuf`
244    let path = std::fs::canonicalize(local_dir_path)?;
245    let mut url = Url::from_file_path(&path)
246        .map_err(|_| anyhow!("Failed to parse input path: {}", path.display()))?;
247    url.path_segments_mut()
248        .map_err(|_| anyhow!("Failed to get path segments: {}", path.display()))?
249        .pop_if_empty()
250        .extend(location.parts());
251    let new_path = url
252        .to_file_path()
253        .map_err(|_| anyhow!("Failed to convert url to path: {}", url.as_str()))?;
254    Ok(new_path)
255}
256
257/// This function will find all child directories in the input store which are
258/// of the form "epoch_num" and return a map of epoch number to the directory
259/// path
260pub async fn find_all_dirs_with_epoch_prefix(
261    store: &Arc<DynObjectStore>,
262    prefix: Option<&Path>,
263) -> anyhow::Result<BTreeMap<u64, Path>> {
264    let mut dirs = BTreeMap::new();
265    let entries = store.list_with_delimiter(prefix).await?;
266    for entry in entries.common_prefixes {
267        if let Some(filename) = entry.filename() {
268            if !filename.starts_with("epoch_") || filename.ends_with(".tmp") {
269                continue;
270            }
271            let epoch = filename
272                .split_once('_')
273                .context("Failed to split dir name")
274                .map(|(_, epoch)| epoch.parse::<u64>())??;
275            dirs.insert(epoch, entry);
276        }
277    }
278    Ok(dirs)
279}
280
281/// Finds all epochs in the store and returns them as a sorted list, paired
282/// with each epoch's end timestamp in ms when its metadata file is present, or
283/// 0 otherwise.
284pub async fn list_all_epochs(object_store: Arc<DynObjectStore>) -> Result<Vec<(u64, u64)>> {
285    let remote_epoch_dirs = find_all_dirs_with_epoch_prefix(&object_store, None).await?;
286    let mut out = vec![];
287    let mut success_marker_found = false;
288    for (epoch, path) in remote_epoch_dirs.iter().sorted() {
289        let success_marker = path.child("_SUCCESS");
290        let get_result = object_store.get(&success_marker).await;
291        match get_result {
292            Err(_) => {
293                if !success_marker_found {
294                    error!("No success marker found for epoch: {epoch}");
295                }
296            }
297            Ok(_) => {
298                let metadata_path = path.child(EPOCH_METADATA_FILENAME);
299                let epoch_end_timestamp_ms = match object_store.get_bytes(&metadata_path).await {
300                    Ok(bytes) => match serde_json::from_slice::<EpochMetadata>(&bytes) {
301                        Ok(metadata) => metadata.epoch_end_timestamp_ms,
302                        Err(err) => {
303                            warn!("Failed to parse epoch metadata for epoch {epoch}: {err}");
304                            0
305                        }
306                    },
307                    Err(_) => 0,
308                };
309                out.push((*epoch, epoch_end_timestamp_ms));
310                success_marker_found = true;
311            }
312        }
313    }
314    Ok(out)
315}
316
317/// Writes the epochs existed in the store to the root MANIFEST (contains only a
318/// list of epochs in the store) every 300 seconds.
319// TODO: Is 300 seconds too frequent? Or should this be triggered by other
320// events?
321pub async fn run_manifest_update_loop(
322    store: Arc<DynObjectStore>,
323    mut recv: tokio::sync::broadcast::Receiver<()>,
324) -> Result<()> {
325    let mut update_interval = tokio::time::interval(Duration::from_secs(300));
326    loop {
327        tokio::select! {
328            _now = update_interval.tick() => {
329                if let Ok(available_epochs) = list_all_epochs(store.clone()).await {
330                    let manifest_path = Path::from(MANIFEST_FILENAME);
331                    let manifest = RootManifest { available_epochs };
332                    put(&store, &manifest_path, Bytes::from(manifest.to_bytes()?)).await?;
333                }
334            },
335             _ = recv.recv() => break,
336        }
337    }
338    Ok(())
339}
340
341/// This function will find all child directories in the input store which are
342/// of the form "epoch_num" and return a map of epoch number to the directory
343/// path
344pub async fn find_all_files_with_epoch_prefix(
345    store: &Arc<DynObjectStore>,
346    prefix: Option<&Path>,
347) -> anyhow::Result<Vec<Range<u64>>> {
348    let mut ranges = Vec::new();
349    let entries = store.list_with_delimiter(prefix).await?;
350    for entry in entries.objects {
351        let checkpoint_seq_range = entry
352            .location
353            .filename()
354            .ok_or(anyhow!("Illegal file name"))?
355            .split_once('.')
356            .context("Failed to split dir name")?
357            .0
358            .split_once('_')
359            .context("Failed to split dir name")
360            .map(|(start, end)| Range {
361                start: start.parse::<u64>().unwrap(),
362                end: end.parse::<u64>().unwrap(),
363            })?;
364
365        ranges.push(checkpoint_seq_range);
366    }
367    Ok(ranges)
368}
369
370/// This function will find missing epoch directories in the input store and
371/// return a list of such epoch numbers. If the highest epoch directory in the
372/// store is `epoch_N` then it is expected that the store will have all epoch
373/// directories from `epoch_0` to `epoch_N`. Additionally, any epoch directory
374/// should have the passed in marker file present or else that epoch number is
375/// already considered as missing.
376/// The returned list will contain epoch_N+1.
377pub async fn find_missing_epochs_dirs(
378    store: &Arc<DynObjectStore>,
379    success_marker: &str,
380) -> anyhow::Result<Vec<u64>> {
381    let remote_checkpoints_by_epoch = find_all_dirs_with_epoch_prefix(store, None).await?;
382    let mut dirs: Vec<_> = remote_checkpoints_by_epoch.iter().collect();
383    dirs.sort_by_key(|(epoch_num, _path)| *epoch_num);
384    let mut candidate_epoch: u64 = 0;
385    let mut missing_epochs = Vec::new();
386    for (epoch_num, path) in dirs {
387        while candidate_epoch < *epoch_num {
388            // The whole epoch directory is missing
389            missing_epochs.push(candidate_epoch);
390            candidate_epoch += 1;
391            continue;
392        }
393        let success_marker = path.child(success_marker);
394        let get_result = store.get(&success_marker).await;
395        match get_result {
396            Err(Error::NotFound { .. }) => {
397                error!("No success marker found in db checkpoint for epoch: {epoch_num}");
398                missing_epochs.push(*epoch_num);
399            }
400            Err(_) => {
401                // Probably a transient error
402                warn!(
403                    "Failed while trying to read success marker in db checkpoint for epoch: {epoch_num}"
404                );
405            }
406            Ok(_) => {
407                // Nothing to do
408            }
409        }
410        candidate_epoch += 1
411    }
412    missing_epochs.push(candidate_epoch);
413    Ok(missing_epochs)
414}
415
416pub fn get_path(prefix: &str) -> Path {
417    Path::from(prefix)
418}
419
420// Snapshot MANIFEST file is very simple. Just a newline delimited list of all
421// paths in the snapshot directory this simplicity enables easy parsing for
422// scripts to download snapshots
423pub async fn write_snapshot_manifest<S: ObjectStoreListExt + ObjectStorePutExt>(
424    dir: &Path,
425    store: &S,
426    epoch_prefix: String,
427) -> Result<()> {
428    let mut file_names = vec![];
429    let mut paths = store.list_objects(Some(dir)).await;
430    while let Some(res) = paths.next().await {
431        if let Ok(object_metadata) = res {
432            // trim the "epoch_XX/" dir prefix here
433            let mut path_str = object_metadata.location.to_string();
434            if path_str.starts_with(&epoch_prefix) {
435                path_str = String::from(&path_str[epoch_prefix.len()..]);
436                file_names.push(path_str);
437            } else {
438                warn!("{path_str}, should be coming from the files in the {epoch_prefix} dir",)
439            }
440        } else {
441            return Err(res.err().unwrap().into());
442        }
443    }
444
445    let epoch_manifest = PerEpochManifest::new(file_names);
446    let bytes = Bytes::from(epoch_manifest.serialize_as_newline_delimited());
447    put(
448        store,
449        &Path::from(format!("{dir}/{MANIFEST_FILENAME}")),
450        bytes,
451    )
452    .await?;
453
454    Ok(())
455}
456
457#[cfg(test)]
458mod tests {
459    use std::{fs, num::NonZeroUsize};
460
461    use iota_config::object_storage_config::{ObjectStoreConfig, ObjectStoreType};
462    use object_store::path::Path;
463    use tempfile::TempDir;
464
465    use crate::object_store::util::{
466        MANIFEST_FILENAME, copy_recursively, delete_recursively, write_snapshot_manifest,
467    };
468
469    #[tokio::test]
470    pub async fn test_copy_recursively() -> anyhow::Result<()> {
471        let input = TempDir::new()?;
472        let input_path = input.path();
473        let child = input_path.join("child");
474        fs::create_dir(&child)?;
475        let file1 = child.join("file1");
476        fs::write(file1, b"Lorem ipsum")?;
477        let grandchild = child.join("grand_child");
478        fs::create_dir(&grandchild)?;
479        let file2 = grandchild.join("file2");
480        fs::write(file2, b"Lorem ipsum")?;
481
482        let output = TempDir::new()?;
483        let output_path = output.path();
484
485        let input_store = ObjectStoreConfig {
486            object_store: Some(ObjectStoreType::File),
487            directory: Some(input_path.to_path_buf()),
488            ..Default::default()
489        }
490        .make()?;
491
492        let output_store = ObjectStoreConfig {
493            object_store: Some(ObjectStoreType::File),
494            directory: Some(output_path.to_path_buf()),
495            ..Default::default()
496        }
497        .make()?;
498
499        copy_recursively(
500            &Path::from("child"),
501            &input_store,
502            &output_store,
503            NonZeroUsize::new(1).unwrap(),
504        )
505        .await?;
506
507        assert!(output_path.join("child").exists());
508        assert!(output_path.join("child").join("file1").exists());
509        assert!(output_path.join("child").join("grand_child").exists());
510        assert!(
511            output_path
512                .join("child")
513                .join("grand_child")
514                .join("file2")
515                .exists()
516        );
517        let content = fs::read_to_string(output_path.join("child").join("file1"))?;
518        assert_eq!(content, "Lorem ipsum");
519        let content =
520            fs::read_to_string(output_path.join("child").join("grand_child").join("file2"))?;
521        assert_eq!(content, "Lorem ipsum");
522        Ok(())
523    }
524
525    #[tokio::test]
526    pub async fn test_write_snapshot_manifest() -> anyhow::Result<()> {
527        let input = TempDir::new()?;
528        let input_path = input.path();
529        let epoch_0 = input_path.join("epoch_0");
530        fs::create_dir(&epoch_0)?;
531        let file1 = epoch_0.join("file1");
532        fs::write(file1, b"Lorem ipsum")?;
533        let file2 = epoch_0.join("file2");
534        fs::write(file2, b"Lorem ipsum")?;
535        let grandchild = epoch_0.join("grand_child");
536        fs::create_dir(&grandchild)?;
537        let file3 = grandchild.join("file2.tar.gz");
538        fs::write(file3, b"Lorem ipsum")?;
539
540        let input_store = ObjectStoreConfig {
541            object_store: Some(ObjectStoreType::File),
542            directory: Some(input_path.to_path_buf()),
543            ..Default::default()
544        }
545        .make()?;
546
547        write_snapshot_manifest(
548            &Path::from("epoch_0"),
549            &input_store,
550            String::from("epoch_0/"),
551        )
552        .await?;
553
554        assert!(input_path.join("epoch_0").join(MANIFEST_FILENAME).exists());
555        let content = fs::read_to_string(input_path.join("epoch_0").join(MANIFEST_FILENAME))?;
556        assert!(content.contains("file2"));
557        assert!(content.contains("file1"));
558        assert!(content.contains("grand_child/file2.tar.gz"));
559        Ok(())
560    }
561
562    #[tokio::test]
563    pub async fn test_delete_recursively() -> anyhow::Result<()> {
564        let input = TempDir::new()?;
565        let input_path = input.path();
566        let child = input_path.join("child");
567        fs::create_dir(&child)?;
568        let file1 = child.join("file1");
569        fs::write(file1, b"Lorem ipsum")?;
570        let grandchild = child.join("grand_child");
571        fs::create_dir(&grandchild)?;
572        let file2 = grandchild.join("file2");
573        fs::write(file2, b"Lorem ipsum")?;
574
575        let input_store = ObjectStoreConfig {
576            object_store: Some(ObjectStoreType::File),
577            directory: Some(input_path.to_path_buf()),
578            ..Default::default()
579        }
580        .make()?;
581
582        delete_recursively(
583            &Path::from("child"),
584            &input_store,
585            NonZeroUsize::new(1).unwrap(),
586        )
587        .await?;
588
589        assert!(!input_path.join("child").join("file1").exists());
590        assert!(
591            !input_path
592                .join("child")
593                .join("grand_child")
594                .join("file2")
595                .exists()
596        );
597        Ok(())
598    }
599}
iota_storage/object_store/util.rs

iota_storage/object_store/
util.rs
