coven_replication/sync/store/blob/
eager_cache.rs

1use std::sync::atomic::{AtomicU64, Ordering};
2use std::sync::Arc;
3
4use futures_util::StreamExt;
5use tokio::sync::watch;
6use tracing::{info, warn};
7
8use super::{BlobCacheError, RemoteStoreBlobAccess};
9
10#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub struct EagerCacheFillProgress {
12    pub files_done: u64,
13    pub files_total: u64,
14    pub bytes_done: u64,
15    pub bytes_total: u64,
16}
17
18impl EagerCacheFillProgress {
19    pub(crate) fn empty() -> Self {
20        Self {
21            files_done: 0,
22            files_total: 0,
23            bytes_done: 0,
24            bytes_total: 0,
25        }
26    }
27}
28
29#[derive(Debug, Clone)]
30pub enum EagerCacheFillStatus {
31    NotRunning,
32    Scanning,
33    Downloading(EagerCacheFillProgress),
34    Complete {
35        files_total: u64,
36        bytes_total: u64,
37    },
38    Cancelled(EagerCacheFillProgress),
39    Failed {
40        progress: EagerCacheFillProgress,
41        error: Arc<EagerCacheFillError>,
42    },
43}
44
45#[derive(Debug, thiserror::Error)]
46pub enum EagerCacheFillError {
47    #[error("read eager cache bindings: {0}")]
48    Database(#[from] coven_database::DbError),
49    #[error("inspect eager cache row {table}/{row_id}: {source}")]
50    Inspect {
51        table: String,
52        row_id: String,
53        #[source]
54        source: BlobCacheError,
55    },
56    #[error("eager cache byte total exceeds the supported byte count")]
57    ByteCountOverflow,
58    #[error("remote eager cache row {table}/{row_id} has no exact stored reference")]
59    MissingStoredReference { table: String, row_id: String },
60    #[error("eager cache downloads ended before every admitted byte completed")]
61    IncompleteProgress,
62    #[error("download eager cache row {table}/{row_id}: {source}")]
63    Download {
64        table: String,
65        row_id: String,
66        #[source]
67        source: BlobCacheError,
68    },
69}
70
71struct EagerDownload {
72    reference: coven_protocol::blob::RowBlobRef,
73    stored_size: u64,
74}
75
76struct FileProgress {
77    bytes: AtomicU64,
78    total: u64,
79}
80
81impl FileProgress {
82    fn new(total: u64) -> Self {
83        Self {
84            bytes: AtomicU64::new(0),
85            total,
86        }
87    }
88}
89
90fn aggregate_progress(
91    files: &[Arc<FileProgress>],
92    completed_files: &AtomicU64,
93    files_total: u64,
94    bytes_total: u64,
95) -> Result<EagerCacheFillProgress, EagerCacheFillError> {
96    let mut bytes_done = 0_u64;
97    for file in files {
98        let current = file.bytes.load(Ordering::Relaxed);
99        if current > file.total {
100            return Err(EagerCacheFillError::ByteCountOverflow);
101        }
102        bytes_done = bytes_done
103            .checked_add(current)
104            .ok_or(EagerCacheFillError::ByteCountOverflow)?;
105    }
106    Ok(EagerCacheFillProgress {
107        files_done: completed_files.load(Ordering::Relaxed),
108        files_total,
109        bytes_done,
110        bytes_total,
111    })
112}
113
114fn fail(
115    status: &watch::Sender<EagerCacheFillStatus>,
116    progress: EagerCacheFillProgress,
117    error: EagerCacheFillError,
118) -> Arc<EagerCacheFillError> {
119    let error = Arc::new(error);
120    status.send_replace(EagerCacheFillStatus::Failed {
121        progress,
122        error: Arc::clone(&error),
123    });
124    error
125}
126
127pub(crate) async fn run(
128    database: &coven_database::StoreDatabase,
129    access: &RemoteStoreBlobAccess,
130    mut cancel: watch::Receiver<bool>,
131    status: &watch::Sender<EagerCacheFillStatus>,
132) -> Result<(), Arc<EagerCacheFillError>> {
133    status.send_replace(EagerCacheFillStatus::Scanning);
134    if *cancel.borrow() {
135        status.send_replace(EagerCacheFillStatus::Cancelled(
136            EagerCacheFillProgress::empty(),
137        ));
138        return Ok(());
139    }
140
141    let references = match database.eager_row_blob_refs().await {
142        Ok(references) => references,
143        Err(error) => return Err(fail(status, EagerCacheFillProgress::empty(), error.into())),
144    };
145    let mut downloads = Vec::new();
146    for reference in references {
147        if !matches!(
148            reference.authority(),
149            coven_protocol::blob::RowBlobAuthority::Remote(_)
150        ) {
151            continue;
152        }
153        let materialized = match access.is_materialized(&reference).await {
154            Ok(materialized) => materialized,
155            Err(source) => {
156                let error = EagerCacheFillError::Inspect {
157                    table: reference.table().to_string(),
158                    row_id: reference.row_id().to_string(),
159                    source,
160                };
161                return Err(fail(status, EagerCacheFillProgress::empty(), error));
162            }
163        };
164        if materialized {
165            continue;
166        }
167        let Some(stored) = reference.stored() else {
168            let error = EagerCacheFillError::MissingStoredReference {
169                table: reference.table().to_string(),
170                row_id: reference.row_id().to_string(),
171            };
172            return Err(fail(status, EagerCacheFillProgress::empty(), error));
173        };
174        let stored_size = stored.object().stored_size();
175        downloads.push(EagerDownload {
176            reference,
177            stored_size,
178        });
179    }
180
181    let files_total = downloads.len() as u64;
182    let bytes_total = match downloads.iter().try_fold(0_u64, |total, download| {
183        total.checked_add(download.stored_size)
184    }) {
185        Some(total) => total,
186        None => {
187            return Err(fail(
188                status,
189                EagerCacheFillProgress::empty(),
190                EagerCacheFillError::ByteCountOverflow,
191            ))
192        }
193    };
194    if downloads.is_empty() {
195        status.send_replace(EagerCacheFillStatus::Complete {
196            files_total,
197            bytes_total,
198        });
199        return Ok(());
200    }
201
202    let file_progress = downloads
203        .iter()
204        .map(|download| Arc::new(FileProgress::new(download.stored_size)))
205        .collect::<Vec<_>>();
206    let completed_files = Arc::new(AtomicU64::new(0));
207    let initial = EagerCacheFillProgress {
208        files_done: 0,
209        files_total,
210        bytes_done: 0,
211        bytes_total,
212    };
213    status.send_replace(EagerCacheFillStatus::Downloading(initial));
214
215    let limit = database.transfer_limits().downloads.get();
216    let stream = futures_util::stream::iter(downloads.into_iter().enumerate())
217        .map(|(index, download)| {
218            let progress = Arc::clone(&file_progress[index]);
219            let completed_files = Arc::clone(&completed_files);
220            async move {
221                let callback_progress = Arc::clone(&progress);
222                let callback: coven_storage::cloud::DownloadProgress =
223                    Arc::new(move |bytes_done| {
224                        callback_progress.bytes.store(bytes_done, Ordering::Relaxed);
225                    });
226                let table = download.reference.table().to_string();
227                let row_id = download.reference.row_id().to_string();
228                match access
229                    .materialize_with_progress(&download.reference, callback)
230                    .await
231                {
232                    Ok(()) => {
233                        progress.bytes.store(progress.total, Ordering::Relaxed);
234                        completed_files.fetch_add(1, Ordering::Relaxed);
235                        Ok(())
236                    }
237                    Err(source) => Err(EagerCacheFillError::Download {
238                        table,
239                        row_id,
240                        source,
241                    }),
242                }
243            }
244        })
245        .buffer_unordered(limit);
246    tokio::pin!(stream);
247    let mut ticker = tokio::time::interval(crate::blob::progress::TRANSFER_PROGRESS_TICK);
248    ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
249    ticker.tick().await;
250    let mut last_reported = initial;
251
252    loop {
253        tokio::select! {
254            next = stream.next() => {
255                match next {
256                    Some(Ok(())) => {}
257                    Some(Err(error)) => {
258                        let progress = aggregate_progress(
259                            &file_progress,
260                            &completed_files,
261                            files_total,
262                            bytes_total,
263                        ).map_err(|error| fail(status, last_reported, error))?;
264                        return Err(fail(status, progress, error));
265                    }
266                    None => break,
267                }
268            }
269            _ = ticker.tick() => {
270                let current = aggregate_progress(
271                    &file_progress,
272                    &completed_files,
273                    files_total,
274                    bytes_total,
275                ).map_err(|error| fail(status, last_reported, error))?;
276                if current != last_reported {
277                    last_reported = current;
278                    status.send_replace(EagerCacheFillStatus::Downloading(current));
279                }
280            }
281            changed = cancel.changed() => {
282                if changed.is_err() || *cancel.borrow() {
283                    let current = aggregate_progress(
284                        &file_progress,
285                        &completed_files,
286                        files_total,
287                        bytes_total,
288                    ).map_err(|error| fail(status, last_reported, error))?;
289                    status.send_replace(EagerCacheFillStatus::Cancelled(current));
290                    info!(files_done = current.files_done, files_total, "eager cache fill cancelled");
291                    return Ok(());
292                }
293            }
294        }
295    }
296
297    let completed = aggregate_progress(&file_progress, &completed_files, files_total, bytes_total)
298        .map_err(|error| fail(status, last_reported, error))?;
299    if completed.files_done != files_total || completed.bytes_done != bytes_total {
300        let error = EagerCacheFillError::IncompleteProgress;
301        warn!(
302            ?completed,
303            "eager cache fill finished with incomplete counters"
304        );
305        return Err(fail(status, completed, error));
306    }
307    status.send_replace(EagerCacheFillStatus::Complete {
308        files_total,
309        bytes_total,
310    });
311    info!(files_total, bytes_total, "eager cache fill complete");
312    Ok(())
313}
coven_replication/sync/store/blob/eager_cache.rs

coven_replication/sync/store/blob/
eager_cache.rs
