Merge pull request #42 from ilikepi63/feat/make-partition-key-array-of-bytes

ilikepi63 · web-flow · commit 27f1e0dd576d · 2025-05-20T11:10:26.000+02:00
feat: Make Partition key array of bytes
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "riskless"
-version = "0.6.3"
+version = "0.7.1"
 edition = "2024"
 description = "A pure Rust implementation of Diskless Topics"
 license = "MIT / Apache-2.0"
diff --git a/examples/concurrent.rs b/examples/concurrent.rs
@@ -24,7 +24,7 @@ async fn main() {
             .collect(ProduceRequest {
                 request_id: 1,
                 topic: "example-topic".to_string(),
-                partition: 1,
+                partition: Vec::from(&1_u8.to_be_bytes()),
                 data: "hello".as_bytes().to_vec(),
             })
             .expect("");
@@ -55,7 +55,7 @@ async fn main() {
     let consume_response = consume(
         ConsumeRequest {
             topic: "example-topic".to_string(),
-            partition: 1,
+            partition: Vec::from(&1_u8.to_be_bytes()),
             offset: 0,
             max_partition_fetch_bytes: 0,
         },
diff --git a/examples/simple.rs b/examples/simple.rs
@@ -17,7 +17,7 @@ async fn main() {
     col.collect(ProduceRequest {
         request_id: 1,
         topic: "example-topic".to_string(),
-        partition: 1,
+        partition: Vec::from(&1_u8.to_be_bytes()),
         data: "hello".as_bytes().to_vec(),
     })
     .expect("");
@@ -31,7 +31,7 @@ async fn main() {
     let consume_response = consume(
         ConsumeRequest {
             topic: "example-topic".to_string(),
-            partition: 1,
+            partition: Vec::from(&1_u8.to_be_bytes()),
             offset: 0,
             max_partition_fetch_bytes: 0,
         },
diff --git a/src/batch_coordinator/mod.rs b/src/batch_coordinator/mod.rs
@@ -14,7 +14,7 @@ use crate::messages::CommitBatchRequest;
 
 /// Merged Topic/Partition identification struc .
 #[derive(Debug, Hash, PartialEq, Eq, Clone, Default)]
-pub struct TopicIdPartition(pub String, pub u64);
+pub struct TopicIdPartition(pub String, pub Vec<u8>);
 
 /// The type of the timestamp given.
 #[derive(Debug, Default, Clone)]
diff --git a/src/batch_coordinator/simple/mod.rs b/src/batch_coordinator/simple/mod.rs
@@ -55,11 +55,14 @@ impl SimpleBatchCoordinator {
     }
 
     /// Retrieves the directory of the given topic/partition.
-    fn partition_index_file_from_topic_dir(
-        topic_dir: &mut PathBuf,
-        partition: u64,
-    ) -> &mut PathBuf {
-        topic_dir.push(format!("{:0>20}.index", partition.to_string()));
+    fn partition_index_file_from_topic_dir<'a>(
+        topic_dir: &'a mut PathBuf,
+        partition: &'a [u8],
+    ) -> &'a mut PathBuf {
+        topic_dir.push(format!(
+            "{:0>20}.index",
+            partition.iter().map(|b| b.to_string()).collect::<String>()
+        ));
 
         (topic_dir) as _
     }
@@ -105,7 +108,7 @@ impl CommitFile for SimpleBatchCoordinator {
 
             let current_partition_file = Self::partition_index_file_from_topic_dir(
                 &mut current_topic_dir,
-                batch.topic_id_partition.1,
+                &batch.topic_id_partition.1,
             );
 
             let file = Self::open_or_create_file(current_partition_file);
@@ -162,7 +165,7 @@ impl FindBatches for SimpleBatchCoordinator {
 
             let current_partition_file = Self::partition_index_file_from_topic_dir(
                 &mut current_topic_dir,
-                request.topic_id_partition.1,
+                &request.topic_id_partition.1,
             );
 
             let file = Self::open_file(current_partition_file);
@@ -383,10 +386,10 @@ mod tests {
     #[test]
     fn test_partition_index_file_from_topic_dir() {
         let mut topic_dir = PathBuf::from("test_topic");
-        let partition = 42;
+        let partition = Vec::from(&42_u8.to_be_bytes());
 
         let result =
-            SimpleBatchCoordinator::partition_index_file_from_topic_dir(&mut topic_dir, partition);
+            SimpleBatchCoordinator::partition_index_file_from_topic_dir(&mut topic_dir, &partition);
 
         assert_eq!(
             result.to_str().expect(""),
@@ -461,11 +464,11 @@ mod tests {
         let whole_dir = temp_dir.clone();
 
         let topic = "test_topic".to_string();
-        let partition = 1;
+        let partition = Vec::from(&1_u8.to_be_bytes());
 
         let object_key = Uuid::new_v4().into_bytes();
         let batches = vec![CommitBatchRequest {
-            topic_id_partition: TopicIdPartition(topic.clone(), partition),
+            topic_id_partition: TopicIdPartition(topic.clone(), partition.clone()),
             byte_offset: 0,
             size: 100,
             request_id: 1,
@@ -479,9 +482,10 @@ mod tests {
             last_sequence: 0,
         }];
 
-        let expected_file_path = temp_dir
-            .join(&topic)
-            .join(format!("{:0>20}.index", partition));
+        let expected_file_path = temp_dir.join(&topic).join(format!(
+            "{:0>20}.index",
+            partition.iter().map(|b| b.to_string()).collect::<String>()
+        ));
 
         // File shouldn't exist yet
         assert!(!expected_file_path.exists());
@@ -518,15 +522,15 @@ mod tests {
         let _ = std::fs::create_dir(index_path);
 
         let topic = "test_topic".to_string();
-        let partition = 1;
+        let partition = Vec::from(&1_u8.to_be_bytes());
 
         // First, create an index file with some data
         let object_key = Uuid::new_v4().into_bytes();
         let offset = 0;
         let size = 100;
 
         let batches = vec![CommitBatchRequest {
-            topic_id_partition: TopicIdPartition(topic.clone(), partition),
+            topic_id_partition: TopicIdPartition(topic.clone(), partition.clone()),
             byte_offset: offset,
             size,
             request_id: 1,
@@ -582,7 +586,7 @@ mod tests {
         let _ = std::fs::create_dir(&index_path);
 
         let topic = "test_topic".to_string();
-        let partition = 1;
+        let partition = Vec::from(&1_u8.to_be_bytes());
 
         // First, create an index file with some data
         let object_key = Uuid::new_v4().into_bytes();
@@ -591,7 +595,7 @@ mod tests {
         let size = 100;
 
         let batches = vec![CommitBatchRequest {
-            topic_id_partition: TopicIdPartition(topic.clone(), partition),
+            topic_id_partition: TopicIdPartition(topic.clone(), partition.clone()),
             byte_offset: offset,
             size,
             request_id: 1,
@@ -608,7 +612,7 @@ mod tests {
         coordinator.commit_file(object_key, 1, 100, batches).await;
 
         let batches = vec![CommitBatchRequest {
-            topic_id_partition: TopicIdPartition(topic.clone(), partition),
+            topic_id_partition: TopicIdPartition(topic.clone(), partition.clone()),
             byte_offset: 1,
             size,
             request_id: 1,
@@ -627,7 +631,7 @@ mod tests {
             .await;
 
         let batches = vec![CommitBatchRequest {
-            topic_id_partition: TopicIdPartition(topic.clone(), partition),
+            topic_id_partition: TopicIdPartition(topic.clone(), partition.clone()),
             byte_offset: 1,
             size,
             request_id: 2,
@@ -649,15 +653,18 @@ mod tests {
 
         index_path.push(&topic);
 
-        index_path.push(format!("{:0>20}.index", partition.to_string()));
+        index_path.push(format!(
+            "{:0>20}.index",
+            partition.iter().map(|b| b.to_string()).collect::<String>()
+        ));
 
         let data = std::fs::read(index_path)?;
 
         assert_eq!(data.len(), Index::packed_size() * 3);
 
         // Now try to find the batch
         let find_requests = vec![FindBatchRequest {
-            topic_id_partition: TopicIdPartition(topic.clone(), partition),
+            topic_id_partition: TopicIdPartition(topic.clone(), partition.clone()),
             offset,
             max_partition_fetch_bytes: 1024,
         }];
@@ -675,7 +682,7 @@ mod tests {
         assert_eq!(batch.metadata.byte_size, size);
 
         let find_requests = vec![FindBatchRequest {
-            topic_id_partition: TopicIdPartition(topic.clone(), partition),
+            topic_id_partition: TopicIdPartition(topic.clone(), partition.clone()),
             offset: 1,
             max_partition_fetch_bytes: 1024,
         }];
@@ -706,7 +713,10 @@ mod tests {
         let coordinator = SimpleBatchCoordinator::new(temp_dir.to_str().expect("").to_string());
 
         let find_requests = vec![FindBatchRequest {
-            topic_id_partition: TopicIdPartition("nonexistent_topic".to_string(), 1),
+            topic_id_partition: TopicIdPartition(
+                "nonexistent_topic".to_string(),
+                Vec::from(&1_u8.to_be_bytes()),
+            ),
             offset: 0,
             max_partition_fetch_bytes: 1024,
         }];
diff --git a/src/messages/batch_coordinate.rs b/src/messages/batch_coordinate.rs
@@ -6,7 +6,7 @@ pub struct BatchCoordinate {
     /// The topic for this record.
     pub topic: String,
     /// The partition for this record.
-    pub partition: u64,
+    pub partition: Vec<u8>,
     /// The base_offset inside of the file at which this record exists.
     pub base_offset: u64,
     /// The offset at which this record exists.
diff --git a/src/messages/commit_batch_request.rs b/src/messages/commit_batch_request.rs
@@ -39,7 +39,7 @@ impl From<&BatchCoordinate> for CommitBatchRequest {
         // Everything that is defaulted is unknown for now.
         CommitBatchRequest {
             request_id: value.request.request_id,
-            topic_id_partition: TopicIdPartition(value.topic.clone(), value.partition),
+            topic_id_partition: TopicIdPartition(value.topic.clone(), value.partition.clone()),
             byte_offset: value.offset,
             size: value.size,
             base_offset: value.base_offset,
diff --git a/src/messages/consume_request.rs b/src/messages/consume_request.rs
@@ -9,7 +9,7 @@ pub struct ConsumeRequest {
     /// The topic that this request consumes from.
     pub topic: String,
     /// The partition that this request consumes from.
-    pub partition: u64,
+    pub partition: Vec<u8>,
     /// The offset from which this request consumes from.
     pub offset: u64,
     /// The maximum amount of bytes to retrieve from a partition.
diff --git a/src/messages/consume_response.rs b/src/messages/consume_response.rs
@@ -11,7 +11,7 @@ pub struct ConsumeBatch {
     /// The topic of the record.
     pub topic: String,
     /// Partition of a record.
-    pub partition: u64,
+    pub partition: Vec<u8>,
     /// The offset of this record.
     pub offset: u64,
     /// The maximum amount of bytes to retrieve from each partition.
@@ -46,7 +46,7 @@ impl TryFrom<(FindBatchResponse, &BatchInfo, &Bytes)> for ConsumeBatch {
 
         let batch = ConsumeBatch {
             topic: batch_info.metadata.topic_id_partition.0.clone(),
-            partition: batch_info.metadata.topic_id_partition.1,
+            partition: batch_info.metadata.topic_id_partition.1.clone(),
             offset: find_batch_response.log_start_offset,
             max_partition_fetch_bytes: 0,
             data,
diff --git a/src/messages/delete_record_request.rs b/src/messages/delete_record_request.rs
@@ -6,7 +6,7 @@ pub struct DeleteRecordsRequest {
     /// The topic of the record.
     pub topic: String,
     /// Partition of a record.
-    pub partition: u64,
+    pub partition: Vec<u8>,
     /// The offset of this record.
     pub offset: u64,
 }
diff --git a/src/messages/produce_request.rs b/src/messages/produce_request.rs
@@ -12,7 +12,7 @@ pub struct ProduceRequest {
     /// The topic to which this data will be produced.
     pub topic: String,
     /// The partition on which this data will be produced.
-    pub partition: u64,
+    pub partition: Vec<u8>,
     /// The data that will be added to the partition.
     pub data: Vec<u8>,
 }
@@ -60,7 +60,7 @@ impl ProduceRequestCollection {
 
     /// Collect a produce request into this struct.
     pub fn collect(&self, req: ProduceRequest) -> RisklessResult<()> {
-        let topic_id_partition = TopicIdPartition(req.topic.clone(), req.partition);
+        let topic_id_partition = TopicIdPartition(req.topic.clone(), req.partition.clone());
 
         let entry = self.inner.entry(topic_id_partition);
 
diff --git a/src/shared_log_segment/mod.rs b/src/shared_log_segment/mod.rs
@@ -100,7 +100,7 @@ impl TryFrom<ProduceRequestCollection> for SharedLogSegment {
 
                 batch_coords.push(BatchCoordinate {
                     topic: req.topic.clone(),
-                    partition: req.partition,
+                    partition: req.partition.clone(),
                     base_offset,
                     offset,
                     size: size.try_into()?,
@@ -149,7 +149,7 @@ mod tests {
         collection.collect(ProduceRequest {
             request_id: 1,
             topic: "test".to_string(),
-            partition: 0,
+            partition: Vec::from(&0_u8.to_be_bytes()),
             data: vec![1, 2, 3],
         })?;
 
@@ -162,7 +162,7 @@ mod tests {
 
         let coord = &segment.0[0];
         assert_eq!(coord.topic, "test");
-        assert_eq!(coord.partition, 0);
+        assert_eq!(coord.partition, Vec::from(&0_u8.to_be_bytes()));
         assert_eq!(
             coord.base_offset,
             SharedLogSegmentHeaderV1::size().try_into().expect("")
@@ -191,22 +191,22 @@ mod tests {
         collection.collect(ProduceRequest {
             request_id: 1,
             topic: "test".to_string(),
-            partition: 0,
+            partition: Vec::from(&0_u8.to_be_bytes()),
             data: vec![1, 2, 3],
         })?;
 
         collection.collect(ProduceRequest {
             request_id: 2,
             topic: "test".to_string(),
-            partition: 0,
+            partition: Vec::from(&0_u8.to_be_bytes()),
             data: vec![4, 5],
         })?;
 
         // Partition 1
         collection.collect(ProduceRequest {
             request_id: 3,
             topic: "test".to_string(),
-            partition: 1,
+            partition: Vec::from(&1_u8.to_be_bytes()),
             data: vec![6, 7, 8, 9],
         })?;
 
@@ -250,7 +250,7 @@ mod tests {
         let request = ProduceRequest {
             request_id: 1,
             topic: "large".to_string(),
-            partition: 0,
+            partition: Vec::from(&0_u8.to_be_bytes()),
             data: large_data.clone(),
         };
         collection.collect(request)?;
diff --git a/tests/broker.rs b/tests/broker.rs
