Introduce TimerKind to support different types of timers

crates/partition-store/src/keys.rs

@@ -12,6 +12,7 @@ use anyhow::anyhow;
 use bytes::{Buf, BufMut, Bytes, BytesMut};
 use bytestring::ByteString;
 use prost::encoding::encoded_len_varint;
+use std::mem;
 use strum_macros::EnumIter;
 
 /// Every table key needs to have a key kind. This allows to multiplex different keys in the same
@@ -305,6 +306,7 @@ macro_rules! define_table_key {
 use crate::TableKind;
 pub(crate) use define_table_key;
 use restate_storage_api::deduplication_table::ProducerId;
+use restate_storage_api::timer_table::TimerKeyKind;
 use restate_storage_api::StorageError;
 use restate_types::identifiers::{InvocationUuid, PartitionId};
 
@@ -485,6 +487,85 @@ impl KeyCodec for ProducerId {
  }
 }
 
+impl KeyCodec for TimerKeyKind {
+ fn encode<B: BufMut>(&self, target: &mut B) {
+ assert!(
+ self.serialized_length() <= target.remaining_mut(),
+ "serialization buffer has not enough space to serialize TimerKind: '{}' bytes required",
+ self.serialized_length()
+ );
+ match self {
+ TimerKeyKind::Invoke { invocation_uuid } => {
+ target.put_u8(0);
+ invocation_uuid.encode(target);
+ }
+ TimerKeyKind::CompleteJournalEntry {
+ invocation_uuid,
+ journal_index,
+ } => {
+ target.put_u8(1);
+ invocation_uuid.encode(target);
+ journal_index.encode(target);
+ }
+ TimerKeyKind::CleanInvocationStatus { invocation_uuid } => {
+ target.put_u8(2);
+ invocation_uuid.encode(target);
+ }
+ }
+ }
+
+ fn decode<B: Buf>(source: &mut B) -> crate::partition_store::Result<Self> {
+ if source.remaining() < mem::size_of::<u8>() {
+ return Err(StorageError::Generic(anyhow!(
+ "TimerKind discriminator byte is missing"
+ )));
+ }
+
+ Ok(match source.get_u8() {
+ 0 => {
+ let invocation_uuid = InvocationUuid::decode(source)?;
+ TimerKeyKind::Invoke { invocation_uuid }
+ }
+ 1 => {
+ let invocation_uuid = InvocationUuid::decode(source)?;
+ let journal_index = u32::decode(source)?;
+ TimerKeyKind::CompleteJournalEntry {
+ invocation_uuid,
+ journal_index,
+ }
+ }
+ 2 => {
+ let invocation_uuid = InvocationUuid::decode(source)?;
+ TimerKeyKind::CleanInvocationStatus { invocation_uuid }
+ }
+ i => {
+ return Err(StorageError::Generic(anyhow!(
+ "Unknown discriminator for TimerKind: '{}'",
+ i
+ )))
+ }
+ })
+ }
+
+ fn serialized_length(&self) -> usize {
+ 1 + match self {
+ TimerKeyKind::Invoke { invocation_uuid } => {
+ KeyCodec::serialized_length(invocation_uuid)
+ }
+ TimerKeyKind::CompleteJournalEntry {
+ invocation_uuid,
+ journal_index,
+ } => {
+ KeyCodec::serialized_length(invocation_uuid)
+ + KeyCodec::serialized_length(journal_index)
+ }
+ TimerKeyKind::CleanInvocationStatus { invocation_uuid } => {
+ KeyCodec::serialized_length(invocation_uuid)
+ }
+ }
+ }
+}
+
 #[inline]
 fn write_delimited<B: BufMut>(source: impl AsRef<[u8]>, target: &mut B) {
  let source = source.as_ref();