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_secs(30 * 60);
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        #[serde(serialize_with = "serialize_password")]
73        password: String,
74    },
75    #[serde(rename = "SCRAM-SHA-256")]
76    ScramSha256 {
77        username: String,
78        #[serde(serialize_with = "serialize_password")]
79        password: String,
80    },
81    #[serde(rename = "SCRAM-SHA-512")]
82    ScramSha512 {
83        username: String,
84        #[serde(serialize_with = "serialize_password")]
85        password: String,
86    },
87}
88
89fn serialize_password<S>(_: &str, serializer: S) -> std::result::Result<S::Ok, S::Error>
90where
91    S: serde::Serializer,
92{
93    serializer.serialize_str("<REDACTED>")
94}
95
96impl fmt::Debug for KafkaClientSaslConfig {
97    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
98        match self {
99            KafkaClientSaslConfig::Plain { username, .. } => f
100                .debug_struct("Plain")
101                .field("username", username)
102                .field("password", &"<REDACTED>")
103                .finish(),
104            KafkaClientSaslConfig::ScramSha256 { username, .. } => f
105                .debug_struct("ScramSha256")
106                .field("username", username)
107                .field("password", &"<REDACTED>")
108                .finish(),
109            KafkaClientSaslConfig::ScramSha512 { username, .. } => f
110                .debug_struct("ScramSha512")
111                .field("username", username)
112                .field("password", &"<REDACTED>")
113                .finish(),
114        }
115    }
116}
117
118impl KafkaClientSaslConfig {
119    /// Converts to [`SaslConfig`].
120    pub fn into_sasl_config(self) -> SaslConfig {
121        match self {
122            KafkaClientSaslConfig::Plain { username, password } => {
123                SaslConfig::Plain(Credentials::new(username, password))
124            }
125            KafkaClientSaslConfig::ScramSha256 { username, password } => {
126                SaslConfig::ScramSha256(Credentials::new(username, password))
127            }
128            KafkaClientSaslConfig::ScramSha512 { username, password } => {
129                SaslConfig::ScramSha512(Credentials::new(username, password))
130            }
131        }
132    }
133}
134
135/// The TLS configurations for kafka client.
136#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
137pub struct KafkaClientTls {
138    pub server_ca_cert_path: Option<String>,
139    pub client_cert_path: Option<String>,
140    pub client_key_path: Option<String>,
141}
142
143impl KafkaClientTls {
144    /// Builds the [`ClientConfig`].
145    pub async fn to_tls_config(&self) -> Result<Arc<ClientConfig>> {
146        let builder = ClientConfig::builder();
147        let mut roots = RootCertStore::empty();
148
149        if let Some(server_ca_cert_path) = &self.server_ca_cert_path {
150            let root_cert_bytes =
151                tokio::fs::read(&server_ca_cert_path)
152                    .await
153                    .context(error::ReadFileSnafu {
154                        path: server_ca_cert_path,
155                    })?;
156            let mut cursor = Cursor::new(root_cert_bytes);
157            for cert in rustls_pemfile::certs(&mut cursor)
158                .collect::<std::result::Result<Vec<_>, _>>()
159                .context(error::ReadCertsSnafu {
160                    path: server_ca_cert_path,
161                })?
162            {
163                roots.add(cert).context(error::AddCertSnafu)?;
164            }
165        };
166        roots.add_parsable_certificates(
167            rustls_native_certs::load_native_certs().context(error::LoadSystemCertsSnafu)?,
168        );
169
170        let builder = builder.with_root_certificates(roots);
171        let config = if let (Some(cert_path), Some(key_path)) =
172            (&self.client_cert_path, &self.client_key_path)
173        {
174            let cert_bytes = tokio::fs::read(cert_path)
175                .await
176                .context(error::ReadFileSnafu { path: cert_path })?;
177            let client_certs = rustls_pemfile::certs(&mut Cursor::new(cert_bytes))
178                .collect::<std::result::Result<Vec<_>, _>>()
179                .context(error::ReadCertsSnafu { path: cert_path })?;
180            let key_bytes = tokio::fs::read(key_path)
181                .await
182                .context(error::ReadFileSnafu { path: key_path })?;
183            let client_key = rustls_pemfile::private_key(&mut Cursor::new(key_bytes))
184                .context(error::ReadKeySnafu { path: key_path })?
185                .context(error::KeyNotFoundSnafu { path: key_path })?;
186
187            builder
188                .with_client_auth_cert(client_certs, client_key)
189                .context(error::SetClientAuthCertSnafu)?
190        } else {
191            builder.with_no_client_auth()
192        };
193
194        Ok(Arc::new(config))
195    }
196}
197
198/// The connection configurations for kafka clients.
199#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
200#[serde(default)]
201pub struct KafkaConnectionConfig {
202    /// The broker endpoints of the Kafka cluster.
203    pub broker_endpoints: Vec<String>,
204    /// Client SASL.
205    pub sasl: Option<KafkaClientSasl>,
206    /// Client TLS config
207    pub tls: Option<KafkaClientTls>,
208    /// The connect timeout for kafka client.
209    #[serde(with = "humantime_serde")]
210    pub connect_timeout: Duration,
211    /// The total request timeout for kafka client.
212    #[serde(with = "humantime_serde")]
213    pub timeout: Duration,
214}
215
216impl Default for KafkaConnectionConfig {
217    fn default() -> Self {
218        Self {
219            broker_endpoints: vec![BROKER_ENDPOINT.to_string()],
220            sasl: None,
221            tls: None,
222            connect_timeout: Duration::from_secs(3),
223            timeout: Duration::from_secs(5),
224        }
225    }
226}
227
228/// Topic configurations for kafka clients.
229#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
230#[serde(default)]
231pub struct KafkaTopicConfig {
232    /// Number of topics.
233    pub num_topics: usize,
234    /// Number of partitions per topic.
235    pub num_partitions: i32,
236    /// The type of the topic selector with which to select a topic for a region.
237    pub selector_type: TopicSelectorType,
238    /// The replication factor of each topic.
239    pub replication_factor: i16,
240    /// The timeout of topic creation.
241    #[serde(with = "humantime_serde")]
242    pub create_topic_timeout: Duration,
243    /// Topic name prefix.
244    pub topic_name_prefix: String,
245}
246
247impl Default for KafkaTopicConfig {
248    fn default() -> Self {
249        Self {
250            num_topics: 64,
251            num_partitions: 1,
252            selector_type: TopicSelectorType::RoundRobin,
253            replication_factor: 1,
254            create_topic_timeout: Duration::from_secs(30),
255            topic_name_prefix: TOPIC_NAME_PREFIX.to_string(),
256        }
257    }
258}
common_wal/config/kafka/common.rs

common_wal/config/kafka/
common.rs
