xet_data/processing/
file_cleaner.rs

1use std::future::{self, Future};
2use std::pin::Pin;
3use std::sync::Arc;
4
5use bytes::Bytes;
6use chrono::{DateTime, Utc};
7#[cfg(target_family = "wasm")]
8use tokio_with_wasm::alias as tokio;
9use tracing::{Instrument, debug_span, info, instrument};
10use xet_core_structures::merklehash::ChunkHashList;
11use xet_core_structures::metadata_shard::Sha256;
12use xet_core_structures::metadata_shard::file_structs::{FileMetadataExt, MDBFileInfo};
13use xet_runtime::core::XetContext;
14
15use super::XetFileInfo;
16use super::deduplication_interface::UploadSessionDataManager;
17use super::file_upload_session::FileUploadSession;
18use super::sha256::Sha256Generator;
19use crate::deduplication::{Chunk, Chunker, DeduplicationMetrics, FileDeduper};
20use crate::error::Result;
21use crate::progress_tracking::upload_tracking::CompletionTrackerFileId;
22
23/// Controls how SHA-256 is handled during file cleaning.
24#[derive(Clone, Copy)]
25pub enum Sha256Policy {
26    /// Compute SHA-256 from the file data.
27    Compute,
28    /// Use a pre-computed SHA-256 value.
29    Provided(Sha256),
30    /// Skip SHA-256 entirely; no metadata_ext is written to the shard.
31    Skip,
32}
33
34impl Sha256Policy {
35    /// Returns `Skip` when `true`, `Compute` when `false`.
36    pub fn from_skip(skip: bool) -> Self {
37        if skip { Self::Skip } else { Self::Compute }
38    }
39
40    /// Parses a hex-encoded SHA-256 string into a policy.
41    ///
42    /// Returns `Provided(hash)` if the hex is valid, `Compute` otherwise.
43    pub fn from_hex(hex: &str) -> Self {
44        Sha256::from_hex(hex).ok().into()
45    }
46}
47
48impl From<Option<Sha256>> for Sha256Policy {
49    fn from(sha256: Option<Sha256>) -> Self {
50        match sha256 {
51            Some(hash) => Self::Provided(hash),
52            None => Self::Compute,
53        }
54    }
55}
56
57/// A class that encapsulates the clean and data task around a single file.
58pub struct SingleFileCleaner {
59    ctx: XetContext,
60
61    // File name, if known.
62    file_name: Option<Arc<str>>,
63
64    // Completion id
65    file_id: CompletionTrackerFileId,
66
67    // Common state.
68    session: Arc<FileUploadSession>,
69
70    // The chunker.
71    chunker: Chunker,
72
73    // The deduplication interface.  Use a future that always returns the dedup manager
74    // on await so that we can background this part.
75    dedup_manager_fut: Pin<Box<dyn Future<Output = Result<FileDeduper<UploadSessionDataManager>>> + Send + 'static>>,
76
77    // SHA-256 generator, present only when computing from file data.
78    sha_generator: Option<Sha256Generator>,
79
80    // Pre-computed or finalized SHA-256 value.
81    provided_sha256: Option<Sha256>,
82
83    // Start time
84    start_time: DateTime<Utc>,
85}
86
87impl SingleFileCleaner {
88    pub(crate) fn new(
89        file_name: Option<Arc<str>>,
90        file_id: CompletionTrackerFileId,
91        sha256: Sha256Policy,
92        session: Arc<FileUploadSession>,
93    ) -> Self {
94        let ctx = session.ctx.clone();
95        let deduper = FileDeduper::new(UploadSessionDataManager::new(session.clone()), file_id, ctx.clone());
96
97        let (sha_generator, provided_sha256) = match sha256 {
98            Sha256Policy::Compute => (Some(Sha256Generator::new(ctx.clone())), None),
99            Sha256Policy::Provided(hash) => (None, Some(hash)),
100            Sha256Policy::Skip => (None, None),
101        };
102
103        Self {
104            ctx,
105            file_name,
106            file_id,
107            dedup_manager_fut: Box::pin(async move { Ok(deduper) }),
108            session,
109            chunker: crate::deduplication::Chunker::default(),
110            sha_generator,
111            provided_sha256,
112            start_time: Utc::now(),
113        }
114    }
115
116    /// Gets the dedupe manager to process new chunks, by first
117    /// waiting for background operations to complete, then triggering a
118    /// new background task.
119    async fn deduper_process_chunks(&mut self, chunks: Arc<[Chunk]>) -> Result<()> {
120        // Handle the move out by replacing it with a dummy future discarded below.
121        let mut deduper = std::mem::replace(&mut self.dedup_manager_fut, Box::pin(future::pending())).await?;
122
123        let num_chunks = chunks.len();
124
125        let dedup_background = tokio::spawn(
126            async move {
127                deduper.process_chunks(&chunks).await?;
128                Ok(deduper)
129            }
130            .instrument(debug_span!("deduper::process_chunks_task", num_chunks).or_current()),
131        );
132
133        self.dedup_manager_fut = Box::pin(async move { dedup_background.await? });
134
135        Ok(())
136    }
137
138    pub async fn add_data(&mut self, data: &[u8]) -> Result<()> {
139        self.add_data_from_bytes(Bytes::copy_from_slice(data)).await
140    }
141
142    pub async fn add_data_from_bytes(&mut self, data: Bytes) -> Result<()> {
143        let block_size = usize::try_from(*self.ctx.config.data.ingestion_block_size)
144            .expect("ingestion_block_size exceeds usize::MAX on this target");
145        if data.len() > block_size {
146            let mut pos = 0;
147            while pos < data.len() {
148                let next_pos = usize::min(pos + block_size, data.len());
149                self.add_data_chunk_impl(data.slice(pos..next_pos)).await?;
150                pos = next_pos;
151            }
152        } else {
153            self.add_data_chunk_impl(data).await?;
154        }
155
156        Ok(())
157    }
158
159    #[instrument(skip_all, level="debug", name = "FileCleaner::add_data", fields(file_name=self.file_name.as_ref().map(|s|s.to_string()), len=data.len()))]
160    async fn add_data_chunk_impl(&mut self, data: Bytes) -> Result<()> {
161        // If the file size was not specified at the beginning, then incrementally update tho total size with
162        // how much data we know about.
163        self.session
164            .completion_tracker
165            .increment_file_size(self.file_id, data.len() as u64);
166
167        // Put the chunking on a compute thread so it doesn't tie up the async schedulers
168        let chunk_data_jh = {
169            let mut chunker = std::mem::take(&mut self.chunker);
170            let data = data.clone();
171            let runtime = self.ctx.runtime.clone();
172
173            runtime.spawn_blocking(move || {
174                let chunks: Arc<[Chunk]> = Arc::from(chunker.next_block_bytes(&data, false));
175                (chunks, chunker)
176            })
177        };
178
179        // Update the sha256 hasher, which hands this off to be done in the background.
180        if let Some(ref mut generator) = self.sha_generator {
181            generator.update(data.clone()).await?;
182        }
183
184        // Get the chunk data and start processing it.
185        let (chunks, chunker) = chunk_data_jh.await?;
186
187        // Restore the chunker state.
188        self.chunker = chunker;
189
190        // It's possible this didn't actually add any data in.
191        if chunks.is_empty() {
192            return Ok(());
193        }
194
195        // Run the deduplication interface here.
196        self.deduper_process_chunks(chunks).await?;
197
198        Ok(())
199    }
200
201    /// Ensures all current background work is completed.
202    pub async fn checkpoint(&mut self) -> Result<()> {
203        // Flush the background process by sending it a dummy bit of data.
204        self.deduper_process_chunks(Arc::new([])).await
205    }
206
207    /// Return the representation of the file after clean as a pointer file instance.
208    pub async fn finish(self) -> Result<(XetFileInfo, DeduplicationMetrics)> {
209        let (info, _chunks, metrics) = self.finish_with_chunks().await?;
210        Ok((info, metrics))
211    }
212
213    /// Same as [`finish`], but also returns the per-chunk hash list produced during CDC.
214    /// Only needed by composition flows (e.g. `upload_ranges`) that build partial
215    /// `MerkleHashSubtree` nodes for newly-uploaded windows; regular uploads should call
216    /// [`finish`] instead.
217    #[instrument(skip_all, name = "FileCleaner::finish_with_chunks", fields(file_name=self.file_name.as_ref().map(|s|s.to_string())))]
218    pub async fn finish_with_chunks(self) -> Result<(XetFileInfo, ChunkHashList, DeduplicationMetrics)> {
219        let (file_info, chunk_hashes, _, deduplication_metrics) = Self::finish_inner(self, true).await?;
220        Ok((file_info, chunk_hashes, deduplication_metrics))
221    }
222
223    /// Like `finish_with_chunks`, but does NOT register the file's MDBFileInfo in the
224    /// session shard. Returns the MDBFileInfo directly so the caller can compose it into a
225    /// larger file without creating orphan shard entries.
226    pub async fn finish_with_chunks_detached(
227        self,
228    ) -> Result<(XetFileInfo, ChunkHashList, MDBFileInfo, DeduplicationMetrics)> {
229        Self::finish_inner(self, false).await
230    }
231
232    async fn finish_inner(
233        mut self,
234        register: bool,
235    ) -> Result<(XetFileInfo, ChunkHashList, MDBFileInfo, DeduplicationMetrics)> {
236        if let Some(chunk) = self.chunker.finish() {
237            let data = Arc::new([chunk]);
238            self.deduper_process_chunks(data).await?;
239        }
240
241        let sha256 = if let Some(generator) = self.sha_generator.take() {
242            Some(generator.finalize().await?)
243        } else {
244            self.provided_sha256
245        };
246        let metadata_ext = sha256.map(FileMetadataExt::new);
247
248        let (file_hash, chunk_hashes, remaining_file_data, deduplication_metrics) =
249            self.dedup_manager_fut.await?.finalize(metadata_ext);
250
251        let file_info = XetFileInfo {
252            hash: file_hash.hex(),
253            file_size: Some(deduplication_metrics.total_bytes),
254            sha256: sha256.map(|s| s.hex()),
255        };
256
257        #[cfg(debug_assertions)]
258        {
259            debug_assert_eq!(remaining_file_data.pending_file_info.len(), 1);
260            debug_assert_eq!(remaining_file_data.pending_file_info[0].0.file_size(), deduplication_metrics.total_bytes)
261        }
262
263        let mdb_file_info = if register {
264            self.session
265                .register_single_file_clean_completion(remaining_file_data, &deduplication_metrics)
266                .await?;
267            MDBFileInfo::default()
268        } else {
269            self.session
270                .register_single_file_clean_completion_detached(remaining_file_data, &deduplication_metrics)
271                .await?
272        };
273
274        info!(
275            target: "client_telemetry",
276            action = "clean",
277            file_name = self.file_name.as_deref().unwrap_or_default().to_string(),
278            file_size_count = deduplication_metrics.total_bytes,
279            new_bytes_count = deduplication_metrics.new_bytes,
280            start_ts = self.start_time.to_rfc3339(),
281            end_processing_ts = Utc::now().to_rfc3339(),
282        );
283
284        Ok((file_info, chunk_hashes, mdb_file_info, deduplication_metrics))
285    }
286}
xet_data/processing/file_cleaner.rs

xet_data/processing/
file_cleaner.rs
