1use anyhow::Result;
2use heed::{CompactionOption, EnvOpenOptions};
3use std::collections::HashSet;
4use std::path::{Path, PathBuf};
5#[cfg(feature = "s3")]
6use std::sync::Arc;
7#[cfg(feature = "s3")]
8use std::time::{Duration, SystemTime, UNIX_EPOCH};
9
10use super::{GcStats, HashtreeStore};
11
12#[cfg(feature = "s3")]
13use futures::{stream::FuturesUnordered, StreamExt};
14#[cfg(feature = "s3")]
15use hashtree_core::from_hex;
16use hashtree_core::{sha256, to_hex};
17use serde::{Deserialize, Serialize};
18
19#[derive(Debug, Clone)]
21pub struct VerifyResult {
22 pub total: usize,
23 pub valid: usize,
24 pub corrupted: usize,
25 pub deleted: usize,
26}
27
28#[derive(Debug, Clone)]
29pub struct CompactResult {
30 pub env_dir: PathBuf,
31 pub before_bytes: u64,
32 pub after_bytes: u64,
33}
34
35#[derive(Debug, Clone, Default)]
36pub struct R2ImportOptions {
37 pub concurrency: usize,
38 pub check_only: bool,
39 pub resume: bool,
40 pub fast_list: bool,
41 pub stream_merge: bool,
42 pub keys: Vec<String>,
43 pub keys_file: Option<PathBuf>,
44 pub start_after: Option<String>,
45 pub scan_prefix: Option<String>,
46 pub state_file: Option<PathBuf>,
47 pub max_objects: Option<usize>,
48 pub progress_every: usize,
49 pub scan_delay_ms: u64,
50}
51
52#[derive(Debug, Clone, Default, Serialize, Deserialize)]
53pub struct R2ImportResult {
54 pub listed: usize,
55 pub skipped: usize,
56 pub missing: usize,
57 pub imported: usize,
58 pub corrupted: usize,
59 pub failed: usize,
60 pub bytes_imported: u64,
61 pub last_key: Option<String>,
62 pub completed: bool,
63}
64
65#[cfg(feature = "s3")]
66#[derive(Debug, Clone, Default, Serialize, Deserialize)]
67struct R2ImportState {
68 #[serde(flatten)]
69 result: R2ImportResult,
70 updated_at_unix: u64,
71}
72
73#[cfg(feature = "s3")]
74#[derive(Debug, Clone)]
75struct R2ObjectCandidate {
76 key: String,
77 hash: hashtree_core::types::Hash,
78}
79
80#[cfg(feature = "s3")]
81#[derive(Debug, Clone, Default)]
82struct R2ObjectImportOutcome {
83 skipped: bool,
84 missing: bool,
85 imported: bool,
86 corrupted: bool,
87 failed: bool,
88 bytes_imported: u64,
89 message: Option<String>,
90}
91
92#[cfg(feature = "s3")]
93const R2_IMPORT_OBJECT_READ_ATTEMPTS: usize = 4;
94#[cfg(feature = "s3")]
95const R2_IMPORT_OBJECT_RETRY_BASE_DELAY_MS: u64 = 250;
96
97const COMPACT_MAX_DBS: u32 = 64;
98const COMPACT_MAX_READERS: u32 = 2048;
99const COMPACT_OPEN_MAP_SIZE_BYTES: usize = 10 * 1024 * 1024;
100const COMPACT_PAGE_SIZE_BYTES: u64 = 4096;
101
102#[cfg(feature = "s3")]
103fn unix_timestamp_now() -> u64 {
104 SystemTime::now()
105 .duration_since(UNIX_EPOCH)
106 .unwrap_or_default()
107 .as_secs()
108}
109
110#[cfg(feature = "s3")]
111fn r2_import_key_hash(prefix: &str, key: &str) -> Option<hashtree_core::types::Hash> {
112 let filename = key.strip_prefix(prefix).unwrap_or(key);
113 let hash_hex = filename.strip_suffix(".bin")?;
114 if hash_hex.contains('/') {
115 return None;
116 }
117 if hash_hex.len() != 64 {
118 return None;
119 }
120 from_hex(hash_hex).ok()
121}
122
123#[cfg(feature = "s3")]
124fn r2_import_key_candidate(prefix: &str, input: &str) -> Option<R2ObjectCandidate> {
125 let input = input.trim();
126 if input.is_empty() {
127 return None;
128 }
129
130 let key = if input.len() == 64 && input.chars().all(|ch| ch.is_ascii_hexdigit()) {
131 format!("{prefix}{input}.bin")
132 } else if !prefix.is_empty() && !input.starts_with(prefix) && !input.contains('/') {
133 format!("{prefix}{input}")
134 } else {
135 input.to_string()
136 };
137
138 let hash = r2_import_key_hash(prefix, &key)?;
139 Some(R2ObjectCandidate { key, hash })
140}
141
142#[cfg(feature = "s3")]
143fn existing_r2_candidates(
144 local: &super::LocalStore,
145 candidates: &[R2ObjectCandidate],
146) -> Result<Vec<bool>> {
147 let mut indexed_hashes: Vec<(usize, hashtree_core::types::Hash)> = candidates
148 .iter()
149 .enumerate()
150 .map(|(index, candidate)| (index, candidate.hash))
151 .collect();
152 indexed_hashes.sort_unstable_by(|left, right| left.1.cmp(&right.1).then(left.0.cmp(&right.0)));
153
154 let sorted_hashes: Vec<hashtree_core::types::Hash> =
155 indexed_hashes.iter().map(|(_, hash)| *hash).collect();
156 let sorted_existing = local
157 .existing_hashes_in_sorted_candidates(&sorted_hashes)
158 .map_err(|err| anyhow::anyhow!("Failed to compare local hashes: {err}"))?;
159
160 let mut existing = vec![false; candidates.len()];
161 for ((candidate_index, _), exists) in indexed_hashes.into_iter().zip(sorted_existing) {
162 existing[candidate_index] = exists;
163 }
164 Ok(existing)
165}
166
167#[cfg(feature = "s3")]
168fn read_r2_import_keys_file(path: &Path) -> Result<Vec<String>> {
169 let raw = std::fs::read_to_string(path)?;
170 Ok(raw
171 .lines()
172 .map(str::trim)
173 .filter(|line| !line.is_empty() && !line.starts_with('#'))
174 .map(ToOwned::to_owned)
175 .collect())
176}
177
178#[cfg(feature = "s3")]
179fn read_r2_import_state(path: &Path) -> Option<R2ImportState> {
180 let raw = std::fs::read_to_string(path).ok()?;
181 serde_json::from_str(&raw).ok()
182}
183
184#[cfg(feature = "s3")]
185async fn fetch_r2_object_body_with_retries(
186 client: &aws_sdk_s3::Client,
187 bucket: &str,
188 key: &str,
189) -> Result<Vec<u8>, String> {
190 let mut last_error = None;
191 for attempt in 1..=R2_IMPORT_OBJECT_READ_ATTEMPTS {
192 let output = match client.get_object().bucket(bucket).key(key).send().await {
193 Ok(output) => output,
194 Err(err) => {
195 last_error = Some(format!("fetch failed for {key}: {err}"));
196 if attempt < R2_IMPORT_OBJECT_READ_ATTEMPTS {
197 let delay_ms = R2_IMPORT_OBJECT_RETRY_BASE_DELAY_MS << (attempt - 1);
198 tokio::time::sleep(Duration::from_millis(delay_ms)).await;
199 }
200 continue;
201 }
202 };
203
204 match output.body.collect().await {
205 Ok(body) => return Ok(body.into_bytes().to_vec()),
206 Err(err) => {
207 last_error = Some(format!("read failed for {key}: {err}"));
208 if attempt < R2_IMPORT_OBJECT_READ_ATTEMPTS {
209 let delay_ms = R2_IMPORT_OBJECT_RETRY_BASE_DELAY_MS << (attempt - 1);
210 tokio::time::sleep(Duration::from_millis(delay_ms)).await;
211 }
212 }
213 }
214 }
215
216 Err(format!(
217 "{} after {} attempt(s)",
218 last_error.unwrap_or_else(|| format!("fetch failed for {key}: unknown error")),
219 R2_IMPORT_OBJECT_READ_ATTEMPTS
220 ))
221}
222
223#[cfg(feature = "s3")]
224fn write_r2_import_state(path: &Path, result: &R2ImportResult) -> Result<()> {
225 if let Some(parent) = path.parent() {
226 std::fs::create_dir_all(parent)?;
227 }
228 let state = R2ImportState {
229 result: result.clone(),
230 updated_at_unix: unix_timestamp_now(),
231 };
232 std::fs::write(path, serde_json::to_vec_pretty(&state)?)?;
233 Ok(())
234}
235
236#[cfg(feature = "s3")]
237async fn import_r2_object_to_local(
238 client: Arc<aws_sdk_s3::Client>,
239 bucket: Arc<String>,
240 local: Arc<super::LocalStore>,
241 candidate: R2ObjectCandidate,
242 check_only: bool,
243 prechecked_missing: bool,
244) -> R2ObjectImportOutcome {
245 if !prechecked_missing {
246 match local.exists(&candidate.hash) {
247 Ok(true) => {
248 return R2ObjectImportOutcome {
249 skipped: true,
250 ..Default::default()
251 };
252 }
253 Ok(false) => {}
254 Err(err) => {
255 return R2ObjectImportOutcome {
256 failed: true,
257 message: Some(format!("local exists failed for {}: {err}", candidate.key)),
258 ..Default::default()
259 };
260 }
261 }
262 }
263
264 if check_only {
265 return R2ObjectImportOutcome {
266 missing: true,
267 ..Default::default()
268 };
269 }
270
271 let body =
272 match fetch_r2_object_body_with_retries(client.as_ref(), bucket.as_str(), &candidate.key)
273 .await
274 {
275 Ok(body) => body,
276 Err(err) => {
277 return R2ObjectImportOutcome {
278 missing: true,
279 failed: true,
280 message: Some(err),
281 ..Default::default()
282 };
283 }
284 };
285 let data = body.as_slice();
286 let actual_hash = sha256(data);
287 if actual_hash != candidate.hash {
288 return R2ObjectImportOutcome {
289 missing: true,
290 corrupted: true,
291 message: Some(format!(
292 "hash mismatch for {}: actual {}",
293 candidate.key,
294 to_hex(&actual_hash)
295 )),
296 ..Default::default()
297 };
298 }
299
300 match local.put_sync(candidate.hash, data) {
301 Ok(inserted) => R2ObjectImportOutcome {
302 missing: true,
303 imported: inserted,
304 skipped: !inserted,
305 bytes_imported: if inserted { data.len() as u64 } else { 0 },
306 ..Default::default()
307 },
308 Err(err) => R2ObjectImportOutcome {
309 missing: true,
310 failed: true,
311 message: Some(format!("local put failed for {}: {err}", candidate.key)),
312 ..Default::default()
313 },
314 }
315}
316
317#[cfg(feature = "s3")]
318async fn settle_one_r2_import(
319 pending: &mut FuturesUnordered<impl std::future::Future<Output = R2ObjectImportOutcome>>,
320 result: &mut R2ImportResult,
321) {
322 if let Some(outcome) = pending.next().await {
323 if outcome.skipped {
324 result.skipped += 1;
325 }
326 if outcome.missing {
327 result.missing += 1;
328 }
329 if outcome.imported {
330 result.imported += 1;
331 result.bytes_imported = result.bytes_imported.saturating_add(outcome.bytes_imported);
332 }
333 if outcome.corrupted {
334 result.corrupted += 1;
335 }
336 if outcome.failed {
337 result.failed += 1;
338 }
339 if let Some(message) = outcome.message {
340 println!(" {message}");
341 }
342 }
343}
344
345impl HashtreeStore {
346 pub fn gc(&self) -> Result<GcStats> {
348 let rtxn = self.env.read_txn()?;
349
350 let pinned: HashSet<[u8; 32]> = self
352 .pins
353 .iter(&rtxn)?
354 .filter_map(|item| item.ok())
355 .filter_map(|(hash_bytes, _)| {
356 if hash_bytes.len() == 32 {
357 let mut hash = [0u8; 32];
358 hash.copy_from_slice(hash_bytes);
359 Some(hash)
360 } else {
361 None
362 }
363 })
364 .collect();
365
366 drop(rtxn);
367
368 let all_hashes = self
371 .router
372 .list_writable()
373 .map_err(|e| anyhow::anyhow!("Failed to list writable hashes: {}", e))?;
374
375 let mut deleted = 0;
377 let mut freed_bytes = 0u64;
378
379 for hash in all_hashes {
380 if !pinned.contains(&hash) {
381 if let Ok(Some(size)) = self.router.blob_size_sync(&hash) {
382 freed_bytes += size;
383 let _ = self.router.delete_local_only(&hash);
385 deleted += 1;
386 }
387 }
388 }
389
390 Ok(GcStats {
391 deleted_dags: deleted,
392 freed_bytes,
393 })
394 }
395
396 pub fn verify_lmdb_integrity(&self, delete: bool) -> Result<VerifyResult> {
399 let all_hashes = self
400 .router
401 .list()
402 .map_err(|e| anyhow::anyhow!("Failed to list hashes: {}", e))?;
403
404 let total = all_hashes.len();
405 let mut valid = 0;
406 let mut corrupted = 0;
407 let mut deleted = 0;
408 let mut corrupted_hashes = Vec::new();
409
410 for hash in &all_hashes {
411 let hash_hex = to_hex(hash);
412
413 match self.router.get_sync(hash) {
414 Ok(Some(data)) => {
415 let actual_hash = sha256(&data);
416
417 if actual_hash == *hash {
418 valid += 1;
419 } else {
420 corrupted += 1;
421 let actual_hex = to_hex(&actual_hash);
422 println!(
423 " CORRUPTED: key={} actual={} size={}",
424 &hash_hex[..16],
425 &actual_hex[..16],
426 data.len()
427 );
428 corrupted_hashes.push(*hash);
429 }
430 }
431 Ok(None) => {
432 corrupted += 1;
433 println!(" MISSING: key={}", &hash_hex[..16]);
434 corrupted_hashes.push(*hash);
435 }
436 Err(e) => {
437 corrupted += 1;
438 println!(" ERROR: key={} err={}", &hash_hex[..16], e);
439 corrupted_hashes.push(*hash);
440 }
441 }
442 }
443
444 if delete {
445 for hash in &corrupted_hashes {
446 match self.router.delete_sync(hash) {
447 Ok(true) => deleted += 1,
448 Ok(false) => {}
449 Err(e) => {
450 let hash_hex = to_hex(hash);
451 println!(" Failed to delete {}: {}", &hash_hex[..16], e);
452 }
453 }
454 }
455 }
456
457 Ok(VerifyResult {
458 total,
459 valid,
460 corrupted,
461 deleted,
462 })
463 }
464
465 #[cfg(feature = "s3")]
468 pub async fn verify_r2_integrity(&self, delete: bool) -> Result<VerifyResult> {
469 use aws_sdk_s3::Client as S3Client;
470
471 let config = crate::config::Config::load()?;
472 let s3_config = config
473 .storage
474 .s3
475 .ok_or_else(|| anyhow::anyhow!("S3 not configured"))?;
476
477 let aws_config = aws_config::from_env()
478 .region(aws_sdk_s3::config::Region::new(s3_config.region.clone()))
479 .load()
480 .await;
481
482 let s3_client = S3Client::from_conf(
483 aws_sdk_s3::config::Builder::from(&aws_config)
484 .endpoint_url(&s3_config.endpoint)
485 .force_path_style(true)
486 .build(),
487 );
488
489 let bucket = &s3_config.bucket;
490 let prefix = s3_config.prefix.as_deref().unwrap_or("");
491
492 let mut total = 0;
493 let mut valid = 0;
494 let mut corrupted = 0;
495 let mut deleted = 0;
496 let mut corrupted_keys = Vec::new();
497
498 let mut continuation_token: Option<String> = None;
499
500 loop {
501 let mut list_req = s3_client.list_objects_v2().bucket(bucket).prefix(prefix);
502
503 if let Some(ref token) = continuation_token {
504 list_req = list_req.continuation_token(token);
505 }
506
507 let list_resp = list_req
508 .send()
509 .await
510 .map_err(|e| anyhow::anyhow!("Failed to list S3 objects: {}", e))?;
511
512 for object in list_resp.contents() {
513 let key = object.key().unwrap_or("");
514
515 if !key.ends_with(".bin") {
516 continue;
517 }
518
519 total += 1;
520
521 let filename = key.strip_prefix(prefix).unwrap_or(key);
522 let expected_hash_hex = filename.strip_suffix(".bin").unwrap_or(filename);
523
524 if expected_hash_hex.len() != 64 {
525 corrupted += 1;
526 println!(" INVALID KEY: {}", key);
527 corrupted_keys.push(key.to_string());
528 continue;
529 }
530
531 let expected_hash = match from_hex(expected_hash_hex) {
532 Ok(h) => h,
533 Err(_) => {
534 corrupted += 1;
535 println!(" INVALID HEX: {}", key);
536 corrupted_keys.push(key.to_string());
537 continue;
538 }
539 };
540
541 match s3_client.get_object().bucket(bucket).key(key).send().await {
542 Ok(resp) => match resp.body.collect().await {
543 Ok(bytes) => {
544 let data = bytes.into_bytes();
545 let actual_hash = sha256(&data);
546
547 if actual_hash == expected_hash {
548 valid += 1;
549 } else {
550 corrupted += 1;
551 let actual_hex = to_hex(&actual_hash);
552 println!(
553 " CORRUPTED: key={} actual={} size={}",
554 &expected_hash_hex[..16],
555 &actual_hex[..16],
556 data.len()
557 );
558 corrupted_keys.push(key.to_string());
559 }
560 }
561 Err(e) => {
562 corrupted += 1;
563 println!(" READ ERROR: {} - {}", key, e);
564 corrupted_keys.push(key.to_string());
565 }
566 },
567 Err(e) => {
568 corrupted += 1;
569 println!(" FETCH ERROR: {} - {}", key, e);
570 corrupted_keys.push(key.to_string());
571 }
572 }
573
574 if total % 100 == 0 {
575 println!(
576 " Progress: {} objects checked, {} corrupted so far",
577 total, corrupted
578 );
579 }
580 }
581
582 if list_resp.is_truncated() == Some(true) {
583 continuation_token = list_resp.next_continuation_token().map(|s| s.to_string());
584 } else {
585 break;
586 }
587 }
588
589 if delete {
590 for key in &corrupted_keys {
591 match s3_client
592 .delete_object()
593 .bucket(bucket)
594 .key(key)
595 .send()
596 .await
597 {
598 Ok(_) => deleted += 1,
599 Err(e) => {
600 println!(" Failed to delete {}: {}", key, e);
601 }
602 }
603 }
604 }
605
606 Ok(VerifyResult {
607 total,
608 valid,
609 corrupted,
610 deleted,
611 })
612 }
613
614 #[cfg(feature = "s3")]
621 pub async fn import_r2_to_local(&self, options: R2ImportOptions) -> Result<R2ImportResult> {
622 use aws_sdk_s3::Client as S3Client;
623
624 let config = crate::config::Config::load()?;
625 let s3_config = config
626 .storage
627 .s3
628 .ok_or_else(|| anyhow::anyhow!("S3 not configured"))?;
629
630 let aws_config = aws_config::from_env()
631 .region(aws_sdk_s3::config::Region::new(s3_config.region.clone()))
632 .load()
633 .await;
634
635 let s3_client = S3Client::from_conf(
636 aws_sdk_s3::config::Builder::from(&aws_config)
637 .endpoint_url(&s3_config.endpoint)
638 .force_path_style(true)
639 .build(),
640 );
641
642 let bucket = Arc::new(s3_config.bucket);
643 let prefix = s3_config.prefix.unwrap_or_default();
644 let list_prefix = options
645 .scan_prefix
646 .as_ref()
647 .map(|scan_prefix| format!("{prefix}{scan_prefix}"))
648 .unwrap_or_else(|| prefix.clone());
649 let mut explicit_keys = options.keys.clone();
650 if let Some(keys_file) = options.keys_file.as_ref() {
651 explicit_keys.extend(read_r2_import_keys_file(keys_file)?);
652 }
653
654 let local = self.router.local_store();
655 let client = Arc::new(s3_client);
656 let concurrency = options.concurrency.max(1);
657 let mut pending = FuturesUnordered::new();
658
659 if !explicit_keys.is_empty() {
660 let mut result = R2ImportResult {
661 completed: false,
662 ..Default::default()
663 };
664
665 println!(
666 "R2 import {} targeted: bucket={}, prefix={}, requested_keys={}, state_file={}",
667 if options.check_only { "check" } else { "sync" },
668 bucket.as_str(),
669 prefix,
670 explicit_keys.len(),
671 options
672 .state_file
673 .as_ref()
674 .map(|path| path.display().to_string())
675 .unwrap_or_else(|| "<none>".to_string()),
676 );
677 for key in explicit_keys {
678 let Some(candidate) = r2_import_key_candidate(&prefix, &key) else {
679 result.failed += 1;
680 println!(" invalid R2 blob key/hash: {key}");
681 continue;
682 };
683
684 result.last_key = Some(candidate.key.clone());
685 result.listed += 1;
686 pending.push(import_r2_object_to_local(
687 client.clone(),
688 bucket.clone(),
689 local.clone(),
690 candidate,
691 options.check_only,
692 false,
693 ));
694
695 while pending.len() >= concurrency {
696 settle_one_r2_import(&mut pending, &mut result).await;
697 }
698 }
699
700 while !pending.is_empty() {
701 settle_one_r2_import(&mut pending, &mut result).await;
702 }
703
704 result.completed = true;
705 if let Some(state_file) = options.state_file.as_ref() {
706 write_r2_import_state(state_file, &result)?;
707 }
708 return Ok(result);
709 }
710
711 let state_file = options
712 .state_file
713 .unwrap_or_else(|| self.base_path().join("r2-import-state.json"));
714 let saved_state = read_r2_import_state(&state_file);
715 let saved_incomplete = saved_state
716 .as_ref()
717 .is_some_and(|state| !state.result.completed && state.result.last_key.is_some());
718 let start_after = options.start_after.clone().or_else(|| {
719 if options.resume && saved_incomplete {
720 saved_state
721 .as_ref()
722 .and_then(|state| state.result.last_key.clone())
723 } else {
724 None
725 }
726 });
727 let mut result = if options.resume && options.start_after.is_none() && saved_incomplete {
728 saved_state.map(|state| state.result).unwrap_or_default()
729 } else {
730 R2ImportResult::default()
731 };
732 result.completed = false;
733
734 println!(
735 "R2 import {}: bucket={}, prefix={}, list_prefix={}, start_after={}, state_file={}",
736 if options.check_only { "check" } else { "sync" },
737 bucket.as_str(),
738 prefix,
739 list_prefix,
740 start_after.as_deref().unwrap_or("<beginning>"),
741 state_file.display(),
742 );
743
744 if options.stream_merge && options.fast_list {
745 println!(" Stream merge enabled; skipping in-memory --fast-list index");
746 }
747
748 let local_hashes = if options.fast_list && !options.stream_merge {
749 println!(" Loading local hash index...");
750 let mut local_hashes = self
751 .router
752 .list()
753 .map_err(|err| anyhow::anyhow!("Failed to list local blobs: {err}"))?;
754 local_hashes.sort_unstable();
755 println!(" Local hash index loaded: {} blobs", local_hashes.len());
756 Some(local_hashes)
757 } else {
758 None
759 };
760
761 let progress_every = options.progress_every.max(1);
762 let mut continuation_token: Option<String> = None;
763 let mut listed_since_progress = 0usize;
764 let mut listed_this_run = 0usize;
765 let mut first_page = true;
766 let mut hit_max_objects = false;
767
768 loop {
769 let mut list_req = client
770 .list_objects_v2()
771 .bucket(bucket.as_str())
772 .prefix(&list_prefix);
773
774 if let Some(ref token) = continuation_token {
775 list_req = list_req.continuation_token(token);
776 } else if first_page {
777 if let Some(ref start_after) = start_after {
778 list_req = list_req.start_after(start_after);
779 }
780 }
781 first_page = false;
782
783 let list_resp = list_req
784 .send()
785 .await
786 .map_err(|err| anyhow::anyhow!("Failed to list S3 objects: {err}"))?;
787
788 let mut page_candidates = Vec::new();
789 let mut page_last_key = None;
790 for object in list_resp.contents() {
791 if options
792 .max_objects
793 .is_some_and(|max_objects| listed_this_run >= max_objects)
794 {
795 hit_max_objects = true;
796 break;
797 }
798
799 let key = object.key().unwrap_or("").to_string();
800 page_last_key = Some(key.clone());
801 if !key.ends_with(".bin") {
802 continue;
803 }
804
805 let Some(hash) = r2_import_key_hash(&prefix, &key) else {
806 continue;
807 };
808
809 result.listed += 1;
810 listed_this_run += 1;
811 listed_since_progress += 1;
812 page_candidates.push(R2ObjectCandidate { key, hash });
813 }
814
815 let page_existing = if options.stream_merge && !page_candidates.is_empty() {
816 Some(existing_r2_candidates(local.as_ref(), &page_candidates)?)
817 } else {
818 None
819 };
820
821 for (candidate_index, candidate) in page_candidates.into_iter().enumerate() {
822 let already_exists = page_existing
823 .as_ref()
824 .is_some_and(|existing| existing[candidate_index]);
825
826 if options.scan_delay_ms > 0 {
827 tokio::time::sleep(Duration::from_millis(options.scan_delay_ms)).await;
828 }
829
830 if already_exists {
831 result.skipped += 1;
832 continue;
833 }
834
835 if let Some(local_hashes) = &local_hashes {
836 if local_hashes.binary_search(&candidate.hash).is_ok() {
837 result.skipped += 1;
838 continue;
839 }
840 }
841
842 pending.push(import_r2_object_to_local(
843 client.clone(),
844 bucket.clone(),
845 local.clone(),
846 candidate,
847 options.check_only,
848 page_existing.is_some(),
849 ));
850
851 while pending.len() >= concurrency {
852 settle_one_r2_import(&mut pending, &mut result).await;
853 }
854 }
855
856 while !pending.is_empty() {
857 settle_one_r2_import(&mut pending, &mut result).await;
858 }
859 if let Some(last_key) = page_last_key {
860 result.last_key = Some(last_key);
861 }
862 if listed_since_progress >= progress_every {
863 listed_since_progress = 0;
864 println!(
865 " Progress: {} listed, {} imported, {} skipped, {} missing, {} corrupted, {} failed, {:.2} GB imported",
866 result.listed,
867 result.imported,
868 result.skipped,
869 result.missing,
870 result.corrupted,
871 result.failed,
872 result.bytes_imported as f64 / 1024.0 / 1024.0 / 1024.0,
873 );
874 }
875 write_r2_import_state(&state_file, &result)?;
876
877 if hit_max_objects {
878 break;
879 }
880 if list_resp.is_truncated() == Some(true) {
881 continuation_token = list_resp.next_continuation_token().map(|s| s.to_string());
882 } else {
883 result.completed = true;
884 break;
885 }
886 }
887
888 write_r2_import_state(&state_file, &result)?;
889 Ok(result)
890 }
891
892 #[cfg(not(feature = "s3"))]
894 pub async fn verify_r2_integrity(&self, _delete: bool) -> Result<VerifyResult> {
895 Err(anyhow::anyhow!("S3 feature not enabled"))
896 }
897
898 pub fn compact_lmdb_environments(
899 &self,
900 env_dirs: &[PathBuf],
901 keep_backup: bool,
902 ) -> Result<Vec<CompactResult>> {
903 compact_lmdb_environments_under(self.base_path(), env_dirs, keep_backup)
904 }
905}
906
907pub fn compact_lmdb_environments_under(
908 base_path: &Path,
909 env_dirs: &[PathBuf],
910 keep_backup: bool,
911) -> Result<Vec<CompactResult>> {
912 let targets = if env_dirs.is_empty() {
913 discover_lmdb_environment_dirs(base_path)?
914 } else {
915 env_dirs
916 .iter()
917 .map(|path| {
918 if path.is_absolute() {
919 path.clone()
920 } else {
921 base_path.join(path)
922 }
923 })
924 .collect()
925 };
926
927 let mut results = Vec::new();
928 for env_dir in targets {
929 results.push(compact_lmdb_environment_dir(&env_dir, keep_backup)?);
930 }
931 Ok(results)
932}
933
934fn discover_lmdb_environment_dirs(root: &Path) -> Result<Vec<PathBuf>> {
935 let mut dirs = Vec::new();
936 collect_lmdb_environment_dirs(root, &mut dirs)?;
937 dirs.sort();
938 Ok(dirs)
939}
940
941fn collect_lmdb_environment_dirs(root: &Path, dirs: &mut Vec<PathBuf>) -> Result<()> {
942 if root.join("data.mdb").exists() {
943 dirs.push(root.to_path_buf());
944 }
945
946 for entry in std::fs::read_dir(root)? {
947 let entry = entry?;
948 let path = entry.path();
949 if path.is_dir() {
950 collect_lmdb_environment_dirs(&path, dirs)?;
951 }
952 }
953
954 Ok(())
955}
956
957fn compact_lmdb_environment_dir(env_dir: &Path, keep_backup: bool) -> Result<CompactResult> {
958 let data_path = env_dir.join("data.mdb");
959 if !data_path.exists() {
960 anyhow::bail!("No data.mdb found in {}", env_dir.display());
961 }
962
963 let before_bytes = std::fs::metadata(&data_path)?.len();
964 let compact_path = env_dir.join("data.mdb.compact");
965 let backup_path = env_dir.join("data.mdb.bak");
966
967 if compact_path.exists() {
968 std::fs::remove_file(&compact_path)?;
969 }
970 if !keep_backup && backup_path.exists() {
971 std::fs::remove_file(&backup_path)?;
972 }
973
974 let open_map_size = existing_lmdb_map_size_bytes(&data_path)?;
975
976 {
977 let env = unsafe {
978 EnvOpenOptions::new()
979 .map_size(open_map_size)
980 .max_dbs(COMPACT_MAX_DBS)
981 .max_readers(COMPACT_MAX_READERS)
982 .open(env_dir)
983 }?;
984 env.force_sync()?;
985 env.copy_to_file(&compact_path, CompactionOption::Enabled)?;
986 }
987
988 let after_bytes = std::fs::metadata(&compact_path)?.len();
989
990 if backup_path.exists() {
991 std::fs::remove_file(&backup_path)?;
992 }
993
994 std::fs::rename(&data_path, &backup_path)?;
995 if let Err(error) = std::fs::rename(&compact_path, &data_path) {
996 let _ = std::fs::rename(&backup_path, &data_path);
997 return Err(error.into());
998 }
999
1000 if !keep_backup {
1001 std::fs::remove_file(&backup_path)?;
1002 }
1003
1004 Ok(CompactResult {
1005 env_dir: env_dir.to_path_buf(),
1006 before_bytes,
1007 after_bytes,
1008 })
1009}
1010
1011fn existing_lmdb_map_size_bytes(data_path: &Path) -> Result<usize> {
1012 let file_bytes = std::fs::metadata(data_path)?.len();
1013 let aligned_bytes = if file_bytes == 0 {
1014 COMPACT_OPEN_MAP_SIZE_BYTES as u64
1015 } else {
1016 let remainder = file_bytes % COMPACT_PAGE_SIZE_BYTES;
1017 if remainder == 0 {
1018 file_bytes
1019 } else {
1020 file_bytes.saturating_add(COMPACT_PAGE_SIZE_BYTES - remainder)
1021 }
1022 };
1023
1024 Ok(usize::try_from(aligned_bytes)
1025 .unwrap_or(usize::MAX)
1026 .max(COMPACT_OPEN_MAP_SIZE_BYTES))
1027}
1028
1029#[cfg(all(test, feature = "s3"))]
1030mod tests {
1031 use super::{r2_import_key_candidate, r2_import_key_hash};
1032
1033 const HASH: &str = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef";
1034
1035 #[test]
1036 fn r2_import_key_hash_accepts_only_root_blob_keys() {
1037 assert!(r2_import_key_hash("", &format!("{HASH}.bin")).is_some());
1038 assert!(r2_import_key_hash("legacy/", &format!("legacy/{HASH}.bin")).is_some());
1039
1040 assert!(r2_import_key_hash("", &format!("hot/{HASH}.bin")).is_none());
1041 assert!(
1042 r2_import_key_hash("", &format!("site-bytes/pubkey/tree/root/{HASH}.bin")).is_none()
1043 );
1044 assert!(r2_import_key_hash("", "roots/pubkey/tree.json").is_none());
1045 assert!(r2_import_key_hash("", &format!("{HASH}.png")).is_none());
1046 assert!(r2_import_key_hash("", "not-a-hash.bin").is_none());
1047 }
1048
1049 #[test]
1050 fn r2_import_key_candidate_accepts_hash_or_canonical_key() {
1051 let bare = r2_import_key_candidate("", HASH).expect("bare hash");
1052 assert_eq!(bare.key, format!("{HASH}.bin"));
1053
1054 let explicit = r2_import_key_candidate("", &format!("{HASH}.bin")).expect("hash key");
1055 assert_eq!(explicit.key, format!("{HASH}.bin"));
1056
1057 assert!(r2_import_key_candidate("", &format!("hot/{HASH}.bin")).is_none());
1058 }
1059
1060 #[test]
1061 fn r2_import_key_candidate_applies_configured_prefix() {
1062 let bare = r2_import_key_candidate("legacy/", HASH).expect("prefixed bare hash");
1063 assert_eq!(bare.key, format!("legacy/{HASH}.bin"));
1064
1065 let explicit =
1066 r2_import_key_candidate("legacy/", &format!("{HASH}.bin")).expect("prefixed key");
1067 assert_eq!(explicit.key, format!("legacy/{HASH}.bin"));
1068
1069 let already_prefixed =
1070 r2_import_key_candidate("legacy/", &format!("legacy/{HASH}.bin")).expect("key");
1071 assert_eq!(already_prefixed.key, format!("legacy/{HASH}.bin"));
1072 }
1073}