feat(source): support private link for kafka connector

Cargo.toml

@@ -58,6 +58,7 @@ repository = "https://github.com/risingwavelabs/risingwave"
 aws-config = { version = "0.51", default-features = false, features = ["rt-tokio", "native-tls"] }
 aws-sdk-kinesis = { version = "0.21", default-features = false, features = ["rt-tokio", "native-tls"] }
 aws-sdk-s3 = { version = "0.21", default-features = false, features = ["rt-tokio","native-tls"] }
+aws-sdk-ec2 = { version = "0.21", default-features = false, features = ["rt-tokio","native-tls"] }
 aws-sdk-sqs = { version = "0.21", default-features = false, features = ["rt-tokio", "native-tls"] }
 aws-smithy-http = "0.51"
 aws-smithy-types = "0.51"
@@ -123,3 +124,4 @@ tokio-stream = { git = "https://github.com/madsim-rs/tokio.git", rev = "0c25710"
 tokio-retry = { git = "https://github.com/madsim-rs/rust-tokio-retry.git", rev = "95e2fd3" }
 tokio-postgres = { git = "https://github.com/madsim-rs/rust-postgres.git", rev = "87ca1dc" }
 postgres-types = { git = "https://github.com/madsim-rs/rust-postgres.git", rev = "87ca1dc" }
+madsim-rdkafka = { git = "https://github.com/madsim-rs/madsim.git", rev = "06a7af8" }

proto/catalog.proto

@@ -90,6 +90,21 @@ message Sink {
   string definition = 13;
 }
 
+
+message Connection {
+  message PrivateLinkService {
+    string endpoint_id = 1;
+    map<string, string> dns_entries = 2;
+  }
+
+  uint32 id = 1;
+  string name = 2;
+  oneof info {
+    PrivateLinkService private_link_service = 3;
+  }
+
+}
+
 message Index {
   uint32 id = 1;
   uint32 schema_id = 2;

proto/ddl_service.proto

@@ -237,6 +237,34 @@ message GetDdlProgressResponse {
   repeated DdlProgress ddl_progress = 1;
 }
 
+
+message CreateConnectionRequest {
+  message PrivateLink {
+    string provider = 1;
+    string service_name = 2;
+    repeated string availability_zones = 3;
+  }
+
+  oneof payload {
+    PrivateLink private_link = 1;
+  }
+}
+
+message CreateConnectionResponse {
+  uint32 connection_id = 1;
+  // global catalog version
+  uint64 version = 2;
+}
+
+message  ListConnectionRequest {
+}
+
+message  ListConnectionResponse {
+  repeated catalog.Connection connections = 1;
+  // global catalog version
+  uint64 version = 2;
+}
+
 service DdlService {
   rpc CreateDatabase(CreateDatabaseRequest) returns (CreateDatabaseResponse);
   rpc DropDatabase(DropDatabaseRequest) returns (DropDatabaseResponse);
@@ -260,4 +288,6 @@ service DdlService {
   rpc ReplaceTablePlan(ReplaceTablePlanRequest) returns (ReplaceTablePlanResponse);
   rpc GetTable(GetTableRequest) returns (GetTableResponse);
   rpc GetDdlProgress(GetDdlProgressRequest) returns (GetDdlProgressResponse);
+  rpc CreateConnection(CreateConnectionRequest) returns (CreateConnectionResponse);
+  rpc ListConnections(ListConnectionRequest) returns (ListConnectionResponse);
 }

src/connector/Cargo.toml

@@ -18,8 +18,9 @@ anyhow = "1"
 apache-avro = { git = "https://github.com/risingwavelabs/avro", branch = "waruto/modify-decimal", features = ["snappy", "zstandard", "bzip", "xz"] }
 async-trait = "0.1"
 aws-config = { workspace = true }
+aws-sdk-ec2 = { workspace = true }
 aws-sdk-kinesis = { workspace = true } 
-aws-sdk-s3 = { workspace = true } 
+aws-sdk-s3 = { workspace = true }
 aws-smithy-http = { workspace = true }
 aws-types = { workspace = true }
 bincode = "1"

src/connector/src/common.rs

@@ -24,11 +24,27 @@ use crate::source::kinesis::config::AwsConfigInfo;
 // The file describes the common abstractions for each connector and can be used in both source and
 // sink.
 
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct AwsPrivateLinkItem {
+    pub service_name: String,
+    pub availability_zones: Vec<String>,
+    pub port: u16,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct AwsPrivateLinks {
+    pub provider: String,
+    pub infos: Vec<AwsPrivateLinkItem>,
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct KafkaCommon {
     #[serde(rename = "properties.bootstrap.server", alias = "kafka.brokers")]
     pub brokers: String,
 
+    #[serde(rename = "private.links.dns.names")]
+    pub private_link_dns_names: Option<String>,
+
     #[serde(rename = "topic", alias = "kafka.topic")]
     pub topic: String,
 

src/connector/src/source/kafka/enumerator/client.rs

@@ -16,13 +16,13 @@ use std::collections::HashMap;
 
 use anyhow::anyhow;
 use async_trait::async_trait;
-use rdkafka::consumer::{BaseConsumer, Consumer, DefaultConsumerContext};
+use rdkafka::consumer::{BaseConsumer, Consumer};
 use rdkafka::error::KafkaResult;
 use rdkafka::{Offset, TopicPartitionList};
 
 use crate::source::base::SplitEnumerator;
 use crate::source::kafka::split::KafkaSplit;
-use crate::source::kafka::{KafkaProperties, KAFKA_SYNC_CALL_TIMEOUT};
+use crate::source::kafka::{KafkaProperties, PrivateLinkConsumerContext, KAFKA_SYNC_CALL_TIMEOUT};
 
 #[derive(Debug, Copy, Clone, Eq, PartialEq)]
 pub enum KafkaEnumeratorOffset {
@@ -35,7 +35,7 @@ pub enum KafkaEnumeratorOffset {
 pub struct KafkaSplitEnumerator {
     broker_address: String,
     topic: String,
-    client: BaseConsumer,
+    client: BaseConsumer<PrivateLinkConsumerContext>,
     start_offset: KafkaEnumeratorOffset,
 
     // maybe used in the future for batch processing
@@ -54,10 +54,12 @@ impl SplitEnumerator for KafkaSplitEnumerator {
         let common_props = &properties.common;
 
         let broker_address = common_props.brokers.clone();
+        let private_links = common_props.private_link_dns_names.clone();
         let topic = common_props.topic.clone();
         config.set("bootstrap.servers", &broker_address);
         common_props.set_security_properties(&mut config);
 
+        tracing::info!("private_links: {:?}", private_links);
         let mut scan_start_offset = match properties
             .scan_startup_mode
             .as_ref()
@@ -79,7 +81,9 @@ impl SplitEnumerator for KafkaSplitEnumerator {
             scan_start_offset = KafkaEnumeratorOffset::Timestamp(time_offset)
         }
 
-        let client: BaseConsumer = config.create_with_context(DefaultConsumerContext).await?;
+        let client_ctx = PrivateLinkConsumerContext::new(&broker_address, &private_links);
+        let client: BaseConsumer<PrivateLinkConsumerContext> =
+            config.create_with_context(client_ctx).await?;
 
         Ok(Self {
             broker_address,

src/connector/src/source/kafka/mod.rs

@@ -17,16 +17,19 @@ use std::time::Duration;
 use serde::Deserialize;
 
 pub mod enumerator;
+pub mod private_link;
 pub mod source;
 pub mod split;
 
 pub use enumerator::*;
+pub use private_link::*;
 pub use source::*;
 pub use split::*;
 
 use crate::common::KafkaCommon;
-
 pub const KAFKA_CONNECTOR: &str = "kafka";
+pub const KAFKA_PROPS_BROKER_KEY: &str = "properties.bootstrap.server";
+pub const KAFKA_PROPS_BROKER_KEY_ALIAS: &str = "kafka.brokers";
 
 #[derive(Clone, Debug, Deserialize)]
 pub struct KafkaProperties {

src/connector/src/source/kafka/private_link.rs

@@ -0,0 +1,69 @@
+// Copyright 2023 RisingWave Labs
+//
+// 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::BTreeMap;
+
+use itertools::Itertools;
+use rdkafka::client::BrokerAddr;
+use rdkafka::consumer::ConsumerContext;
+use rdkafka::ClientContext;
+use risingwave_common::util::iter_util::ZipEqFast;
+
+pub struct PrivateLinkConsumerContext {
+    rewrite_map: BTreeMap<BrokerAddr, BrokerAddr>,
+}
+
+impl PrivateLinkConsumerContext {
+    pub fn new(brokers: &str, private_links: &Option<String>) -> Self {
+        let mut rewrite_map = BTreeMap::new();
+        if let Some(private_links) = private_links {
+            let dns_names = private_links.split(',').collect_vec();
+            let broker_adds = brokers.split(',').collect_vec();
+
+            broker_adds
+                .into_iter()
+                .zip_eq_fast(dns_names.into_iter())
+                .for_each(|(broker_addr, dns_name)| {
+                    let broker_addr = broker_addr.split(':').collect_vec();
+                    let dns_name = dns_name.split(':').collect_vec();
+                    let old_addr = BrokerAddr {
+                        host: broker_addr[0].to_string(),
+                        port: broker_addr[1].to_string(),
+                    };
+                    let new_addr = BrokerAddr {
+                        host: dns_name[0].to_string(),
+                        port: dns_name[1].to_string(),
+                    };
+                    rewrite_map.insert(old_addr, new_addr);
+                });
+        }
+        tracing::info!("broker addr rewrite map {:?}", rewrite_map);
+        Self { rewrite_map }
+    }
+}
+
+impl ClientContext for PrivateLinkConsumerContext {
+    fn rewrite_broker_addr(&self, addr: BrokerAddr) -> BrokerAddr {
+        match self.rewrite_map.get(&addr) {
+            None => addr,
+            Some(new_addr) => {
+                tracing::debug!("broker addr {:?} rewrote to {:?}", addr, new_addr);
+                new_addr.clone()
+            }
+        }
+    }
+}
+
+// required by the trait bound of BaseConsumer
+impl ConsumerContext for PrivateLinkConsumerContext {}

src/connector/src/source/kafka/source/reader.rs

@@ -20,13 +20,13 @@ use async_trait::async_trait;
 use futures::{StreamExt, TryStreamExt};
 use futures_async_stream::try_stream;
 use rdkafka::config::RDKafkaLogLevel;
-use rdkafka::consumer::{Consumer, DefaultConsumerContext, StreamConsumer};
+use rdkafka::consumer::{Consumer, StreamConsumer};
 use rdkafka::{ClientConfig, Message, Offset, TopicPartitionList};
 
 use crate::impl_common_split_reader_logic;
 use crate::parser::ParserConfig;
 use crate::source::base::{SourceMessage, MAX_CHUNK_SIZE};
-use crate::source::kafka::KafkaProperties;
+use crate::source::kafka::{KafkaProperties, PrivateLinkConsumerContext};
 use crate::source::{
     BoxSourceWithStateStream, Column, SourceContextRef, SplitId, SplitImpl, SplitMetaData,
     SplitReader,
@@ -35,7 +35,7 @@ use crate::source::{
 impl_common_split_reader_logic!(KafkaSplitReader, KafkaProperties);
 
 pub struct KafkaSplitReader {
-    consumer: StreamConsumer<DefaultConsumerContext>,
+    consumer: StreamConsumer<PrivateLinkConsumerContext>,
     start_offset: Option<i64>,
     stop_offset: Option<i64>,
     bytes_per_second: usize,
@@ -61,6 +61,7 @@ impl SplitReader for KafkaSplitReader {
         let mut config = ClientConfig::new();
 
         let bootstrap_servers = &properties.common.brokers;
+        let private_links = properties.common.private_link_dns_names.clone();
 
         // disable partition eof
         config.set("enable.partition.eof", "false");
@@ -83,9 +84,10 @@ impl SplitReader for KafkaSplitReader {
             );
         }
 
-        let consumer: StreamConsumer = config
+        let client_ctx = PrivateLinkConsumerContext::new(bootstrap_servers, &private_links);
+        let consumer: StreamConsumer<PrivateLinkConsumerContext> = config
             .set_log_level(RDKafkaLogLevel::Info)
-            .create_with_context(DefaultConsumerContext)
+            .create_with_context(client_ctx)
             .await
             .map_err(|e| anyhow!("failed to create kafka consumer: {}", e))?;
 

src/ctl/src/cmd_impl/meta.rs

@@ -14,10 +14,12 @@
 
 mod backup_meta;
 mod cluster_info;
+mod connection;
 mod pause_resume;
 mod reschedule;
 
 pub use backup_meta::*;
 pub use cluster_info::*;
+pub use connection::*;
 pub use pause_resume::*;
 pub use reschedule::*;