Defer transient zerobus failures to source via EventStatus::Errored

When the Tower retry budget is exhausted on a retryable failure (UC 5xx/408/429,
network blip, transient SDK error), the driver previously mapped the resulting
Err to EventStatus::Rejected — a permanent drop, indistinguishable from a UC
404 or an EncodingError. Ack-enabled sources and disk buffers had no way to
know a replay would have helped.

This change wires a small Tower layer outside the existing retry stack:
- Ok responses pass through unchanged.
- Err that downcasts to a retryable ZerobusSinkError (or doesn't downcast at
  all — conservative) becomes Ok(ZerobusResponse::errored()), so the driver
  sets Errored and the source / disk buffer may replay.
- Err that is permanent (EncodingError, SchemaError { retryable: false },
  MissingAckOffset, non-retryable SDK errors) still propagates so the driver
  maps it to Rejected.

ZerobusResponse now carries an EventStatus field with delivered() / errored()
constructors; ingest() returns ZerobusResponse::delivered(...).

Also brings unity_catalog_schema::status_is_retryable into line with Vector's
canonical RetryStrategy::Default by delegating to it. The visible effect: UC
501 (Not Implemented) is now classified as permanent (was retryable under the
prior blanket 5xx rule).

Co-authored-by: Isaac

Author: flaviofcruz

src/sinks/databricks_zerobus/service.rs

@@ -9,7 +9,7 @@ use databricks_zerobus_ingest_sdk::{ConnectorFactory, ProxyConnector, ZerobusSdk
 use futures::future::BoxFuture;
 use std::sync::Arc;
 use tokio::sync::{Mutex, OnceCell, RwLock};
-use tower::Service;
+use tower::{Layer, Service};
 use tracing::warn;
 use vector_lib::codecs::encoding::{
     BatchEncoder, BatchOutput, BatchSerializerConfig, ProtoBatchSerializerConfig,
@@ -71,14 +71,40 @@ pub struct ZerobusRequest {
 }
 
 /// Response type for the Zerobus service.
+///
+/// Carries the final `EventStatus` so the driver can mark finalizers correctly:
+/// `Delivered` on success, `Errored` when the retry budget was exhausted on a
+/// transient failure (asking the source / disk buffer to replay), and `Err`
+/// from `Service::call` reserved for permanent failures (driver maps to
+/// `Rejected`).
 #[derive(Debug)]
 pub struct ZerobusResponse {
     pub events_byte_size: GroupedCountByteSize,
+    pub status: vector_lib::event::EventStatus,
+}
+
+impl ZerobusResponse {
+    const fn delivered(events_byte_size: GroupedCountByteSize) -> Self {
+        Self {
+            events_byte_size,
+            status: vector_lib::event::EventStatus::Delivered,
+        }
+    }
+
+    /// Synthesize a response signalling a transient failure that exhausted the
+    /// retry budget. Carries a telemetry-aware zero `events_byte_size` because
+    /// the driver only consumes `events_sent()` on the `Delivered` path.
+    fn errored() -> Self {
+        Self {
+            events_byte_size: vector_lib::config::telemetry().create_request_count_byte_size(),
+            status: vector_lib::event::EventStatus::Errored,
+        }
+    }
 }
 
 impl DriverResponse for ZerobusResponse {
     fn event_status(&self) -> vector_lib::event::EventStatus {
-        vector_lib::event::EventStatus::Delivered
+        self.status
     }
 
     fn events_sent(&self) -> &GroupedCountByteSize {
@@ -441,7 +467,7 @@ impl ZerobusService {
         };
 
         match result {
-            Ok(()) => Ok(ZerobusResponse { events_byte_size }),
+            Ok(()) => Ok(ZerobusResponse::delivered(events_byte_size)),
             Err(e) => {
                 if e.is_retryable() {
                     // Clear the slot so the next attempt creates a fresh stream,
@@ -559,6 +585,80 @@ impl RetryLogic for ZerobusRetryLogic {
     }
 }
 
+/// Tower layer that converts retry-budget-exhausted retryable errors into a
+/// successful `ZerobusResponse` carrying `EventStatus::Errored`.
+///
+/// Wraps the retry layer from the outside. When the retry layer returns:
+/// - `Ok(resp)` — pass through unchanged.
+/// - `Err(e)` where `e.is_retryable()` — convert to `Ok(ZerobusResponse::errored())`
+///   so the driver marks finalizers `Errored` (transient — source / disk
+///   buffer may replay) rather than `Rejected` (permanent drop).
+/// - `Err(e)` permanent — propagate so the driver maps to `Rejected`.
+///
+/// Without this layer the driver maps every `Err` from `Service::call` to
+/// `EventStatus::Rejected`, which would drop transient-but-exhausted failures
+/// as if they were permanent.
+#[derive(Clone, Debug, Default)]
+pub struct RetryableErrorAsErroredLayer;
+
+impl<S> Layer<S> for RetryableErrorAsErroredLayer {
+    type Service = RetryableErrorAsErrored<S>;
+
+    fn layer(&self, inner: S) -> Self::Service {
+        RetryableErrorAsErrored { inner }
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct RetryableErrorAsErrored<S> {
+    inner: S,
+}
+
+impl<S> Service<ZerobusRequest> for RetryableErrorAsErrored<S>
+where
+    S: Service<ZerobusRequest, Response = ZerobusResponse, Error = crate::Error>,
+    S::Future: Send + 'static,
+{
+    type Response = ZerobusResponse;
+    type Error = crate::Error;
+    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
+
+    fn poll_ready(
+        &mut self,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Result<(), Self::Error>> {
+        self.inner.poll_ready(cx)
+    }
+
+    fn call(&mut self, req: ZerobusRequest) -> Self::Future {
+        let fut = self.inner.call(req);
+        Box::pin(async move {
+            match fut.await {
+                Ok(resp) => Ok(resp),
+                Err(e) => {
+                    // The Tower stack boxes errors above us (retry, timeout,
+                    // adaptive-concurrency). Downcast to inspect retryability;
+                    // anything that isn't a `ZerobusSinkError` (e.g. a timeout
+                    // `Elapsed`) is conservatively treated as transient.
+                    let retryable = match e.downcast_ref::<ZerobusSinkError>() {
+                        Some(zb) => zb.is_retryable(),
+                        None => true,
+                    };
+                    if retryable {
+                        warn!(
+                            message = "Zerobus retry budget exhausted on transient error; signaling Errored so source or buffer may replay.",
+                            error = %e,
+                        );
+                        Ok(ZerobusResponse::errored())
+                    } else {
+                        Err(e)
+                    }
+                }
+            }
+        })
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -809,4 +909,102 @@ mod tests {
         // And the slot was cleared so the next ingest creates a fresh stream.
         assert!(!service.has_active_stream().await);
     }
+
+    fn dummy_request() -> ZerobusRequest {
+        ZerobusRequest {
+            events: Arc::new(vec![]),
+            metadata: RequestMetadata::default(),
+            finalizers: EventFinalizers::default(),
+        }
+    }
+
+    #[tokio::test]
+    async fn retryable_err_after_exhaustion_becomes_ok_errored() {
+        use tower::ServiceExt;
+        let inner = tower::service_fn(|_req: ZerobusRequest| async move {
+            let err: crate::Error = Box::new(ZerobusSinkError::SchemaError {
+                message: "UC 503".to_string(),
+                retryable: true,
+            });
+            Err::<ZerobusResponse, _>(err)
+        });
+        let mut svc = RetryableErrorAsErrored { inner };
+        let resp = svc
+            .ready()
+            .await
+            .unwrap()
+            .call(dummy_request())
+            .await
+            .unwrap();
+        assert_eq!(resp.status, vector_lib::event::EventStatus::Errored);
+    }
+
+    #[tokio::test]
+    async fn non_retryable_err_propagates() {
+        use tower::ServiceExt;
+        let inner = tower::service_fn(|_req: ZerobusRequest| async move {
+            let err: crate::Error = Box::new(ZerobusSinkError::EncodingError {
+                message: "bad".to_string(),
+            });
+            Err::<ZerobusResponse, _>(err)
+        });
+        let mut svc = RetryableErrorAsErrored { inner };
+        let err = svc
+            .ready()
+            .await
+            .unwrap()
+            .call(dummy_request())
+            .await
+            .unwrap_err();
+        let zb = err.downcast_ref::<ZerobusSinkError>().unwrap();
+        assert!(matches!(zb, ZerobusSinkError::EncodingError { .. }));
+    }
+
+    #[tokio::test]
+    async fn unknown_err_treated_as_transient() {
+        use tower::ServiceExt;
+        // Simulate a Tower-layer error that isn't a ZerobusSinkError (e.g.
+        // timeout `Elapsed`): conservatively becomes Errored, not Rejected.
+        #[derive(Debug)]
+        struct Other;
+        impl std::fmt::Display for Other {
+            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+                write!(f, "other")
+            }
+        }
+        impl std::error::Error for Other {}
+
+        let inner = tower::service_fn(|_req: ZerobusRequest| async move {
+            let err: crate::Error = Box::new(Other);
+            Err::<ZerobusResponse, _>(err)
+        });
+        let mut svc = RetryableErrorAsErrored { inner };
+        let resp = svc
+            .ready()
+            .await
+            .unwrap()
+            .call(dummy_request())
+            .await
+            .unwrap();
+        assert_eq!(resp.status, vector_lib::event::EventStatus::Errored);
+    }
+
+    #[tokio::test]
+    async fn ok_response_passes_through() {
+        use tower::ServiceExt;
+        let inner = tower::service_fn(|_req: ZerobusRequest| async move {
+            Ok::<_, crate::Error>(ZerobusResponse::delivered(
+                GroupedCountByteSize::new_untagged(),
+            ))
+        });
+        let mut svc = RetryableErrorAsErrored { inner };
+        let resp = svc
+            .ready()
+            .await
+            .unwrap()
+            .call(dummy_request())
+            .await
+            .unwrap();
+        assert_eq!(resp.status, vector_lib::event::EventStatus::Delivered);
+    }
 }

src/sinks/databricks_zerobus/sink.rs

@@ -12,7 +12,9 @@ use crate::sinks::prelude::*;
 use crate::sinks::util::metadata::RequestMetadataBuilder;
 use crate::sinks::util::{RealtimeSizeBasedDefaultBatchSettings, TowerRequestSettings};
 
-use super::service::{ZerobusRequest, ZerobusRetryLogic, ZerobusService};
+use super::service::{
+    RetryableErrorAsErroredLayer, ZerobusRequest, ZerobusRetryLogic, ZerobusService,
+};
 
 /// The main Zerobus sink.
 pub struct ZerobusSink {
@@ -39,6 +41,7 @@ impl ZerobusSink {
     async fn run_inner(self: Box<Self>, input: BoxStream<'_, Event>) -> Result<(), ()> {
         let result = {
             let tower_service = ServiceBuilder::new()
+                .layer(RetryableErrorAsErroredLayer)
                 .settings(self.request_limits, ZerobusRetryLogic)
                 .service(self.service.clone());
 

src/sinks/databricks_zerobus/unity_catalog_schema.rs

@@ -18,15 +18,17 @@ use crate::http::HttpClient;
 
 /// Whether a Unity Catalog HTTP response status should be retried.
 ///
-/// 5xx, 408 (Request Timeout), and 429 (Too Many Requests) are transient;
-/// other 4xx statuses (404, 401, 403, ...) indicate permanent configuration
-/// problems that won't fix themselves on retry. The canonical list of
-/// retryable statuses for HTTP-based sinks lives at
-/// [`crate::sinks::util::http::RetryStrategy::Default`].
+/// Delegates to the canonical Vector HTTP retry policy
+/// [`crate::sinks::util::http::RetryStrategy::Default`] so this sink stays in
+/// lock-step with other HTTP-based sinks: 5xx (except 501 Not Implemented),
+/// 408 (Request Timeout), and 429 (Too Many Requests) are transient; 4xx
+/// otherwise (404, 401, 403, ...) and 501 are permanent.
 fn status_is_retryable(status: StatusCode) -> bool {
-    status.is_server_error()
-        || status == StatusCode::REQUEST_TIMEOUT
-        || status == StatusCode::TOO_MANY_REQUESTS
+    use crate::sinks::util::{http::RetryStrategy, retries::RetryAction};
+    matches!(
+        RetryStrategy::Default.retry_action::<()>(status),
+        RetryAction::Retry(_) | RetryAction::RetryPartial(_)
+    )
 }
 
 // Alias the SDK types under the names the rest of the sink already uses.
@@ -390,6 +392,24 @@ fn sanitize_package_name(name: &str) -> String {
 mod tests {
     use super::*;
 
+    #[test]
+    fn status_is_retryable_matches_canonical_policy() {
+        // Transient — must retry.
+        assert!(status_is_retryable(StatusCode::INTERNAL_SERVER_ERROR));
+        assert!(status_is_retryable(StatusCode::BAD_GATEWAY));
+        assert!(status_is_retryable(StatusCode::SERVICE_UNAVAILABLE));
+        assert!(status_is_retryable(StatusCode::GATEWAY_TIMEOUT));
+        assert!(status_is_retryable(StatusCode::REQUEST_TIMEOUT));
+        assert!(status_is_retryable(StatusCode::TOO_MANY_REQUESTS));
+        // Permanent — must not retry. 501 in particular: the server doesn't
+        // support the requested functionality; retry won't change that.
+        assert!(!status_is_retryable(StatusCode::NOT_IMPLEMENTED));
+        assert!(!status_is_retryable(StatusCode::NOT_FOUND));
+        assert!(!status_is_retryable(StatusCode::UNAUTHORIZED));
+        assert!(!status_is_retryable(StatusCode::FORBIDDEN));
+        assert!(!status_is_retryable(StatusCode::BAD_REQUEST));
+    }
+
     /// Smoke test: the wrapper calls into the SDK and builds a usable
     /// `MessageDescriptor` via the descriptor pool.
     #[test]