core: add lease epoch fencing for holder commands (#90)

* core: add lease epoch fencing for holder commands

* Address CodeRabbit feedback

Author: skel84

crates/allocdb-bench/src/lib.rs

@@ -1128,6 +1128,7 @@ fn release_request(
         command: Command::Release {
             reservation_id,
             holder_id,
+            lease_epoch: 1,
         },
     }
 }

crates/allocdb-core/src/command.rs

@@ -37,10 +37,12 @@ pub enum Command {
     Confirm {
         reservation_id: ReservationId,
         holder_id: HolderId,
+        lease_epoch: u64,
     },
     Release {
         reservation_id: ReservationId,
         holder_id: HolderId,
+        lease_epoch: u64,
     },
     Expire {
         reservation_id: ReservationId,
@@ -98,18 +100,22 @@ impl Command {
             Self::Confirm {
                 reservation_id,
                 holder_id,
+                lease_epoch,
             } => {
                 state = mix(state, 4);
                 state = mix(state, reservation_id.get());
-                mix(state, holder_id.get())
+                state = mix(state, holder_id.get());
+                mix(state, u128::from(*lease_epoch))
             }
             Self::Release {
                 reservation_id,
                 holder_id,
+                lease_epoch,
             } => {
                 state = mix(state, 5);
                 state = mix(state, reservation_id.get());
-                mix(state, holder_id.get())
+                state = mix(state, holder_id.get());
+                mix(state, u128::from(*lease_epoch))
             }
             Self::Expire {
                 reservation_id,

crates/allocdb-core/src/command_codec.rs

@@ -92,18 +92,22 @@ fn encode_command(bytes: &mut Vec<u8>, command: &Command) {
         Command::Confirm {
             reservation_id,
             holder_id,
+            lease_epoch,
         } => {
             bytes.push(4);
             bytes.extend_from_slice(&reservation_id.get().to_le_bytes());
             bytes.extend_from_slice(&holder_id.get().to_le_bytes());
+            bytes.extend_from_slice(&lease_epoch.to_le_bytes());
         }
         Command::Release {
             reservation_id,
             holder_id,
+            lease_epoch,
         } => {
             bytes.push(5);
             bytes.extend_from_slice(&reservation_id.get().to_le_bytes());
             bytes.extend_from_slice(&holder_id.get().to_le_bytes());
+            bytes.extend_from_slice(&lease_epoch.to_le_bytes());
         }
         Command::Expire {
             reservation_id,
@@ -142,10 +146,12 @@ fn decode_command(cursor: &mut Cursor<'_>) -> Result<Command, CommandCodecError>
         4 => Ok(Command::Confirm {
             reservation_id: ReservationId(cursor.read_u128()?),
             holder_id: HolderId(cursor.read_u128()?),
+            lease_epoch: decode_optional_legacy_epoch(cursor)?,
         }),
         5 => Ok(Command::Release {
             reservation_id: ReservationId(cursor.read_u128()?),
             holder_id: HolderId(cursor.read_u128()?),
+            lease_epoch: decode_optional_legacy_epoch(cursor)?,
         }),
         6 => Ok(Command::Expire {
             reservation_id: ReservationId(cursor.read_u128()?),
@@ -199,6 +205,18 @@ impl<'a> Cursor<'a> {
     fn read_u128(&mut self) -> Result<u128, CommandCodecError> {
         Ok(u128::from_le_bytes(self.read_exact::<16>()?))
     }
+
+    fn remaining(&self) -> usize {
+        self.bytes.len().saturating_sub(self.offset)
+    }
+}
+
+fn decode_optional_legacy_epoch(cursor: &mut Cursor<'_>) -> Result<u64, CommandCodecError> {
+    match cursor.remaining() {
+        0 => Ok(1),
+        8 => cursor.read_u64(),
+        _ => Err(CommandCodecError::InvalidLayout),
+    }
 }
 
 #[cfg(test)]
@@ -254,6 +272,59 @@ mod tests {
         assert_eq!(decoded, request);
     }
 
+    #[test]
+    fn holder_commands_round_trip_with_lease_epoch() {
+        let requests = [
+            ClientRequest {
+                operation_id: OperationId(11),
+                client_id: ClientId(12),
+                command: Command::Confirm {
+                    reservation_id: ReservationId(13),
+                    holder_id: HolderId(14),
+                    lease_epoch: 15,
+                },
+            },
+            ClientRequest {
+                operation_id: OperationId(21),
+                client_id: ClientId(22),
+                command: Command::Release {
+                    reservation_id: ReservationId(23),
+                    holder_id: HolderId(24),
+                    lease_epoch: 25,
+                },
+            },
+        ];
+
+        for request in requests {
+            let decoded = decode_client_request(&encode_client_request(&request)).unwrap();
+            assert_eq!(decoded, request);
+        }
+    }
+
+    #[test]
+    fn decoder_accepts_legacy_holder_command_without_lease_epoch() {
+        let mut bytes = Vec::new();
+        bytes.extend_from_slice(&OperationId(11).get().to_le_bytes());
+        bytes.extend_from_slice(&ClientId(12).get().to_le_bytes());
+        bytes.push(4);
+        bytes.extend_from_slice(&ReservationId(13).get().to_le_bytes());
+        bytes.extend_from_slice(&HolderId(14).get().to_le_bytes());
+
+        let decoded = decode_client_request(&bytes).unwrap();
+        assert_eq!(
+            decoded,
+            ClientRequest {
+                operation_id: OperationId(11),
+                client_id: ClientId(12),
+                command: Command::Confirm {
+                    reservation_id: ReservationId(13),
+                    holder_id: HolderId(14),
+                    lease_epoch: 1,
+                },
+            }
+        );
+    }
+
     #[test]
     fn decoder_rejects_truncated_payload() {
         let request = ClientRequest {

crates/allocdb-core/src/recovery_tests.rs

@@ -133,6 +133,7 @@ fn recover_allocdb_skips_frames_covered_by_snapshot() {
         command: Command::Confirm {
             reservation_id: ReservationId(2),
             holder_id: HolderId(5),
+            lease_epoch: 1,
         },
     };
 

crates/allocdb-core/src/result.rs

@@ -19,13 +19,15 @@ pub enum ResultCode {
     OperationConflict,
     InvalidState,
     HolderMismatch,
+    StaleEpoch,
     SlotOverflow,
 }
 
 #[derive(Clone, Copy, Debug, Eq, PartialEq)]
 pub struct CommandOutcome {
     pub result_code: ResultCode,
     pub reservation_id: Option<ReservationId>,
+    pub lease_epoch: Option<u64>,
     pub deadline_slot: Option<Slot>,
 }
 
@@ -35,6 +37,7 @@ impl CommandOutcome {
         Self {
             result_code,
             reservation_id: None,
+            lease_epoch: None,
             deadline_slot: None,
         }
     }
@@ -48,6 +51,22 @@ impl CommandOutcome {
         Self {
             result_code,
             reservation_id: Some(reservation_id),
+            lease_epoch: None,
+            deadline_slot: Some(deadline_slot),
+        }
+    }
+
+    #[must_use]
+    pub const fn with_reservation_epoch(
+        result_code: ResultCode,
+        reservation_id: ReservationId,
+        lease_epoch: u64,
+        deadline_slot: Slot,
+    ) -> Self {
+        Self {
+            result_code,
+            reservation_id: Some(reservation_id),
+            lease_epoch: Some(lease_epoch),
             deadline_slot: Some(deadline_slot),
         }
     }

crates/allocdb-core/src/snapshot.rs

@@ -18,7 +18,7 @@ mod issue_30_tests;
 mod tests;
 
 const MAGIC: u32 = 0x4144_4253;
-const VERSION: u16 = 3;
+const VERSION: u16 = 4;
 
 #[derive(Clone, Debug, Eq, PartialEq)]
 pub struct Snapshot {

crates/allocdb-core/src/snapshot_codec.rs

@@ -54,6 +54,7 @@ pub(super) fn encode_snapshot(snapshot: &Snapshot) -> Vec<u8> {
         bytes.extend_from_slice(&reservation.reservation_id.get().to_le_bytes());
         bytes.extend_from_slice(&reservation.resource_id.get().to_le_bytes());
         bytes.extend_from_slice(&reservation.holder_id.get().to_le_bytes());
+        bytes.extend_from_slice(&reservation.lease_epoch.to_le_bytes());
         bytes.push(encode_reservation_state(reservation.state));
         bytes.extend_from_slice(&reservation.created_lsn.get().to_le_bytes());
         bytes.extend_from_slice(&reservation.deadline_slot.get().to_le_bytes());
@@ -76,6 +77,7 @@ pub(super) fn encode_snapshot(snapshot: &Snapshot) -> Vec<u8> {
             &mut bytes,
             operation.result_reservation_id.map(ReservationId::get),
         );
+        encode_optional_u64(&mut bytes, operation.result_lease_epoch);
         encode_optional_u64(&mut bytes, operation.result_deadline_slot.map(Slot::get));
         bytes.extend_from_slice(&operation.applied_lsn.get().to_le_bytes());
         bytes.extend_from_slice(&operation.retire_after_slot.get().to_le_bytes());
@@ -103,7 +105,7 @@ pub(super) fn decode_snapshot(bytes: &[u8]) -> Result<Snapshot, SnapshotError> {
     }
 
     let version = cursor.read_u16()?;
-    if version != 1 && version != 2 && version != VERSION {
+    if version != 1 && version != 2 && version != 3 && version != VERSION {
         return Err(SnapshotError::InvalidVersion(version));
     }
 
@@ -143,6 +145,11 @@ pub(super) fn decode_snapshot(bytes: &[u8]) -> Result<Snapshot, SnapshotError> {
             command_fingerprint: cursor.read_u128()?,
             result_code: decode_result_code(cursor.read_u8()?)?,
             result_reservation_id: cursor.read_optional_u128()?.map(ReservationId),
+            result_lease_epoch: if version >= 4 {
+                cursor.read_optional_u64()?
+            } else {
+                None
+            },
             result_deadline_slot: cursor.read_optional_u64()?.map(Slot),
             applied_lsn: Lsn(cursor.read_u64()?),
             retire_after_slot: Slot(cursor.read_u64()?),
@@ -235,6 +242,7 @@ fn decode_snapshot_reservations(
             reservation_id: ReservationId(cursor.read_u128()?),
             resource_id: ResourceId(cursor.read_u128()?),
             holder_id: HolderId(cursor.read_u128()?),
+            lease_epoch: if version >= 4 { cursor.read_u64()? } else { 1 },
             state: decode_reservation_state(cursor.read_u8()?)?,
             created_lsn: Lsn(cursor.read_u64()?),
             deadline_slot: Slot(cursor.read_u64()?),
@@ -358,7 +366,8 @@ fn decode_result_code(tag: u8) -> Result<ResultCode, SnapshotError> {
         14 => Ok(ResultCode::OperationConflict),
         15 => Ok(ResultCode::InvalidState),
         16 => Ok(ResultCode::HolderMismatch),
-        17 => Ok(ResultCode::SlotOverflow),
+        17 => Ok(ResultCode::StaleEpoch),
+        18 => Ok(ResultCode::SlotOverflow),
         _ => Err(SnapshotError::InvalidStateTag(tag)),
     }
 }
@@ -382,6 +391,7 @@ fn encode_result_code(code: ResultCode) -> u8 {
         ResultCode::OperationConflict => 14,
         ResultCode::InvalidState => 15,
         ResultCode::HolderMismatch => 16,
-        ResultCode::SlotOverflow => 17,
+        ResultCode::StaleEpoch => 17,
+        ResultCode::SlotOverflow => 18,
     }
 }

crates/allocdb-core/src/snapshot_issue_30_tests.rs

@@ -36,6 +36,7 @@ fn expired_reservation(reservation_id: u128) -> ReservationRecord {
         reservation_id: ReservationId(reservation_id),
         resource_id: ResourceId(11),
         holder_id: HolderId(5),
+        lease_epoch: 2,
         state: ReservationState::Expired,
         created_lsn: Lsn(reservation_lsn),
         deadline_slot: Slot(5),
@@ -53,6 +54,7 @@ fn operation_record(operation_id: u128) -> OperationRecord {
         command_fingerprint: operation_id + 10,
         result_code: crate::result::ResultCode::Ok,
         result_reservation_id: None,
+        result_lease_epoch: None,
         result_deadline_slot: None,
         applied_lsn: Lsn(operation_lsn),
         retire_after_slot: Slot(operation_lsn + 8),