generated code

Author: anil-db

Cargo.lock

@@ -111,6 +111,7 @@ members = [
   "lib/k8s-e2e-tests",
   "lib/k8s-test-framework",
   "lib/loki-logproto",
+  "lib/my-vector-proto",
   "lib/portpicker",
   "lib/prometheus-parser",
   "lib/opentelemetry-proto",
@@ -761,6 +762,7 @@ sinks-logs = [
   "sinks-socket",
   "sinks-splunk_hec",
   "sinks-vector",
+  "sinks-my_vector",
   "sinks-webhdfs",
   "sinks-websocket",
   "sinks-websocket-server",
@@ -779,6 +781,7 @@ sinks-metrics = [
   "sinks-sematext",
   "sinks-statsd",
   "sinks-vector",
+  "sinks-my_vector",
   "sinks-splunk_hec"
 ]
 
@@ -833,6 +836,7 @@ sinks-splunk_hec = []
 sinks-statsd = ["sinks-utils-udp", "tokio-util/net"]
 sinks-utils-udp = []
 sinks-vector = ["sinks-utils-udp", "dep:tonic", "protobuf-build", "dep:prost"]
+sinks-my_vector = ["sinks-utils-udp", "dep:tonic", "protobuf-build", "dep:prost"]
 sinks-websocket = ["dep:tokio-tungstenite"]
 sinks-websocket-server = ["dep:tokio-tungstenite", "sources-utils-http-auth", "sources-utils-http-error", "sources-utils-http-prelude"]
 sinks-webhdfs = ["dep:opendal"]

Cargo.toml

@@ -0,0 +1,21 @@
+[package]
+name = "my-vector-proto"
+version = "0.1.0"
+authors = ["Vector Contributors <vector@datadoghq.com>"]
+edition = "2021"
+publish = false
+
+[build-dependencies]
+prost-build.workspace = true
+tonic-build.workspace = true
+
+[dependencies]
+bytes = { version = "1.10.1", default-features = false, features = ["serde"] }
+chrono.workspace = true
+hex = { version = "0.4.3", default-features = false, features = ["std"] }
+lookup = { package = "vector-lookup", path = "../vector-lookup", default-features = false }
+ordered-float = { version = "4.6.0", default-features = false }
+prost .workspace = true
+tonic.workspace = true
+vrl.workspace = true
+vector-core = { path = "../vector-core", default-features = false }

lib/my-vector-proto/Cargo.toml

@@ -0,0 +1,15 @@
+use std::io::Error;
+
+fn main() -> Result<(), Error> {
+    tonic_build::configure()
+        .build_client(true)
+        .build_server(true)
+        .compile(
+            &[
+                "src/proto/kafkaproducerproxy/kafkaproducerproxy.proto",
+            ],
+            &["src/proto/kafkaproducerproxy"],
+        )?;
+
+    Ok(())
+}

lib/my-vector-proto/build.rs

@@ -0,0 +1,3 @@
+pub mod kafkaproducerproxy {
+    tonic::include_proto!("kafkaproducerproxy");
+}

lib/my-vector-proto/src/lib.rs

@@ -0,0 +1,35 @@
+syntax = "proto2";
+
+package kafkaproducerproxy;
+
+service KafkaProducerProxyService {
+
+  rpc produceMessage(KafkaMessage) returns (KafkaMessage.Response) {
+  }
+
+  rpc produceMessages(KafkaMessages) returns (KafkaMessages.Response) {
+  }
+
+}
+
+message KafkaMessage {
+  optional string topic_name = 1;
+  optional string key = 2;
+  optional bytes data = 3;
+  optional bytes log_entry = 4;
+
+  message Response {
+    optional int32 error_code = 1;
+    optional string error_msg = 2;
+  }
+}
+
+message KafkaMessages {
+  repeated KafkaMessage messages = 1;
+
+  message Response {
+    repeated KafkaMessage.Response responses = 1;
+    optional bool all_succeed = 2;
+  }
+}
+

lib/my-vector-proto/src/proto/kafkaproducerproxy/kafkaproducerproxy.proto

@@ -108,6 +108,8 @@ pub mod splunk_hec;
 pub mod statsd;
 #[cfg(feature = "sinks-vector")]
 pub mod vector;
+#[cfg(feature = "sinks-my_vector")]
+pub mod my_vector;
 #[cfg(feature = "sinks-webhdfs")]
 pub mod webhdfs;
 #[cfg(feature = "sinks-websocket")]
@@ -137,6 +139,14 @@ pub enum BuildError {
 /// Common healthcheck errors
 #[derive(Debug, Snafu)]
 pub enum HealthcheckError {
-    #[snafu(display("Unexpected status: {}", status))]
-    UnexpectedStatus { status: ::http::StatusCode },
+    #[snafu(display("Unable to resolve DNS for {:?}", address))]
+    DnsFailure { address: String },
+    #[snafu(display("DNS errored {}", source))]
+    DnsError { source: crate::dns::DnsError },
+    #[snafu(display("Socket address problem: {}", source))]
+    SocketAddressError { source: std::io::Error },
+    #[snafu(display("URI parse error: {}", source))]
+    UriParseError { source: ::http::uri::InvalidUri },
+    #[snafu(display("HTTP request build error: {}", source))]
+    HTTPRequestBuilderError { source: ::http::Error },
 }

src/sinks/mod.rs

@@ -0,0 +1,152 @@
+use std::{
+    sync::{Arc, Mutex},
+    time::{Duration, Instant},
+};
+use tower::{Layer, Service};
+use vector_lib::stream::DriverResponse;
+
+/// A layer that blocks all requests for 30 seconds when any clone receives a response code 555
+#[derive(Clone)]
+pub struct BlockingLayer {
+    blocked_until: Arc<Mutex<Option<Instant>>>,
+}
+
+impl BlockingLayer {
+    pub fn new() -> Self {
+        Self {
+            blocked_until: Arc::new(Mutex::new(None)),
+        }
+    }
+}
+
+impl<S> Layer<S> for BlockingLayer {
+    type Service = BlockingService<S>;
+
+    fn layer(&self, service: S) -> Self::Service {
+        BlockingService {
+            service,
+            blocked_until: self.blocked_until.clone(),
+        }
+    }
+}
+
+#[derive(Clone)]
+pub struct BlockingService<S> {
+    service: S,
+    blocked_until: Arc<Mutex<Option<Instant>>>,
+}
+
+impl<S> BlockingService<S> {
+    pub fn new(service: S) -> Self {
+        Self {
+            service,
+            blocked_until: Arc::new(Mutex::new(None)),
+        }
+    }
+}
+
+impl<S, Request> Service<Request> for BlockingService<S>
+where
+    S: Service<Request> + Clone + Send + 'static,
+    S::Future: Send + 'static,
+    S::Response: DriverResponse + Send + 'static,
+{
+    type Response = S::Response;
+    type Error = S::Error;
+    type Future = std::pin::Pin<Box<dyn std::future::Future<Output = Result<Self::Response, Self::Error>> + Send>>;
+
+    fn poll_ready(&mut self, cx: &mut std::task::Context<'_>) -> std::task::Poll<Result<(), Self::Error>> {
+        // Check if we're currently blocked
+        if let Some(blocked_until) = *self.blocked_until.lock().unwrap() {
+            if blocked_until > Instant::now() {
+                // Still blocked, return NotReady
+                return std::task::Poll::Pending;
+            } else {
+                // Block has expired, clear it
+                *self.blocked_until.lock().unwrap() = None;
+            }
+        }
+
+        // Not blocked, check if inner service is ready
+        self.service.poll_ready(cx)
+    }
+
+    fn call(&mut self, req: Request) -> Self::Future {
+        let mut service = self.service.clone();
+        let blocked_until = self.blocked_until.clone();
+
+        Box::pin(async move {
+            // Check if we're blocked before making the request
+            if let Some(blocked_until) = *blocked_until.lock().unwrap() {
+                if blocked_until > Instant::now() {
+                    // Calculate remaining time
+                    let remaining = blocked_until.duration_since(Instant::now());
+                    tokio::time::sleep(remaining).await;
+                } else {
+                    // Block has expired, clear it
+                    *blocked_until.lock().unwrap() = None;
+                }
+            }
+
+            // Make the request
+            let response = service.call(req).await?;
+
+            // Check if we got a 555 response
+            if response.event_status().is_rejected() {
+                // Set the block for all clones
+                *blocked_until.lock().unwrap() = Some(Instant::now() + Duration::from_secs(30));
+            }
+
+            Ok(response)
+        })
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::time::Duration;
+    use tower::ServiceBuilder;
+    use vector_lib::stream::EventStatus;
+
+    struct MockResponse;
+    impl DriverResponse for MockResponse {
+        fn event_status(&self) -> EventStatus {
+            EventStatus::Rejected
+        }
+    }
+
+    struct MockService;
+    impl Service<()> for MockService {
+        type Response = MockResponse;
+        type Error = ();
+        type Future = std::pin::Pin<Box<dyn std::future::Future<Output = Result<Self::Response, Self::Error>> + Send>>;
+
+        fn poll_ready(&mut self, _: &mut std::task::Context<'_>) -> std::task::Poll<Result<(), Self::Error>> {
+            std::task::Poll::Ready(Ok(()))
+        }
+
+        fn call(&mut self, _: ()) -> Self::Future {
+            Box::pin(async move { Ok(MockResponse) })
+        }
+    }
+
+    #[tokio::test]
+    async fn test_blocking_layer() {
+        let layer = BlockingLayer::new();
+        let service = ServiceBuilder::new()
+            .layer(layer)
+            .service(MockService);
+
+        // First request should succeed but trigger blocking
+        let response = service.clone().call(()).await.unwrap();
+        assert!(response.event_status().is_rejected());
+
+        // Second request should be blocked
+        let start = Instant::now();
+        let response = service.clone().call(()).await.unwrap();
+        let duration = start.elapsed();
+        assert!(duration >= Duration::from_secs(30));
+        assert!(response.event_status().is_rejected());
+    }
+}