From 81804ebdc1168e55b5be6038de2d5b1093d110ae Mon Sep 17 00:00:00 2001 From: niebayes Date: Mon, 25 Dec 2023 21:05:52 +0800 Subject: [PATCH 01/30] test: add unit tests --- Cargo.lock | 2 + src/log-store/src/kafka.rs | 10 +- tests-integration/Cargo.toml | 2 + tests-integration/tests/wal.rs | 248 +++++++++++++++++++++++++++++++++ 4 files changed, 257 insertions(+), 5 deletions(-) create mode 100644 tests-integration/tests/wal.rs diff --git a/Cargo.lock b/Cargo.lock index ac08a49f6baa..885ca3e2e0e3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9481,6 +9481,7 @@ dependencies = [ "frontend", "futures", "itertools 0.10.5", + "log-store", "meta-client", "meta-srv", "mysql_async", @@ -9494,6 +9495,7 @@ dependencies = [ "prost 0.12.3", "query", "rand", + "rskafka", "rstest", "rstest_reuse", "script", diff --git a/src/log-store/src/kafka.rs b/src/log-store/src/kafka.rs index fefa7823c5c2..a08afc508eac 100644 --- a/src/log-store/src/kafka.rs +++ b/src/log-store/src/kafka.rs @@ -29,8 +29,8 @@ use crate::error::Error; /// Kafka Namespace implementation. #[derive(Debug, PartialEq, Eq, Hash, Clone, Serialize, Deserialize)] pub struct NamespaceImpl { - region_id: u64, - topic: Topic, + pub region_id: u64, + pub topic: Topic, } impl Namespace for NamespaceImpl { @@ -49,11 +49,11 @@ impl Display for NamespaceImpl { #[derive(Debug, PartialEq, Clone)] pub struct EntryImpl { /// Entry payload. - data: Vec, + pub data: Vec, /// 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 { diff --git a/tests-integration/Cargo.toml b/tests-integration/Cargo.toml index 1ba6f8e05d36..8d311d3f7ea2 100644 --- a/tests-integration/Cargo.toml +++ b/tests-integration/Cargo.toml @@ -34,6 +34,7 @@ datatypes.workspace = true dotenv = "0.15" frontend = { workspace = true, features = ["testing"] } futures.workspace = true +log-store.workspace = true meta-client.workspace = true meta-srv = { workspace = true, features = ["mock"] } mysql_async = { version = "0.33", default-features = false, features = [ @@ -44,6 +45,7 @@ once_cell.workspace = true operator.workspace = true query.workspace = true rand.workspace = true +rskafka.workspace = true rstest = "0.17" rstest_reuse = "0.5" secrecy = "0.8" diff --git a/tests-integration/tests/wal.rs b/tests-integration/tests/wal.rs new file mode 100644 index 000000000000..30695217087d --- /dev/null +++ b/tests-integration/tests/wal.rs @@ -0,0 +1,248 @@ +// 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 std::collections::HashMap; +use std::sync::Arc; + +use common_config::wal::KafkaConfig as DatanodeKafkaConfig; +use common_config::{KafkaWalOptions, WalOptions}; +use common_meta::kv_backend::memory::MemoryKvBackend; +use common_meta::kv_backend::KvBackendRef; +use common_meta::wal::kafka::{ + KafkaConfig as MetaSrvKafkaConfig, TopicManager as KafkaTopicManager, +}; +use common_meta::wal::{allocate_region_wal_options, WalConfig, WalOptionsAllocator}; +use futures::StreamExt; +use log_store::kafka::log_store::KafkaLogStore; +use log_store::kafka::{EntryImpl, NamespaceImpl}; +use rskafka::client::controller::ControllerClient; +use rskafka::client::ClientBuilder; +use store_api::logstore::entry::Id as EntryId; +use store_api::logstore::LogStore; + +// Notice: the following tests are literally unit tests. They are placed at here since +// it seems too heavy to start a Kafka cluster for each unit test. + +// The key of an env variable that stores a series of Kafka broker endpoints. +const BROKER_ENDPOINTS_KEY: &str = "GT_KAFKA_ENDPOINTS"; + +// Tests that the TopicManager allocates topics in a round-robin mannar. +#[tokio::test] +async fn test_kafka_alloc_topics() { + let broker_endpoints = std::env::var(BROKER_ENDPOINTS_KEY) + .unwrap() + .split(',') + .map(ToString::to_string) + .collect::>(); + let config = MetaSrvKafkaConfig { + topic_name_prefix: "__test_kafka_alloc_topics".to_string(), + replication_factor: broker_endpoints.len() as i16, + broker_endpoints, + ..Default::default() + }; + let kv_backend = Arc::new(MemoryKvBackend::new()) as KvBackendRef; + let manager = KafkaTopicManager::new(config.clone(), kv_backend); + manager.start().await.unwrap(); + + // Topics should be created. + let topics = (0..config.num_topics) + .map(|topic_id| format!("{}_{topic_id}", config.topic_name_prefix)) + .collect::>(); + + // Selects exactly the number of `num_topics` topics one by one. + for expected in topics.iter() { + let got = manager.select().unwrap(); + assert_eq!(got, expected); + } + + // Selects exactly the number of `num_topics` topics in a batching manner. + let got = manager + .select_batch(config.num_topics) + .unwrap() + .into_iter() + .map(ToString::to_string) + .collect::>(); + assert_eq!(got, topics); + + // Selects none. + let got = manager.select_batch(config.num_topics).unwrap(); + assert!(got.is_empty()); + + // Selects more than the number of `num_topics` topics. + let got = manager + .select_batch(2 * config.num_topics) + .unwrap() + .into_iter() + .map(ToString::to_string) + .collect::>(); + let expected = vec![topics.clone(); 2] + .into_iter() + .flatten() + .collect::>(); + assert_eq!(got, expected); +} + +// Tests that the wal options allocator could successfully allocate Kafka wal options. +#[tokio::test] +async fn test_kafka_options_allocator() { + let broker_endpoints = std::env::var(BROKER_ENDPOINTS_KEY) + .unwrap() + .split(',') + .map(ToString::to_string) + .collect::>(); + let config = MetaSrvKafkaConfig { + topic_name_prefix: "__test_kafka_options_allocator".to_string(), + replication_factor: broker_endpoints.len() as i16, + broker_endpoints, + ..Default::default() + }; + let wal_config = WalConfig::Kafka(config.clone()); + let kv_backend = Arc::new(MemoryKvBackend::new()) as KvBackendRef; + let allocator = WalOptionsAllocator::new(wal_config, kv_backend); + allocator.start().await.unwrap(); + + let num_regions = 32; + let regions = (0..num_regions).collect::>(); + let got = allocate_region_wal_options(regions.clone(), &allocator).unwrap(); + + // Topics should be allocated. + let topics = (0..num_regions) + .map(|topic_id| format!("{}_{topic_id}", config.topic_name_prefix)) + .collect::>(); + // 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::>(); + assert_eq!(got, expected); +} + +fn new_test_entry>(data: D, entry_id: EntryId, ns: NamespaceImpl) -> EntryImpl { + EntryImpl { + data: data.as_ref().to_vec(), + id: entry_id, + ns, + } +} + +async fn create_topic(topic: &str, replication_factor: i16, client: &ControllerClient) { + client + .create_topic(topic, 1, replication_factor, 5_000) + .await + .unwrap(); +} + +async fn check_entries( + ns: &NamespaceImpl, + start_offset: EntryId, + expected: Vec, + logstore: &KafkaLogStore, +) { + let mut stream = logstore.read(ns, start_offset).await.unwrap(); + for entry in expected { + let got = stream.next().await.unwrap().unwrap(); + assert_eq!(entry, got[0]); + } +} + +// Tests that the Kafka log store is able to write and read log entries from Kafka. +#[tokio::test] +async fn test_kafka_log_store() { + let broker_endpoints = std::env::var(BROKER_ENDPOINTS_KEY) + .unwrap() + .split(',') + .map(ToString::to_string) + .collect::>(); + let config = DatanodeKafkaConfig { + broker_endpoints, + ..Default::default() + }; + let logstore = KafkaLogStore::try_new(&config).await.unwrap(); + + let client = ClientBuilder::new(config.broker_endpoints.clone()) + .build() + .await + .unwrap() + .controller_client() + .unwrap(); + + // Appends one entry. + let topic = "__test_kafka_log_store_topic_append"; + create_topic(topic, config.broker_endpoints.len() as i16, &client).await; + let ns = NamespaceImpl { + region_id: 0, + topic: topic.to_string(), + }; + let entry = new_test_entry(b"0", 0, ns.clone()); + let last_entry_id = logstore.append(entry.clone()).await.unwrap().last_entry_id; + check_entries(&ns, last_entry_id, vec![entry], &logstore).await; + + // Appends a batch of entries. + // Region 1, 2 are mapped to topic 1, + let topic = "__test_kafka_log_store_topic_append_batch_1"; + create_topic(topic, config.broker_endpoints.len() as i16, &client).await; + let ns_1 = NamespaceImpl { + region_id: 1, + topic: topic.to_string(), + }; + let ns_2 = NamespaceImpl { + region_id: 2, + topic: topic.to_string(), + }; + + // Region 3 is mapped to topic 2. + let topic = "__test_kafka_log_store_topic_append_batch_2"; + create_topic(topic, config.broker_endpoints.len() as i16, &client).await; + let ns_3 = NamespaceImpl { + region_id: 3, + topic: topic.to_string(), + }; + + // Constructs a batch of entries. + let entries_1 = vec![ + new_test_entry(b"1", 0, ns_1.clone()), + new_test_entry(b"1", 1, ns_1.clone()), + ]; + let entries_2 = vec![ + new_test_entry(b"2", 2, ns_2.clone()), + new_test_entry(b"2", 3, ns_2.clone()), + ]; + let entries_3 = vec![ + new_test_entry(b"3", 7, ns_3.clone()), + new_test_entry(b"3", 8, ns_3.clone()), + ]; + let entries = vec![entries_1.clone(), entries_2.clone(), entries_3.clone()] + .into_iter() + .flatten() + .collect::>(); + + let last_entry_ids = logstore + .append_batch(entries.clone()) + .await + .unwrap() + .last_entry_ids; + + // Reads entries for region 1. + check_entries(&ns_1, last_entry_ids[&1], entries_1, &logstore).await; + // Reads entries from region 2. + check_entries(&ns_2, last_entry_ids[&2], entries_2, &logstore).await; + // Reads entries from region 3. + check_entries(&ns_3, last_entry_ids[&3], entries_3, &logstore).await; +} + +// TODO(niebayes): add more integration tests. From a7f639ff50f51656a1697e10be5a54831603ee29 Mon Sep 17 00:00:00 2001 From: niebayes Date: Wed, 27 Dec 2023 19:04:24 +0800 Subject: [PATCH 02/30] feat: introduce kafka runtime backed by testcontainers --- Cargo.lock | 93 ++++++++++++++++++- tests-integration/Cargo.toml | 1 + tests-integration/src/lib.rs | 1 + tests-integration/src/wal_util.rs | 15 +++ tests-integration/src/wal_util/kafka.rs | 17 ++++ .../src/wal_util/kafka/config.rs | 93 +++++++++++++++++++ tests-integration/src/wal_util/kafka/image.rs | 77 +++++++++++++++ .../src/wal_util/kafka/runtime.rs | 38 ++++++++ 8 files changed, 331 insertions(+), 4 deletions(-) create mode 100644 tests-integration/src/wal_util.rs create mode 100644 tests-integration/src/wal_util/kafka.rs create mode 100644 tests-integration/src/wal_util/kafka/config.rs create mode 100644 tests-integration/src/wal_util/kafka/image.rs create mode 100644 tests-integration/src/wal_util/kafka/runtime.rs diff --git a/Cargo.lock b/Cargo.lock index 885ca3e2e0e3..d8b15f1bb700 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -972,6 +972,16 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bollard-stubs" +version = "1.42.0-rc.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed59b5c00048f48d7af971b71f800fdf23e858844a6f9e4d32ca72e9399e7864" +dependencies = [ + "serde", + "serde_with 1.14.0", +] + [[package]] name = "borsh" version = "1.3.0" @@ -1629,7 +1639,7 @@ dependencies = [ "rskafka", "serde", "serde_json", - "serde_with", + "serde_with 3.4.0", "toml 0.8.8", ] @@ -1841,7 +1851,7 @@ dependencies = [ "rskafka", "serde", "serde_json", - "serde_with", + "serde_with 3.4.0", "snafu", "store-api", "strum 0.25.0", @@ -2332,6 +2342,16 @@ dependencies = [ "memchr", ] +[[package]] +name = "darling" +version = "0.13.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a01d95850c592940db9b8194bc39f4bc0e89dee5c4265e4b1807c34a9aba453c" +dependencies = [ + "darling_core 0.13.4", + "darling_macro 0.13.4", +] + [[package]] name = "darling" version = "0.14.4" @@ -2352,6 +2372,20 @@ dependencies = [ "darling_macro 0.20.3", ] +[[package]] +name = "darling_core" +version = "0.13.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "859d65a907b6852c9361e3185c862aae7fafd2887876799fa55f5f99dc40d610" +dependencies = [ + "fnv", + "ident_case", + "proc-macro2", + "quote", + "strsim 0.10.0", + "syn 1.0.109", +] + [[package]] name = "darling_core" version = "0.14.4" @@ -2380,6 +2414,17 @@ dependencies = [ "syn 2.0.43", ] +[[package]] +name = "darling_macro" +version = "0.13.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9c972679f83bdf9c42bd905396b6c3588a843a17f0f16dfcfa3e2c5d57441835" +dependencies = [ + "darling_core 0.13.4", + "quote", + "syn 1.0.109", +] + [[package]] name = "darling_macro" version = "0.14.4" @@ -4988,7 +5033,7 @@ dependencies = [ "regex", "serde", "serde_json", - "serde_with", + "serde_with 3.4.0", "smallvec", "snafu", "store-api", @@ -8391,6 +8436,16 @@ dependencies = [ "serde", ] +[[package]] +name = "serde_with" +version = "1.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "678b5a069e50bf00ecd22d0cd8ddf7c236f68581b03db652061ed5eb13a312ff" +dependencies = [ + "serde", + "serde_with_macros 1.5.2", +] + [[package]] name = "serde_with" version = "3.4.0" @@ -8404,10 +8459,22 @@ dependencies = [ "indexmap 2.1.0", "serde", "serde_json", - "serde_with_macros", + "serde_with_macros 3.4.0", "time", ] +[[package]] +name = "serde_with_macros" +version = "1.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e182d6ec6f05393cc0e5ed1bf81ad6db3a8feedf8ee515ecdd369809bcce8082" +dependencies = [ + "darling 0.13.4", + "proc-macro2", + "quote", + "syn 1.0.109", +] + [[package]] name = "serde_with_macros" version = "3.4.0" @@ -9448,6 +9515,23 @@ version = "0.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3369f5ac52d5eb6ab48c6b4ffdc8efbcad6b89c765749064ba298f2c68a16a76" +[[package]] +name = "testcontainers" +version = "0.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f83d2931d7f521af5bae989f716c3fa43a6af9af7ec7a5e21b59ae40878cec00" +dependencies = [ + "bollard-stubs", + "futures", + "hex", + "hmac", + "log", + "rand", + "serde", + "serde_json", + "sha2", +] + [[package]] name = "tests-integration" version = "0.5.0" @@ -9511,6 +9595,7 @@ dependencies = [ "substrait 0.5.0", "table", "tempfile", + "testcontainers", "time", "tokio", "tokio-postgres", diff --git a/tests-integration/Cargo.toml b/tests-integration/Cargo.toml index 8d311d3f7ea2..377c428c2a55 100644 --- a/tests-integration/Cargo.toml +++ b/tests-integration/Cargo.toml @@ -64,6 +64,7 @@ sqlx = { version = "0.6", features = [ substrait.workspace = true table.workspace = true tempfile.workspace = true +testcontainers = "0.15.0" time = "0.3" tokio.workspace = true tonic.workspace = true diff --git a/tests-integration/src/lib.rs b/tests-integration/src/lib.rs index b0b28ba4651c..8dc4bc5a9829 100644 --- a/tests-integration/src/lib.rs +++ b/tests-integration/src/lib.rs @@ -20,6 +20,7 @@ mod opentsdb; mod otlp; mod prom_store; pub mod test_util; +pub mod wal_util; mod standalone; #[cfg(test)] diff --git a/tests-integration/src/wal_util.rs b/tests-integration/src/wal_util.rs new file mode 100644 index 000000000000..28c04633b640 --- /dev/null +++ b/tests-integration/src/wal_util.rs @@ -0,0 +1,15 @@ +// 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. + +pub mod kafka; diff --git a/tests-integration/src/wal_util/kafka.rs b/tests-integration/src/wal_util/kafka.rs new file mode 100644 index 000000000000..ee71c6410b1c --- /dev/null +++ b/tests-integration/src/wal_util/kafka.rs @@ -0,0 +1,17 @@ +// 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. + +pub mod config; +mod image; +pub mod runtime; diff --git a/tests-integration/src/wal_util/kafka/config.rs b/tests-integration/src/wal_util/kafka/config.rs new file mode 100644 index 000000000000..4da810dfa93b --- /dev/null +++ b/tests-integration/src/wal_util/kafka/config.rs @@ -0,0 +1,93 @@ +// 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 std::collections::HashMap; + +use testcontainers::core::WaitFor; + +/// Through which port the Zookeeper node listens for external traffics, i.e. traffics from the Kafka node. +pub const ZOOKEEPER_PORT: u16 = 2181; +/// Through which port the Kafka node listens for internal traffics, i.e. traffics between Kafka nodes in the same Kafka cluster. +pub const KAFAK_LISTENER_PORT: u16 = 19092; +/// Through which port the Kafka node listens for external traffics, e.g. traffics from Kafka clients. +pub const KAFKA_ADVERTISED_LISTENER_PORT: u16 = 9092; + +/// Configurations for a Kafka runtime. +/// Since the runtime corresponds to a cluster with a single Kafka node and a single Zookeeper node, the ports are all singletons. +pub struct Config { + /// The name of the Kafka image hosted in the docker hub. + pub image_name: String, + /// The tag of the kafka image hosted in the docker hub. + /// Warning: please use a tag with long-term support. Do not use `latest` or any other tags that + /// the underlying image may suddenly change. + pub image_tag: String, + /// Through which port clients could connect with the runtime. + pub exposed_port: u16, + /// The runtime is regarded ready to be used if all ready conditions are met. + /// Warning: be sure to update the conditions when necessary if the image is altered. + pub ready_conditions: Vec, + /// The environment variables required to run the runtime. + /// Warning: be sure to update the environment variables when necessary if the image is altered. + pub env_vars: HashMap, +} + +impl Default for Config { + fn default() -> Self { + Self { + image_name: "confluentinc/cp-kafka".to_string(), + image_tag: "7.4.3".to_string(), + exposed_port: KAFKA_ADVERTISED_LISTENER_PORT, + // The runtime is safe to be used as long as this message is printed on stdout. + ready_conditions: vec![WaitFor::message_on_stdout( + "started (kafka.server.KafkaServer)", + )], + env_vars: build_env_vars( + ZOOKEEPER_PORT, + KAFAK_LISTENER_PORT, + KAFKA_ADVERTISED_LISTENER_PORT, + ), + } + } +} + +fn build_env_vars( + zookeeper_port: u16, + kafka_listener_port: u16, + kafka_advertised_listener_port: u16, +) -> HashMap { + [ + ( + "KAFKA_ZOOKEEPER_CONNECT".to_string(), + format!("localhost:{zookeeper_port}"), + ), + ( + "KAFKA_LISTENERS".to_string(), + format!("PLAINTEXT://0.0.0.0:{kafka_advertised_listener_port},PLAINTEXT://0.0.0.0:{kafka_listener_port}"), + ), + ( + "KAFKA_ADVERTISED_LISTENERS".to_string(), + format!("PLAINTEXT://localhost:{kafka_advertised_listener_port},PLAINTEXT://localhost:{kafka_listener_port}",), + ), + ( + "KAFKA_INTER_BROKER_LISTENER_NAME".to_string(), + "BROKER".to_string(), + ), + ("KAFKA_BROKER_ID".to_string(), "1".to_string()), + ( + "KAFKA_OFFSETS_TOPIC_REPLICATION_FACTOR".to_string(), + "1".to_string(), + ), + ] + .into() +} diff --git a/tests-integration/src/wal_util/kafka/image.rs b/tests-integration/src/wal_util/kafka/image.rs new file mode 100644 index 000000000000..90cb164efa4b --- /dev/null +++ b/tests-integration/src/wal_util/kafka/image.rs @@ -0,0 +1,77 @@ +// 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 testcontainers::core::{ContainerState, ExecCommand, WaitFor}; + +use crate::wal_util::kafka::config::{Config, ZOOKEEPER_PORT}; + +#[derive(Debug, Clone, Default)] +pub struct ImageArgs; + +impl testcontainers::ImageArgs for ImageArgs { + fn into_iterator(self) -> Box> { + Box::new( + vec![ + "/bin/bash".to_string(), + "-c".to_string(), + format!( + r#" + echo 'clientPort={}' > zookeeper.properties; + echo 'dataDir=/var/lib/zookeeper/data' >> zookeeper.properties; + echo 'dataLogDir=/var/lib/zookeeper/log' >> zookeeper.properties; + zookeeper-server-start zookeeper.properties & + . /etc/confluent/docker/bash-config && + /etc/confluent/docker/configure && + /etc/confluent/docker/launch + "#, + ZOOKEEPER_PORT, + ), + ] + .into_iter(), + ) + } +} + +#[derive(Default)] +pub struct Image { + config: Config, +} + +impl testcontainers::Image for Image { + type Args = ImageArgs; + + fn name(&self) -> String { + self.config.image_name.clone() + } + + fn tag(&self) -> String { + self.config.image_tag.clone() + } + + fn ready_conditions(&self) -> Vec { + self.config.ready_conditions.clone() + } + + fn env_vars(&self) -> Box + '_> { + Box::new(self.config.env_vars.iter()) + } + + fn expose_ports(&self) -> Vec { + vec![self.config.exposed_port] + } + + fn exec_after_start(&self, _cs: ContainerState) -> Vec { + vec![] + } +} diff --git a/tests-integration/src/wal_util/kafka/runtime.rs b/tests-integration/src/wal_util/kafka/runtime.rs new file mode 100644 index 000000000000..24ff66de8c7e --- /dev/null +++ b/tests-integration/src/wal_util/kafka/runtime.rs @@ -0,0 +1,38 @@ +// 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 testcontainers::clients::Cli as DockerCli; +use testcontainers::Container; + +use crate::wal_util::kafka::image::Image; + +/// A runtime running a cluster consisting of a single Kafka node and a single ZooKeeper node. +#[derive(Default)] +pub struct Runtime { + docker: DockerCli, +} + +impl Runtime { + /// Starts the runtime. The runtime terminates when the returned container is dropped. + pub async fn start(&self) -> Container { + self.docker.run(Image::default()) + } +} + +#[macro_export] +macro_rules! start_kafka { + () => { + let _ = Runtime::default().start(); + }; +} From 1e3cafe39fb6481187e65000fba6f3d7b4aacc69 Mon Sep 17 00:00:00 2001 From: niebayes Date: Wed, 27 Dec 2023 19:40:52 +0800 Subject: [PATCH 03/30] test: add test for kafka runtime --- tests-integration/src/wal_util/kafka.rs | 7 +++ .../src/wal_util/kafka/runtime.rs | 54 +++++++++++++++++-- 2 files changed, 56 insertions(+), 5 deletions(-) diff --git a/tests-integration/src/wal_util/kafka.rs b/tests-integration/src/wal_util/kafka.rs index ee71c6410b1c..2e14f90fc499 100644 --- a/tests-integration/src/wal_util/kafka.rs +++ b/tests-integration/src/wal_util/kafka.rs @@ -15,3 +15,10 @@ pub mod config; mod image; pub mod runtime; + +#[macro_export] +macro_rules! start_kafka { + () => { + let _ = $crate::wal_util::kafka::runtime::Runtime::default().start().await; + }; +} diff --git a/tests-integration/src/wal_util/kafka/runtime.rs b/tests-integration/src/wal_util/kafka/runtime.rs index 24ff66de8c7e..f593d6cf733f 100644 --- a/tests-integration/src/wal_util/kafka/runtime.rs +++ b/tests-integration/src/wal_util/kafka/runtime.rs @@ -30,9 +30,53 @@ impl Runtime { } } -#[macro_export] -macro_rules! start_kafka { - () => { - let _ = Runtime::default().start(); - }; +#[cfg(test)] +mod tests { + use rskafka::chrono::Utc; + use rskafka::client::partition::UnknownTopicHandling; + use rskafka::client::ClientBuilder; + use rskafka::record::Record; + + use crate::start_kafka; + + #[tokio::test] + async fn test_runtime() { + start_kafka!(); + + let bootstrap_brokers = vec![9092.to_string()]; + let client = ClientBuilder::new(bootstrap_brokers).build().await.unwrap(); + + // Creates a topic. + let topic = "test_topic"; + client + .controller_client() + .unwrap() + .create_topic(topic, 1, 1, 500) + .await + .unwrap(); + + // Produces a record. + let partition_client = client + .partition_client(topic, 0, UnknownTopicHandling::Error) + .await + .unwrap(); + let produced = vec![Record { + key: Some(b"111".to_vec()), + value: Some(b"222".to_vec()), + timestamp: Utc::now(), + headers: Default::default(), + }]; + let offset = partition_client + .produce(produced.clone(), Default::default()) + .await + .unwrap()[0]; + + // Consumes the record. + let consumed = partition_client + .fetch_records(offset, 1..4096, 500) + .await + .unwrap() + .0; + assert_eq!(produced[0], consumed[0].record); + } } From 90fd40a577d149b5192827d5a8b93a416e7dd710 Mon Sep 17 00:00:00 2001 From: niebayes Date: Wed, 27 Dec 2023 19:43:22 +0800 Subject: [PATCH 04/30] fix: format --- tests-integration/src/wal_util/kafka.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests-integration/src/wal_util/kafka.rs b/tests-integration/src/wal_util/kafka.rs index 2e14f90fc499..abde462d720d 100644 --- a/tests-integration/src/wal_util/kafka.rs +++ b/tests-integration/src/wal_util/kafka.rs @@ -19,6 +19,8 @@ pub mod runtime; #[macro_export] macro_rules! start_kafka { () => { - let _ = $crate::wal_util::kafka::runtime::Runtime::default().start().await; + let _ = $crate::wal_util::kafka::runtime::Runtime::default() + .start() + .await; }; } From 796aeae9e73574ae2d4f84d03af786e0c0626421 Mon Sep 17 00:00:00 2001 From: niebayes Date: Thu, 28 Dec 2023 15:25:28 +0800 Subject: [PATCH 05/30] chore: make kafka image ready to be used --- tests-integration/src/wal_util/kafka.rs | 10 --- .../src/wal_util/kafka/config.rs | 30 +++++-- tests-integration/src/wal_util/kafka/image.rs | 87 ++++++++++++++++++- .../src/wal_util/kafka/runtime.rs | 82 ----------------- 4 files changed, 107 insertions(+), 102 deletions(-) delete mode 100644 tests-integration/src/wal_util/kafka/runtime.rs diff --git a/tests-integration/src/wal_util/kafka.rs b/tests-integration/src/wal_util/kafka.rs index abde462d720d..d4d0f8773d91 100644 --- a/tests-integration/src/wal_util/kafka.rs +++ b/tests-integration/src/wal_util/kafka.rs @@ -14,13 +14,3 @@ pub mod config; mod image; -pub mod runtime; - -#[macro_export] -macro_rules! start_kafka { - () => { - let _ = $crate::wal_util::kafka::runtime::Runtime::default() - .start() - .await; - }; -} diff --git a/tests-integration/src/wal_util/kafka/config.rs b/tests-integration/src/wal_util/kafka/config.rs index 4da810dfa93b..ca5136aec1ae 100644 --- a/tests-integration/src/wal_util/kafka/config.rs +++ b/tests-integration/src/wal_util/kafka/config.rs @@ -16,15 +16,15 @@ use std::collections::HashMap; use testcontainers::core::WaitFor; -/// Through which port the Zookeeper node listens for external traffics, i.e. traffics from the Kafka node. +/// Through which port the Zookeeper node listens for external traffics, e.g. traffics from the Kafka node. pub const ZOOKEEPER_PORT: u16 = 2181; /// Through which port the Kafka node listens for internal traffics, i.e. traffics between Kafka nodes in the same Kafka cluster. -pub const KAFAK_LISTENER_PORT: u16 = 19092; +pub const KAFKA_LISTENER_PORT: u16 = 19092; /// Through which port the Kafka node listens for external traffics, e.g. traffics from Kafka clients. pub const KAFKA_ADVERTISED_LISTENER_PORT: u16 = 9092; /// Configurations for a Kafka runtime. -/// Since the runtime corresponds to a cluster with a single Kafka node and a single Zookeeper node, the ports are all singletons. +#[derive(Debug, Clone)] pub struct Config { /// The name of the Kafka image hosted in the docker hub. pub image_name: String, @@ -32,7 +32,10 @@ pub struct Config { /// Warning: please use a tag with long-term support. Do not use `latest` or any other tags that /// the underlying image may suddenly change. pub image_tag: String, - /// Through which port clients could connect with the runtime. + /// The runtime is running in a docker container and has its own network. In order to be used by the host machine, + /// the runtime must expose an internal port. For e.g. assume the runtime has an internal port 9092, + /// and the `exposed_port` is set to 9092, then the host machine can get a mapped external port with + /// `container.get_host_port_ipv4(exposed_port)`. With the mapped port, the host machine could connect with the runtime. pub exposed_port: u16, /// The runtime is regarded ready to be used if all ready conditions are met. /// Warning: be sure to update the conditions when necessary if the image is altered. @@ -42,6 +45,15 @@ pub struct Config { pub env_vars: HashMap, } +impl Config { + pub fn with_exposed_port(port: u16) -> Self { + Self { + exposed_port: port, + ..Default::default() + } + } +} + impl Default for Config { fn default() -> Self { Self { @@ -54,7 +66,7 @@ impl Default for Config { )], env_vars: build_env_vars( ZOOKEEPER_PORT, - KAFAK_LISTENER_PORT, + KAFKA_LISTENER_PORT, KAFKA_ADVERTISED_LISTENER_PORT, ), } @@ -73,11 +85,15 @@ fn build_env_vars( ), ( "KAFKA_LISTENERS".to_string(), - format!("PLAINTEXT://0.0.0.0:{kafka_advertised_listener_port},PLAINTEXT://0.0.0.0:{kafka_listener_port}"), + format!("PLAINTEXT://0.0.0.0:{kafka_advertised_listener_port},BROKER://0.0.0.0:{kafka_listener_port}"), ), ( "KAFKA_ADVERTISED_LISTENERS".to_string(), - format!("PLAINTEXT://localhost:{kafka_advertised_listener_port},PLAINTEXT://localhost:{kafka_listener_port}",), + format!("PLAINTEXT://localhost:{kafka_advertised_listener_port},BROKER://localhost:{kafka_listener_port}",), + ), + ( + "KAFKA_LISTENER_SECURITY_PROTOCOL_MAP".to_string(), + "BROKER:PLAINTEXT,PLAINTEXT:PLAINTEXT".to_string(), ), ( "KAFKA_INTER_BROKER_LISTENER_NAME".to_string(), diff --git a/tests-integration/src/wal_util/kafka/image.rs b/tests-integration/src/wal_util/kafka/image.rs index 90cb164efa4b..78f513dd8406 100644 --- a/tests-integration/src/wal_util/kafka/image.rs +++ b/tests-integration/src/wal_util/kafka/image.rs @@ -14,7 +14,9 @@ use testcontainers::core::{ContainerState, ExecCommand, WaitFor}; -use crate::wal_util::kafka::config::{Config, ZOOKEEPER_PORT}; +use crate::wal_util::kafka::config::{ + Config, KAFKA_ADVERTISED_LISTENER_PORT, KAFKA_LISTENER_PORT, ZOOKEEPER_PORT, +}; #[derive(Debug, Clone, Default)] pub struct ImageArgs; @@ -48,6 +50,13 @@ pub struct Image { config: Config, } +impl Image { + #[allow(unused)] + pub fn new(config: Config) -> Self { + Self { config } + } +} + impl testcontainers::Image for Image { type Args = ImageArgs; @@ -71,7 +80,79 @@ impl testcontainers::Image for Image { vec![self.config.exposed_port] } - fn exec_after_start(&self, _cs: ContainerState) -> Vec { - vec![] + fn exec_after_start(&self, cs: ContainerState) -> Vec { + let mut commands = vec![]; + let cmd = format!( + "kafka-configs --alter --bootstrap-server 0.0.0.0:{} --entity-type brokers --entity-name 1 --add-config advertised.listeners=[PLAINTEXT://127.0.0.1:{},BROKER://localhost:9092]", + KAFKA_LISTENER_PORT, + cs.host_port_ipv4(KAFKA_ADVERTISED_LISTENER_PORT), + ); + let ready_conditions = vec![WaitFor::message_on_stdout( + "Checking need to trigger auto leader balancing", + )]; + commands.push(ExecCommand { + cmd, + ready_conditions, + }); + commands + } +} + +#[cfg(test)] +mod tests { + use chrono::TimeZone; + use rskafka::chrono::Utc; + use rskafka::client::partition::UnknownTopicHandling; + use rskafka::client::ClientBuilder; + use rskafka::record::Record; + use testcontainers::clients::Cli as DockerCli; + + use crate::wal_util::kafka::config::{Config, KAFKA_ADVERTISED_LISTENER_PORT}; + use crate::wal_util::kafka::image::Image; + + #[tokio::test] + async fn test_image() { + // Starts a Kafka container. + let port = KAFKA_ADVERTISED_LISTENER_PORT; + let config = Config::with_exposed_port(port); + let docker = DockerCli::default(); + let container = docker.run(Image::new(config)); + + // Creates a Kafka client. + let bootstrap_brokers = vec![format!("127.0.0.1:{}", container.get_host_port_ipv4(port))]; + let client = ClientBuilder::new(bootstrap_brokers).build().await.unwrap(); + + // Creates a topic. + let topic = "test_topic"; + client + .controller_client() + .unwrap() + .create_topic(topic, 1, 1, 500) + .await + .unwrap(); + + // Produces a record. + let partition_client = client + .partition_client(topic, 0, UnknownTopicHandling::Error) + .await + .unwrap(); + let produced = vec![Record { + key: Some(b"111".to_vec()), + value: Some(b"222".to_vec()), + timestamp: Utc.timestamp_millis_opt(42).unwrap(), + headers: Default::default(), + }]; + let offset = partition_client + .produce(produced.clone(), Default::default()) + .await + .unwrap()[0]; + + // Consumes the record. + let consumed = partition_client + .fetch_records(offset, 1..4096, 500) + .await + .unwrap() + .0; + assert_eq!(produced[0], consumed[0].record); } } diff --git a/tests-integration/src/wal_util/kafka/runtime.rs b/tests-integration/src/wal_util/kafka/runtime.rs deleted file mode 100644 index f593d6cf733f..000000000000 --- a/tests-integration/src/wal_util/kafka/runtime.rs +++ /dev/null @@ -1,82 +0,0 @@ -// 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 testcontainers::clients::Cli as DockerCli; -use testcontainers::Container; - -use crate::wal_util::kafka::image::Image; - -/// A runtime running a cluster consisting of a single Kafka node and a single ZooKeeper node. -#[derive(Default)] -pub struct Runtime { - docker: DockerCli, -} - -impl Runtime { - /// Starts the runtime. The runtime terminates when the returned container is dropped. - pub async fn start(&self) -> Container { - self.docker.run(Image::default()) - } -} - -#[cfg(test)] -mod tests { - use rskafka::chrono::Utc; - use rskafka::client::partition::UnknownTopicHandling; - use rskafka::client::ClientBuilder; - use rskafka::record::Record; - - use crate::start_kafka; - - #[tokio::test] - async fn test_runtime() { - start_kafka!(); - - let bootstrap_brokers = vec![9092.to_string()]; - let client = ClientBuilder::new(bootstrap_brokers).build().await.unwrap(); - - // Creates a topic. - let topic = "test_topic"; - client - .controller_client() - .unwrap() - .create_topic(topic, 1, 1, 500) - .await - .unwrap(); - - // Produces a record. - let partition_client = client - .partition_client(topic, 0, UnknownTopicHandling::Error) - .await - .unwrap(); - let produced = vec![Record { - key: Some(b"111".to_vec()), - value: Some(b"222".to_vec()), - timestamp: Utc::now(), - headers: Default::default(), - }]; - let offset = partition_client - .produce(produced.clone(), Default::default()) - .await - .unwrap()[0]; - - // Consumes the record. - let consumed = partition_client - .fetch_records(offset, 1..4096, 500) - .await - .unwrap() - .0; - assert_eq!(produced[0], consumed[0].record); - } -} From 18513936a5c9df64baabeb52bebd17d1f7532bab Mon Sep 17 00:00:00 2001 From: niebayes Date: Fri, 29 Dec 2023 01:23:00 +0800 Subject: [PATCH 06/30] feat: add entry builder --- Cargo.lock | 2 + src/common/test-util/Cargo.toml | 2 + src/common/test-util/src/lib.rs | 1 + src/common/test-util/src/wal.rs | 15 +++ src/common/test-util/src/wal/kafka.rs | 17 +++ .../test-util/src/wal/kafka/entry_builder.rs | 110 ++++++++++++++++++ src/log-store/src/kafka.rs | 2 +- src/store-api/src/logstore.rs | 6 +- tests-integration/src/wal_util.rs | 5 + tests-integration/src/wal_util/kafka.rs | 2 +- .../src/wal_util/kafka/config.rs | 9 -- tests-integration/src/wal_util/kafka/image.rs | 12 +- 12 files changed, 165 insertions(+), 18 deletions(-) create mode 100644 src/common/test-util/src/wal.rs create mode 100644 src/common/test-util/src/wal/kafka.rs create mode 100644 src/common/test-util/src/wal/kafka/entry_builder.rs diff --git a/Cargo.lock b/Cargo.lock index d8b15f1bb700..09fbe9275128 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1983,8 +1983,10 @@ dependencies = [ name = "common-test-util" version = "0.5.0" dependencies = [ + "log-store", "once_cell", "rand", + "store-api", "tempfile", ] diff --git a/src/common/test-util/Cargo.toml b/src/common/test-util/Cargo.toml index 60e854740643..a1f8def92d01 100644 --- a/src/common/test-util/Cargo.toml +++ b/src/common/test-util/Cargo.toml @@ -5,6 +5,8 @@ edition.workspace = true license.workspace = true [dependencies] +log-store.workspace = true once_cell.workspace = true rand.workspace = true tempfile.workspace = true +store-api.workspace = true diff --git a/src/common/test-util/src/lib.rs b/src/common/test-util/src/lib.rs index ef6ff4696861..f4a9c747e518 100644 --- a/src/common/test-util/src/lib.rs +++ b/src/common/test-util/src/lib.rs @@ -20,6 +20,7 @@ use std::sync::LazyLock; pub mod ports; pub mod temp_dir; +pub mod wal; // Rust is working on an env possibly named `CARGO_WORKSPACE_DIR` to find the root path to the // workspace, see https://github.com/rust-lang/cargo/issues/3946. diff --git a/src/common/test-util/src/wal.rs b/src/common/test-util/src/wal.rs new file mode 100644 index 000000000000..28c04633b640 --- /dev/null +++ b/src/common/test-util/src/wal.rs @@ -0,0 +1,15 @@ +// 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. + +pub mod kafka; diff --git a/src/common/test-util/src/wal/kafka.rs b/src/common/test-util/src/wal/kafka.rs new file mode 100644 index 000000000000..bd615f8a4af3 --- /dev/null +++ b/src/common/test-util/src/wal/kafka.rs @@ -0,0 +1,17 @@ +// 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. + +pub mod entry_builder; + +pub use crate::wal::kafka::entry_builder::EntryBuilder; diff --git a/src/common/test-util/src/wal/kafka/entry_builder.rs b/src/common/test-util/src/wal/kafka/entry_builder.rs new file mode 100644 index 000000000000..09910c6c7077 --- /dev/null +++ b/src/common/test-util/src/wal/kafka/entry_builder.rs @@ -0,0 +1,110 @@ +// 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 std::sync::atomic::{AtomicU64 as AtomicEntryId, Ordering}; +use std::sync::Mutex; + +use log_store::kafka::{EntryImpl, NamespaceImpl}; +use rand::rngs::ThreadRng; +use rand::seq::SliceRandom; +use rand::{thread_rng, Rng}; +use store_api::logstore::EntryId; + +const DEFAULT_DATA: &[u8; 10] = b"[greptime]"; + +/// A builder for building entries for a namespace. +pub struct EntryBuilder { + /// The namespace of the entries. + ns: NamespaceImpl, + /// The next entry id to allocate. It starts from 0 by default. + next_entry_id: AtomicEntryId, + /// A generator for supporting random data generation. + /// Wrapped with Mutex> to provide interior mutability. + rng: Mutex>, + /// The data pool from which random data is constructed. + data_pool: Vec, +} + +impl EntryBuilder { + /// Creates an EntryBuilder for the given namespace. + pub fn new(ns: NamespaceImpl) -> Self { + // Makes a data pool with alphabets and numbers. + let data_pool = ('a'..='z') + .chain('A'..='Z') + .chain((0..=9).map(|digit| char::from_digit(digit, 10).unwrap())) + .map(|c| c as u8) + .collect::>(); + Self { + ns, + next_entry_id: AtomicEntryId::new(0), + rng: Mutex::new(Some(thread_rng())), + data_pool, + } + } + + /// Sets the next entry id to the given entry id. + pub fn next_entry_id(self, entry_id: EntryId) -> Self { + Self { + next_entry_id: AtomicEntryId::new(entry_id), + ..self + } + } + + /// Skips the next `step` entry ids and returns the next entry id after the stepping. + pub fn skip(&mut self, step: EntryId) -> EntryId { + let old = self.next_entry_id.fetch_add(step, Ordering::Relaxed); + old + step + } + + /// Builds an entry with the given data. + pub fn with_data>(&self, data: D) -> EntryImpl { + EntryImpl { + data: data.as_ref().to_vec(), + id: self.alloc_entry_id(), + ns: self.ns.clone(), + } + } + + /// Builds an entry with the default data. + pub fn with_default_data(&self) -> EntryImpl { + EntryImpl { + data: DEFAULT_DATA.to_vec(), + id: self.alloc_entry_id(), + ns: self.ns.clone(), + } + } + + /// Builds an entry with random data. + pub fn with_random_data(&self) -> EntryImpl { + EntryImpl { + data: self.make_random_data(), + id: self.alloc_entry_id(), + ns: self.ns.clone(), + } + } + + fn alloc_entry_id(&self) -> EntryId { + self.next_entry_id.fetch_add(1, Ordering::Relaxed) + } + + fn make_random_data(&self) -> Vec { + let mut rng_guard = self.rng.lock().unwrap(); + let mut rng = rng_guard.as_mut().unwrap(); + let amount = rng.gen_range(0..self.data_pool.len()); + self.data_pool + .choose_multiple(&mut rng, amount) + .map(|x| *x) + .collect() + } +} diff --git a/src/log-store/src/kafka.rs b/src/log-store/src/kafka.rs index a08afc508eac..01448ce91cdb 100644 --- a/src/log-store/src/kafka.rs +++ b/src/log-store/src/kafka.rs @@ -41,7 +41,7 @@ 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) } } diff --git a/src/store-api/src/logstore.rs b/src/store-api/src/logstore.rs index 16809c26b1a1..903de87ee4f0 100644 --- a/src/store-api/src/logstore.rs +++ b/src/store-api/src/logstore.rs @@ -19,9 +19,11 @@ use std::collections::HashMap; use common_config::wal::WalOptions; use common_error::ext::ErrorExt; -use crate::logstore::entry::{Entry, Id as EntryId}; +use crate::logstore::entry::Entry; +pub use crate::logstore::entry::Id as EntryId; use crate::logstore::entry_stream::SendableEntryStream; -use crate::logstore::namespace::{Id as NamespaceId, Namespace}; +pub use crate::logstore::namespace::Id as NamespaceId; +use crate::logstore::namespace::Namespace; pub mod entry; pub mod entry_stream; diff --git a/tests-integration/src/wal_util.rs b/tests-integration/src/wal_util.rs index 28c04633b640..150cc065a471 100644 --- a/tests-integration/src/wal_util.rs +++ b/tests-integration/src/wal_util.rs @@ -13,3 +13,8 @@ // limitations under the License. pub mod kafka; + +pub use testcontainers::clients::Cli as DockerCli; + +pub use crate::wal_util::kafka::config::KAFKA_ADVERTISED_LISTENER_PORT as DEFAULT_EXPOSED_PORT; +pub use crate::wal_util::kafka::image::Image as KafkaImage; diff --git a/tests-integration/src/wal_util/kafka.rs b/tests-integration/src/wal_util/kafka.rs index d4d0f8773d91..b4289768ed83 100644 --- a/tests-integration/src/wal_util/kafka.rs +++ b/tests-integration/src/wal_util/kafka.rs @@ -13,4 +13,4 @@ // limitations under the License. pub mod config; -mod image; +pub mod image; diff --git a/tests-integration/src/wal_util/kafka/config.rs b/tests-integration/src/wal_util/kafka/config.rs index ca5136aec1ae..fad01474888c 100644 --- a/tests-integration/src/wal_util/kafka/config.rs +++ b/tests-integration/src/wal_util/kafka/config.rs @@ -45,15 +45,6 @@ pub struct Config { pub env_vars: HashMap, } -impl Config { - pub fn with_exposed_port(port: u16) -> Self { - Self { - exposed_port: port, - ..Default::default() - } - } -} - impl Default for Config { fn default() -> Self { Self { diff --git a/tests-integration/src/wal_util/kafka/image.rs b/tests-integration/src/wal_util/kafka/image.rs index 78f513dd8406..f509d5f60a00 100644 --- a/tests-integration/src/wal_util/kafka/image.rs +++ b/tests-integration/src/wal_util/kafka/image.rs @@ -51,8 +51,11 @@ pub struct Image { } impl Image { - #[allow(unused)] - pub fn new(config: Config) -> Self { + pub fn with_exposed_port(port: u16) -> Self { + let config = Config { + exposed_port: port, + ..Default::default() + }; Self { config } } } @@ -107,16 +110,15 @@ mod tests { use rskafka::record::Record; use testcontainers::clients::Cli as DockerCli; - use crate::wal_util::kafka::config::{Config, KAFKA_ADVERTISED_LISTENER_PORT}; + use crate::wal_util::kafka::config::KAFKA_ADVERTISED_LISTENER_PORT; use crate::wal_util::kafka::image::Image; #[tokio::test] async fn test_image() { // Starts a Kafka container. let port = KAFKA_ADVERTISED_LISTENER_PORT; - let config = Config::with_exposed_port(port); let docker = DockerCli::default(); - let container = docker.run(Image::new(config)); + let container = docker.run(Image::with_exposed_port(port)); // Creates a Kafka client. let bootstrap_brokers = vec![format!("127.0.0.1:{}", container.get_host_port_ipv4(port))]; From ef031eebaada9979e3c957e0767111ba24ecef8d Mon Sep 17 00:00:00 2001 From: niebayes Date: Fri, 29 Dec 2023 17:11:41 +0800 Subject: [PATCH 07/30] tmp --- Cargo.lock | 7 +- src/common/test-util/Cargo.toml | 5 +- src/common/test-util/src/wal.rs | 5 + src/common/test-util/src/wal/kafka.rs | 5 +- .../common/test-util/src/wal}/kafka/config.rs | 0 .../common/test-util/src/wal}/kafka/image.rs | 9 +- src/log-store/Cargo.toml | 4 +- src/log-store/src/kafka.rs | 2 +- src/log-store/src/kafka/client_manager.rs | 103 +++- src/log-store/src/kafka/log_store.rs | 55 +- src/log-store/src/kafka/record_utils.rs | 5 + src/log-store/src/test_util.rs | 1 + .../log-store/src/test_util/kafka.rs | 15 +- .../src/test_util}/kafka/entry_builder.rs | 5 +- .../src/test_util/kafka/topic_builder.rs | 94 ++++ tests-integration/src/lib.rs | 1 - tests-integration/src/wal_util/kafka.rs | 16 - tests-integration/tests/main.rs | 2 - tests-integration/tests/wal.rs | 521 +++++++++--------- 19 files changed, 540 insertions(+), 315 deletions(-) rename {tests-integration/src/wal_util => src/common/test-util/src/wal}/kafka/config.rs (100%) rename {tests-integration/src/wal_util => src/common/test-util/src/wal}/kafka/image.rs (95%) rename tests-integration/src/wal_util.rs => src/log-store/src/test_util/kafka.rs (53%) rename src/{common/test-util/src/wal => log-store/src/test_util}/kafka/entry_builder.rs (97%) create mode 100644 src/log-store/src/test_util/kafka/topic_builder.rs delete mode 100644 tests-integration/src/wal_util/kafka.rs diff --git a/Cargo.lock b/Cargo.lock index 09fbe9275128..072d2a5559d8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1983,11 +1983,12 @@ dependencies = [ name = "common-test-util" version = "0.5.0" dependencies = [ - "log-store", "once_cell", "rand", - "store-api", + "rskafka", "tempfile", + "testcontainers", + "tokio", ] [[package]] @@ -4542,6 +4543,7 @@ dependencies = [ "async-trait", "byteorder", "bytes", + "chrono", "common-base", "common-config", "common-error", @@ -4550,7 +4552,6 @@ dependencies = [ "common-runtime", "common-telemetry", "common-test-util", - "dashmap", "futures", "futures-util", "protobuf", diff --git a/src/common/test-util/Cargo.toml b/src/common/test-util/Cargo.toml index a1f8def92d01..e3caa859d21f 100644 --- a/src/common/test-util/Cargo.toml +++ b/src/common/test-util/Cargo.toml @@ -5,8 +5,9 @@ edition.workspace = true license.workspace = true [dependencies] -log-store.workspace = true once_cell.workspace = true rand.workspace = true +rskafka.workspace = true tempfile.workspace = true -store-api.workspace = true +testcontainers = "0.15.0" +tokio.workspace = true diff --git a/src/common/test-util/src/wal.rs b/src/common/test-util/src/wal.rs index 28c04633b640..c1085a8f1e17 100644 --- a/src/common/test-util/src/wal.rs +++ b/src/common/test-util/src/wal.rs @@ -13,3 +13,8 @@ // limitations under the License. pub mod kafka; + +pub use testcontainers::clients::Cli as DockerCli; + +pub use crate::wal::kafka::config::KAFKA_ADVERTISED_LISTENER_PORT as DEFAULT_EXPOSED_PORT; +pub use crate::wal::kafka::image::Image as KafkaImage; diff --git a/src/common/test-util/src/wal/kafka.rs b/src/common/test-util/src/wal/kafka.rs index bd615f8a4af3..427cbea89a71 100644 --- a/src/common/test-util/src/wal/kafka.rs +++ b/src/common/test-util/src/wal/kafka.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -pub mod entry_builder; +pub mod config; +pub mod image; -pub use crate::wal::kafka::entry_builder::EntryBuilder; +pub const BROKER_ENDPOINTS_KEY: &str = "GT_KAFKA_ENDPOINTS"; diff --git a/tests-integration/src/wal_util/kafka/config.rs b/src/common/test-util/src/wal/kafka/config.rs similarity index 100% rename from tests-integration/src/wal_util/kafka/config.rs rename to src/common/test-util/src/wal/kafka/config.rs diff --git a/tests-integration/src/wal_util/kafka/image.rs b/src/common/test-util/src/wal/kafka/image.rs similarity index 95% rename from tests-integration/src/wal_util/kafka/image.rs rename to src/common/test-util/src/wal/kafka/image.rs index f509d5f60a00..3428549a4c4b 100644 --- a/tests-integration/src/wal_util/kafka/image.rs +++ b/src/common/test-util/src/wal/kafka/image.rs @@ -14,7 +14,7 @@ use testcontainers::core::{ContainerState, ExecCommand, WaitFor}; -use crate::wal_util::kafka::config::{ +use crate::wal::kafka::config::{ Config, KAFKA_ADVERTISED_LISTENER_PORT, KAFKA_LISTENER_PORT, ZOOKEEPER_PORT, }; @@ -103,15 +103,14 @@ impl testcontainers::Image for Image { #[cfg(test)] mod tests { - use chrono::TimeZone; - use rskafka::chrono::Utc; + use rskafka::chrono::{TimeZone, Utc}; use rskafka::client::partition::UnknownTopicHandling; use rskafka::client::ClientBuilder; use rskafka::record::Record; use testcontainers::clients::Cli as DockerCli; - use crate::wal_util::kafka::config::KAFKA_ADVERTISED_LISTENER_PORT; - use crate::wal_util::kafka::image::Image; + use crate::wal::kafka::config::KAFKA_ADVERTISED_LISTENER_PORT; + use crate::wal::kafka::image::Image; #[tokio::test] async fn test_image() { diff --git a/src/log-store/Cargo.toml b/src/log-store/Cargo.toml index 7941cd11e918..afc3ddbab934 100644 --- a/src/log-store/Cargo.toml +++ b/src/log-store/Cargo.toml @@ -14,6 +14,7 @@ async-stream.workspace = true async-trait.workspace = true byteorder = "1.4" bytes.workspace = true +chrono.workspace = true common-base.workspace = true common-config.workspace = true common-error.workspace = true @@ -21,11 +22,12 @@ common-macro.workspace = true common-meta.workspace = true common-runtime.workspace = true common-telemetry.workspace = true -dashmap.workspace = true +common-test-util.workspace = true futures-util.workspace = true futures.workspace = true protobuf = { version = "2", features = ["bytes"] } raft-engine.workspace = true +rand.workspace = true rskafka.workspace = true serde.workspace = true serde_json.workspace = true diff --git a/src/log-store/src/kafka.rs b/src/log-store/src/kafka.rs index 01448ce91cdb..9b9493d4fe9e 100644 --- a/src/log-store/src/kafka.rs +++ b/src/log-store/src/kafka.rs @@ -77,7 +77,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() diff --git a/src/log-store/src/kafka/client_manager.rs b/src/log-store/src/kafka/client_manager.rs index e272840201bb..1e49f977de4d 100644 --- a/src/log-store/src/kafka/client_manager.rs +++ b/src/log-store/src/kafka/client_manager.rs @@ -12,17 +12,19 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::hash_map::Entry; +use std::collections::HashMap; use std::sync::Arc; use common_config::wal::{KafkaConfig, KafkaWalTopic as Topic}; -use dashmap::mapref::entry::Entry as DashMapEntry; -use dashmap::DashMap; +use common_telemetry::debug; use rskafka::client::partition::{PartitionClient, UnknownTopicHandling}; use rskafka::client::producer::aggregator::RecordAggregator; use rskafka::client::producer::{BatchProducer, BatchProducerBuilder}; use rskafka::client::{Client as RsKafkaClient, ClientBuilder}; use rskafka::BackoffConfig; use snafu::ResultExt; +use tokio::sync::{Mutex as TokioMutex, RwLock as TokioRwLock}; use crate::error::{BuildClientSnafu, BuildPartitionClientSnafu, Result}; @@ -67,7 +69,7 @@ pub(crate) struct ClientManager { client_factory: RsKafkaClient, /// A pool maintaining a collection of clients. /// Key: a topic. Value: the associated client of the topic. - client_pool: DashMap, + client_pool: TokioMutex>, } impl ClientManager { @@ -88,26 +90,28 @@ impl ClientManager { broker_endpoints: config.broker_endpoints.clone(), })?; + debug!("Created a ClientManager"); + Ok(Self { config: config.clone(), client_factory: client, - client_pool: DashMap::new(), + client_pool: TokioMutex::new(HashMap::new()), }) } /// Gets the client associated with the topic. If the client does not exist, a new one will /// be created and returned. pub(crate) async fn get_or_insert(&self, topic: &Topic) -> Result { - match self.client_pool.entry(topic.to_string()) { - DashMapEntry::Occupied(entry) => Ok(entry.get().clone()), - DashMapEntry::Vacant(entry) => { - let topic_client = self.try_create_client(topic).await?; - Ok(entry.insert(topic_client).clone()) - } + let mut client_pool = self.client_pool.lock().await; + if let Entry::Vacant(entry) = client_pool.entry(topic.to_string()) { + entry.insert(self.try_create_client(topic).await?); } + Ok(client_pool[topic].clone()) } async fn try_create_client(&self, topic: &Topic) -> Result { + debug!("Try to create client for topic {}", topic); + // Sets to Retry to retry connecting if the kafka cluter replies with an UnknownTopic error. // That's because the topic is believed to exist as the metasrv is expected to create required topics upon start. // The reconnecting won't stop until succeed or a different error returns. @@ -121,6 +125,85 @@ impl ClientManager { }) .map(Arc::new)?; + debug!("Created a client for topic {}", topic); + Ok(Client::new(raw_client, &self.config)) } } + +#[cfg(test)] +mod tests { + use common_test_util::wal::kafka::BROKER_ENDPOINTS_KEY; + + use super::*; + use crate::test_util::kafka::topic_builder::{Affix, TopicBuilder}; + + /// Checks clients for the given topics are created. + async fn ensure_topics_exist(topics: &[Topic], client_manager: &ClientManager) { + let client_pool = client_manager.client_pool.lock().await; + let all_exist = topics.iter().all(|topic| client_pool.contains_key(topic)); + assert_eq!(all_exist, true); + } + + /// Sends `get_or_insert` requests sequentially to the client manager, and checks if it could handle them correctly. + #[tokio::test] + async fn test_sequential() { + let broker_endpoints = std::env::var(BROKER_ENDPOINTS_KEY) + .unwrap() + .split(',') + .map(ToString::to_string) + .collect::>(); + if broker_endpoints.is_empty() { + return; + } + + let client = ClientBuilder::new(broker_endpoints.clone()) + .build() + .await + .unwrap(); + let ctrl_client = client.controller_client().unwrap(); + + let topic_builder = TopicBuilder::default() + .with_prefix(Affix::Fixed("test_sequential".to_string())) + .with_suffix(Affix::TimeNow); + + + let config = KafkaConfig { + broker_endpoints, + ..Default::default() + }; + let manager = ClientManager::try_new(&config).await.unwrap(); + + // Constructs a collection of mock topics. + let num_topics = 256; + let topics = (0..num_topics) + .map(|i| format!("topic_{i}")) + .collect::>(); + + // Gets all clients sequentially. + for topic in topics.iter() { + manager.get_or_insert(topic).await.unwrap(); + } + ensure_topics_exist(&topics, &manager).await; + } + + /// Sends `get_or_insert` requests in parallel to the client manager, and checks if it could handle them correctly. + #[tokio::test] + async fn test_parallel() { + let manager = ClientManager::try_new(&config).await.unwrap(); + + // Constructs a collection of mock topics. + let num_topics = 256; + let topics = (0..num_topics) + .map(|i| format!("topic_{i}")) + .collect::>(); + + // Gets all clients in parallel. + let tasks = topics + .iter() + .map(|topic| manager.get_or_insert(topic)) + .collect::>(); + futures::future::try_join_all(tasks).await.unwrap(); + ensure_topics_exist(&topics, &manager).await; + } +} diff --git a/src/log-store/src/kafka/log_store.rs b/src/log-store/src/kafka/log_store.rs index 73b0fe1de2a9..f0170245e5e2 100644 --- a/src/log-store/src/kafka/log_store.rs +++ b/src/log-store/src/kafka/log_store.rs @@ -85,7 +85,7 @@ impl LogStore for KafkaLogStore { /// Appends a batch of entries and returns a response containing a map where the key is a region id /// while the value is the id of the last successfully written entry of the region. async fn append_batch(&self, entries: Vec) -> Result { - debug!("LogStore handles append_batch with entries {:?}", entries); + println!("LogStore handles append_batch with entries {:?}", entries); if entries.is_empty() { return Ok(AppendBatchResponse::default()); @@ -101,12 +101,10 @@ impl LogStore for KafkaLogStore { } // Builds a record from entries belong to a region and produces them to kafka server. - let region_ids = producers.keys().cloned().collect::>(); - - let tasks = producers - .into_values() - .map(|producer| producer.produce(&self.client_manager)) - .collect::>(); + let (region_ids, tasks): (Vec<_>, Vec<_>) = producers + .into_iter() + .map(|(id, producer)| (id, producer.produce(&self.client_manager))) + .unzip(); // Each produce operation returns a kafka offset of the produced record. // The offsets are then converted to entry ids. let entry_ids = futures::future::try_join_all(tasks) @@ -114,11 +112,20 @@ impl LogStore for KafkaLogStore { .into_iter() .map(TryInto::try_into) .collect::>>()?; - debug!("The entries are appended at offsets {:?}", entry_ids); + println!("The entries are appended at offsets {:?}", entry_ids); - Ok(AppendBatchResponse { - last_entry_ids: region_ids.into_iter().zip(entry_ids).collect(), - }) + let last_entry_ids = region_ids + .into_iter() + .zip(entry_ids) + .collect::>(); + for (region, last_entry_id) in last_entry_ids.iter() { + println!( + "The entries for region {} are appended at offset {}", + region, last_entry_id + ); + } + + Ok(AppendBatchResponse { last_entry_ids }) } /// Creates a new `EntryStream` to asynchronously generates `Entry` with entry ids @@ -148,14 +155,19 @@ impl LogStore for KafkaLogStore { .await .context(GetOffsetSnafu { ns: ns.clone() })? - 1; - // Reads entries with offsets in the range [start_offset, end_offset). + // Reads entries with offsets in the range [start_offset, end_offset]. let start_offset = Offset::try_from(entry_id)?.0; + println!( + "Start reading entries in range [{}, {}] for ns {}", + start_offset, end_offset, ns + ); + // Abort if there're no new entries. // FIXME(niebayes): how come this case happens? if start_offset > end_offset { - warn!( - "No new entries for ns {} in range [{}, {})", + println!( + "No new entries for ns {} in range [{}, {}]", ns, start_offset, end_offset ); return Ok(futures_util::stream::empty().boxed()); @@ -166,8 +178,8 @@ impl LogStore for KafkaLogStore { .with_max_wait_ms(self.config.produce_record_timeout.as_millis() as i32) .build(); - debug!( - "Built a stream consumer for ns {} to consume entries in range [{}, {})", + println!( + "Built a stream consumer for ns {} to consume entries in range [{}, {}]", ns, start_offset, end_offset ); @@ -181,7 +193,7 @@ impl LogStore for KafkaLogStore { ns: ns_clone.clone(), })?; let record_offset = record.offset; - debug!( + println!( "Read a record at offset {} for ns {}, high watermark: {}", record_offset, ns_clone, high_watermark ); @@ -192,14 +204,19 @@ impl LogStore for KafkaLogStore { if let Some(entry) = entries.first() && entry.ns.region_id == region_id { + println!("{} entries are yielded for ns {}", entries.len(), ns_clone); yield Ok(entries); } else { - yield Ok(vec![]); + println!( + "{} entries are filtered out for ns {}", + entries.len(), + ns_clone + ); } // Terminates the stream if the entry with the end offset was read. if record_offset >= end_offset { - debug!( + println!( "Stream consumer for ns {} terminates at offset {}", ns_clone, record_offset ); diff --git a/src/log-store/src/kafka/record_utils.rs b/src/log-store/src/kafka/record_utils.rs index 3707b873f3e3..2ef5b2db81b2 100644 --- a/src/log-store/src/kafka/record_utils.rs +++ b/src/log-store/src/kafka/record_utils.rs @@ -89,6 +89,8 @@ impl RecordProducer { pub(crate) async fn produce(self, client_manager: &ClientManagerRef) -> Result { ensure!(!self.entries.is_empty(), EmptyEntriesSnafu); + println!("Start producing a record for ns {}", self.ns); + // Produces the record through a client. The client determines when to send the record to kafka server. let client = client_manager .get_or_insert(&self.ns.topic) @@ -100,6 +102,9 @@ impl RecordProducer { } .build() })?; + + println!("Got the client for ns {}", self.ns); + client .producer .produce(encode_to_record(self.ns.clone(), self.entries)?) diff --git a/src/log-store/src/test_util.rs b/src/log-store/src/test_util.rs index 973d6d3f9720..52ffcca39370 100644 --- a/src/log-store/src/test_util.rs +++ b/src/log-store/src/test_util.rs @@ -12,4 +12,5 @@ // See the License for the specific language governing permissions and // limitations under the License. +pub mod kafka; pub mod log_store_util; diff --git a/tests-integration/src/wal_util.rs b/src/log-store/src/test_util/kafka.rs similarity index 53% rename from tests-integration/src/wal_util.rs rename to src/log-store/src/test_util/kafka.rs index 150cc065a471..b66e692836d7 100644 --- a/tests-integration/src/wal_util.rs +++ b/src/log-store/src/test_util/kafka.rs @@ -12,9 +12,16 @@ // See the License for the specific language governing permissions and // limitations under the License. -pub mod kafka; +pub mod entry_builder; +pub mod topic_builder; -pub use testcontainers::clients::Cli as DockerCli; +use common_config::wal::KafkaWalTopic as Topic; -pub use crate::wal_util::kafka::config::KAFKA_ADVERTISED_LISTENER_PORT as DEFAULT_EXPOSED_PORT; -pub use crate::wal_util::kafka::image::Image as KafkaImage; +pub use crate::test_util::kafka::entry_builder::EntryBuilder; +pub use crate::test_util::kafka::topic_builder::TopicBuilder; + +/// Creates `num_topiocs` number of topics with the given TopicBuilder. +/// Requests for creating these topics on the Kafka cluster if the `broker_endpoints` is not empty. +pub fn create_topics(num_topics: usize, builder: TopicBuilder, broker_endpoints: Vec) -> Vec { + +} diff --git a/src/common/test-util/src/wal/kafka/entry_builder.rs b/src/log-store/src/test_util/kafka/entry_builder.rs similarity index 97% rename from src/common/test-util/src/wal/kafka/entry_builder.rs rename to src/log-store/src/test_util/kafka/entry_builder.rs index 09910c6c7077..3e234c78d6e4 100644 --- a/src/common/test-util/src/wal/kafka/entry_builder.rs +++ b/src/log-store/src/test_util/kafka/entry_builder.rs @@ -15,12 +15,13 @@ use std::sync::atomic::{AtomicU64 as AtomicEntryId, Ordering}; use std::sync::Mutex; -use log_store::kafka::{EntryImpl, NamespaceImpl}; use rand::rngs::ThreadRng; use rand::seq::SliceRandom; use rand::{thread_rng, Rng}; use store_api::logstore::EntryId; +use crate::kafka::{EntryImpl, NamespaceImpl}; + const DEFAULT_DATA: &[u8; 10] = b"[greptime]"; /// A builder for building entries for a namespace. @@ -108,3 +109,5 @@ impl EntryBuilder { .collect() } } + +// TODO(niebayes): add tests for EntryBuilder. diff --git a/src/log-store/src/test_util/kafka/topic_builder.rs b/src/log-store/src/test_util/kafka/topic_builder.rs new file mode 100644 index 000000000000..485b846cc6d4 --- /dev/null +++ b/src/log-store/src/test_util/kafka/topic_builder.rs @@ -0,0 +1,94 @@ +// 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 std::collections::HashSet; + +use common_config::wal::KafkaWalTopic as Topic; + +/// Things need to bo inserted at the front or the back of the topic. +#[derive(Debug, Default)] +pub enum Affix { + /// Inserts a provided string to each topic. + Fixed(String), + /// Computes the current time for each topic and inserts it into the topic. + TimeNow, + /// Nothing to be inserted. + #[default] + Nothing, +} + +impl ToString for Affix { + fn to_string(&self) -> String { + match self { + Affix::Fixed(s) => s.to_string(), + Affix::TimeNow => chrono::Local::now().to_string(), + Affix::Nothing => String::default(), + } + } +} + +/// A builder for building topics on demand. +pub struct TopicBuilder { + /// A prefix to be inserted at the front of each topic. + prefix: Affix, + /// A suffix to be inserted at the back of each topic. + suffix: Affix, + /// Topics built so far. Used to filter out duplicate topics. + created: HashSet, +} + +impl Default for TopicBuilder { + fn default() -> Self { + Self { + prefix: Affix::Nothing, + suffix: Affix::Nothing, + created: HashSet::with_capacity(256), + } + } +} + +impl TopicBuilder { + /// Overrides the current prefix with the given prefix. + pub fn with_prefix(self, prefix: Affix) -> Self { + Self { prefix, ..self } + } + + /// Overrides the current suffix with the given suffix. + pub fn with_suffix(self, suffix: Affix) -> Self { + Self { suffix, ..self } + } + + /// Builds a topic by inserting a prefix and a suffix into the given topic. + pub fn build(&mut self, topic: &Topic) -> Topic { + const ITERS: usize = 24; + for _ in 0..ITERS { + let topic = format!( + "{}_{}_{}", + self.prefix.to_string(), + topic, + self.suffix.to_string() + ); + if !self.created.contains(&topic) { + self.created.insert(topic.clone()); + return topic; + } + } + unreachable!( + "Building a topic should be completed within iterations {}", + ITERS + ) + } +} + +// TODO(niebayes): add tests for TopicBuilder. diff --git a/tests-integration/src/lib.rs b/tests-integration/src/lib.rs index 8dc4bc5a9829..b0b28ba4651c 100644 --- a/tests-integration/src/lib.rs +++ b/tests-integration/src/lib.rs @@ -20,7 +20,6 @@ mod opentsdb; mod otlp; mod prom_store; pub mod test_util; -pub mod wal_util; mod standalone; #[cfg(test)] diff --git a/tests-integration/src/wal_util/kafka.rs b/tests-integration/src/wal_util/kafka.rs deleted file mode 100644 index b4289768ed83..000000000000 --- a/tests-integration/src/wal_util/kafka.rs +++ /dev/null @@ -1,16 +0,0 @@ -// 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. - -pub mod config; -pub mod image; diff --git a/tests-integration/tests/main.rs b/tests-integration/tests/main.rs index 65dfcc8cf7c4..8b1e064579c8 100644 --- a/tests-integration/tests/main.rs +++ b/tests-integration/tests/main.rs @@ -25,5 +25,3 @@ grpc_tests!(File, S3, S3WithCache, Oss, Azblob, Gcs); http_tests!(File, S3, S3WithCache, Oss, Azblob, Gcs); // region_failover_tests!(File, S3, S3WithCache, Oss, Azblob); sql_tests!(File); - -// TODO(niebayes): add integration tests for remote wal. diff --git a/tests-integration/tests/wal.rs b/tests-integration/tests/wal.rs index 30695217087d..94c6c3b12cb3 100644 --- a/tests-integration/tests/wal.rs +++ b/tests-integration/tests/wal.rs @@ -1,248 +1,273 @@ -// 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 std::collections::HashMap; -use std::sync::Arc; - -use common_config::wal::KafkaConfig as DatanodeKafkaConfig; -use common_config::{KafkaWalOptions, WalOptions}; -use common_meta::kv_backend::memory::MemoryKvBackend; -use common_meta::kv_backend::KvBackendRef; -use common_meta::wal::kafka::{ - KafkaConfig as MetaSrvKafkaConfig, TopicManager as KafkaTopicManager, -}; -use common_meta::wal::{allocate_region_wal_options, WalConfig, WalOptionsAllocator}; -use futures::StreamExt; -use log_store::kafka::log_store::KafkaLogStore; -use log_store::kafka::{EntryImpl, NamespaceImpl}; -use rskafka::client::controller::ControllerClient; -use rskafka::client::ClientBuilder; -use store_api::logstore::entry::Id as EntryId; -use store_api::logstore::LogStore; - -// Notice: the following tests are literally unit tests. They are placed at here since -// it seems too heavy to start a Kafka cluster for each unit test. - -// The key of an env variable that stores a series of Kafka broker endpoints. -const BROKER_ENDPOINTS_KEY: &str = "GT_KAFKA_ENDPOINTS"; - -// Tests that the TopicManager allocates topics in a round-robin mannar. -#[tokio::test] -async fn test_kafka_alloc_topics() { - let broker_endpoints = std::env::var(BROKER_ENDPOINTS_KEY) - .unwrap() - .split(',') - .map(ToString::to_string) - .collect::>(); - let config = MetaSrvKafkaConfig { - topic_name_prefix: "__test_kafka_alloc_topics".to_string(), - replication_factor: broker_endpoints.len() as i16, - broker_endpoints, - ..Default::default() - }; - let kv_backend = Arc::new(MemoryKvBackend::new()) as KvBackendRef; - let manager = KafkaTopicManager::new(config.clone(), kv_backend); - manager.start().await.unwrap(); - - // Topics should be created. - let topics = (0..config.num_topics) - .map(|topic_id| format!("{}_{topic_id}", config.topic_name_prefix)) - .collect::>(); - - // Selects exactly the number of `num_topics` topics one by one. - for expected in topics.iter() { - let got = manager.select().unwrap(); - assert_eq!(got, expected); - } - - // Selects exactly the number of `num_topics` topics in a batching manner. - let got = manager - .select_batch(config.num_topics) - .unwrap() - .into_iter() - .map(ToString::to_string) - .collect::>(); - assert_eq!(got, topics); - - // Selects none. - let got = manager.select_batch(config.num_topics).unwrap(); - assert!(got.is_empty()); - - // Selects more than the number of `num_topics` topics. - let got = manager - .select_batch(2 * config.num_topics) - .unwrap() - .into_iter() - .map(ToString::to_string) - .collect::>(); - let expected = vec![topics.clone(); 2] - .into_iter() - .flatten() - .collect::>(); - assert_eq!(got, expected); -} - -// Tests that the wal options allocator could successfully allocate Kafka wal options. -#[tokio::test] -async fn test_kafka_options_allocator() { - let broker_endpoints = std::env::var(BROKER_ENDPOINTS_KEY) - .unwrap() - .split(',') - .map(ToString::to_string) - .collect::>(); - let config = MetaSrvKafkaConfig { - topic_name_prefix: "__test_kafka_options_allocator".to_string(), - replication_factor: broker_endpoints.len() as i16, - broker_endpoints, - ..Default::default() - }; - let wal_config = WalConfig::Kafka(config.clone()); - let kv_backend = Arc::new(MemoryKvBackend::new()) as KvBackendRef; - let allocator = WalOptionsAllocator::new(wal_config, kv_backend); - allocator.start().await.unwrap(); - - let num_regions = 32; - let regions = (0..num_regions).collect::>(); - let got = allocate_region_wal_options(regions.clone(), &allocator).unwrap(); - - // Topics should be allocated. - let topics = (0..num_regions) - .map(|topic_id| format!("{}_{topic_id}", config.topic_name_prefix)) - .collect::>(); - // 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::>(); - assert_eq!(got, expected); -} - -fn new_test_entry>(data: D, entry_id: EntryId, ns: NamespaceImpl) -> EntryImpl { - EntryImpl { - data: data.as_ref().to_vec(), - id: entry_id, - ns, - } -} - -async fn create_topic(topic: &str, replication_factor: i16, client: &ControllerClient) { - client - .create_topic(topic, 1, replication_factor, 5_000) - .await - .unwrap(); -} - -async fn check_entries( - ns: &NamespaceImpl, - start_offset: EntryId, - expected: Vec, - logstore: &KafkaLogStore, -) { - let mut stream = logstore.read(ns, start_offset).await.unwrap(); - for entry in expected { - let got = stream.next().await.unwrap().unwrap(); - assert_eq!(entry, got[0]); - } -} - -// Tests that the Kafka log store is able to write and read log entries from Kafka. -#[tokio::test] -async fn test_kafka_log_store() { - let broker_endpoints = std::env::var(BROKER_ENDPOINTS_KEY) - .unwrap() - .split(',') - .map(ToString::to_string) - .collect::>(); - let config = DatanodeKafkaConfig { - broker_endpoints, - ..Default::default() - }; - let logstore = KafkaLogStore::try_new(&config).await.unwrap(); - - let client = ClientBuilder::new(config.broker_endpoints.clone()) - .build() - .await - .unwrap() - .controller_client() - .unwrap(); - - // Appends one entry. - let topic = "__test_kafka_log_store_topic_append"; - create_topic(topic, config.broker_endpoints.len() as i16, &client).await; - let ns = NamespaceImpl { - region_id: 0, - topic: topic.to_string(), - }; - let entry = new_test_entry(b"0", 0, ns.clone()); - let last_entry_id = logstore.append(entry.clone()).await.unwrap().last_entry_id; - check_entries(&ns, last_entry_id, vec![entry], &logstore).await; - - // Appends a batch of entries. - // Region 1, 2 are mapped to topic 1, - let topic = "__test_kafka_log_store_topic_append_batch_1"; - create_topic(topic, config.broker_endpoints.len() as i16, &client).await; - let ns_1 = NamespaceImpl { - region_id: 1, - topic: topic.to_string(), - }; - let ns_2 = NamespaceImpl { - region_id: 2, - topic: topic.to_string(), - }; - - // Region 3 is mapped to topic 2. - let topic = "__test_kafka_log_store_topic_append_batch_2"; - create_topic(topic, config.broker_endpoints.len() as i16, &client).await; - let ns_3 = NamespaceImpl { - region_id: 3, - topic: topic.to_string(), - }; - - // Constructs a batch of entries. - let entries_1 = vec![ - new_test_entry(b"1", 0, ns_1.clone()), - new_test_entry(b"1", 1, ns_1.clone()), - ]; - let entries_2 = vec![ - new_test_entry(b"2", 2, ns_2.clone()), - new_test_entry(b"2", 3, ns_2.clone()), - ]; - let entries_3 = vec![ - new_test_entry(b"3", 7, ns_3.clone()), - new_test_entry(b"3", 8, ns_3.clone()), - ]; - let entries = vec![entries_1.clone(), entries_2.clone(), entries_3.clone()] - .into_iter() - .flatten() - .collect::>(); - - let last_entry_ids = logstore - .append_batch(entries.clone()) - .await - .unwrap() - .last_entry_ids; - - // Reads entries for region 1. - check_entries(&ns_1, last_entry_ids[&1], entries_1, &logstore).await; - // Reads entries from region 2. - check_entries(&ns_2, last_entry_ids[&2], entries_2, &logstore).await; - // Reads entries from region 3. - check_entries(&ns_3, last_entry_ids[&3], entries_3, &logstore).await; -} - -// TODO(niebayes): add more integration tests. +// // 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 std::collections::HashMap; +// use std::sync::Arc; + +// use common_config::wal::KafkaConfig as DatanodeKafkaConfig; +// use common_config::{KafkaWalOptions, WalOptions}; +// use common_meta::kv_backend::memory::MemoryKvBackend; +// use common_meta::kv_backend::KvBackendRef; +// use common_meta::wal::kafka::{ +// KafkaConfig as MetaSrvKafkaConfig, TopicManager as KafkaTopicManager, +// }; +// use common_meta::wal::{allocate_region_wal_options, WalConfig, WalOptionsAllocator}; +// use futures::StreamExt; +// use log_store::error::Result as LogStoreResult; +// use log_store::kafka::log_store::KafkaLogStore; +// use log_store::kafka::{EntryImpl, NamespaceImpl}; +// use rskafka::client::controller::ControllerClient; +// use rskafka::client::ClientBuilder; +// use store_api::logstore::entry::Id as EntryId; +// use store_api::logstore::LogStore; +// use tests_integration::wal_util::{DockerCli, KafkaImage, DEFAULT_EXPOSED_PORT}; + +// // Notice: the following tests are literally unit tests. They are placed at here since +// // it seems too heavy to start a Kafka cluster for each unit test. + +// // The key of an env variable that stores a series of Kafka broker endpoints. +// const BROKER_ENDPOINTS_KEY: &str = "GT_KAFKA_ENDPOINTS"; + +// // Tests that the TopicManager allocates topics in a round-robin mannar. +// #[tokio::test] +// async fn test_kafka_alloc_topics() { +// let broker_endpoints = std::env::var(BROKER_ENDPOINTS_KEY) +// .unwrap() +// .split(',') +// .map(ToString::to_string) +// .collect::>(); +// let config = MetaSrvKafkaConfig { +// topic_name_prefix: "__test_kafka_alloc_topics".to_string(), +// replication_factor: broker_endpoints.len() as i16, +// broker_endpoints, +// ..Default::default() +// }; +// let kv_backend = Arc::new(MemoryKvBackend::new()) as KvBackendRef; +// let manager = KafkaTopicManager::new(config.clone(), kv_backend); +// manager.start().await.unwrap(); + +// // Topics should be created. +// let topics = (0..config.num_topics) +// .map(|topic_id| format!("{}_{topic_id}", config.topic_name_prefix)) +// .collect::>(); + +// // Selects exactly the number of `num_topics` topics one by one. +// for expected in topics.iter() { +// let got = manager.select().unwrap(); +// assert_eq!(got, expected); +// } + +// // Selects exactly the number of `num_topics` topics in a batching manner. +// let got = manager +// .select_batch(config.num_topics) +// .unwrap() +// .into_iter() +// .map(ToString::to_string) +// .collect::>(); +// assert_eq!(got, topics); + +// // Selects none. +// let got = manager.select_batch(config.num_topics).unwrap(); +// assert!(got.is_empty()); + +// // Selects more than the number of `num_topics` topics. +// let got = manager +// .select_batch(2 * config.num_topics) +// .unwrap() +// .into_iter() +// .map(ToString::to_string) +// .collect::>(); +// let expected = vec![topics.clone(); 2] +// .into_iter() +// .flatten() +// .collect::>(); +// assert_eq!(got, expected); +// } + +// // Tests that the wal options allocator could successfully allocate Kafka wal options. +// #[tokio::test] +// async fn test_kafka_options_allocator() { +// let broker_endpoints = std::env::var(BROKER_ENDPOINTS_KEY) +// .unwrap() +// .split(',') +// .map(ToString::to_string) +// .collect::>(); +// let config = MetaSrvKafkaConfig { +// topic_name_prefix: "__test_kafka_options_allocator".to_string(), +// replication_factor: broker_endpoints.len() as i16, +// broker_endpoints, +// ..Default::default() +// }; +// let wal_config = WalConfig::Kafka(config.clone()); +// let kv_backend = Arc::new(MemoryKvBackend::new()) as KvBackendRef; +// let allocator = WalOptionsAllocator::new(wal_config, kv_backend); +// allocator.start().await.unwrap(); + +// let num_regions = 32; +// let regions = (0..num_regions).collect::>(); +// let got = allocate_region_wal_options(regions.clone(), &allocator).unwrap(); + +// // Topics should be allocated. +// let topics = (0..num_regions) +// .map(|topic_id| format!("{}_{topic_id}", config.topic_name_prefix)) +// .collect::>(); +// // 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::>(); +// assert_eq!(got, expected); +// } + +// async fn create_topic(topic: &str, replication_factor: i16, client: &ControllerClient) { +// client +// .create_topic(topic, 1, replication_factor, 500) +// .await +// .unwrap(); +// } + +// async fn check_entries( +// ns: &NamespaceImpl, +// start_offset: EntryId, +// expected: Vec, +// logstore: &KafkaLogStore, +// ) { +// let stream = logstore.read(ns, start_offset).await.unwrap(); +// let got = stream +// .collect::>() +// .await +// .into_iter() +// .flat_map(|x| x.unwrap()) +// .collect::>(); +// assert_eq!(expected, got); +// // for entry in expected { +// // let got = stream.next().await.unwrap().unwrap(); +// // } +// } + +// // Tests that the Kafka log store is able to write and read log entries from Kafka. +// // #[tokio::test] +// // async fn test_kafka_log_store() { +// // println!("Start running test"); + +// // // Starts a Kafka container. +// // let docker = DockerCli::default(); +// // let container = docker.run(KafkaImage::default()); + +// // println!("Started the container"); + +// // let broker_endpoints = std::env::var(BROKER_ENDPOINTS_KEY) +// // .unwrap_or(format!( +// // "localhost:{}", +// // container.get_host_port_ipv4(DEFAULT_EXPOSED_PORT) +// // )) +// // .split(',') +// // .map(ToString::to_string) +// // .collect::>(); +// // let config = DatanodeKafkaConfig { +// // broker_endpoints, +// // ..Default::default() +// // }; +// // let logstore = KafkaLogStore::try_new(&config).await.unwrap(); + +// // println!("Started the log store"); + +// // let client = ClientBuilder::new(config.broker_endpoints.clone()) +// // .build() +// // .await +// // .unwrap() +// // .controller_client() +// // .unwrap(); + +// // println!("Created a client"); + +// // // Appends one entry. +// // let topic = "__test_kafka_log_store_topic_append"; +// // create_topic(topic, config.broker_endpoints.len() as i16, &client).await; + +// // println!("Created a topic"); + +// // let ns = NamespaceImpl { +// // region_id: 0, +// // topic: topic.to_string(), +// // }; +// // let entry = new_test_entry(b"0", 0, ns.clone()); +// // let last_entry_id = logstore.append(entry.clone()).await.unwrap().last_entry_id; + +// // println!("Appended an entry"); + +// // check_entries(&ns, last_entry_id, vec![entry], &logstore).await; + +// // // Appends a batch of entries. +// // // Region 1, 2 are mapped to topic 1, +// // let topic = "__test_kafka_log_store_topic_append_batch_1"; +// // create_topic(topic, config.broker_endpoints.len() as i16, &client).await; + +// // println!("Created a topic"); + +// // let ns_1 = NamespaceImpl { +// // region_id: 1, +// // topic: topic.to_string(), +// // }; +// // let ns_2 = NamespaceImpl { +// // region_id: 2, +// // topic: topic.to_string(), +// // }; + +// // // Region 3 is mapped to topic 2. +// // let topic = "__test_kafka_log_store_topic_append_batch_2"; +// // create_topic(topic, config.broker_endpoints.len() as i16, &client).await; + +// // println!("Created a topic"); + +// // let ns_3 = NamespaceImpl { +// // region_id: 3, +// // topic: topic.to_string(), +// // }; + +// // // Constructs a batch of entries. +// // let entries_1 = vec![ +// // new_test_entry(b"1", 0, ns_1.clone()), +// // new_test_entry(b"1", 1, ns_1.clone()), +// // ]; +// // let entries_2 = vec![ +// // new_test_entry(b"2", 2, ns_2.clone()), +// // new_test_entry(b"2", 3, ns_2.clone()), +// // ]; +// // let entries_3 = vec![ +// // new_test_entry(b"3", 7, ns_3.clone()), +// // new_test_entry(b"3", 8, ns_3.clone()), +// // ]; +// // let entries = vec![entries_1.clone(), entries_2.clone(), entries_3.clone()] +// // .into_iter() +// // .flatten() +// // .collect::>(); + +// // let last_entry_ids = logstore +// // .append_batch(entries.clone()) +// // .await +// // .unwrap() +// // .last_entry_ids; + +// // // Reads entries for region 1. +// // check_entries(&ns_1, last_entry_ids[&1], entries_1, &logstore).await; +// // // Reads entries from region 2. +// // check_entries(&ns_2, last_entry_ids[&2], entries_2, &logstore).await; +// // // Reads entries from region 3. +// // check_entries(&ns_3, last_entry_ids[&3], entries_3, &logstore).await; +// // } From f7954dcc5e0b47a348a65230506b0c40540f0cd1 Mon Sep 17 00:00:00 2001 From: niebayes Date: Sat, 30 Dec 2023 18:42:37 +0800 Subject: [PATCH 08/30] test: add unit tests for client manager --- src/log-store/src/kafka/client_manager.rs | 111 ++++++++---------- src/log-store/src/kafka/log_store.rs | 1 - src/log-store/src/test_util/kafka.rs | 42 ++++++- .../src/test_util/kafka/topic_builder.rs | 6 +- 4 files changed, 93 insertions(+), 67 deletions(-) diff --git a/src/log-store/src/kafka/client_manager.rs b/src/log-store/src/kafka/client_manager.rs index 1e49f977de4d..72d0b839a755 100644 --- a/src/log-store/src/kafka/client_manager.rs +++ b/src/log-store/src/kafka/client_manager.rs @@ -12,19 +12,17 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::hash_map::Entry; use std::collections::HashMap; use std::sync::Arc; use common_config::wal::{KafkaConfig, KafkaWalTopic as Topic}; -use common_telemetry::debug; use rskafka::client::partition::{PartitionClient, UnknownTopicHandling}; use rskafka::client::producer::aggregator::RecordAggregator; use rskafka::client::producer::{BatchProducer, BatchProducerBuilder}; use rskafka::client::{Client as RsKafkaClient, ClientBuilder}; use rskafka::BackoffConfig; use snafu::ResultExt; -use tokio::sync::{Mutex as TokioMutex, RwLock as TokioRwLock}; +use tokio::sync::RwLock; use crate::error::{BuildClientSnafu, BuildPartitionClientSnafu, Result}; @@ -69,7 +67,7 @@ pub(crate) struct ClientManager { client_factory: RsKafkaClient, /// A pool maintaining a collection of clients. /// Key: a topic. Value: the associated client of the topic. - client_pool: TokioMutex>, + client_pool: RwLock>, } impl ClientManager { @@ -90,28 +88,36 @@ impl ClientManager { broker_endpoints: config.broker_endpoints.clone(), })?; - debug!("Created a ClientManager"); - Ok(Self { config: config.clone(), client_factory: client, - client_pool: TokioMutex::new(HashMap::new()), + client_pool: RwLock::new(HashMap::new()), }) } /// Gets the client associated with the topic. If the client does not exist, a new one will /// be created and returned. pub(crate) async fn get_or_insert(&self, topic: &Topic) -> Result { - let mut client_pool = self.client_pool.lock().await; - if let Entry::Vacant(entry) = client_pool.entry(topic.to_string()) { - entry.insert(self.try_create_client(topic).await?); + let client_pool = self.client_pool.read().await; + if let Some(client) = client_pool.get(topic) { + return Ok(client.clone()); + } + // Manullay releases the read lock. + drop(client_pool); + + // Acquires the write lock. + let mut client_pool = self.client_pool.write().await; + match client_pool.get(topic) { + Some(client) => Ok(client.clone()), + None => { + let client = self.try_create_client(topic).await?; + client_pool.insert(topic.clone(), client.clone()); + Ok(client) + } } - Ok(client_pool[topic].clone()) } async fn try_create_client(&self, topic: &Topic) -> Result { - debug!("Try to create client for topic {}", topic); - // Sets to Retry to retry connecting if the kafka cluter replies with an UnknownTopic error. // That's because the topic is believed to exist as the metasrv is expected to create required topics upon start. // The reconnecting won't stop until succeed or a different error returns. @@ -125,8 +131,6 @@ impl ClientManager { }) .map(Arc::new)?; - debug!("Created a client for topic {}", topic); - Ok(Client::new(raw_client, &self.config)) } } @@ -136,37 +140,24 @@ mod tests { use common_test_util::wal::kafka::BROKER_ENDPOINTS_KEY; use super::*; + use crate::get_broker_endpoints_from_env; + use crate::test_util::kafka::create_topics; use crate::test_util::kafka::topic_builder::{Affix, TopicBuilder}; /// Checks clients for the given topics are created. async fn ensure_topics_exist(topics: &[Topic], client_manager: &ClientManager) { - let client_pool = client_manager.client_pool.lock().await; + let client_pool = client_manager.client_pool.read().await; let all_exist = topics.iter().all(|topic| client_pool.contains_key(topic)); assert_eq!(all_exist, true); } - /// Sends `get_or_insert` requests sequentially to the client manager, and checks if it could handle them correctly. - #[tokio::test] - async fn test_sequential() { - let broker_endpoints = std::env::var(BROKER_ENDPOINTS_KEY) - .unwrap() - .split(',') - .map(ToString::to_string) - .collect::>(); - if broker_endpoints.is_empty() { - return; - } - - let client = ClientBuilder::new(broker_endpoints.clone()) - .build() - .await - .unwrap(); - let ctrl_client = client.controller_client().unwrap(); - + async fn test_which(test_name: &str) { + // Creates a collection of topics in Kafka. + let broker_endpoints = get_broker_endpoints_from_env!(BROKER_ENDPOINTS_KEY); let topic_builder = TopicBuilder::default() - .with_prefix(Affix::Fixed("test_sequential".to_string())) + .with_prefix(Affix::Fixed(test_name.to_string())) .with_suffix(Affix::TimeNow); - + let topics = create_topics(256, topic_builder, &broker_endpoints).await; let config = KafkaConfig { broker_endpoints, @@ -174,36 +165,36 @@ mod tests { }; let manager = ClientManager::try_new(&config).await.unwrap(); - // Constructs a collection of mock topics. - let num_topics = 256; - let topics = (0..num_topics) - .map(|i| format!("topic_{i}")) - .collect::>(); - - // Gets all clients sequentially. - for topic in topics.iter() { - manager.get_or_insert(topic).await.unwrap(); + match test_name { + "test_sequential" => { + // Gets all clients sequentially. + for topic in topics.iter() { + manager.get_or_insert(topic).await.unwrap(); + } + } + "test_parallel" => { + // Gets all clients in parallel. + let tasks = topics + .iter() + .map(|topic| manager.get_or_insert(topic)) + .collect::>(); + futures::future::try_join_all(tasks).await.unwrap(); + } + _ => unreachable!(), } + ensure_topics_exist(&topics, &manager).await; } + /// Sends `get_or_insert` requests sequentially to the client manager, and checks if it could handle them correctly. + #[tokio::test] + async fn test_sequential() { + test_which("test_sequential").await; + } + /// Sends `get_or_insert` requests in parallel to the client manager, and checks if it could handle them correctly. #[tokio::test] async fn test_parallel() { - let manager = ClientManager::try_new(&config).await.unwrap(); - - // Constructs a collection of mock topics. - let num_topics = 256; - let topics = (0..num_topics) - .map(|i| format!("topic_{i}")) - .collect::>(); - - // Gets all clients in parallel. - let tasks = topics - .iter() - .map(|topic| manager.get_or_insert(topic)) - .collect::>(); - futures::future::try_join_all(tasks).await.unwrap(); - ensure_topics_exist(&topics, &manager).await; + test_which("test_parallel").await; } } diff --git a/src/log-store/src/kafka/log_store.rs b/src/log-store/src/kafka/log_store.rs index f0170245e5e2..d78f846f9277 100644 --- a/src/log-store/src/kafka/log_store.rs +++ b/src/log-store/src/kafka/log_store.rs @@ -16,7 +16,6 @@ use std::collections::HashMap; use std::sync::Arc; use common_config::wal::{KafkaConfig, WalOptions}; -use common_telemetry::{debug, warn}; use futures_util::StreamExt; use rskafka::client::consumer::{StartOffset, StreamConsumerBuilder}; use rskafka::client::partition::OffsetAt; diff --git a/src/log-store/src/test_util/kafka.rs b/src/log-store/src/test_util/kafka.rs index b66e692836d7..f2af612e1f58 100644 --- a/src/log-store/src/test_util/kafka.rs +++ b/src/log-store/src/test_util/kafka.rs @@ -16,12 +16,48 @@ pub mod entry_builder; pub mod topic_builder; use common_config::wal::KafkaWalTopic as Topic; +use rskafka::client::ClientBuilder; pub use crate::test_util::kafka::entry_builder::EntryBuilder; pub use crate::test_util::kafka::topic_builder::TopicBuilder; +/// Gets broker endpoints from environment variables with the given key. +#[macro_export] +macro_rules! get_broker_endpoints_from_env { + ($key:expr) => {{ + let broker_endpoints = std::env::var($key) + .unwrap() + .split(',') + .map(ToString::to_string) + .collect::>(); + assert!(!broker_endpoints.is_empty()); + broker_endpoints + }}; +} + /// Creates `num_topiocs` number of topics with the given TopicBuilder. -/// Requests for creating these topics on the Kafka cluster if the `broker_endpoints` is not empty. -pub fn create_topics(num_topics: usize, builder: TopicBuilder, broker_endpoints: Vec) -> Vec { - +/// Requests for creating these topics on the Kafka cluster. +pub async fn create_topics( + num_topics: usize, + mut builder: TopicBuilder, + broker_endpoints: &[String], +) -> Vec { + assert!(!broker_endpoints.is_empty()); + + let client = ClientBuilder::new(broker_endpoints.to_vec()) + .build() + .await + .unwrap(); + let ctrl_client = client.controller_client().unwrap(); + + let (topics, tasks): (Vec<_>, Vec<_>) = (0..num_topics) + .map(|i| { + let topic = builder.build(&format!("topic_{i}")); + let task = ctrl_client.create_topic(topic.clone(), 1, 1, 500); + (topic, task) + }) + .unzip(); + futures::future::try_join_all(tasks).await.unwrap(); + + topics } diff --git a/src/log-store/src/test_util/kafka/topic_builder.rs b/src/log-store/src/test_util/kafka/topic_builder.rs index 485b846cc6d4..be51d8c10f49 100644 --- a/src/log-store/src/test_util/kafka/topic_builder.rs +++ b/src/log-store/src/test_util/kafka/topic_builder.rs @@ -32,7 +32,7 @@ impl ToString for Affix { fn to_string(&self) -> String { match self { Affix::Fixed(s) => s.to_string(), - Affix::TimeNow => chrono::Local::now().to_string(), + Affix::TimeNow => chrono::Local::now().timestamp_millis().to_string(), Affix::Nothing => String::default(), } } @@ -70,13 +70,13 @@ impl TopicBuilder { } /// Builds a topic by inserting a prefix and a suffix into the given topic. - pub fn build(&mut self, topic: &Topic) -> Topic { + pub fn build(&mut self, body: &Topic) -> Topic { const ITERS: usize = 24; for _ in 0..ITERS { let topic = format!( "{}_{}_{}", self.prefix.to_string(), - topic, + body, self.suffix.to_string() ); if !self.created.contains(&topic) { From 0b503a5851efaff5cc09f0f7b94b0f3a4d0ab4b4 Mon Sep 17 00:00:00 2001 From: niebayes Date: Sat, 30 Dec 2023 23:27:24 +0800 Subject: [PATCH 09/30] test: add some unit tests for kafka log store --- src/log-store/src/kafka/client_manager.rs | 2 +- src/log-store/src/kafka/log_store.rs | 132 +++++++++++++++++- src/log-store/src/test_util/kafka.rs | 4 +- .../src/test_util/kafka/entry_builder.rs | 16 ++- .../src/test_util/kafka/topic_builder.rs | 2 +- tests-integration/tests/wal.rs | 131 ----------------- 6 files changed, 149 insertions(+), 138 deletions(-) diff --git a/src/log-store/src/kafka/client_manager.rs b/src/log-store/src/kafka/client_manager.rs index 72d0b839a755..ef4361b3d700 100644 --- a/src/log-store/src/kafka/client_manager.rs +++ b/src/log-store/src/kafka/client_manager.rs @@ -148,7 +148,7 @@ mod tests { async fn ensure_topics_exist(topics: &[Topic], client_manager: &ClientManager) { let client_pool = client_manager.client_pool.read().await; let all_exist = topics.iter().all(|topic| client_pool.contains_key(topic)); - assert_eq!(all_exist, true); + assert!(all_exist); } async fn test_which(test_name: &str) { diff --git a/src/log-store/src/kafka/log_store.rs b/src/log-store/src/kafka/log_store.rs index d78f846f9277..75f34b95caf0 100644 --- a/src/log-store/src/kafka/log_store.rs +++ b/src/log-store/src/kafka/log_store.rs @@ -15,7 +15,7 @@ use std::collections::HashMap; use std::sync::Arc; -use common_config::wal::{KafkaConfig, WalOptions}; +use common_config::wal::{KafkaConfig, KafkaWalTopic as Topic, WalOptions}; use futures_util::StreamExt; use rskafka::client::consumer::{StartOffset, StreamConsumerBuilder}; use rskafka::client::partition::OffsetAt; @@ -47,6 +47,26 @@ impl KafkaLogStore { config: config.clone(), }) } + + /// Gets the end offset of the last record in a Kafka topic. + /// Warning: this method is intended to be used only in testing. + // TODO(niebayes): use this to test that the initial offset is 1 for a Kafka log store in that + // a no-op record is successfully appended into each topic. + #[allow(unused)] + pub async fn get_offset(&self, topic: &Topic) -> EntryId { + let client = self + .client_manager + .get_or_insert(topic) + .await + .unwrap() + .raw_client; + client + .get_offset(OffsetAt::Latest) + .await + .map(TryInto::try_into) + .unwrap() + .unwrap() + } } #[async_trait::async_trait] @@ -267,3 +287,113 @@ impl LogStore for KafkaLogStore { Ok(()) } } + +#[cfg(test)] +mod tests { + use common_test_util::wal::kafka::BROKER_ENDPOINTS_KEY; + + use super::*; + use crate::get_broker_endpoints_from_env; + use crate::test_util::kafka::topic_builder::Affix; + use crate::test_util::kafka::{create_topics, EntryBuilder, TopicBuilder}; + + fn new_namespace(topic: &str, region_id: u64) -> NamespaceImpl { + NamespaceImpl { + region_id, + topic: topic.to_string(), + } + } + + // TODO(niebayes): change `expected` to &[EntryImpl]. + async fn check_entries( + ns: &NamespaceImpl, + start_offset: EntryId, + expected: Vec, + logstore: &KafkaLogStore, + ) { + let stream = logstore.read(ns, start_offset).await.unwrap(); + let got = stream + .collect::>() + .await + .into_iter() + .flat_map(|x| x.unwrap()) + .collect::>(); + assert_eq!(expected, got); + } + + /// Appends entries for one region and checks all entries can be read successfully. + #[tokio::test] + async fn test_one_region() { + let broker_endpoints = get_broker_endpoints_from_env!(BROKER_ENDPOINTS_KEY); + let topic_builder = TopicBuilder::default() + .with_prefix(Affix::Fixed("test_one_region".to_string())) + .with_suffix(Affix::TimeNow); + let topic = create_topics(1, topic_builder, &broker_endpoints).await[0].clone(); + + let config = KafkaConfig { + broker_endpoints, + ..Default::default() + }; + let logstore = KafkaLogStore::try_new(&config).await.unwrap(); + + let ns = new_namespace(&topic, 0); + let entry_builder = EntryBuilder::new(ns.clone()); + let entry = entry_builder.with_random_data(); + + let last_entry_id = logstore.append(entry.clone()).await.unwrap().last_entry_id; + check_entries(&ns, last_entry_id, vec![entry], &logstore).await; + + let entries = (0..10) + .map(|_| entry_builder.with_random_data()) + .collect::>(); + let last_entry_id = logstore + .append_batch(entries.clone()) + .await + .unwrap() + .last_entry_ids[&ns.region_id]; + check_entries(&ns, last_entry_id, entries, &logstore).await; + } + + /// Appends entries for multiple regions and checks entries for each region can be read successfully. + /// A topic is assigned only a single region. + #[tokio::test] + async fn test_multi_regions_disjoint() { + let broker_endpoints = get_broker_endpoints_from_env!(BROKER_ENDPOINTS_KEY); + let topic_builder = TopicBuilder::default() + .with_prefix(Affix::Fixed("test_multi_regions_disjoint".to_string())) + .with_suffix(Affix::TimeNow); + let topics = create_topics(10, topic_builder, &broker_endpoints).await; + + let config = KafkaConfig { + broker_endpoints, + ..Default::default() + }; + let logstore = KafkaLogStore::try_new(&config).await.unwrap(); + + let (region_namespaces, mut entry_builders): (Vec<_>, Vec<_>) = topics + .iter() + .enumerate() + .map(|(i, topic)| { + let ns = new_namespace(topic, i as u64); + let entry_builder = EntryBuilder::new(ns.clone()); + (ns, entry_builder) + }) + .unzip(); + let region_entries = entry_builders + .iter_mut() + .map(|builder| builder.with_random_data_batch(5)) + .collect::>(); + let entries = region_entries.iter().flatten().cloned().collect::>(); + let last_entry_ids = logstore.append_batch(entries).await.unwrap().last_entry_ids; + + for (i, ns) in region_namespaces.iter().enumerate() { + let expected = region_entries[i].clone(); + check_entries(ns, last_entry_ids[&ns.region_id], expected, &logstore).await; + } + } + + /// Appends entries for multiple regions and checks entries for each region can be read successfully. + /// A topic may be assigned multiple regions. + #[tokio::test] + async fn test_multi_regions_overlapped() {} +} diff --git a/src/log-store/src/test_util/kafka.rs b/src/log-store/src/test_util/kafka.rs index f2af612e1f58..d80eb838d1d0 100644 --- a/src/log-store/src/test_util/kafka.rs +++ b/src/log-store/src/test_util/kafka.rs @@ -26,7 +26,7 @@ pub use crate::test_util::kafka::topic_builder::TopicBuilder; macro_rules! get_broker_endpoints_from_env { ($key:expr) => {{ let broker_endpoints = std::env::var($key) - .unwrap() + .unwrap_or("localhost:9092".to_string()) .split(',') .map(ToString::to_string) .collect::>(); @@ -36,7 +36,7 @@ macro_rules! get_broker_endpoints_from_env { } /// Creates `num_topiocs` number of topics with the given TopicBuilder. -/// Requests for creating these topics on the Kafka cluster. +/// Requests for creating these topics on the Kafka cluster specified with the `broker_endpoints`. pub async fn create_topics( num_topics: usize, mut builder: TopicBuilder, diff --git a/src/log-store/src/test_util/kafka/entry_builder.rs b/src/log-store/src/test_util/kafka/entry_builder.rs index 3e234c78d6e4..9bfb41027abe 100644 --- a/src/log-store/src/test_util/kafka/entry_builder.rs +++ b/src/log-store/src/test_util/kafka/entry_builder.rs @@ -27,7 +27,7 @@ const DEFAULT_DATA: &[u8; 10] = b"[greptime]"; /// A builder for building entries for a namespace. pub struct EntryBuilder { /// The namespace of the entries. - ns: NamespaceImpl, + pub ns: NamespaceImpl, /// The next entry id to allocate. It starts from 0 by default. next_entry_id: AtomicEntryId, /// A generator for supporting random data generation. @@ -78,6 +78,7 @@ impl EntryBuilder { } /// Builds an entry with the default data. + // TODO(niebayes): may be remove this method since it's not used. pub fn with_default_data(&self) -> EntryImpl { EntryImpl { data: DEFAULT_DATA.to_vec(), @@ -86,6 +87,17 @@ impl EntryBuilder { } } + /// Builds a batch of entries each with random data. + pub fn with_random_data_batch(&self, batch_size: usize) -> Vec { + (0..batch_size) + .map(|_| EntryImpl { + data: self.make_random_data(), + id: self.alloc_entry_id(), + ns: self.ns.clone(), + }) + .collect() + } + /// Builds an entry with random data. pub fn with_random_data(&self) -> EntryImpl { EntryImpl { @@ -105,7 +117,7 @@ impl EntryBuilder { let amount = rng.gen_range(0..self.data_pool.len()); self.data_pool .choose_multiple(&mut rng, amount) - .map(|x| *x) + .copied() .collect() } } diff --git a/src/log-store/src/test_util/kafka/topic_builder.rs b/src/log-store/src/test_util/kafka/topic_builder.rs index be51d8c10f49..037390914c8b 100644 --- a/src/log-store/src/test_util/kafka/topic_builder.rs +++ b/src/log-store/src/test_util/kafka/topic_builder.rs @@ -70,7 +70,7 @@ impl TopicBuilder { } /// Builds a topic by inserting a prefix and a suffix into the given topic. - pub fn build(&mut self, body: &Topic) -> Topic { + pub fn build(&mut self, body: &str) -> Topic { const ITERS: usize = 24; for _ in 0..ITERS { let topic = format!( diff --git a/tests-integration/tests/wal.rs b/tests-integration/tests/wal.rs index 94c6c3b12cb3..8bf09f55c550 100644 --- a/tests-integration/tests/wal.rs +++ b/tests-integration/tests/wal.rs @@ -140,134 +140,3 @@ // .await // .unwrap(); // } - -// async fn check_entries( -// ns: &NamespaceImpl, -// start_offset: EntryId, -// expected: Vec, -// logstore: &KafkaLogStore, -// ) { -// let stream = logstore.read(ns, start_offset).await.unwrap(); -// let got = stream -// .collect::>() -// .await -// .into_iter() -// .flat_map(|x| x.unwrap()) -// .collect::>(); -// assert_eq!(expected, got); -// // for entry in expected { -// // let got = stream.next().await.unwrap().unwrap(); -// // } -// } - -// // Tests that the Kafka log store is able to write and read log entries from Kafka. -// // #[tokio::test] -// // async fn test_kafka_log_store() { -// // println!("Start running test"); - -// // // Starts a Kafka container. -// // let docker = DockerCli::default(); -// // let container = docker.run(KafkaImage::default()); - -// // println!("Started the container"); - -// // let broker_endpoints = std::env::var(BROKER_ENDPOINTS_KEY) -// // .unwrap_or(format!( -// // "localhost:{}", -// // container.get_host_port_ipv4(DEFAULT_EXPOSED_PORT) -// // )) -// // .split(',') -// // .map(ToString::to_string) -// // .collect::>(); -// // let config = DatanodeKafkaConfig { -// // broker_endpoints, -// // ..Default::default() -// // }; -// // let logstore = KafkaLogStore::try_new(&config).await.unwrap(); - -// // println!("Started the log store"); - -// // let client = ClientBuilder::new(config.broker_endpoints.clone()) -// // .build() -// // .await -// // .unwrap() -// // .controller_client() -// // .unwrap(); - -// // println!("Created a client"); - -// // // Appends one entry. -// // let topic = "__test_kafka_log_store_topic_append"; -// // create_topic(topic, config.broker_endpoints.len() as i16, &client).await; - -// // println!("Created a topic"); - -// // let ns = NamespaceImpl { -// // region_id: 0, -// // topic: topic.to_string(), -// // }; -// // let entry = new_test_entry(b"0", 0, ns.clone()); -// // let last_entry_id = logstore.append(entry.clone()).await.unwrap().last_entry_id; - -// // println!("Appended an entry"); - -// // check_entries(&ns, last_entry_id, vec![entry], &logstore).await; - -// // // Appends a batch of entries. -// // // Region 1, 2 are mapped to topic 1, -// // let topic = "__test_kafka_log_store_topic_append_batch_1"; -// // create_topic(topic, config.broker_endpoints.len() as i16, &client).await; - -// // println!("Created a topic"); - -// // let ns_1 = NamespaceImpl { -// // region_id: 1, -// // topic: topic.to_string(), -// // }; -// // let ns_2 = NamespaceImpl { -// // region_id: 2, -// // topic: topic.to_string(), -// // }; - -// // // Region 3 is mapped to topic 2. -// // let topic = "__test_kafka_log_store_topic_append_batch_2"; -// // create_topic(topic, config.broker_endpoints.len() as i16, &client).await; - -// // println!("Created a topic"); - -// // let ns_3 = NamespaceImpl { -// // region_id: 3, -// // topic: topic.to_string(), -// // }; - -// // // Constructs a batch of entries. -// // let entries_1 = vec![ -// // new_test_entry(b"1", 0, ns_1.clone()), -// // new_test_entry(b"1", 1, ns_1.clone()), -// // ]; -// // let entries_2 = vec![ -// // new_test_entry(b"2", 2, ns_2.clone()), -// // new_test_entry(b"2", 3, ns_2.clone()), -// // ]; -// // let entries_3 = vec![ -// // new_test_entry(b"3", 7, ns_3.clone()), -// // new_test_entry(b"3", 8, ns_3.clone()), -// // ]; -// // let entries = vec![entries_1.clone(), entries_2.clone(), entries_3.clone()] -// // .into_iter() -// // .flatten() -// // .collect::>(); - -// // let last_entry_ids = logstore -// // .append_batch(entries.clone()) -// // .await -// // .unwrap() -// // .last_entry_ids; - -// // // Reads entries for region 1. -// // check_entries(&ns_1, last_entry_ids[&1], entries_1, &logstore).await; -// // // Reads entries from region 2. -// // check_entries(&ns_2, last_entry_ids[&2], entries_2, &logstore).await; -// // // Reads entries from region 3. -// // check_entries(&ns_3, last_entry_ids[&3], entries_3, &logstore).await; -// // } From 3d489ff3e285064438811e621b86878124a3eaf9 Mon Sep 17 00:00:00 2001 From: niebayes Date: Sun, 31 Dec 2023 08:35:14 +0800 Subject: [PATCH 10/30] chore: resolve some todos --- src/common/meta/src/wal/kafka/topic.rs | 1 + src/log-store/src/kafka/client_manager.rs | 7 +++-- src/log-store/src/kafka/log_store.rs | 19 ++++++------- src/log-store/src/test_util/kafka.rs | 21 ++++++++++----- .../src/test_util/kafka/entry_builder.rs | 27 +------------------ .../{topic_builder.rs => topic_decorator.rs} | 22 +++++++-------- 6 files changed, 40 insertions(+), 57 deletions(-) rename src/log-store/src/test_util/kafka/{topic_builder.rs => topic_decorator.rs} (85%) diff --git a/src/common/meta/src/wal/kafka/topic.rs b/src/common/meta/src/wal/kafka/topic.rs index 67223637f1ae..34e15ad22450 100644 --- a/src/common/meta/src/wal/kafka/topic.rs +++ b/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; diff --git a/src/log-store/src/kafka/client_manager.rs b/src/log-store/src/kafka/client_manager.rs index ef4361b3d700..c6201b864d93 100644 --- a/src/log-store/src/kafka/client_manager.rs +++ b/src/log-store/src/kafka/client_manager.rs @@ -141,8 +141,7 @@ mod tests { use super::*; use crate::get_broker_endpoints_from_env; - use crate::test_util::kafka::create_topics; - use crate::test_util::kafka::topic_builder::{Affix, TopicBuilder}; + use crate::test_util::kafka::{create_topics, Affix, TopicDecorator}; /// Checks clients for the given topics are created. async fn ensure_topics_exist(topics: &[Topic], client_manager: &ClientManager) { @@ -154,10 +153,10 @@ mod tests { async fn test_which(test_name: &str) { // Creates a collection of topics in Kafka. let broker_endpoints = get_broker_endpoints_from_env!(BROKER_ENDPOINTS_KEY); - let topic_builder = TopicBuilder::default() + let decorator = TopicDecorator::default() .with_prefix(Affix::Fixed(test_name.to_string())) .with_suffix(Affix::TimeNow); - let topics = create_topics(256, topic_builder, &broker_endpoints).await; + let topics = create_topics(256, decorator, &broker_endpoints, None).await; let config = KafkaConfig { broker_endpoints, diff --git a/src/log-store/src/kafka/log_store.rs b/src/log-store/src/kafka/log_store.rs index 75f34b95caf0..7d3c8bde36cd 100644 --- a/src/log-store/src/kafka/log_store.rs +++ b/src/log-store/src/kafka/log_store.rs @@ -294,8 +294,9 @@ mod tests { use super::*; use crate::get_broker_endpoints_from_env; - use crate::test_util::kafka::topic_builder::Affix; - use crate::test_util::kafka::{create_topics, EntryBuilder, TopicBuilder}; + use crate::test_util::kafka::{ + create_topics, entries_with_random_data, Affix, EntryBuilder, TopicDecorator, + }; fn new_namespace(topic: &str, region_id: u64) -> NamespaceImpl { NamespaceImpl { @@ -325,10 +326,10 @@ mod tests { #[tokio::test] async fn test_one_region() { let broker_endpoints = get_broker_endpoints_from_env!(BROKER_ENDPOINTS_KEY); - let topic_builder = TopicBuilder::default() + let decorator = TopicDecorator::default() .with_prefix(Affix::Fixed("test_one_region".to_string())) .with_suffix(Affix::TimeNow); - let topic = create_topics(1, topic_builder, &broker_endpoints).await[0].clone(); + let topic = create_topics(1, decorator, &broker_endpoints, None).await[0].clone(); let config = KafkaConfig { broker_endpoints, @@ -359,10 +360,10 @@ mod tests { #[tokio::test] async fn test_multi_regions_disjoint() { let broker_endpoints = get_broker_endpoints_from_env!(BROKER_ENDPOINTS_KEY); - let topic_builder = TopicBuilder::default() + let decorator = TopicDecorator::default() .with_prefix(Affix::Fixed("test_multi_regions_disjoint".to_string())) .with_suffix(Affix::TimeNow); - let topics = create_topics(10, topic_builder, &broker_endpoints).await; + let topics = create_topics(10, decorator, &broker_endpoints, None).await; let config = KafkaConfig { broker_endpoints, @@ -370,7 +371,7 @@ mod tests { }; let logstore = KafkaLogStore::try_new(&config).await.unwrap(); - let (region_namespaces, mut entry_builders): (Vec<_>, Vec<_>) = topics + let (region_namespaces, entry_builders): (Vec<_>, Vec<_>) = topics .iter() .enumerate() .map(|(i, topic)| { @@ -380,8 +381,8 @@ mod tests { }) .unzip(); let region_entries = entry_builders - .iter_mut() - .map(|builder| builder.with_random_data_batch(5)) + .iter() + .map(|builder| entries_with_random_data(25, builder)) .collect::>(); let entries = region_entries.iter().flatten().cloned().collect::>(); let last_entry_ids = logstore.append_batch(entries).await.unwrap().last_entry_ids; diff --git a/src/log-store/src/test_util/kafka.rs b/src/log-store/src/test_util/kafka.rs index d80eb838d1d0..aa431551c91e 100644 --- a/src/log-store/src/test_util/kafka.rs +++ b/src/log-store/src/test_util/kafka.rs @@ -13,13 +13,14 @@ // limitations under the License. pub mod entry_builder; -pub mod topic_builder; +pub mod topic_decorator; use common_config::wal::KafkaWalTopic as Topic; use rskafka::client::ClientBuilder; +use crate::kafka::EntryImpl; pub use crate::test_util::kafka::entry_builder::EntryBuilder; -pub use crate::test_util::kafka::topic_builder::TopicBuilder; +pub use crate::test_util::kafka::topic_decorator::{Affix, TopicDecorator}; /// Gets broker endpoints from environment variables with the given key. #[macro_export] @@ -35,12 +36,12 @@ macro_rules! get_broker_endpoints_from_env { }}; } -/// Creates `num_topiocs` number of topics with the given TopicBuilder. -/// Requests for creating these topics on the Kafka cluster specified with the `broker_endpoints`. +/// Creates `num_topiocs` number of topics from the seed topic which are going to be decorated with the given TopicDecorator. pub async fn create_topics( num_topics: usize, - mut builder: TopicBuilder, + mut decorator: TopicDecorator, broker_endpoints: &[String], + seed: Option<&str>, ) -> Vec { assert!(!broker_endpoints.is_empty()); @@ -50,9 +51,10 @@ pub async fn create_topics( .unwrap(); let ctrl_client = client.controller_client().unwrap(); + let seed = seed.unwrap_or("topic"); let (topics, tasks): (Vec<_>, Vec<_>) = (0..num_topics) .map(|i| { - let topic = builder.build(&format!("topic_{i}")); + let topic = decorator.decorate(&format!("{seed}_{i}")); let task = ctrl_client.create_topic(topic.clone(), 1, 1, 500); (topic, task) }) @@ -61,3 +63,10 @@ pub async fn create_topics( topics } + +/// Builds a batch of entries each with random data. +pub fn entries_with_random_data(batch_size: usize, builder: &EntryBuilder) -> Vec { + (0..batch_size) + .map(|_| builder.with_random_data()) + .collect() +} diff --git a/src/log-store/src/test_util/kafka/entry_builder.rs b/src/log-store/src/test_util/kafka/entry_builder.rs index 9bfb41027abe..2adf85c13345 100644 --- a/src/log-store/src/test_util/kafka/entry_builder.rs +++ b/src/log-store/src/test_util/kafka/entry_builder.rs @@ -22,12 +22,10 @@ use store_api::logstore::EntryId; use crate::kafka::{EntryImpl, NamespaceImpl}; -const DEFAULT_DATA: &[u8; 10] = b"[greptime]"; - /// A builder for building entries for a namespace. pub struct EntryBuilder { /// The namespace of the entries. - pub ns: NamespaceImpl, + ns: NamespaceImpl, /// The next entry id to allocate. It starts from 0 by default. next_entry_id: AtomicEntryId, /// A generator for supporting random data generation. @@ -77,27 +75,6 @@ impl EntryBuilder { } } - /// Builds an entry with the default data. - // TODO(niebayes): may be remove this method since it's not used. - pub fn with_default_data(&self) -> EntryImpl { - EntryImpl { - data: DEFAULT_DATA.to_vec(), - id: self.alloc_entry_id(), - ns: self.ns.clone(), - } - } - - /// Builds a batch of entries each with random data. - pub fn with_random_data_batch(&self, batch_size: usize) -> Vec { - (0..batch_size) - .map(|_| EntryImpl { - data: self.make_random_data(), - id: self.alloc_entry_id(), - ns: self.ns.clone(), - }) - .collect() - } - /// Builds an entry with random data. pub fn with_random_data(&self) -> EntryImpl { EntryImpl { @@ -121,5 +98,3 @@ impl EntryBuilder { .collect() } } - -// TODO(niebayes): add tests for EntryBuilder. diff --git a/src/log-store/src/test_util/kafka/topic_builder.rs b/src/log-store/src/test_util/kafka/topic_decorator.rs similarity index 85% rename from src/log-store/src/test_util/kafka/topic_builder.rs rename to src/log-store/src/test_util/kafka/topic_decorator.rs index 037390914c8b..7b24774141e0 100644 --- a/src/log-store/src/test_util/kafka/topic_builder.rs +++ b/src/log-store/src/test_util/kafka/topic_decorator.rs @@ -38,8 +38,8 @@ impl ToString for Affix { } } -/// A builder for building topics on demand. -pub struct TopicBuilder { +/// Decorates a topic with the given prefix and suffix. +pub struct TopicDecorator { /// A prefix to be inserted at the front of each topic. prefix: Affix, /// A suffix to be inserted at the back of each topic. @@ -48,7 +48,7 @@ pub struct TopicBuilder { created: HashSet, } -impl Default for TopicBuilder { +impl Default for TopicDecorator { fn default() -> Self { Self { prefix: Affix::Nothing, @@ -58,7 +58,7 @@ impl Default for TopicBuilder { } } -impl TopicBuilder { +impl TopicDecorator { /// Overrides the current prefix with the given prefix. pub fn with_prefix(self, prefix: Affix) -> Self { Self { prefix, ..self } @@ -70,18 +70,18 @@ impl TopicBuilder { } /// Builds a topic by inserting a prefix and a suffix into the given topic. - pub fn build(&mut self, body: &str) -> Topic { + pub fn decorate(&mut self, topic: &str) -> Topic { const ITERS: usize = 24; for _ in 0..ITERS { - let topic = format!( + let decorated = format!( "{}_{}_{}", self.prefix.to_string(), - body, + topic, self.suffix.to_string() ); - if !self.created.contains(&topic) { - self.created.insert(topic.clone()); - return topic; + if !self.created.contains(&decorated) { + self.created.insert(decorated.clone()); + return decorated; } } unreachable!( @@ -90,5 +90,3 @@ impl TopicBuilder { ) } } - -// TODO(niebayes): add tests for TopicBuilder. From 21fbddcd3eeb83ae0931b1d6a7af5605a516b5eb Mon Sep 17 00:00:00 2001 From: niebayes Date: Sun, 31 Dec 2023 08:36:45 +0800 Subject: [PATCH 11/30] chore: resolve some todos --- src/log-store/src/kafka/log_store.rs | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/src/log-store/src/kafka/log_store.rs b/src/log-store/src/kafka/log_store.rs index 7d3c8bde36cd..21597f31609f 100644 --- a/src/log-store/src/kafka/log_store.rs +++ b/src/log-store/src/kafka/log_store.rs @@ -305,11 +305,10 @@ mod tests { } } - // TODO(niebayes): change `expected` to &[EntryImpl]. async fn check_entries( ns: &NamespaceImpl, start_offset: EntryId, - expected: Vec, + expected: &[EntryImpl], logstore: &KafkaLogStore, ) { let stream = logstore.read(ns, start_offset).await.unwrap(); @@ -319,7 +318,7 @@ mod tests { .into_iter() .flat_map(|x| x.unwrap()) .collect::>(); - assert_eq!(expected, got); + assert_eq!(expected, &got); } /// Appends entries for one region and checks all entries can be read successfully. @@ -342,7 +341,7 @@ mod tests { let entry = entry_builder.with_random_data(); let last_entry_id = logstore.append(entry.clone()).await.unwrap().last_entry_id; - check_entries(&ns, last_entry_id, vec![entry], &logstore).await; + check_entries(&ns, last_entry_id, &[entry], &logstore).await; let entries = (0..10) .map(|_| entry_builder.with_random_data()) @@ -352,7 +351,7 @@ mod tests { .await .unwrap() .last_entry_ids[&ns.region_id]; - check_entries(&ns, last_entry_id, entries, &logstore).await; + check_entries(&ns, last_entry_id, &entries, &logstore).await; } /// Appends entries for multiple regions and checks entries for each region can be read successfully. @@ -389,7 +388,7 @@ mod tests { for (i, ns) in region_namespaces.iter().enumerate() { let expected = region_entries[i].clone(); - check_entries(ns, last_entry_ids[&ns.region_id], expected, &logstore).await; + check_entries(ns, last_entry_ids[&ns.region_id], &expected, &logstore).await; } } From 84541ed298378e1b1ab2209b3d811facaf4328bf Mon Sep 17 00:00:00 2001 From: niebayes Date: Mon, 1 Jan 2024 11:29:28 +0800 Subject: [PATCH 12/30] test: add unit tests for kafka log store --- src/log-store/src/kafka/client_manager.rs | 4 +- src/log-store/src/kafka/log_store.rs | 193 +++++++++++------- src/log-store/src/test_util/kafka.rs | 2 + .../src/test_util/kafka/topic_decorator.rs | 2 +- 4 files changed, 125 insertions(+), 76 deletions(-) diff --git a/src/log-store/src/kafka/client_manager.rs b/src/log-store/src/kafka/client_manager.rs index c6201b864d93..3eb6312b2f72 100644 --- a/src/log-store/src/kafka/client_manager.rs +++ b/src/log-store/src/kafka/client_manager.rs @@ -144,7 +144,7 @@ mod tests { use crate::test_util::kafka::{create_topics, Affix, TopicDecorator}; /// Checks clients for the given topics are created. - async fn ensure_topics_exist(topics: &[Topic], client_manager: &ClientManager) { + async fn ensure_clients_exist(topics: &[Topic], client_manager: &ClientManager) { let client_pool = client_manager.client_pool.read().await; let all_exist = topics.iter().all(|topic| client_pool.contains_key(topic)); assert!(all_exist); @@ -182,7 +182,7 @@ mod tests { _ => unreachable!(), } - ensure_topics_exist(&topics, &manager).await; + ensure_clients_exist(&topics, &manager).await; } /// Sends `get_or_insert` requests sequentially to the client manager, and checks if it could handle them correctly. diff --git a/src/log-store/src/kafka/log_store.rs b/src/log-store/src/kafka/log_store.rs index 21597f31609f..12d7339e8e25 100644 --- a/src/log-store/src/kafka/log_store.rs +++ b/src/log-store/src/kafka/log_store.rs @@ -291,6 +291,7 @@ impl LogStore for KafkaLogStore { #[cfg(test)] mod tests { use common_test_util::wal::kafka::BROKER_ENDPOINTS_KEY; + use rand::seq::IteratorRandom; use super::*; use crate::get_broker_endpoints_from_env; @@ -298,37 +299,20 @@ mod tests { create_topics, entries_with_random_data, Affix, EntryBuilder, TopicDecorator, }; - fn new_namespace(topic: &str, region_id: u64) -> NamespaceImpl { - NamespaceImpl { - region_id, - topic: topic.to_string(), - } - } - - async fn check_entries( - ns: &NamespaceImpl, - start_offset: EntryId, - expected: &[EntryImpl], - logstore: &KafkaLogStore, - ) { - let stream = logstore.read(ns, start_offset).await.unwrap(); - let got = stream - .collect::>() - .await - .into_iter() - .flat_map(|x| x.unwrap()) - .collect::>(); - assert_eq!(expected, &got); + // Stores test context for a region. + struct RegionContext { + ns: NamespaceImpl, + entry_builder: EntryBuilder, + expected: Vec, } - /// Appends entries for one region and checks all entries can be read successfully. - #[tokio::test] - async fn test_one_region() { + /// Prepares for a test in that a log store is constructed and a collection of topics is created. + async fn prepare(test_name: &str, num_topics: usize) -> (KafkaLogStore, Vec) { let broker_endpoints = get_broker_endpoints_from_env!(BROKER_ENDPOINTS_KEY); let decorator = TopicDecorator::default() - .with_prefix(Affix::Fixed("test_one_region".to_string())) + .with_prefix(Affix::Fixed(test_name.to_string())) .with_suffix(Affix::TimeNow); - let topic = create_topics(1, decorator, &broker_endpoints, None).await[0].clone(); + let topics = create_topics(num_topics, decorator, &broker_endpoints, None).await; let config = KafkaConfig { broker_endpoints, @@ -336,64 +320,127 @@ mod tests { }; let logstore = KafkaLogStore::try_new(&config).await.unwrap(); - let ns = new_namespace(&topic, 0); - let entry_builder = EntryBuilder::new(ns.clone()); - let entry = entry_builder.with_random_data(); + (logstore, topics) + } + + /// Builds a context for each region. + fn build_contexts(num_regions: usize, topics: &[Topic]) -> Vec { + (0..num_regions) + .map(|i| { + let topic = &topics[i % topics.len()]; + let ns = NamespaceImpl { + region_id: i as u64, + topic: topic.to_string(), + }; + let entry_builder = EntryBuilder::new(ns.clone()); + RegionContext { + ns, + entry_builder, + expected: Vec::new(), + } + }) + .collect() + } - let last_entry_id = logstore.append(entry.clone()).await.unwrap().last_entry_id; - check_entries(&ns, last_entry_id, &[entry], &logstore).await; + /// Creates a vector containing indexes of all regions if the `all` is true. + /// Otherwise, creates a subset of the indexes. The cardinality of the subset + /// is nearly a quarter of that of the universe set. + fn all_or_subset(all: bool, num_regions: usize) -> Vec { + assert!(num_regions > 0); + let amount = if all { + num_regions + } else { + (num_regions / 4).max(1) + }; + (0..num_regions).choose_multiple(&mut rand::thread_rng(), amount) + } - let entries = (0..10) - .map(|_| entry_builder.with_random_data()) - .collect::>(); - let last_entry_id = logstore - .append_batch(entries.clone()) - .await - .unwrap() - .last_entry_ids[&ns.region_id]; - check_entries(&ns, last_entry_id, &entries, &logstore).await; + /// Builds entries for regions specified with `which`. + /// For example, assume `which[i] = x`, then entries for region with id = x will be built. + fn build_entries(region_contexts: &mut [RegionContext], which: Vec) -> Vec { + let mut entries = Vec::with_capacity(which.len()); + for i in which { + let ctx = &mut region_contexts[i]; + ctx.expected = entries_with_random_data(25, &ctx.entry_builder); + entries.push(ctx.expected.clone()); + } + entries.into_iter().flatten().collect() + } + + /// Reads entries for regions and checks for each region that the gotten entries are identical with the expected ones. + async fn check_entries( + region_contexts: &[RegionContext], + last_entry_ids: HashMap, + logstore: &KafkaLogStore, + which: Vec, + ) { + for i in which { + let ctx = ®ion_contexts[i]; + let start_offset = last_entry_ids[&ctx.ns.region_id]; + let stream = logstore.read(&ctx.ns, start_offset).await.unwrap(); + let got = stream + .collect::>() + .await + .into_iter() + .flat_map(|x| x.unwrap()) + .collect::>(); + assert_eq!(ctx.expected, got); + } + } + + /// Starts a test with: + /// * `test_name` - The name of the test. + /// * `num_topics` - Number of topics to be created in the preparation phase. + /// * `num_regions` - Number of regions involved in the test. + /// * `num_appends` - Number of append operations to be performed. + /// * `all` - All regions will be involved in an append operation if `all` is true. Otherwise, + /// an append operation will only randomly choose a subset of regions. + async fn test_with( + test_name: &str, + num_topics: usize, + num_regions: usize, + num_appends: usize, + all: bool, + ) { + let (logstore, topics) = prepare(test_name, num_topics).await; + let mut region_contexts = build_contexts(num_regions, &topics); + for _ in 0..num_appends { + let which = all_or_subset(all, num_regions); + let entries = build_entries(&mut region_contexts, which.clone()); + let last_entry_ids = logstore.append_batch(entries).await.unwrap().last_entry_ids; + check_entries(®ion_contexts, last_entry_ids, &logstore, which).await; + } + } + + /// Appends entries for one region and checks all entries can be read successfully. + #[tokio::test] + async fn test_one_region() { + test_with("test_one_region", 1, 1, 1, true).await; } /// Appends entries for multiple regions and checks entries for each region can be read successfully. /// A topic is assigned only a single region. #[tokio::test] async fn test_multi_regions_disjoint() { - let broker_endpoints = get_broker_endpoints_from_env!(BROKER_ENDPOINTS_KEY); - let decorator = TopicDecorator::default() - .with_prefix(Affix::Fixed("test_multi_regions_disjoint".to_string())) - .with_suffix(Affix::TimeNow); - let topics = create_topics(10, decorator, &broker_endpoints, None).await; - - let config = KafkaConfig { - broker_endpoints, - ..Default::default() - }; - let logstore = KafkaLogStore::try_new(&config).await.unwrap(); + test_with("test_multi_regions_disjoint", 10, 10, 1, true).await; + } - let (region_namespaces, entry_builders): (Vec<_>, Vec<_>) = topics - .iter() - .enumerate() - .map(|(i, topic)| { - let ns = new_namespace(topic, i as u64); - let entry_builder = EntryBuilder::new(ns.clone()); - (ns, entry_builder) - }) - .unzip(); - let region_entries = entry_builders - .iter() - .map(|builder| entries_with_random_data(25, builder)) - .collect::>(); - let entries = region_entries.iter().flatten().cloned().collect::>(); - let last_entry_ids = logstore.append_batch(entries).await.unwrap().last_entry_ids; - - for (i, ns) in region_namespaces.iter().enumerate() { - let expected = region_entries[i].clone(); - check_entries(ns, last_entry_ids[&ns.region_id], &expected, &logstore).await; - } + /// Appends entries for multiple regions and checks entries for each region can be read successfully. + /// A topic is assigned multiple regions. + #[tokio::test] + async fn test_multi_regions_overlapped() { + test_with("test_multi_regions_overlapped", 10, 50, 1, true).await; } /// Appends entries for multiple regions and checks entries for each region can be read successfully. - /// A topic may be assigned multiple regions. + /// A topic may be assigned multiple regions. The append operation repeats for a several iterations. + /// Each append operation will only append entries for a subset of randomly chosen regions. #[tokio::test] - async fn test_multi_regions_overlapped() {} + async fn test_multi_appends() { + test_with("test_multi_appends", 32, 256, 16, false).await; + } + + // / Appends a way-too large entry and checks the log store could handle it correctly. + // #[tokio::test] + // async fn test_way_too_large() {} } diff --git a/src/log-store/src/test_util/kafka.rs b/src/log-store/src/test_util/kafka.rs index aa431551c91e..3f926e6c1c8b 100644 --- a/src/log-store/src/test_util/kafka.rs +++ b/src/log-store/src/test_util/kafka.rs @@ -23,6 +23,7 @@ pub use crate::test_util::kafka::entry_builder::EntryBuilder; pub use crate::test_util::kafka::topic_decorator::{Affix, TopicDecorator}; /// Gets broker endpoints from environment variables with the given key. +/// Returns ["localhost:9092"] if no environment variables set for broker endpoints. #[macro_export] macro_rules! get_broker_endpoints_from_env { ($key:expr) => {{ @@ -37,6 +38,7 @@ macro_rules! get_broker_endpoints_from_env { } /// Creates `num_topiocs` number of topics from the seed topic which are going to be decorated with the given TopicDecorator. +/// A default seed `topic` will be used if the provided seed is None. pub async fn create_topics( num_topics: usize, mut decorator: TopicDecorator, diff --git a/src/log-store/src/test_util/kafka/topic_decorator.rs b/src/log-store/src/test_util/kafka/topic_decorator.rs index 7b24774141e0..e40db5a02227 100644 --- a/src/log-store/src/test_util/kafka/topic_decorator.rs +++ b/src/log-store/src/test_util/kafka/topic_decorator.rs @@ -32,7 +32,7 @@ impl ToString for Affix { fn to_string(&self) -> String { match self { Affix::Fixed(s) => s.to_string(), - Affix::TimeNow => chrono::Local::now().timestamp_millis().to_string(), + Affix::TimeNow => chrono::Local::now().timestamp_micros().to_string(), Affix::Nothing => String::default(), } } From a676ba8eee8ce2c7cb3fa378af9fd4e5b1233030 Mon Sep 17 00:00:00 2001 From: Ruihang Xia Date: Thu, 28 Dec 2023 11:47:19 +0800 Subject: [PATCH 13/30] chore: add deprecate develop branch warning Signed-off-by: Ruihang Xia --- README.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/README.md b/README.md index 415608ad3c83..9dd8c54c7e78 100644 --- a/README.md +++ b/README.md @@ -27,6 +27,9 @@ slack

+> [!WARNING] +> Our default branch has changed from `develop` to `main` (issue [#3025](https://github.com/GreptimeTeam/greptimedb/issues/3025)). Please update your local repository to use the `main` branch. + ## What is GreptimeDB GreptimeDB is an open-source time-series database with a special focus on From 68980c1d12cf66ad94718940bdbd38e90b47f2c8 Mon Sep 17 00:00:00 2001 From: niebayes Date: Mon, 1 Jan 2024 18:34:19 +0800 Subject: [PATCH 14/30] tmp: ready to move unit tests to an indie dir --- .../meta/src/wal/kafka/topic_manager.rs | 56 +++++++ src/common/meta/src/wal/options_allocator.rs | 39 +++++ src/log-store/src/kafka/client_manager.rs | 10 +- src/log-store/src/kafka/log_store.rs | 7 +- src/log-store/src/kafka/record_utils.rs | 31 +--- src/log-store/src/test_util/kafka.rs | 10 +- tests-integration/tests/wal.rs | 142 ------------------ 7 files changed, 116 insertions(+), 179 deletions(-) delete mode 100644 tests-integration/tests/wal.rs diff --git a/src/common/meta/src/wal/kafka/topic_manager.rs b/src/common/meta/src/wal/kafka/topic_manager.rs index 80aaa90d402f..a7519493ea63 100644 --- a/src/common/meta/src/wal/kafka/topic_manager.rs +++ b/src/common/meta/src/wal/kafka/topic_manager.rs @@ -295,4 +295,60 @@ mod tests { let manager = TopicManager::new(config, kv_backend); manager.start().await.unwrap(); } + + // Tests that the TopicManager allocates topics in a round-robin mannar. + // #[tokio::test] + // async fn test_kafka_alloc_topics() { + // let broker_endpoints = std::env::var(BROKER_ENDPOINTS_KEY) + // .unwrap() + // .split(',') + // .map(ToString::to_string) + // .collect::>(); + // let config = MetaSrvKafkaConfig { + // topic_name_prefix: "__test_kafka_alloc_topics".to_string(), + // replication_factor: broker_endpoints.len() as i16, + // broker_endpoints, + // ..Default::default() + // }; + // let kv_backend = Arc::new(MemoryKvBackend::new()) as KvBackendRef; + // let manager = KafkaTopicManager::new(config.clone(), kv_backend); + // manager.start().await.unwrap(); + + // // Topics should be created. + // let topics = (0..config.num_topics) + // .map(|topic_id| format!("{}_{topic_id}", config.topic_name_prefix)) + // .collect::>(); + + // // Selects exactly the number of `num_topics` topics one by one. + // for expected in topics.iter() { + // let got = manager.select().unwrap(); + // assert_eq!(got, expected); + // } + + // // Selects exactly the number of `num_topics` topics in a batching manner. + // let got = manager + // .select_batch(config.num_topics) + // .unwrap() + // .into_iter() + // .map(ToString::to_string) + // .collect::>(); + // assert_eq!(got, topics); + + // // Selects none. + // let got = manager.select_batch(config.num_topics).unwrap(); + // assert!(got.is_empty()); + + // // Selects more than the number of `num_topics` topics. + // let got = manager + // .select_batch(2 * config.num_topics) + // .unwrap() + // .into_iter() + // .map(ToString::to_string) + // .collect::>(); + // let expected = vec![topics.clone(); 2] + // .into_iter() + // .flatten() + // .collect::>(); + // assert_eq!(got, expected); + // } } diff --git a/src/common/meta/src/wal/options_allocator.rs b/src/common/meta/src/wal/options_allocator.rs index 6c2702053b87..b6f6a99ba089 100644 --- a/src/common/meta/src/wal/options_allocator.rs +++ b/src/common/meta/src/wal/options_allocator.rs @@ -128,4 +128,43 @@ mod tests { .collect(); assert_eq!(got, expected); } + + // Tests that the wal options allocator could successfully allocate Kafka wal options. + // #[tokio::test] + // async fn test_kafka_options_allocator() { + // let broker_endpoints = std::env::var(BROKER_ENDPOINTS_KEY) + // .unwrap() + // .split(',') + // .map(ToString::to_string) + // .collect::>(); + // let config = MetaSrvKafkaConfig { + // topic_name_prefix: "__test_kafka_options_allocator".to_string(), + // replication_factor: broker_endpoints.len() as i16, + // broker_endpoints, + // ..Default::default() + // }; + // let wal_config = WalConfig::Kafka(config.clone()); + // let kv_backend = Arc::new(MemoryKvBackend::new()) as KvBackendRef; + // let allocator = WalOptionsAllocator::new(wal_config, kv_backend); + // allocator.start().await.unwrap(); + + // let num_regions = 32; + // let regions = (0..num_regions).collect::>(); + // let got = allocate_region_wal_options(regions.clone(), &allocator).unwrap(); + + // // Topics should be allocated. + // let topics = (0..num_regions) + // .map(|topic_id| format!("{}_{topic_id}", config.topic_name_prefix)) + // .collect::>(); + // // 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::>(); + // assert_eq!(got, expected); + // } } diff --git a/src/log-store/src/kafka/client_manager.rs b/src/log-store/src/kafka/client_manager.rs index 3eb6312b2f72..27b01035c710 100644 --- a/src/log-store/src/kafka/client_manager.rs +++ b/src/log-store/src/kafka/client_manager.rs @@ -98,12 +98,12 @@ impl ClientManager { /// Gets the client associated with the topic. If the client does not exist, a new one will /// be created and returned. pub(crate) async fn get_or_insert(&self, topic: &Topic) -> Result { - let client_pool = self.client_pool.read().await; - if let Some(client) = client_pool.get(topic) { - return Ok(client.clone()); + { + let client_pool = self.client_pool.read().await; + if let Some(client) = client_pool.get(topic) { + return Ok(client.clone()); + } } - // Manullay releases the read lock. - drop(client_pool); // Acquires the write lock. let mut client_pool = self.client_pool.write().await; diff --git a/src/log-store/src/kafka/log_store.rs b/src/log-store/src/kafka/log_store.rs index 30ea652c51f7..9177ca423247 100644 --- a/src/log-store/src/kafka/log_store.rs +++ b/src/log-store/src/kafka/log_store.rs @@ -293,7 +293,7 @@ mod tests { use super::*; use crate::get_broker_endpoints_from_env; use crate::test_util::kafka::{ - create_topics, entries_with_random_data, Affix, EntryBuilder, TopicDecorator, + create_topics, entries_with_random_data, new_namespace, Affix, EntryBuilder, TopicDecorator, }; // Stores test context for a region. @@ -325,10 +325,7 @@ mod tests { (0..num_regions) .map(|i| { let topic = &topics[i % topics.len()]; - let ns = NamespaceImpl { - region_id: i as u64, - topic: topic.to_string(), - }; + let ns = new_namespace(topic, i as u64); let entry_builder = EntryBuilder::new(ns.clone()); RegionContext { ns, diff --git a/src/log-store/src/kafka/record_utils.rs b/src/log-store/src/kafka/record_utils.rs index 547e2c7b3112..2307020fa37d 100644 --- a/src/log-store/src/kafka/record_utils.rs +++ b/src/log-store/src/kafka/record_utils.rs @@ -145,26 +145,12 @@ pub(crate) fn decode_from_record(record: Record) -> Result> { #[cfg(test)] mod tests { use super::*; - - fn new_test_entry>(data: D, entry_id: EntryId, ns: NamespaceImpl) -> EntryImpl { - EntryImpl { - data: data.as_ref().to_vec(), - id: entry_id, - ns, - } - } + use crate::test_util::kafka::{entries_with_random_data, new_namespace, EntryBuilder}; #[test] fn test_serde_record_meta() { - let ns = NamespaceImpl { - region_id: 1, - topic: "test_topic".to_string(), - }; - let entries = vec![ - new_test_entry(b"111", 1, ns.clone()), - new_test_entry(b"2222", 2, ns.clone()), - new_test_entry(b"33333", 3, ns.clone()), - ]; + let ns = new_namespace("test_topic", 1); + let entries = entries_with_random_data(3, &EntryBuilder::new(ns.clone())); let meta = RecordMeta::new(ns, &entries); let encoded = serde_json::to_vec(&meta).unwrap(); let decoded: RecordMeta = serde_json::from_slice(&encoded).unwrap(); @@ -173,15 +159,8 @@ mod tests { #[test] fn test_encdec_record() { - let ns = NamespaceImpl { - region_id: 1, - topic: "test_topic".to_string(), - }; - let entries = vec![ - new_test_entry(b"111", 1, ns.clone()), - new_test_entry(b"2222", 2, ns.clone()), - new_test_entry(b"33333", 3, ns.clone()), - ]; + let ns = new_namespace("test_topic", 1); + let entries = entries_with_random_data(3, &EntryBuilder::new(ns.clone())); let record = encode_to_record(ns, entries.clone()).unwrap(); let decoded_entries = decode_from_record(record).unwrap(); assert_eq!(entries, decoded_entries); diff --git a/src/log-store/src/test_util/kafka.rs b/src/log-store/src/test_util/kafka.rs index 3f926e6c1c8b..1bd9a997c203 100644 --- a/src/log-store/src/test_util/kafka.rs +++ b/src/log-store/src/test_util/kafka.rs @@ -18,7 +18,7 @@ pub mod topic_decorator; use common_config::wal::KafkaWalTopic as Topic; use rskafka::client::ClientBuilder; -use crate::kafka::EntryImpl; +use crate::kafka::{EntryImpl, NamespaceImpl}; pub use crate::test_util::kafka::entry_builder::EntryBuilder; pub use crate::test_util::kafka::topic_decorator::{Affix, TopicDecorator}; @@ -66,6 +66,14 @@ pub async fn create_topics( topics } +/// Creates a new namespace with the given topic and region id. +pub fn new_namespace(topic: &str, region_id: u64) -> NamespaceImpl { + NamespaceImpl { + topic: topic.to_string(), + region_id, + } +} + /// Builds a batch of entries each with random data. pub fn entries_with_random_data(batch_size: usize, builder: &EntryBuilder) -> Vec { (0..batch_size) diff --git a/tests-integration/tests/wal.rs b/tests-integration/tests/wal.rs deleted file mode 100644 index 8bf09f55c550..000000000000 --- a/tests-integration/tests/wal.rs +++ /dev/null @@ -1,142 +0,0 @@ -// // 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 std::collections::HashMap; -// use std::sync::Arc; - -// use common_config::wal::KafkaConfig as DatanodeKafkaConfig; -// use common_config::{KafkaWalOptions, WalOptions}; -// use common_meta::kv_backend::memory::MemoryKvBackend; -// use common_meta::kv_backend::KvBackendRef; -// use common_meta::wal::kafka::{ -// KafkaConfig as MetaSrvKafkaConfig, TopicManager as KafkaTopicManager, -// }; -// use common_meta::wal::{allocate_region_wal_options, WalConfig, WalOptionsAllocator}; -// use futures::StreamExt; -// use log_store::error::Result as LogStoreResult; -// use log_store::kafka::log_store::KafkaLogStore; -// use log_store::kafka::{EntryImpl, NamespaceImpl}; -// use rskafka::client::controller::ControllerClient; -// use rskafka::client::ClientBuilder; -// use store_api::logstore::entry::Id as EntryId; -// use store_api::logstore::LogStore; -// use tests_integration::wal_util::{DockerCli, KafkaImage, DEFAULT_EXPOSED_PORT}; - -// // Notice: the following tests are literally unit tests. They are placed at here since -// // it seems too heavy to start a Kafka cluster for each unit test. - -// // The key of an env variable that stores a series of Kafka broker endpoints. -// const BROKER_ENDPOINTS_KEY: &str = "GT_KAFKA_ENDPOINTS"; - -// // Tests that the TopicManager allocates topics in a round-robin mannar. -// #[tokio::test] -// async fn test_kafka_alloc_topics() { -// let broker_endpoints = std::env::var(BROKER_ENDPOINTS_KEY) -// .unwrap() -// .split(',') -// .map(ToString::to_string) -// .collect::>(); -// let config = MetaSrvKafkaConfig { -// topic_name_prefix: "__test_kafka_alloc_topics".to_string(), -// replication_factor: broker_endpoints.len() as i16, -// broker_endpoints, -// ..Default::default() -// }; -// let kv_backend = Arc::new(MemoryKvBackend::new()) as KvBackendRef; -// let manager = KafkaTopicManager::new(config.clone(), kv_backend); -// manager.start().await.unwrap(); - -// // Topics should be created. -// let topics = (0..config.num_topics) -// .map(|topic_id| format!("{}_{topic_id}", config.topic_name_prefix)) -// .collect::>(); - -// // Selects exactly the number of `num_topics` topics one by one. -// for expected in topics.iter() { -// let got = manager.select().unwrap(); -// assert_eq!(got, expected); -// } - -// // Selects exactly the number of `num_topics` topics in a batching manner. -// let got = manager -// .select_batch(config.num_topics) -// .unwrap() -// .into_iter() -// .map(ToString::to_string) -// .collect::>(); -// assert_eq!(got, topics); - -// // Selects none. -// let got = manager.select_batch(config.num_topics).unwrap(); -// assert!(got.is_empty()); - -// // Selects more than the number of `num_topics` topics. -// let got = manager -// .select_batch(2 * config.num_topics) -// .unwrap() -// .into_iter() -// .map(ToString::to_string) -// .collect::>(); -// let expected = vec![topics.clone(); 2] -// .into_iter() -// .flatten() -// .collect::>(); -// assert_eq!(got, expected); -// } - -// // Tests that the wal options allocator could successfully allocate Kafka wal options. -// #[tokio::test] -// async fn test_kafka_options_allocator() { -// let broker_endpoints = std::env::var(BROKER_ENDPOINTS_KEY) -// .unwrap() -// .split(',') -// .map(ToString::to_string) -// .collect::>(); -// let config = MetaSrvKafkaConfig { -// topic_name_prefix: "__test_kafka_options_allocator".to_string(), -// replication_factor: broker_endpoints.len() as i16, -// broker_endpoints, -// ..Default::default() -// }; -// let wal_config = WalConfig::Kafka(config.clone()); -// let kv_backend = Arc::new(MemoryKvBackend::new()) as KvBackendRef; -// let allocator = WalOptionsAllocator::new(wal_config, kv_backend); -// allocator.start().await.unwrap(); - -// let num_regions = 32; -// let regions = (0..num_regions).collect::>(); -// let got = allocate_region_wal_options(regions.clone(), &allocator).unwrap(); - -// // Topics should be allocated. -// let topics = (0..num_regions) -// .map(|topic_id| format!("{}_{topic_id}", config.topic_name_prefix)) -// .collect::>(); -// // 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::>(); -// assert_eq!(got, expected); -// } - -// async fn create_topic(topic: &str, replication_factor: i16, client: &ControllerClient) { -// client -// .create_topic(topic, 1, replication_factor, 500) -// .await -// .unwrap(); -// } From d6f2b34807a9af5a239e3619f6bc389e472586ac Mon Sep 17 00:00:00 2001 From: niebayes Date: Tue, 2 Jan 2024 10:38:08 +0800 Subject: [PATCH 15/30] test: update unit tests for client manager --- README.md | 3 -- src/log-store/src/kafka/client_manager.rs | 43 ++++++++++++++--------- src/log-store/src/kafka/log_store.rs | 4 +-- 3 files changed, 29 insertions(+), 21 deletions(-) diff --git a/README.md b/README.md index 9dd8c54c7e78..415608ad3c83 100644 --- a/README.md +++ b/README.md @@ -27,9 +27,6 @@ slack

-> [!WARNING] -> Our default branch has changed from `develop` to `main` (issue [#3025](https://github.com/GreptimeTeam/greptimedb/issues/3025)). Please update your local repository to use the `main` branch. - ## What is GreptimeDB GreptimeDB is an open-source time-series database with a special focus on diff --git a/src/log-store/src/kafka/client_manager.rs b/src/log-store/src/kafka/client_manager.rs index 27b01035c710..dbee6dc872e0 100644 --- a/src/log-store/src/kafka/client_manager.rs +++ b/src/log-store/src/kafka/client_manager.rs @@ -143,20 +143,13 @@ mod tests { use crate::get_broker_endpoints_from_env; use crate::test_util::kafka::{create_topics, Affix, TopicDecorator}; - /// Checks clients for the given topics are created. - async fn ensure_clients_exist(topics: &[Topic], client_manager: &ClientManager) { - let client_pool = client_manager.client_pool.read().await; - let all_exist = topics.iter().all(|topic| client_pool.contains_key(topic)); - assert!(all_exist); - } - - async fn test_which(test_name: &str) { - // Creates a collection of topics in Kafka. + /// Prepares for a test in that a collection of topics and a client manager are created. + async fn prepare(test_name: &str, num_topics: usize) -> (ClientManager, Vec) { let broker_endpoints = get_broker_endpoints_from_env!(BROKER_ENDPOINTS_KEY); let decorator = TopicDecorator::default() .with_prefix(Affix::Fixed(test_name.to_string())) .with_suffix(Affix::TimeNow); - let topics = create_topics(256, decorator, &broker_endpoints, None).await; + let topics = create_topics(num_topics, decorator, &broker_endpoints, None).await; let config = KafkaConfig { broker_endpoints, @@ -164,36 +157,54 @@ mod tests { }; let manager = ClientManager::try_new(&config).await.unwrap(); + (manager, topics) + } + + /// Checks clients for the given topics are created. + async fn ensure_clients_exist(topics: &[Topic], client_manager: &ClientManager) { + let client_pool = client_manager.client_pool.read().await; + let all_exist = topics.iter().all(|topic| client_pool.contains_key(topic)); + assert!(all_exist); + } + + async fn test_with(test_name: &str) { + let (manager, topics) = prepare(test_name, 128).await; + // Assigns multiple regions to a topic. + let region_topic = (0..512) + .map(|region_id| (region_id, &topics[region_id % topics.len()])) + .collect::>(); + + // Gets the assigned topic for each region and then gets the associated client. match test_name { "test_sequential" => { // Gets all clients sequentially. - for topic in topics.iter() { + for (_, topic) in region_topic { manager.get_or_insert(topic).await.unwrap(); } } "test_parallel" => { // Gets all clients in parallel. - let tasks = topics - .iter() + let tasks = region_topic + .values() .map(|topic| manager.get_or_insert(topic)) .collect::>(); futures::future::try_join_all(tasks).await.unwrap(); } _ => unreachable!(), } - + // Ensures all clients are created successfully. ensure_clients_exist(&topics, &manager).await; } /// Sends `get_or_insert` requests sequentially to the client manager, and checks if it could handle them correctly. #[tokio::test] async fn test_sequential() { - test_which("test_sequential").await; + test_with("test_sequential").await; } /// Sends `get_or_insert` requests in parallel to the client manager, and checks if it could handle them correctly. #[tokio::test] async fn test_parallel() { - test_which("test_parallel").await; + test_with("test_parallel").await; } } diff --git a/src/log-store/src/kafka/log_store.rs b/src/log-store/src/kafka/log_store.rs index 9177ca423247..4703c5bd2b96 100644 --- a/src/log-store/src/kafka/log_store.rs +++ b/src/log-store/src/kafka/log_store.rs @@ -16,7 +16,7 @@ use std::collections::HashMap; use std::sync::Arc; use common_config::wal::{KafkaConfig, KafkaWalTopic as Topic, WalOptions}; -use common_telemetry::debug; +use common_telemetry::{debug, warn}; use futures_util::StreamExt; use rskafka::client::consumer::{StartOffset, StreamConsumerBuilder}; use rskafka::client::partition::OffsetAt; @@ -178,7 +178,7 @@ impl LogStore for KafkaLogStore { // Abort if there're no new entries. // FIXME(niebayes): how come this case happens? if start_offset > end_offset { - debug!( + warn!( "No new entries for ns {} in range [{}, {}]", ns, start_offset, end_offset ); From b06dc94af1e6e720215bcbd524097018c464f22c Mon Sep 17 00:00:00 2001 From: niebayes Date: Wed, 3 Jan 2024 01:31:49 +0800 Subject: [PATCH 16/30] test: add unit tests for meta srv remote wal --- Cargo.lock | 5 + src/common/meta/Cargo.toml | 1 + .../meta/src/wal/kafka/topic_manager.rs | 123 +++++++---------- .../meta/src/wal/kafka/topic_selector.rs | 8 ++ src/common/meta/src/wal/options_allocator.rs | 83 ++++++------ src/common/test-util/Cargo.toml | 4 + src/common/test-util/src/wal.rs | 5 - src/common/test-util/src/wal/kafka.rs | 50 +++++++ .../src/wal}/kafka/topic_decorator.rs | 0 src/log-store/src/kafka/client_manager.rs | 8 +- src/log-store/src/kafka/log_store.rs | 12 +- src/log-store/src/test_util/kafka.rs | 128 +++++++++++------- .../src/test_util/kafka/entry_builder.rs | 100 -------------- 13 files changed, 251 insertions(+), 276 deletions(-) rename src/{log-store/src/test_util => common/test-util/src/wal}/kafka/topic_decorator.rs (100%) delete mode 100644 src/log-store/src/test_util/kafka/entry_builder.rs diff --git a/Cargo.lock b/Cargo.lock index 663fc1788c0b..355b828073e9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1837,6 +1837,7 @@ dependencies = [ "common-recordbatch", "common-runtime", "common-telemetry", + "common-test-util", "common-time", "datatypes", "derive_builder 0.12.0", @@ -1984,9 +1985,13 @@ dependencies = [ name = "common-test-util" version = "0.5.0" dependencies = [ + "chrono", + "common-config", + "futures", "once_cell", "rand", "rskafka", + "store-api", "tempfile", "testcontainers", "tokio", diff --git a/src/common/meta/Cargo.toml b/src/common/meta/Cargo.toml index 5a15581f41c6..5f34915f9728 100644 --- a/src/common/meta/Cargo.toml +++ b/src/common/meta/Cargo.toml @@ -23,6 +23,7 @@ common-procedure.workspace = true common-recordbatch.workspace = true common-runtime.workspace = true common-telemetry.workspace = true +common-test-util.workspace = true common-time.workspace = true datatypes.workspace = true derive_builder.workspace = true diff --git a/src/common/meta/src/wal/kafka/topic_manager.rs b/src/common/meta/src/wal/kafka/topic_manager.rs index a7519493ea63..64f0e7b31e93 100644 --- a/src/common/meta/src/wal/kafka/topic_manager.rs +++ b/src/common/meta/src/wal/kafka/topic_manager.rs @@ -47,8 +47,7 @@ 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, + pub(crate) topic_pool: Vec, topic_selector: TopicSelectorRef, kv_backend: KvBackendRef, } @@ -242,7 +241,11 @@ impl TopicManager { mod tests { use std::env; + use chrono::format::Fixed; use common_telemetry::info; + use common_test_util::get_broker_endpoints; + use common_test_util::wal::kafka::topic_decorator::{Affix, TopicDecorator}; + use common_test_util::wal::kafka::{create_topics, BROKER_ENDPOINTS_KEY}; use super::*; use crate::kv_backend::memory::MemoryKvBackend; @@ -273,82 +276,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!(BROKER_ENDPOINTS_KEY); + // Constructs topics that should be created. + let mut decorator = TopicDecorator::default() + .with_prefix(Affix::Fixed("test_alloc_topics".to_string())) + .with_suffix(Affix::TimeNow); + let topics = (0..256) + .map(|i| decorator.decorate(&format!("topic_{i}"))) + .collect::>(); - 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::>(), + 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(); manager.start().await.unwrap(); - } - // Tests that the TopicManager allocates topics in a round-robin mannar. - // #[tokio::test] - // async fn test_kafka_alloc_topics() { - // let broker_endpoints = std::env::var(BROKER_ENDPOINTS_KEY) - // .unwrap() - // .split(',') - // .map(ToString::to_string) - // .collect::>(); - // let config = MetaSrvKafkaConfig { - // topic_name_prefix: "__test_kafka_alloc_topics".to_string(), - // replication_factor: broker_endpoints.len() as i16, - // broker_endpoints, - // ..Default::default() - // }; - // let kv_backend = Arc::new(MemoryKvBackend::new()) as KvBackendRef; - // let manager = KafkaTopicManager::new(config.clone(), kv_backend); - // manager.start().await.unwrap(); - - // // Topics should be created. - // let topics = (0..config.num_topics) - // .map(|topic_id| format!("{}_{topic_id}", config.topic_name_prefix)) - // .collect::>(); - - // // Selects exactly the number of `num_topics` topics one by one. - // for expected in topics.iter() { - // let got = manager.select().unwrap(); - // assert_eq!(got, expected); - // } - - // // Selects exactly the number of `num_topics` topics in a batching manner. - // let got = manager - // .select_batch(config.num_topics) - // .unwrap() - // .into_iter() - // .map(ToString::to_string) - // .collect::>(); - // assert_eq!(got, topics); - - // // Selects none. - // let got = manager.select_batch(config.num_topics).unwrap(); - // assert!(got.is_empty()); - - // // Selects more than the number of `num_topics` topics. - // let got = manager - // .select_batch(2 * config.num_topics) - // .unwrap() - // .into_iter() - // .map(ToString::to_string) - // .collect::>(); - // let expected = vec![topics.clone(); 2] - // .into_iter() - // .flatten() - // .collect::>(); - // assert_eq!(got, expected); - // } + // Selects exactly the number of `num_topics` topics one by one. + for expected in topics.iter() { + let got = manager.select().unwrap(); + assert_eq!(got, expected); + } + + // 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::>(); + 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::>(); + let expected = vec![topics.clone(); 2] + .into_iter() + .flatten() + .collect::>(); + assert_eq!(got, expected); + } } diff --git a/src/common/meta/src/wal/kafka/topic_selector.rs b/src/common/meta/src/wal/kafka/topic_selector.rs index fe7517bfd0b5..df4feb8b2e12 100644 --- a/src/common/meta/src/wal/kafka/topic_selector.rs +++ b/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(); diff --git a/src/common/meta/src/wal/options_allocator.rs b/src/common/meta/src/wal/options_allocator.rs index b6f6a99ba089..44a15c69e3fd 100644 --- a/src/common/meta/src/wal/options_allocator.rs +++ b/src/common/meta/src/wal/options_allocator.rs @@ -105,11 +105,15 @@ pub fn allocate_region_wal_options( #[cfg(test)] mod tests { + use common_test_util::get_broker_endpoints; + use common_test_util::wal::kafka::topic_decorator::{Affix, TopicDecorator}; + use common_test_util::wal::kafka::BROKER_ENDPOINTS_KEY; + use super::*; use crate::kv_backend::memory::MemoryKvBackend; + 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; @@ -130,41 +134,44 @@ mod tests { } // Tests that the wal options allocator could successfully allocate Kafka wal options. - // #[tokio::test] - // async fn test_kafka_options_allocator() { - // let broker_endpoints = std::env::var(BROKER_ENDPOINTS_KEY) - // .unwrap() - // .split(',') - // .map(ToString::to_string) - // .collect::>(); - // let config = MetaSrvKafkaConfig { - // topic_name_prefix: "__test_kafka_options_allocator".to_string(), - // replication_factor: broker_endpoints.len() as i16, - // broker_endpoints, - // ..Default::default() - // }; - // let wal_config = WalConfig::Kafka(config.clone()); - // let kv_backend = Arc::new(MemoryKvBackend::new()) as KvBackendRef; - // let allocator = WalOptionsAllocator::new(wal_config, kv_backend); - // allocator.start().await.unwrap(); - - // let num_regions = 32; - // let regions = (0..num_regions).collect::>(); - // let got = allocate_region_wal_options(regions.clone(), &allocator).unwrap(); - - // // Topics should be allocated. - // let topics = (0..num_regions) - // .map(|topic_id| format!("{}_{topic_id}", config.topic_name_prefix)) - // .collect::>(); - // // 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::>(); - // assert_eq!(got, expected); - // } + #[tokio::test] + async fn test_allocator_with_kafka() { + let broker_endpoints = get_broker_endpoints!(BROKER_ENDPOINTS_KEY); + // Constructs topics that should be created. + let mut decorator = TopicDecorator::default() + .with_prefix(Affix::Fixed("test_allocator_with_kafka".to_string())) + .with_suffix(Affix::TimeNow); + let topics = (0..256) + .map(|i| decorator.decorate(&format!("topic_{i}"))) + .collect::>(); + + // 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); + + // 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::>(); + 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::>(); + assert_eq!(got, expected); + } } diff --git a/src/common/test-util/Cargo.toml b/src/common/test-util/Cargo.toml index e3caa859d21f..0884c626c2cf 100644 --- a/src/common/test-util/Cargo.toml +++ b/src/common/test-util/Cargo.toml @@ -5,9 +5,13 @@ edition.workspace = true license.workspace = true [dependencies] +chrono.workspace = true +common-config.workspace = true +futures.workspace = true once_cell.workspace = true rand.workspace = true rskafka.workspace = true +store-api.workspace = true tempfile.workspace = true testcontainers = "0.15.0" tokio.workspace = true diff --git a/src/common/test-util/src/wal.rs b/src/common/test-util/src/wal.rs index c1085a8f1e17..28c04633b640 100644 --- a/src/common/test-util/src/wal.rs +++ b/src/common/test-util/src/wal.rs @@ -13,8 +13,3 @@ // limitations under the License. pub mod kafka; - -pub use testcontainers::clients::Cli as DockerCli; - -pub use crate::wal::kafka::config::KAFKA_ADVERTISED_LISTENER_PORT as DEFAULT_EXPOSED_PORT; -pub use crate::wal::kafka::image::Image as KafkaImage; diff --git a/src/common/test-util/src/wal/kafka.rs b/src/common/test-util/src/wal/kafka.rs index 427cbea89a71..db8fde3ea4f4 100644 --- a/src/common/test-util/src/wal/kafka.rs +++ b/src/common/test-util/src/wal/kafka.rs @@ -14,5 +14,55 @@ pub mod config; pub mod image; +pub mod topic_decorator; + +use common_config::wal::KafkaWalTopic as Topic; +use rskafka::client::ClientBuilder; + +use crate::wal::kafka::topic_decorator::TopicDecorator; pub const BROKER_ENDPOINTS_KEY: &str = "GT_KAFKA_ENDPOINTS"; + +/// Gets broker endpoints from environment variables with the given key. +/// Returns ["localhost:9092"] if no environment variables set for broker endpoints. +#[macro_export] +macro_rules! get_broker_endpoints { + ($key:expr) => {{ + let broker_endpoints = std::env::var($key) + .unwrap_or("localhost:9092".to_string()) + .split(',') + .map(ToString::to_string) + .collect::>(); + assert!(!broker_endpoints.is_empty()); + broker_endpoints + }}; +} + +/// Creates `num_topiocs` number of topics from the seed topic which are going to be decorated with the given TopicDecorator. +/// A default seed `topic` will be used if the provided seed is None. +pub async fn create_topics( + num_topics: usize, + mut decorator: TopicDecorator, + broker_endpoints: &[String], + seed: Option<&str>, +) -> Vec { + assert!(!broker_endpoints.is_empty()); + + let client = ClientBuilder::new(broker_endpoints.to_vec()) + .build() + .await + .unwrap(); + let ctrl_client = client.controller_client().unwrap(); + + let seed = seed.unwrap_or("topic"); + let (topics, tasks): (Vec<_>, Vec<_>) = (0..num_topics) + .map(|i| { + let topic = decorator.decorate(&format!("{seed}_{i}")); + let task = ctrl_client.create_topic(topic.clone(), 1, 1, 500); + (topic, task) + }) + .unzip(); + futures::future::try_join_all(tasks).await.unwrap(); + + topics +} diff --git a/src/log-store/src/test_util/kafka/topic_decorator.rs b/src/common/test-util/src/wal/kafka/topic_decorator.rs similarity index 100% rename from src/log-store/src/test_util/kafka/topic_decorator.rs rename to src/common/test-util/src/wal/kafka/topic_decorator.rs diff --git a/src/log-store/src/kafka/client_manager.rs b/src/log-store/src/kafka/client_manager.rs index dbee6dc872e0..d114bffa0813 100644 --- a/src/log-store/src/kafka/client_manager.rs +++ b/src/log-store/src/kafka/client_manager.rs @@ -137,15 +137,15 @@ impl ClientManager { #[cfg(test)] mod tests { - use common_test_util::wal::kafka::BROKER_ENDPOINTS_KEY; + use common_test_util::get_broker_endpoints; + use common_test_util::wal::kafka::topic_decorator::{Affix, TopicDecorator}; + use common_test_util::wal::kafka::{create_topics, BROKER_ENDPOINTS_KEY}; use super::*; - use crate::get_broker_endpoints_from_env; - use crate::test_util::kafka::{create_topics, Affix, TopicDecorator}; /// Prepares for a test in that a collection of topics and a client manager are created. async fn prepare(test_name: &str, num_topics: usize) -> (ClientManager, Vec) { - let broker_endpoints = get_broker_endpoints_from_env!(BROKER_ENDPOINTS_KEY); + let broker_endpoints = get_broker_endpoints!(BROKER_ENDPOINTS_KEY); let decorator = TopicDecorator::default() .with_prefix(Affix::Fixed(test_name.to_string())) .with_suffix(Affix::TimeNow); diff --git a/src/log-store/src/kafka/log_store.rs b/src/log-store/src/kafka/log_store.rs index 4703c5bd2b96..c1d663a62804 100644 --- a/src/log-store/src/kafka/log_store.rs +++ b/src/log-store/src/kafka/log_store.rs @@ -287,14 +287,13 @@ impl LogStore for KafkaLogStore { #[cfg(test)] mod tests { - use common_test_util::wal::kafka::BROKER_ENDPOINTS_KEY; + use common_test_util::get_broker_endpoints; + use common_test_util::wal::kafka::topic_decorator::{Affix, TopicDecorator}; + use common_test_util::wal::kafka::{create_topics, BROKER_ENDPOINTS_KEY}; use rand::seq::IteratorRandom; use super::*; - use crate::get_broker_endpoints_from_env; - use crate::test_util::kafka::{ - create_topics, entries_with_random_data, new_namespace, Affix, EntryBuilder, TopicDecorator, - }; + use crate::test_util::kafka::{entries_with_random_data, new_namespace, EntryBuilder}; // Stores test context for a region. struct RegionContext { @@ -305,7 +304,7 @@ mod tests { /// Prepares for a test in that a log store is constructed and a collection of topics is created. async fn prepare(test_name: &str, num_topics: usize) -> (KafkaLogStore, Vec) { - let broker_endpoints = get_broker_endpoints_from_env!(BROKER_ENDPOINTS_KEY); + let broker_endpoints = get_broker_endpoints!(BROKER_ENDPOINTS_KEY); let decorator = TopicDecorator::default() .with_prefix(Affix::Fixed(test_name.to_string())) .with_suffix(Affix::TimeNow); @@ -434,6 +433,7 @@ mod tests { test_with("test_multi_appends", 32, 256, 16, false).await; } + // TODO(niebayes): implement. // / Appends a way-too large entry and checks the log store could handle it correctly. // #[tokio::test] // async fn test_way_too_large() {} diff --git a/src/log-store/src/test_util/kafka.rs b/src/log-store/src/test_util/kafka.rs index 1bd9a997c203..cdce23d8b156 100644 --- a/src/log-store/src/test_util/kafka.rs +++ b/src/log-store/src/test_util/kafka.rs @@ -11,66 +11,86 @@ // 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 std::sync::atomic::{AtomicU64 as AtomicEntryId, Ordering}; +use std::sync::Mutex; -pub mod entry_builder; -pub mod topic_decorator; - -use common_config::wal::KafkaWalTopic as Topic; -use rskafka::client::ClientBuilder; +use rand::rngs::ThreadRng; +use rand::seq::SliceRandom; +use rand::{thread_rng, Rng}; +use store_api::logstore::EntryId; use crate::kafka::{EntryImpl, NamespaceImpl}; -pub use crate::test_util::kafka::entry_builder::EntryBuilder; -pub use crate::test_util::kafka::topic_decorator::{Affix, TopicDecorator}; -/// Gets broker endpoints from environment variables with the given key. -/// Returns ["localhost:9092"] if no environment variables set for broker endpoints. -#[macro_export] -macro_rules! get_broker_endpoints_from_env { - ($key:expr) => {{ - let broker_endpoints = std::env::var($key) - .unwrap_or("localhost:9092".to_string()) - .split(',') - .map(ToString::to_string) - .collect::>(); - assert!(!broker_endpoints.is_empty()); - broker_endpoints - }}; +/// A builder for building entries for a namespace. +pub struct EntryBuilder { + /// The namespace of the entries. + ns: NamespaceImpl, + /// The next entry id to allocate. It starts from 0 by default. + next_entry_id: AtomicEntryId, + /// A generator for supporting random data generation. + /// Wrapped with Mutex> to provide interior mutability. + rng: Mutex>, + /// The data pool from which random data is constructed. + data_pool: Vec, } -/// Creates `num_topiocs` number of topics from the seed topic which are going to be decorated with the given TopicDecorator. -/// A default seed `topic` will be used if the provided seed is None. -pub async fn create_topics( - num_topics: usize, - mut decorator: TopicDecorator, - broker_endpoints: &[String], - seed: Option<&str>, -) -> Vec { - assert!(!broker_endpoints.is_empty()); +impl EntryBuilder { + /// Creates an EntryBuilder for the given namespace. + pub fn new(ns: NamespaceImpl) -> Self { + // Makes a data pool with alphabets and numbers. + let data_pool = ('a'..='z') + .chain('A'..='Z') + .chain((0..=9).map(|digit| char::from_digit(digit, 10).unwrap())) + .map(|c| c as u8) + .collect::>(); + Self { + ns, + next_entry_id: AtomicEntryId::new(0), + rng: Mutex::new(Some(thread_rng())), + data_pool, + } + } - let client = ClientBuilder::new(broker_endpoints.to_vec()) - .build() - .await - .unwrap(); - let ctrl_client = client.controller_client().unwrap(); + /// Sets the next entry id to the given entry id. + pub fn next_entry_id(self, entry_id: EntryId) -> Self { + Self { + next_entry_id: AtomicEntryId::new(entry_id), + ..self + } + } - let seed = seed.unwrap_or("topic"); - let (topics, tasks): (Vec<_>, Vec<_>) = (0..num_topics) - .map(|i| { - let topic = decorator.decorate(&format!("{seed}_{i}")); - let task = ctrl_client.create_topic(topic.clone(), 1, 1, 500); - (topic, task) - }) - .unzip(); - futures::future::try_join_all(tasks).await.unwrap(); + /// Skips the next `step` entry ids and returns the next entry id after the stepping. + pub fn skip(&mut self, step: EntryId) -> EntryId { + let old = self.next_entry_id.fetch_add(step, Ordering::Relaxed); + old + step + } - topics -} + /// Builds an entry with the given data. + pub fn with_data>(&self, data: D) -> EntryImpl { + EntryImpl { + data: data.as_ref().to_vec(), + id: self.alloc_entry_id(), + ns: self.ns.clone(), + } + } -/// Creates a new namespace with the given topic and region id. -pub fn new_namespace(topic: &str, region_id: u64) -> NamespaceImpl { - NamespaceImpl { - topic: topic.to_string(), - region_id, + /// Builds an entry with random data. + pub fn with_random_data(&self) -> EntryImpl { + self.with_data(self.make_random_data()) + } + + fn alloc_entry_id(&self) -> EntryId { + self.next_entry_id.fetch_add(1, Ordering::Relaxed) + } + + fn make_random_data(&self) -> Vec { + let mut rng_guard = self.rng.lock().unwrap(); + let mut rng = rng_guard.as_mut().unwrap(); + let amount = rng.gen_range(0..self.data_pool.len()); + self.data_pool + .choose_multiple(&mut rng, amount) + .copied() + .collect() } } @@ -80,3 +100,11 @@ pub fn entries_with_random_data(batch_size: usize, builder: &EntryBuilder) -> Ve .map(|_| builder.with_random_data()) .collect() } + +/// Creates a new Kafka namespace with the given topic and region id. +pub fn new_namespace(topic: &str, region_id: u64) -> NamespaceImpl { + NamespaceImpl { + topic: topic.to_string(), + region_id, + } +} diff --git a/src/log-store/src/test_util/kafka/entry_builder.rs b/src/log-store/src/test_util/kafka/entry_builder.rs deleted file mode 100644 index 2adf85c13345..000000000000 --- a/src/log-store/src/test_util/kafka/entry_builder.rs +++ /dev/null @@ -1,100 +0,0 @@ -// 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 std::sync::atomic::{AtomicU64 as AtomicEntryId, Ordering}; -use std::sync::Mutex; - -use rand::rngs::ThreadRng; -use rand::seq::SliceRandom; -use rand::{thread_rng, Rng}; -use store_api::logstore::EntryId; - -use crate::kafka::{EntryImpl, NamespaceImpl}; - -/// A builder for building entries for a namespace. -pub struct EntryBuilder { - /// The namespace of the entries. - ns: NamespaceImpl, - /// The next entry id to allocate. It starts from 0 by default. - next_entry_id: AtomicEntryId, - /// A generator for supporting random data generation. - /// Wrapped with Mutex> to provide interior mutability. - rng: Mutex>, - /// The data pool from which random data is constructed. - data_pool: Vec, -} - -impl EntryBuilder { - /// Creates an EntryBuilder for the given namespace. - pub fn new(ns: NamespaceImpl) -> Self { - // Makes a data pool with alphabets and numbers. - let data_pool = ('a'..='z') - .chain('A'..='Z') - .chain((0..=9).map(|digit| char::from_digit(digit, 10).unwrap())) - .map(|c| c as u8) - .collect::>(); - Self { - ns, - next_entry_id: AtomicEntryId::new(0), - rng: Mutex::new(Some(thread_rng())), - data_pool, - } - } - - /// Sets the next entry id to the given entry id. - pub fn next_entry_id(self, entry_id: EntryId) -> Self { - Self { - next_entry_id: AtomicEntryId::new(entry_id), - ..self - } - } - - /// Skips the next `step` entry ids and returns the next entry id after the stepping. - pub fn skip(&mut self, step: EntryId) -> EntryId { - let old = self.next_entry_id.fetch_add(step, Ordering::Relaxed); - old + step - } - - /// Builds an entry with the given data. - pub fn with_data>(&self, data: D) -> EntryImpl { - EntryImpl { - data: data.as_ref().to_vec(), - id: self.alloc_entry_id(), - ns: self.ns.clone(), - } - } - - /// Builds an entry with random data. - pub fn with_random_data(&self) -> EntryImpl { - EntryImpl { - data: self.make_random_data(), - id: self.alloc_entry_id(), - ns: self.ns.clone(), - } - } - - fn alloc_entry_id(&self) -> EntryId { - self.next_entry_id.fetch_add(1, Ordering::Relaxed) - } - - fn make_random_data(&self) -> Vec { - let mut rng_guard = self.rng.lock().unwrap(); - let mut rng = rng_guard.as_mut().unwrap(); - let amount = rng.gen_range(0..self.data_pool.len()); - self.data_pool - .choose_multiple(&mut rng, amount) - .copied() - .collect() - } -} From 070404bb06c0984879f44da546048efffcf5ad7c Mon Sep 17 00:00:00 2001 From: niebayes Date: Wed, 3 Jan 2024 01:34:12 +0800 Subject: [PATCH 17/30] fix: license --- src/log-store/src/test_util/kafka.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/src/log-store/src/test_util/kafka.rs b/src/log-store/src/test_util/kafka.rs index cdce23d8b156..409b7fc3acac 100644 --- a/src/log-store/src/test_util/kafka.rs +++ b/src/log-store/src/test_util/kafka.rs @@ -11,6 +11,7 @@ // 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 std::sync::atomic::{AtomicU64 as AtomicEntryId, Ordering}; use std::sync::Mutex; From 40eccc9f61d34210b77b12cb478b398d84774fee Mon Sep 17 00:00:00 2001 From: niebayes Date: Wed, 3 Jan 2024 11:34:11 +0800 Subject: [PATCH 18/30] fix: test --- src/common/meta/src/wal/kafka/topic_manager.rs | 16 +++++++++------- src/common/meta/src/wal/options_allocator.rs | 7 ++++++- src/common/test-util/src/wal/kafka.rs | 2 +- src/common/test-util/src/wal/kafka/image.rs | 4 ++-- .../test-util/src/wal/kafka/topic_decorator.rs | 14 ++++++++------ src/log-store/src/test_util/kafka.rs | 11 +++++------ 6 files changed, 31 insertions(+), 23 deletions(-) diff --git a/src/common/meta/src/wal/kafka/topic_manager.rs b/src/common/meta/src/wal/kafka/topic_manager.rs index 64f0e7b31e93..4307d5b9f8b1 100644 --- a/src/common/meta/src/wal/kafka/topic_manager.rs +++ b/src/common/meta/src/wal/kafka/topic_manager.rs @@ -48,7 +48,7 @@ const DEFAULT_PARTITION: i32 = 0; pub struct TopicManager { config: KafkaConfig, pub(crate) topic_pool: Vec, - topic_selector: TopicSelectorRef, + pub(crate) topic_selector: TopicSelectorRef, kv_backend: KvBackendRef, } @@ -249,7 +249,6 @@ mod tests { use super::*; 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] @@ -281,7 +280,7 @@ mod tests { async fn test_alloc_topics() { let broker_endpoints = get_broker_endpoints!(BROKER_ENDPOINTS_KEY); // Constructs topics that should be created. - let mut decorator = TopicDecorator::default() + let decorator = TopicDecorator::default() .with_prefix(Affix::Fixed("test_alloc_topics".to_string())) .with_suffix(Affix::TimeNow); let topics = (0..256) @@ -298,13 +297,16 @@ mod tests { 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. - for expected in topics.iter() { - let got = manager.select().unwrap(); - assert_eq!(got, expected); - } + let got = (0..topics.len()) + .map(|_| manager.select().unwrap()) + .cloned() + .collect::>(); + assert_eq!(got, topics); // Selects exactly the number of `num_topics` topics in a batching manner. let got = manager diff --git a/src/common/meta/src/wal/options_allocator.rs b/src/common/meta/src/wal/options_allocator.rs index 44a15c69e3fd..4c85d987889a 100644 --- a/src/common/meta/src/wal/options_allocator.rs +++ b/src/common/meta/src/wal/options_allocator.rs @@ -111,6 +111,7 @@ mod tests { use super::*; 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. @@ -138,7 +139,7 @@ mod tests { async fn test_allocator_with_kafka() { let broker_endpoints = get_broker_endpoints!(BROKER_ENDPOINTS_KEY); // Constructs topics that should be created. - let mut decorator = TopicDecorator::default() + let decorator = TopicDecorator::default() .with_prefix(Affix::Fixed("test_allocator_with_kafka".to_string())) .with_suffix(Affix::TimeNow); let topics = (0..256) @@ -153,6 +154,10 @@ mod tests { }; 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()); diff --git a/src/common/test-util/src/wal/kafka.rs b/src/common/test-util/src/wal/kafka.rs index db8fde3ea4f4..84c7989d56b3 100644 --- a/src/common/test-util/src/wal/kafka.rs +++ b/src/common/test-util/src/wal/kafka.rs @@ -42,7 +42,7 @@ macro_rules! get_broker_endpoints { /// A default seed `topic` will be used if the provided seed is None. pub async fn create_topics( num_topics: usize, - mut decorator: TopicDecorator, + decorator: TopicDecorator, broker_endpoints: &[String], seed: Option<&str>, ) -> Vec { diff --git a/src/common/test-util/src/wal/kafka/image.rs b/src/common/test-util/src/wal/kafka/image.rs index 3428549a4c4b..48a2449fab96 100644 --- a/src/common/test-util/src/wal/kafka/image.rs +++ b/src/common/test-util/src/wal/kafka/image.rs @@ -120,8 +120,8 @@ mod tests { let container = docker.run(Image::with_exposed_port(port)); // Creates a Kafka client. - let bootstrap_brokers = vec![format!("127.0.0.1:{}", container.get_host_port_ipv4(port))]; - let client = ClientBuilder::new(bootstrap_brokers).build().await.unwrap(); + let broker_endpoints = vec![format!("127.0.0.1:{}", container.get_host_port_ipv4(port))]; + let client = ClientBuilder::new(broker_endpoints).build().await.unwrap(); // Creates a topic. let topic = "test_topic"; diff --git a/src/common/test-util/src/wal/kafka/topic_decorator.rs b/src/common/test-util/src/wal/kafka/topic_decorator.rs index e40db5a02227..60f4ac1c0993 100644 --- a/src/common/test-util/src/wal/kafka/topic_decorator.rs +++ b/src/common/test-util/src/wal/kafka/topic_decorator.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::cell::RefCell; use std::collections::HashSet; use common_config::wal::KafkaWalTopic as Topic; @@ -45,7 +46,7 @@ pub struct TopicDecorator { /// A suffix to be inserted at the back of each topic. suffix: Affix, /// Topics built so far. Used to filter out duplicate topics. - created: HashSet, + created: RefCell>, } impl Default for TopicDecorator { @@ -53,7 +54,7 @@ impl Default for TopicDecorator { Self { prefix: Affix::Nothing, suffix: Affix::Nothing, - created: HashSet::with_capacity(256), + created: RefCell::new(HashSet::with_capacity(256)), } } } @@ -70,7 +71,7 @@ impl TopicDecorator { } /// Builds a topic by inserting a prefix and a suffix into the given topic. - pub fn decorate(&mut self, topic: &str) -> Topic { + pub fn decorate(&self, topic: &str) -> Topic { const ITERS: usize = 24; for _ in 0..ITERS { let decorated = format!( @@ -79,13 +80,14 @@ impl TopicDecorator { topic, self.suffix.to_string() ); - if !self.created.contains(&decorated) { - self.created.insert(decorated.clone()); + let mut created_topics = self.created.borrow_mut(); + if !created_topics.contains(&decorated) { + created_topics.insert(decorated.clone()); return decorated; } } unreachable!( - "Building a topic should be completed within iterations {}", + "Topic decoration should be completed within iterations {}", ITERS ) } diff --git a/src/log-store/src/test_util/kafka.rs b/src/log-store/src/test_util/kafka.rs index 409b7fc3acac..7fc25e477c0a 100644 --- a/src/log-store/src/test_util/kafka.rs +++ b/src/log-store/src/test_util/kafka.rs @@ -12,8 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::cell::RefCell; use std::sync::atomic::{AtomicU64 as AtomicEntryId, Ordering}; -use std::sync::Mutex; use rand::rngs::ThreadRng; use rand::seq::SliceRandom; @@ -30,7 +30,7 @@ pub struct EntryBuilder { next_entry_id: AtomicEntryId, /// A generator for supporting random data generation. /// Wrapped with Mutex> to provide interior mutability. - rng: Mutex>, + rng: RefCell, /// The data pool from which random data is constructed. data_pool: Vec, } @@ -47,7 +47,7 @@ impl EntryBuilder { Self { ns, next_entry_id: AtomicEntryId::new(0), - rng: Mutex::new(Some(thread_rng())), + rng: RefCell::new(thread_rng()), data_pool, } } @@ -85,11 +85,10 @@ impl EntryBuilder { } fn make_random_data(&self) -> Vec { - let mut rng_guard = self.rng.lock().unwrap(); - let mut rng = rng_guard.as_mut().unwrap(); + let mut rng = self.rng.borrow_mut(); let amount = rng.gen_range(0..self.data_pool.len()); self.data_pool - .choose_multiple(&mut rng, amount) + .choose_multiple(&mut *rng, amount) .copied() .collect() } From 815e03a1d97416083e1641f2b9188b1eaad8b028 Mon Sep 17 00:00:00 2001 From: niebayes Date: Wed, 3 Jan 2024 11:51:56 +0800 Subject: [PATCH 19/30] refactor: kafka image --- src/common/test-util/src/wal/kafka.rs | 1 - src/common/test-util/src/wal/kafka/config.rs | 100 ------------------- src/common/test-util/src/wal/kafka/image.rs | 89 +++++++++++++---- 3 files changed, 69 insertions(+), 121 deletions(-) delete mode 100644 src/common/test-util/src/wal/kafka/config.rs diff --git a/src/common/test-util/src/wal/kafka.rs b/src/common/test-util/src/wal/kafka.rs index 84c7989d56b3..db9d49b4be29 100644 --- a/src/common/test-util/src/wal/kafka.rs +++ b/src/common/test-util/src/wal/kafka.rs @@ -12,7 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -pub mod config; pub mod image; pub mod topic_decorator; diff --git a/src/common/test-util/src/wal/kafka/config.rs b/src/common/test-util/src/wal/kafka/config.rs deleted file mode 100644 index fad01474888c..000000000000 --- a/src/common/test-util/src/wal/kafka/config.rs +++ /dev/null @@ -1,100 +0,0 @@ -// 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 std::collections::HashMap; - -use testcontainers::core::WaitFor; - -/// Through which port the Zookeeper node listens for external traffics, e.g. traffics from the Kafka node. -pub const ZOOKEEPER_PORT: u16 = 2181; -/// Through which port the Kafka node listens for internal traffics, i.e. traffics between Kafka nodes in the same Kafka cluster. -pub const KAFKA_LISTENER_PORT: u16 = 19092; -/// Through which port the Kafka node listens for external traffics, e.g. traffics from Kafka clients. -pub const KAFKA_ADVERTISED_LISTENER_PORT: u16 = 9092; - -/// Configurations for a Kafka runtime. -#[derive(Debug, Clone)] -pub struct Config { - /// The name of the Kafka image hosted in the docker hub. - pub image_name: String, - /// The tag of the kafka image hosted in the docker hub. - /// Warning: please use a tag with long-term support. Do not use `latest` or any other tags that - /// the underlying image may suddenly change. - pub image_tag: String, - /// The runtime is running in a docker container and has its own network. In order to be used by the host machine, - /// the runtime must expose an internal port. For e.g. assume the runtime has an internal port 9092, - /// and the `exposed_port` is set to 9092, then the host machine can get a mapped external port with - /// `container.get_host_port_ipv4(exposed_port)`. With the mapped port, the host machine could connect with the runtime. - pub exposed_port: u16, - /// The runtime is regarded ready to be used if all ready conditions are met. - /// Warning: be sure to update the conditions when necessary if the image is altered. - pub ready_conditions: Vec, - /// The environment variables required to run the runtime. - /// Warning: be sure to update the environment variables when necessary if the image is altered. - pub env_vars: HashMap, -} - -impl Default for Config { - fn default() -> Self { - Self { - image_name: "confluentinc/cp-kafka".to_string(), - image_tag: "7.4.3".to_string(), - exposed_port: KAFKA_ADVERTISED_LISTENER_PORT, - // The runtime is safe to be used as long as this message is printed on stdout. - ready_conditions: vec![WaitFor::message_on_stdout( - "started (kafka.server.KafkaServer)", - )], - env_vars: build_env_vars( - ZOOKEEPER_PORT, - KAFKA_LISTENER_PORT, - KAFKA_ADVERTISED_LISTENER_PORT, - ), - } - } -} - -fn build_env_vars( - zookeeper_port: u16, - kafka_listener_port: u16, - kafka_advertised_listener_port: u16, -) -> HashMap { - [ - ( - "KAFKA_ZOOKEEPER_CONNECT".to_string(), - format!("localhost:{zookeeper_port}"), - ), - ( - "KAFKA_LISTENERS".to_string(), - format!("PLAINTEXT://0.0.0.0:{kafka_advertised_listener_port},BROKER://0.0.0.0:{kafka_listener_port}"), - ), - ( - "KAFKA_ADVERTISED_LISTENERS".to_string(), - format!("PLAINTEXT://localhost:{kafka_advertised_listener_port},BROKER://localhost:{kafka_listener_port}",), - ), - ( - "KAFKA_LISTENER_SECURITY_PROTOCOL_MAP".to_string(), - "BROKER:PLAINTEXT,PLAINTEXT:PLAINTEXT".to_string(), - ), - ( - "KAFKA_INTER_BROKER_LISTENER_NAME".to_string(), - "BROKER".to_string(), - ), - ("KAFKA_BROKER_ID".to_string(), "1".to_string()), - ( - "KAFKA_OFFSETS_TOPIC_REPLICATION_FACTOR".to_string(), - "1".to_string(), - ), - ] - .into() -} diff --git a/src/common/test-util/src/wal/kafka/image.rs b/src/common/test-util/src/wal/kafka/image.rs index 48a2449fab96..8e4462135f4b 100644 --- a/src/common/test-util/src/wal/kafka/image.rs +++ b/src/common/test-util/src/wal/kafka/image.rs @@ -12,11 +12,18 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::HashMap; + use testcontainers::core::{ContainerState, ExecCommand, WaitFor}; -use crate::wal::kafka::config::{ - Config, KAFKA_ADVERTISED_LISTENER_PORT, KAFKA_LISTENER_PORT, ZOOKEEPER_PORT, -}; +const IMAGE_NAME: &str = "confluentinc/cp-kafka"; +const IMAGE_TAG: &str = "7.4.3"; +/// Through which port the Zookeeper node listens for external traffics, e.g. traffics from the Kafka node. +const ZOOKEEPER_PORT: u16 = 2181; +/// Through which port the Kafka node listens for internal traffics, i.e. traffics between Kafka nodes in the same Kafka cluster. +const KAFKA_LISTENER_PORT: u16 = 19092; +/// Through which port the Kafka node listens for external traffics, e.g. traffics from Kafka clients. +const KAFKA_ADVERTISED_LISTENER_PORT: u16 = 9092; #[derive(Debug, Clone, Default)] pub struct ImageArgs; @@ -45,44 +52,56 @@ impl testcontainers::ImageArgs for ImageArgs { } } -#[derive(Default)] pub struct Image { - config: Config, + env_vars: HashMap, } -impl Image { - pub fn with_exposed_port(port: u16) -> Self { - let config = Config { - exposed_port: port, - ..Default::default() - }; - Self { config } +impl Default for Image { + fn default() -> Self { + Self { + env_vars: build_env_vars(), + } } } impl testcontainers::Image for Image { type Args = ImageArgs; + /// The name of the Kafka image hosted in the docker hub. fn name(&self) -> String { - self.config.image_name.clone() + IMAGE_NAME.to_string() } + /// The tag of the kafka image hosted in the docker hub. + /// Warning: please use a tag with long-term support. Do not use `latest` or any other tags that + /// the underlying image may suddenly change. fn tag(&self) -> String { - self.config.image_tag.clone() + IMAGE_TAG.to_string() } + /// The runtime is regarded ready to be used if all ready conditions are met. + /// Warning: be sure to update the conditions when necessary if the image is altered. fn ready_conditions(&self) -> Vec { - self.config.ready_conditions.clone() + vec![WaitFor::message_on_stdout( + "started (kafka.server.KafkaServer)", + )] } + /// The environment variables required to run the runtime. + /// Warning: be sure to update the environment variables when necessary if the image is altered. fn env_vars(&self) -> Box + '_> { - Box::new(self.config.env_vars.iter()) + Box::new(self.env_vars.iter()) } + /// The runtime is running in a docker container and has its own network. In order to be used by the host machine, + /// the runtime must expose an internal port. For e.g. assume the runtime has an internal port 9092, + /// and the `exposed_port` is set to 9092, then the host machine can get a mapped external port with + /// `container.get_host_port_ipv4(exposed_port)`. With the mapped port, the host machine could connect with the runtime. fn expose_ports(&self) -> Vec { - vec![self.config.exposed_port] + vec![KAFKA_ADVERTISED_LISTENER_PORT] } + /// Specifies a collection of commands to be executed when the container is started. fn exec_after_start(&self, cs: ContainerState) -> Vec { let mut commands = vec![]; let cmd = format!( @@ -101,6 +120,37 @@ impl testcontainers::Image for Image { } } +fn build_env_vars() -> HashMap { + [ + ( + "KAFKA_ZOOKEEPER_CONNECT".to_string(), + format!("localhost:{ZOOKEEPER_PORT}"), + ), + ( + "KAFKA_LISTENERS".to_string(), + format!("PLAINTEXT://0.0.0.0:{KAFKA_ADVERTISED_LISTENER_PORT},BROKER://0.0.0.0:{KAFKA_LISTENER_PORT}"), + ), + ( + "KAFKA_ADVERTISED_LISTENERS".to_string(), + format!("PLAINTEXT://localhost:{KAFKA_ADVERTISED_LISTENER_PORT},BROKER://localhost:{KAFKA_LISTENER_PORT}",), + ), + ( + "KAFKA_LISTENER_SECURITY_PROTOCOL_MAP".to_string(), + "BROKER:PLAINTEXT,PLAINTEXT:PLAINTEXT".to_string(), + ), + ( + "KAFKA_INTER_BROKER_LISTENER_NAME".to_string(), + "BROKER".to_string(), + ), + ("KAFKA_BROKER_ID".to_string(), "1".to_string()), + ( + "KAFKA_OFFSETS_TOPIC_REPLICATION_FACTOR".to_string(), + "1".to_string(), + ), + ] + .into() +} + #[cfg(test)] mod tests { use rskafka::chrono::{TimeZone, Utc}; @@ -109,15 +159,14 @@ mod tests { use rskafka::record::Record; use testcontainers::clients::Cli as DockerCli; - use crate::wal::kafka::config::KAFKA_ADVERTISED_LISTENER_PORT; - use crate::wal::kafka::image::Image; + use super::*; #[tokio::test] async fn test_image() { // Starts a Kafka container. let port = KAFKA_ADVERTISED_LISTENER_PORT; let docker = DockerCli::default(); - let container = docker.run(Image::with_exposed_port(port)); + let container = docker.run(Image::default()); // Creates a Kafka client. let broker_endpoints = vec![format!("127.0.0.1:{}", container.get_host_port_ipv4(port))]; From 1c504aa18065a5ffedc01963f9409d74eecdfc1a Mon Sep 17 00:00:00 2001 From: niebayes Date: Wed, 3 Jan 2024 11:56:57 +0800 Subject: [PATCH 20/30] doc: add doc example for kafka image --- src/common/test-util/src/wal/kafka/image.rs | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/src/common/test-util/src/wal/kafka/image.rs b/src/common/test-util/src/wal/kafka/image.rs index 8e4462135f4b..ff801ea3dfec 100644 --- a/src/common/test-util/src/wal/kafka/image.rs +++ b/src/common/test-util/src/wal/kafka/image.rs @@ -52,6 +52,17 @@ impl testcontainers::ImageArgs for ImageArgs { } } +/// # Example +/// ```rust +/// // Starts a Kafka container. +/// let port = KAFKA_ADVERTISED_LISTENER_PORT; +/// let docker = DockerCli::default(); +/// let container = docker.run(Image::default()); +/// // Gets the broker endpoints of the containerized Kafka node. +/// let broker_endpoints = vec![format!("127.0.0.1:{}", container.get_host_port_ipv4(port))]; +/// // Do something with the broker endpoints, for e.g. building a Kafka client. +/// let client = ClientBuilder::new(broker_endpoints).build().await.unwrap(); +/// ``` pub struct Image { env_vars: HashMap, } From fd31088526edada39d48e2eefc3cba1903a16c5c Mon Sep 17 00:00:00 2001 From: niebayes Date: Wed, 3 Jan 2024 12:01:10 +0800 Subject: [PATCH 21/30] chore: migrate kafka image to an indie PR --- Cargo.lock | 94 +-------- src/common/test-util/Cargo.toml | 1 - src/common/test-util/src/wal/kafka.rs | 1 - src/common/test-util/src/wal/kafka/image.rs | 219 -------------------- tests-integration/Cargo.toml | 1 - 5 files changed, 4 insertions(+), 312 deletions(-) delete mode 100644 src/common/test-util/src/wal/kafka/image.rs diff --git a/Cargo.lock b/Cargo.lock index 355b828073e9..4161fd5e1353 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -972,16 +972,6 @@ dependencies = [ "generic-array", ] -[[package]] -name = "bollard-stubs" -version = "1.42.0-rc.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ed59b5c00048f48d7af971b71f800fdf23e858844a6f9e4d32ca72e9399e7864" -dependencies = [ - "serde", - "serde_with 1.14.0", -] - [[package]] name = "borsh" version = "1.3.0" @@ -1640,7 +1630,7 @@ dependencies = [ "rskafka", "serde", "serde_json", - "serde_with 3.4.0", + "serde_with", "toml 0.8.8", ] @@ -1853,7 +1843,7 @@ dependencies = [ "rskafka", "serde", "serde_json", - "serde_with 3.4.0", + "serde_with", "snafu", "store-api", "strum 0.25.0", @@ -1993,7 +1983,6 @@ dependencies = [ "rskafka", "store-api", "tempfile", - "testcontainers", "tokio", ] @@ -2352,16 +2341,6 @@ dependencies = [ "memchr", ] -[[package]] -name = "darling" -version = "0.13.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a01d95850c592940db9b8194bc39f4bc0e89dee5c4265e4b1807c34a9aba453c" -dependencies = [ - "darling_core 0.13.4", - "darling_macro 0.13.4", -] - [[package]] name = "darling" version = "0.14.4" @@ -2382,20 +2361,6 @@ dependencies = [ "darling_macro 0.20.3", ] -[[package]] -name = "darling_core" -version = "0.13.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "859d65a907b6852c9361e3185c862aae7fafd2887876799fa55f5f99dc40d610" -dependencies = [ - "fnv", - "ident_case", - "proc-macro2", - "quote", - "strsim 0.10.0", - "syn 1.0.109", -] - [[package]] name = "darling_core" version = "0.14.4" @@ -2424,17 +2389,6 @@ dependencies = [ "syn 2.0.43", ] -[[package]] -name = "darling_macro" -version = "0.13.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9c972679f83bdf9c42bd905396b6c3588a843a17f0f16dfcfa3e2c5d57441835" -dependencies = [ - "darling_core 0.13.4", - "quote", - "syn 1.0.109", -] - [[package]] name = "darling_macro" version = "0.14.4" @@ -5044,7 +4998,7 @@ dependencies = [ "regex", "serde", "serde_json", - "serde_with 3.4.0", + "serde_with", "smallvec", "snafu", "store-api", @@ -8458,16 +8412,6 @@ dependencies = [ "serde", ] -[[package]] -name = "serde_with" -version = "1.14.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "678b5a069e50bf00ecd22d0cd8ddf7c236f68581b03db652061ed5eb13a312ff" -dependencies = [ - "serde", - "serde_with_macros 1.5.2", -] - [[package]] name = "serde_with" version = "3.4.0" @@ -8481,22 +8425,10 @@ dependencies = [ "indexmap 2.1.0", "serde", "serde_json", - "serde_with_macros 3.4.0", + "serde_with_macros", "time", ] -[[package]] -name = "serde_with_macros" -version = "1.5.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e182d6ec6f05393cc0e5ed1bf81ad6db3a8feedf8ee515ecdd369809bcce8082" -dependencies = [ - "darling 0.13.4", - "proc-macro2", - "quote", - "syn 1.0.109", -] - [[package]] name = "serde_with_macros" version = "3.4.0" @@ -9538,23 +9470,6 @@ version = "0.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3369f5ac52d5eb6ab48c6b4ffdc8efbcad6b89c765749064ba298f2c68a16a76" -[[package]] -name = "testcontainers" -version = "0.15.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f83d2931d7f521af5bae989f716c3fa43a6af9af7ec7a5e21b59ae40878cec00" -dependencies = [ - "bollard-stubs", - "futures", - "hex", - "hmac", - "log", - "rand", - "serde", - "serde_json", - "sha2", -] - [[package]] name = "tests-integration" version = "0.5.0" @@ -9618,7 +9533,6 @@ dependencies = [ "substrait 0.5.0", "table", "tempfile", - "testcontainers", "time", "tokio", "tokio-postgres", diff --git a/src/common/test-util/Cargo.toml b/src/common/test-util/Cargo.toml index 0884c626c2cf..be5b5a5275d8 100644 --- a/src/common/test-util/Cargo.toml +++ b/src/common/test-util/Cargo.toml @@ -13,5 +13,4 @@ rand.workspace = true rskafka.workspace = true store-api.workspace = true tempfile.workspace = true -testcontainers = "0.15.0" tokio.workspace = true diff --git a/src/common/test-util/src/wal/kafka.rs b/src/common/test-util/src/wal/kafka.rs index db9d49b4be29..728c6978757f 100644 --- a/src/common/test-util/src/wal/kafka.rs +++ b/src/common/test-util/src/wal/kafka.rs @@ -12,7 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -pub mod image; pub mod topic_decorator; use common_config::wal::KafkaWalTopic as Topic; diff --git a/src/common/test-util/src/wal/kafka/image.rs b/src/common/test-util/src/wal/kafka/image.rs deleted file mode 100644 index ff801ea3dfec..000000000000 --- a/src/common/test-util/src/wal/kafka/image.rs +++ /dev/null @@ -1,219 +0,0 @@ -// 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 std::collections::HashMap; - -use testcontainers::core::{ContainerState, ExecCommand, WaitFor}; - -const IMAGE_NAME: &str = "confluentinc/cp-kafka"; -const IMAGE_TAG: &str = "7.4.3"; -/// Through which port the Zookeeper node listens for external traffics, e.g. traffics from the Kafka node. -const ZOOKEEPER_PORT: u16 = 2181; -/// Through which port the Kafka node listens for internal traffics, i.e. traffics between Kafka nodes in the same Kafka cluster. -const KAFKA_LISTENER_PORT: u16 = 19092; -/// Through which port the Kafka node listens for external traffics, e.g. traffics from Kafka clients. -const KAFKA_ADVERTISED_LISTENER_PORT: u16 = 9092; - -#[derive(Debug, Clone, Default)] -pub struct ImageArgs; - -impl testcontainers::ImageArgs for ImageArgs { - fn into_iterator(self) -> Box> { - Box::new( - vec![ - "/bin/bash".to_string(), - "-c".to_string(), - format!( - r#" - echo 'clientPort={}' > zookeeper.properties; - echo 'dataDir=/var/lib/zookeeper/data' >> zookeeper.properties; - echo 'dataLogDir=/var/lib/zookeeper/log' >> zookeeper.properties; - zookeeper-server-start zookeeper.properties & - . /etc/confluent/docker/bash-config && - /etc/confluent/docker/configure && - /etc/confluent/docker/launch - "#, - ZOOKEEPER_PORT, - ), - ] - .into_iter(), - ) - } -} - -/// # Example -/// ```rust -/// // Starts a Kafka container. -/// let port = KAFKA_ADVERTISED_LISTENER_PORT; -/// let docker = DockerCli::default(); -/// let container = docker.run(Image::default()); -/// // Gets the broker endpoints of the containerized Kafka node. -/// let broker_endpoints = vec![format!("127.0.0.1:{}", container.get_host_port_ipv4(port))]; -/// // Do something with the broker endpoints, for e.g. building a Kafka client. -/// let client = ClientBuilder::new(broker_endpoints).build().await.unwrap(); -/// ``` -pub struct Image { - env_vars: HashMap, -} - -impl Default for Image { - fn default() -> Self { - Self { - env_vars: build_env_vars(), - } - } -} - -impl testcontainers::Image for Image { - type Args = ImageArgs; - - /// The name of the Kafka image hosted in the docker hub. - fn name(&self) -> String { - IMAGE_NAME.to_string() - } - - /// The tag of the kafka image hosted in the docker hub. - /// Warning: please use a tag with long-term support. Do not use `latest` or any other tags that - /// the underlying image may suddenly change. - fn tag(&self) -> String { - IMAGE_TAG.to_string() - } - - /// The runtime is regarded ready to be used if all ready conditions are met. - /// Warning: be sure to update the conditions when necessary if the image is altered. - fn ready_conditions(&self) -> Vec { - vec![WaitFor::message_on_stdout( - "started (kafka.server.KafkaServer)", - )] - } - - /// The environment variables required to run the runtime. - /// Warning: be sure to update the environment variables when necessary if the image is altered. - fn env_vars(&self) -> Box + '_> { - Box::new(self.env_vars.iter()) - } - - /// The runtime is running in a docker container and has its own network. In order to be used by the host machine, - /// the runtime must expose an internal port. For e.g. assume the runtime has an internal port 9092, - /// and the `exposed_port` is set to 9092, then the host machine can get a mapped external port with - /// `container.get_host_port_ipv4(exposed_port)`. With the mapped port, the host machine could connect with the runtime. - fn expose_ports(&self) -> Vec { - vec![KAFKA_ADVERTISED_LISTENER_PORT] - } - - /// Specifies a collection of commands to be executed when the container is started. - fn exec_after_start(&self, cs: ContainerState) -> Vec { - let mut commands = vec![]; - let cmd = format!( - "kafka-configs --alter --bootstrap-server 0.0.0.0:{} --entity-type brokers --entity-name 1 --add-config advertised.listeners=[PLAINTEXT://127.0.0.1:{},BROKER://localhost:9092]", - KAFKA_LISTENER_PORT, - cs.host_port_ipv4(KAFKA_ADVERTISED_LISTENER_PORT), - ); - let ready_conditions = vec![WaitFor::message_on_stdout( - "Checking need to trigger auto leader balancing", - )]; - commands.push(ExecCommand { - cmd, - ready_conditions, - }); - commands - } -} - -fn build_env_vars() -> HashMap { - [ - ( - "KAFKA_ZOOKEEPER_CONNECT".to_string(), - format!("localhost:{ZOOKEEPER_PORT}"), - ), - ( - "KAFKA_LISTENERS".to_string(), - format!("PLAINTEXT://0.0.0.0:{KAFKA_ADVERTISED_LISTENER_PORT},BROKER://0.0.0.0:{KAFKA_LISTENER_PORT}"), - ), - ( - "KAFKA_ADVERTISED_LISTENERS".to_string(), - format!("PLAINTEXT://localhost:{KAFKA_ADVERTISED_LISTENER_PORT},BROKER://localhost:{KAFKA_LISTENER_PORT}",), - ), - ( - "KAFKA_LISTENER_SECURITY_PROTOCOL_MAP".to_string(), - "BROKER:PLAINTEXT,PLAINTEXT:PLAINTEXT".to_string(), - ), - ( - "KAFKA_INTER_BROKER_LISTENER_NAME".to_string(), - "BROKER".to_string(), - ), - ("KAFKA_BROKER_ID".to_string(), "1".to_string()), - ( - "KAFKA_OFFSETS_TOPIC_REPLICATION_FACTOR".to_string(), - "1".to_string(), - ), - ] - .into() -} - -#[cfg(test)] -mod tests { - use rskafka::chrono::{TimeZone, Utc}; - use rskafka::client::partition::UnknownTopicHandling; - use rskafka::client::ClientBuilder; - use rskafka::record::Record; - use testcontainers::clients::Cli as DockerCli; - - use super::*; - - #[tokio::test] - async fn test_image() { - // Starts a Kafka container. - let port = KAFKA_ADVERTISED_LISTENER_PORT; - let docker = DockerCli::default(); - let container = docker.run(Image::default()); - - // Creates a Kafka client. - let broker_endpoints = vec![format!("127.0.0.1:{}", container.get_host_port_ipv4(port))]; - let client = ClientBuilder::new(broker_endpoints).build().await.unwrap(); - - // Creates a topic. - let topic = "test_topic"; - client - .controller_client() - .unwrap() - .create_topic(topic, 1, 1, 500) - .await - .unwrap(); - - // Produces a record. - let partition_client = client - .partition_client(topic, 0, UnknownTopicHandling::Error) - .await - .unwrap(); - let produced = vec![Record { - key: Some(b"111".to_vec()), - value: Some(b"222".to_vec()), - timestamp: Utc.timestamp_millis_opt(42).unwrap(), - headers: Default::default(), - }]; - let offset = partition_client - .produce(produced.clone(), Default::default()) - .await - .unwrap()[0]; - - // Consumes the record. - let consumed = partition_client - .fetch_records(offset, 1..4096, 500) - .await - .unwrap() - .0; - assert_eq!(produced[0], consumed[0].record); - } -} diff --git a/tests-integration/Cargo.toml b/tests-integration/Cargo.toml index 377c428c2a55..8d311d3f7ea2 100644 --- a/tests-integration/Cargo.toml +++ b/tests-integration/Cargo.toml @@ -64,7 +64,6 @@ sqlx = { version = "0.6", features = [ substrait.workspace = true table.workspace = true tempfile.workspace = true -testcontainers = "0.15.0" time = "0.3" tokio.workspace = true tonic.workspace = true From d38a16f63a75ddb829e3778edb4d17c04f993c9c Mon Sep 17 00:00:00 2001 From: niebayes Date: Fri, 5 Jan 2024 21:08:10 +0800 Subject: [PATCH 22/30] fix: CR --- Cargo.lock | 2 -- .../src/wal/kafka/topic_decorator.rs | 28 ++++--------------- src/log-store/src/kafka/log_store.rs | 23 ++------------- src/log-store/src/test_util/kafka.rs | 11 ++++---- tests-integration/Cargo.toml | 2 -- 5 files changed, 13 insertions(+), 53 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 52c738521745..2cbb30eba7d2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9506,7 +9506,6 @@ dependencies = [ "frontend", "futures", "itertools 0.10.5", - "log-store", "meta-client", "meta-srv", "mysql_async", @@ -9520,7 +9519,6 @@ dependencies = [ "prost 0.12.3", "query", "rand", - "rskafka", "rstest", "rstest_reuse", "script", diff --git a/src/common/test-util/src/wal/kafka/topic_decorator.rs b/src/common/test-util/src/wal/kafka/topic_decorator.rs index 60f4ac1c0993..afa7719d20a2 100644 --- a/src/common/test-util/src/wal/kafka/topic_decorator.rs +++ b/src/common/test-util/src/wal/kafka/topic_decorator.rs @@ -12,9 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::cell::RefCell; -use std::collections::HashSet; - use common_config::wal::KafkaWalTopic as Topic; /// Things need to bo inserted at the front or the back of the topic. @@ -45,8 +42,6 @@ pub struct TopicDecorator { prefix: Affix, /// A suffix to be inserted at the back of each topic. suffix: Affix, - /// Topics built so far. Used to filter out duplicate topics. - created: RefCell>, } impl Default for TopicDecorator { @@ -54,7 +49,6 @@ impl Default for TopicDecorator { Self { prefix: Affix::Nothing, suffix: Affix::Nothing, - created: RefCell::new(HashSet::with_capacity(256)), } } } @@ -72,23 +66,11 @@ impl TopicDecorator { /// Builds a topic by inserting a prefix and a suffix into the given topic. pub fn decorate(&self, topic: &str) -> Topic { - const ITERS: usize = 24; - for _ in 0..ITERS { - let decorated = format!( - "{}_{}_{}", - self.prefix.to_string(), - topic, - self.suffix.to_string() - ); - let mut created_topics = self.created.borrow_mut(); - if !created_topics.contains(&decorated) { - created_topics.insert(decorated.clone()); - return decorated; - } - } - unreachable!( - "Topic decoration should be completed within iterations {}", - ITERS + format!( + "{}_{}_{}", + self.prefix.to_string(), + topic, + self.suffix.to_string() ) } } diff --git a/src/log-store/src/kafka/log_store.rs b/src/log-store/src/kafka/log_store.rs index 7f1ec863cd75..2eda64473148 100644 --- a/src/log-store/src/kafka/log_store.rs +++ b/src/log-store/src/kafka/log_store.rs @@ -15,7 +15,7 @@ use std::collections::HashMap; use std::sync::Arc; -use common_config::wal::{KafkaConfig, KafkaWalTopic as Topic, WalOptions}; +use common_config::wal::{KafkaConfig, WalOptions}; use common_telemetry::{debug, warn}; use futures_util::StreamExt; use rskafka::client::consumer::{StartOffset, StreamConsumerBuilder}; @@ -48,26 +48,6 @@ impl KafkaLogStore { config: config.clone(), }) } - - /// Gets the end offset of the last record in a Kafka topic. - /// Warning: this method is intended to be used only in testing. - // TODO(niebayes): use this to test that the initial offset is 1 for a Kafka log store in that - // a no-op record is successfully appended into each topic. - #[allow(unused)] - pub async fn get_offset(&self, topic: &Topic) -> EntryId { - let client = self - .client_manager - .get_or_insert(topic) - .await - .unwrap() - .raw_client; - client - .get_offset(OffsetAt::Latest) - .await - .map(TryInto::try_into) - .unwrap() - .unwrap() - } } #[async_trait::async_trait] @@ -287,6 +267,7 @@ impl LogStore for KafkaLogStore { #[cfg(test)] mod tests { + use common_config::wal::KafkaWalTopic as Topic; use common_test_util::get_broker_endpoints; use common_test_util::wal::kafka::topic_decorator::{Affix, TopicDecorator}; use common_test_util::wal::kafka::{create_topics, BROKER_ENDPOINTS_KEY}; diff --git a/src/log-store/src/test_util/kafka.rs b/src/log-store/src/test_util/kafka.rs index 7fc25e477c0a..eaf6a0e5a38b 100644 --- a/src/log-store/src/test_util/kafka.rs +++ b/src/log-store/src/test_util/kafka.rs @@ -12,8 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::cell::RefCell; use std::sync::atomic::{AtomicU64 as AtomicEntryId, Ordering}; +use std::sync::Mutex; use rand::rngs::ThreadRng; use rand::seq::SliceRandom; @@ -30,7 +30,7 @@ pub struct EntryBuilder { next_entry_id: AtomicEntryId, /// A generator for supporting random data generation. /// Wrapped with Mutex> to provide interior mutability. - rng: RefCell, + rng: Mutex>, /// The data pool from which random data is constructed. data_pool: Vec, } @@ -47,7 +47,7 @@ impl EntryBuilder { Self { ns, next_entry_id: AtomicEntryId::new(0), - rng: RefCell::new(thread_rng()), + rng: Mutex::new(Some(thread_rng())), data_pool, } } @@ -85,10 +85,11 @@ impl EntryBuilder { } fn make_random_data(&self) -> Vec { - let mut rng = self.rng.borrow_mut(); + let mut guard = self.rng.lock().unwrap(); + let rng = guard.as_mut().unwrap(); let amount = rng.gen_range(0..self.data_pool.len()); self.data_pool - .choose_multiple(&mut *rng, amount) + .choose_multiple(rng, amount) .copied() .collect() } diff --git a/tests-integration/Cargo.toml b/tests-integration/Cargo.toml index 8d311d3f7ea2..1ba6f8e05d36 100644 --- a/tests-integration/Cargo.toml +++ b/tests-integration/Cargo.toml @@ -34,7 +34,6 @@ datatypes.workspace = true dotenv = "0.15" frontend = { workspace = true, features = ["testing"] } futures.workspace = true -log-store.workspace = true meta-client.workspace = true meta-srv = { workspace = true, features = ["mock"] } mysql_async = { version = "0.33", default-features = false, features = [ @@ -45,7 +44,6 @@ once_cell.workspace = true operator.workspace = true query.workspace = true rand.workspace = true -rskafka.workspace = true rstest = "0.17" rstest_reuse = "0.5" secrecy = "0.8" From cf68a87539e55ec1020d81b7414fe6b3008e588a Mon Sep 17 00:00:00 2001 From: niebayes Date: Fri, 5 Jan 2024 21:35:42 +0800 Subject: [PATCH 23/30] fix: CR --- Cargo.lock | 1 + src/common/meta/Cargo.toml | 2 +- src/log-store/Cargo.toml | 1 + src/log-store/src/test_util.rs | 1 + src/log-store/src/test_util/kafka.rs | 17 ++--------------- 5 files changed, 6 insertions(+), 16 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index f13c198f91ae..2eafe6763912 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4521,6 +4521,7 @@ dependencies = [ "protobuf-build", "raft-engine", "rand", + "rand_distr", "rskafka", "serde", "serde_json", diff --git a/src/common/meta/Cargo.toml b/src/common/meta/Cargo.toml index 5f34915f9728..bcb562437134 100644 --- a/src/common/meta/Cargo.toml +++ b/src/common/meta/Cargo.toml @@ -23,7 +23,6 @@ common-procedure.workspace = true common-recordbatch.workspace = true common-runtime.workspace = true common-telemetry.workspace = true -common-test-util.workspace = true common-time.workspace = true datatypes.workspace = true derive_builder.workspace = true @@ -50,5 +49,6 @@ tonic.workspace = true [dev-dependencies] chrono.workspace = true common-procedure = { workspace = true, features = ["testing"] } +common-test-util.workspace = true datatypes.workspace = true hyper = { version = "0.14", features = ["full"] } diff --git a/src/log-store/Cargo.toml b/src/log-store/Cargo.toml index 3b36d00c7bc5..15ef7b4ee66a 100644 --- a/src/log-store/Cargo.toml +++ b/src/log-store/Cargo.toml @@ -29,6 +29,7 @@ itertools.workspace = true protobuf = { version = "2", features = ["bytes"] } raft-engine.workspace = true rand.workspace = true +rand_distr = "0.4" rskafka.workspace = true serde.workspace = true serde_json.workspace = true diff --git a/src/log-store/src/test_util.rs b/src/log-store/src/test_util.rs index 52ffcca39370..ce5ba3eb854f 100644 --- a/src/log-store/src/test_util.rs +++ b/src/log-store/src/test_util.rs @@ -12,5 +12,6 @@ // See the License for the specific language governing permissions and // limitations under the License. +#[cfg(test)] pub mod kafka; pub mod log_store_util; diff --git a/src/log-store/src/test_util/kafka.rs b/src/log-store/src/test_util/kafka.rs index eaf6a0e5a38b..67a71b23004c 100644 --- a/src/log-store/src/test_util/kafka.rs +++ b/src/log-store/src/test_util/kafka.rs @@ -15,8 +15,8 @@ use std::sync::atomic::{AtomicU64 as AtomicEntryId, Ordering}; use std::sync::Mutex; +use rand::distributions::Alphanumeric; use rand::rngs::ThreadRng; -use rand::seq::SliceRandom; use rand::{thread_rng, Rng}; use store_api::logstore::EntryId; @@ -31,24 +31,15 @@ pub struct EntryBuilder { /// A generator for supporting random data generation. /// Wrapped with Mutex> to provide interior mutability. rng: Mutex>, - /// The data pool from which random data is constructed. - data_pool: Vec, } impl EntryBuilder { /// Creates an EntryBuilder for the given namespace. pub fn new(ns: NamespaceImpl) -> Self { - // Makes a data pool with alphabets and numbers. - let data_pool = ('a'..='z') - .chain('A'..='Z') - .chain((0..=9).map(|digit| char::from_digit(digit, 10).unwrap())) - .map(|c| c as u8) - .collect::>(); Self { ns, next_entry_id: AtomicEntryId::new(0), rng: Mutex::new(Some(thread_rng())), - data_pool, } } @@ -87,11 +78,7 @@ impl EntryBuilder { fn make_random_data(&self) -> Vec { let mut guard = self.rng.lock().unwrap(); let rng = guard.as_mut().unwrap(); - let amount = rng.gen_range(0..self.data_pool.len()); - self.data_pool - .choose_multiple(rng, amount) - .copied() - .collect() + (0..42).map(|_| rng.sample(Alphanumeric)).collect() } } From 5de92f2ae0cbb27e74daa64ff48063ab50c8e5a4 Mon Sep 17 00:00:00 2001 From: niebayes Date: Sat, 6 Jan 2024 15:59:18 +0800 Subject: [PATCH 24/30] fix: test --- src/common/config/src/wal/kafka.rs | 1 - src/common/meta/Cargo.toml | 1 + .../meta/src/wal/kafka/topic_manager.rs | 2 +- src/log-store/src/kafka/log_store.rs | 115 +++++++++++++----- src/log-store/src/kafka/util/record.rs | 1 + 5 files changed, 88 insertions(+), 32 deletions(-) diff --git a/src/common/config/src/wal/kafka.rs b/src/common/config/src/wal/kafka.rs index 884e3f3f0a99..e7179f96f106 100644 --- a/src/common/config/src/wal/kafka.rs +++ b/src/common/config/src/wal/kafka.rs @@ -40,7 +40,6 @@ pub struct KafkaConfig { pub broker_endpoints: Vec, /// 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, diff --git a/src/common/meta/Cargo.toml b/src/common/meta/Cargo.toml index bcb562437134..e32443ea26f6 100644 --- a/src/common/meta/Cargo.toml +++ b/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 diff --git a/src/common/meta/src/wal/kafka/topic_manager.rs b/src/common/meta/src/wal/kafka/topic_manager.rs index 4307d5b9f8b1..0a3ea92d3c6d 100644 --- a/src/common/meta/src/wal/kafka/topic_manager.rs +++ b/src/common/meta/src/wal/kafka/topic_manager.rs @@ -167,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, diff --git a/src/log-store/src/kafka/log_store.rs b/src/log-store/src/kafka/log_store.rs index b6d64aef8575..736ece995d02 100644 --- a/src/log-store/src/kafka/log_store.rs +++ b/src/log-store/src/kafka/log_store.rs @@ -282,6 +282,7 @@ fn check_termination( #[cfg(test)] mod tests { + use common_base::readable_size::ReadableSize; use common_config::wal::KafkaWalTopic as Topic; use common_test_util::get_broker_endpoints; use common_test_util::wal::kafka::topic_decorator::{Affix, TopicDecorator}; @@ -296,6 +297,7 @@ mod tests { ns: NamespaceImpl, entry_builder: EntryBuilder, expected: Vec, + flushed_entry_id: EntryId, } /// Prepares for a test in that a log store is constructed and a collection of topics is created. @@ -308,25 +310,44 @@ mod tests { let config = KafkaConfig { broker_endpoints, + max_batch_size: ReadableSize::kb(32), ..Default::default() }; let logstore = KafkaLogStore::try_new(&config).await.unwrap(); + // Appends a no-op record to each topic. + for topic in topics.iter() { + let last_entry_id = logstore + .append(EntryImpl { + data: vec![], + id: 0, + ns: new_namespace(topic, 0), + }) + .await + .unwrap() + .last_entry_id; + assert_eq!(last_entry_id, 0); + } + (logstore, topics) } /// Builds a context for each region. - fn build_contexts(num_regions: usize, topics: &[Topic]) -> Vec { + fn build_contexts(num_regions: usize, topics: &[Topic]) -> HashMap { (0..num_regions) .map(|i| { let topic = &topics[i % topics.len()]; let ns = new_namespace(topic, i as u64); let entry_builder = EntryBuilder::new(ns.clone()); - RegionContext { - ns, - entry_builder, - expected: Vec::new(), - } + ( + i as u64, + RegionContext { + ns, + entry_builder, + expected: Vec::new(), + flushed_entry_id: 0, + }, + ) }) .collect() } @@ -334,39 +355,52 @@ mod tests { /// Creates a vector containing indexes of all regions if the `all` is true. /// Otherwise, creates a subset of the indexes. The cardinality of the subset /// is nearly a quarter of that of the universe set. - fn all_or_subset(all: bool, num_regions: usize) -> Vec { + fn all_or_subset(all: bool, num_regions: usize) -> Vec { assert!(num_regions > 0); let amount = if all { num_regions } else { (num_regions / 4).max(1) }; - (0..num_regions).choose_multiple(&mut rand::thread_rng(), amount) + (0..num_regions as u64).choose_multiple(&mut rand::thread_rng(), amount) } - /// Builds entries for regions specified with `which`. - /// For example, assume `which[i] = x`, then entries for region with id = x will be built. - fn build_entries(region_contexts: &mut [RegionContext], which: Vec) -> Vec { - let mut entries = Vec::with_capacity(which.len()); - for i in which { - let ctx = &mut region_contexts[i]; - ctx.expected = entries_with_random_data(25, &ctx.entry_builder); - entries.push(ctx.expected.clone()); + /// Builds entries for regions specified by `which`. Builds large entries if `large` is true. + /// Returns the aggregated entries. + fn build_entries( + region_contexts: &mut HashMap, + which: &[u64], + large: bool, + ) -> Vec { + let mut aggregated = Vec::with_capacity(which.len()); + for region_id in which { + let ctx = region_contexts.get_mut(region_id).unwrap(); + // Builds entries for the region. + ctx.expected = if !large { + entries_with_random_data(3, &ctx.entry_builder) + } else { + // Builds a large entry of size 256KB which is way greater than the configured `max_batch_size` which is 32KB. + let large_entry = ctx.entry_builder.with_data([b'1'; 256 * 1024]); + vec![large_entry] + }; + // Aggregates entries of all regions. + aggregated.push(ctx.expected.clone()); } - entries.into_iter().flatten().collect() + aggregated.into_iter().flatten().collect() } /// Reads entries for regions and checks for each region that the gotten entries are identical with the expected ones. async fn check_entries( - region_contexts: &[RegionContext], - last_entry_ids: HashMap, + region_contexts: &HashMap, logstore: &KafkaLogStore, - which: Vec, + which: &[u64], ) { - for i in which { - let ctx = ®ion_contexts[i]; - let start_offset = last_entry_ids[&ctx.ns.region_id]; - let stream = logstore.read(&ctx.ns, start_offset).await.unwrap(); + for region_id in which { + let ctx = ®ion_contexts[region_id]; + let stream = logstore + .read(&ctx.ns, ctx.flushed_entry_id + 1) + .await + .unwrap(); let got = stream .collect::>() .await @@ -377,6 +411,17 @@ mod tests { } } + /// Simulates a flush for regions. + fn flush( + region_contexts: &mut HashMap, + last_entry_ids: &HashMap, + ) { + for (region_id, last_entry_id) in last_entry_ids { + let ctx = region_contexts.get_mut(region_id).unwrap(); + ctx.flushed_entry_id = *last_entry_id; + } + } + /// Starts a test with: /// * `test_name` - The name of the test. /// * `num_topics` - Number of topics to be created in the preparation phase. @@ -384,41 +429,44 @@ mod tests { /// * `num_appends` - Number of append operations to be performed. /// * `all` - All regions will be involved in an append operation if `all` is true. Otherwise, /// an append operation will only randomly choose a subset of regions. + /// * `large` - Builds large entries for each region is `large` is true. async fn test_with( test_name: &str, num_topics: usize, num_regions: usize, num_appends: usize, all: bool, + large: bool, ) { let (logstore, topics) = prepare(test_name, num_topics).await; let mut region_contexts = build_contexts(num_regions, &topics); for _ in 0..num_appends { let which = all_or_subset(all, num_regions); - let entries = build_entries(&mut region_contexts, which.clone()); + let entries = build_entries(&mut region_contexts, &which, large); let last_entry_ids = logstore.append_batch(entries).await.unwrap().last_entry_ids; - check_entries(®ion_contexts, last_entry_ids, &logstore, which).await; + check_entries(®ion_contexts, &logstore, &which).await; + flush(&mut region_contexts, &last_entry_ids); } } /// Appends entries for one region and checks all entries can be read successfully. #[tokio::test] async fn test_one_region() { - test_with("test_one_region", 1, 1, 1, true).await; + test_with("test_one_region", 1, 1, 1, true, false).await; } /// Appends entries for multiple regions and checks entries for each region can be read successfully. /// A topic is assigned only a single region. #[tokio::test] async fn test_multi_regions_disjoint() { - test_with("test_multi_regions_disjoint", 10, 10, 1, true).await; + test_with("test_multi_regions_disjoint", 5, 5, 1, true, false).await; } /// Appends entries for multiple regions and checks entries for each region can be read successfully. /// A topic is assigned multiple regions. #[tokio::test] async fn test_multi_regions_overlapped() { - test_with("test_multi_regions_overlapped", 10, 50, 1, true).await; + test_with("test_multi_regions_overlapped", 5, 20, 1, true, false).await; } /// Appends entries for multiple regions and checks entries for each region can be read successfully. @@ -426,6 +474,13 @@ mod tests { /// Each append operation will only append entries for a subset of randomly chosen regions. #[tokio::test] async fn test_multi_appends() { - test_with("test_multi_appends", 32, 256, 16, false).await; + test_with("test_multi_appends", 5, 20, 3, false, false).await; + } + + /// Appends large entries for multiple regions and checks entries for each region can be read successfully. + /// A topic may be assigned multiple regions. + #[tokio::test] + async fn test_append_large_entries() { + test_with("test_append_large_entries", 5, 20, 3, true, true).await; } } diff --git a/src/log-store/src/kafka/util/record.rs b/src/log-store/src/kafka/util/record.rs index 7d45165514a7..0193e234ad0a 100644 --- a/src/log-store/src/kafka/util/record.rs +++ b/src/log-store/src/kafka/util/record.rs @@ -149,6 +149,7 @@ impl RecordProducer { } /// Produces the buffered entries to Kafka sever. Those entries may span several Kafka records. + /// Returns the offset of the last successfully produced record. pub(crate) async fn produce(self, client_manager: &ClientManagerRef) -> Result { ensure!(!self.entries.is_empty(), EmptyEntriesSnafu); From ae057d022d1fb61ed4bc9366415ce50f21377356 Mon Sep 17 00:00:00 2001 From: niebayes Date: Mon, 8 Jan 2024 11:41:11 +0800 Subject: [PATCH 25/30] fix: CR --- Cargo.lock | 3 +- src/common/meta/Cargo.toml | 2 +- src/common/meta/src/wal/kafka.rs | 2 + src/common/meta/src/wal/kafka/test_util.rs | 32 +++++ .../meta/src/wal/kafka/topic_manager.rs | 11 +- src/common/meta/src/wal/options_allocator.rs | 13 +- src/common/test-util/src/lib.rs | 1 - src/common/test-util/src/wal.rs | 15 --- src/common/test-util/src/wal/kafka.rs | 66 ---------- .../src/wal/kafka/topic_decorator.rs | 76 ------------ src/log-store/Cargo.toml | 1 + src/log-store/src/kafka/client_manager.rs | 82 ++++++------- src/log-store/src/kafka/log_store.rs | 116 ++++++++---------- src/log-store/src/test_util/kafka.rs | 59 +++++++-- 14 files changed, 184 insertions(+), 295 deletions(-) create mode 100644 src/common/meta/src/wal/kafka/test_util.rs delete mode 100644 src/common/test-util/src/wal.rs delete mode 100644 src/common/test-util/src/wal/kafka.rs delete mode 100644 src/common/test-util/src/wal/kafka/topic_decorator.rs diff --git a/Cargo.lock b/Cargo.lock index 2eafe6763912..f9208f025fae 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1828,7 +1828,6 @@ dependencies = [ "common-recordbatch", "common-runtime", "common-telemetry", - "common-test-util", "common-time", "datatypes", "derive_builder 0.12.0", @@ -1852,6 +1851,7 @@ dependencies = [ "tokio", "toml 0.8.8", "tonic 0.10.2", + "uuid", ] [[package]] @@ -4529,6 +4529,7 @@ dependencies = [ "store-api", "tokio", "tokio-util", + "uuid", ] [[package]] diff --git a/src/common/meta/Cargo.toml b/src/common/meta/Cargo.toml index e32443ea26f6..724b508ebc6a 100644 --- a/src/common/meta/Cargo.toml +++ b/src/common/meta/Cargo.toml @@ -50,6 +50,6 @@ tonic.workspace = true [dev-dependencies] chrono.workspace = true common-procedure = { workspace = true, features = ["testing"] } -common-test-util.workspace = true datatypes.workspace = true hyper = { version = "0.14", features = ["full"] } +uuid.workspace = true diff --git a/src/common/meta/src/wal/kafka.rs b/src/common/meta/src/wal/kafka.rs index 703dfa7e3de0..625b45fce22f 100644 --- a/src/common/meta/src/wal/kafka.rs +++ b/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; diff --git a/src/common/meta/src/wal/kafka/test_util.rs b/src/common/meta/src/wal/kafka/test_util.rs new file mode 100644 index 000000000000..81014bc56cbe --- /dev/null +++ b/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::>(); + assert!(!broker_endpoints.is_empty()); + broker_endpoints + }}; +} diff --git a/src/common/meta/src/wal/kafka/topic_manager.rs b/src/common/meta/src/wal/kafka/topic_manager.rs index 0a3ea92d3c6d..9765fe14e2a6 100644 --- a/src/common/meta/src/wal/kafka/topic_manager.rs +++ b/src/common/meta/src/wal/kafka/topic_manager.rs @@ -243,11 +243,9 @@ mod tests { use chrono::format::Fixed; use common_telemetry::info; - use common_test_util::get_broker_endpoints; - use common_test_util::wal::kafka::topic_decorator::{Affix, TopicDecorator}; - use common_test_util::wal::kafka::{create_topics, BROKER_ENDPOINTS_KEY}; use super::*; + use crate::get_broker_endpoints; use crate::kv_backend::memory::MemoryKvBackend; // Tests that topics can be successfully persisted into the kv backend and can be successfully restored from the kv backend. @@ -278,13 +276,10 @@ mod tests { /// Tests that the topic manager could allocate topics correctly. #[tokio::test] async fn test_alloc_topics() { - let broker_endpoints = get_broker_endpoints!(BROKER_ENDPOINTS_KEY); + let broker_endpoints = get_broker_endpoints!(); // Constructs topics that should be created. - let decorator = TopicDecorator::default() - .with_prefix(Affix::Fixed("test_alloc_topics".to_string())) - .with_suffix(Affix::TimeNow); let topics = (0..256) - .map(|i| decorator.decorate(&format!("topic_{i}"))) + .map(|i| format!("test_alloc_topics_{}_{}", i, uuid::Uuid::new_v4())) .collect::>(); // Creates a topic manager. diff --git a/src/common/meta/src/wal/options_allocator.rs b/src/common/meta/src/wal/options_allocator.rs index 4c85d987889a..c9df0cf44bcb 100644 --- a/src/common/meta/src/wal/options_allocator.rs +++ b/src/common/meta/src/wal/options_allocator.rs @@ -105,11 +105,8 @@ pub fn allocate_region_wal_options( #[cfg(test)] mod tests { - use common_test_util::get_broker_endpoints; - use common_test_util::wal::kafka::topic_decorator::{Affix, TopicDecorator}; - use common_test_util::wal::kafka::BROKER_ENDPOINTS_KEY; - 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; @@ -137,13 +134,9 @@ mod tests { // 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!(BROKER_ENDPOINTS_KEY); - // Constructs topics that should be created. - let decorator = TopicDecorator::default() - .with_prefix(Affix::Fixed("test_allocator_with_kafka".to_string())) - .with_suffix(Affix::TimeNow); + let broker_endpoints = get_broker_endpoints!(); let topics = (0..256) - .map(|i| decorator.decorate(&format!("topic_{i}"))) + .map(|i| format!("test_allocator_with_kafka_{}_{}", i, uuid::Uuid::new_v4())) .collect::>(); // Creates a topic manager. diff --git a/src/common/test-util/src/lib.rs b/src/common/test-util/src/lib.rs index f4a9c747e518..ef6ff4696861 100644 --- a/src/common/test-util/src/lib.rs +++ b/src/common/test-util/src/lib.rs @@ -20,7 +20,6 @@ use std::sync::LazyLock; pub mod ports; pub mod temp_dir; -pub mod wal; // Rust is working on an env possibly named `CARGO_WORKSPACE_DIR` to find the root path to the // workspace, see https://github.com/rust-lang/cargo/issues/3946. diff --git a/src/common/test-util/src/wal.rs b/src/common/test-util/src/wal.rs deleted file mode 100644 index 28c04633b640..000000000000 --- a/src/common/test-util/src/wal.rs +++ /dev/null @@ -1,15 +0,0 @@ -// 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. - -pub mod kafka; diff --git a/src/common/test-util/src/wal/kafka.rs b/src/common/test-util/src/wal/kafka.rs deleted file mode 100644 index 728c6978757f..000000000000 --- a/src/common/test-util/src/wal/kafka.rs +++ /dev/null @@ -1,66 +0,0 @@ -// 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. - -pub mod topic_decorator; - -use common_config::wal::KafkaWalTopic as Topic; -use rskafka::client::ClientBuilder; - -use crate::wal::kafka::topic_decorator::TopicDecorator; - -pub const BROKER_ENDPOINTS_KEY: &str = "GT_KAFKA_ENDPOINTS"; - -/// Gets broker endpoints from environment variables with the given key. -/// Returns ["localhost:9092"] if no environment variables set for broker endpoints. -#[macro_export] -macro_rules! get_broker_endpoints { - ($key:expr) => {{ - let broker_endpoints = std::env::var($key) - .unwrap_or("localhost:9092".to_string()) - .split(',') - .map(ToString::to_string) - .collect::>(); - assert!(!broker_endpoints.is_empty()); - broker_endpoints - }}; -} - -/// Creates `num_topiocs` number of topics from the seed topic which are going to be decorated with the given TopicDecorator. -/// A default seed `topic` will be used if the provided seed is None. -pub async fn create_topics( - num_topics: usize, - decorator: TopicDecorator, - broker_endpoints: &[String], - seed: Option<&str>, -) -> Vec { - assert!(!broker_endpoints.is_empty()); - - let client = ClientBuilder::new(broker_endpoints.to_vec()) - .build() - .await - .unwrap(); - let ctrl_client = client.controller_client().unwrap(); - - let seed = seed.unwrap_or("topic"); - let (topics, tasks): (Vec<_>, Vec<_>) = (0..num_topics) - .map(|i| { - let topic = decorator.decorate(&format!("{seed}_{i}")); - let task = ctrl_client.create_topic(topic.clone(), 1, 1, 500); - (topic, task) - }) - .unzip(); - futures::future::try_join_all(tasks).await.unwrap(); - - topics -} diff --git a/src/common/test-util/src/wal/kafka/topic_decorator.rs b/src/common/test-util/src/wal/kafka/topic_decorator.rs deleted file mode 100644 index afa7719d20a2..000000000000 --- a/src/common/test-util/src/wal/kafka/topic_decorator.rs +++ /dev/null @@ -1,76 +0,0 @@ -// 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 common_config::wal::KafkaWalTopic as Topic; - -/// Things need to bo inserted at the front or the back of the topic. -#[derive(Debug, Default)] -pub enum Affix { - /// Inserts a provided string to each topic. - Fixed(String), - /// Computes the current time for each topic and inserts it into the topic. - TimeNow, - /// Nothing to be inserted. - #[default] - Nothing, -} - -impl ToString for Affix { - fn to_string(&self) -> String { - match self { - Affix::Fixed(s) => s.to_string(), - Affix::TimeNow => chrono::Local::now().timestamp_micros().to_string(), - Affix::Nothing => String::default(), - } - } -} - -/// Decorates a topic with the given prefix and suffix. -pub struct TopicDecorator { - /// A prefix to be inserted at the front of each topic. - prefix: Affix, - /// A suffix to be inserted at the back of each topic. - suffix: Affix, -} - -impl Default for TopicDecorator { - fn default() -> Self { - Self { - prefix: Affix::Nothing, - suffix: Affix::Nothing, - } - } -} - -impl TopicDecorator { - /// Overrides the current prefix with the given prefix. - pub fn with_prefix(self, prefix: Affix) -> Self { - Self { prefix, ..self } - } - - /// Overrides the current suffix with the given suffix. - pub fn with_suffix(self, suffix: Affix) -> Self { - Self { suffix, ..self } - } - - /// Builds a topic by inserting a prefix and a suffix into the given topic. - pub fn decorate(&self, topic: &str) -> Topic { - format!( - "{}_{}_{}", - self.prefix.to_string(), - topic, - self.suffix.to_string() - ) - } -} diff --git a/src/log-store/Cargo.toml b/src/log-store/Cargo.toml index 15ef7b4ee66a..6eeb3ff190ff 100644 --- a/src/log-store/Cargo.toml +++ b/src/log-store/Cargo.toml @@ -37,6 +37,7 @@ snafu.workspace = true store-api.workspace = true tokio-util.workspace = true tokio.workspace = true +uuid.workspace = true [dev-dependencies] common-meta = { workspace = true, features = ["testing"] } diff --git a/src/log-store/src/kafka/client_manager.rs b/src/log-store/src/kafka/client_manager.rs index 2186244227ae..ef986236e354 100644 --- a/src/log-store/src/kafka/client_manager.rs +++ b/src/log-store/src/kafka/client_manager.rs @@ -105,7 +105,6 @@ impl ClientManager { } } - // Acquires the write lock. let mut client_pool = self.client_pool.write().await; match client_pool.get(topic) { Some(client) => Ok(client.clone()), @@ -137,19 +136,20 @@ impl ClientManager { #[cfg(test)] mod tests { - use common_test_util::get_broker_endpoints; - use common_test_util::wal::kafka::topic_decorator::{Affix, TopicDecorator}; - use common_test_util::wal::kafka::{create_topics, BROKER_ENDPOINTS_KEY}; use super::*; + use crate::get_broker_endpoints; + use crate::test_util::kafka::create_topics; /// Prepares for a test in that a collection of topics and a client manager are created. async fn prepare(test_name: &str, num_topics: usize) -> (ClientManager, Vec) { - let broker_endpoints = get_broker_endpoints!(BROKER_ENDPOINTS_KEY); - let decorator = TopicDecorator::default() - .with_prefix(Affix::Fixed(test_name.to_string())) - .with_suffix(Affix::TimeNow); - let topics = create_topics(num_topics, decorator, &broker_endpoints, None).await; + let broker_endpoints = get_broker_endpoints!(); + let topics = create_topics( + num_topics, + |i| format!("{test_name}_{}_{}", i, uuid::Uuid::new_v4()), + &broker_endpoints, + ) + .await; let config = KafkaConfig { broker_endpoints, @@ -160,51 +160,45 @@ mod tests { (manager, topics) } - /// Checks clients for the given topics are created. - async fn ensure_clients_exist(topics: &[Topic], client_manager: &ClientManager) { - let client_pool = client_manager.client_pool.read().await; - let all_exist = topics.iter().all(|topic| client_pool.contains_key(topic)); - assert!(all_exist); - } - - async fn test_with(test_name: &str) { - let (manager, topics) = prepare(test_name, 128).await; + /// Sends `get_or_insert` requests sequentially to the client manager, and checks if it could handle them correctly. + #[tokio::test] + async fn test_sequential() { + let (manager, topics) = prepare("test_sequential", 128).await; // Assigns multiple regions to a topic. let region_topic = (0..512) .map(|region_id| (region_id, &topics[region_id % topics.len()])) .collect::>(); - // Gets the assigned topic for each region and then gets the associated client. - match test_name { - "test_sequential" => { - // Gets all clients sequentially. - for (_, topic) in region_topic { - manager.get_or_insert(topic).await.unwrap(); - } - } - "test_parallel" => { - // Gets all clients in parallel. - let tasks = region_topic - .values() - .map(|topic| manager.get_or_insert(topic)) - .collect::>(); - futures::future::try_join_all(tasks).await.unwrap(); - } - _ => unreachable!(), + // Gets all clients sequentially. + for (_, topic) in region_topic { + manager.get_or_insert(topic).await.unwrap(); } - // Ensures all clients are created successfully. - ensure_clients_exist(&topics, &manager).await; - } - /// Sends `get_or_insert` requests sequentially to the client manager, and checks if it could handle them correctly. - #[tokio::test] - async fn test_sequential() { - test_with("test_sequential").await; + // Ensures all clients exist. + let client_pool = manager.client_pool.read().await; + let all_exist = topics.iter().all(|topic| client_pool.contains_key(topic)); + assert!(all_exist); } /// Sends `get_or_insert` requests in parallel to the client manager, and checks if it could handle them correctly. - #[tokio::test] + #[tokio::test(flavor = "multi_thread")] async fn test_parallel() { - test_with("test_parallel").await; + let (manager, topics) = prepare("test_parallel", 128).await; + // Assigns multiple regions to a topic. + let region_topic = (0..512) + .map(|region_id| (region_id, &topics[region_id % topics.len()])) + .collect::>(); + + // Gets all clients in parallel. + let tasks = region_topic + .values() + .map(|topic| manager.get_or_insert(topic)) + .collect::>(); + futures::future::try_join_all(tasks).await.unwrap(); + + // Ensures all clients exist. + let client_pool = manager.client_pool.read().await; + let all_exist = topics.iter().all(|topic| client_pool.contains_key(topic)); + assert!(all_exist); } } diff --git a/src/log-store/src/kafka/log_store.rs b/src/log-store/src/kafka/log_store.rs index 736ece995d02..71a426172fbe 100644 --- a/src/log-store/src/kafka/log_store.rs +++ b/src/log-store/src/kafka/log_store.rs @@ -284,13 +284,13 @@ fn check_termination( mod tests { use common_base::readable_size::ReadableSize; use common_config::wal::KafkaWalTopic as Topic; - use common_test_util::get_broker_endpoints; - use common_test_util::wal::kafka::topic_decorator::{Affix, TopicDecorator}; - use common_test_util::wal::kafka::{create_topics, BROKER_ENDPOINTS_KEY}; use rand::seq::IteratorRandom; use super::*; - use crate::test_util::kafka::{entries_with_random_data, new_namespace, EntryBuilder}; + use crate::get_broker_endpoints; + use crate::test_util::kafka::{ + create_topics, entries_with_random_data, new_namespace, EntryBuilder, + }; // Stores test context for a region. struct RegionContext { @@ -302,11 +302,13 @@ mod tests { /// Prepares for a test in that a log store is constructed and a collection of topics is created. async fn prepare(test_name: &str, num_topics: usize) -> (KafkaLogStore, Vec) { - let broker_endpoints = get_broker_endpoints!(BROKER_ENDPOINTS_KEY); - let decorator = TopicDecorator::default() - .with_prefix(Affix::Fixed(test_name.to_string())) - .with_suffix(Affix::TimeNow); - let topics = create_topics(num_topics, decorator, &broker_endpoints, None).await; + let broker_endpoints = get_broker_endpoints!(); + let topics = create_topics( + num_topics, + |i| format!("{test_name}_{}_{}", i, uuid::Uuid::new_v4()), + &broker_endpoints, + ) + .await; let config = KafkaConfig { broker_endpoints, @@ -332,26 +334,6 @@ mod tests { (logstore, topics) } - /// Builds a context for each region. - fn build_contexts(num_regions: usize, topics: &[Topic]) -> HashMap { - (0..num_regions) - .map(|i| { - let topic = &topics[i % topics.len()]; - let ns = new_namespace(topic, i as u64); - let entry_builder = EntryBuilder::new(ns.clone()); - ( - i as u64, - RegionContext { - ns, - entry_builder, - expected: Vec::new(), - flushed_entry_id: 0, - }, - ) - }) - .collect() - } - /// Creates a vector containing indexes of all regions if the `all` is true. /// Otherwise, creates a subset of the indexes. The cardinality of the subset /// is nearly a quarter of that of the universe set. @@ -389,39 +371,6 @@ mod tests { aggregated.into_iter().flatten().collect() } - /// Reads entries for regions and checks for each region that the gotten entries are identical with the expected ones. - async fn check_entries( - region_contexts: &HashMap, - logstore: &KafkaLogStore, - which: &[u64], - ) { - for region_id in which { - let ctx = ®ion_contexts[region_id]; - let stream = logstore - .read(&ctx.ns, ctx.flushed_entry_id + 1) - .await - .unwrap(); - let got = stream - .collect::>() - .await - .into_iter() - .flat_map(|x| x.unwrap()) - .collect::>(); - assert_eq!(ctx.expected, got); - } - } - - /// Simulates a flush for regions. - fn flush( - region_contexts: &mut HashMap, - last_entry_ids: &HashMap, - ) { - for (region_id, last_entry_id) in last_entry_ids { - let ctx = region_contexts.get_mut(region_id).unwrap(); - ctx.flushed_entry_id = *last_entry_id; - } - } - /// Starts a test with: /// * `test_name` - The name of the test. /// * `num_topics` - Number of topics to be created in the preparation phase. @@ -439,13 +388,50 @@ mod tests { large: bool, ) { let (logstore, topics) = prepare(test_name, num_topics).await; - let mut region_contexts = build_contexts(num_regions, &topics); + let mut region_contexts = (0..num_regions) + .map(|i| { + let topic = &topics[i % topics.len()]; + let ns = new_namespace(topic, i as u64); + let entry_builder = EntryBuilder::new(ns.clone()); + ( + i as u64, + RegionContext { + ns, + entry_builder, + expected: Vec::new(), + flushed_entry_id: 0, + }, + ) + }) + .collect(); + for _ in 0..num_appends { + // Appends entries for a subset of regions. let which = all_or_subset(all, num_regions); let entries = build_entries(&mut region_contexts, &which, large); let last_entry_ids = logstore.append_batch(entries).await.unwrap().last_entry_ids; - check_entries(®ion_contexts, &logstore, &which).await; - flush(&mut region_contexts, &last_entry_ids); + + // Reads entries for regions and checks for each region that the gotten entries are identical with the expected ones. + for region_id in which { + let ctx = ®ion_contexts[®ion_id]; + let stream = logstore + .read(&ctx.ns, ctx.flushed_entry_id + 1) + .await + .unwrap(); + let got = stream + .collect::>() + .await + .into_iter() + .flat_map(|x| x.unwrap()) + .collect::>(); + assert_eq!(ctx.expected, got); + } + + // Simulates a flush for regions. + for (region_id, last_entry_id) in last_entry_ids { + let ctx = region_contexts.get_mut(®ion_id).unwrap(); + ctx.flushed_entry_id = last_entry_id; + } } } diff --git a/src/log-store/src/test_util/kafka.rs b/src/log-store/src/test_util/kafka.rs index 67a71b23004c..c448d7f60848 100644 --- a/src/log-store/src/test_util/kafka.rs +++ b/src/log-store/src/test_util/kafka.rs @@ -15,13 +15,64 @@ use std::sync::atomic::{AtomicU64 as AtomicEntryId, Ordering}; use std::sync::Mutex; +use common_meta::wal::KafkaWalTopic as Topic; use rand::distributions::Alphanumeric; use rand::rngs::ThreadRng; use rand::{thread_rng, Rng}; +use rskafka::client::ClientBuilder; use store_api::logstore::EntryId; use crate::kafka::{EntryImpl, NamespaceImpl}; +/// 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::>(); + assert!(!broker_endpoints.is_empty()); + broker_endpoints + }}; +} + +/// Creates `num_topiocs` number of topics each will be decorated by the given decorator. +pub async fn create_topics( + num_topics: usize, + decorator: F, + broker_endpoints: &[String], +) -> Vec +where + F: Fn(usize) -> String, +{ + assert!(!broker_endpoints.is_empty()); + let client = ClientBuilder::new(broker_endpoints.to_vec()) + .build() + .await + .unwrap(); + let ctrl_client = client.controller_client().unwrap(); + let (topics, tasks): (Vec<_>, Vec<_>) = (0..num_topics) + .map(|i| { + let topic = decorator(i); + let task = ctrl_client.create_topic(topic.clone(), 1, 1, 500); + (topic, task) + }) + .unzip(); + futures::future::try_join_all(tasks).await.unwrap(); + topics +} + +/// Creates a new Kafka namespace with the given topic and region id. +pub fn new_namespace(topic: &str, region_id: u64) -> NamespaceImpl { + NamespaceImpl { + topic: topic.to_string(), + region_id, + } +} + /// A builder for building entries for a namespace. pub struct EntryBuilder { /// The namespace of the entries. @@ -88,11 +139,3 @@ pub fn entries_with_random_data(batch_size: usize, builder: &EntryBuilder) -> Ve .map(|_| builder.with_random_data()) .collect() } - -/// Creates a new Kafka namespace with the given topic and region id. -pub fn new_namespace(topic: &str, region_id: u64) -> NamespaceImpl { - NamespaceImpl { - topic: topic.to_string(), - region_id, - } -} From b017551d0f796eb367539c4bbda03fc1b2a81db5 Mon Sep 17 00:00:00 2001 From: niebayes Date: Mon, 8 Jan 2024 11:49:58 +0800 Subject: [PATCH 26/30] fix: update Cargo.toml --- Cargo.lock | 6 ------ src/common/meta/src/wal/kafka/topic_manager.rs | 1 - src/common/test-util/Cargo.toml | 6 ------ src/log-store/Cargo.toml | 8 +++----- 4 files changed, 3 insertions(+), 18 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index f9208f025fae..1b5c41fe0d65 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1976,15 +1976,9 @@ dependencies = [ name = "common-test-util" version = "0.5.0" dependencies = [ - "chrono", - "common-config", - "futures", "once_cell", "rand", - "rskafka", - "store-api", "tempfile", - "tokio", ] [[package]] diff --git a/src/common/meta/src/wal/kafka/topic_manager.rs b/src/common/meta/src/wal/kafka/topic_manager.rs index 9765fe14e2a6..593d4209835c 100644 --- a/src/common/meta/src/wal/kafka/topic_manager.rs +++ b/src/common/meta/src/wal/kafka/topic_manager.rs @@ -241,7 +241,6 @@ impl TopicManager { mod tests { use std::env; - use chrono::format::Fixed; use common_telemetry::info; use super::*; diff --git a/src/common/test-util/Cargo.toml b/src/common/test-util/Cargo.toml index be5b5a5275d8..60e854740643 100644 --- a/src/common/test-util/Cargo.toml +++ b/src/common/test-util/Cargo.toml @@ -5,12 +5,6 @@ edition.workspace = true license.workspace = true [dependencies] -chrono.workspace = true -common-config.workspace = true -futures.workspace = true once_cell.workspace = true rand.workspace = true -rskafka.workspace = true -store-api.workspace = true tempfile.workspace = true -tokio.workspace = true diff --git a/src/log-store/Cargo.toml b/src/log-store/Cargo.toml index 6eeb3ff190ff..9f2c8b9bc728 100644 --- a/src/log-store/Cargo.toml +++ b/src/log-store/Cargo.toml @@ -22,14 +22,10 @@ common-macro.workspace = true common-meta.workspace = true common-runtime.workspace = true common-telemetry.workspace = true -common-test-util.workspace = true futures-util.workspace = true futures.workspace = true -itertools.workspace = true protobuf = { version = "2", features = ["bytes"] } raft-engine.workspace = true -rand.workspace = true -rand_distr = "0.4" rskafka.workspace = true serde.workspace = true serde_json.workspace = true @@ -37,9 +33,11 @@ snafu.workspace = true store-api.workspace = true tokio-util.workspace = true tokio.workspace = true -uuid.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 From 48d967c5e6d82cb49f5d98d5ad058ee8eb21a1e6 Mon Sep 17 00:00:00 2001 From: niebayes Date: Mon, 8 Jan 2024 16:52:00 +0800 Subject: [PATCH 27/30] fix: CR --- src/common/meta/src/lib.rs | 1 - src/common/meta/src/wal.rs | 5 +---- src/common/meta/src/wal/kafka.rs | 3 +-- src/common/meta/src/wal/kafka/test_util.rs | 4 ---- src/common/meta/src/wal/kafka/topic_manager.rs | 9 ++------- src/common/meta/src/wal/kafka/topic_selector.rs | 1 - src/common/meta/src/wal/options_allocator.rs | 3 +-- src/log-store/src/kafka/client_manager.rs | 3 ++- src/log-store/src/kafka/log_store.rs | 2 +- src/log-store/src/test_util/kafka.rs | 15 --------------- 10 files changed, 8 insertions(+), 38 deletions(-) diff --git a/src/common/meta/src/lib.rs b/src/common/meta/src/lib.rs index 4132a58839db..45b4db5b0865 100644 --- a/src/common/meta/src/lib.rs +++ b/src/common/meta/src/lib.rs @@ -36,7 +36,6 @@ pub mod sequence; pub mod state_store; pub mod table_name; pub mod util; -#[allow(unused)] pub mod wal; pub type ClusterId = u64; diff --git a/src/common/meta/src/wal.rs b/src/common/meta/src/wal.rs index c7af1d64d306..2c018145ab42 100644 --- a/src/common/meta/src/wal.rs +++ b/src/common/meta/src/wal.rs @@ -19,12 +19,9 @@ use std::collections::HashMap; use common_config::wal::StandaloneWalConfig; use common_config::WAL_OPTIONS_KEY; -use common_telemetry::warn; use serde::{Deserialize, Serialize}; -use serde_with::with_prefix; use store_api::storage::{RegionId, RegionNumber}; -use crate::error::Result; use crate::wal::kafka::KafkaConfig; pub use crate::wal::kafka::Topic as KafkaWalTopic; pub use crate::wal::options_allocator::{ @@ -43,7 +40,7 @@ pub enum WalConfig { impl From for WalConfig { fn from(value: StandaloneWalConfig) -> Self { match value { - StandaloneWalConfig::RaftEngine(config) => WalConfig::RaftEngine, + StandaloneWalConfig::RaftEngine(_) => WalConfig::RaftEngine, StandaloneWalConfig::Kafka(config) => WalConfig::Kafka(KafkaConfig { broker_endpoints: config.base.broker_endpoints, num_topics: config.num_topics, diff --git a/src/common/meta/src/wal/kafka.rs b/src/common/meta/src/wal/kafka.rs index 625b45fce22f..18cde0fdaa52 100644 --- a/src/common/meta/src/wal/kafka.rs +++ b/src/common/meta/src/wal/kafka.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -#[cfg(test)] +#[cfg(any(test, feature = "testing"))] pub mod test_util; pub mod topic; pub mod topic_manager; @@ -21,7 +21,6 @@ pub mod topic_selector; use std::time::Duration; use common_config::wal::kafka::{kafka_backoff, KafkaBackoffConfig, TopicSelectorType}; -use common_config::wal::StandaloneWalConfig; use serde::{Deserialize, Serialize}; pub use crate::wal::kafka::topic::Topic; diff --git a/src/common/meta/src/wal/kafka/test_util.rs b/src/common/meta/src/wal/kafka/test_util.rs index 81014bc56cbe..788884444352 100644 --- a/src/common/meta/src/wal/kafka/test_util.rs +++ b/src/common/meta/src/wal/kafka/test_util.rs @@ -12,10 +12,6 @@ // 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] diff --git a/src/common/meta/src/wal/kafka/topic_manager.rs b/src/common/meta/src/wal/kafka/topic_manager.rs index 593d4209835c..5551f5ae83bb 100644 --- a/src/common/meta/src/wal/kafka/topic_manager.rs +++ b/src/common/meta/src/wal/kafka/topic_manager.rs @@ -14,10 +14,9 @@ use std::collections::HashSet; use std::sync::Arc; -use std::time::Duration; use common_config::wal::kafka::TopicSelectorType; -use common_telemetry::{debug, error, info}; +use common_telemetry::{error, info}; use rskafka::client::controller::ControllerClient; use rskafka::client::error::Error as RsKafkaError; use rskafka::client::error::ProtocolError::TopicAlreadyExists; @@ -25,7 +24,7 @@ use rskafka::client::partition::{Compression, UnknownTopicHandling}; use rskafka::client::{Client, ClientBuilder}; use rskafka::record::Record; use rskafka::BackoffConfig; -use snafu::{ensure, AsErrorSource, ResultExt}; +use snafu::{ensure, ResultExt}; use crate::error::{ BuildKafkaClientSnafu, BuildKafkaCtrlClientSnafu, BuildKafkaPartitionClientSnafu, @@ -239,10 +238,6 @@ impl TopicManager { #[cfg(test)] mod tests { - use std::env; - - use common_telemetry::info; - use super::*; use crate::get_broker_endpoints; use crate::kv_backend::memory::MemoryKvBackend; diff --git a/src/common/meta/src/wal/kafka/topic_selector.rs b/src/common/meta/src/wal/kafka/topic_selector.rs index df4feb8b2e12..432900ebacc3 100644 --- a/src/common/meta/src/wal/kafka/topic_selector.rs +++ b/src/common/meta/src/wal/kafka/topic_selector.rs @@ -16,7 +16,6 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; use rand::Rng; -use serde::{Deserialize, Serialize}; use snafu::ensure; use crate::error::{EmptyTopicPoolSnafu, Result}; diff --git a/src/common/meta/src/wal/options_allocator.rs b/src/common/meta/src/wal/options_allocator.rs index c9df0cf44bcb..062e8b6da696 100644 --- a/src/common/meta/src/wal/options_allocator.rs +++ b/src/common/meta/src/wal/options_allocator.rs @@ -116,7 +116,7 @@ mod tests { async fn test_allocator_with_raft_engine() { let kv_backend = Arc::new(MemoryKvBackend::new()) as KvBackendRef; let wal_config = WalConfig::RaftEngine; - let mut allocator = WalOptionsAllocator::new(wal_config, kv_backend); + let allocator = WalOptionsAllocator::new(wal_config, kv_backend); allocator.start().await.unwrap(); let num_regions = 32; @@ -153,7 +153,6 @@ mod tests { 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(); diff --git a/src/log-store/src/kafka/client_manager.rs b/src/log-store/src/kafka/client_manager.rs index ef986236e354..36402b5a1079 100644 --- a/src/log-store/src/kafka/client_manager.rs +++ b/src/log-store/src/kafka/client_manager.rs @@ -137,8 +137,9 @@ impl ClientManager { #[cfg(test)] mod tests { + use common_meta::get_broker_endpoints; + use super::*; - use crate::get_broker_endpoints; use crate::test_util::kafka::create_topics; /// Prepares for a test in that a collection of topics and a client manager are created. diff --git a/src/log-store/src/kafka/log_store.rs b/src/log-store/src/kafka/log_store.rs index 71a426172fbe..108b7a2ad515 100644 --- a/src/log-store/src/kafka/log_store.rs +++ b/src/log-store/src/kafka/log_store.rs @@ -284,10 +284,10 @@ fn check_termination( mod tests { use common_base::readable_size::ReadableSize; use common_config::wal::KafkaWalTopic as Topic; + use common_meta::get_broker_endpoints; use rand::seq::IteratorRandom; use super::*; - use crate::get_broker_endpoints; use crate::test_util::kafka::{ create_topics, entries_with_random_data, new_namespace, EntryBuilder, }; diff --git a/src/log-store/src/test_util/kafka.rs b/src/log-store/src/test_util/kafka.rs index c448d7f60848..7107c6e5c3f1 100644 --- a/src/log-store/src/test_util/kafka.rs +++ b/src/log-store/src/test_util/kafka.rs @@ -24,21 +24,6 @@ use store_api::logstore::EntryId; use crate::kafka::{EntryImpl, NamespaceImpl}; -/// 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::>(); - assert!(!broker_endpoints.is_empty()); - broker_endpoints - }}; -} - /// Creates `num_topiocs` number of topics each will be decorated by the given decorator. pub async fn create_topics( num_topics: usize, From c117dd8ce922f86292aeecca980d51bbcc4862fb Mon Sep 17 00:00:00 2001 From: niebayes Date: Mon, 8 Jan 2024 17:35:28 +0800 Subject: [PATCH 28/30] feat: skip test if no endpoints env --- Cargo.lock | 1 + src/common/meta/Cargo.toml | 2 + src/common/meta/src/wal/kafka/test_util.rs | 31 +++--- .../meta/src/wal/kafka/topic_manager.rs | 102 +++++++++--------- src/common/meta/src/wal/options_allocator.rs | 76 ++++++------- src/log-store/src/kafka/client_manager.rs | 84 ++++++++------- src/log-store/src/kafka/log_store.rs | 19 +++- 7 files changed, 177 insertions(+), 138 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index e19a08964241..92c636190f20 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1833,6 +1833,7 @@ dependencies = [ "derive_builder 0.12.0", "etcd-client", "futures", + "futures-util", "humantime-serde", "hyper", "lazy_static", diff --git a/src/common/meta/Cargo.toml b/src/common/meta/Cargo.toml index 724b508ebc6a..0aa97ec01fc0 100644 --- a/src/common/meta/Cargo.toml +++ b/src/common/meta/Cargo.toml @@ -29,6 +29,7 @@ datatypes.workspace = true derive_builder.workspace = true etcd-client.workspace = true futures.workspace = true +futures-util.workspace = true humantime-serde.workspace = true lazy_static.workspace = true prometheus.workspace = true @@ -51,5 +52,6 @@ tonic.workspace = true chrono.workspace = true common-procedure = { workspace = true, features = ["testing"] } datatypes.workspace = true +futures-util.workspace = true hyper = { version = "0.14", features = ["full"] } uuid.workspace = true diff --git a/src/common/meta/src/wal/kafka/test_util.rs b/src/common/meta/src/wal/kafka/test_util.rs index 788884444352..9dab26e83738 100644 --- a/src/common/meta/src/wal/kafka/test_util.rs +++ b/src/common/meta/src/wal/kafka/test_util.rs @@ -12,17 +12,22 @@ // See the License for the specific language governing permissions and // limitations under the License. -/// 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::>(); - assert!(!broker_endpoints.is_empty()); - broker_endpoints - }}; +use common_telemetry::warn; +use futures_util::future::BoxFuture; + +pub async fn run_test_with_kafka_wal(test: F) +where + F: FnOnce(Vec) -> BoxFuture<'static, ()>, +{ + let Ok(endpoints) = std::env::var("GT_KAFKA_ENDPOINTS") else { + warn!("The endpoints is empty, skipping the test"); + return; + }; + + let endpoints = endpoints + .split(',') + .map(|s| s.trim().to_string()) + .collect::>(); + + test(endpoints).await } diff --git a/src/common/meta/src/wal/kafka/topic_manager.rs b/src/common/meta/src/wal/kafka/topic_manager.rs index 5551f5ae83bb..f57d2044862b 100644 --- a/src/common/meta/src/wal/kafka/topic_manager.rs +++ b/src/common/meta/src/wal/kafka/topic_manager.rs @@ -239,8 +239,8 @@ impl TopicManager { #[cfg(test)] mod tests { use super::*; - use crate::get_broker_endpoints; use crate::kv_backend::memory::MemoryKvBackend; + use crate::wal::kafka::test_util::run_test_with_kafka_wal; // Tests that topics can be successfully persisted into the kv backend and can be successfully restored from the kv backend. #[tokio::test] @@ -270,53 +270,57 @@ mod tests { /// Tests that the topic manager could allocate topics correctly. #[tokio::test] 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::>(); - - // 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 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::>(); - 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::>(); - 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::>(); - let expected = vec![topics.clone(); 2] - .into_iter() - .flatten() - .collect::>(); - assert_eq!(got, expected); + run_test_with_kafka_wal(|broker_endpoints| { + Box::pin(async { + // Constructs topics that should be created. + let topics = (0..256) + .map(|i| format!("test_alloc_topics_{}_{}", i, uuid::Uuid::new_v4())) + .collect::>(); + + // 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 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::>(); + 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::>(); + 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::>(); + let expected = vec![topics.clone(); 2] + .into_iter() + .flatten() + .collect::>(); + assert_eq!(got, expected); + }) + }) + .await; } } diff --git a/src/common/meta/src/wal/options_allocator.rs b/src/common/meta/src/wal/options_allocator.rs index 062e8b6da696..58724b64573b 100644 --- a/src/common/meta/src/wal/options_allocator.rs +++ b/src/common/meta/src/wal/options_allocator.rs @@ -106,8 +106,8 @@ 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::test_util::run_test_with_kafka_wal; use crate::wal::kafka::topic_selector::RoundRobinTopicSelector; use crate::wal::kafka::KafkaConfig; @@ -134,41 +134,45 @@ mod tests { // 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::>(); - - // 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 allocator = WalOptionsAllocator::Kafka(topic_manager); - allocator.start().await.unwrap(); - - let num_regions = 32; - let regions = (0..num_regions).collect::>(); - 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()) + run_test_with_kafka_wal(|broker_endpoints| { + Box::pin(async { + let topics = (0..256) + .map(|i| format!("test_allocator_with_kafka_{}_{}", i, uuid::Uuid::new_v4())) + .collect::>(); + + // 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 allocator = WalOptionsAllocator::Kafka(topic_manager); + allocator.start().await.unwrap(); + + let num_regions = 32; + let regions = (0..num_regions).collect::>(); + 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::>(); + assert_eq!(got, expected); }) - .collect::>(); - assert_eq!(got, expected); + }) + .await; } } diff --git a/src/log-store/src/kafka/client_manager.rs b/src/log-store/src/kafka/client_manager.rs index 36402b5a1079..a16a17e03380 100644 --- a/src/log-store/src/kafka/client_manager.rs +++ b/src/log-store/src/kafka/client_manager.rs @@ -136,15 +136,17 @@ impl ClientManager { #[cfg(test)] mod tests { - - use common_meta::get_broker_endpoints; + use common_meta::wal::kafka::test_util::run_test_with_kafka_wal; use super::*; use crate::test_util::kafka::create_topics; /// Prepares for a test in that a collection of topics and a client manager are created. - async fn prepare(test_name: &str, num_topics: usize) -> (ClientManager, Vec) { - let broker_endpoints = get_broker_endpoints!(); + async fn prepare( + test_name: &str, + num_topics: usize, + broker_endpoints: Vec, + ) -> (ClientManager, Vec) { let topics = create_topics( num_topics, |i| format!("{test_name}_{}_{}", i, uuid::Uuid::new_v4()), @@ -164,42 +166,52 @@ mod tests { /// Sends `get_or_insert` requests sequentially to the client manager, and checks if it could handle them correctly. #[tokio::test] async fn test_sequential() { - let (manager, topics) = prepare("test_sequential", 128).await; - // Assigns multiple regions to a topic. - let region_topic = (0..512) - .map(|region_id| (region_id, &topics[region_id % topics.len()])) - .collect::>(); - - // Gets all clients sequentially. - for (_, topic) in region_topic { - manager.get_or_insert(topic).await.unwrap(); - } - - // Ensures all clients exist. - let client_pool = manager.client_pool.read().await; - let all_exist = topics.iter().all(|topic| client_pool.contains_key(topic)); - assert!(all_exist); + run_test_with_kafka_wal(|broker_endpoints| { + Box::pin(async { + let (manager, topics) = prepare("test_sequential", 128, broker_endpoints).await; + // Assigns multiple regions to a topic. + let region_topic = (0..512) + .map(|region_id| (region_id, &topics[region_id % topics.len()])) + .collect::>(); + + // Gets all clients sequentially. + for (_, topic) in region_topic { + manager.get_or_insert(topic).await.unwrap(); + } + + // Ensures all clients exist. + let client_pool = manager.client_pool.read().await; + let all_exist = topics.iter().all(|topic| client_pool.contains_key(topic)); + assert!(all_exist); + }) + }) + .await; } /// Sends `get_or_insert` requests in parallel to the client manager, and checks if it could handle them correctly. #[tokio::test(flavor = "multi_thread")] async fn test_parallel() { - let (manager, topics) = prepare("test_parallel", 128).await; - // Assigns multiple regions to a topic. - let region_topic = (0..512) - .map(|region_id| (region_id, &topics[region_id % topics.len()])) - .collect::>(); - - // Gets all clients in parallel. - let tasks = region_topic - .values() - .map(|topic| manager.get_or_insert(topic)) - .collect::>(); - futures::future::try_join_all(tasks).await.unwrap(); - - // Ensures all clients exist. - let client_pool = manager.client_pool.read().await; - let all_exist = topics.iter().all(|topic| client_pool.contains_key(topic)); - assert!(all_exist); + run_test_with_kafka_wal(|broker_endpoints| { + Box::pin(async { + let (manager, topics) = prepare("test_parallel", 128, broker_endpoints).await; + // Assigns multiple regions to a topic. + let region_topic = (0..512) + .map(|region_id| (region_id, &topics[region_id % topics.len()])) + .collect::>(); + + // Gets all clients in parallel. + let tasks = region_topic + .values() + .map(|topic| manager.get_or_insert(topic)) + .collect::>(); + futures::future::try_join_all(tasks).await.unwrap(); + + // Ensures all clients exist. + let client_pool = manager.client_pool.read().await; + let all_exist = topics.iter().all(|topic| client_pool.contains_key(topic)); + assert!(all_exist); + }) + }) + .await; } } diff --git a/src/log-store/src/kafka/log_store.rs b/src/log-store/src/kafka/log_store.rs index 108b7a2ad515..2e05683fe246 100644 --- a/src/log-store/src/kafka/log_store.rs +++ b/src/log-store/src/kafka/log_store.rs @@ -284,7 +284,6 @@ fn check_termination( mod tests { use common_base::readable_size::ReadableSize; use common_config::wal::KafkaWalTopic as Topic; - use common_meta::get_broker_endpoints; use rand::seq::IteratorRandom; use super::*; @@ -301,8 +300,11 @@ mod tests { } /// Prepares for a test in that a log store is constructed and a collection of topics is created. - async fn prepare(test_name: &str, num_topics: usize) -> (KafkaLogStore, Vec) { - let broker_endpoints = get_broker_endpoints!(); + async fn prepare( + test_name: &str, + num_topics: usize, + broker_endpoints: Vec, + ) -> (KafkaLogStore, Vec) { let topics = create_topics( num_topics, |i| format!("{test_name}_{}_{}", i, uuid::Uuid::new_v4()), @@ -387,7 +389,16 @@ mod tests { all: bool, large: bool, ) { - let (logstore, topics) = prepare(test_name, num_topics).await; + let Ok(broker_endpoints) = std::env::var("GT_KAFKA_ENDPOINTS") else { + warn!("The endpoints is empty, skipping the test {test_name}"); + return; + }; + let broker_endpoints = broker_endpoints + .split(',') + .map(|s| s.trim().to_string()) + .collect::>(); + + let (logstore, topics) = prepare(test_name, num_topics, broker_endpoints).await; let mut region_contexts = (0..num_regions) .map(|i| { let topic = &topics[i % topics.len()]; From 80760388aac518591aa49624490b838d94411472 Mon Sep 17 00:00:00 2001 From: niebayes Date: Mon, 8 Jan 2024 17:54:13 +0800 Subject: [PATCH 29/30] fix: format --- src/common/meta/Cargo.toml | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/common/meta/Cargo.toml b/src/common/meta/Cargo.toml index 0aa97ec01fc0..609fe259d2e2 100644 --- a/src/common/meta/Cargo.toml +++ b/src/common/meta/Cargo.toml @@ -28,8 +28,8 @@ common-time.workspace = true datatypes.workspace = true derive_builder.workspace = true etcd-client.workspace = true -futures.workspace = true futures-util.workspace = true +futures.workspace = true humantime-serde.workspace = true lazy_static.workspace = true prometheus.workspace = true @@ -52,6 +52,5 @@ tonic.workspace = true chrono.workspace = true common-procedure = { workspace = true, features = ["testing"] } datatypes.workspace = true -futures-util.workspace = true hyper = { version = "0.14", features = ["full"] } uuid.workspace = true From 341ae676ca93e4ea435ac268ada9a62fffb4bb5c Mon Sep 17 00:00:00 2001 From: niebayes Date: Mon, 8 Jan 2024 18:17:08 +0800 Subject: [PATCH 30/30] test: rewrite parallel test with barrier --- src/log-store/src/kafka/client_manager.rs | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/src/log-store/src/kafka/client_manager.rs b/src/log-store/src/kafka/client_manager.rs index a16a17e03380..d5402a42b1d1 100644 --- a/src/log-store/src/kafka/client_manager.rs +++ b/src/log-store/src/kafka/client_manager.rs @@ -137,6 +137,7 @@ impl ClientManager { #[cfg(test)] mod tests { use common_meta::wal::kafka::test_util::run_test_with_kafka_wal; + use tokio::sync::Barrier; use super::*; use crate::test_util::kafka::create_topics; @@ -196,13 +197,22 @@ mod tests { let (manager, topics) = prepare("test_parallel", 128, broker_endpoints).await; // Assigns multiple regions to a topic. let region_topic = (0..512) - .map(|region_id| (region_id, &topics[region_id % topics.len()])) + .map(|region_id| (region_id, topics[region_id % topics.len()].clone())) .collect::>(); // Gets all clients in parallel. + let manager = Arc::new(manager); + let barrier = Arc::new(Barrier::new(region_topic.len())); let tasks = region_topic - .values() - .map(|topic| manager.get_or_insert(topic)) + .into_values() + .map(|topic| { + let manager = manager.clone(); + let barrier = barrier.clone(); + tokio::spawn(async move { + barrier.wait().await; + assert!(manager.get_or_insert(&topic).await.is_ok()); + }) + }) .collect::>(); futures::future::try_join_all(tasks).await.unwrap();