Merge pull request #37 from ilikepi63/feat/decompose-batch-coordinator-trait

feat: Decompose BatchCoordinator Trait

Author: ilikepi63

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "riskless"
-version = "0.5.2"
+version = "0.6.1"
 edition = "2024"
 description = "A pure Rust implementation of Diskless Topics"
 license = "MIT / Apache-2.0"

Cargo.toml

@@ -213,62 +213,59 @@ pub struct DeleteFilesRequest {
     pub object_key_paths: HashSet<String>,
 }
 
-/// The BatchCoordinator trait.
-///
-/// This structure is responsible for handling the indexing of offsets within topic's partitions.
-/// It is designed to be reimplementable for custom usecases depending on what the specific need is.
+/// Trait for enabling the creation of topics and partitions.
 #[async_trait::async_trait]
-pub trait BatchCoordinator
-where
-    Self: Send + Sync + std::fmt::Debug,
-{
+pub trait CreateTopicAndPartitions {
     /// This operation is called when a Diskless partition
     /// (or a topic with one or more partitions) is created in the cluster.
     /// The Batch Coordinator initializes the corresponding logs.
     ///
     /// # Errors
     /// Returns an error if an unexpected error occurs.
     async fn create_topic_and_partitions(&self, requests: HashSet<CreateTopicAndPartitionsRequest>);
+}
 
+/// This operation is called by a broker after uploading the
+/// shared log segment object to the object storage.
+#[async_trait::async_trait]
+pub trait CommitFile
+where
+    Self: Send + Sync + std::fmt::Debug,
+{
     /// This operation is called by a broker after uploading the
     /// shared log segment object to the object storage.
-    ///
-    /// The Batch Coordinator:
-    /// 1. Performs the necessary checks for idempotent produce.
-    /// 2. Accordingly increases the high watermark of the affected logs.
-    /// 3. Assigns offsets to the batches.
-    /// 4. Saves the batch and object metadata.
-    /// 5. Returns the result to the broker.
-    ///
-    /// # Errors
-    /// Returns an error if an unexpected error occurs.
     async fn commit_file(
         &self,
         object_key: [u8; 16],
         uploader_broker_id: u32,
         file_size: u64,
         batches: Vec<CommitBatchRequest>,
     ) -> Vec<CommitBatchResponse>;
+}
 
+/// Trait for Implementing fn find_batches.
+#[async_trait::async_trait]
+pub trait FindBatches
+where
+    Self: Send + Sync + std::fmt::Debug,
+{
     /// This operation is called by a broker when it needs to serve a Fetch request.
-    /// The Batch Coordinator collects the batch coordinates to satisfy
-    /// this request and sends the response back to the broker.
-    ///
-    /// # Errors
-    /// Returns an error if an unexpected error occurs.
     async fn find_batches(
         &self,
         find_batch_requests: Vec<FindBatchRequest>,
         fetch_max_bytes: u32,
     ) -> Vec<FindBatchResponse>;
+}
 
-    /// This operation allows the broker to get the information about log offsets:
-    /// earliest, latest, etc. The operation is a read-only operation.
-    ///
-    /// # Errors
-    /// Returns an error if an unexpected error occurs.
-    async fn list_offsets(&self, requests: Vec<ListOffsetsRequest>) -> Vec<ListOffsetsResponse>;
-
+/// The BatchCoordinator trait.
+///
+/// This structure is responsible for handling the indexing of offsets within topic's partitions.
+/// It is designed to be reimplementable for custom usecases depending on what the specific need is.
+#[async_trait::async_trait]
+pub trait DeleteFiles
+where
+    Self: Send + Sync + std::fmt::Debug,
+{
     /// This operation is called when a partition needs to be truncated by the user.
     /// The Batch Coordinator:
     /// 1. Modifies the log start offset for the affected partitions (logs).

src/batch_coordinator/mod.rs

@@ -15,11 +15,12 @@ use uuid::Uuid;
 use crate::{batch_coordinator::BatchInfo, error::RisklessResult, messages::CommitBatchRequest};
 
 use crate::batch_coordinator::{
-    BatchCoordinator, BatchMetadata, CommitBatchResponse, CreateTopicAndPartitionsRequest,
-    DeleteFilesRequest, DeleteRecordsRequest, DeleteRecordsResponse, FileToDelete,
-    FindBatchRequest, FindBatchResponse, ListOffsetsRequest, ListOffsetsResponse,
+    BatchMetadata, CommitBatchResponse, DeleteFilesRequest, DeleteRecordsRequest,
+    DeleteRecordsResponse, FileToDelete, FindBatchRequest, FindBatchResponse,
 };
 
+use super::{CommitFile, DeleteFiles, FindBatches};
+
 /// The SimpleBatchCoordinator is a default implementation that is
 ///
 /// - Single Node
@@ -95,15 +96,7 @@ impl SimpleBatchCoordinator {
 }
 
 #[async_trait::async_trait]
-impl BatchCoordinator for SimpleBatchCoordinator {
-    async fn create_topic_and_partitions(
-        &self,
-        _requests: HashSet<CreateTopicAndPartitionsRequest>,
-    ) {
-        // This is not implemented for SimpleBatchCoordinator as the topics + partitions get
-        // created as they have data produced to them.
-    }
-
+impl CommitFile for SimpleBatchCoordinator {
     async fn commit_file(
         &self,
         object_key: [u8; 16],
@@ -163,7 +156,10 @@ impl BatchCoordinator for SimpleBatchCoordinator {
 
         commit_batch_responses
     }
+}
 
+#[async_trait::async_trait]
+impl FindBatches for SimpleBatchCoordinator {
     async fn find_batches(
         &self,
         find_batch_requests: Vec<FindBatchRequest>,
@@ -295,11 +291,10 @@ impl BatchCoordinator for SimpleBatchCoordinator {
 
         results
     }
+}
 
-    async fn list_offsets(&self, _requests: Vec<ListOffsetsRequest>) -> Vec<ListOffsetsResponse> {
-        todo!()
-    }
-
+#[async_trait::async_trait]
+impl DeleteFiles for SimpleBatchCoordinator {
     /// Simply returns errors as this implementation does not support this operation.
     async fn delete_records(
         &self,

src/batch_coordinator/simple/mod.rs

@@ -55,7 +55,9 @@ mod shared_log_segment;
 
 use std::{collections::HashSet, sync::Arc};
 
-use batch_coordinator::{BatchCoordinator, DeleteFilesRequest, FindBatchRequest, TopicIdPartition};
+use batch_coordinator::{
+    CommitFile, DeleteFiles, DeleteFilesRequest, FindBatchRequest, FindBatches, TopicIdPartition,
+};
 use bytes::Bytes;
 use messages::{
     CommitBatchRequest, ConsumeBatch, ConsumeRequest, ConsumeResponse, ProduceRequestCollection,
@@ -72,7 +74,7 @@ use shared_log_segment::SharedLogSegment;
 pub async fn flush(
     reqs: ProduceRequestCollection,
     object_storage: Arc<dyn ObjectStore>,
-    batch_coordinator: Arc<dyn BatchCoordinator>,
+    batch_coordinator: Arc<dyn CommitFile>,
 ) -> RisklessResult<Vec<ProduceResponse>> {
     tracing::info!("Produce Requests: {:#?}", reqs);
 
@@ -120,7 +122,7 @@ pub async fn flush(
 pub async fn consume(
     request: ConsumeRequest,
     object_storage: Arc<dyn ObjectStore>,
-    batch_coordinator: Arc<dyn BatchCoordinator>,
+    batch_coordinator: Arc<dyn FindBatches>,
 ) -> RisklessResult<tokio::sync::mpsc::Receiver<ConsumeResponse>> {
     let batch_responses = batch_coordinator
         .find_batches(
@@ -207,7 +209,7 @@ pub async fn consume(
 #[tracing::instrument(skip_all, name = "delete_records")]
 pub async fn delete_record(
     request: crate::messages::DeleteRecordsRequest,
-    batch_coordinator: Arc<dyn BatchCoordinator>,
+    batch_coordinator: Arc<dyn DeleteFiles>,
 ) -> RisklessResult<crate::messages::DeleteRecordsResponse> {
     let result = batch_coordinator
         .delete_records(vec![request.try_into().map_err(|e| {
@@ -232,7 +234,7 @@ pub async fn delete_record(
 /// The interval at which this happens is delegated to the implementor.
 #[tracing::instrument(skip_all, name = "heartbeat_permanent_delete")]
 pub async fn scan_and_permanently_delete_records(
-    batch_coordinator: Arc<dyn BatchCoordinator>,
+    batch_coordinator: Arc<dyn DeleteFiles>,
     object_store: Arc<dyn ObjectStore>,
 ) -> RisklessResult<()> {
     let files_to_delete = batch_coordinator.get_files_to_delete().await;