From 3397bc79e3e20231c6601c0ec1cff3017674fdc8 Mon Sep 17 00:00:00 2001
From: aradwann <aradwann@proton.me>
Date: Wed, 6 Nov 2024 23:44:22 +0200
Subject: [PATCH] add serialized_size to StorageEncode

---
 crates/storage-api/src/storage.rs |  3 +++
 crates/types/src/storage.rs       | 33 +++++++++++++++++++++++++++++--
 2 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/crates/storage-api/src/storage.rs b/crates/storage-api/src/storage.rs
index 56c051446..b8ac4ac04 100644
--- a/crates/storage-api/src/storage.rs
+++ b/crates/storage-api/src/storage.rs
@@ -31,6 +31,9 @@ macro_rules! protobuf_storage_encode_decode {
                     restate_types::storage::StorageEncodeError::EncodeValue(err.into())
                 })
             }
+            fn serialized_size(&self) -> usize {
+                bytes::BytesMut::default().len() // todo: to be calculated
+            }
         }
 
         impl restate_types::storage::StorageDecode for $ty {
diff --git a/crates/types/src/storage.rs b/crates/types/src/storage.rs
index f97e5408b..38d32aa62 100644
--- a/crates/types/src/storage.rs
+++ b/crates/types/src/storage.rs
@@ -77,9 +77,16 @@ impl StorageCodec {
         value: &T,
         buf: &mut BytesMut,
     ) -> Result<(), StorageEncodeError> {
-        // write codec
+        // Write codec version byte
         buf.put_u8(value.default_codec().into());
-        // encode value
+
+        // Calculate the exact size required for serialization
+        let size = value.serialized_size();
+
+        // Ensure the buffer has enough capacity
+        buf.reserve(size);
+
+        // Encode the value into the buffer
         value.encode(buf)
     }
 
@@ -119,6 +126,9 @@ pub trait StorageEncode: DowncastSync {
 
     /// Codec which is used when encode new values.
     fn default_codec(&self) -> StorageCodecKind;
+
+    /// Returns the expected serialized size in bytes of the encoded value.
+    fn serialized_size(&self) -> usize;
 }
 impl_downcast!(sync StorageEncode);
 
@@ -153,6 +163,12 @@ macro_rules! flexbuffers_storage_encode_decode {
                 $crate::storage::encode_as_flexbuffers(self, buf)
                     .map_err(|err| $crate::storage::StorageEncodeError::EncodeValue(err.into()))
             }
+
+
+            fn serialized_size(&self) -> usize {
+                bytes::BytesMut::default().len()  // todo: to be calculated
+            }
+
         }
 
         impl $crate::storage::StorageDecode for $name {
@@ -204,6 +220,12 @@ impl StorageEncode for PolyBytes {
     fn default_codec(&self) -> StorageCodecKind {
         StorageCodecKind::FlexbuffersSerde
     }
+    fn serialized_size(&self) -> usize {
+        match self {
+            PolyBytes::Bytes(bytes) => bytes.len(),
+            PolyBytes::Typed(typed) => typed.serialized_size(),
+        }
+    }
 }
 
 /// SerializeAs/DeserializeAs to implement ser/de trait for [`PolyBytes`]
@@ -266,6 +288,10 @@ impl StorageEncode for String {
         buf.put_slice(my_bytes);
         Ok(())
     }
+
+    fn serialized_size(&self) -> usize {
+        4 + self.len() // 4 bytes for length + string byte length
+    }
 }
 impl StorageDecode for String {
     fn decode<B: ::bytes::Buf>(
@@ -332,6 +358,9 @@ impl StorageEncode for bytes::Bytes {
         buf.put_slice(&self[..]);
         Ok(())
     }
+    fn serialized_size(&self) -> usize {
+        4 + self.len() // 4 bytes for length + byte slice length
+    }
 }
 
 /// Utility method to encode a [`Serialize`] type as flexbuffers using serde.