1#![deny(unsafe_code)]
2#![cfg_attr(
3 test,
4 allow(
5 clippy::unwrap_used,
6 clippy::expect_used,
7 clippy::indexing_slicing,
8 clippy::arithmetic_side_effects,
9 clippy::shadow_unrelated,
10 clippy::let_underscore_must_use,
11 clippy::format_push_string
12 )
13)]
14
15use std::{
21 collections::{HashMap, HashSet},
22 hash::Hash,
23 io,
24 num::TryFromIntError,
25};
26
27use serde_json::Error as JsonError;
28use shardline_index::{
29 AsyncIndexStore, DedupeShardMapping, FileId, LocalIndexStoreError, MemoryIndexStoreError,
30 MemoryRecordStoreError, PostgresMetadataStoreError, RecordMutation, RecordTraversal,
31 parse_xet_hash_hex, xet_hash_hex_string,
32};
33use shardline_protocol::HashParseError;
34use shardline_server_core::{
35 InvalidSerializedShardError, OpsRecordStore, ServerObjectStore, ServerObjectStoreError,
36 ShardMetadataLimits, checked_increment, read_full_object,
37};
38use shardline_storage::{
39 LocalObjectStoreError, ObjectPrefix, ObjectPrefixError, S3ObjectStoreError,
40};
41use shardline_xet_adapter::{XetAdapterError, retained_shard_chunk_hashes};
42use thiserror::Error;
43
44mod candidates;
45use candidates::{VersionCandidate, collect_candidate};
46
47#[derive(Debug, Error)]
49pub enum RebuildError {
50 #[error("local storage operation failed")]
52 Io(#[from] io::Error),
53 #[error("json operation failed")]
55 Json(#[from] JsonError),
56 #[error("numeric conversion exceeded supported bounds")]
58 NumericConversion(#[from] TryFromIntError),
59 #[error("content hash must be 64 hexadecimal characters")]
61 InvalidContentHash,
62 #[error(
64 "file identifier must be relative and must not contain traversal or control characters"
65 )]
66 InvalidFileId,
67 #[error("arithmetic overflow")]
69 Overflow,
70 #[error("object storage prefix validation failed")]
72 ObjectPrefix(#[from] ObjectPrefixError),
73 #[error("local storage adapter operation failed")]
75 LocalObjectStore(#[from] LocalObjectStoreError),
76 #[error("s3 object storage adapter operation failed")]
78 S3ObjectStore(#[from] S3ObjectStoreError),
79 #[error("xet adapter operation failed")]
81 XetAdapter(#[from] XetAdapterError),
82 #[error("index adapter operation failed")]
84 IndexStore(#[from] LocalIndexStoreError),
85 #[error("memory index adapter operation failed")]
87 MemoryIndexStore(#[from] MemoryIndexStoreError),
88 #[error("memory record adapter operation failed")]
90 MemoryRecordStore(#[from] MemoryRecordStoreError),
91 #[error("postgres metadata adapter operation failed")]
93 PostgresMetadata(#[from] PostgresMetadataStoreError),
94 #[error("hash parsing failed")]
96 HashParse(#[from] HashParseError),
97 #[error("stored file metadata exceeded the bounded parser ceiling")]
99 StoredFileMetadataTooLarge {
100 observed_bytes: u64,
102 maximum_bytes: u64,
104 },
105}
106
107impl From<shardline_server_core::ParseStoredFileRecordError> for RebuildError {
108 fn from(value: shardline_server_core::ParseStoredFileRecordError) -> Self {
109 match value {
110 shardline_server_core::ParseStoredFileRecordError::StoredFileMetadataTooLarge {
111 observed_bytes,
112 maximum_bytes,
113 } => Self::StoredFileMetadataTooLarge {
114 observed_bytes,
115 maximum_bytes,
116 },
117 shardline_server_core::ParseStoredFileRecordError::Json(e) => Self::Json(e),
118 }
119 }
120}
121
122impl From<shardline_server_core::ValidateIdentifierError> for RebuildError {
123 fn from(_: shardline_server_core::ValidateIdentifierError) -> Self {
124 Self::InvalidFileId
125 }
126}
127
128impl From<shardline_server_core::ValidateContentHashError> for RebuildError {
129 fn from(_: shardline_server_core::ValidateContentHashError) -> Self {
130 Self::InvalidContentHash
131 }
132}
133
134impl From<shardline_server_core::RebuildOverflowError> for RebuildError {
135 fn from(_: shardline_server_core::RebuildOverflowError) -> Self {
136 Self::Overflow
137 }
138}
139
140impl From<InvalidSerializedShardError> for RebuildError {
141 fn from(value: InvalidSerializedShardError) -> Self {
142 Self::XetAdapter(XetAdapterError::InvalidSerializedShard(value))
143 }
144}
145
146impl From<ServerObjectStoreError> for RebuildError {
147 fn from(value: ServerObjectStoreError) -> Self {
148 match value {
149 ServerObjectStoreError::NotFound => Self::Overflow,
150 ServerObjectStoreError::Overflow => Self::Overflow,
151 ServerObjectStoreError::InvalidContentHash => Self::Overflow,
152 ServerObjectStoreError::StoredObjectLengthMismatch => Self::Overflow,
153 ServerObjectStoreError::Local(e) => Self::LocalObjectStore(e),
154 ServerObjectStoreError::S3(e) => Self::S3ObjectStore(e),
155 ServerObjectStoreError::Io(e) => Self::Io(e),
156 ServerObjectStoreError::NumericConversion(e) => Self::NumericConversion(e),
157 }
158 }
159}
160
161#[derive(Debug, Clone, PartialEq, Eq)]
163pub struct IndexRebuildReport {
164 pub scanned_version_records: u64,
166 pub scanned_retained_shards: u64,
168 pub rebuilt_latest_records: u64,
170 pub unchanged_latest_records: u64,
172 pub removed_stale_latest_records: u64,
174 pub scanned_reconstructions: u64,
176 pub unchanged_reconstructions: u64,
178 pub removed_stale_reconstructions: u64,
180 pub rebuilt_dedupe_shard_mappings: u64,
182 pub unchanged_dedupe_shard_mappings: u64,
184 pub removed_stale_dedupe_shard_mappings: u64,
186 pub issues: Vec<IndexRebuildIssue>,
188}
189
190impl IndexRebuildReport {
191 #[must_use]
193 pub const fn issue_count(&self) -> usize {
194 self.issues.len()
195 }
196
197 #[must_use]
199 pub const fn is_clean(&self) -> bool {
200 self.issues.is_empty()
201 }
202}
203
204#[derive(Debug, Clone, PartialEq, Eq)]
206pub struct IndexRebuildIssue {
207 pub kind: IndexRebuildIssueKind,
209 pub location: String,
211 pub detail: IndexRebuildIssueDetail,
213}
214
215#[derive(Debug, Clone, Error, PartialEq, Eq)]
217pub enum IndexRebuildIssueDetail {
218 #[error("record metadata exceeded the bounded parser ceiling")]
220 OversizedVersionRecordMetadata,
221 #[error("record json was invalid")]
223 RecordJsonInvalid,
224 #[error("record file_id `{file_id}` is invalid")]
226 InvalidFileId {
227 file_id: String,
229 },
230 #[error("record content hash `{content_hash}` is invalid")]
232 InvalidContentHash {
233 content_hash: String,
235 },
236 #[error("record repository scope is invalid")]
238 InvalidRepositoryScope,
239 #[error("expected version record at {expected_locator}")]
241 VersionPathMismatch {
242 expected_locator: String,
244 },
245 #[error("{0}")]
247 InvalidReconstructionPlan(IndexRebuildReconstructionPlanDetail),
248 #[error("{0}")]
250 InvalidRetainedShard(InvalidSerializedShardError),
251}
252
253#[derive(Debug, Clone, Copy, Error, PartialEq, Eq)]
255pub enum IndexRebuildReconstructionPlanDetail {
256 #[error("record chunk hash is invalid")]
258 ChunkHashInvalid,
259 #[error("record contains an empty chunk")]
261 EmptyChunk,
262 #[error("record chunks are not contiguous")]
264 NonContiguousChunkOffsets,
265 #[error("record chunk range is invalid")]
267 InvalidChunkRange,
268 #[error("record packed range is invalid")]
270 InvalidPackedRange,
271 #[error("record length overflowed")]
273 LengthOverflow,
274 #[error("record total byte count did not match chunks")]
276 TotalBytesMismatch,
277}
278
279#[derive(Debug, Clone, Copy, PartialEq, Eq)]
281pub enum IndexRebuildIssueKind {
282 OversizedVersionRecordMetadata,
284 InvalidVersionRecordJson,
286 InvalidVersionFileId,
288 InvalidVersionContentHash,
290 InvalidVersionRepositoryScope,
292 VersionPathMismatch,
294 InvalidVersionReconstructionPlan,
296 InvalidRetainedShard,
298}
299
300impl IndexRebuildIssueKind {
301 #[must_use]
303 pub const fn as_str(self) -> &'static str {
304 match self {
305 Self::OversizedVersionRecordMetadata => "oversized_version_record_metadata",
306 Self::InvalidVersionRecordJson => "invalid_version_record_json",
307 Self::InvalidVersionFileId => "invalid_version_file_id",
308 Self::InvalidVersionContentHash => "invalid_version_content_hash",
309 Self::InvalidVersionRepositoryScope => "invalid_version_repository_scope",
310 Self::VersionPathMismatch => "version_path_mismatch",
311 Self::InvalidVersionReconstructionPlan => "invalid_version_reconstruction_plan",
312 Self::InvalidRetainedShard => "invalid_retained_shard",
313 }
314 }
315}
316
317pub type LocalIndexRebuildReport = IndexRebuildReport;
319
320pub type LocalIndexRebuildIssue = IndexRebuildIssue;
322
323pub type LocalIndexRebuildIssueKind = IndexRebuildIssueKind;
325
326pub async fn run_index_rebuild_with_stores<RecordAdapter, IndexAdapter>(
333 record_store: &RecordAdapter,
334 index_store: &IndexAdapter,
335 object_store: &ServerObjectStore,
336 shard_metadata_limits: ShardMetadataLimits,
337) -> Result<IndexRebuildReport, RebuildError>
338where
339 RecordAdapter: OpsRecordStore + Sync,
340 RecordAdapter::Error: Into<RebuildError>,
341 RecordAdapter::Locator: Hash,
342 IndexAdapter: AsyncIndexStore + Sync,
343 IndexAdapter::Error: Into<RebuildError>,
344{
345 let mut report = IndexRebuildReport {
346 scanned_version_records: 0,
347 scanned_retained_shards: 0,
348 rebuilt_latest_records: 0,
349 unchanged_latest_records: 0,
350 removed_stale_latest_records: 0,
351 scanned_reconstructions: 0,
352 unchanged_reconstructions: 0,
353 removed_stale_reconstructions: 0,
354 rebuilt_dedupe_shard_mappings: 0,
355 unchanged_dedupe_shard_mappings: 0,
356 removed_stale_dedupe_shard_mappings: 0,
357 issues: Vec::new(),
358 };
359 let mut candidates = HashMap::new();
360 RecordTraversal::visit_version_records(record_store, |entry| {
361 report.scanned_version_records = checked_increment(report.scanned_version_records)?;
362 collect_candidate(record_store, entry, &mut candidates, &mut report)
363 })
364 .await?;
365
366 let mut desired_latest_paths = HashSet::new();
367 for candidate in candidates.values() {
368 let latest_path = RecordTraversal::latest_record_locator(record_store, &candidate.record);
369 desired_latest_paths.insert(latest_path.clone());
370
371 let record_bytes = serde_json::to_vec(&candidate.record)?;
372 let existing_bytes =
373 RecordTraversal::read_latest_record_bytes(record_store, &candidate.record)
374 .await
375 .map_err(Into::into)?;
376
377 if existing_bytes.as_deref() == Some(record_bytes.as_slice()) {
378 report.unchanged_latest_records = checked_increment(report.unchanged_latest_records)?;
379 continue;
380 }
381
382 RecordMutation::write_latest_record(record_store, &candidate.record)
383 .await
384 .map_err(Into::into)?;
385 report.rebuilt_latest_records = checked_increment(report.rebuilt_latest_records)?;
386 }
387
388 let mut stale_latest_paths = Vec::new();
389 RecordTraversal::visit_latest_record_locators(record_store, |path| {
390 if !desired_latest_paths.contains(&path) {
391 stale_latest_paths.push(path);
392 }
393
394 Ok::<(), RebuildError>(())
395 })
396 .await?;
397 for path in stale_latest_paths {
398 RecordMutation::delete_record_locator(record_store, &path)
399 .await
400 .map_err(Into::into)?;
401 report.removed_stale_latest_records =
402 checked_increment(report.removed_stale_latest_records)?;
403 }
404
405 RecordMutation::prune_empty_latest_records(record_store)
406 .await
407 .map_err(Into::into)?;
408
409 let desired_reconstructions = desired_reconstruction_file_ids(candidates.values());
410 prune_stale_reconstructions(index_store, &desired_reconstructions, &mut report).await?;
411
412 rebuild_dedupe_shard_mappings(
413 index_store,
414 object_store,
415 shard_metadata_limits,
416 &mut report,
417 )
418 .await?;
419
420 Ok(report)
421}
422
423fn desired_reconstruction_file_ids<'record, Locator, Records>(records: Records) -> HashSet<String>
424where
425 Records: IntoIterator<Item = &'record VersionCandidate<Locator>>,
426 Locator: 'record,
427{
428 records
429 .into_iter()
430 .filter_map(|candidate| {
431 parse_xet_hash_hex(&candidate.record.file_id)
432 .ok()
433 .map(xet_hash_hex_string)
434 })
435 .collect::<HashSet<_>>()
436}
437
438async fn prune_stale_reconstructions<IndexAdapter>(
439 index_store: &IndexAdapter,
440 desired_reconstructions: &HashSet<String>,
441 report: &mut IndexRebuildReport,
442) -> Result<(), RebuildError>
443where
444 IndexAdapter: AsyncIndexStore + Sync,
445 IndexAdapter::Error: Into<RebuildError>,
446{
447 if !report.is_clean() {
448 return Ok(());
449 }
450
451 let existing_file_ids = index_store
452 .list_reconstruction_file_ids()
453 .await
454 .map_err(Into::into)?;
455 for file_id in existing_file_ids {
456 report.scanned_reconstructions = checked_increment(report.scanned_reconstructions)?;
457 let file_id_hex = xet_hash_hex_string(file_id.hash());
458 if desired_reconstructions.contains(&file_id_hex) {
459 report.unchanged_reconstructions = checked_increment(report.unchanged_reconstructions)?;
460 continue;
461 }
462
463 delete_reconstruction(index_store, &file_id).await?;
464 report.removed_stale_reconstructions =
465 checked_increment(report.removed_stale_reconstructions)?;
466 }
467
468 Ok(())
469}
470
471async fn delete_reconstruction<IndexAdapter>(
472 index_store: &IndexAdapter,
473 file_id: &FileId,
474) -> Result<(), RebuildError>
475where
476 IndexAdapter: AsyncIndexStore + Sync,
477 IndexAdapter::Error: Into<RebuildError>,
478{
479 let _deleted = index_store
480 .delete_reconstruction(file_id)
481 .await
482 .map_err(Into::into)?;
483 Ok(())
484}
485
486async fn rebuild_dedupe_shard_mappings<IndexAdapter>(
487 index_store: &IndexAdapter,
488 object_store: &ServerObjectStore,
489 shard_metadata_limits: ShardMetadataLimits,
490 report: &mut IndexRebuildReport,
491) -> Result<(), RebuildError>
492where
493 IndexAdapter: AsyncIndexStore + Sync,
494 IndexAdapter::Error: Into<RebuildError>,
495{
496 let prefix =
497 ObjectPrefix::parse("shards/").map_err(|_error| RebuildError::InvalidContentHash)?;
498 let mut desired = HashMap::<String, DedupeShardMapping>::new();
499 let issue_count_before_scan = report.issue_count();
500
501 object_store.visit_prefix(&prefix, |metadata| -> Result<(), RebuildError> {
502 report.scanned_retained_shards = checked_increment(report.scanned_retained_shards)?;
503 let shard_key = metadata.key().clone();
504 let shard_location = shard_key.as_str().to_owned();
505 let shard_bytes = read_full_object(object_store, &shard_key, metadata.length())
506 .map_err(RebuildError::from)?;
507 let chunk_hashes = match retained_shard_chunk_hashes(&shard_bytes, shard_metadata_limits) {
508 Ok(chunk_hashes) => chunk_hashes,
509 Err(XetAdapterError::InvalidSerializedShard(detail)) => {
510 push_issue(
511 report,
512 IndexRebuildIssueKind::InvalidRetainedShard,
513 shard_location,
514 IndexRebuildIssueDetail::InvalidRetainedShard(detail),
515 )?;
516 return Ok(());
517 }
518 Err(error) => return Err(error.into()),
519 };
520
521 for chunk_hash_hex in chunk_hashes {
522 let mapping =
523 DedupeShardMapping::new(parse_xet_hash_hex(&chunk_hash_hex)?, shard_key.clone());
524 match desired.get(&chunk_hash_hex) {
525 Some(existing)
526 if existing.shard_object_key().as_str()
527 <= mapping.shard_object_key().as_str() => {}
528 _ => {
529 desired.insert(chunk_hash_hex, mapping);
530 }
531 }
532 }
533 Ok(())
534 })?;
535
536 if report.issue_count() != issue_count_before_scan {
537 return Ok(());
538 }
539
540 let mut existing = HashMap::new();
541 index_store
542 .visit_dedupe_shard_mappings(|mapping| {
543 existing.insert(xet_hash_hex_string(mapping.chunk_hash()), mapping);
544 Ok::<(), RebuildError>(())
545 })
546 .await?;
547
548 for (chunk_hash_hex, mapping) in &desired {
549 match existing.get(chunk_hash_hex) {
550 Some(existing_mapping)
551 if existing_mapping.shard_object_key() == mapping.shard_object_key() =>
552 {
553 report.unchanged_dedupe_shard_mappings =
554 checked_increment(report.unchanged_dedupe_shard_mappings)?;
555 }
556 _ => {
557 index_store
558 .upsert_dedupe_shard_mapping(mapping)
559 .await
560 .map_err(Into::into)?;
561 report.rebuilt_dedupe_shard_mappings =
562 checked_increment(report.rebuilt_dedupe_shard_mappings)?;
563 }
564 }
565 }
566
567 for (chunk_hash_hex, _mapping) in existing {
568 if desired.contains_key(&chunk_hash_hex) {
569 continue;
570 }
571
572 let chunk_hash = parse_xet_hash_hex(&chunk_hash_hex)?;
573 let _deleted = index_store
574 .delete_dedupe_shard_mapping(&chunk_hash)
575 .await
576 .map_err(Into::into)?;
577 report.removed_stale_dedupe_shard_mappings =
578 checked_increment(report.removed_stale_dedupe_shard_mappings)?;
579 }
580
581 Ok(())
582}
583
584fn push_issue(
585 report: &mut IndexRebuildReport,
586 kind: IndexRebuildIssueKind,
587 location: String,
588 detail: IndexRebuildIssueDetail,
589) -> Result<(), RebuildError> {
590 let _count = u64::try_from(report.issues.len())?;
591 report.issues.push(IndexRebuildIssue {
592 kind,
593 location,
594 detail,
595 });
596 Ok(())
597}
598
599#[cfg(test)]
600mod tests {
601 use super::*;
602
603 fn empty_report() -> IndexRebuildReport {
604 IndexRebuildReport {
605 scanned_version_records: 0,
606 scanned_retained_shards: 0,
607 rebuilt_latest_records: 0,
608 unchanged_latest_records: 0,
609 removed_stale_latest_records: 0,
610 scanned_reconstructions: 0,
611 unchanged_reconstructions: 0,
612 removed_stale_reconstructions: 0,
613 rebuilt_dedupe_shard_mappings: 0,
614 unchanged_dedupe_shard_mappings: 0,
615 removed_stale_dedupe_shard_mappings: 0,
616 issues: Vec::new(),
617 }
618 }
619
620 #[test]
621 fn report_is_clean_when_no_issues() {
622 let report = empty_report();
623 assert!(report.is_clean());
624 assert_eq!(report.issue_count(), 0);
625 }
626
627 #[test]
628 fn report_is_not_clean_with_issues() {
629 let mut report = empty_report();
630 report.issues.push(IndexRebuildIssue {
631 kind: IndexRebuildIssueKind::InvalidVersionRecordJson,
632 location: "test/path".to_owned(),
633 detail: IndexRebuildIssueDetail::RecordJsonInvalid,
634 });
635 assert!(!report.is_clean());
636 assert_eq!(report.issue_count(), 1);
637 }
638
639 #[test]
640 fn issue_count_matches_vec_len() {
641 let mut report = empty_report();
642 for i in 0..5 {
643 report.issues.push(IndexRebuildIssue {
644 kind: IndexRebuildIssueKind::InvalidVersionFileId,
645 location: format!("loc/{i}"),
646 detail: IndexRebuildIssueDetail::InvalidFileId {
647 file_id: format!("fid-{i}"),
648 },
649 });
650 }
651 assert_eq!(report.issue_count(), 5);
652 }
653
654 #[test]
655 fn index_rebuild_report_equality() {
656 let a = empty_report();
657 let mut b = empty_report();
658 b.scanned_version_records = 42;
659 assert_ne!(a, b);
660
661 b.scanned_version_records = 0;
662 assert_eq!(a, b);
663 }
664
665 #[test]
666 fn issue_kind_as_str_returns_expected_labels() {
667 assert_eq!(
668 IndexRebuildIssueKind::OversizedVersionRecordMetadata.as_str(),
669 "oversized_version_record_metadata"
670 );
671 assert_eq!(
672 IndexRebuildIssueKind::InvalidVersionRecordJson.as_str(),
673 "invalid_version_record_json"
674 );
675 assert_eq!(
676 IndexRebuildIssueKind::InvalidVersionFileId.as_str(),
677 "invalid_version_file_id"
678 );
679 assert_eq!(
680 IndexRebuildIssueKind::InvalidVersionContentHash.as_str(),
681 "invalid_version_content_hash"
682 );
683 assert_eq!(
684 IndexRebuildIssueKind::InvalidVersionRepositoryScope.as_str(),
685 "invalid_version_repository_scope"
686 );
687 assert_eq!(
688 IndexRebuildIssueKind::VersionPathMismatch.as_str(),
689 "version_path_mismatch"
690 );
691 assert_eq!(
692 IndexRebuildIssueKind::InvalidVersionReconstructionPlan.as_str(),
693 "invalid_version_reconstruction_plan"
694 );
695 assert_eq!(
696 IndexRebuildIssueKind::InvalidRetainedShard.as_str(),
697 "invalid_retained_shard"
698 );
699 }
700
701 #[test]
702 fn issue_detail_display_messages() {
703 let detail = IndexRebuildIssueDetail::OversizedVersionRecordMetadata;
704 assert!(!detail.to_string().is_empty());
705
706 let detail = IndexRebuildIssueDetail::RecordJsonInvalid;
707 assert!(!detail.to_string().is_empty());
708
709 let detail = IndexRebuildIssueDetail::InvalidFileId {
710 file_id: "bad-id".to_owned(),
711 };
712 assert!(detail.to_string().contains("bad-id"));
713
714 let detail = IndexRebuildIssueDetail::InvalidContentHash {
715 content_hash: "abc123".to_owned(),
716 };
717 assert!(detail.to_string().contains("abc123"));
718
719 let detail = IndexRebuildIssueDetail::InvalidRepositoryScope;
720 assert!(!detail.to_string().is_empty());
721
722 let detail = IndexRebuildIssueDetail::VersionPathMismatch {
723 expected_locator: "/expected/path".to_owned(),
724 };
725 assert!(detail.to_string().contains("/expected/path"));
726
727 let detail = IndexRebuildIssueDetail::InvalidReconstructionPlan(
728 IndexRebuildReconstructionPlanDetail::ChunkHashInvalid,
729 );
730 assert!(!detail.to_string().is_empty());
731
732 let detail = IndexRebuildIssueDetail::InvalidRetainedShard(
733 shardline_server_core::InvalidSerializedShardError::ParserRejectedMetadata,
734 );
735 assert!(!detail.to_string().is_empty());
736 }
737
738 #[test]
739 fn push_issue_increments_count() {
740 let mut report = empty_report();
741 push_issue(
742 &mut report,
743 IndexRebuildIssueKind::InvalidVersionFileId,
744 "loc".to_owned(),
745 IndexRebuildIssueDetail::InvalidFileId {
746 file_id: "x".to_owned(),
747 },
748 )
749 .unwrap();
750 assert_eq!(report.issue_count(), 1);
751 }
752}