iota_data_ingestion_core/progress_store/
file.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::{io::SeekFrom, path::PathBuf};
6
7use async_trait::async_trait;
8use iota_types::messages_checkpoint::CheckpointSequenceNumber;
9use serde_json::{Number, Value};
10use tokio::{
11    fs::File,
12    io::{AsyncReadExt, AsyncSeekExt, AsyncWriteExt},
13};
14
15use crate::{IngestionError, IngestionResult, progress_store::ProgressStore};
16
17/// Manages persistent progress information stored in a JSON file.
18///
19/// This struct encapsulates file operations for reading, writing, and
20/// synchronizing progress data to disk. It uses asynchronous I/O provided by
21/// [`tokio::fs`](tokio::fs) for efficient operation within a Tokio runtime.
22///
23/// # Example
24/// ```
25/// use iota_data_ingestion_core::{FileProgressStore, ProgressStore};
26///
27/// #[tokio::main]
28/// async fn main() {
29///     let mut store = FileProgressStore::new("progress.json").await.unwrap();
30///     store.save("task1".into(), 42).await.unwrap();
31///     # tokio::time::sleep(std::time::Duration::from_secs(1)).await;
32///     let checkpoint = store.load("task1".into()).await.unwrap();
33///     # tokio::fs::remove_file("progress.json").await.unwrap();
34///     assert_eq!(checkpoint, 42);
35/// }
36/// ```
37pub struct FileProgressStore {
38    /// The path to the progress file.
39    path: PathBuf,
40    /// The [`File`] handle used to interact with the progress file.
41    file: File,
42}
43
44impl FileProgressStore {
45    /// Creates a new `FileProgressStore` by opening or creating the file at the
46    /// specified path.
47    pub async fn new(path: impl Into<PathBuf>) -> IngestionResult<Self> {
48        let path = path.into();
49        Self::open_or_create_file(&path)
50            .await
51            .map(|file| Self { file, path })
52    }
53
54    /// Open or create the file at the specified path.
55    async fn open_or_create_file(path: &PathBuf) -> IngestionResult<File> {
56        Ok(File::options()
57            .read(true)
58            .write(true)
59            .create(true)
60            .truncate(false)
61            .open(path)
62            .await?)
63    }
64
65    /// Returns an empty JSON object.
66    fn empty_json_object() -> Value {
67        Value::Object(serde_json::Map::new())
68    }
69
70    /// Checks if the file is empty.
71    async fn is_file_empty(&self) -> IngestionResult<bool> {
72        Ok(self.file.metadata().await.map(|m| m.len() == 0)?)
73    }
74
75    /// Reads the file content and parses it as a JSON [`Value`].
76    ///
77    /// - If the file is empty, this function *avoids reading the file* and immediately returns an
78    ///   empty JSON object.
79    /// - If the file is not empty, it reads the entire file content and parses into a JSON
80    ///   [`Value`].
81    /// - If JSON parsing fails (indicating a corrupted or invalid JSON file), it also returns an
82    ///   empty JSON object. This ensures that the progress store starts with a clean state in case
83    ///   of file corruption. Later, the `ProgressStore::load` method will interpret this empty JSON
84    ///   object as a default checkpoint sequence number of 0.
85    async fn read_file_to_json_value(&mut self) -> IngestionResult<Value> {
86        if self.is_file_empty().await? {
87            return Ok(Self::empty_json_object());
88        }
89        // before reading seek to the start of the file
90        self.file.seek(SeekFrom::Start(0)).await?;
91        let mut buf = Vec::new();
92        self.file.read_to_end(&mut buf).await?;
93        Ok(serde_json::from_slice::<Value>(buf.as_slice())
94            .inspect_err(|err| tracing::warn!("corrupted or invalid JSON file: {err}"))
95            .unwrap_or_else(|_| Self::empty_json_object()))
96    }
97
98    /// Writes the given data to the file, overwriting any existing content.
99    async fn write_to_file(&mut self, data: impl AsRef<[u8]>) -> IngestionResult<()> {
100        let tmp_path = self.path.with_extension("tmp");
101
102        {
103            let mut tmp_file = File::options()
104                .write(true)
105                .create(true)
106                .truncate(true)
107                .open(&tmp_path)
108                .await?;
109            tmp_file.write_all(data.as_ref()).await?;
110            tmp_file.sync_data().await?;
111
112            // only for testing add a small delay, useful for simulate crashes
113            if cfg!(test) {
114                tokio::time::sleep(std::time::Duration::from_nanos(10)).await;
115            }
116        }
117
118        // Atomically replace the original file
119        tokio::fs::rename(&tmp_path, &self.path).await?;
120
121        // Re-open the file handle for further reads
122        self.file = File::open(&self.path).await?;
123
124        Ok(())
125    }
126}
127
128#[async_trait]
129impl ProgressStore for FileProgressStore {
130    type Error = IngestionError;
131
132    async fn load(&mut self, task_name: String) -> Result<CheckpointSequenceNumber, Self::Error> {
133        let content = self.read_file_to_json_value().await?;
134        Ok(content
135            .get(&task_name)
136            .and_then(|v| v.as_u64())
137            .unwrap_or_default())
138    }
139    async fn save(
140        &mut self,
141        task_name: String,
142        checkpoint_number: CheckpointSequenceNumber,
143    ) -> Result<(), Self::Error> {
144        let mut content = self.read_file_to_json_value().await?;
145        content[task_name] = Value::Number(Number::from(checkpoint_number));
146        self.write_to_file(serde_json::to_string_pretty(&content)?)
147            .await?;
148        Ok(())
149    }
150}
iota_data_ingestion_core/progress_store/file.rs

iota_data_ingestion_core/progress_store/
file.rs
