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