Skip to main content

common_wal/config/kafka/
common.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::fmt;
16use std::io::Cursor;
17use std::sync::Arc;
18use std::time::Duration;
19
20use common_base::readable_size::ReadableSize;
21use rskafka::BackoffConfig;
22use rskafka::client::{Credentials, SaslConfig};
23use rustls::{ClientConfig, RootCertStore};
24use serde::{Deserialize, Serialize};
25use snafu::{OptionExt, ResultExt};
26
27/// The default backoff config for kafka client.
28///
29/// If the operation fails, the client will retry 3 times.
30/// The backoff time is 100ms, 300ms, 900ms.
31pub const DEFAULT_BACKOFF_CONFIG: BackoffConfig = BackoffConfig {
32    init_backoff: Duration::from_millis(100),
33    max_backoff: Duration::from_secs(1),
34    base: 3.0,
35    // The deadline shouldn't be too long,
36    // otherwise the client will block the worker loop for a long time.
37    deadline: Some(Duration::from_secs(3)),
38};
39
40/// Default interval for auto WAL pruning.
41pub const DEFAULT_AUTO_PRUNE_INTERVAL: Duration = Duration::from_mins(30);
42/// Default mode for auto WAL pruning.
43pub const DEFAULT_AUTO_PRUNE_LOGICAL_DELETE: bool = false;
44/// Default limit for concurrent auto pruning tasks.
45pub const DEFAULT_AUTO_PRUNE_PARALLELISM: usize = 10;
46/// Default size of WAL to trigger flush.
47pub const DEFAULT_FLUSH_TRIGGER_SIZE: ReadableSize = ReadableSize::mb(512);
48/// Default checkpoint trigger size.
49pub const DEFAULT_CHECKPOINT_TRIGGER_SIZE: ReadableSize = ReadableSize::mb(128);
50/// Default interval for remote WAL region flush trigger.
51pub const DEFAULT_REGION_FLUSH_TRIGGER_INTERVAL: Duration = Duration::from_secs(60);
52/// Default interval to periodically persist remote WAL checkpoints.
53pub const DEFAULT_PERIODIC_CHECKPOINT_PERSIST_INTERVAL: Duration = Duration::from_secs(60 * 60);
54/// Default interval for fetching latest Kafka topic offsets.
55pub const DEFAULT_TOPIC_LATEST_OFFSET_FETCH_INTERVAL: Duration = Duration::from_secs(60);
56
57use crate::error::{self, Result};
58use crate::{BROKER_ENDPOINT, TOPIC_NAME_PREFIX, TopicSelectorType};
59
60/// The SASL configurations for kafka client.
61#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
62pub struct KafkaClientSasl {
63    #[serde(flatten)]
64    pub config: KafkaClientSaslConfig,
65}
66
67#[derive(Clone, Serialize, Deserialize, PartialEq, Eq)]
68#[serde(tag = "type", rename_all = "SCREAMING-KEBAB-CASE")]
69pub enum KafkaClientSaslConfig {
70    Plain {
71        username: String,
72        password: String,
73    },
74    #[serde(rename = "SCRAM-SHA-256")]
75    ScramSha256 {
76        username: String,
77        password: String,
78    },
79    #[serde(rename = "SCRAM-SHA-512")]
80    ScramSha512 {
81        username: String,
82        password: String,
83    },
84}
85
86impl fmt::Debug for KafkaClientSaslConfig {
87    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
88        match self {
89            KafkaClientSaslConfig::Plain { username, .. } => f
90                .debug_struct("Plain")
91                .field("username", username)
92                .field("password", &"<REDACTED>")
93                .finish(),
94            KafkaClientSaslConfig::ScramSha256 { username, .. } => f
95                .debug_struct("ScramSha256")
96                .field("username", username)
97                .field("password", &"<REDACTED>")
98                .finish(),
99            KafkaClientSaslConfig::ScramSha512 { username, .. } => f
100                .debug_struct("ScramSha512")
101                .field("username", username)
102                .field("password", &"<REDACTED>")
103                .finish(),
104        }
105    }
106}
107
108impl KafkaClientSaslConfig {
109    /// Converts to [`SaslConfig`].
110    pub fn into_sasl_config(self) -> SaslConfig {
111        match self {
112            KafkaClientSaslConfig::Plain { username, password } => {
113                SaslConfig::Plain(Credentials::new(username, password))
114            }
115            KafkaClientSaslConfig::ScramSha256 { username, password } => {
116                SaslConfig::ScramSha256(Credentials::new(username, password))
117            }
118            KafkaClientSaslConfig::ScramSha512 { username, password } => {
119                SaslConfig::ScramSha512(Credentials::new(username, password))
120            }
121        }
122    }
123}
124
125/// The TLS configurations for kafka client.
126#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
127pub struct KafkaClientTls {
128    pub server_ca_cert_path: Option<String>,
129    pub client_cert_path: Option<String>,
130    pub client_key_path: Option<String>,
131}
132
133impl KafkaClientTls {
134    /// Builds the [`ClientConfig`].
135    pub async fn to_tls_config(&self) -> Result<Arc<ClientConfig>> {
136        let builder = ClientConfig::builder();
137        let mut roots = RootCertStore::empty();
138
139        if let Some(server_ca_cert_path) = &self.server_ca_cert_path {
140            let root_cert_bytes =
141                tokio::fs::read(&server_ca_cert_path)
142                    .await
143                    .context(error::ReadFileSnafu {
144                        path: server_ca_cert_path,
145                    })?;
146            let mut cursor = Cursor::new(root_cert_bytes);
147            for cert in rustls_pemfile::certs(&mut cursor)
148                .collect::<std::result::Result<Vec<_>, _>>()
149                .context(error::ReadCertsSnafu {
150                    path: server_ca_cert_path,
151                })?
152            {
153                roots.add(cert).context(error::AddCertSnafu)?;
154            }
155        };
156        roots.add_parsable_certificates(
157            rustls_native_certs::load_native_certs().context(error::LoadSystemCertsSnafu)?,
158        );
159
160        let builder = builder.with_root_certificates(roots);
161        let config = if let (Some(cert_path), Some(key_path)) =
162            (&self.client_cert_path, &self.client_key_path)
163        {
164            let cert_bytes = tokio::fs::read(cert_path)
165                .await
166                .context(error::ReadFileSnafu { path: cert_path })?;
167            let client_certs = rustls_pemfile::certs(&mut Cursor::new(cert_bytes))
168                .collect::<std::result::Result<Vec<_>, _>>()
169                .context(error::ReadCertsSnafu { path: cert_path })?;
170            let key_bytes = tokio::fs::read(key_path)
171                .await
172                .context(error::ReadFileSnafu { path: key_path })?;
173            let client_key = rustls_pemfile::private_key(&mut Cursor::new(key_bytes))
174                .context(error::ReadKeySnafu { path: key_path })?
175                .context(error::KeyNotFoundSnafu { path: key_path })?;
176
177            builder
178                .with_client_auth_cert(client_certs, client_key)
179                .context(error::SetClientAuthCertSnafu)?
180        } else {
181            builder.with_no_client_auth()
182        };
183
184        Ok(Arc::new(config))
185    }
186}
187
188/// The connection configurations for kafka clients.
189#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
190#[serde(default)]
191pub struct KafkaConnectionConfig {
192    /// The broker endpoints of the Kafka cluster.
193    pub broker_endpoints: Vec<String>,
194    /// Client SASL.
195    pub sasl: Option<KafkaClientSasl>,
196    /// Client TLS config
197    pub tls: Option<KafkaClientTls>,
198    /// The connect timeout for kafka client.
199    #[serde(with = "humantime_serde")]
200    pub connect_timeout: Duration,
201    /// The timeout for kafka client.
202    #[serde(with = "humantime_serde")]
203    pub timeout: Duration,
204}
205
206impl Default for KafkaConnectionConfig {
207    fn default() -> Self {
208        Self {
209            broker_endpoints: vec![BROKER_ENDPOINT.to_string()],
210            sasl: None,
211            tls: None,
212            connect_timeout: Duration::from_secs(3),
213            timeout: Duration::from_secs(3),
214        }
215    }
216}
217
218/// Topic configurations for kafka clients.
219#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
220#[serde(default)]
221pub struct KafkaTopicConfig {
222    /// Number of topics.
223    pub num_topics: usize,
224    /// Number of partitions per topic.
225    pub num_partitions: i32,
226    /// The type of the topic selector with which to select a topic for a region.
227    pub selector_type: TopicSelectorType,
228    /// The replication factor of each topic.
229    pub replication_factor: i16,
230    /// The timeout of topic creation.
231    #[serde(with = "humantime_serde")]
232    pub create_topic_timeout: Duration,
233    /// Topic name prefix.
234    pub topic_name_prefix: String,
235}
236
237impl Default for KafkaTopicConfig {
238    fn default() -> Self {
239        Self {
240            num_topics: 64,
241            num_partitions: 1,
242            selector_type: TopicSelectorType::RoundRobin,
243            replication_factor: 1,
244            create_topic_timeout: Duration::from_secs(30),
245            topic_name_prefix: TOPIC_NAME_PREFIX.to_string(),
246        }
247    }
248}