1use crate::object_store::{
5 ObjectStoreDeleteExt, ObjectStoreGetExt, ObjectStoreListExt, ObjectStorePutExt,
6};
7use anyhow::{Context, Result, anyhow};
8use backoff::ExponentialBackoff;
9use backoff::future::retry;
10use bytes::Bytes;
11use futures::StreamExt;
12use futures::TryStreamExt;
13use indicatif::ProgressBar;
14use itertools::Itertools;
15use mysten_common::ZipDebugEqIteratorExt;
16use object_store::aws::{AmazonS3Builder, AmazonS3ConfigKey};
17use object_store::gcp::{GoogleCloudStorageBuilder, GoogleConfigKey};
18use object_store::http::HttpBuilder;
19use object_store::local::LocalFileSystem;
20use object_store::path::Path;
21use object_store::{
22 ClientOptions, DynObjectStore, Error, ObjectStore, ObjectStoreExt, RetryConfig,
23};
24use prost::Message;
25use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
26use serde::{Deserialize, Serialize};
27use serde_json;
28use std::collections::BTreeMap;
29use std::num::NonZeroUsize;
30use std::ops::Range;
31use std::path::PathBuf;
32use std::str::FromStr;
33use std::sync::Arc;
34use std::time::Duration;
35use sui_rpc::proto::sui::rpc::v2 as proto;
36use sui_types::full_checkpoint_content::Checkpoint;
37use sui_types::messages_checkpoint::CheckpointSequenceNumber;
38use tracing::{error, warn};
39use url::Url;
40
41pub const MANIFEST_FILENAME: &str = "MANIFEST";
42
43#[derive(Serialize, Deserialize)]
44
45pub struct Manifest {
46 pub available_epochs: Vec<u64>,
47}
48
49impl Manifest {
50 pub fn new(available_epochs: Vec<u64>) -> Self {
51 Manifest { available_epochs }
52 }
53
54 pub fn epoch_exists(&self, epoch: u64) -> bool {
55 self.available_epochs.contains(&epoch)
56 }
57}
58
59#[derive(Debug, Clone)]
60pub struct PerEpochManifest {
61 pub lines: Vec<String>,
62}
63
64impl PerEpochManifest {
65 pub fn new(lines: Vec<String>) -> Self {
66 PerEpochManifest { lines }
67 }
68
69 pub fn serialize_as_newline_delimited(&self) -> String {
70 self.lines.join("\n")
71 }
72
73 pub fn deserialize_from_newline_delimited(s: &str) -> PerEpochManifest {
74 PerEpochManifest {
75 lines: s.lines().map(String::from).collect(),
76 }
77 }
78
79 pub fn filter_by_prefix(&self, prefix: &str) -> PerEpochManifest {
81 let filtered_lines = self
82 .lines
83 .iter()
84 .filter(|line| line.starts_with(prefix))
85 .cloned()
86 .collect();
87
88 PerEpochManifest {
89 lines: filtered_lines,
90 }
91 }
92}
93
94pub async fn get<S: ObjectStoreGetExt>(store: &S, src: &Path) -> Result<Bytes> {
95 let bytes = retry(backoff::ExponentialBackoff::default(), || async {
96 store.get_bytes(src).await.map_err(|e| {
97 error!("Failed to read file from object store with error: {:?}", &e);
98 backoff::Error::transient(e)
99 })
100 })
101 .await?;
102 Ok(bytes)
103}
104
105pub async fn exists<S: ObjectStoreGetExt>(store: &S, src: &Path) -> bool {
106 store.get_bytes(src).await.is_ok()
107}
108
109pub async fn put<S: ObjectStorePutExt>(store: &S, src: &Path, bytes: Bytes) -> Result<()> {
110 retry(backoff::ExponentialBackoff::default(), || async {
111 if !bytes.is_empty() {
112 store.put_bytes(src, bytes.clone()).await.map_err(|e| {
113 error!("Failed to write file to object store with error: {:?}", &e);
114 backoff::Error::transient(e)
115 })
116 } else {
117 warn!("Not copying empty file: {:?}", src);
118 Ok(())
119 }
120 })
121 .await?;
122 Ok(())
123}
124
125pub async fn copy_file<S: ObjectStoreGetExt, D: ObjectStorePutExt>(
126 src: &Path,
127 dest: &Path,
128 src_store: &S,
129 dest_store: &D,
130) -> Result<()> {
131 let bytes = get(src_store, src).await?;
132 if !bytes.is_empty() {
133 put(dest_store, dest, bytes).await
134 } else {
135 warn!("Not copying empty file: {:?}", src);
136 Ok(())
137 }
138}
139
140pub async fn copy_files<S: ObjectStoreGetExt, D: ObjectStorePutExt>(
141 src: &[Path],
142 dest: &[Path],
143 src_store: &S,
144 dest_store: &D,
145 concurrency: NonZeroUsize,
146 progress_bar: Option<ProgressBar>,
147) -> Result<Vec<()>> {
148 futures::stream::iter(src.iter().zip_debug_eq(dest.iter()))
149 .map(|(path_in, path_out)| {
150 let progress_bar = progress_bar.clone();
151 async move {
152 copy_file(path_in, path_out, src_store, dest_store)
153 .await
154 .with_context(|| format!("Failed to copy {path_in} to {path_out}"))?;
155 if let Some(progress_bar) = progress_bar {
156 progress_bar.inc(1);
157 progress_bar.set_message(format!("file: {path_out}"));
158 }
159 Ok(())
160 }
161 })
162 .boxed()
163 .buffer_unordered(concurrency.get())
164 .try_collect()
165 .await
166}
167
168pub async fn copy_recursively<S: ObjectStoreGetExt + ObjectStoreListExt, D: ObjectStorePutExt>(
169 dir: &Path,
170 src_store: &S,
171 dest_store: &D,
172 concurrency: NonZeroUsize,
173) -> Result<Vec<()>> {
174 let mut input_paths = vec![];
175 let mut output_paths = vec![];
176 let mut paths = src_store.list_objects(Some(dir)).await;
177 while let Some(res) = paths.next().await {
178 if let Ok(object_metadata) = res {
179 input_paths.push(object_metadata.location.clone());
180 output_paths.push(object_metadata.location);
181 } else {
182 return Err(res.err().unwrap().into());
183 }
184 }
185 copy_files(
186 &input_paths,
187 &output_paths,
188 src_store,
189 dest_store,
190 concurrency,
191 None,
192 )
193 .await
194}
195
196pub async fn delete_files<S: ObjectStoreDeleteExt>(
197 files: &[Path],
198 store: &S,
199 concurrency: NonZeroUsize,
200) -> Result<Vec<()>> {
201 let results: Vec<Result<()>> = futures::stream::iter(files)
202 .map(|f| {
203 retry(backoff::ExponentialBackoff::default(), || async {
204 store.delete_object(f).await.map_err(|e| {
205 error!("Failed to delete file on object store with error: {:?}", &e);
206 backoff::Error::transient(e)
207 })
208 })
209 })
210 .boxed()
211 .buffer_unordered(concurrency.get())
212 .collect()
213 .await;
214 results.into_iter().collect()
215}
216
217pub async fn delete_recursively<S: ObjectStoreDeleteExt + ObjectStoreListExt>(
218 path: &Path,
219 store: &S,
220 concurrency: NonZeroUsize,
221) -> Result<Vec<()>> {
222 let mut paths_to_delete = vec![];
223 let mut paths = store.list_objects(Some(path)).await;
224 while let Some(res) = paths.next().await {
225 if let Ok(object_metadata) = res {
226 paths_to_delete.push(object_metadata.location);
227 } else {
228 return Err(res.err().unwrap().into());
229 }
230 }
231 delete_files(&paths_to_delete, store, concurrency).await
232}
233
234pub fn path_to_filesystem(local_dir_path: PathBuf, location: &Path) -> anyhow::Result<PathBuf> {
235 let path = std::fs::canonicalize(local_dir_path)?;
237 let mut url = Url::from_file_path(&path)
238 .map_err(|_| anyhow!("Failed to parse input path: {}", path.display()))?;
239 url.path_segments_mut()
240 .map_err(|_| anyhow!("Failed to get path segments: {}", path.display()))?
241 .pop_if_empty()
242 .extend(location.parts());
243 let new_path = url
244 .to_file_path()
245 .map_err(|_| anyhow!("Failed to convert url to path: {}", url.as_str()))?;
246 Ok(new_path)
247}
248
249pub async fn find_all_dirs_with_epoch_prefix(
252 store: &Arc<DynObjectStore>,
253 prefix: Option<&Path>,
254) -> anyhow::Result<BTreeMap<u64, Path>> {
255 let mut dirs = BTreeMap::new();
256 let entries = store.list_with_delimiter(prefix).await?;
257 for entry in entries.common_prefixes {
258 if let Some(filename) = entry.filename() {
259 if !filename.starts_with("epoch_") || filename.ends_with(".tmp") {
260 continue;
261 }
262 let epoch = filename
263 .split_once('_')
264 .context("Failed to split dir name")
265 .map(|(_, epoch)| epoch.parse::<u64>())??;
266 dirs.insert(epoch, entry);
267 }
268 }
269 Ok(dirs)
270}
271
272pub async fn list_all_epochs(object_store: Arc<DynObjectStore>) -> Result<Vec<u64>> {
273 let remote_epoch_dirs = find_all_dirs_with_epoch_prefix(&object_store, None).await?;
274 let mut out = vec![];
275 let mut success_marker_found = false;
276 for (epoch, path) in remote_epoch_dirs.iter().sorted() {
277 let success_marker = path.child("_SUCCESS");
278 let get_result = object_store.get(&success_marker).await;
279 match get_result {
280 Err(_) => {
281 if !success_marker_found {
282 error!("No success marker found for epoch: {epoch}");
283 }
284 }
285 Ok(_) => {
286 out.push(*epoch);
287 success_marker_found = true;
288 }
289 }
290 }
291
292 let archive_prefix = Path::from("archive");
294 if let Ok(archive_epoch_dirs) =
295 find_all_dirs_with_epoch_prefix(&object_store, Some(&archive_prefix)).await
296 {
297 for (epoch, path) in archive_epoch_dirs.iter().sorted() {
298 let success_marker = path.child("_SUCCESS");
299 let get_result = object_store.get(&success_marker).await;
300 if get_result.is_ok() && !out.contains(epoch) {
301 out.push(*epoch);
302 }
303 }
304 }
305
306 Ok(out)
307}
308
309pub async fn run_manifest_update_loop(
310 store: Arc<DynObjectStore>,
311 mut recv: tokio::sync::broadcast::Receiver<()>,
312) -> Result<()> {
313 let mut update_interval = tokio::time::interval(Duration::from_secs(300));
314 loop {
315 tokio::select! {
316 _now = update_interval.tick() => {
317 if let Ok(epochs) = list_all_epochs(store.clone()).await {
318 let manifest_path = Path::from(MANIFEST_FILENAME);
319 let manifest = Manifest::new(epochs);
320 let bytes = serde_json::to_string(&manifest)?;
321 put(&store, &manifest_path, Bytes::from(bytes)).await?;
322 }
323 },
324 _ = recv.recv() => break,
325 }
326 }
327 Ok(())
328}
329
330pub async fn find_all_files_with_epoch_prefix(
333 store: &Arc<DynObjectStore>,
334 prefix: Option<&Path>,
335) -> anyhow::Result<Vec<Range<u64>>> {
336 let mut ranges = Vec::new();
337 let entries = store.list_with_delimiter(prefix).await?;
338 for entry in entries.objects {
339 let checkpoint_seq_range = entry
340 .location
341 .filename()
342 .ok_or(anyhow!("Illegal file name"))?
343 .split_once('.')
344 .context("Failed to split dir name")?
345 .0
346 .split_once('_')
347 .context("Failed to split dir name")
348 .map(|(start, end)| Range {
349 start: start.parse::<u64>().unwrap(),
350 end: end.parse::<u64>().unwrap(),
351 })?;
352
353 ranges.push(checkpoint_seq_range);
354 }
355 Ok(ranges)
356}
357
358pub async fn find_missing_epochs_dirs(
364 store: &Arc<DynObjectStore>,
365 success_marker: &str,
366) -> anyhow::Result<Vec<u64>> {
367 let remote_checkpoints_by_epoch = find_all_dirs_with_epoch_prefix(store, None).await?;
368 let mut dirs: Vec<_> = remote_checkpoints_by_epoch.iter().collect();
369 dirs.sort_by_key(|(epoch_num, _path)| *epoch_num);
370 let mut candidate_epoch: u64 = 0;
371 let mut missing_epochs = Vec::new();
372 for (epoch_num, path) in dirs {
373 while candidate_epoch < *epoch_num {
374 missing_epochs.push(candidate_epoch);
376 candidate_epoch += 1;
377 continue;
378 }
379 let success_marker = path.child(success_marker);
380 let get_result = store.get(&success_marker).await;
381 match get_result {
382 Err(Error::NotFound { .. }) => {
383 error!("No success marker found in db checkpoint for epoch: {epoch_num}");
384 missing_epochs.push(*epoch_num);
385 }
386 Err(_) => {
387 warn!(
389 "Failed while trying to read success marker in db checkpoint for epoch: {epoch_num}"
390 );
391 }
392 Ok(_) => {
393 }
395 }
396 candidate_epoch += 1
397 }
398 missing_epochs.push(candidate_epoch);
399 Ok(missing_epochs)
400}
401
402pub fn get_path(prefix: &str) -> Path {
403 Path::from(prefix)
404}
405
406pub async fn write_snapshot_manifest<S: ObjectStoreListExt + ObjectStorePutExt>(
409 dir: &Path,
410 store: &S,
411 epoch_prefix: String,
412) -> Result<()> {
413 let mut file_names = vec![];
414 let mut paths = store.list_objects(Some(dir)).await;
415 while let Some(res) = paths.next().await {
416 if let Ok(object_metadata) = res {
417 let mut path_str = object_metadata.location.to_string();
419 if path_str.starts_with(&epoch_prefix) {
420 path_str = String::from(&path_str[epoch_prefix.len()..]);
421 file_names.push(path_str);
422 } else {
423 warn!("{path_str}, should be coming from the files in the {epoch_prefix} dir",)
424 }
425 } else {
426 return Err(res.err().unwrap().into());
427 }
428 }
429
430 let epoch_manifest = PerEpochManifest::new(file_names);
431 let bytes = Bytes::from(epoch_manifest.serialize_as_newline_delimited());
432 put(
433 store,
434 &Path::from(format!("{}/{}", dir, MANIFEST_FILENAME)),
435 bytes,
436 )
437 .await?;
438
439 Ok(())
440}
441
442pub fn build_object_store(
443 ingestion_url: &str,
444 remote_store_options: Vec<(String, String)>,
445 remote_store_headers: Vec<(String, String)>,
446) -> Arc<dyn ObjectStore> {
447 let timeout_secs = 5;
448 let mut client_options = ClientOptions::new()
449 .with_timeout(Duration::from_secs(timeout_secs))
450 .with_allow_http(true);
451 if !remote_store_headers.is_empty() {
452 let mut headers = HeaderMap::new();
453 for (name, value) in &remote_store_headers {
454 headers.insert(
455 HeaderName::from_bytes(name.as_bytes()).expect("invalid remote store header name"),
456 HeaderValue::from_str(value).expect("invalid remote store header value"),
457 );
458 }
459 client_options = client_options.with_default_headers(headers);
460 }
461 let retry_config = RetryConfig {
462 max_retries: 10,
463 retry_timeout: Duration::from_secs(timeout_secs + 1),
464 ..Default::default()
465 };
466 let url = ingestion_url
467 .parse::<Url>()
468 .expect("archival ingestion url must be valid");
469 if url.scheme() == "file" {
470 Arc::new(
471 LocalFileSystem::new_with_prefix(
472 url.to_file_path()
473 .expect("archival ingestion url must have a valid file path"),
474 )
475 .expect("failed to create local file system store"),
476 )
477 } else if url.scheme() == "gs" {
478 let mut builder = GoogleCloudStorageBuilder::new()
479 .with_client_options(client_options)
480 .with_retry(retry_config)
481 .with_url(ingestion_url);
482 for (key, value) in &remote_store_options {
483 builder = builder.with_config(
484 GoogleConfigKey::from_str(key).expect("invalid GCS config key"),
485 value.clone(),
486 );
487 }
488 Arc::new(builder.build().expect("failed to build GCS store"))
489 } else if url.host_str().unwrap_or_default().starts_with("s3") {
490 let mut builder = AmazonS3Builder::new()
491 .with_client_options(client_options)
492 .with_retry(retry_config)
493 .with_imdsv1_fallback()
494 .with_url(ingestion_url);
495 for (key, value) in &remote_store_options {
496 builder = builder.with_config(
497 AmazonS3ConfigKey::from_str(key).expect("invalid S3 config key"),
498 value.clone(),
499 );
500 }
501 Arc::new(builder.build().expect("failed to build S3 store"))
502 } else {
503 Arc::new(
504 HttpBuilder::new()
505 .with_url(url.to_string())
506 .with_client_options(client_options)
507 .with_retry(retry_config)
508 .build()
509 .expect("failed to build HTTP store"),
510 )
511 }
512}
513
514pub async fn fetch_checkpoint(
515 store: &Arc<dyn ObjectStore>,
516 seq: u64,
517) -> anyhow::Result<Checkpoint> {
518 let store = store.clone();
519 let request = move || {
520 let store = store.clone();
521 async move {
522 use backoff::Error as BE;
523 let path = Path::from(format!("{seq}.binpb.zst"));
524 let bytes = store
525 .get(&path)
526 .await
527 .map_err(|e| match e {
528 object_store::Error::NotFound { .. } => {
529 BE::permanent(anyhow!("Checkpoint {seq} not found in archive"))
530 }
531 e => BE::transient(anyhow::Error::from(e)),
532 })?
533 .bytes()
534 .await
535 .map_err(|e| BE::transient(anyhow::Error::from(e)))?;
536 let decompressed =
537 zstd::decode_all(&bytes[..]).map_err(|e| BE::transient(anyhow::Error::from(e)))?;
538 let proto_checkpoint = proto::Checkpoint::decode(&decompressed[..])
539 .map_err(|e| BE::transient(anyhow::Error::from(e)))?;
540 Checkpoint::try_from(&proto_checkpoint).map_err(|e| BE::transient(anyhow!(e)))
541 }
542 };
543 let backoff = ExponentialBackoff {
544 max_elapsed_time: Some(Duration::from_secs(60)),
545 multiplier: 1.0,
546 ..Default::default()
547 };
548 backoff::future::retry(backoff, request).await
549}
550
551pub async fn end_of_epoch_data(
552 url: &str,
553 remote_store_options: Vec<(String, String)>,
554) -> anyhow::Result<Vec<CheckpointSequenceNumber>> {
555 let store = build_object_store(url, remote_store_options, vec![]);
556 let response = store.get(&Path::from("epochs.json")).await?;
557 let bytes = response.bytes().await?;
558 Ok(serde_json::from_slice(&bytes)?)
559}
560
561#[cfg(test)]
562mod tests {
563 use crate::object_store::util::{
564 MANIFEST_FILENAME, copy_recursively, delete_recursively, write_snapshot_manifest,
565 };
566 use object_store::path::Path;
567 use std::fs;
568 use std::num::NonZeroUsize;
569 use sui_config::object_storage_config::{ObjectStoreConfig, ObjectStoreType};
570 use tempfile::TempDir;
571
572 #[tokio::test]
573 pub async fn test_copy_recursively() -> anyhow::Result<()> {
574 let input = TempDir::new()?;
575 let input_path = input.path();
576 let child = input_path.join("child");
577 fs::create_dir(&child)?;
578 let file1 = child.join("file1");
579 fs::write(file1, b"Lorem ipsum")?;
580 let grandchild = child.join("grand_child");
581 fs::create_dir(&grandchild)?;
582 let file2 = grandchild.join("file2");
583 fs::write(file2, b"Lorem ipsum")?;
584
585 let output = TempDir::new()?;
586 let output_path = output.path();
587
588 let input_store = ObjectStoreConfig {
589 object_store: Some(ObjectStoreType::File),
590 directory: Some(input_path.to_path_buf()),
591 ..Default::default()
592 }
593 .make()?;
594
595 let output_store = ObjectStoreConfig {
596 object_store: Some(ObjectStoreType::File),
597 directory: Some(output_path.to_path_buf()),
598 ..Default::default()
599 }
600 .make()?;
601
602 copy_recursively(
603 &Path::from("child"),
604 &input_store,
605 &output_store,
606 NonZeroUsize::new(1).unwrap(),
607 )
608 .await?;
609
610 assert!(output_path.join("child").exists());
611 assert!(output_path.join("child").join("file1").exists());
612 assert!(output_path.join("child").join("grand_child").exists());
613 assert!(
614 output_path
615 .join("child")
616 .join("grand_child")
617 .join("file2")
618 .exists()
619 );
620 let content = fs::read_to_string(output_path.join("child").join("file1"))?;
621 assert_eq!(content, "Lorem ipsum");
622 let content =
623 fs::read_to_string(output_path.join("child").join("grand_child").join("file2"))?;
624 assert_eq!(content, "Lorem ipsum");
625 Ok(())
626 }
627
628 #[tokio::test]
629 pub async fn test_write_snapshot_manifest() -> anyhow::Result<()> {
630 let input = TempDir::new()?;
631 let input_path = input.path();
632 let epoch_0 = input_path.join("epoch_0");
633 fs::create_dir(&epoch_0)?;
634 let file1 = epoch_0.join("file1");
635 fs::write(file1, b"Lorem ipsum")?;
636 let file2 = epoch_0.join("file2");
637 fs::write(file2, b"Lorem ipsum")?;
638 let grandchild = epoch_0.join("grand_child");
639 fs::create_dir(&grandchild)?;
640 let file3 = grandchild.join("file2.tar.gz");
641 fs::write(file3, b"Lorem ipsum")?;
642
643 let input_store = ObjectStoreConfig {
644 object_store: Some(ObjectStoreType::File),
645 directory: Some(input_path.to_path_buf()),
646 ..Default::default()
647 }
648 .make()?;
649
650 write_snapshot_manifest(
651 &Path::from("epoch_0"),
652 &input_store,
653 String::from("epoch_0/"),
654 )
655 .await?;
656
657 assert!(input_path.join("epoch_0").join(MANIFEST_FILENAME).exists());
658 let content = fs::read_to_string(input_path.join("epoch_0").join(MANIFEST_FILENAME))?;
659 assert!(content.contains("file2"));
660 assert!(content.contains("file1"));
661 assert!(content.contains("grand_child/file2.tar.gz"));
662 Ok(())
663 }
664
665 #[tokio::test]
666 pub async fn test_delete_recursively() -> anyhow::Result<()> {
667 let input = TempDir::new()?;
668 let input_path = input.path();
669 let child = input_path.join("child");
670 fs::create_dir(&child)?;
671 let file1 = child.join("file1");
672 fs::write(file1, b"Lorem ipsum")?;
673 let grandchild = child.join("grand_child");
674 fs::create_dir(&grandchild)?;
675 let file2 = grandchild.join("file2");
676 fs::write(file2, b"Lorem ipsum")?;
677
678 let input_store = ObjectStoreConfig {
679 object_store: Some(ObjectStoreType::File),
680 directory: Some(input_path.to_path_buf()),
681 ..Default::default()
682 }
683 .make()?;
684
685 delete_recursively(
686 &Path::from("child"),
687 &input_store,
688 NonZeroUsize::new(1).unwrap(),
689 )
690 .await?;
691
692 assert!(!input_path.join("child").join("file1").exists());
693 assert!(
694 !input_path
695 .join("child")
696 .join("grand_child")
697 .join("file2")
698 .exists()
699 );
700 Ok(())
701 }
702}