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
130pub async fn delete_recursively<S: ObjectStoreDeleteExt + ObjectStoreListExt>(
131    path: &Path,
132    store: &S,
133    concurrency: NonZeroUsize,
134) -> Result<Vec<()>> {
135    let mut paths_to_delete = vec![];
136    let mut paths = store.list_objects(Some(path)).await;
137    while let Some(res) = paths.next().await {
138        if let Ok(object_metadata) = res {
139            paths_to_delete.push(object_metadata.location);
140        } else {
141            return Err(res.err().unwrap().into());
142        }
143    }
144    delete_files(&paths_to_delete, store, concurrency).await
145}
146
147pub fn path_to_filesystem(local_dir_path: PathBuf, location: &Path) -> anyhow::Result<PathBuf> {
148    // Convert an `object_store::path::Path` to `std::path::PathBuf`
149    let path = std::fs::canonicalize(local_dir_path)?;
150    let mut url = Url::from_file_path(&path)
151        .map_err(|_| anyhow!("Failed to parse input path: {}", path.display()))?;
152    url.path_segments_mut()
153        .map_err(|_| anyhow!("Failed to get path segments: {}", path.display()))?
154        .pop_if_empty()
155        .extend(location.parts());
156    let new_path = url
157        .to_file_path()
158        .map_err(|_| anyhow!("Failed to convert url to path: {}", url.as_str()))?;
159    Ok(new_path)
160}
161
162/// This function will find all child directories in the input store which are
163/// of the form "epoch_num" and return a map of epoch number to the directory
164/// path
165pub async fn find_all_dirs_with_epoch_prefix(
166    store: &Arc<DynObjectStore>,
167    prefix: Option<&Path>,
168) -> anyhow::Result<BTreeMap<u64, Path>> {
169    let mut dirs = BTreeMap::new();
170    let entries = store.list_with_delimiter(prefix).await?;
171    for entry in entries.common_prefixes {
172        if let Some(filename) = entry.filename() {
173            if !filename.starts_with("epoch_") || filename.ends_with(".tmp") {
174                continue;
175            }
176            let epoch = filename
177                .split_once('_')
178                .context("Failed to split dir name")
179                .map(|(_, epoch)| epoch.parse::<u64>())??;
180            dirs.insert(epoch, entry);
181        }
182    }
183    Ok(dirs)
184}
185
186/// Finds all epochs in the store and returns them as a sorted list, paired
187/// with each epoch's end timestamp in ms when its metadata file is present, or
188/// 0 otherwise.
189pub async fn list_all_epochs(object_store: Arc<DynObjectStore>) -> Result<Vec<(u64, u64)>> {
190    let remote_epoch_dirs = find_all_dirs_with_epoch_prefix(&object_store, None).await?;
191    let mut out = vec![];
192    let mut success_marker_found = false;
193    for (epoch, path) in remote_epoch_dirs.iter().sorted() {
194        let success_marker = path.child(SUCCESS_MARKER);
195        let get_result = object_store.get(&success_marker).await;
196        match get_result {
197            Err(_) => {
198                if !success_marker_found {
199                    error!("No success marker found for epoch: {epoch}");
200                }
201            }
202            Ok(_) => {
203                let metadata_path = path.child(EPOCH_METADATA_FILENAME);
204                let epoch_end_timestamp_ms = match object_store.get_bytes(&metadata_path).await {
205                    Ok(bytes) => match serde_json::from_slice::<EpochMetadata>(&bytes) {
206                        Ok(metadata) => metadata.epoch_end_timestamp_ms,
207                        Err(err) => {
208                            warn!("Failed to parse epoch metadata for epoch {epoch}: {err}");
209                            0
210                        }
211                    },
212                    Err(_) => 0,
213                };
214                out.push((*epoch, epoch_end_timestamp_ms));
215                success_marker_found = true;
216            }
217        }
218    }
219    Ok(out)
220}
221
222/// Writes the epochs existed in the store to the root MANIFEST (contains only a
223/// list of epochs in the store) every 300 seconds.
224// TODO: Is 300 seconds too frequent? Or should this be triggered by other
225// events?
226pub async fn run_manifest_update_loop(
227    store: Arc<DynObjectStore>,
228    mut recv: tokio::sync::broadcast::Receiver<()>,
229) -> Result<()> {
230    let mut update_interval = tokio::time::interval(Duration::from_secs(300));
231    loop {
232        tokio::select! {
233            _now = update_interval.tick() => {
234                if let Ok(available_epochs) = list_all_epochs(store.clone()).await {
235                    let manifest_path = Path::from(MANIFEST_FILENAME);
236                    let manifest = RootManifest { available_epochs };
237                    put(&store, &manifest_path, Bytes::from(manifest.to_bytes()?)).await?;
238                }
239            },
240             _ = recv.recv() => break,
241        }
242    }
243    Ok(())
244}
245
246/// This function will find all child directories in the input store which are
247/// of the form "epoch_num" and return a map of epoch number to the directory
248/// path
249pub async fn find_all_files_with_epoch_prefix(
250    store: &Arc<DynObjectStore>,
251    prefix: Option<&Path>,
252) -> anyhow::Result<Vec<Range<u64>>> {
253    let mut ranges = Vec::new();
254    let entries = store.list_with_delimiter(prefix).await?;
255    for entry in entries.objects {
256        let checkpoint_seq_range = entry
257            .location
258            .filename()
259            .ok_or(anyhow!("Illegal file name"))?
260            .split_once('.')
261            .context("Failed to split dir name")?
262            .0
263            .split_once('_')
264            .context("Failed to split dir name")
265            .map(|(start, end)| Range {
266                start: start.parse::<u64>().unwrap(),
267                end: end.parse::<u64>().unwrap(),
268            })?;
269
270        ranges.push(checkpoint_seq_range);
271    }
272    Ok(ranges)
273}
274
275/// This function will find missing epoch directories in the input store and
276/// return a list of such epoch numbers. If the highest epoch directory in the
277/// store is `epoch_N` then it is expected that the store will have all epoch
278/// directories from `epoch_0` to `epoch_N`. Additionally, any epoch directory
279/// should have the passed in marker file present or else that epoch number is
280/// already considered as missing.
281/// The returned list will contain epoch_N+1.
282pub async fn find_missing_epochs_dirs(
283    store: &Arc<DynObjectStore>,
284    success_marker: &str,
285) -> anyhow::Result<Vec<u64>> {
286    let remote_checkpoints_by_epoch = find_all_dirs_with_epoch_prefix(store, None).await?;
287    let mut dirs: Vec<_> = remote_checkpoints_by_epoch.iter().collect();
288    dirs.sort_by_key(|(epoch_num, _path)| *epoch_num);
289    let mut candidate_epoch: u64 = 0;
290    let mut missing_epochs = Vec::new();
291    for (epoch_num, path) in dirs {
292        while candidate_epoch < *epoch_num {
293            // The whole epoch directory is missing
294            missing_epochs.push(candidate_epoch);
295            candidate_epoch += 1;
296            continue;
297        }
298        let success_marker = path.child(success_marker);
299        let get_result = store.get(&success_marker).await;
300        match get_result {
301            Err(Error::NotFound { .. }) => {
302                error!("No success marker found in remote store for epoch: {epoch_num}");
303                missing_epochs.push(*epoch_num);
304            }
305            Err(_) => {
306                // Probably a transient error
307                warn!(
308                    "Failed while trying to read success marker in remote store for epoch: {epoch_num}"
309                );
310            }
311            Ok(_) => {
312                // Nothing to do
313            }
314        }
315        candidate_epoch += 1
316    }
317    missing_epochs.push(candidate_epoch);
318    Ok(missing_epochs)
319}
320
321pub fn get_path(prefix: &str) -> Path {
322    Path::from(prefix)
323}
324
325#[cfg(test)]
326mod tests {
327    use std::{fs, num::NonZeroUsize};
328
329    use iota_config::object_storage_config::{ObjectStoreConfig, ObjectStoreType};
330    use object_store::path::Path;
331    use tempfile::TempDir;
332
333    use crate::object_store::util::delete_recursively;
334
335    #[tokio::test]
336    pub async fn test_delete_recursively() -> anyhow::Result<()> {
337        let input = TempDir::new()?;
338        let input_path = input.path();
339        let child = input_path.join("child");
340        fs::create_dir(&child)?;
341        let file1 = child.join("file1");
342        fs::write(file1, b"Lorem ipsum")?;
343        let grandchild = child.join("grand_child");
344        fs::create_dir(&grandchild)?;
345        let file2 = grandchild.join("file2");
346        fs::write(file2, b"Lorem ipsum")?;
347
348        let input_store = ObjectStoreConfig {
349            object_store: Some(ObjectStoreType::File),
350            directory: Some(input_path.to_path_buf()),
351            ..Default::default()
352        }
353        .make()?;
354
355        delete_recursively(
356            &Path::from("child"),
357            &input_store,
358            NonZeroUsize::new(1).unwrap(),
359        )
360        .await?;
361
362        assert!(!input_path.join("child").join("file1").exists());
363        assert!(
364            !input_path
365                .join("child")
366                .join("grand_child")
367                .join("file2")
368                .exists()
369        );
370        Ok(())
371    }
372}
iota_storage/object_store/util.rs

iota_storage/object_store/
util.rs
