log_store/kafka/
client_manager.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::HashMap;
16use std::sync::Arc;
17
18use common_wal::config::kafka::DatanodeKafkaConfig;
19use common_wal::config::kafka::common::DEFAULT_BACKOFF_CONFIG;
20use dashmap::DashMap;
21use rskafka::client::ClientBuilder;
22use rskafka::client::partition::{Compression, PartitionClient, UnknownTopicHandling};
23use snafu::ResultExt;
24use store_api::logstore::provider::KafkaProvider;
25use tokio::sync::{Mutex, RwLock};
26
27use crate::error::{BuildClientSnafu, BuildPartitionClientSnafu, Result, TlsConfigSnafu};
28use crate::kafka::log_store::TopicStat;
29use crate::kafka::producer::{OrderedBatchProducer, OrderedBatchProducerRef};
30
31// Each topic only has one partition for now.
32// The `DEFAULT_PARTITION` refers to the index of the partition.
33pub const DEFAULT_PARTITION: i32 = 0;
34
35/// Arc wrapper of ClientManager.
36pub(crate) type ClientManagerRef = Arc<ClientManager>;
37
38/// Topic client.
39#[derive(Debug, Clone)]
40pub(crate) struct Client {
41    client: Arc<PartitionClient>,
42    producer: OrderedBatchProducerRef,
43}
44
45impl Client {
46    pub(crate) fn client(&self) -> &Arc<PartitionClient> {
47        &self.client
48    }
49
50    pub(crate) fn producer(&self) -> &OrderedBatchProducerRef {
51        &self.producer
52    }
53}
54
55/// Manages client construction and accesses.
56#[derive(Debug)]
57pub(crate) struct ClientManager {
58    client: rskafka::client::Client,
59    /// Used to initialize a new [Client].
60    mutex: Mutex<()>,
61    instances: RwLock<HashMap<Arc<KafkaProvider>, Client>>,
62
63    flush_batch_size: usize,
64    compression: Compression,
65
66    /// The stats of each topic.
67    topic_stats: Arc<DashMap<Arc<KafkaProvider>, TopicStat>>,
68}
69
70impl ClientManager {
71    /// Tries to create a ClientManager.
72    pub(crate) async fn try_new(
73        config: &DatanodeKafkaConfig,
74        topic_stats: Arc<DashMap<Arc<KafkaProvider>, TopicStat>>,
75    ) -> Result<Self> {
76        // Sets backoff config for the top-level kafka client and all clients constructed by it.
77        let mut builder = ClientBuilder::new(config.connection.broker_endpoints.clone())
78            .backoff_config(DEFAULT_BACKOFF_CONFIG)
79            .connect_timeout(Some(config.connection.connect_timeout))
80            .timeout(Some(config.connection.timeout));
81        if let Some(sasl) = &config.connection.sasl {
82            builder = builder.sasl_config(sasl.config.clone().into_sasl_config());
83        };
84        if let Some(tls) = &config.connection.tls {
85            builder = builder.tls_config(tls.to_tls_config().await.context(TlsConfigSnafu)?)
86        };
87
88        let client = builder.build().await.with_context(|_| BuildClientSnafu {
89            broker_endpoints: config.connection.broker_endpoints.clone(),
90        })?;
91
92        Ok(Self {
93            client,
94            mutex: Mutex::new(()),
95            instances: RwLock::new(HashMap::new()),
96            flush_batch_size: config.max_batch_bytes.as_bytes() as usize,
97            compression: Compression::Lz4,
98            topic_stats,
99        })
100    }
101
102    async fn try_insert(&self, provider: &Arc<KafkaProvider>) -> Result<Client> {
103        let _guard = self.mutex.lock().await;
104
105        let client = self.instances.read().await.get(provider).cloned();
106        match client {
107            Some(client) => Ok(client),
108            None => {
109                let client = self.try_create_client(provider).await?;
110                self.instances
111                    .write()
112                    .await
113                    .insert(provider.clone(), client.clone());
114                self.topic_stats
115                    .insert(provider.clone(), TopicStat::default());
116                Ok(client)
117            }
118        }
119    }
120
121    /// Gets the client associated with the topic. If the client does not exist, a new one will
122    /// be created and returned.
123    pub(crate) async fn get_or_insert(&self, provider: &Arc<KafkaProvider>) -> Result<Client> {
124        let client = self.instances.read().await.get(provider).cloned();
125        match client {
126            Some(client) => Ok(client),
127            None => self.try_insert(provider).await,
128        }
129    }
130
131    async fn try_create_client(&self, provider: &Arc<KafkaProvider>) -> Result<Client> {
132        // Sets to Retry to retry connecting if the kafka cluster replies with an UnknownTopic error.
133        // That's because the topic is believed to exist as the metasrv is expected to create required topics upon start.
134        // The reconnecting won't stop until succeed or a different error returns.
135        let client = self
136            .client
137            .partition_client(
138                provider.topic.as_str(),
139                DEFAULT_PARTITION,
140                UnknownTopicHandling::Retry,
141            )
142            .await
143            .context(BuildPartitionClientSnafu {
144                topic: &provider.topic,
145                partition: DEFAULT_PARTITION,
146            })
147            .map(Arc::new)?;
148
149        let (tx, rx) = OrderedBatchProducer::channel();
150        let producer = Arc::new(OrderedBatchProducer::new(
151            (tx, rx),
152            provider.clone(),
153            client.clone(),
154            self.compression,
155            self.flush_batch_size,
156            self.topic_stats.clone(),
157        ));
158
159        Ok(Client { client, producer })
160    }
161
162    /// Lists all topics.
163    pub(crate) async fn list_topics(&self) -> Vec<Arc<KafkaProvider>> {
164        self.instances.read().await.keys().cloned().collect()
165    }
166
167    #[cfg(test)]
168    pub(crate) fn topic_stats(&self) -> &Arc<DashMap<Arc<KafkaProvider>, TopicStat>> {
169        &self.topic_stats
170    }
171}
172
173#[cfg(test)]
174impl ClientManager {
175    /// Returns the controller client.
176    pub(crate) fn controller_client(&self) -> rskafka::client::controller::ControllerClient {
177        self.client.controller_client().unwrap()
178    }
179}
180
181#[cfg(test)]
182mod tests {
183    use common_wal::maybe_skip_kafka_integration_test;
184    use common_wal::test_util::get_kafka_endpoints;
185    use tokio::sync::Barrier;
186
187    use super::*;
188    use crate::kafka::test_util::prepare;
189
190    /// Sends `get_or_insert` requests sequentially to the client manager, and checks if it could handle them correctly.
191    #[tokio::test]
192    async fn test_sequential() {
193        maybe_skip_kafka_integration_test!();
194        let broker_endpoints = get_kafka_endpoints();
195        let (manager, topics) = prepare("test_sequential", 128, broker_endpoints).await;
196        // Assigns multiple regions to a topic.
197        let region_topic = (0..512)
198            .map(|region_id| (region_id, &topics[region_id % topics.len()]))
199            .collect::<HashMap<_, _>>();
200
201        // Gets all clients sequentially.
202        for (_, topic) in region_topic {
203            let provider = Arc::new(KafkaProvider::new(topic.clone()));
204            manager.get_or_insert(&provider).await.unwrap();
205        }
206
207        // Ensures all clients exist.
208        let client_pool = manager.instances.read().await;
209        let all_exist = topics.iter().all(|topic| {
210            let provider = Arc::new(KafkaProvider::new(topic.clone()));
211            client_pool.contains_key(&provider)
212        });
213        assert!(all_exist);
214    }
215
216    /// Sends `get_or_insert` requests in parallel to the client manager, and checks if it could handle them correctly.
217    #[tokio::test(flavor = "multi_thread")]
218    async fn test_parallel() {
219        maybe_skip_kafka_integration_test!();
220        let broker_endpoints = get_kafka_endpoints();
221        let (manager, topics) = prepare("test_parallel", 128, broker_endpoints).await;
222        // Assigns multiple regions to a topic.
223        let region_topic = (0..512)
224            .map(|region_id| (region_id, topics[region_id % topics.len()].clone()))
225            .collect::<HashMap<_, _>>();
226
227        // Gets all clients in parallel.
228        let manager = Arc::new(manager);
229        let barrier = Arc::new(Barrier::new(region_topic.len()));
230        let tasks = region_topic
231            .into_values()
232            .map(|topic| {
233                let manager = manager.clone();
234                let barrier = barrier.clone();
235
236                tokio::spawn(async move {
237                    barrier.wait().await;
238                    let provider = Arc::new(KafkaProvider::new(topic));
239                    assert!(manager.get_or_insert(&provider).await.is_ok());
240                })
241            })
242            .collect::<Vec<_>>();
243        futures::future::try_join_all(tasks).await.unwrap();
244
245        // Ensures all clients exist.
246        let client_pool = manager.instances.read().await;
247        let all_exist = topics.iter().all(|topic| {
248            let provider = Arc::new(KafkaProvider::new(topic.clone()));
249            client_pool.contains_key(&provider)
250        });
251        assert!(all_exist);
252    }
253}
log_store/kafka/client_manager.rs

log_store/kafka/
client_manager.rs
