Skip to content

Commit d2c4de2

Browse files
committed
fix: preserve metadata lifecycle freshness
1 parent e3eda32 commit d2c4de2

3 files changed

Lines changed: 249 additions & 46 deletions

File tree

operations/src/process_placements.rs

Lines changed: 154 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ use aruna_core::storage_entries::{
1818
};
1919
use aruna_core::structs::{MetadataRegistryRecord, PlacementRef, RealmConfigDocument, RealmId};
2020
use aruna_core::task::TaskEvent;
21-
use aruna_core::types::{Effects, Key};
21+
use aruna_core::types::{Effects, Key, TxnId};
2222
use smallvec::smallvec;
2323
use 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

Comments
 (0)