common_meta/
wal_provider.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
15mod selector;
16pub(crate) mod topic_creator;
17mod topic_manager;
18pub(crate) mod topic_pool;
19
20use std::collections::{BTreeSet, HashMap};
21use std::sync::Arc;
22
23use async_trait::async_trait;
24use common_procedure::Context as ProcedureContext;
25use common_procedure::local::DynamicKeyLockGuard;
26use common_wal::config::MetasrvWalConfig;
27use common_wal::options::{KafkaWalOptions, ObjectStoreWalOptions, WAL_OPTIONS_KEY, WalOptions};
28use serde::{Deserialize, Deserializer, Serialize, Serializer};
29use snafu::ensure;
30use store_api::storage::{RegionId, RegionNumber};
31
32use crate::ddl::allocator::wal_options::WalOptionsAllocator;
33use crate::error::{InvalidTopicNamePrefixSnafu, Result};
34use crate::key::topic_name::TopicNameKey;
35use crate::key::{TOPIC_NAME_PATTERN_REGEX, TableMetadataManagerRef};
36use crate::kv_backend::KvBackendRef;
37use crate::leadership_notifier::LeadershipChangeListener;
38use crate::lock_key::RemoteWalLock;
39pub use crate::wal_provider::topic_creator::{build_kafka_client, build_kafka_topic_creator};
40use crate::wal_provider::topic_pool::KafkaTopicPool;
41
42/// WAL options allocated for each region.
43pub type RegionWalOptions = HashMap<RegionNumber, WalOptions>;
44
45/// Returns remote WAL topics referenced by region WAL options.
46pub fn remote_wal_topics(region_wal_options: &RegionWalOptions) -> Vec<&str> {
47    region_wal_options
48        .values()
49        .filter_map(|wal_options| match wal_options {
50            WalOptions::Kafka(kafka_options) => Some(kafka_options.topic.as_str()),
51            _ => None,
52        })
53        .collect::<BTreeSet<_>>()
54        .into_iter()
55        .collect()
56}
57
58/// Acquires per-topic read locks for remote WAL topics.
59pub async fn acquire_remote_wal_read_locks(
60    ctx: &ProcedureContext,
61    region_wal_options: &RegionWalOptions,
62) -> Vec<DynamicKeyLockGuard> {
63    let topics = remote_wal_topics(region_wal_options)
64        .into_iter()
65        .map(str::to_string)
66        .collect::<Vec<_>>();
67    let mut guards = Vec::with_capacity(topics.len());
68    for topic in topics {
69        let guard = ctx
70            .provider
71            .acquire_lock(&(RemoteWalLock::Read(topic).into()))
72            .await;
73        guards.push(guard);
74    }
75    guards
76}
77
78/// Refreshes initial pruned entry ids for Kafka WAL options.
79pub async fn refresh_initial_pruned_entry_ids(
80    table_metadata_manager: &TableMetadataManagerRef,
81    region_wal_options: &mut RegionWalOptions,
82) -> Result<()> {
83    let topics = remote_wal_topics(region_wal_options);
84    if topics.is_empty() {
85        return Ok(());
86    }
87
88    let topic_values = table_metadata_manager
89        .topic_name_manager()
90        .batch_get(
91            topics
92                .iter()
93                .map(|topic| TopicNameKey::new(topic))
94                .collect(),
95        )
96        .await?;
97
98    for wal_options in region_wal_options.values_mut() {
99        let WalOptions::Kafka(kafka_options) = wal_options else {
100            continue;
101        };
102        kafka_options.initial_pruned_entry_id = Some(
103            topic_values
104                .get(&kafka_options.topic)
105                .map(|value| value.pruned_entry_id)
106                .unwrap_or_default(),
107        );
108    }
109
110    Ok(())
111}
112
113#[derive(Deserialize)]
114#[serde(untagged)]
115enum WalOptionsCompat {
116    Encoded(String),
117    Structured(WalOptions),
118}
119
120fn deserialize_region_wal_options<E>(
121    values: HashMap<String, WalOptionsCompat>,
122) -> std::result::Result<RegionWalOptions, E>
123where
124    E: serde::de::Error,
125{
126    values
127        .into_iter()
128        .map(|(region_number, wal_options)| {
129            let region_number = region_number.parse::<RegionNumber>().map_err(|err| {
130                E::custom(format!(
131                    "invalid region number in region_wal_options: {region_number}, err: {err}"
132                ))
133            })?;
134            let wal_options = match wal_options {
135                WalOptionsCompat::Encoded(encoded) => serde_json::from_str(&encoded).map_err(|err| {
136                    E::custom(format!(
137                        "failed to decode legacy wal options for region {region_number}: {encoded}, err: {err}"
138                    ))
139                })?,
140                WalOptionsCompat::Structured(wal_options) => wal_options,
141            };
142            Ok((region_number, wal_options))
143        })
144        .collect()
145}
146
147fn serialize_region_wal_options<S>(
148    value: &RegionWalOptions,
149    serializer: S,
150) -> std::result::Result<S::Ok, S::Error>
151where
152    S: Serializer,
153{
154    let values = value
155        .iter()
156        .map(|(region_number, wal_options)| {
157            serde_json::to_string(wal_options).map(|encoded| (*region_number, encoded))
158        })
159        .collect::<std::result::Result<HashMap<_, _>, _>>()
160        .map_err(serde::ser::Error::custom)?;
161    values.serialize(serializer)
162}
163
164/// Serde helpers for [`RegionWalOptions`] persisted in metadata.
165///
166/// Metadata writes WAL options as JSON strings encoded from [`WalOptions`] for
167/// compatibility. The deserializer also accepts structured JSON objects.
168pub mod region_wal_options_serde {
169    use super::*;
170
171    /// Serializes region WAL options as encoded JSON strings.
172    pub fn serialize<S>(
173        value: &RegionWalOptions,
174        serializer: S,
175    ) -> std::result::Result<S::Ok, S::Error>
176    where
177        S: Serializer,
178    {
179        serialize_region_wal_options(value, serializer)
180    }
181
182    /// Deserializes region WAL options from either structured or legacy encoded form.
183    pub fn deserialize<'de, D>(deserializer: D) -> std::result::Result<RegionWalOptions, D::Error>
184    where
185        D: Deserializer<'de>,
186    {
187        let values = HashMap::<String, WalOptionsCompat>::deserialize(deserializer)?;
188        deserialize_region_wal_options(values)
189    }
190}
191
192/// Serde helpers for optional [`RegionWalOptions`] persisted in procedure state.
193pub mod optional_region_wal_options_serde {
194    use super::*;
195
196    /// Serializes optional region WAL options as encoded JSON strings.
197    pub fn serialize<S>(
198        value: &Option<RegionWalOptions>,
199        serializer: S,
200    ) -> std::result::Result<S::Ok, S::Error>
201    where
202        S: Serializer,
203    {
204        match value {
205            Some(value) => serialize_region_wal_options(value, serializer),
206            None => serializer.serialize_none(),
207        }
208    }
209
210    /// Deserializes optional region WAL options from structured or legacy encoded form.
211    pub fn deserialize<'de, D>(
212        deserializer: D,
213    ) -> std::result::Result<Option<RegionWalOptions>, D::Error>
214    where
215        D: Deserializer<'de>,
216    {
217        let Some(values) = Option::<HashMap<String, WalOptionsCompat>>::deserialize(deserializer)?
218        else {
219            return Ok(None);
220        };
221
222        deserialize_region_wal_options(values).map(Some)
223    }
224}
225
226/// Provides wal options in region granularity.
227#[derive(Default, Debug)]
228pub enum WalProvider {
229    #[default]
230    RaftEngine,
231    Kafka(KafkaTopicPool),
232    ObjectStore {
233        prefix: String,
234    },
235}
236
237/// Arc wrapper of WalProvider.
238pub type WalProviderRef = Arc<WalProvider>;
239
240#[async_trait::async_trait]
241impl WalOptionsAllocator for WalProvider {
242    async fn allocate(
243        &self,
244        region_numbers: &[RegionNumber],
245        skip_wal: bool,
246    ) -> Result<RegionWalOptions> {
247        let wal_options = self.alloc_batch(region_numbers.len(), skip_wal).await?;
248
249        Ok(region_numbers.iter().copied().zip(wal_options).collect())
250    }
251}
252
253impl WalProvider {
254    /// Tries to start the provider.
255    pub async fn start(&self) -> Result<()> {
256        match self {
257            Self::RaftEngine | Self::ObjectStore { .. } => Ok(()),
258            Self::Kafka(kafka_topic_manager) => kafka_topic_manager.activate().await,
259        }
260    }
261
262    /// Allocates a batch of wal options where each wal options goes to a region.
263    /// If skip_wal is true, the wal options will be set to Noop regardless of the provider type.
264    pub async fn alloc_batch(&self, num_regions: usize, skip_wal: bool) -> Result<Vec<WalOptions>> {
265        if skip_wal {
266            return Ok(vec![WalOptions::Noop; num_regions]);
267        }
268        match self {
269            WalProvider::RaftEngine => Ok(vec![WalOptions::RaftEngine; num_regions]),
270            WalProvider::Kafka(topic_manager) => {
271                let options_batch = topic_manager
272                    .select_batch(num_regions)?
273                    .into_iter()
274                    .map(|topic| WalOptions::Kafka(KafkaWalOptions::new(topic.clone())))
275                    .collect();
276                Ok(options_batch)
277            }
278            WalProvider::ObjectStore { prefix } => Ok(vec![
279                WalOptions::ObjectStore(
280                    ObjectStoreWalOptions::new(prefix.clone())
281                );
282                num_regions
283            ]),
284        }
285    }
286
287    /// Returns true if it's the remote WAL.
288    pub fn is_remote_wal(&self) -> bool {
289        matches!(&self, WalProvider::Kafka(_))
290    }
291}
292
293#[async_trait]
294impl LeadershipChangeListener for WalProvider {
295    fn name(&self) -> &str {
296        "WalProvider"
297    }
298
299    async fn on_leader_start(&self) -> Result<()> {
300        self.start().await
301    }
302
303    async fn on_leader_stop(&self) -> Result<()> {
304        Ok(())
305    }
306}
307
308/// Builds a wal provider based on the given configuration.
309pub async fn build_wal_provider(
310    config: &MetasrvWalConfig,
311    kv_backend: KvBackendRef,
312) -> Result<WalProvider> {
313    match config {
314        MetasrvWalConfig::RaftEngine => Ok(WalProvider::RaftEngine),
315        MetasrvWalConfig::Kafka(kafka_config) => {
316            let prefix = &kafka_config.kafka_topic.topic_name_prefix;
317            ensure!(
318                TOPIC_NAME_PATTERN_REGEX.is_match(prefix),
319                InvalidTopicNamePrefixSnafu { prefix }
320            );
321            let topic_creator =
322                build_kafka_topic_creator(&kafka_config.connection, &kafka_config.kafka_topic)
323                    .await?;
324            let topic_pool = KafkaTopicPool::new(kafka_config, kv_backend, topic_creator);
325            Ok(WalProvider::Kafka(topic_pool))
326        }
327    }
328}
329
330/// Serializes and inserts WAL options into the region options.
331pub fn serialize_wal_options(
332    options: &mut HashMap<String, String>,
333    region_id: RegionId,
334    region_wal_options: &RegionWalOptions,
335) -> std::result::Result<(), serde_json::Error> {
336    if let Some(wal_options) = region_wal_options.get(&region_id.region_number()) {
337        let encoded = serde_json::to_string(wal_options)?;
338        options.insert(WAL_OPTIONS_KEY.to_string(), encoded);
339    }
340    Ok(())
341}
342
343/// Extracts the topic from the wal options.
344pub fn extract_topic_from_wal_options(
345    region_id: RegionId,
346    region_options: &RegionWalOptions,
347) -> Option<String> {
348    region_options
349        .get(&region_id.region_number())
350        .and_then(|wal_options| match wal_options {
351            WalOptions::Kafka(kafka_wal_option) => Some(kafka_wal_option.topic.clone()),
352            _ => None,
353        })
354}
355
356#[cfg(test)]
357mod tests {
358    use std::assert_matches;
359
360    use common_wal::config::kafka::MetasrvKafkaConfig;
361    use common_wal::config::kafka::common::KafkaTopicConfig;
362    use common_wal::maybe_skip_kafka_integration_test;
363    use common_wal::test_util::get_kafka_endpoints;
364
365    use super::*;
366    use crate::error::Error;
367    use crate::kv_backend::memory::MemoryKvBackend;
368    use crate::test_util::test_kafka_topic_pool;
369    use crate::wal_provider::selector::RoundRobinTopicSelector;
370
371    // Tests that the wal provider could successfully allocate raft-engine wal options.
372    #[tokio::test]
373    async fn test_provider_with_raft_engine() {
374        let kv_backend = Arc::new(MemoryKvBackend::new()) as KvBackendRef;
375        let wal_config = MetasrvWalConfig::RaftEngine;
376        let provider = build_wal_provider(&wal_config, kv_backend).await.unwrap();
377        provider.start().await.unwrap();
378
379        let num_regions = 32;
380        let regions = (0..num_regions).collect::<Vec<_>>();
381        let got = provider.allocate(&regions, false).await.unwrap();
382
383        let expected = regions
384            .into_iter()
385            .zip(vec![WalOptions::RaftEngine; num_regions as usize])
386            .collect();
387        assert_eq!(got, expected);
388    }
389
390    #[tokio::test]
391    async fn test_refuse_invalid_topic_name_prefix() {
392        let kv_backend = Arc::new(MemoryKvBackend::new()) as KvBackendRef;
393        let wal_config = MetasrvWalConfig::Kafka(MetasrvKafkaConfig {
394            kafka_topic: KafkaTopicConfig {
395                topic_name_prefix: "``````".to_string(),
396                ..Default::default()
397            },
398            ..Default::default()
399        });
400        let got = build_wal_provider(&wal_config, kv_backend)
401            .await
402            .unwrap_err();
403        assert_matches!(got, Error::InvalidTopicNamePrefix { .. });
404    }
405
406    #[tokio::test]
407    async fn test_provider_with_kafka_allocate_wal_options() {
408        common_telemetry::init_default_ut_logging();
409        maybe_skip_kafka_integration_test!();
410        let num_topics = 5;
411        let mut topic_pool = test_kafka_topic_pool(
412            get_kafka_endpoints(),
413            num_topics,
414            true,
415            Some("test_allocator_with_kafka"),
416        )
417        .await;
418        topic_pool.selector = Arc::new(RoundRobinTopicSelector::default());
419        let topics = topic_pool.topics.clone();
420        // clean up the topics before test
421        let topic_creator = topic_pool.topic_creator();
422        topic_creator.delete_topics(&topics).await.unwrap();
423
424        // Creates an options provider.
425        let provider = WalProvider::Kafka(topic_pool);
426        provider.start().await.unwrap();
427
428        let num_regions = 3;
429        let regions = (0..num_regions).collect::<Vec<_>>();
430        let got = provider.allocate(&regions, false).await.unwrap();
431
432        // Check the allocated wal options contain the expected topics.
433        let expected = (0..num_regions)
434            .map(|i| {
435                let options = WalOptions::Kafka(KafkaWalOptions::new(topics[i as usize].clone()));
436                (i, options)
437            })
438            .collect::<HashMap<_, _>>();
439        assert_eq!(got, expected);
440    }
441
442    #[derive(Debug, PartialEq, Serialize, Deserialize)]
443    struct RegionWalOptionsWrapper {
444        #[serde(with = "region_wal_options_serde")]
445        region_wal_options: RegionWalOptions,
446    }
447
448    #[derive(Debug, PartialEq, Serialize, Deserialize)]
449    struct OptionalRegionWalOptionsWrapper {
450        #[serde(with = "optional_region_wal_options_serde")]
451        region_wal_options: Option<RegionWalOptions>,
452    }
453
454    #[test]
455    fn test_deserialize_legacy_region_wal_options_from_encoded_map() {
456        let legacy_region_wal_options = HashMap::from([
457            (1, serde_json::to_string(&WalOptions::RaftEngine).unwrap()),
458            (
459                2,
460                serde_json::to_string(&WalOptions::Kafka(KafkaWalOptions::new(
461                    "topic_a".to_string(),
462                )))
463                .unwrap(),
464            ),
465        ]);
466        let legacy_json = serde_json::json!({
467            "region_wal_options": legacy_region_wal_options,
468        });
469
470        assert_eq!(
471            legacy_json.to_string(),
472            r#"{"region_wal_options":{"1":"{\"wal.provider\":\"raft_engine\"}","2":"{\"wal.provider\":\"kafka\",\"wal.kafka.topic\":\"topic_a\"}"}}"#
473        );
474
475        let decoded: RegionWalOptionsWrapper = serde_json::from_value(legacy_json).unwrap();
476
477        assert_eq!(
478            decoded.region_wal_options,
479            HashMap::from([
480                (1, WalOptions::RaftEngine),
481                (
482                    2,
483                    WalOptions::Kafka(KafkaWalOptions::new("topic_a".to_string())),
484                ),
485            ])
486        );
487    }
488
489    #[test]
490    fn test_deserialize_structured_region_wal_options() {
491        let json = r#"{
492            "region_wal_options": {
493                "1": {"wal.provider":"raft_engine"},
494                "2": {"wal.provider":"noop"}
495            }
496        }"#;
497
498        let decoded: RegionWalOptionsWrapper = serde_json::from_str(json).unwrap();
499
500        assert_eq!(
501            decoded.region_wal_options,
502            HashMap::from([(1, WalOptions::RaftEngine), (2, WalOptions::Noop)])
503        );
504    }
505
506    #[test]
507    fn test_serialize_legacy_region_wal_options() {
508        let wrapper = RegionWalOptionsWrapper {
509            region_wal_options: HashMap::from([(1, WalOptions::RaftEngine)]),
510        };
511
512        let encoded = serde_json::to_string(&wrapper).unwrap();
513
514        assert_eq!(
515            encoded,
516            r#"{"region_wal_options":{"1":"{\"wal.provider\":\"raft_engine\"}"}}"#
517        );
518    }
519
520    #[test]
521    fn test_serialize_optional_region_wal_options() {
522        let wrapper = OptionalRegionWalOptionsWrapper {
523            region_wal_options: Some(HashMap::from([(1, WalOptions::RaftEngine)])),
524        };
525
526        let encoded = serde_json::to_string(&wrapper).unwrap();
527
528        assert_eq!(
529            encoded,
530            r#"{"region_wal_options":{"1":"{\"wal.provider\":\"raft_engine\"}"}}"#
531        );
532        assert_eq!(
533            serde_json::from_str::<OptionalRegionWalOptionsWrapper>(&encoded).unwrap(),
534            wrapper
535        );
536
537        let none = OptionalRegionWalOptionsWrapper {
538            region_wal_options: None,
539        };
540        assert_eq!(
541            serde_json::to_string(&none).unwrap(),
542            r#"{"region_wal_options":null}"#
543        );
544    }
545
546    #[tokio::test]
547    async fn test_provider_with_object_store() {
548        let provider = WalProvider::ObjectStore {
549            prefix: "cluster-a/wal".to_string(),
550        };
551        provider.start().await.unwrap();
552
553        let regions = (0..4).collect::<Vec<_>>();
554        let got = provider.allocate(&regions, false).await.unwrap();
555        let expected = regions
556            .iter()
557            .map(|region| {
558                (
559                    *region,
560                    WalOptions::ObjectStore(ObjectStoreWalOptions::new(
561                        "cluster-a/wal".to_string(),
562                    )),
563                )
564            })
565            .collect::<HashMap<_, _>>();
566        assert_eq!(got, expected);
567
568        let got = provider.allocate(&regions, true).await.unwrap();
569        assert!(got.values().all(|options| options == &WalOptions::Noop));
570    }
571
572    #[tokio::test]
573    async fn test_provider_with_skip_wal() {
574        let provider = WalProvider::RaftEngine;
575        provider.start().await.unwrap();
576
577        let num_regions = 32;
578        let regions = (0..num_regions).collect::<Vec<_>>();
579        let got = provider.allocate(&regions, true).await.unwrap();
580        assert_eq!(got.len(), num_regions as usize);
581        for wal_options in got.values() {
582            assert_eq!(wal_options, &WalOptions::Noop);
583        }
584    }
585}
common_meta/wal_provider.rs

common_meta/
wal_provider.rs
