feat(kafka): kafka api command.

Author: echooymxq

Cargo.lock

@@ -144,13 +144,14 @@ impl BrokerServer {
             broker_cache.clone(),
             rocksdb_engine_handler.clone(),
             connection_manager.clone(),
-            message_storage_adapter,
+            message_storage_adapter.clone(),
         );
 
         let kafka_params = BrokerServer::build_kafka_server_params(
             client_pool.clone(),
             broker_cache.clone(),
             connection_manager.clone(),
+            message_storage_adapter.clone(),
         );
 
         let journal_params = BrokerServer::build_journal_server(client_pool.clone());
@@ -404,6 +405,7 @@ impl BrokerServer {
         client_pool: Arc<ClientPool>,
         broker_cache: Arc<BrokerCacheManager>,
         connection_manager: Arc<NetworkConnectionManager>,
+        message_storage_adapter: ArcStorageAdapter,
     ) -> KafkaBrokerServerParams {
         let conf = broker_config();
         let proc_config = ProcessorConfig {
@@ -414,6 +416,7 @@ impl BrokerServer {
         };
         KafkaBrokerServerParams {
             connection_manager,
+            message_storage_adapter,
             client_pool,
             proc_config,
             broker_cache,

src/broker-server/src/lib.rs

@@ -21,6 +21,7 @@ license.workspace = true
 
 [dependencies]
 axum.workspace = true
+common-base.workspace = true
 common-config.workspace = true
 tokio.workspace = true
 tokio-util.workspace = true
@@ -31,4 +32,7 @@ grpc-clients.workspace = true
 tracing.workspace = true
 protocol.workspace = true
 kafka-protocol.workspace = true
-thiserror.workspace = true
+thiserror.workspace = true
+storage-adapter.workspace = true
+bytes = "1.10.1"
+serde = { version = "1.0.228", features = ["derive"] }

src/kafka-broker/Cargo.toml

@@ -24,10 +24,12 @@ use network_server::tcp::server::TcpServer;
 use std::sync::Arc;
 use tokio::sync::broadcast::{self};
 use tracing::{error, info};
+use storage_adapter::storage::ArcStorageAdapter;
 
 #[derive(Clone)]
 pub struct KafkaBrokerServerParams {
     pub connection_manager: Arc<ConnectionManager>,
+    pub message_storage_adapter: ArcStorageAdapter,
     pub client_pool: Arc<ClientPool>,
     pub proc_config: ProcessorConfig,
     pub broker_cache: Arc<BrokerCacheManager>,
@@ -43,7 +45,7 @@ pub struct KafkaBrokerServer {
 impl KafkaBrokerServer {
     pub fn new(params: KafkaBrokerServerParams, main_stop: broadcast::Sender<bool>) -> Self {
         let (inner_stop, _) = broadcast::channel(2);
-        let command = create_command();
+        let command = create_command(params.message_storage_adapter);
         let context = ServerContext {
             connection_manager: params.connection_manager.clone(),
             client_pool: params.client_pool.clone(),

src/kafka-broker/src/broker.rs

@@ -17,4 +17,6 @@ use thiserror::Error;
 #[derive(Error, Debug)]
 pub enum KafkaBrokerError {}
 
-pub type ResultKafkaBrokerError = Result<(), KafkaBrokerError>;
+pub type ResultKafkaBrokerError = Result<()>;
+
+pub type Result<T> = std::result::Result<T, KafkaBrokerError>;

src/kafka-broker/src/common/error.rs

@@ -12,4 +12,4 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-pub(crate) mod error;
+pub mod error;

src/kafka-broker/src/common/mod.rs

@@ -14,19 +14,41 @@
 
 use axum::async_trait;
 use kafka_protocol::messages::api_versions_response::ApiVersion;
-use kafka_protocol::messages::{ApiKey, ApiVersionsResponse};
+use kafka_protocol::messages::{ApiKey, ApiVersionsResponse, ProduceRequest, ProduceResponse};
+use kafka_protocol::messages::produce_response::{PartitionProduceResponse, TopicProduceResponse};
+use kafka_protocol::records::{RecordBatchDecoder, RecordBatchEncoder};
 use metadata_struct::connection::NetworkConnection;
 use network_server::{command::Command, common::packet::ResponsePackage};
 use protocol::kafka::packet::KafkaPacket;
 use protocol::robust::RobustMQPacket;
 use std::net::SocketAddr;
 use std::sync::Arc;
+use tracing::log::{info, warn};
+use storage_adapter::storage::ArcStorageAdapter;
+use crate::{
+    common::error::Result,
+    storage::message::{Message, TopicPartition},
+};
+use crate::manager::offset::OffsetManager;
+use crate::storage::log_reader::Reader;
+use crate::storage::log_writer::Writer;
 
-#[derive(Clone, Default)]
-pub struct KafkaCommand {}
+pub struct KafkaCommand {
+    pub writer: Writer,
+    pub reader: Reader,
+    pub offset_manager: OffsetManager,
+}
 
 impl KafkaCommand {
-    fn handle_api_versions(tcp_connection: &NetworkConnection) -> ResponsePackage {
+    pub fn new(storage_adapter: ArcStorageAdapter) -> Self {
+        Self {
+            writer: Writer::new(storage_adapter.clone()),
+            reader: Reader::new(storage_adapter.clone()),
+            offset_manager: OffsetManager::new(),
+        }
+    }
+
+    fn handle_api_versions() -> Result<KafkaPacket> {
         let api_versions: Vec<ApiVersion> = ApiKey::iter()
             .map(|k| {
                 let range = ApiKey::valid_versions(&k);
@@ -40,10 +62,79 @@ impl KafkaCommand {
             .with_error_code(0)
             .with_api_keys(api_versions)
             .with_throttle_time_ms(0);
-        ResponsePackage::build(
-            tcp_connection.connection_id,
-            RobustMQPacket::KAFKA(KafkaPacket::ApiVersionResponse(res)),
-        )
+        Ok(KafkaPacket::ApiVersionResponse(res))
+    }
+
+    async fn handle_produce(
+        &self,
+        request: ProduceRequest,
+    ) -> Result<KafkaPacket> {
+        let mut topic_responses = vec![];
+        let mut message_batch = vec![];
+
+        for topic_data in &request.topic_data {
+            let topic_name = topic_data.name.as_str();
+            let mut partition_responses = vec![];
+            for partition_data in &topic_data.partition_data {
+                let partition_index = partition_data.index;
+                let mut base_offset = None;
+                let mut partition_error_code = 0i16;
+                if let Some(records_bytes) = &partition_data.records {
+                    let mut records_buf = records_bytes.clone();
+                    let result = RecordBatchDecoder::decode(&mut records_buf);
+                    match result {
+                        Ok(record_set) => {
+                            let mut log_start_offset = 0i64;
+                            if record_set.records.len() > 0 && base_offset.is_none() {
+                                log_start_offset = self.offset_manager.next_offset(topic_name, partition_index)
+                                    .await
+                                    .unwrap_or(0);
+                                base_offset = Some(log_start_offset);
+                            }
+                            for record in &record_set.records {
+                                log_start_offset += 1;
+                                message_batch.push(Message {
+                                    topic_partition: TopicPartition {
+                                        topic: topic_name.to_string(),
+                                        partition: partition_index,
+                                    },
+                                    offset: log_start_offset,
+                                    record: Message::encode(record.clone()),
+                                });
+                            }
+                        }
+                        Err(e) => {
+                            warn!("Failed to decode record batch: {:?}", e);
+                            partition_error_code = 1;
+                        }
+                    }
+                }
+
+                let partition_resp = PartitionProduceResponse::default()
+                    .with_index(partition_index)
+                    .with_error_code(partition_error_code)
+                    .with_base_offset(base_offset.unwrap_or_default())
+                    .with_log_append_time_ms(-1) // message.timestamp.type
+                    .with_log_start_offset(0);
+
+                partition_responses.push(partition_resp);
+            }
+
+            let topic_resp = TopicProduceResponse::default()
+                .with_name(topic_data.name.clone())
+                .with_partition_responses(partition_responses);
+
+            topic_responses.push(topic_resp);
+        }
+
+        for msg in message_batch.into_iter() {
+            self.writer.write(&msg).await?;
+        }
+
+        let response = ProduceResponse::default()
+            .with_responses(topic_responses)
+            .with_throttle_time_ms(0);
+        Ok(KafkaPacket::ProduceResponse(response))
     }
 }
 
@@ -55,17 +146,46 @@ impl Command for KafkaCommand {
         _addr: &SocketAddr,
         robust_packet: &RobustMQPacket,
     ) -> Option<ResponsePackage> {
-        let packet = robust_packet.get_kafka_packet().unwrap();
-        match packet.clone() {
-            KafkaPacket::ApiVersionReq(_) => Some(Self::handle_api_versions(tcp_connection)),
-            KafkaPacket::ProduceReq(_) => None,
-            KafkaPacket::FetchReq(_) => None,
-            _ => None,
-        }
+        let packet = match robust_packet.get_kafka_packet() {
+            Some(p) => p,
+            None => {
+                warn!("No Kafka packet found in RobustMQPacket");
+                return None;
+            }
+        };
+        let kafka_response = match packet {
+            KafkaPacket::ApiVersionReq(_) => match Self::handle_api_versions() {
+                Ok(resp) => resp,
+                Err(e) => {
+                    warn!("Failed to build ApiVersionsResponse: {:?}", e);
+                    return None;
+                }
+            },
+            KafkaPacket::ProduceReq(req) => match self.handle_produce(req).await {
+                Ok(resp) => resp,
+                Err(e) => {
+                    warn!("Produce handler failed: {:?}", e);
+                    return None;
+                }
+            },
+            KafkaPacket::FetchReq(_) => {
+                warn!("Received Fetch request but Fetch handling is not implemented");
+                return None;
+            }
+            other => {
+                warn!("Unsupported or unhandled Kafka packet: {:?}", other);
+                return None;
+            }
+        };
+
+        Some(ResponsePackage::build(
+            tcp_connection.connection_id,
+            RobustMQPacket::KAFKA(kafka_response),
+        ))
     }
 }
 
-pub fn create_command() -> Arc<Box<dyn Command + Send + Sync>> {
-    let storage: Box<dyn Command + Send + Sync> = Box::new(KafkaCommand::default());
+pub fn create_command(message_storage_adapter: ArcStorageAdapter) -> Arc<Box<dyn Command + Send + Sync>> {
+    let storage: Box<dyn Command + Send + Sync> = Box::new(KafkaCommand::new(message_storage_adapter));
     Arc::new(storage)
 }

src/kafka-broker/src/handler/command.rs

@@ -16,3 +16,5 @@ pub mod broker;
 pub mod common;
 pub mod handler;
 pub mod server;
+pub mod storage;
+pub mod manager;

src/kafka-broker/src/lib.rs

@@ -0,0 +1,14 @@
+// Copyright 2023 RobustMQ Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+pub mod offset;

src/kafka-broker/src/manager/mod.rs

@@ -0,0 +1,69 @@
+// Copyright 2023 RobustMQ Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
+use std::sync::{Arc, Mutex};
+use tracing::debug;
+use crate::{
+    common::error::Result,
+};
+#[derive(Debug, Clone, Default, Serialize, Deserialize)]
+pub struct LogOffsetMetadata {
+    pub message_offset: i64,
+    pub segment_start_offset: i64,
+}
+
+pub struct OffsetManager {
+    store_cache: Arc<Mutex<HashMap<(String, i32), LogOffsetMetadata>>>,
+}
+
+impl OffsetManager {
+
+    pub fn new() -> Self {
+        Self {
+            store_cache: Arc::new(Mutex::new(HashMap::new())),
+        }
+    }
+    pub async fn next_offset(&self, topic: &str, partition: i32) -> Result<i64> {
+        let mut cache = self.store_cache.lock().await;
+        let key = (topic.to_string(), partition);
+
+        let offsets = cache.entry(key.clone()).or_insert_with(|| {
+            LogOffsetMetadata::default()
+        });
+
+        let next_offset = offsets.next_offset;
+        offsets.next_offset += 1;
+
+        if next_offset % 100 == 0 {
+            let topic_clone = topic.to_string();
+            tokio::spawn(async move {
+                if let Err(e) = Self::persist_offsets(
+                    self,
+                    &topic_clone,
+                    partition,
+                    0,
+                ).await {
+                    debug!("Failed to persist offsets for topic {} partition {}: {:?}", topic_clone, partition, e);
+                }
+            });
+        }
+        Ok(next_offset)
+    }
+
+    pub async fn persist_offsets(
+        &self, topic: &str, partition: i32, offset: i64) -> Result<()> {
+        Ok(())
+    }
+}