feat(aws_kinesis): aggregate records, stop dropping partial failures

changelog.d/aws_kinesis_streams_aggregation.feature.md

@@ -0,0 +1,3 @@
+Added optional record aggregation to the `aws_kinesis_streams` sink: many events are packed into one record as newline-delimited JSON and compressed together, measured at roughly 5x versus 2x on real CloudTrail data since Kinesis bills each record rounded up to 1 KB. Off by default; requires a compact `json` or `native_json` codec, replaces `partition_key_field` with a random key per record, and needs a consumer that splits on newlines. Also fixes both Kinesis sinks acking a partially failed batch as fully delivered, which dropped records that were never written.
+
+authors: smolaon

src/sinks/aws_kinesis/request_builder.rs

@@ -42,6 +42,29 @@ where
     metadata: RequestMetadata,
 }
 
+impl<R> KinesisRequest<R>
+where
+    R: Record,
+{
+    /// Assembles a request from an already-encoded payload.
+    ///
+    /// `metadata` is private, so this is how sibling modules (such as the
+    /// aggregating request builder) construct one.
+    pub(crate) fn new(
+        key: KinesisKey,
+        record: R,
+        finalizers: EventFinalizers,
+        metadata: RequestMetadata,
+    ) -> Self {
+        Self {
+            key,
+            record,
+            finalizers,
+            metadata,
+        }
+    }
+}
+
 impl<R> Finalizable for KinesisRequest<R>
 where
     R: Record,

src/sinks/aws_kinesis/service.rs

@@ -53,7 +53,16 @@ pub struct RecordResult {
 
 impl DriverResponse for KinesisResponse {
     fn event_status(&self) -> EventStatus {
-        EventStatus::Delivered
+        if self.failure_count > 0 {
+            // A partial failure is a 200 response carrying per-record errors --
+            // overwhelmingly throttling, which is transient. Reporting
+            // `Delivered` here acked records that were never written, so a
+            // source with acknowledgements enabled deleted its message and the
+            // records were lost outright.
+            EventStatus::Errored
+        } else {
+            EventStatus::Delivered
+        }
     }
 
     fn events_sent(&self) -> &GroupedCountByteSize {
@@ -100,3 +109,31 @@ where
         })
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn response(failure_count: usize) -> KinesisResponse {
+        KinesisResponse {
+            failure_count,
+            events_byte_size: GroupedCountByteSize::new_untagged(),
+            #[cfg(feature = "sinks-aws_kinesis_streams")]
+            failed_records: vec![],
+        }
+    }
+
+    #[test]
+    fn partial_failure_is_not_reported_as_delivered() {
+        // Acking a partial failure as `Delivered` let a source with
+        // acknowledgements delete its message for records that were never
+        // written. `Errored` marks them retriable so the message survives.
+        assert_eq!(response(1).event_status(), EventStatus::Errored);
+        assert_eq!(response(500).event_status(), EventStatus::Errored);
+    }
+
+    #[test]
+    fn full_success_is_delivered() {
+        assert_eq!(response(0).event_status(), EventStatus::Delivered);
+    }
+}

src/sinks/aws_kinesis/streams/aggregation/build.rs

@@ -0,0 +1,93 @@
+use std::marker::PhantomData;
+
+use vector_lib::codecs::encoding::{Framer, NewlineDelimitedEncoder, SerializerConfig};
+
+use super::{request_builder::AggregateRequestBuilder, sink::AggregatedKinesisSink};
+use crate::sinks::{
+    aws_kinesis::{
+        config::KinesisSinkBaseConfig,
+        record::{Record, SendRecord},
+        service::{KinesisResponse, KinesisService},
+        sink::BatchKinesisRequest,
+    },
+    prelude::*,
+};
+
+/// Builds an aggregating `aws_kinesis_streams` sink.
+///
+/// Mirrors `aws_kinesis::config::build_sink`, differing only in the encoder and
+/// request builder: a newline-delimited `Encoder<Framer>` over `Vec<Event>`
+/// instead of an unframed `Encoder<()>` over a single `Event`.
+pub fn build_aggregated_sink<C, R, RR, E, RT>(
+    config: &KinesisSinkBaseConfig,
+    batch_settings: BatcherSettings,
+    aggregate_settings: BatcherSettings,
+    client: C,
+    retry_logic: RT,
+) -> crate::Result<VectorSink>
+where
+    C: SendRecord + Clone + Send + Sync + 'static,
+    <C as SendRecord>::T: Send,
+    <C as SendRecord>::E: Send + Sync + snafu::Error,
+    Vec<<C as SendRecord>::T>: FromIterator<R>,
+    R: Send + 'static,
+    RR: Record + Record<T = R> + Clone + Send + Sync + Unpin + 'static,
+    E: Send + 'static,
+    RT: RetryLogic<Request = BatchKinesisRequest<RR>, Response = KinesisResponse> + Default,
+{
+    let request_limits = config.request.into_settings();
+
+    let region = config.region.region();
+    let service = ServiceBuilder::new()
+        .settings::<RT, BatchKinesisRequest<RR>>(request_limits, retry_logic)
+        .service(KinesisService::<C, R, E> {
+            client,
+            stream_name: config.stream_name.clone(),
+            region,
+            _phantom_t: PhantomData,
+            _phantom_e: PhantomData,
+        });
+
+    // Newline framing is what makes an aggregate splittable by the consumer,
+    // which is only sound if no event can contain a literal newline. Compact
+    // JSON escapes them as the two characters `\n`, so a raw 0x0A never
+    // reaches the payload. `pretty` is the trap: it emits real newlines, which
+    // would split one event into several unparseable fragments.
+    match config.encoding.config() {
+        SerializerConfig::NativeJson => {}
+        SerializerConfig::Json(json) if !json.options.pretty => {}
+        SerializerConfig::Json(_) => {
+            return Err("aggregation is incompatible with `encoding.json.pretty`: \
+                        pretty-printed JSON contains literal newlines, which \
+                        would split one event into several on the consumer"
+                .into());
+        }
+        _ => {
+            return Err("aggregation requires `encoding.codec` to be `json` or \
+                        `native_json`: events are newline-delimited within a \
+                        record, and any other codec may emit a literal newline \
+                        that would split one event into two on the consumer"
+                .into());
+        }
+    }
+
+    let transformer = config.encoding.transformer();
+    let serializer = config.encoding.build()?;
+
+    let encoder = Encoder::<Framer>::new(NewlineDelimitedEncoder::default().into(), serializer);
+
+    let request_builder = AggregateRequestBuilder::<RR> {
+        compression: config.compression,
+        encoder: (transformer, encoder),
+        _phantom: PhantomData,
+    };
+
+    let sink = AggregatedKinesisSink {
+        batch_settings,
+        aggregate_settings,
+        service,
+        request_builder,
+        _phantom: PhantomData,
+    };
+    Ok(VectorSink::from_event_streamsink(sink))
+}

src/sinks/aws_kinesis/streams/aggregation/config.rs

@@ -0,0 +1,103 @@
+use std::{num::NonZeroUsize, time::Duration};
+
+use vector_lib::configurable::configurable_component;
+use vector_lib::stream::BatcherSettings;
+
+/// Largest aggregate we will ever build, measured BEFORE compression.
+///
+/// A Kinesis record is capped at 1 MB on the wire. We bound the uncompressed
+/// input instead, because the compressed size is not known until after the
+/// batch is closed and `RequestBuilder` can only emit one request per batch.
+/// Staying well under 1 MB means even wholly incompressible input still fits.
+pub const MAX_AGGREGATE_BYTES: usize = 900_000;
+
+/// Aggregation settings for the `aws_kinesis_streams` sink.
+///
+/// When enabled, many events are concatenated into a single Kinesis record
+/// (newline-delimited) and compressed as one unit. Kinesis bills each record
+/// rounded up to 1 KB, so aggregation both amortizes that rounding away and
+/// lets the compressor exploit redundancy across events instead of restarting
+/// per record.
+#[configurable_component]
+#[derive(Clone, Copy, Debug)]
+#[serde(deny_unknown_fields)]
+pub struct KinesisAggregationConfig {
+    /// Whether to aggregate multiple events into each Kinesis record.
+    #[serde(default)]
+    pub enabled: bool,
+
+    /// Maximum number of events to place in a single Kinesis record.
+    #[serde(default = "default_max_events")]
+    pub max_events: usize,
+
+    /// Maximum uncompressed size, in bytes, of a single Kinesis record.
+    ///
+    /// This bounds the input to the compressor, not the resulting record, so it
+    /// must leave headroom under the 1 MB Kinesis record limit for input that
+    /// does not compress.
+    #[serde(default = "default_max_bytes")]
+    pub max_bytes: usize,
+
+    /// Maximum age, in seconds, of an aggregate before it is flushed.
+    ///
+    /// Only reached by sources that cannot fill a record within the window, so
+    /// it trades latency for aggregation on low-volume streams.
+    #[serde(default = "default_timeout_secs")]
+    pub timeout_secs: f64,
+}
+
+const fn default_max_events() -> usize {
+    500
+}
+
+const fn default_max_bytes() -> usize {
+    262_144
+}
+
+const fn default_timeout_secs() -> f64 {
+    5.0
+}
+
+impl Default for KinesisAggregationConfig {
+    fn default() -> Self {
+        Self {
+            enabled: false,
+            max_events: default_max_events(),
+            max_bytes: default_max_bytes(),
+            timeout_secs: default_timeout_secs(),
+        }
+    }
+}
+
+impl KinesisAggregationConfig {
+    /// Builds the inner batcher settings, which group events into one record.
+    ///
+    /// Deliberately does not go through `BatchConfig`: that type is clamped to
+    /// the `PutRecords` limits (500 records / 5 MB per API call), which are the
+    /// units of the *outer* batch. These are events per record.
+    pub fn into_batcher_settings(self) -> crate::Result<BatcherSettings> {
+        if self.max_events == 0 {
+            return Err("aggregation.max_events must be greater than 0".into());
+        }
+        if self.max_bytes == 0 {
+            return Err("aggregation.max_bytes must be greater than 0".into());
+        }
+        if self.max_bytes > MAX_AGGREGATE_BYTES {
+            return Err(format!(
+                "aggregation.max_bytes must be at most {MAX_AGGREGATE_BYTES} to stay \
+                 under the 1 MB Kinesis record limit once framing is added, got {}",
+                self.max_bytes
+            )
+            .into());
+        }
+        if !(self.timeout_secs.is_finite() && self.timeout_secs > 0.0) {
+            return Err("aggregation.timeout_secs must be a positive number".into());
+        }
+
+        Ok(BatcherSettings::new(
+            Duration::from_secs_f64(self.timeout_secs),
+            NonZeroUsize::new(self.max_bytes).expect("checked above"),
+            NonZeroUsize::new(self.max_events).expect("checked above"),
+        ))
+    }
+}

src/sinks/aws_kinesis/streams/aggregation/mod.rs

@@ -0,0 +1,16 @@
+//! Record aggregation for the `aws_kinesis_streams` sink.
+//!
+//! Kinesis bills each record rounded up to 1 KB, and the default sink encodes
+//! and compresses one event per record, so the compressor restarts every time
+//! and cannot exploit redundancy between events. Aggregation packs many events
+//! into one newline-delimited record compressed as a single frame, which
+//! amortizes the rounding away and materially improves the ratio.
+//!
+//! Off by default; when disabled the sink takes its original code path.
+
+mod build;
+mod config;
+mod request_builder;
+mod sink;
+
+pub use self::{build::build_aggregated_sink, config::KinesisAggregationConfig};

src/sinks/aws_kinesis/streams/aggregation/request_builder.rs

@@ -0,0 +1,91 @@
+use std::{io, marker::PhantomData};
+
+use bytes::Bytes;
+use uuid::Uuid;
+use vector_lib::{codecs::encoding::Framer, request_metadata::RequestMetadata};
+
+use crate::{
+    codecs::{Encoder, Transformer},
+    event::{Event, Finalizable},
+    sinks::{
+        aws_kinesis::{
+            record::Record,
+            request_builder::{KinesisMetadata, KinesisRequest},
+            sink::KinesisKey,
+        },
+        util::{
+            Compression, RequestBuilder, metadata::RequestMetadataBuilder,
+            request_builder::EncodeResult,
+        },
+    },
+};
+
+/// Builds one Kinesis record from *many* events.
+///
+/// The non-aggregating builder is `RequestBuilder<KinesisProcessedEvent>` with
+/// `Events = Event`, so `encode_events` constructs a fresh compressor per event
+/// and every record ends up its own compression frame. Taking `Vec<Event>` here
+/// is what collapses that to one frame per record.
+#[derive(Clone)]
+pub struct AggregateRequestBuilder<R> {
+    pub compression: Compression,
+    pub encoder: (Transformer, Encoder<Framer>),
+    pub _phantom: PhantomData<R>,
+}
+
+impl<R> RequestBuilder<Vec<Event>> for AggregateRequestBuilder<R>
+where
+    R: Record,
+{
+    type Metadata = KinesisMetadata;
+    type Events = Vec<Event>;
+    type Encoder = (Transformer, Encoder<Framer>);
+    type Payload = Bytes;
+    type Request = KinesisRequest<R>;
+    type Error = io::Error;
+
+    fn compression(&self) -> Compression {
+        self.compression
+    }
+
+    fn encoder(&self) -> &Self::Encoder {
+        &self.encoder
+    }
+
+    fn split_input(
+        &self,
+        mut events: Vec<Event>,
+    ) -> (Self::Metadata, RequestMetadataBuilder, Self::Events) {
+        let builder = RequestMetadataBuilder::from_events(&events);
+
+        // One key per record rather than per event. A random key keeps the MD5
+        // hash distribution — and therefore shard distribution — uniform. The
+        // sink's `partition_key_field` cannot be honoured here because the
+        // events in one record may disagree on it.
+        let kinesis_metadata = KinesisMetadata {
+            finalizers: events.take_finalizers(),
+            partition_key: Uuid::new_v4().to_string(),
+        };
+
+        (kinesis_metadata, builder, events)
+    }
+
+    fn build_request(
+        &self,
+        kinesis_metadata: Self::Metadata,
+        metadata: RequestMetadata,
+        payload: EncodeResult<Self::Payload>,
+    ) -> Self::Request {
+        let payload_bytes = payload.into_payload();
+        let record = R::new(&payload_bytes, &kinesis_metadata.partition_key);
+
+        KinesisRequest::new(
+            KinesisKey {
+                partition_key: kinesis_metadata.partition_key,
+            },
+            record,
+            kinesis_metadata.finalizers,
+            metadata,
+        )
+    }
+}