1use std::sync::Arc;
9
10use async_trait::async_trait;
11use bytes::Bytes;
12use futures::StreamExt;
13use lance_core::utils::tracing::{
14 AUDIT_MODE_CREATE, AUDIT_MODE_DELETE, AUDIT_TYPE_MANIFEST, TRACE_FILE_AUDIT,
15};
16use lance_core::{Error, Result};
17use lance_io::object_store::ObjectStore;
18use log::warn;
19use object_store::ObjectMeta;
20use object_store::ObjectStoreExt;
21use object_store::{Error as ObjectStoreError, ObjectStore as OSObjectStore, path::Path};
22use tracing::info;
23
24use super::{
25 MANIFEST_EXTENSION, ManifestLocation, ManifestNamingScheme, current_manifest_path,
26 default_resolve_version, make_staging_manifest_path, write_version_hint,
27};
28use crate::format::{IndexMetadata, Manifest, Transaction};
29use crate::io::commit::{CommitError, CommitHandler};
30
31#[async_trait]
76pub trait ExternalManifestStore: std::fmt::Debug + Send + Sync {
77 async fn get(&self, base_uri: &str, version: u64) -> Result<String>;
79
80 async fn get_manifest_location(
81 &self,
82 base_uri: &str,
83 version: u64,
84 ) -> Result<ManifestLocation> {
85 let path = self.get(base_uri, version).await?;
86 let path = Path::parse(&path).map_err(|e| Error::invalid_input(e.to_string()))?;
87 let naming_scheme = detect_naming_scheme_from_path(&path)?;
88 Ok(ManifestLocation {
89 version,
90 path,
91 size: None,
92 naming_scheme,
93 e_tag: None,
94 })
95 }
96
97 async fn get_latest_version(&self, base_uri: &str) -> Result<Option<(u64, String)>>;
101
102 async fn get_latest_manifest_location(
108 &self,
109 base_uri: &str,
110 ) -> Result<Option<ManifestLocation>> {
111 self.get_latest_version(base_uri).await.and_then(|res| {
112 res.map(|(version, uri)| {
113 let path = Path::parse(&uri).map_err(|e| Error::invalid_input(e.to_string()))?;
114 let naming_scheme = detect_naming_scheme_from_path(&path)?;
115 Ok(ManifestLocation {
116 version,
117 path,
118 size: None,
119 naming_scheme,
120 e_tag: None,
121 })
122 })
123 .transpose()
124 })
125 }
126
127 #[allow(clippy::too_many_arguments)]
136 async fn put(
137 &self,
138 base_path: &Path,
139 version: u64,
140 staging_path: &Path,
141 size: u64,
142 _e_tag: Option<String>,
143 object_store: &dyn OSObjectStore,
144 naming_scheme: ManifestNamingScheme,
145 ) -> Result<ManifestLocation> {
146 self.put_if_not_exists(
155 base_path.as_ref(),
156 version,
157 staging_path.as_ref(),
158 size,
159 None,
160 )
161 .await?;
162
163 let final_path = naming_scheme.manifest_path(base_path, version);
165 let final_e_tag =
166 copy_or_verify_final_manifest(object_store, staging_path, &final_path, version, size)
167 .await?;
168
169 let location = ManifestLocation {
170 version,
171 path: final_path.clone(),
172 size: Some(size),
173 naming_scheme,
174 e_tag: final_e_tag,
175 };
176
177 let published = self
184 .put_if_exists(base_path.as_ref(), version, final_path.as_ref(), size, None)
185 .await;
186
187 if let Err(error) = published {
188 warn!(
193 "Final manifest '{}' is committed, but the external manifest index could not be updated; retaining staging manifest '{}' for repair: {}",
194 final_path, staging_path, error
195 );
196 return Ok(location);
197 }
198
199 match object_store.delete(staging_path).await {
201 Ok(_) => {}
202 Err(ObjectStoreError::NotFound { .. }) => {}
203 Err(error) => {
204 warn!(
208 "Failed to delete finalized staging manifest '{}': {}",
209 staging_path, error
210 );
211 return Ok(location);
212 }
213 }
214 info!(target: TRACE_FILE_AUDIT, mode=AUDIT_MODE_DELETE, r#type=AUDIT_TYPE_MANIFEST, path = staging_path.as_ref());
215
216 Ok(location)
217 }
218
219 async fn put_if_not_exists(
226 &self,
227 base_uri: &str,
228 version: u64,
229 path: &str,
230 size: u64,
231 e_tag: Option<String>,
232 ) -> Result<()>;
233
234 async fn put_if_exists(
238 &self,
239 base_uri: &str,
240 version: u64,
241 path: &str,
242 size: u64,
243 e_tag: Option<String>,
244 ) -> Result<()>;
245
246 async fn delete(&self, _base_uri: &str) -> Result<()> {
248 Ok(())
249 }
250}
251
252pub(crate) fn detect_naming_scheme_from_path(path: &Path) -> Result<ManifestNamingScheme> {
253 path.filename()
254 .and_then(|name| {
255 ManifestNamingScheme::detect_scheme(name)
256 .or_else(|| Some(ManifestNamingScheme::detect_scheme_staging(name)))
257 })
258 .ok_or_else(|| {
259 Error::corrupt_file(
260 path.clone(),
261 "Path does not follow known manifest naming convention.",
262 )
263 })
264}
265
266const MAX_SERVER_SIDE_COPY_BYTES: u64 = 5 * 1024 * 1024 * 1024;
275
276const COPY_REWRITE_PART_SIZE: usize = 100 * 1024 * 1024;
282
283async fn copy_size_aware(
304 store: &dyn OSObjectStore,
305 from: &Path,
306 to: &Path,
307 size: u64,
308) -> std::result::Result<(), ObjectStoreError> {
309 if size < MAX_SERVER_SIDE_COPY_BYTES {
310 store.copy(from, to).await
311 } else {
312 copy_via_read_rewrite(store, from, to).await
313 }
314}
315
316async fn copy_or_verify_final_manifest(
332 object_store: &dyn OSObjectStore,
333 staging_path: &Path,
334 final_path: &Path,
335 version: u64,
336 selected_size: u64,
337) -> Result<Option<String>> {
338 match copy_size_aware(object_store, staging_path, final_path, selected_size).await {
339 Ok(()) => {
340 info!(target: TRACE_FILE_AUDIT, mode=AUDIT_MODE_CREATE, r#type=AUDIT_TYPE_MANIFEST, path = final_path.as_ref());
341 let final_meta = object_store.head(final_path).await?;
342 if final_meta.size != selected_size {
343 return Err(Error::corrupt_file(
344 final_path.clone(),
345 format!(
346 "Manifest size mismatch for version {}: selected staging manifest had {}, object store returned {}",
347 version, selected_size, final_meta.size
348 ),
349 ));
350 }
351 Ok(final_meta.e_tag)
352 }
353 Err(ObjectStoreError::NotFound { .. }) => match object_store.head(final_path).await {
354 Ok(final_meta) if final_meta.size == selected_size => Ok(final_meta.e_tag),
355 Ok(final_meta) => Err(Error::corrupt_file(
356 final_path.clone(),
357 format!(
358 "Manifest size mismatch for version {}: selected staging manifest had {}, object store returned {}",
359 version, selected_size, final_meta.size
360 ),
361 )),
362 Err(error) => Err(error.into()),
363 },
364 Err(error) => Err(error.into()),
365 }
366}
367
368async fn copy_via_read_rewrite(
375 store: &dyn OSObjectStore,
376 from: &Path,
377 to: &Path,
378) -> std::result::Result<(), ObjectStoreError> {
379 let mut stream = store.get(from).await?.into_stream();
381
382 let mut upload = store.put_multipart(to).await?;
392 let mut part_buf: Vec<u8> = Vec::with_capacity(COPY_REWRITE_PART_SIZE);
393
394 while let Some(chunk) = stream.next().await {
395 let chunk = match chunk {
396 Ok(b) => b,
397 Err(e) => {
398 let _ = upload.abort().await;
399 return Err(e);
400 }
401 };
402 let mut offset = 0;
408 while offset < chunk.len() {
409 let want = COPY_REWRITE_PART_SIZE - part_buf.len();
410 let take = want.min(chunk.len() - offset);
411 part_buf.extend_from_slice(&chunk[offset..offset + take]);
412 offset += take;
413
414 if part_buf.len() >= COPY_REWRITE_PART_SIZE {
415 let payload =
416 std::mem::replace(&mut part_buf, Vec::with_capacity(COPY_REWRITE_PART_SIZE));
417 if let Err(e) = upload.put_part(Bytes::from(payload).into()).await {
418 let _ = upload.abort().await;
419 return Err(e);
420 }
421 }
422 }
423 }
424
425 if !part_buf.is_empty()
428 && let Err(e) = upload.put_part(Bytes::from(part_buf).into()).await
429 {
430 let _ = upload.abort().await;
431 return Err(e);
432 }
433
434 if let Err(e) = upload.complete().await {
435 let _ = upload.abort().await;
436 return Err(e);
437 }
438 Ok(())
439}
440
441#[derive(Debug)]
445pub struct ExternalManifestCommitHandler {
446 pub external_manifest_store: Arc<dyn ExternalManifestStore>,
447}
448
449impl ExternalManifestCommitHandler {
450 async fn verify_finalized_manifest_location(
451 &self,
452 base_path: &Path,
453 location: ManifestLocation,
454 object_store: &dyn OSObjectStore,
455 ) -> std::result::Result<ManifestLocation, Error> {
456 match object_store.head(&location.path).await {
457 Ok(ObjectMeta { size, e_tag, .. }) => {
458 let ManifestLocation {
459 version,
460 path,
461 size: expected_size,
462 naming_scheme,
463 e_tag: _,
464 } = location;
465
466 let size = match expected_size {
467 Some(expected_size) if expected_size != size => {
468 return Err(Error::corrupt_file(
469 path,
470 format!(
471 "Manifest size mismatch for version {}: external store expected {}, object store returned {}",
472 version, expected_size, size
473 ),
474 ));
475 }
476 Some(expected_size) => Some(expected_size),
477 None => Some(size),
478 };
479
480 Ok(ManifestLocation {
487 version,
488 path,
489 size,
490 naming_scheme,
491 e_tag,
492 })
493 }
494 Err(ObjectStoreError::NotFound { .. }) => {
495 default_resolve_version(base_path, location.version, object_store).await
498 }
499 Err(e) => Err(e.into()),
500 }
501 }
502
503 #[allow(clippy::too_many_arguments)]
511 async fn finalize_manifest(
512 &self,
513 base_path: &Path,
514 staging_manifest_path: &Path,
515 version: u64,
516 size: u64,
517 store: &dyn OSObjectStore,
518 naming_scheme: ManifestNamingScheme,
519 ) -> std::result::Result<ManifestLocation, Error> {
520 let final_manifest_path = naming_scheme.manifest_path(base_path, version);
522
523 let final_e_tag = copy_or_verify_final_manifest(
524 store,
525 staging_manifest_path,
526 &final_manifest_path,
527 version,
528 size,
529 )
530 .await?;
531
532 let location = ManifestLocation {
533 version,
534 path: final_manifest_path,
535 size: Some(size),
536 naming_scheme,
537 e_tag: final_e_tag,
538 };
539
540 let published = self
547 .external_manifest_store
548 .put_if_exists(
549 base_path.as_ref(),
550 version,
551 location.path.as_ref(),
552 size,
553 None,
554 )
555 .await;
556
557 if let Err(error) = published {
558 warn!(
562 "Final manifest '{}' is committed, but the external manifest index could not be updated; retaining staging manifest '{}' for repair: {}",
563 location.path, staging_manifest_path, error
564 );
565 return Ok(location);
566 }
567
568 match store.delete(staging_manifest_path).await {
570 Ok(_) => {}
571 Err(ObjectStoreError::NotFound { .. }) => {}
572 Err(error) => {
573 warn!(
574 "Failed to delete finalized staging manifest '{}': {}",
575 staging_manifest_path, error
576 );
577 return Ok(location);
578 }
579 }
580 info!(target: TRACE_FILE_AUDIT, mode=AUDIT_MODE_DELETE, r#type=AUDIT_TYPE_MANIFEST, path = staging_manifest_path.as_ref());
581
582 Ok(location)
583 }
584}
585
586#[async_trait]
587impl CommitHandler for ExternalManifestCommitHandler {
588 async fn resolve_latest_location(
589 &self,
590 base_path: &Path,
591 object_store: &ObjectStore,
592 ) -> std::result::Result<ManifestLocation, Error> {
593 let location = self
594 .external_manifest_store
595 .get_latest_manifest_location(base_path.as_ref())
596 .await?;
597
598 match location {
599 Some(location) => {
600 if location.path.extension() == Some(MANIFEST_EXTENSION) {
601 return self
602 .verify_finalized_manifest_location(
603 base_path,
604 location,
605 object_store.inner.as_ref(),
606 )
607 .await;
608 }
609
610 let ManifestLocation {
611 version,
612 path,
613 size,
614 naming_scheme,
615 e_tag: _,
616 } = location;
617
618 let size = if let Some(size) = size {
619 size
620 } else {
621 match object_store.inner.head(&path).await {
622 Ok(meta) => meta.size,
623 Err(ObjectStoreError::NotFound { .. }) => {
624 let new_location = self
626 .external_manifest_store
627 .get_manifest_location(base_path.as_ref(), version)
628 .await?;
629 return Ok(new_location);
630 }
631 Err(e) => return Err(e.into()),
632 }
633 };
634
635 let final_location = self
636 .finalize_manifest(
637 base_path,
638 &path,
639 version,
640 size,
641 &object_store.inner,
642 naming_scheme,
643 )
644 .await?;
645
646 Ok(final_location)
647 }
648 None => current_manifest_path(object_store, base_path).await,
651 }
652 }
653
654 async fn resolve_version_location(
655 &self,
656 base_path: &Path,
657 version: u64,
658 object_store: &dyn OSObjectStore,
659 ) -> std::result::Result<ManifestLocation, Error> {
660 let location_res = self
661 .external_manifest_store
662 .get_manifest_location(base_path.as_ref(), version)
663 .await;
664
665 let location = match location_res {
666 Ok(p) => p,
667 Err(Error::NotFound { .. }) => {
669 let path = default_resolve_version(base_path, version, object_store)
670 .await
671 .map_err(|_| Error::not_found(format!("{}@{}", base_path, version)))?
672 .path;
673 match object_store.head(&path).await {
674 Ok(ObjectMeta { size, e_tag, .. }) => {
675 let res = self
676 .external_manifest_store
677 .put_if_not_exists(
678 base_path.as_ref(),
679 version,
680 path.as_ref(),
681 size,
682 None,
683 )
684 .await;
685 if let Err(e) = res {
686 warn!(
687 "could not update external manifest store during load, with error: {}",
688 e
689 );
690 }
691 let naming_scheme =
692 ManifestNamingScheme::detect_scheme_staging(path.filename().unwrap());
693 return Ok(ManifestLocation {
694 version,
695 path,
696 size: Some(size),
697 naming_scheme,
698 e_tag,
699 });
700 }
701 Err(ObjectStoreError::NotFound { .. }) => {
702 return Err(Error::not_found(path.to_string()));
703 }
704 Err(e) => return Err(e.into()),
705 }
706 }
707 Err(e) => return Err(e),
708 };
709
710 if location.path.extension() == Some(MANIFEST_EXTENSION) {
711 return self
712 .verify_finalized_manifest_location(base_path, location, object_store)
713 .await;
714 }
715
716 let naming_scheme =
717 ManifestNamingScheme::detect_scheme_staging(location.path.filename().unwrap());
718
719 let size = if let Some(size) = location.size {
720 size
721 } else {
722 let meta = object_store.head(&location.path).await?;
723 meta.size
724 };
725
726 self.finalize_manifest(
727 base_path,
728 &location.path,
729 version,
730 size,
731 object_store,
732 naming_scheme,
733 )
734 .await
735 }
736
737 async fn version_exists(
738 &self,
739 base_path: &Path,
740 version: u64,
741 object_store: &dyn OSObjectStore,
742 naming_scheme: ManifestNamingScheme,
743 ) -> Result<bool> {
744 match self
745 .external_manifest_store
746 .get_manifest_location(base_path.as_ref(), version)
747 .await
748 {
749 Ok(_) => Ok(true),
750 Err(Error::NotFound { .. }) => {
751 let path = naming_scheme.manifest_path(base_path, version);
752 match object_store.head(&path).await {
753 Ok(_) => Ok(true),
754 Err(ObjectStoreError::NotFound { .. }) => Ok(false),
755 Err(e) => Err(e.into()),
756 }
757 }
758 Err(e) => Err(e),
759 }
760 }
761
762 async fn commit(
763 &self,
764 manifest: &mut Manifest,
765 indices: Option<Vec<IndexMetadata>>,
766 base_path: &Path,
767 object_store: &ObjectStore,
768 manifest_writer: super::ManifestWriter,
769 naming_scheme: ManifestNamingScheme,
770 transaction: Option<Transaction>,
771 ) -> std::result::Result<ManifestLocation, CommitError> {
772 let path = naming_scheme.manifest_path(base_path, manifest.version);
777 let staging_path = make_staging_manifest_path(&path)?;
778 let write_res =
779 manifest_writer(object_store, manifest, indices, &staging_path, transaction).await?;
780
781 let result = self
783 .external_manifest_store
784 .put(
785 base_path,
786 manifest.version,
787 &staging_path,
788 write_res.size as u64,
789 write_res.e_tag,
790 &object_store.inner,
791 naming_scheme,
792 )
793 .await;
794
795 match result {
796 Ok(location) => {
797 write_version_hint(object_store, base_path, manifest.version).await;
798 Ok(location)
799 }
800 Err(error) => {
801 let recorded_location = self
806 .external_manifest_store
807 .get_manifest_location(base_path.as_ref(), manifest.version)
808 .await;
809 if matches!(
810 &recorded_location,
811 Ok(location) if location.path != staging_path
812 ) {
813 match object_store.inner.delete(&staging_path).await {
814 Ok(()) => {
815 info!(target: TRACE_FILE_AUDIT, mode=AUDIT_MODE_DELETE, r#type=AUDIT_TYPE_MANIFEST, path = staging_path.as_ref());
816 }
817 Err(ObjectStoreError::NotFound { .. }) => {}
818 Err(delete_error) => {
819 warn!(
820 "Failed to delete losing staging manifest '{}': {}",
821 staging_path, delete_error
822 );
823 }
824 }
825 return Err(CommitError::CommitConflict);
826 }
827 warn!(
828 "External manifest commit for version {} failed; retaining staging manifest \
829 '{}' until the commit outcome is resolved: {}",
830 manifest.version, staging_path, error
831 );
832 Err(CommitError::CommitConflict)
833 }
834 }
835 }
836
837 async fn delete(&self, base_path: &Path) -> Result<()> {
838 self.external_manifest_store
839 .delete(base_path.as_ref())
840 .await
841 }
842}
843
844#[cfg(test)]
845mod tests {
846 use std::collections::HashMap;
847 use std::sync::Mutex;
848 use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
849
850 use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema};
851 use lance_core::datatypes::Schema;
852 use lance_core::utils::testing::{ProxyObjectStore, ProxyObjectStorePolicy};
853 use lance_file::version::LanceFileVersion;
854 use tokio::sync::Notify;
855
856 use super::*;
857 use crate::format::DataStorageFormat;
858 use crate::io::commit::write_manifest_file_to_path;
859
860 #[derive(Debug, Clone)]
861 struct StoredManifest {
862 path: String,
863 size: u64,
864 e_tag: Option<String>,
865 }
866
867 #[derive(Debug)]
868 struct TestExternalManifestStore {
869 manifests: Mutex<HashMap<(String, u64), StoredManifest>>,
870 fail_next_put_response: AtomicBool,
871 fail_next_final_publish: AtomicBool,
872 block_first_final_publish: bool,
873 final_publish_calls: AtomicUsize,
874 first_final_publish_started: Notify,
875 release_first_final_publish: Notify,
876 }
877
878 impl TestExternalManifestStore {
879 fn new(fail_next_put_response: bool) -> Self {
880 Self {
881 manifests: Mutex::new(HashMap::new()),
882 fail_next_put_response: AtomicBool::new(fail_next_put_response),
883 fail_next_final_publish: AtomicBool::new(false),
884 block_first_final_publish: false,
885 final_publish_calls: AtomicUsize::new(0),
886 first_final_publish_started: Notify::new(),
887 release_first_final_publish: Notify::new(),
888 }
889 }
890
891 fn failing_final_publish_once() -> Self {
892 Self {
893 fail_next_final_publish: AtomicBool::new(true),
894 ..Self::new(false)
895 }
896 }
897
898 fn blocking_first_final_publish() -> Self {
899 Self {
900 block_first_final_publish: true,
901 ..Self::new(false)
902 }
903 }
904 }
905
906 #[async_trait]
907 impl ExternalManifestStore for TestExternalManifestStore {
908 async fn get(&self, base_uri: &str, version: u64) -> Result<String> {
909 self.manifests
910 .lock()
911 .unwrap()
912 .get(&(base_uri.to_string(), version))
913 .map(|manifest| manifest.path.clone())
914 .ok_or_else(|| Error::not_found(format!("{base_uri}@{version}")))
915 }
916
917 async fn get_manifest_location(
918 &self,
919 base_uri: &str,
920 version: u64,
921 ) -> Result<ManifestLocation> {
922 let stored = self
923 .manifests
924 .lock()
925 .unwrap()
926 .get(&(base_uri.to_string(), version))
927 .cloned()
928 .ok_or_else(|| Error::not_found(format!("{base_uri}@{version}")))?;
929 let path = Path::from(stored.path);
930 Ok(ManifestLocation {
931 version,
932 naming_scheme: detect_naming_scheme_from_path(&path)?,
933 path,
934 size: Some(stored.size),
935 e_tag: stored.e_tag,
936 })
937 }
938
939 async fn get_latest_version(&self, base_uri: &str) -> Result<Option<(u64, String)>> {
940 Ok(self
941 .manifests
942 .lock()
943 .unwrap()
944 .iter()
945 .filter(|((stored_base, _), _)| stored_base == base_uri)
946 .max_by_key(|((_, version), _)| *version)
947 .map(|((_, version), manifest)| (*version, manifest.path.clone())))
948 }
949
950 async fn put_if_not_exists(
951 &self,
952 base_uri: &str,
953 version: u64,
954 path: &str,
955 size: u64,
956 e_tag: Option<String>,
957 ) -> Result<()> {
958 let key = (base_uri.to_string(), version);
959 let mut manifests = self.manifests.lock().unwrap();
960 if manifests.contains_key(&key) {
961 return Err(Error::commit_conflict_source(
962 version,
963 "manifest already exists".to_string().into(),
964 ));
965 }
966 manifests.insert(
967 key,
968 StoredManifest {
969 path: path.to_string(),
970 size,
971 e_tag,
972 },
973 );
974 drop(manifests);
975 if self.fail_next_put_response.swap(false, Ordering::SeqCst) {
976 Err(Error::io("simulated lost external-store response"))
977 } else {
978 Ok(())
979 }
980 }
981
982 async fn put_if_exists(
983 &self,
984 base_uri: &str,
985 version: u64,
986 path: &str,
987 size: u64,
988 e_tag: Option<String>,
989 ) -> Result<()> {
990 if self.block_first_final_publish
991 && self.final_publish_calls.fetch_add(1, Ordering::SeqCst) == 0
992 {
993 self.first_final_publish_started.notify_one();
994 self.release_first_final_publish.notified().await;
995 }
996 if self.fail_next_final_publish.swap(false, Ordering::SeqCst) {
997 return Err(Error::io("simulated final index update failure"));
998 }
999 let key = (base_uri.to_string(), version);
1000 let mut manifests = self.manifests.lock().unwrap();
1001 let manifest = manifests
1002 .get_mut(&key)
1003 .ok_or_else(|| Error::not_found(format!("{base_uri}@{version}")))?;
1004 *manifest = StoredManifest {
1005 path: path.to_string(),
1006 size,
1007 e_tag,
1008 };
1009 Ok(())
1010 }
1011 }
1012
1013 fn test_manifest() -> Manifest {
1014 let arrow_schema = ArrowSchema::new(vec![ArrowField::new("id", DataType::Int32, false)]);
1015 Manifest::new(
1016 Schema::try_from(&arrow_schema).unwrap(),
1017 Arc::new(vec![]),
1018 DataStorageFormat::new(LanceFileVersion::Stable.resolve()),
1019 HashMap::new(),
1020 )
1021 }
1022
1023 #[tokio::test]
1024 async fn test_finalized_manifest_ignores_legacy_external_store_etag() {
1025 let external_store = Arc::new(TestExternalManifestStore::new(false));
1026 let handler = ExternalManifestCommitHandler {
1027 external_manifest_store: external_store.clone(),
1028 };
1029 let object_store = ObjectStore::memory();
1030 let base_path = Path::from("dataset");
1031 let final_path = ManifestNamingScheme::V2.manifest_path(&base_path, 1);
1032
1033 object_store
1034 .inner
1035 .put(
1036 &final_path,
1037 object_store::PutPayload::from_static(b"manifest"),
1038 )
1039 .await
1040 .unwrap();
1041 let final_meta = object_store.inner.head(&final_path).await.unwrap();
1042
1043 external_store
1044 .put_if_not_exists(
1045 base_path.as_ref(),
1046 1,
1047 final_path.as_ref(),
1048 final_meta.size,
1049 Some("expected-generation".to_string()),
1050 )
1051 .await
1052 .unwrap();
1053
1054 let resolved = handler
1055 .resolve_version_location(&base_path, 1, object_store.inner.as_ref())
1056 .await
1057 .expect("a legacy external-store ETag must not override object storage");
1058 assert_eq!(resolved.path, final_path);
1059 assert_eq!(resolved.size, Some(final_meta.size));
1060 assert_eq!(resolved.e_tag, final_meta.e_tag);
1061 }
1062
1063 #[tokio::test]
1064 async fn test_finalized_manifest_without_external_store_etag_uses_current_etag() {
1065 let external_store = Arc::new(TestExternalManifestStore::new(false));
1066 let handler = ExternalManifestCommitHandler {
1067 external_manifest_store: external_store.clone(),
1068 };
1069 let object_store = ObjectStore::memory();
1070 let base_path = Path::from("dataset");
1071 let final_path = ManifestNamingScheme::V2.manifest_path(&base_path, 1);
1072
1073 object_store
1074 .inner
1075 .put(
1076 &final_path,
1077 object_store::PutPayload::from_static(b"manifest"),
1078 )
1079 .await
1080 .unwrap();
1081 let final_meta = object_store.inner.head(&final_path).await.unwrap();
1082 external_store
1083 .put_if_not_exists(
1084 base_path.as_ref(),
1085 1,
1086 final_path.as_ref(),
1087 final_meta.size,
1088 None,
1089 )
1090 .await
1091 .unwrap();
1092
1093 let resolved = handler
1094 .resolve_version_location(&base_path, 1, object_store.inner.as_ref())
1095 .await
1096 .expect("an absent external-store ETag must opt out of comparison");
1097 assert_eq!(resolved.path, final_path);
1098 assert_eq!(resolved.size, Some(final_meta.size));
1099 assert_eq!(resolved.e_tag, final_meta.e_tag);
1100 }
1101
1102 #[tokio::test]
1103 async fn test_default_store_returns_but_does_not_persist_etag() {
1104 let external_store = Arc::new(TestExternalManifestStore::new(false));
1105 let handler = ExternalManifestCommitHandler {
1106 external_manifest_store: external_store.clone(),
1107 };
1108 let object_store = ObjectStore::memory();
1109 let base_path = Path::from("dataset");
1110 let mut manifest = test_manifest();
1111
1112 let committed = handler
1113 .commit(
1114 &mut manifest,
1115 None,
1116 &base_path,
1117 &object_store,
1118 write_manifest_file_to_path,
1119 ManifestNamingScheme::V2,
1120 None,
1121 )
1122 .await
1123 .expect("the default store should finalize the selected manifest");
1124 let original = object_store.inner.head(&committed.path).await.unwrap();
1125 assert_eq!(committed.e_tag, original.e_tag);
1126
1127 let indexed = external_store
1128 .get_manifest_location(base_path.as_ref(), committed.version)
1129 .await
1130 .unwrap();
1131 assert_eq!(indexed.e_tag, None);
1132
1133 object_store
1134 .inner
1135 .put(
1136 &committed.path,
1137 object_store::PutPayload::from(vec![0_u8; original.size as usize]),
1138 )
1139 .await
1140 .unwrap();
1141
1142 let replacement = object_store.inner.head(&committed.path).await.unwrap();
1143 assert_ne!(replacement.e_tag, original.e_tag);
1144
1145 let resolved = handler
1146 .resolve_version_location(&base_path, committed.version, object_store.inner.as_ref())
1147 .await
1148 .expect("the external index must not reject a new physical generation");
1149 assert_eq!(resolved.e_tag, replacement.e_tag);
1150 }
1151
1152 #[tokio::test]
1153 async fn test_helping_finalizer_returns_but_does_not_persist_etag() {
1154 let external_store = Arc::new(TestExternalManifestStore::new(false));
1155 let handler = ExternalManifestCommitHandler {
1156 external_manifest_store: external_store.clone(),
1157 };
1158 let object_store = ObjectStore::memory();
1159 let base_path = Path::from("dataset");
1160 let version = 1;
1161 let final_path = ManifestNamingScheme::V2.manifest_path(&base_path, version);
1162 let staging_path = make_staging_manifest_path(&final_path).unwrap();
1163 let manifest_bytes = Bytes::from_static(b"immutable manifest bytes");
1164
1165 object_store
1166 .inner
1167 .put(&staging_path, manifest_bytes.clone().into())
1168 .await
1169 .unwrap();
1170 let staging_meta = object_store.inner.head(&staging_path).await.unwrap();
1171 external_store
1172 .put_if_not_exists(
1173 base_path.as_ref(),
1174 version,
1175 staging_path.as_ref(),
1176 staging_meta.size,
1177 staging_meta.e_tag,
1178 )
1179 .await
1180 .unwrap();
1181
1182 let finalized = handler
1183 .resolve_version_location(&base_path, version, object_store.inner.as_ref())
1184 .await
1185 .expect("a reader should finalize the selected staging manifest");
1186 let final_meta = object_store.inner.head(&final_path).await.unwrap();
1187 assert_eq!(finalized.e_tag, final_meta.e_tag);
1188
1189 let indexed = external_store
1190 .get_manifest_location(base_path.as_ref(), version)
1191 .await
1192 .unwrap();
1193 assert_eq!(indexed.e_tag, None);
1194 }
1195
1196 #[tokio::test]
1197 async fn test_onboarding_returns_but_does_not_persist_etag() {
1198 let external_store = Arc::new(TestExternalManifestStore::new(false));
1199 let handler = ExternalManifestCommitHandler {
1200 external_manifest_store: external_store.clone(),
1201 };
1202 let object_store = ObjectStore::memory();
1203 let base_path = Path::from("dataset");
1204 let version = 1;
1205 let final_path = ManifestNamingScheme::V2.manifest_path(&base_path, version);
1206
1207 object_store
1208 .inner
1209 .put(
1210 &final_path,
1211 object_store::PutPayload::from_static(b"manifest"),
1212 )
1213 .await
1214 .unwrap();
1215 let final_meta = object_store.inner.head(&final_path).await.unwrap();
1216
1217 let resolved = handler
1218 .resolve_version_location(&base_path, version, object_store.inner.as_ref())
1219 .await
1220 .expect("an existing manifest should be indexed during onboarding");
1221 assert_eq!(resolved.e_tag, final_meta.e_tag);
1222
1223 let indexed = external_store
1224 .get_manifest_location(base_path.as_ref(), version)
1225 .await
1226 .unwrap();
1227 assert_eq!(indexed.e_tag, None);
1228 }
1229
1230 #[tokio::test]
1231 async fn test_finalized_manifest_size_mismatch_remains_corruption() {
1232 let external_store = Arc::new(TestExternalManifestStore::new(false));
1233 let handler = ExternalManifestCommitHandler {
1234 external_manifest_store: external_store.clone(),
1235 };
1236 let object_store = ObjectStore::memory();
1237 let base_path = Path::from("dataset");
1238 let final_path = ManifestNamingScheme::V2.manifest_path(&base_path, 1);
1239
1240 object_store
1241 .inner
1242 .put(
1243 &final_path,
1244 object_store::PutPayload::from_static(b"manifest"),
1245 )
1246 .await
1247 .unwrap();
1248 let final_meta = object_store.inner.head(&final_path).await.unwrap();
1249 external_store
1250 .put_if_not_exists(
1251 base_path.as_ref(),
1252 1,
1253 final_path.as_ref(),
1254 final_meta.size + 1,
1255 None,
1256 )
1257 .await
1258 .unwrap();
1259
1260 let error = handler
1261 .resolve_version_location(&base_path, 1, object_store.inner.as_ref())
1262 .await
1263 .expect_err("copies of the selected staging object must preserve its size");
1264 assert!(matches!(error, Error::CorruptFile { .. }));
1265 assert!(error.to_string().contains("Manifest size mismatch"));
1266 }
1267
1268 #[tokio::test]
1269 async fn test_canonical_manifest_commits_before_index_repair() {
1270 let external_store = Arc::new(TestExternalManifestStore::failing_final_publish_once());
1271 let handler = ExternalManifestCommitHandler {
1272 external_manifest_store: external_store.clone(),
1273 };
1274 let object_store = ObjectStore::memory();
1275 let base_path = Path::from("dataset");
1276 let mut manifest = test_manifest();
1277 let version = manifest.version;
1278 let final_path = ManifestNamingScheme::V2.manifest_path(&base_path, version);
1279
1280 let committed = handler
1281 .commit(
1282 &mut manifest,
1283 None,
1284 &base_path,
1285 &object_store,
1286 write_manifest_file_to_path,
1287 ManifestNamingScheme::V2,
1288 None,
1289 )
1290 .await
1291 .expect("a failed index update must not overturn a canonical S3 commit");
1292 assert_eq!(committed.path, final_path);
1293 assert!(
1294 committed.e_tag.is_some(),
1295 "the caller must retain the canonical generation even when index repair fails"
1296 );
1297 object_store
1298 .inner
1299 .head(&final_path)
1300 .await
1301 .expect("the canonical manifest is the durable commit point");
1302
1303 let pending = external_store
1304 .get_manifest_location(base_path.as_ref(), version)
1305 .await
1306 .unwrap();
1307 assert_ne!(pending.path, final_path);
1308 object_store
1309 .inner
1310 .head(&pending.path)
1311 .await
1312 .expect("staging must remain until the external index is repaired");
1313
1314 let repaired = handler
1315 .resolve_version_location(&base_path, version, object_store.inner.as_ref())
1316 .await
1317 .expect("a reader must be able to repair the pending external index");
1318 assert_eq!(repaired.path, final_path);
1319 assert!(
1320 repaired.e_tag.is_some(),
1321 "a helping reader must receive the generation it observed"
1322 );
1323 let indexed = external_store
1324 .get_manifest_location(base_path.as_ref(), version)
1325 .await
1326 .unwrap();
1327 assert_eq!(indexed.path, final_path);
1328 assert_eq!(indexed.size, repaired.size);
1329 assert_eq!(
1330 indexed.e_tag, None,
1331 "the repaired index must not retain a physical object generation"
1332 );
1333 let staging_error = object_store
1334 .inner
1335 .head(&pending.path)
1336 .await
1337 .expect_err("repair should garbage-collect the retained staging object");
1338 assert!(matches!(staging_error, ObjectStoreError::NotFound { .. }));
1339 }
1340
1341 #[tokio::test]
1342 async fn test_concurrent_finalizers_return_but_do_not_persist_generations() {
1343 let external_store = Arc::new(TestExternalManifestStore::blocking_first_final_publish());
1344 let handler = ExternalManifestCommitHandler {
1345 external_manifest_store: external_store.clone(),
1346 };
1347 let object_store = ObjectStore::memory();
1348 let base_path = Path::from("dataset");
1349 let version = 1;
1350 let final_path = ManifestNamingScheme::V2.manifest_path(&base_path, version);
1351 let staging_path = make_staging_manifest_path(&final_path).unwrap();
1352 let manifest_bytes = Bytes::from_static(b"immutable manifest bytes");
1353
1354 object_store
1355 .inner
1356 .put(&staging_path, manifest_bytes.clone().into())
1357 .await
1358 .unwrap();
1359 let staging_meta = object_store.inner.head(&staging_path).await.unwrap();
1360
1361 let writer_store = object_store.inner.clone();
1362 let writer_external_store = external_store.clone();
1363 let writer_base_path = base_path.clone();
1364 let writer_staging_path = staging_path.clone();
1365 let writer_e_tag = staging_meta.e_tag.clone();
1366 let writer = tokio::spawn(async move {
1367 writer_external_store
1368 .put(
1369 &writer_base_path,
1370 version,
1371 &writer_staging_path,
1372 staging_meta.size,
1373 writer_e_tag,
1374 writer_store.as_ref(),
1375 ManifestNamingScheme::V2,
1376 )
1377 .await
1378 });
1379
1380 tokio::time::timeout(
1381 std::time::Duration::from_secs(5),
1382 external_store.first_final_publish_started.notified(),
1383 )
1384 .await
1385 .expect("the direct finalizer should pause after COPY");
1386
1387 let first_generation = object_store.inner.head(&final_path).await.unwrap();
1388 let reservation = external_store
1389 .get_manifest_location(base_path.as_ref(), version)
1390 .await
1391 .unwrap();
1392 assert_eq!(reservation.path, staging_path);
1393 assert_eq!(reservation.e_tag, None);
1394
1395 let reader_location = handler
1403 .resolve_version_location(&base_path, version, object_store.inner.as_ref())
1404 .await
1405 .unwrap();
1406
1407 external_store.release_first_final_publish.notify_one();
1408 let writer_location = writer.await.unwrap().unwrap();
1409 let final_meta = object_store.inner.head(&final_path).await.unwrap();
1410 let final_bytes = object_store
1411 .inner
1412 .get(&final_path)
1413 .await
1414 .unwrap()
1415 .bytes()
1416 .await
1417 .unwrap();
1418 let indexed = external_store
1419 .get_manifest_location(base_path.as_ref(), version)
1420 .await
1421 .unwrap();
1422
1423 assert_eq!(final_bytes, manifest_bytes);
1424 assert_ne!(
1425 first_generation.e_tag, final_meta.e_tag,
1426 "the deterministic race must create a new physical generation"
1427 );
1428 assert_eq!(writer_location.e_tag, first_generation.e_tag);
1429 assert_eq!(reader_location.e_tag, final_meta.e_tag);
1430 assert_eq!(indexed.path, final_path);
1431 assert_eq!(indexed.size, Some(final_meta.size));
1432 assert_eq!(
1433 indexed.e_tag, None,
1434 "all finalizers must publish the same generation-independent tuple"
1435 );
1436
1437 let resolved = handler
1438 .resolve_version_location(&base_path, version, object_store.inner.as_ref())
1439 .await
1440 .expect("the finalized manifest must remain readable after the race");
1441 assert_eq!(resolved.e_tag, final_meta.e_tag);
1442 }
1443
1444 #[tokio::test]
1445 async fn test_lost_external_store_response_retains_staging_manifest() {
1446 let external_store = Arc::new(TestExternalManifestStore::new(true));
1447 let handler = ExternalManifestCommitHandler {
1448 external_manifest_store: external_store.clone(),
1449 };
1450 let object_store = ObjectStore::memory();
1451 let base_path = Path::from("dataset");
1452 let mut manifest = test_manifest();
1453
1454 let commit_error = handler
1455 .commit(
1456 &mut manifest,
1457 None,
1458 &base_path,
1459 &object_store,
1460 write_manifest_file_to_path,
1461 ManifestNamingScheme::V2,
1462 None,
1463 )
1464 .await
1465 .expect_err("the simulated response loss must be surfaced");
1466 assert!(matches!(commit_error, CommitError::CommitConflict));
1467
1468 let staging_path = Path::from(external_store.get("dataset", 1).await.unwrap());
1469 object_store.inner.head(&staging_path).await.unwrap();
1470
1471 let resolved = handler
1472 .resolve_version_location(&base_path, 1, object_store.inner.as_ref())
1473 .await
1474 .expect("the retained staging manifest must allow finalization");
1475 assert_eq!(
1476 resolved.path,
1477 ManifestNamingScheme::V2.manifest_path(&base_path, 1)
1478 );
1479 object_store.inner.head(&resolved.path).await.unwrap();
1480 }
1481
1482 #[tokio::test]
1483 async fn test_finalization_returns_etag_without_persisting_it() {
1484 let external_store = Arc::new(TestExternalManifestStore::new(false));
1485 let handler = ExternalManifestCommitHandler {
1486 external_manifest_store: external_store.clone(),
1487 };
1488 let object_store = ObjectStore::memory();
1489 let base_path = Path::from("dataset");
1490 let mut manifest = test_manifest();
1491 let version = manifest.version;
1492 let final_path = ManifestNamingScheme::V2.manifest_path(&base_path, version);
1493
1494 let committed = handler
1495 .commit(
1496 &mut manifest,
1497 None,
1498 &base_path,
1499 &object_store,
1500 write_manifest_file_to_path,
1501 ManifestNamingScheme::V2,
1502 None,
1503 )
1504 .await
1505 .expect("the generic workflow should commit the canonical manifest");
1506 assert_eq!(committed.path, final_path);
1507 let final_meta = object_store.inner.head(&final_path).await.unwrap();
1508 assert_eq!(
1509 committed.e_tag, final_meta.e_tag,
1510 "the freshly committed Dataset needs the observed generation for cache separation"
1511 );
1512
1513 let indexed = external_store
1514 .get_manifest_location(base_path.as_ref(), version)
1515 .await
1516 .expect("the external index must advance after the canonical copy");
1517 assert_eq!(indexed.path, final_path);
1518 assert_eq!(
1519 indexed.e_tag, None,
1520 "the external index must remain independent of physical generations"
1521 );
1522 }
1523
1524 #[tokio::test]
1525 async fn test_missing_staging_verifies_existing_final_manifest() {
1526 let object_store = ObjectStore::memory();
1527 let staging_path = Path::from("dataset/_versions/1.manifest-missing");
1528 let final_path = Path::from("dataset/_versions/1.manifest");
1529 let manifest_bytes = Bytes::from_static(b"immutable manifest bytes");
1530 object_store
1531 .inner
1532 .put(&final_path, manifest_bytes.clone().into())
1533 .await
1534 .unwrap();
1535 let final_meta = object_store.inner.head(&final_path).await.unwrap();
1536
1537 let recovered_e_tag = copy_or_verify_final_manifest(
1538 object_store.inner.as_ref(),
1539 &staging_path,
1540 &final_path,
1541 1,
1542 manifest_bytes.len() as u64,
1543 )
1544 .await
1545 .expect("an existing canonical manifest should prove another helper finalized it");
1546
1547 assert_eq!(recovered_e_tag, final_meta.e_tag);
1548 }
1549
1550 #[tokio::test]
1551 async fn test_missing_staging_rejects_missing_final_manifest() {
1552 let object_store = ObjectStore::memory();
1553 let staging_path = Path::from("dataset/_versions/1.manifest-missing");
1554 let final_path = Path::from("dataset/_versions/1.manifest");
1555
1556 let error = copy_or_verify_final_manifest(
1557 object_store.inner.as_ref(),
1558 &staging_path,
1559 &final_path,
1560 1,
1561 42,
1562 )
1563 .await
1564 .expect_err("missing staging and canonical objects cannot establish a commit");
1565
1566 assert!(matches!(error, Error::NotFound { .. }), "{error:?}");
1567 assert!(error.to_string().contains(final_path.as_ref()), "{error}");
1568 }
1569
1570 #[tokio::test]
1571 async fn test_missing_staging_rejects_wrong_final_size() {
1572 let object_store = ObjectStore::memory();
1573 let staging_path = Path::from("dataset/_versions/1.manifest-missing");
1574 let final_path = Path::from("dataset/_versions/1.manifest");
1575 object_store
1576 .inner
1577 .put(&final_path, Bytes::from_static(b"wrong size").into())
1578 .await
1579 .unwrap();
1580
1581 let error = copy_or_verify_final_manifest(
1582 object_store.inner.as_ref(),
1583 &staging_path,
1584 &final_path,
1585 1,
1586 42,
1587 )
1588 .await
1589 .expect_err("a same-path object with the wrong size is not the selected manifest");
1590
1591 assert!(matches!(error, Error::CorruptFile { .. }), "{error:?}");
1592 assert!(
1593 error.to_string().contains("Manifest size mismatch"),
1594 "{error}"
1595 );
1596 }
1597
1598 #[tokio::test]
1599 async fn test_copy_failure_after_external_store_commit_retains_staging_manifest() {
1600 let external_store = Arc::new(TestExternalManifestStore::new(false));
1601 let handler = ExternalManifestCommitHandler {
1602 external_manifest_store: external_store.clone(),
1603 };
1604
1605 let mut object_store = ObjectStore::memory();
1606 let fail_next_copy = Arc::new(AtomicBool::new(true));
1607 let failed_copy_source = Arc::new(Mutex::new(None));
1608 let mut policy = ProxyObjectStorePolicy::new();
1609 let policy_fail_next_copy = fail_next_copy.clone();
1610 let policy_failed_copy_source = failed_copy_source.clone();
1611 policy.set_before_policy(
1612 "fail-copy-once",
1613 Arc::new(move |method, location| {
1614 if method == "copy" && policy_fail_next_copy.swap(false, Ordering::SeqCst) {
1615 *policy_failed_copy_source.lock().unwrap() = Some(location.clone());
1616 return Err(Error::io("simulated copy failure"));
1617 }
1618 Ok(())
1619 }),
1620 );
1621 let policy = Arc::new(Mutex::new(policy));
1622 object_store.inner = Arc::new(ProxyObjectStore::new(
1623 object_store.inner.clone(),
1624 policy.clone(),
1625 ));
1626
1627 let base_path = Path::from("dataset");
1628 let mut manifest = test_manifest();
1629 let version = manifest.version;
1630 let canonical_path = ManifestNamingScheme::V2.manifest_path(&base_path, version);
1631
1632 let commit_error = handler
1633 .commit(
1634 &mut manifest,
1635 None,
1636 &base_path,
1637 &object_store,
1638 write_manifest_file_to_path,
1639 ManifestNamingScheme::V2,
1640 None,
1641 )
1642 .await
1643 .expect_err("the simulated copy failure must be surfaced");
1644 assert!(matches!(commit_error, CommitError::CommitConflict));
1645 assert!(
1646 !fail_next_copy.load(Ordering::SeqCst),
1647 "the one-shot copy failure must be consumed"
1648 );
1649
1650 let recorded_location = external_store
1651 .get_manifest_location(base_path.as_ref(), version)
1652 .await
1653 .expect("the external store must retain the committed staging location");
1654 let staging_path = failed_copy_source
1655 .lock()
1656 .unwrap()
1657 .clone()
1658 .expect("the failure must be injected at copy(staging, canonical)");
1659 assert_eq!(recorded_location.path, staging_path);
1660 object_store
1661 .inner
1662 .head(&staging_path)
1663 .await
1664 .expect("the winning staging manifest must be retained");
1665
1666 let canonical_error = object_store
1667 .inner
1668 .head(&canonical_path)
1669 .await
1670 .expect_err("copy failed before creating the canonical manifest");
1671 assert!(
1672 matches!(canonical_error, ObjectStoreError::NotFound { .. }),
1673 "unexpected canonical manifest error: {canonical_error}"
1674 );
1675
1676 policy.lock().unwrap().clear_before_policy("fail-copy-once");
1677 let resolved = handler
1678 .resolve_version_location(&base_path, version, object_store.inner.as_ref())
1679 .await
1680 .expect("the retained staging manifest must allow finalization");
1681 assert_eq!(resolved.path, canonical_path);
1682
1683 let finalized_location = external_store
1684 .get_manifest_location(base_path.as_ref(), version)
1685 .await
1686 .expect("the external store must publish the canonical location");
1687 assert_eq!(finalized_location.path, canonical_path);
1688 object_store
1689 .inner
1690 .head(&canonical_path)
1691 .await
1692 .expect("the canonical manifest must exist after finalization");
1693
1694 let staging_error = object_store
1695 .inner
1696 .head(&staging_path)
1697 .await
1698 .expect_err("successful finalization must clean up the staging manifest");
1699 assert!(
1700 matches!(staging_error, ObjectStoreError::NotFound { .. }),
1701 "unexpected staging manifest error: {staging_error}"
1702 );
1703 }
1704}