test(remote_wal): add unit tests for kafka remote wal

src/common/config/src/wal/kafka.rs

@@ -40,7 +40,6 @@ pub struct KafkaConfig {
     pub broker_endpoints: Vec<String>,
     /// The compression algorithm used to compress log entries.
     #[serde(skip)]
-    #[serde(default)]
     pub compression: RsKafkaCompression,
     /// The max size of a single producer batch.
     pub max_batch_size: ReadableSize,

src/common/meta/Cargo.toml

@@ -14,6 +14,7 @@ async-stream.workspace = true
 async-trait.workspace = true
 base64.workspace = true
 bytes.workspace = true
+chrono.workspace = true
 common-catalog.workspace = true
 common-config.workspace = true
 common-error.workspace = true
@@ -51,3 +52,4 @@ chrono.workspace = true
 common-procedure = { workspace = true, features = ["testing"] }
 datatypes.workspace = true
 hyper = { version = "0.14", features = ["full"] }
+uuid.workspace = true

src/common/meta/src/wal/kafka.rs

@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#[cfg(test)]
+pub mod test_util;
 pub mod topic;
 pub mod topic_manager;
 pub mod topic_selector;

src/common/meta/src/wal/kafka/test_util.rs

@@ -0,0 +1,32 @@
+// Copyright 2023 Greptime 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 rskafka::client::ClientBuilder;
+
+use crate::wal::kafka::Topic;
+
+/// Gets broker endpoints from environment variables with the given key.
+/// Returns the default ["localhost:9092"] if no environment variables set for broker endpoints.
+#[macro_export]
+macro_rules! get_broker_endpoints {
+    () => {{
+        let broker_endpoints = std::env::var("GT_KAFKA_ENDPOINTS")
+            .unwrap_or("localhost:9092".to_string())
+            .split(',')
+            .map(ToString::to_string)
+            .collect::<Vec<_>>();
+        assert!(!broker_endpoints.is_empty());
+        broker_endpoints
+    }};
+}

src/common/meta/src/wal/kafka/topic.rs

@@ -15,4 +15,5 @@
 /// Kafka wal topic.
 /// Publishers publish log entries to the topic while subscribers pull log entries from the topic.
 /// A topic is simply a string right now. But it may be more complex in the future.
+// TODO(niebayes): remove the Topic alias.
 pub type Topic = String;

src/common/meta/src/wal/kafka/topic_manager.rs

@@ -47,9 +47,8 @@ const DEFAULT_PARTITION: i32 = 0;
 /// Manages topic initialization and selection.
 pub struct TopicManager {
     config: KafkaConfig,
-    // TODO(niebayes): maybe add a guard to ensure all topics in the topic pool are created.
-    topic_pool: Vec<Topic>,
-    topic_selector: TopicSelectorRef,
+    pub(crate) topic_pool: Vec<Topic>,
+    pub(crate) topic_selector: TopicSelectorRef,
     kv_backend: KvBackendRef,
 }
 
@@ -168,7 +167,7 @@ impl TopicManager {
                 vec![Record {
                     key: None,
                     value: None,
-                    timestamp: rskafka::chrono::Utc::now(),
+                    timestamp: chrono::Utc::now(),
                     headers: Default::default(),
                 }],
                 Compression::NoCompression,
@@ -245,8 +244,8 @@ mod tests {
     use common_telemetry::info;
 
     use super::*;
+    use crate::get_broker_endpoints;
     use crate::kv_backend::memory::MemoryKvBackend;
-    use crate::kv_backend::{self};
 
     // Tests that topics can be successfully persisted into the kv backend and can be successfully restored from the kv backend.
     #[tokio::test]
@@ -273,26 +272,56 @@ mod tests {
         assert_eq!(topics, restored_topics);
     }
 
+    /// Tests that the topic manager could allocate topics correctly.
     #[tokio::test]
-    async fn test_topic_manager() {
-        let endpoints = env::var("GT_KAFKA_ENDPOINTS").unwrap_or_default();
-        common_telemetry::init_default_ut_logging();
+    async fn test_alloc_topics() {
+        let broker_endpoints = get_broker_endpoints!();
+        // Constructs topics that should be created.
+        let topics = (0..256)
+            .map(|i| format!("test_alloc_topics_{}_{}", i, uuid::Uuid::new_v4()))
+            .collect::<Vec<_>>();
 
-        if endpoints.is_empty() {
-            info!("The endpoints is empty, skipping the test.");
-            return;
-        }
-        // TODO: supports topic prefix
-        let kv_backend = Arc::new(MemoryKvBackend::new());
+        // Creates a topic manager.
         let config = KafkaConfig {
-            replication_factor: 1,
-            broker_endpoints: endpoints
-                .split(',')
-                .map(|s| s.to_string())
-                .collect::<Vec<_>>(),
+            replication_factor: broker_endpoints.len() as i16,
+            broker_endpoints,
             ..Default::default()
         };
-        let manager = TopicManager::new(config, kv_backend);
+        let kv_backend = Arc::new(MemoryKvBackend::new()) as KvBackendRef;
+        let mut manager = TopicManager::new(config.clone(), kv_backend);
+        // Replaces the default topic pool with the constructed topics.
+        manager.topic_pool = topics.clone();
+        // Replaces the default selector with a round-robin selector without shuffled.
+        manager.topic_selector = Arc::new(RoundRobinTopicSelector::default());
         manager.start().await.unwrap();
+
+        // Selects exactly the number of `num_topics` topics one by one.
+        let got = (0..topics.len())
+            .map(|_| manager.select().unwrap())
+            .cloned()
+            .collect::<Vec<_>>();
+        assert_eq!(got, topics);
+
+        // Selects exactly the number of `num_topics` topics in a batching manner.
+        let got = manager
+            .select_batch(topics.len())
+            .unwrap()
+            .into_iter()
+            .map(ToString::to_string)
+            .collect::<Vec<_>>();
+        assert_eq!(got, topics);
+
+        // Selects more than the number of `num_topics` topics.
+        let got = manager
+            .select_batch(2 * topics.len())
+            .unwrap()
+            .into_iter()
+            .map(ToString::to_string)
+            .collect::<Vec<_>>();
+        let expected = vec![topics.clone(); 2]
+            .into_iter()
+            .flatten()
+            .collect::<Vec<_>>();
+        assert_eq!(got, expected);
     }
 }

src/common/meta/src/wal/kafka/topic_selector.rs

@@ -60,6 +60,14 @@ impl TopicSelector for RoundRobinTopicSelector {
 mod tests {
     use super::*;
 
+    /// Tests that a selector behaves as expected when the given topic pool is empty.
+    #[test]
+    fn test_empty_topic_pool() {
+        let topic_pool = vec![];
+        let selector = RoundRobinTopicSelector::default();
+        assert!(selector.select(&topic_pool).is_err());
+    }
+
     #[test]
     fn test_round_robin_topic_selector() {
         let topic_pool: Vec<_> = [0, 1, 2].into_iter().map(|v| v.to_string()).collect();

src/common/meta/src/wal/options_allocator.rs

@@ -106,10 +106,12 @@ pub fn allocate_region_wal_options(
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::get_broker_endpoints;
     use crate::kv_backend::memory::MemoryKvBackend;
+    use crate::wal::kafka::topic_selector::RoundRobinTopicSelector;
+    use crate::wal::kafka::KafkaConfig;
 
     // Tests the wal options allocator could successfully allocate raft-engine wal options.
-    // Note: tests for allocator with kafka are integration tests.
     #[tokio::test]
     async fn test_allocator_with_raft_engine() {
         let kv_backend = Arc::new(MemoryKvBackend::new()) as KvBackendRef;
@@ -128,4 +130,46 @@ mod tests {
             .collect();
         assert_eq!(got, expected);
     }
+
+    // Tests that the wal options allocator could successfully allocate Kafka wal options.
+    #[tokio::test]
+    async fn test_allocator_with_kafka() {
+        let broker_endpoints = get_broker_endpoints!();
+        let topics = (0..256)
+            .map(|i| format!("test_allocator_with_kafka_{}_{}", i, uuid::Uuid::new_v4()))
+            .collect::<Vec<_>>();
+
+        // Creates a topic manager.
+        let config = KafkaConfig {
+            replication_factor: broker_endpoints.len() as i16,
+            broker_endpoints,
+            ..Default::default()
+        };
+        let kv_backend = Arc::new(MemoryKvBackend::new()) as KvBackendRef;
+        let mut topic_manager = KafkaTopicManager::new(config.clone(), kv_backend);
+        // Replaces the default topic pool with the constructed topics.
+        topic_manager.topic_pool = topics.clone();
+        // Replaces the default selector with a round-robin selector without shuffled.
+        topic_manager.topic_selector = Arc::new(RoundRobinTopicSelector::default());
+
+        // Creates an options allocator.
+        let wal_config = WalConfig::Kafka(config.clone());
+        let allocator = WalOptionsAllocator::Kafka(topic_manager);
+        allocator.start().await.unwrap();
+
+        let num_regions = 32;
+        let regions = (0..num_regions).collect::<Vec<_>>();
+        let got = allocate_region_wal_options(regions.clone(), &allocator).unwrap();
+
+        // Check the allocated wal options contain the expected topics.
+        let expected = (0..num_regions)
+            .map(|i| {
+                let options = WalOptions::Kafka(KafkaWalOptions {
+                    topic: topics[i as usize].clone(),
+                });
+                (i, serde_json::to_string(&options).unwrap())
+            })
+            .collect::<HashMap<_, _>>();
+        assert_eq!(got, expected);
+    }
 }

src/log-store/Cargo.toml

@@ -22,10 +22,8 @@ common-macro.workspace = true
 common-meta.workspace = true
 common-runtime.workspace = true
 common-telemetry.workspace = true
-dashmap.workspace = true
 futures-util.workspace = true
 futures.workspace = true
-itertools.workspace = true
 protobuf = { version = "2", features = ["bytes"] }
 raft-engine.workspace = true
 rskafka.workspace = true
@@ -39,4 +37,7 @@ tokio.workspace = true
 [dev-dependencies]
 common-meta = { workspace = true, features = ["testing"] }
 common-test-util.workspace = true
+itertools.workspace = true
 rand.workspace = true
+rand_distr = "0.4"
+uuid.workspace = true

src/log-store/src/kafka.rs

@@ -40,19 +40,19 @@ impl Namespace for NamespaceImpl {
 
 impl Display for NamespaceImpl {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "{}/{}", self.topic, self.region_id)
+        write!(f, "[topic: {}, region: {}]", self.topic, self.region_id)
     }
 }
 
 /// Kafka Entry implementation.
 #[derive(Debug, PartialEq, Clone)]
 pub struct EntryImpl {
     /// Entry payload.
-    data: Vec<u8>,
+    pub data: Vec<u8>,
     /// The logical entry id.
-    id: EntryId,
+    pub id: EntryId,
     /// The namespace used to identify and isolate log entries from different regions.
-    ns: NamespaceImpl,
+    pub ns: NamespaceImpl,
 }
 
 impl Entry for EntryImpl {
@@ -76,7 +76,7 @@ impl Display for EntryImpl {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         write!(
             f,
-            "Entry (ns: {}, id: {}, data_len: {})",
+            "Entry [ns: {}, id: {}, data_len: {}]",
             self.ns,
             self.id,
             self.data.len()