feat(kafka sink): add rate limit configuration for kafka sink (vectordotdev#23196)

* feature(kafka sink): add rate limit configuration for `kafka` sink

Related: vectordotdev#21879

* Add changelog entry

---------

Co-authored-by: Pavlos Rontidis <pavlos.rontidis@gmail.com>

Author: 2 people

changelog.d/23196_kafka_sink_rate_limit.feature.md

@@ -0,0 +1,3 @@
+Added `rate_limit_num` and `rate_limit_duration_secs` options to `kafka` sink, to enable rate limiting this sink.
+
+authors: esensar Quad9DNS

src/sinks/kafka/config.rs

@@ -96,6 +96,18 @@ pub struct KafkaSinkConfig {
     #[configurable(metadata(docs::advanced))]
     pub message_timeout_ms: Duration,
 
+    /// The time window used for the `rate_limit_num` option.
+    #[configurable(metadata(docs::type_unit = "seconds"))]
+    #[configurable(metadata(docs::human_name = "Rate Limit Duration"))]
+    #[serde(default = "default_rate_limit_duration_secs")]
+    pub rate_limit_duration_secs: u64,
+
+    /// The maximum number of requests allowed within the `rate_limit_duration_secs` time window.
+    #[configurable(metadata(docs::type_unit = "requests"))]
+    #[configurable(metadata(docs::human_name = "Rate Limit Number"))]
+    #[serde(default = "default_rate_limit_num")]
+    pub rate_limit_num: u64,
+
     /// A map of advanced options to pass directly to the underlying `librdkafka` client.
     ///
     /// For more information on configuration options, see [Configuration properties][config_props_docs].
@@ -134,6 +146,14 @@ const fn default_message_timeout_ms() -> Duration {
     Duration::from_millis(300000) // default in librdkafka
 }
 
+const fn default_rate_limit_duration_secs() -> u64 {
+    1
+}
+
+const fn default_rate_limit_num() -> u64 {
+    i64::MAX as u64 // i64 avoids TOML deserialize issue
+}
+
 fn example_librdkafka_options() -> HashMap<String, String> {
     HashMap::<_, _>::from_iter([
         ("client.id".to_string(), "${ENV_VAR}".to_string()),
@@ -245,6 +265,8 @@ impl GenerateConfig for KafkaSinkConfig {
             auth: Default::default(),
             socket_timeout_ms: default_socket_timeout_ms(),
             message_timeout_ms: default_message_timeout_ms(),
+            rate_limit_duration_secs: default_rate_limit_duration_secs(),
+            rate_limit_num: default_rate_limit_num(),
             librdkafka_options: Default::default(),
             headers_key: None,
             acknowledgements: Default::default(),

src/sinks/kafka/sink.rs

@@ -1,9 +1,12 @@
+use std::time::Duration;
+
 use rdkafka::{
     error::KafkaError,
     producer::{BaseProducer, FutureProducer, Producer},
     ClientConfig,
 };
 use snafu::{ResultExt, Snafu};
+use tower::limit::RateLimit;
 use tracing::Span;
 use vrl::path::OwnedTargetPath;
 
@@ -27,7 +30,7 @@ pub(super) enum BuildError {
 pub struct KafkaSink {
     transformer: Transformer,
     encoder: Encoder<()>,
-    service: KafkaService,
+    service: RateLimit<KafkaService>,
     topic: Template,
     key_field: Option<OwnedTargetPath>,
     headers_key: Option<OwnedTargetPath>,
@@ -57,7 +60,12 @@ impl KafkaSink {
             headers_key: config.headers_key.map(|key| key.0),
             transformer,
             encoder,
-            service: KafkaService::new(producer),
+            service: ServiceBuilder::new()
+                .rate_limit(
+                    config.rate_limit_num,
+                    Duration::from_secs(config.rate_limit_duration_secs),
+                )
+                .service(KafkaService::new(producer)),
             topic: config.topic,
             key_field: config.key_field.map(|key| key.0),
         })

src/sinks/kafka/tests.rs

@@ -59,6 +59,8 @@ mod integration_test {
             auth: KafkaAuthConfig::default(),
             socket_timeout_ms: Duration::from_millis(60000),
             message_timeout_ms: Duration::from_millis(300000),
+            rate_limit_duration_secs: 1,
+            rate_limit_num: i64::MAX as u64,
             librdkafka_options: HashMap::new(),
             headers_key: None,
             acknowledgements: Default::default(),
@@ -85,6 +87,8 @@ mod integration_test {
             auth: KafkaAuthConfig::default(),
             socket_timeout_ms: Duration::from_millis(60000),
             message_timeout_ms: Duration::from_millis(300000),
+            rate_limit_duration_secs: 1,
+            rate_limit_num: i64::MAX as u64,
             librdkafka_options: HashMap::new(),
             headers_key: None,
             acknowledgements: Default::default(),
@@ -185,6 +189,8 @@ mod integration_test {
             },
             socket_timeout_ms: Duration::from_millis(60000),
             message_timeout_ms: Duration::from_millis(300000),
+            rate_limit_duration_secs: 1,
+            rate_limit_num: i64::MAX as u64,
             batch,
             librdkafka_options,
             headers_key: None,
@@ -334,6 +340,8 @@ mod integration_test {
             auth: kafka_auth.clone(),
             socket_timeout_ms: Duration::from_millis(60000),
             message_timeout_ms: Duration::from_millis(300000),
+            rate_limit_duration_secs: 1,
+            rate_limit_num: i64::MAX as u64,
             librdkafka_options: HashMap::new(),
             headers_key: Some(headers_key.clone()),
             acknowledgements: Default::default(),

website/cue/reference/components/sinks/base/kafka.cue

@@ -489,6 +489,22 @@ base: components: sinks: kafka: configuration: {
 			unit: "milliseconds"
 		}
 	}
+	rate_limit_duration_secs: {
+		description: "The time window used for the `rate_limit_num` option."
+		required:    false
+		type: uint: {
+			default: 1
+			unit:    "seconds"
+		}
+	}
+	rate_limit_num: {
+		description: "The maximum number of requests allowed within the `rate_limit_duration_secs` time window."
+		required:    false
+		type: uint: {
+			default: 9223372036854775807
+			unit:    "requests"
+		}
+	}
 	sasl: {
 		description: "Configuration for SASL authentication when interacting with Kafka."
 		required:    false