Skip to main content

hashtree_cli/storage/
maintenance.rs

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/// Result of blob integrity verification
20#[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    /// Garbage collect unpinned content
347    pub fn gc(&self) -> Result<GcStats> {
348        let rtxn = self.env.read_txn()?;
349
350        // Get all pinned hashes as raw bytes
351        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        // Get hashes in the writable tier. Cold/legacy tiers are archival in
369        // tiered LMDB mode and must not be reclaimed by local GC.
370        let all_hashes = self
371            .router
372            .list_writable()
373            .map_err(|e| anyhow::anyhow!("Failed to list writable hashes: {}", e))?;
374
375        // Delete unpinned hashes
376        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                    // Delete locally only - keep S3 as archive
384                    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    /// Verify LMDB blob integrity - checks that stored data matches its key hash
397    /// Returns verification statistics and optionally deletes corrupted entries
398    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    /// Verify R2/S3 blob integrity - lists all objects and verifies hash matches filename
466    /// Returns verification statistics and optionally deletes corrupted entries
467    #[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    /// Import missing R2/S3 blobs into local storage without writing back to S3.
615    ///
616    /// This mirrors rclone's shape: list the source, compare each source object
617    /// against the destination by cheap metadata (here the content-addressed key),
618    /// and only transfer missing objects. `--check-only` runs the same comparison
619    /// without downloading object bodies.
620    #[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    /// Fallback for non-S3 builds
893    #[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}