feat(vector sink): add multiple endpoint strategies

Author: fpytloun

changelog.d/vector_sink_multiple_backends.feature.md

@@ -0,0 +1,3 @@
+Add support for configuring multiple endpoints in the `vector` sink via the new `addresses` option, enabling built-in load balancing and failover across downstream Vector instances.
+
+authors: fpytloun

src/sinks/vector/config.rs

@@ -1,9 +1,18 @@
+use std::{
+    sync::{
+        Arc,
+        atomic::{AtomicUsize, Ordering},
+    },
+    task::{Context, Poll},
+};
+
+use futures::{FutureExt, TryFutureExt, future::BoxFuture};
 use http::Uri;
 use hyper::client::HttpConnector;
 use hyper_openssl::HttpsConnector;
 use hyper_proxy::ProxyConnector;
 use tonic::body::BoxBody;
-use tower::ServiceBuilder;
+use tower::{Service, ServiceBuilder};
 use vector_lib::configurable::configurable_component;
 
 use super::{
@@ -22,8 +31,9 @@ use crate::{
     sinks::{
         Healthcheck, VectorSink as VectorSinkType,
         util::{
-            BatchConfig, RealtimeEventBasedDefaultBatchSettings, ServiceBuilderExt,
-            TowerRequestConfig, retries::RetryLogic,
+            BatchConfig, RealtimeEventBasedDefaultBatchSettings, TowerRequestConfig,
+            retries::RetryLogic,
+            service::{HealthConfig, HealthLogic, ServiceBuilderExt},
         },
     },
     tls::{MaybeTlsSettings, TlsEnableableConfig},
@@ -45,10 +55,28 @@ pub struct VectorConfig {
     /// Both IP address and hostname are accepted formats.
     ///
     /// The address _must_ include a port.
+    ///
+    /// This option is mutually exclusive with `addresses`. Set exactly one of
+    /// `address` or `addresses`.
     #[configurable(validation(format = "uri"))]
     #[configurable(metadata(docs::examples = "92.12.333.224:6000"))]
     #[configurable(metadata(docs::examples = "https://somehost:6000"))]
-    address: String,
+    #[serde(default)]
+    address: Option<String>,
+
+    /// The downstream Vector addresses to which to connect.
+    ///
+    /// Both IP addresses and hostnames are accepted formats.
+    ///
+    /// Each address _must_ include a port.
+    ///
+    /// This option is mutually exclusive with `address`. Set exactly one of
+    /// `address` or `addresses`.
+    #[configurable(validation(format = "uri"))]
+    #[configurable(metadata(docs::examples = "92.12.333.224:6000"))]
+    #[configurable(metadata(docs::examples = "https://somehost:6000"))]
+    #[serde(default)]
+    addresses: Vec<String>,
 
     /// Compression algorithm for requests.
     ///
@@ -72,6 +100,17 @@ pub struct VectorConfig {
     #[serde(default)]
     pub request: TowerRequestConfig,
 
+    /// Options for determining the health of Vector endpoints.
+    #[serde(default)]
+    #[configurable(derived)]
+    pub endpoint_health: Option<HealthConfig>,
+
+    /// Strategy for routing requests across multiple configured addresses.
+    ///
+    /// This option is only used when `addresses` is configured.
+    #[serde(default)]
+    pub endpoint_strategy: EndpointStrategy,
+
     #[configurable(derived)]
     #[serde(default)]
     tls: Option<TlsEnableableConfig>,
@@ -102,10 +141,13 @@ impl GenerateConfig for VectorConfig {
 fn default_config(address: &str) -> VectorConfig {
     VectorConfig {
         version: None,
-        address: address.to_owned(),
+        address: Some(address.to_owned()),
+        addresses: Vec::new(),
         compression: VectorCompression::None,
         batch: BatchConfig::default(),
         request: TowerRequestConfig::default(),
+        endpoint_health: None,
+        endpoint_strategy: EndpointStrategy::default(),
         tls: None,
         acknowledgements: Default::default(),
     }
@@ -116,36 +158,73 @@ fn default_config(address: &str) -> VectorConfig {
 impl SinkConfig for VectorConfig {
     async fn build(&self, cx: SinkContext) -> crate::Result<(VectorSinkType, Healthcheck)> {
         let tls = MaybeTlsSettings::from_config(self.tls.as_ref(), false)?;
-        let uri = with_default_scheme(&self.address, tls.is_tls())?;
+        let uris = self.uris(tls.is_tls())?;
 
         let client = new_client(&tls, cx.proxy())?;
 
-        let healthcheck_uri = cx
-            .healthcheck
-            .uri
-            .clone()
-            .map(|uri| uri.uri)
-            .unwrap_or_else(|| uri.clone());
-        let healthcheck_client =
-            VectorService::new(client.clone(), healthcheck_uri, VectorCompression::None);
-        let healthcheck = healthcheck(healthcheck_client, cx.healthcheck);
-        let service = VectorService::new(client, uri, self.compression);
+        let healthcheck = healthchecks(client.clone(), &uris, cx.healthcheck);
         let request_settings = self.request.into_settings();
         let batch_settings = self.batch.into_batcher_settings()?;
 
-        let service = ServiceBuilder::new()
-            .settings(request_settings, VectorGrpcRetryLogic)
-            .service(service);
+        let services = uris
+            .into_iter()
+            .map(|uri| {
+                let endpoint = uri.to_string();
+                let service = VectorService::new(client.clone(), uri, self.compression);
+                (endpoint, service)
+            })
+            .collect::<Vec<_>>();
+
+        let sink = match self.endpoint_strategy {
+            _ if services.len() == 1 => {
+                let service = ServiceBuilder::new()
+                    .settings(request_settings, VectorGrpcRetryLogic)
+                    .service(services.into_iter().next().expect("one service").1);
+
+                VectorSinkType::from_event_streamsink(VectorSink {
+                    batch_settings,
+                    service,
+                })
+            }
+            EndpointStrategy::LoadBalance => {
+                let service = request_settings.distributed_service(
+                    VectorGrpcRetryLogic,
+                    services,
+                    self.endpoint_health.clone().unwrap_or_default(),
+                    VectorGrpcHealthLogic,
+                    1,
+                );
 
-        let sink = VectorSink {
-            batch_settings,
-            service,
+                VectorSinkType::from_event_streamsink(VectorSink {
+                    batch_settings,
+                    service,
+                })
+            }
+            EndpointStrategy::Failover => {
+                let endpoint_timeout = request_settings.timeout;
+                let mut failover_request_settings = request_settings;
+                failover_request_settings.timeout = endpoint_timeout
+                    .checked_mul((services.len() + 1) as u32)
+                    .unwrap_or(endpoint_timeout);
+
+                let service = ServiceBuilder::new()
+                    .settings(failover_request_settings, VectorGrpcRetryLogic)
+                    .service(FailoverVectorService::new(
+                        services
+                            .into_iter()
+                            .map(|(_endpoint, service)| service)
+                            .collect(),
+                        endpoint_timeout,
+                    ));
+
+                VectorSinkType::from_event_streamsink(VectorSink {
+                    batch_settings,
+                    service,
+                })
+            }
         };
 
-        Ok((
-            VectorSinkType::from_event_streamsink(sink),
-            Box::pin(healthcheck),
-        ))
+        Ok((sink, Box::pin(healthcheck)))
     }
 
     fn input(&self) -> Input {
@@ -157,6 +236,124 @@ impl SinkConfig for VectorConfig {
     }
 }
 
+/// Strategy for routing requests across multiple Vector endpoints.
+#[configurable_component]
+#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
+#[serde(rename_all = "snake_case")]
+pub enum EndpointStrategy {
+    /// Distribute requests across healthy endpoints.
+    #[default]
+    LoadBalance,
+    /// Use endpoints in configured order, moving to the next endpoint only when
+    /// the current endpoint fails.
+    Failover,
+}
+
+#[derive(Clone)]
+struct FailoverVectorService {
+    services: Vec<VectorService>,
+    active: Arc<AtomicUsize>,
+    endpoint_timeout: std::time::Duration,
+}
+
+impl FailoverVectorService {
+    fn new(services: Vec<VectorService>, endpoint_timeout: std::time::Duration) -> Self {
+        Self {
+            services,
+            active: Arc::new(AtomicUsize::new(0)),
+            endpoint_timeout,
+        }
+    }
+}
+
+impl Service<VectorRequest> for FailoverVectorService {
+    type Response = VectorResponse;
+    type Error = crate::Error;
+    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
+
+    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        Poll::Ready(Ok(()))
+    }
+
+    fn call(&mut self, request: VectorRequest) -> Self::Future {
+        let services = self.services.clone();
+        let active = Arc::clone(&self.active);
+        let endpoint_timeout = self.endpoint_timeout;
+
+        Box::pin(async move {
+            let start = active.load(Ordering::Acquire);
+            let mut last_error = None;
+
+            for offset in 0..services.len() {
+                let index = (start + offset) % services.len();
+                let next_index = (index + 1) % services.len();
+                let mut service = services[index].clone();
+
+                match tokio::time::timeout(endpoint_timeout, service.call(request.clone())).await {
+                    Ok(Ok(response)) => {
+                        let current = active.load(Ordering::Acquire);
+                        if current == start || current == index {
+                            active.store(index, Ordering::Release);
+                        }
+                        return Ok(response);
+                    }
+                    Ok(Err(error)) => {
+                        if !is_retriable_vector_error(&error) {
+                            return Err(error);
+                        }
+
+                        let _ = active.compare_exchange(
+                            index,
+                            next_index,
+                            Ordering::AcqRel,
+                            Ordering::Acquire,
+                        );
+                        last_error = Some(error);
+                    }
+                    Err(_elapsed) => {
+                        let _ = active.compare_exchange(
+                            index,
+                            next_index,
+                            Ordering::AcqRel,
+                            Ordering::Acquire,
+                        );
+                        last_error = Some(Box::new(VectorSinkError::Request {
+                            source: tonic::Status::deadline_exceeded(
+                                "vector endpoint request timed out",
+                            ),
+                        }) as crate::Error);
+                    }
+                }
+            }
+
+            Err(last_error.expect("failover service should have at least one endpoint"))
+        })
+    }
+}
+
+fn is_retriable_vector_error(error: &crate::Error) -> bool {
+    error
+        .downcast_ref::<VectorSinkError>()
+        .is_none_or(|error| VectorGrpcRetryLogic.is_retriable_error(error))
+}
+
+impl VectorConfig {
+    fn uris(&self, tls: bool) -> crate::Result<Vec<Uri>> {
+        match (self.address.as_ref(), self.addresses.as_slice()) {
+            (Some(_), [_first, ..]) => Err(
+                "`address` and `addresses` options are mutually exclusive. Please use `addresses` for multiple Vector endpoints."
+                    .into(),
+            ),
+            (None, []) => Err("No Vector endpoint configured. Please set `address` or `addresses`.".into()),
+            (Some(address), []) => Ok(vec![with_default_scheme(address, tls)?]),
+            (None, addresses) => addresses
+                .iter()
+                .map(|address| with_default_scheme(address, tls))
+                .collect(),
+        }
+    }
+}
+
 /// Check to see if the remote service accepts new events.
 async fn healthcheck(
     mut service: VectorService,
@@ -183,6 +380,39 @@ async fn healthcheck(
     }
 }
 
+fn healthchecks(
+    client: hyper::Client<ProxyConnector<HttpsConnector<HttpConnector>>, BoxBody>,
+    uris: &[Uri],
+    options: SinkHealthcheckOptions,
+) -> Healthcheck {
+    if !options.enabled {
+        return Box::pin(futures::future::ok(()));
+    }
+
+    let healthcheck_uris = options
+        .uri
+        .clone()
+        .map(|uri| vec![uri.uri])
+        .unwrap_or_else(|| uris.to_vec());
+
+    Box::pin(
+        futures::future::select_ok(healthcheck_uris.into_iter().map(move |uri| {
+            let service = VectorService::new(client.clone(), uri, VectorCompression::None);
+            let timeout = options.timeout;
+            healthcheck(
+                service,
+                SinkHealthcheckOptions {
+                    enabled: true,
+                    uri: None,
+                    timeout,
+                },
+            )
+            .boxed()
+        }))
+        .map_ok(|((), _)| ()),
+    )
+}
+
 /// grpc doesn't like an address without a scheme, so we default to http or https if one isn't
 /// specified in the address.
 pub fn with_default_scheme(address: &str, tls: bool) -> crate::Result<Uri> {
@@ -256,3 +486,15 @@ impl RetryLogic for VectorGrpcRetryLogic {
         }
     }
 }
+
+#[derive(Debug, Clone)]
+struct VectorGrpcHealthLogic;
+
+impl HealthLogic for VectorGrpcHealthLogic {
+    type Error = crate::Error;
+    type Response = VectorResponse;
+
+    fn is_healthy(&self, response: &Result<Self::Response, Self::Error>) -> Option<bool> {
+        Some(response.is_ok())
+    }
+}