@@ -18,7 +18,7 @@ use aruna_core::storage_entries::{
1818} ;
1919use aruna_core:: structs:: { MetadataRegistryRecord , PlacementRef , RealmConfigDocument , RealmId } ;
2020use aruna_core:: task:: TaskEvent ;
21- use aruna_core:: types:: { Effects , Key } ;
21+ use aruna_core:: types:: { Effects , Key , TxnId } ;
2222use smallvec:: smallvec;
2323use thiserror:: Error ;
2424
@@ -54,6 +54,7 @@ pub struct ProcessPlacementsOperation {
5454 records : Vec < PendingDocumentPlacement > ,
5555 next_start_after : Option < Key > ,
5656 current : Option < CurrentPlacement > ,
57+ txn_id : Option < TxnId > ,
5758 retry_needed : bool ,
5859 metadata_refresh_queued : bool ,
5960 rearmed : bool ,
@@ -67,7 +68,10 @@ enum PlacementState {
6768 ListPending ,
6869 Publish ,
6970 StorePlacement ,
71+ StartMetadataTransaction ,
7072 ReadMetadataState ,
73+ WriteMetadataState ,
74+ CommitMetadataTransaction ,
7175 ScheduleMetadataRefresh ,
7276 ScheduleRetry ,
7377 Finish ,
@@ -113,6 +117,7 @@ impl ProcessPlacementsOperation {
113117 records : Vec :: new ( ) ,
114118 next_start_after : None ,
115119 current : None ,
120+ txn_id : None ,
116121 retry_needed : false ,
117122 metadata_refresh_queued : false ,
118123 rearmed : false ,
@@ -121,9 +126,10 @@ impl ProcessPlacementsOperation {
121126 }
122127
123128 fn fail ( & mut self , error : PlacementError ) -> Effects {
129+ let cleanup = self . abort ( ) ;
124130 self . state = PlacementState :: Error ;
125131 self . output = Some ( Err ( error) ) ;
126- smallvec ! [ ]
132+ cleanup
127133 }
128134
129135 fn unexpected_event ( & mut self , expected : & ' static str , got : String ) -> Effects {
@@ -285,26 +291,10 @@ impl ProcessPlacementsOperation {
285291 current. planned . desired_holder_count ,
286292 ) ;
287293 if current. metadata_holders_changed {
288- let document_id = match & current. planned . target {
289- DocumentSyncTarget :: MetadataDocumentLifecycle { document_id } => * document_id,
290- _ => unreachable ! ( "metadata holder changes only apply to lifecycle placements" ) ,
291- } ;
292- let target = current. planned . target . clone ( ) ;
293294 self . current = Some ( current) ;
294- self . state = PlacementState :: ReadMetadataState ;
295- return smallvec ! [ Effect :: Storage ( StorageEffect :: BatchRead {
296- reads: vec![
297- ( target. storage_keyspace( ) . to_string( ) , target. storage_key( ) , ) ,
298- (
299- DOCUMENT_SYNC_REVISION_KEYSPACE . to_string( ) ,
300- document_sync_revision_key( & target) ,
301- ) ,
302- (
303- METADATA_DOCUMENT_INDEX_KEYSPACE . to_string( ) ,
304- metadata_document_key( document_id) ,
305- ) ,
306- ] ,
307- txn_id: None ,
295+ self . state = PlacementState :: StartMetadataTransaction ;
296+ return smallvec ! [ Effect :: Storage ( StorageEffect :: StartTransaction {
297+ read: false ,
308298 } ) ] ;
309299 }
310300 match write_placement_effect ( & current. planned ) {
@@ -322,6 +312,11 @@ impl ProcessPlacementsOperation {
322312 let Some ( current) = self . current . take ( ) else {
323313 return self . emit_next_record ( ) ;
324314 } ;
315+ let Some ( txn_id) = self . txn_id else {
316+ return self . fail ( PlacementError :: Placement (
317+ "missing metadata replan transaction" . to_string ( ) ,
318+ ) ) ;
319+ } ;
325320 let mut writes = Vec :: new ( ) ;
326321 if let Some ( mut registry) = registry {
327322 let DocumentSyncTarget :: MetadataDocumentLifecycle { document_id } =
@@ -425,13 +420,48 @@ impl ProcessPlacementsOperation {
425420 Ok ( entry) => writes. push ( entry) ,
426421 Err ( error) => return self . fail ( error. into ( ) ) ,
427422 }
428- self . state = PlacementState :: StorePlacement ;
423+ self . state = PlacementState :: WriteMetadataState ;
429424 smallvec ! [ Effect :: Storage ( StorageEffect :: BatchWrite {
430425 writes,
431- txn_id: None ,
426+ txn_id: Some ( txn_id ) ,
432427 } ) ]
433428 }
434429
430+ fn emit_metadata_state_read ( & mut self , txn_id : TxnId ) -> Effects {
431+ let Some ( current) = self . current . as_ref ( ) else {
432+ return self . emit_next_record ( ) ;
433+ } ;
434+ let document_id = match & current. planned . target {
435+ DocumentSyncTarget :: MetadataDocumentLifecycle { document_id } => * document_id,
436+ _ => unreachable ! ( "metadata holder changes only apply to lifecycle placements" ) ,
437+ } ;
438+ let target = current. planned . target . clone ( ) ;
439+ self . txn_id = Some ( txn_id) ;
440+ self . state = PlacementState :: ReadMetadataState ;
441+ smallvec ! [ Effect :: Storage ( StorageEffect :: BatchRead {
442+ reads: vec![
443+ ( target. storage_keyspace( ) . to_string( ) , target. storage_key( ) ) ,
444+ (
445+ DOCUMENT_SYNC_REVISION_KEYSPACE . to_string( ) ,
446+ document_sync_revision_key( & target) ,
447+ ) ,
448+ (
449+ METADATA_DOCUMENT_INDEX_KEYSPACE . to_string( ) ,
450+ metadata_document_key( document_id) ,
451+ ) ,
452+ ] ,
453+ txn_id: Some ( txn_id) ,
454+ } ) ]
455+ }
456+
457+ fn retry_metadata_transaction ( & mut self ) -> Effects {
458+ self . txn_id = None ;
459+ self . current = None ;
460+ self . metadata_refresh_queued = false ;
461+ self . retry_needed = true ;
462+ self . finish_placement_store ( )
463+ }
464+
435465 fn finish_placement_store ( & mut self ) -> Effects {
436466 if self . retry_needed {
437467 self . rearmed = true ;
@@ -528,6 +558,18 @@ impl Operation for ProcessPlacementsOperation {
528558 Event :: Storage ( StorageEvent :: Error { error } ) => self . fail ( error. into ( ) ) ,
529559 other => self . unexpected_event ( "placement storage result" , format ! ( "{other:?}" ) ) ,
530560 } ,
561+ PlacementState :: StartMetadataTransaction => match event {
562+ Event :: Storage ( StorageEvent :: TransactionStarted { txn_id } ) => {
563+ self . emit_metadata_state_read ( txn_id)
564+ }
565+ Event :: Storage ( StorageEvent :: Error {
566+ error : StorageError :: TransactionConflict ,
567+ } ) => self . retry_metadata_transaction ( ) ,
568+ Event :: Storage ( StorageEvent :: Error { error } ) => self . fail ( error. into ( ) ) ,
569+ other => {
570+ self . unexpected_event ( "metadata transaction start result" , format ! ( "{other:?}" ) )
571+ }
572+ } ,
531573 PlacementState :: ReadMetadataState => match event {
532574 Event :: Storage ( StorageEvent :: BatchReadResult { values } ) => {
533575 let [ ( _, lifecycle) , ( _, lifecycle_revision) , ( _, registry) ] =
@@ -558,6 +600,41 @@ impl Operation for ProcessPlacementsOperation {
558600 self . unexpected_event ( "metadata state batch read result" , format ! ( "{other:?}" ) )
559601 }
560602 } ,
603+ PlacementState :: WriteMetadataState => match event {
604+ Event :: Storage ( StorageEvent :: BatchWriteResult { .. } ) => {
605+ let Some ( txn_id) = self . txn_id else {
606+ return self . fail ( PlacementError :: Placement (
607+ "missing metadata replan transaction" . to_string ( ) ,
608+ ) ) ;
609+ } ;
610+ self . state = PlacementState :: CommitMetadataTransaction ;
611+ smallvec ! [ Effect :: Storage ( StorageEffect :: CommitTransaction { txn_id } ) ]
612+ }
613+ Event :: Storage ( StorageEvent :: Error { error } ) => self . fail ( error. into ( ) ) ,
614+ other => {
615+ self . unexpected_event ( "metadata transaction write result" , format ! ( "{other:?}" ) )
616+ }
617+ } ,
618+ PlacementState :: CommitMetadataTransaction => match event {
619+ Event :: Storage ( StorageEvent :: TransactionCommitted { .. } ) => {
620+ self . txn_id = None ;
621+ if self . metadata_refresh_queued {
622+ self . state = PlacementState :: ScheduleMetadataRefresh ;
623+ smallvec ! [ schedule_outbox_drain_effect( ) ]
624+ } else {
625+ self . finish_placement_store ( )
626+ }
627+ }
628+ Event :: Storage ( StorageEvent :: Error {
629+ error : StorageError :: TransactionConflict ,
630+ } ) => self . retry_metadata_transaction ( ) ,
631+ Event :: Storage ( StorageEvent :: Error { error } ) => {
632+ self . txn_id = None ;
633+ self . fail ( error. into ( ) )
634+ }
635+ other => self
636+ . unexpected_event ( "metadata transaction commit result" , format ! ( "{other:?}" ) ) ,
637+ } ,
561638 PlacementState :: ScheduleMetadataRefresh => match event {
562639 Event :: Task ( TaskEvent :: TimerScheduled { .. } ) => {
563640 self . metadata_refresh_queued = false ;
@@ -601,7 +678,10 @@ impl Operation for ProcessPlacementsOperation {
601678 }
602679
603680 fn abort ( & mut self ) -> Effects {
604- smallvec ! [ ]
681+ match self . txn_id . take ( ) {
682+ Some ( txn_id) => smallvec ! [ Effect :: Storage ( StorageEffect :: AbortTransaction { txn_id } ) ] ,
683+ None => smallvec ! [ ] ,
684+ }
605685 }
606686}
607687
@@ -654,8 +734,18 @@ mod tests {
654734 ) -> Effects {
655735 assert ! ( matches!(
656736 effects. as_slice( ) ,
657- [ Effect :: Storage ( StorageEffect :: BatchRead { reads, .. } ) ]
658- if reads. len( ) == 3
737+ [ Effect :: Storage ( StorageEffect :: StartTransaction {
738+ read: false
739+ } ) ]
740+ ) ) ;
741+ let txn_id = Ulid :: from_bytes ( [ 14 ; 16 ] ) ;
742+ let effects = operation. step ( Event :: Storage ( StorageEvent :: TransactionStarted { txn_id } ) ) ;
743+ assert ! ( matches!(
744+ effects. as_slice( ) ,
745+ [ Effect :: Storage ( StorageEffect :: BatchRead {
746+ reads,
747+ txn_id: Some ( read_txn_id) ,
748+ } ) ] if reads. len( ) == 3 && * read_txn_id == txn_id
659749 ) ) ;
660750 operation. step ( Event :: Storage ( StorageEvent :: BatchReadResult {
661751 values : vec ! [
@@ -813,7 +903,7 @@ mod tests {
813903 }
814904
815905 #[ test]
816- fn metadata_replan_does_not_claim_replacement_without_transfer_source ( ) {
906+ fn metadata_replan_missing_source_and_transaction_conflict_remain_retryable ( ) {
817907 let realm_id = RealmId :: from_bytes ( [ 8u8 ; 32 ] ) ;
818908 let document_id = Ulid :: from_bytes ( [ 12 ; 16 ] ) ;
819909 let target = DocumentSyncTarget :: MetadataDocumentLifecycle { document_id } ;
@@ -866,8 +956,18 @@ mod tests {
866956 let effects = operation. step ( transfer. finalize ( ) ) ;
867957 assert ! ( matches!(
868958 effects. as_slice( ) ,
869- [ Effect :: Storage ( StorageEffect :: BatchRead { reads, .. } ) ]
870- if reads. len( ) == 3
959+ [ Effect :: Storage ( StorageEffect :: StartTransaction {
960+ read: false
961+ } ) ]
962+ ) ) ;
963+ let txn_id = Ulid :: from_bytes ( [ 15 ; 16 ] ) ;
964+ let effects = operation. step ( Event :: Storage ( StorageEvent :: TransactionStarted { txn_id } ) ) ;
965+ assert ! ( matches!(
966+ effects. as_slice( ) ,
967+ [ Effect :: Storage ( StorageEffect :: BatchRead {
968+ reads,
969+ txn_id: Some ( read_txn_id) ,
970+ } ) ] if reads. len( ) == 3 && * read_txn_id == txn_id
871971 ) ) ;
872972
873973 let registry = MetadataRegistryRecord {
@@ -902,9 +1002,16 @@ mod tests {
9021002 ) ,
9031003 ] ,
9041004 } ) ) ;
905- let [ Effect :: Storage ( StorageEffect :: BatchWrite { writes, .. } ) ] = effects. as_slice ( ) else {
1005+ let [
1006+ Effect :: Storage ( StorageEffect :: BatchWrite {
1007+ writes,
1008+ txn_id : Some ( write_txn_id) ,
1009+ } ) ,
1010+ ] = effects. as_slice ( )
1011+ else {
9061012 panic ! ( "expected atomic registry and placement update, got {effects:?}" ) ;
9071013 } ;
1014+ assert_eq ! ( * write_txn_id, txn_id) ;
9081015 let registry = writes
9091016 . iter ( )
9101017 . find ( |( key_space, _, _) | key_space == aruna_core:: keyspaces:: METADATA_INDEX_KEYSPACE )
@@ -918,6 +1025,23 @@ mod tests {
9181025 . expect ( "placement update is present" ) ;
9191026 assert_eq ! ( placement. selected_holders, vec![ node( 2 ) ] ) ;
9201027 assert ! ( operation. retry_needed) ;
1028+
1029+ let effects = operation. step ( Event :: Storage ( StorageEvent :: BatchWriteResult {
1030+ entries : Vec :: new ( ) ,
1031+ } ) ) ;
1032+ assert ! ( matches!(
1033+ effects. as_slice( ) ,
1034+ [ Effect :: Storage ( StorageEffect :: CommitTransaction { txn_id: commit_txn_id } ) ]
1035+ if * commit_txn_id == txn_id
1036+ ) ) ;
1037+ let effects = operation. step ( Event :: Storage ( StorageEvent :: Error {
1038+ error : StorageError :: TransactionConflict ,
1039+ } ) ) ;
1040+ assert ! ( matches!( effects. as_slice( ) , [ Effect :: Task ( _) ] ) ) ;
1041+ assert ! ( operation. rearmed) ;
1042+ assert ! ( operation. retry_needed) ;
1043+ assert ! ( operation. txn_id. is_none( ) ) ;
1044+ assert ! ( !operation. metadata_refresh_queued) ;
9211045 }
9221046
9231047 #[ test]
0 commit comments