Skip to main content

hashtree_cli/storage/
maintenance.rs

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/// Result of blob integrity verification
24#[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    /// Garbage collect unpinned content
352    pub fn gc(&self) -> Result<GcStats> {
353        let retention = self.active_retention_protection()?;
354        let rtxn = self.env.read_txn()?;
355
356        // Get all pinned hashes as raw bytes
357        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        // Get hashes in the canonical writable store.
375        let all_hashes = self
376            .router
377            .list_writable()
378            .map_err(|e| anyhow::anyhow!("Failed to list writable hashes: {}", e))?;
379
380        // Delete unpinned hashes
381        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                    // Delete locally only - keep S3 as archive
389                    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    /// Verify LMDB blob integrity - checks that stored data matches its key hash
402    /// Returns verification statistics and optionally deletes corrupted entries
403    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    /// Verify R2/S3 blob integrity - lists all objects and verifies hash matches filename
471    /// Returns verification statistics and optionally deletes corrupted entries
472    #[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    /// Import missing R2/S3 blobs into local storage without writing back to S3.
620    ///
621    /// This mirrors rclone's shape: list the source, compare each source object
622    /// against the destination by cheap metadata (here the content-addressed key),
623    /// and only transfer missing objects. `--check-only` runs the same comparison
624    /// without downloading object bodies.
625    #[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    /// Fallback for non-S3 builds
898    #[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}