From 33b7c17b3a7907d525ef01229291b8950cfa0f71 Mon Sep 17 00:00:00 2001
From: Stephen Wakely <fungus.humungus@gmail.com>
Date: Wed, 16 Aug 2023 11:31:10 +0100
Subject: [PATCH 01/32] Refactor prometheus remote write sink

Signed-off-by: Stephen Wakely <fungus.humungus@gmail.com>
---
 src/sinks/prometheus/remote_write.rs          | 791 ------------------
 src/sinks/prometheus/remote_write/config.rs   | 206 +++++
 .../remote_write/integration_tests.rs         | 111 +++
 src/sinks/prometheus/remote_write/mod.rs      |  55 ++
 .../remote_write/request_builder.rs           | 198 +++++
 src/sinks/prometheus/remote_write/service.rs  | 120 +++
 src/sinks/prometheus/remote_write/sink.rs     | 147 ++++
 src/sinks/prometheus/remote_write/tests.rs    | 230 +++++
 8 files changed, 1067 insertions(+), 791 deletions(-)
 delete mode 100644 src/sinks/prometheus/remote_write.rs
 create mode 100644 src/sinks/prometheus/remote_write/config.rs
 create mode 100644 src/sinks/prometheus/remote_write/integration_tests.rs
 create mode 100644 src/sinks/prometheus/remote_write/mod.rs
 create mode 100644 src/sinks/prometheus/remote_write/request_builder.rs
 create mode 100644 src/sinks/prometheus/remote_write/service.rs
 create mode 100644 src/sinks/prometheus/remote_write/sink.rs
 create mode 100644 src/sinks/prometheus/remote_write/tests.rs

diff --git a/src/sinks/prometheus/remote_write.rs b/src/sinks/prometheus/remote_write.rs
deleted file mode 100644
index a1e9b792fed4d..0000000000000
--- a/src/sinks/prometheus/remote_write.rs
+++ /dev/null
@@ -1,791 +0,0 @@
-use std::io::Read;
-use std::sync::Arc;
-use std::task;
-
-use aws_credential_types::provider::SharedCredentialsProvider;
-use aws_types::region::Region;
-use bytes::{Bytes, BytesMut};
-use futures::{future::BoxFuture, stream, FutureExt, SinkExt};
-use http::{Request, Uri};
-use prost::Message;
-use snafu::{ResultExt, Snafu};
-use tower::Service;
-use vector_config::configurable_component;
-use vector_core::{ByteSizeOf, EstimatedJsonEncodedSizeOf};
-
-use super::collector::{self, MetricCollector as _};
-use crate::{
-    aws::RegionOrEndpoint,
-    config::{self, AcknowledgementsConfig, Input, SinkConfig},
-    event::{Event, Metric},
-    http::{Auth, HttpClient},
-    internal_events::{EndpointBytesSent, TemplateRenderingError},
-    sinks::{
-        self,
-        prometheus::PrometheusRemoteWriteAuth,
-        util::{
-            batch::BatchConfig,
-            buffer::metrics::{MetricNormalize, MetricNormalizer, MetricSet, MetricsBuffer},
-            http::HttpRetryLogic,
-            uri, EncodedEvent, PartitionBuffer, PartitionInnerBuffer, SinkBatchSettings,
-            TowerRequestConfig,
-        },
-    },
-    template::Template,
-    tls::{TlsConfig, TlsSettings},
-};
-
-#[derive(Clone, Copy, Debug, Default)]
-pub struct PrometheusRemoteWriteDefaultBatchSettings;
-
-impl SinkBatchSettings for PrometheusRemoteWriteDefaultBatchSettings {
-    const MAX_EVENTS: Option<usize> = Some(1_000);
-    const MAX_BYTES: Option<usize> = None;
-    const TIMEOUT_SECS: f64 = 1.0;
-}
-
-#[derive(Debug, Snafu)]
-enum Errors {
-    #[snafu(display(r#"Prometheus remote_write sink cannot accept "set" metrics"#))]
-    SetMetricInvalid,
-    #[snafu(display("aws.region required when AWS authentication is in use"))]
-    AwsRegionRequired,
-}
-
-/// Configuration for the `prometheus_remote_write` sink.
-#[configurable_component(sink(
-    "prometheus_remote_write",
-    "Deliver metric data to a Prometheus remote write endpoint."
-))]
-#[derive(Clone, Debug, Default)]
-#[serde(deny_unknown_fields)]
-pub struct RemoteWriteConfig {
-    /// The endpoint to send data to.
-    ///
-    /// The endpoint should include the scheme and the path to write to.
-    #[configurable(metadata(docs::examples = "https://localhost:8087/api/v1/write"))]
-    pub endpoint: String,
-
-    /// The default namespace for any metrics sent.
-    ///
-    /// This namespace is only used if a metric has no existing namespace. When a namespace is
-    /// present, it is used as a prefix to the metric name, and separated with an underscore (`_`).
-    ///
-    /// It should follow the Prometheus [naming conventions][prom_naming_docs].
-    ///
-    /// [prom_naming_docs]: https://prometheus.io/docs/practices/naming/#metric-names
-    #[configurable(metadata(docs::examples = "service"))]
-    #[configurable(metadata(docs::advanced))]
-    pub default_namespace: Option<String>,
-
-    /// Default buckets to use for aggregating [distribution][dist_metric_docs] metrics into histograms.
-    ///
-    /// [dist_metric_docs]: https://vector.dev/docs/about/under-the-hood/architecture/data-model/metric/#distribution
-    #[serde(default = "super::default_histogram_buckets")]
-    #[configurable(metadata(docs::advanced))]
-    pub buckets: Vec<f64>,
-
-    /// Quantiles to use for aggregating [distribution][dist_metric_docs] metrics into a summary.
-    ///
-    /// [dist_metric_docs]: https://vector.dev/docs/about/under-the-hood/architecture/data-model/metric/#distribution
-    #[serde(default = "super::default_summary_quantiles")]
-    #[configurable(metadata(docs::advanced))]
-    pub quantiles: Vec<f64>,
-
-    #[configurable(derived)]
-    #[serde(default)]
-    pub batch: BatchConfig<PrometheusRemoteWriteDefaultBatchSettings>,
-
-    #[configurable(derived)]
-    #[serde(default)]
-    pub request: TowerRequestConfig,
-
-    /// The tenant ID to send.
-    ///
-    /// If set, a header named `X-Scope-OrgID` is added to outgoing requests with the value of this setting.
-    ///
-    /// This may be used by Cortex or other remote services to identify the tenant making the request.
-    #[serde(default)]
-    #[configurable(metadata(docs::examples = "my-domain"))]
-    #[configurable(metadata(docs::advanced))]
-    pub tenant_id: Option<Template>,
-
-    #[configurable(derived)]
-    pub tls: Option<TlsConfig>,
-
-    #[configurable(derived)]
-    pub auth: Option<PrometheusRemoteWriteAuth>,
-
-    #[configurable(derived)]
-    #[configurable(metadata(docs::advanced))]
-    pub aws: Option<RegionOrEndpoint>,
-
-    #[configurable(derived)]
-    #[serde(
-        default,
-        deserialize_with = "crate::serde::bool_or_struct",
-        skip_serializing_if = "crate::serde::skip_serializing_if_default"
-    )]
-    pub acknowledgements: AcknowledgementsConfig,
-
-    #[configurable(derived)]
-    #[configurable(metadata(docs::advanced))]
-    #[serde(default)]
-    pub compression: Compression,
-}
-
-impl_generate_config_from_default!(RemoteWriteConfig);
-
-/// Supported compression types for Prometheus Remote Write.
-#[configurable_component]
-#[derive(Clone, Copy, Debug, Derivative)]
-#[derivative(Default)]
-#[serde(rename_all = "lowercase")]
-pub enum Compression {
-    /// Snappy.
-    #[derivative(Default)]
-    Snappy,
-
-    /// Gzip.
-    Gzip,
-
-    /// Zstandard.
-    Zstd,
-}
-
-const fn convert_compression_to_content_encoding(compression: Compression) -> &'static str {
-    match compression {
-        Compression::Snappy => "snappy",
-        Compression::Gzip => "gzip",
-        Compression::Zstd => "zstd",
-    }
-}
-
-#[async_trait::async_trait]
-#[typetag::serde(name = "prometheus_remote_write")]
-impl SinkConfig for RemoteWriteConfig {
-    async fn build(
-        &self,
-        cx: config::SinkContext,
-    ) -> crate::Result<(sinks::VectorSink, sinks::Healthcheck)> {
-        let endpoint = self.endpoint.parse::<Uri>().context(sinks::UriParseSnafu)?;
-        let tls_settings = TlsSettings::from_options(&self.tls)?;
-        let batch = self.batch.into_batch_settings()?;
-        let request_settings = self.request.unwrap_with(&TowerRequestConfig::default());
-        let buckets = self.buckets.clone();
-        let quantiles = self.quantiles.clone();
-
-        let client = HttpClient::new(tls_settings, cx.proxy())?;
-        let tenant_id = self.tenant_id.clone();
-
-        let (http_auth, credentials_provider, aws_region) = match &self.auth {
-            Some(PrometheusRemoteWriteAuth::Basic { user, password }) => (
-                Some(Auth::Basic {
-                    user: user.clone(),
-                    password: password.clone().into(),
-                }),
-                None,
-                None,
-            ),
-            Some(PrometheusRemoteWriteAuth::Bearer { token }) => (
-                Some(Auth::Bearer {
-                    token: token.clone(),
-                }),
-                None,
-                None,
-            ),
-            Some(PrometheusRemoteWriteAuth::Aws(aws_auth)) => {
-                let region = self
-                    .aws
-                    .as_ref()
-                    .map(|config| config.region())
-                    .ok_or(Errors::AwsRegionRequired)?
-                    .ok_or(Errors::AwsRegionRequired)?;
-
-                (
-                    None,
-                    Some(aws_auth.credentials_provider(region.clone()).await?),
-                    Some(region),
-                )
-            }
-            None => (None, None, None),
-        };
-
-        let http_request_builder = Arc::new(HttpRequestBuilder {
-            endpoint: endpoint.clone(),
-            aws_region,
-            credentials_provider,
-            http_auth,
-            compression: self.compression,
-        });
-
-        let healthcheck = healthcheck(client.clone(), Arc::clone(&http_request_builder)).boxed();
-        let service = RemoteWriteService {
-            default_namespace: self.default_namespace.clone(),
-            client,
-            buckets,
-            quantiles,
-            http_request_builder,
-            compression: self.compression,
-        };
-
-        let sink = {
-            let buffer = PartitionBuffer::new(MetricsBuffer::new(batch.size));
-            let mut normalizer = MetricNormalizer::<PrometheusMetricNormalize>::default();
-
-            request_settings
-                .partition_sink(HttpRetryLogic, service, buffer, batch.timeout)
-                .with_flat_map(move |event: Event| {
-                    let byte_size = event.size_of();
-                    let json_size = event.estimated_json_encoded_size_of();
-
-                    stream::iter(normalizer.normalize(event.into_metric()).map(|event| {
-                        let tenant_id = tenant_id.as_ref().and_then(|template| {
-                            template
-                                .render_string(&event)
-                                .map_err(|error| {
-                                    emit!(TemplateRenderingError {
-                                        error,
-                                        field: Some("tenant_id"),
-                                        drop_event: true,
-                                    })
-                                })
-                                .ok()
-                        });
-                        let key = PartitionKey { tenant_id };
-                        Ok(EncodedEvent::new(
-                            PartitionInnerBuffer::new(event, key),
-                            byte_size,
-                            json_size,
-                        ))
-                    }))
-                })
-                .sink_map_err(
-                    |error| error!(message = "Prometheus remote_write sink error.", %error),
-                )
-        };
-
-        #[allow(deprecated)]
-        Ok((sinks::VectorSink::from_event_sink(sink), healthcheck))
-    }
-
-    fn input(&self) -> Input {
-        Input::metric()
-    }
-
-    fn acknowledgements(&self) -> &AcknowledgementsConfig {
-        &self.acknowledgements
-    }
-}
-
-#[derive(Clone, Eq, Hash, PartialEq)]
-struct PartitionKey {
-    tenant_id: Option<String>,
-}
-
-async fn healthcheck(
-    client: HttpClient,
-    http_request_builder: Arc<HttpRequestBuilder>,
-) -> crate::Result<()> {
-    let body = bytes::Bytes::new();
-    let request = http_request_builder
-        .build_request(http::Method::GET, body.into(), None)
-        .await?;
-    let response = client.send(request).await?;
-
-    match response.status() {
-        http::StatusCode::OK => Ok(()),
-        other => Err(sinks::HealthcheckError::UnexpectedStatus { status: other }.into()),
-    }
-}
-
-#[derive(Default)]
-pub struct PrometheusMetricNormalize;
-
-impl MetricNormalize for PrometheusMetricNormalize {
-    fn normalize(&mut self, state: &mut MetricSet, metric: Metric) -> Option<Metric> {
-        state.make_absolute(metric)
-    }
-}
-
-#[derive(Clone)]
-struct RemoteWriteService {
-    default_namespace: Option<String>,
-    client: HttpClient,
-    buckets: Vec<f64>,
-    quantiles: Vec<f64>,
-    http_request_builder: Arc<HttpRequestBuilder>,
-    compression: Compression,
-}
-
-impl RemoteWriteService {
-    fn encode_events(&self, metrics: Vec<Metric>) -> Bytes {
-        let mut time_series = collector::TimeSeries::new();
-        for metric in metrics {
-            time_series.encode_metric(
-                self.default_namespace.as_deref(),
-                &self.buckets,
-                &self.quantiles,
-                &metric,
-            );
-        }
-        let request = time_series.finish();
-
-        let mut out = BytesMut::with_capacity(request.encoded_len());
-        request.encode(&mut out).expect("Out of memory");
-        out.freeze()
-    }
-}
-
-impl Service<PartitionInnerBuffer<Vec<Metric>, PartitionKey>> for RemoteWriteService {
-    type Response = http::Response<Bytes>;
-    type Error = crate::Error;
-    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
-
-    // Emission of an internal event in case of errors is handled upstream by the caller.
-    fn poll_ready(&mut self, _task: &mut task::Context<'_>) -> task::Poll<Result<(), Self::Error>> {
-        task::Poll::Ready(Ok(()))
-    }
-
-    // Emission of internal events for errors and dropped events is handled upstream by the caller.
-    fn call(&mut self, buffer: PartitionInnerBuffer<Vec<Metric>, PartitionKey>) -> Self::Future {
-        let (events, key) = buffer.into_parts();
-        let body = self.encode_events(events);
-        let body = compress_block(self.compression, body);
-
-        let client = self.client.clone();
-        let request_builder = Arc::clone(&self.http_request_builder);
-
-        Box::pin(async move {
-            let request = request_builder
-                .build_request(http::Method::POST, body, key.tenant_id)
-                .await?;
-
-            let (protocol, endpoint) = uri::protocol_endpoint(request.uri().clone());
-
-            let response = client.send(request).await?;
-            let (parts, body) = response.into_parts();
-            let body = hyper::body::to_bytes(body).await?;
-
-            emit!(EndpointBytesSent {
-                byte_size: body.len(),
-                protocol: &protocol,
-                endpoint: &endpoint
-            });
-
-            Ok(hyper::Response::from_parts(parts, body))
-        })
-    }
-}
-
-pub struct HttpRequestBuilder {
-    pub endpoint: Uri,
-    pub aws_region: Option<Region>,
-    pub http_auth: Option<Auth>,
-    pub credentials_provider: Option<SharedCredentialsProvider>,
-    pub compression: Compression,
-}
-
-impl HttpRequestBuilder {
-    pub async fn build_request(
-        &self,
-        method: http::Method,
-        body: Vec<u8>,
-        tenant_id: Option<String>,
-    ) -> Result<Request<hyper::Body>, crate::Error> {
-        let content_encoding = convert_compression_to_content_encoding(self.compression);
-
-        let mut builder = http::Request::builder()
-            .method(method)
-            .uri(self.endpoint.clone())
-            .header("X-Prometheus-Remote-Write-Version", "0.1.0")
-            .header("Content-Encoding", content_encoding)
-            .header("Content-Type", "application/x-protobuf");
-
-        if let Some(tenant_id) = &tenant_id {
-            builder = builder.header("X-Scope-OrgID", tenant_id);
-        }
-
-        let mut request = builder.body(body.into()).unwrap();
-        if let Some(http_auth) = &self.http_auth {
-            http_auth.apply(&mut request);
-        }
-
-        if let Some(credentials_provider) = &self.credentials_provider {
-            sign_request(&mut request, credentials_provider, &self.aws_region).await?;
-        }
-
-        let (parts, body) = request.into_parts();
-        let request: Request<hyper::Body> = hyper::Request::from_parts(parts, body.into());
-
-        Ok(request)
-    }
-}
-
-fn compress_block(compression: Compression, data: Bytes) -> Vec<u8> {
-    match compression {
-        Compression::Snappy => snap::raw::Encoder::new()
-            .compress_vec(&data)
-            .expect("snap compression failed, please report"),
-        Compression::Gzip => {
-            let mut buf = Vec::new();
-            flate2::read::GzEncoder::new(data.as_ref(), flate2::Compression::default())
-                .read_to_end(&mut buf)
-                .expect("gzip compression failed, please report");
-            buf
-        }
-        Compression::Zstd => {
-            zstd::encode_all(data.as_ref(), 0).expect("zstd compression failed, please report")
-        }
-    }
-}
-
-async fn sign_request(
-    request: &mut http::Request<Bytes>,
-    credentials_provider: &SharedCredentialsProvider,
-    region: &Option<Region>,
-) -> crate::Result<()> {
-    crate::aws::sign_request("aps", request, credentials_provider, region).await
-}
-
-#[cfg(test)]
-mod tests {
-    use futures::StreamExt;
-    use http::HeaderMap;
-    use indoc::indoc;
-    use prometheus_parser::proto;
-    use vector_core::metric_tags;
-
-    use super::*;
-    use crate::{
-        config::SinkContext,
-        event::{MetricKind, MetricValue},
-        sinks::util::test::build_test_server,
-        test_util::{
-            self,
-            components::{assert_sink_compliance, HTTP_SINK_TAGS},
-        },
-    };
-
-    #[test]
-    fn generate_config() {
-        crate::test_util::test_generate_config::<RemoteWriteConfig>();
-    }
-
-    macro_rules! labels {
-        ( $( $name:expr => $value:expr ),* ) => {
-            vec![ $( proto::Label {
-                name: $name.to_string(),
-                value: $value.to_string()
-            }, )* ]
-        }
-    }
-
-    #[tokio::test]
-    async fn sends_request() {
-        let outputs = send_request("", vec![create_event("gauge-2".into(), 32.0)]).await;
-        assert_eq!(outputs.len(), 1);
-        let (headers, req) = &outputs[0];
-
-        assert!(!headers.contains_key("x-scope-orgid"));
-
-        assert_eq!(req.timeseries.len(), 1);
-        assert_eq!(
-            req.timeseries[0].labels,
-            labels!("__name__" => "gauge-2", "production" => "true", "region" => "us-west-1")
-        );
-        assert_eq!(req.timeseries[0].samples.len(), 1);
-        assert_eq!(req.timeseries[0].samples[0].value, 32.0);
-        assert_eq!(req.metadata.len(), 1);
-        assert_eq!(req.metadata[0].r#type, proto::MetricType::Gauge as i32);
-        assert_eq!(req.metadata[0].metric_family_name, "gauge-2");
-    }
-
-    #[tokio::test]
-    async fn sends_authenticated_request() {
-        let outputs = send_request(
-            indoc! {r#"
-                tenant_id = "tenant-%Y"
-                [auth]
-                strategy = "basic"
-                user = "user"
-                password = "password"
-            "#},
-            vec![create_event("gauge-2".into(), 32.0)],
-        )
-        .await;
-
-        assert_eq!(outputs.len(), 1);
-        let (_headers, req) = &outputs[0];
-
-        assert_eq!(req.timeseries.len(), 1);
-        assert_eq!(
-            req.timeseries[0].labels,
-            labels!("__name__" => "gauge-2", "production" => "true", "region" => "us-west-1")
-        );
-        assert_eq!(req.timeseries[0].samples.len(), 1);
-        assert_eq!(req.timeseries[0].samples[0].value, 32.0);
-        assert_eq!(req.metadata.len(), 1);
-        assert_eq!(req.metadata[0].r#type, proto::MetricType::Gauge as i32);
-        assert_eq!(req.metadata[0].metric_family_name, "gauge-2");
-    }
-
-    #[tokio::test]
-    async fn sends_authenticated_aws_request() {
-        let outputs = send_request(
-            indoc! {r#"
-                tenant_id = "tenant-%Y"
-                [aws]
-                region = "foo"
-                [auth]
-                strategy = "aws"
-                access_key_id = "foo"
-                secret_access_key = "bar"
-            "#},
-            vec![create_event("gauge-2".into(), 32.0)],
-        )
-        .await;
-
-        assert_eq!(outputs.len(), 1);
-        let (headers, _req) = &outputs[0];
-
-        let auth = headers["authorization"]
-            .to_str()
-            .expect("Missing AWS authorization header");
-        assert!(auth.starts_with("AWS4-HMAC-SHA256"));
-    }
-
-    #[tokio::test]
-    async fn sends_x_scope_orgid_header() {
-        let outputs = send_request(
-            r#"tenant_id = "tenant""#,
-            vec![create_event("gauge-3".into(), 12.0)],
-        )
-        .await;
-
-        assert_eq!(outputs.len(), 1);
-        let (headers, _) = &outputs[0];
-        assert_eq!(headers["x-scope-orgid"], "tenant");
-    }
-
-    #[tokio::test]
-    async fn sends_templated_x_scope_orgid_header() {
-        let outputs = send_request(
-            r#"tenant_id = "tenant-%Y""#,
-            vec![create_event("gauge-3".into(), 12.0)],
-        )
-        .await;
-
-        assert_eq!(outputs.len(), 1);
-        let (headers, _) = &outputs[0];
-        let orgid = headers["x-scope-orgid"]
-            .to_str()
-            .expect("Missing x-scope-orgid header");
-        assert!(orgid.starts_with("tenant-20"));
-        assert_eq!(orgid.len(), 11);
-    }
-
-    #[tokio::test]
-    async fn retains_state_between_requests() {
-        // This sink converts all incremental events to absolute, and
-        // should accumulate their totals between batches.
-        let outputs = send_request(
-            r#"batch.max_events = 1"#,
-            vec![
-                create_inc_event("counter-1".into(), 12.0),
-                create_inc_event("counter-2".into(), 13.0),
-                create_inc_event("counter-1".into(), 14.0),
-            ],
-        )
-        .await;
-
-        assert_eq!(outputs.len(), 3);
-
-        let check_output = |index: usize, name: &str, value: f64| {
-            let (_, req) = &outputs[index];
-            assert_eq!(req.timeseries.len(), 1);
-            assert_eq!(req.timeseries[0].labels, labels!("__name__" => name));
-            assert_eq!(req.timeseries[0].samples.len(), 1);
-            assert_eq!(req.timeseries[0].samples[0].value, value);
-        };
-        check_output(0, "counter-1", 12.0);
-        check_output(1, "counter-2", 13.0);
-        check_output(2, "counter-1", 26.0);
-    }
-
-    async fn send_request(
-        config: &str,
-        events: Vec<Event>,
-    ) -> Vec<(HeaderMap, proto::WriteRequest)> {
-        assert_sink_compliance(&HTTP_SINK_TAGS, async {
-            let addr = test_util::next_addr();
-            let (rx, trigger, server) = build_test_server(addr);
-            tokio::spawn(server);
-
-            let config = format!("endpoint = \"http://{}/write\"\n{}", addr, config);
-            let config: RemoteWriteConfig = toml::from_str(&config).unwrap();
-            let cx = SinkContext::default();
-
-            let (sink, _) = config.build(cx).await.unwrap();
-            sink.run_events(events).await.unwrap();
-
-            drop(trigger);
-
-            rx.map(|(parts, body)| {
-                assert_eq!(parts.method, "POST");
-                assert_eq!(parts.uri.path(), "/write");
-                let headers = parts.headers;
-                assert_eq!(headers["x-prometheus-remote-write-version"], "0.1.0");
-                assert_eq!(headers["content-encoding"], "snappy");
-                assert_eq!(headers["content-type"], "application/x-protobuf");
-
-                if config.auth.is_some() {
-                    assert!(headers.contains_key("authorization"));
-                }
-
-                let decoded = snap::raw::Decoder::new()
-                    .decompress_vec(&body)
-                    .expect("Invalid snappy compressed data");
-                let request =
-                    proto::WriteRequest::decode(Bytes::from(decoded)).expect("Invalid protobuf");
-                (headers, request)
-            })
-            .collect::<Vec<_>>()
-            .await
-        })
-        .await
-    }
-
-    pub(super) fn create_event(name: String, value: f64) -> Event {
-        Metric::new(name, MetricKind::Absolute, MetricValue::Gauge { value })
-            .with_tags(Some(metric_tags!(
-                "region" => "us-west-1",
-                "production" => "true",
-            )))
-            .with_timestamp(Some(chrono::Utc::now()))
-            .into()
-    }
-
-    fn create_inc_event(name: String, value: f64) -> Event {
-        Metric::new(
-            name,
-            MetricKind::Incremental,
-            MetricValue::Counter { value },
-        )
-        .with_timestamp(Some(chrono::Utc::now()))
-        .into()
-    }
-}
-
-#[cfg(all(test, feature = "prometheus-integration-tests"))]
-mod integration_tests {
-    use std::{collections::HashMap, ops::Range};
-
-    use serde_json::Value;
-
-    use super::{tests::*, *};
-    use crate::{
-        config::{SinkConfig, SinkContext},
-        event::{metric::MetricValue, Event},
-        sinks::influxdb::test_util::{cleanup_v1, format_timestamp, onboarding_v1, query_v1},
-        test_util::components::{assert_sink_compliance, HTTP_SINK_TAGS},
-        tls::{self, TlsConfig},
-    };
-
-    const HTTP_URL: &str = "http://influxdb-v1:8086";
-    const HTTPS_URL: &str = "https://influxdb-v1-tls:8087";
-
-    #[tokio::test]
-    async fn insert_metrics_over_http() {
-        insert_metrics(HTTP_URL).await;
-    }
-
-    #[tokio::test]
-    async fn insert_metrics_over_https() {
-        insert_metrics(HTTPS_URL).await;
-    }
-
-    async fn insert_metrics(url: &str) {
-        assert_sink_compliance(&HTTP_SINK_TAGS, async {
-            let database = onboarding_v1(url).await;
-
-            let cx = SinkContext::default();
-
-            let config = RemoteWriteConfig {
-                endpoint: format!("{}/api/v1/prom/write?db={}", url, database),
-                tls: Some(TlsConfig {
-                    ca_file: Some(tls::TEST_PEM_CA_PATH.into()),
-                    ..Default::default()
-                }),
-                ..Default::default()
-            };
-            let events = create_events(0..5, |n| n * 11.0);
-
-            let (sink, _) = config.build(cx).await.expect("error building config");
-            sink.run_events(events.clone()).await.unwrap();
-
-            let result = query(url, &format!("show series on {}", database)).await;
-
-            let values = &result["results"][0]["series"][0]["values"];
-            assert_eq!(values.as_array().unwrap().len(), 5);
-
-            for event in events {
-                let metric = event.into_metric();
-                let result = query(
-                    url,
-                    &format!(r#"SELECT * FROM "{}".."{}""#, database, metric.name()),
-                )
-                .await;
-
-                let metrics = decode_metrics(&result["results"][0]["series"][0]);
-                assert_eq!(metrics.len(), 1);
-                let output = &metrics[0];
-
-                match metric.value() {
-                    MetricValue::Gauge { value } => {
-                        assert_eq!(output["value"], Value::Number((*value as u32).into()))
-                    }
-                    _ => panic!("Unhandled metric value, fix the test"),
-                }
-                for (tag, value) in metric.tags().unwrap().iter_single() {
-                    assert_eq!(output[tag], Value::String(value.to_string()));
-                }
-                let timestamp =
-                    format_timestamp(metric.timestamp().unwrap(), chrono::SecondsFormat::Millis);
-                assert_eq!(output["time"], Value::String(timestamp));
-            }
-
-            cleanup_v1(url, &database).await;
-        })
-        .await
-    }
-
-    async fn query(url: &str, query: &str) -> Value {
-        let result = query_v1(url, query).await;
-        let text = result.text().await.unwrap();
-        serde_json::from_str(&text).expect("error when parsing InfluxDB response JSON")
-    }
-
-    fn decode_metrics(data: &Value) -> Vec<HashMap<String, Value>> {
-        let data = data.as_object().expect("Data is not an object");
-        let columns = data["columns"].as_array().expect("Columns is not an array");
-        data["values"]
-            .as_array()
-            .expect("Values is not an array")
-            .iter()
-            .map(|values| {
-                columns
-                    .iter()
-                    .zip(values.as_array().unwrap().iter())
-                    .map(|(column, value)| (column.as_str().unwrap().to_owned(), value.clone()))
-                    .collect()
-            })
-            .collect()
-    }
-
-    fn create_events(name_range: Range<i32>, value: impl Fn(f64) -> f64) -> Vec<Event> {
-        name_range
-            .map(move |num| create_event(format!("metric_{}", num), value(num as f64)))
-            .collect()
-    }
-}
diff --git a/src/sinks/prometheus/remote_write/config.rs b/src/sinks/prometheus/remote_write/config.rs
new file mode 100644
index 0000000000000..0a3fae3a8b48e
--- /dev/null
+++ b/src/sinks/prometheus/remote_write/config.rs
@@ -0,0 +1,206 @@
+use http::Uri;
+use snafu::prelude::*;
+
+use crate::{
+    aws::RegionOrEndpoint,
+    http::{Auth, HttpClient},
+    sinks::{prelude::*, prometheus::PrometheusRemoteWriteAuth},
+};
+
+use super::{
+    request_builder::build_request,
+    sink::{PrometheusRemoteWriteDefaultBatchSettings, RemoteWriteSink},
+    Errors,
+};
+
+/// Configuration for the `prometheus_remote_write` sink.
+#[configurable_component(sink(
+    "prometheus_remote_write",
+    "Deliver metric data to a Prometheus remote write endpoint."
+))]
+#[derive(Clone, Debug, Default)]
+#[serde(deny_unknown_fields)]
+pub struct RemoteWriteConfig {
+    /// The endpoint to send data to.
+    ///
+    /// The endpoint should include the scheme and the path to write to.
+    #[configurable(metadata(docs::examples = "https://localhost:8087/api/v1/write"))]
+    pub endpoint: String,
+
+    /// The default namespace for any metrics sent.
+    ///
+    /// This namespace is only used if a metric has no existing namespace. When a namespace is
+    /// present, it is used as a prefix to the metric name, and separated with an underscore (`_`).
+    ///
+    /// It should follow the Prometheus [naming conventions][prom_naming_docs].
+    ///
+    /// [prom_naming_docs]: https://prometheus.io/docs/practices/naming/#metric-names
+    #[configurable(metadata(docs::examples = "service"))]
+    #[configurable(metadata(docs::advanced))]
+    pub default_namespace: Option<String>,
+
+    /// Default buckets to use for aggregating [distribution][dist_metric_docs] metrics into histograms.
+    ///
+    /// [dist_metric_docs]: https://vector.dev/docs/about/under-the-hood/architecture/data-model/metric/#distribution
+    #[serde(default = "crate::sinks::prometheus::default_histogram_buckets")]
+    #[configurable(metadata(docs::advanced))]
+    pub buckets: Vec<f64>,
+
+    /// Quantiles to use for aggregating [distribution][dist_metric_docs] metrics into a summary.
+    ///
+    /// [dist_metric_docs]: https://vector.dev/docs/about/under-the-hood/architecture/data-model/metric/#distribution
+    #[serde(default = "crate::sinks::prometheus::default_summary_quantiles")]
+    #[configurable(metadata(docs::advanced))]
+    pub quantiles: Vec<f64>,
+
+    #[configurable(derived)]
+    #[serde(default)]
+    pub batch: BatchConfig<PrometheusRemoteWriteDefaultBatchSettings>,
+
+    #[configurable(derived)]
+    #[serde(default)]
+    pub request: TowerRequestConfig,
+
+    /// The tenant ID to send.
+    ///
+    /// If set, a header named `X-Scope-OrgID` is added to outgoing requests with the value of this setting.
+    ///
+    /// This may be used by Cortex or other remote services to identify the tenant making the request.
+    #[serde(default)]
+    #[configurable(metadata(docs::examples = "my-domain"))]
+    #[configurable(metadata(docs::advanced))]
+    pub tenant_id: Option<Template>,
+
+    #[configurable(derived)]
+    pub tls: Option<TlsConfig>,
+
+    #[configurable(derived)]
+    pub auth: Option<PrometheusRemoteWriteAuth>,
+
+    #[configurable(derived)]
+    #[configurable(metadata(docs::advanced))]
+    pub aws: Option<RegionOrEndpoint>,
+
+    #[configurable(derived)]
+    #[serde(
+        default,
+        deserialize_with = "crate::serde::bool_or_struct",
+        skip_serializing_if = "crate::serde::skip_serializing_if_default"
+    )]
+    pub acknowledgements: AcknowledgementsConfig,
+
+    #[configurable(derived)]
+    #[configurable(metadata(docs::advanced))]
+    #[serde(default)]
+    pub compression: super::Compression,
+}
+
+impl_generate_config_from_default!(RemoteWriteConfig);
+
+#[async_trait::async_trait]
+#[typetag::serde(name = "prometheus_remote_write")]
+impl SinkConfig for RemoteWriteConfig {
+    fn acknowledgements(&self) -> &AcknowledgementsConfig {
+        &self.acknowledgements
+    }
+
+    async fn build(&self, cx: SinkContext) -> crate::Result<(VectorSink, Healthcheck)> {
+        let endpoint = self.endpoint.parse::<Uri>().context(UriParseSnafu)?;
+        let tls_settings = TlsSettings::from_options(&self.tls)?;
+        let batch = self.batch.into_batch_settings()?;
+        let request_settings = self.request.unwrap_with(&TowerRequestConfig::default());
+        let buckets = self.buckets.clone();
+        let quantiles = self.quantiles.clone();
+
+        let client = HttpClient::new(tls_settings, cx.proxy())?;
+        let tenant_id = self.tenant_id.clone();
+
+        let (http_auth, credentials_provider, aws_region) = match &self.auth {
+            Some(PrometheusRemoteWriteAuth::Basic { user, password }) => (
+                Some(Auth::Basic {
+                    user: user.clone(),
+                    password: password.clone().into(),
+                }),
+                None,
+                None,
+            ),
+            Some(PrometheusRemoteWriteAuth::Bearer { token }) => (
+                Some(Auth::Bearer {
+                    token: token.clone(),
+                }),
+                None,
+                None,
+            ),
+            Some(PrometheusRemoteWriteAuth::Aws(aws_auth)) => {
+                let region = self
+                    .aws
+                    .as_ref()
+                    .map(|config| config.region())
+                    .ok_or(Errors::AwsRegionRequired)?
+                    .ok_or(Errors::AwsRegionRequired)?;
+
+                (
+                    None,
+                    Some(aws_auth.credentials_provider(region.clone()).await?),
+                    Some(region),
+                )
+            }
+            None => (None, None, None),
+        };
+
+        // let http_request_builder = Arc::new(RemoteWriteRequestBuilder {
+        //     endpoint: endpoint.clone(),
+        //     aws_region,
+        //     credentials_provider,
+        //     http_auth,
+        //     compression: self.compression,
+        // });
+
+        let healthcheck = healthcheck(
+            client.clone(),
+            endpoint,
+            self.compression.clone(),
+            http_auth.clone(),
+        )
+        .boxed();
+
+        let sink = RemoteWriteSink {
+            tenant_id: self.tenant_id.clone(),
+            compression: self.compression.clone(),
+            batch_settings: self.batch.validate()?.into_batcher_settings()?,
+            endpoint: endpoint.clone(),
+            http_auth,
+        };
+
+        Ok((VectorSink::from_event_streamsink(sink), healthcheck))
+    }
+
+    fn input(&self) -> Input {
+        Input::metric()
+    }
+}
+
+async fn healthcheck(
+    client: HttpClient,
+    endpoint: Uri,
+    compression: super::Compression,
+    auth: Option<Auth>,
+) -> crate::Result<()> {
+    let body = bytes::Bytes::new();
+    // let request = http_request_builder.do_the_thing(http::Method::GET, body.into(), None);
+    let request = build_request(
+        http::Method::GET,
+        endpoint,
+        compression,
+        auth,
+        body.into(),
+        None,
+    );
+    // TODO Sign the request
+    let response = client.send(request).await?;
+
+    match response.status() {
+        http::StatusCode::OK => Ok(()),
+        other => Err(HealthcheckError::UnexpectedStatus { status: other }.into()),
+    }
+}
diff --git a/src/sinks/prometheus/remote_write/integration_tests.rs b/src/sinks/prometheus/remote_write/integration_tests.rs
new file mode 100644
index 0000000000000..c3429bef6db60
--- /dev/null
+++ b/src/sinks/prometheus/remote_write/integration_tests.rs
@@ -0,0 +1,111 @@
+use crate::sinks::prelude::*;
+
+use std::{collections::HashMap, ops::Range};
+
+use serde_json::Value;
+
+use super::{tests::*, *};
+use crate::{
+    config::{SinkConfig, SinkContext},
+    event::{metric::MetricValue, Event},
+    sinks::influxdb::test_util::{cleanup_v1, format_timestamp, onboarding_v1, query_v1},
+    test_util::components::{assert_sink_compliance, HTTP_SINK_TAGS},
+    tls::{self, TlsConfig},
+};
+
+const HTTP_URL: &str = "http://influxdb-v1:8086";
+const HTTPS_URL: &str = "https://influxdb-v1-tls:8087";
+
+#[tokio::test]
+async fn insert_metrics_over_http() {
+    insert_metrics(HTTP_URL).await;
+}
+
+#[tokio::test]
+async fn insert_metrics_over_https() {
+    insert_metrics(HTTPS_URL).await;
+}
+
+async fn insert_metrics(url: &str) {
+    assert_sink_compliance(&HTTP_SINK_TAGS, async {
+        let database = onboarding_v1(url).await;
+
+        let cx = SinkContext::default();
+
+        let config = RemoteWriteConfig {
+            endpoint: format!("{}/api/v1/prom/write?db={}", url, database),
+            tls: Some(TlsConfig {
+                ca_file: Some(tls::TEST_PEM_CA_PATH.into()),
+                ..Default::default()
+            }),
+            ..Default::default()
+        };
+        let events = create_events(0..5, |n| n * 11.0);
+
+        let (sink, _) = config.build(cx).await.expect("error building config");
+        sink.run_events(events.clone()).await.unwrap();
+
+        let result = query(url, &format!("show series on {}", database)).await;
+
+        let values = &result["results"][0]["series"][0]["values"];
+        assert_eq!(values.as_array().unwrap().len(), 5);
+
+        for event in events {
+            let metric = event.into_metric();
+            let result = query(
+                url,
+                &format!(r#"SELECT * FROM "{}".."{}""#, database, metric.name()),
+            )
+            .await;
+
+            let metrics = decode_metrics(&result["results"][0]["series"][0]);
+            assert_eq!(metrics.len(), 1);
+            let output = &metrics[0];
+
+            match metric.value() {
+                MetricValue::Gauge { value } => {
+                    assert_eq!(output["value"], Value::Number((*value as u32).into()))
+                }
+                _ => panic!("Unhandled metric value, fix the test"),
+            }
+            for (tag, value) in metric.tags().unwrap().iter_single() {
+                assert_eq!(output[tag], Value::String(value.to_string()));
+            }
+            let timestamp =
+                format_timestamp(metric.timestamp().unwrap(), chrono::SecondsFormat::Millis);
+            assert_eq!(output["time"], Value::String(timestamp));
+        }
+
+        cleanup_v1(url, &database).await;
+    })
+    .await
+}
+
+async fn query(url: &str, query: &str) -> Value {
+    let result = query_v1(url, query).await;
+    let text = result.text().await.unwrap();
+    serde_json::from_str(&text).expect("error when parsing InfluxDB response JSON")
+}
+
+fn decode_metrics(data: &Value) -> Vec<HashMap<String, Value>> {
+    let data = data.as_object().expect("Data is not an object");
+    let columns = data["columns"].as_array().expect("Columns is not an array");
+    data["values"]
+        .as_array()
+        .expect("Values is not an array")
+        .iter()
+        .map(|values| {
+            columns
+                .iter()
+                .zip(values.as_array().unwrap().iter())
+                .map(|(column, value)| (column.as_str().unwrap().to_owned(), value.clone()))
+                .collect()
+        })
+        .collect()
+}
+
+fn create_events(name_range: Range<i32>, value: impl Fn(f64) -> f64) -> Vec<Event> {
+    name_range
+        .map(move |num| create_event(format!("metric_{}", num), value(num as f64)))
+        .collect()
+}
diff --git a/src/sinks/prometheus/remote_write/mod.rs b/src/sinks/prometheus/remote_write/mod.rs
new file mode 100644
index 0000000000000..5dd531ce2ce80
--- /dev/null
+++ b/src/sinks/prometheus/remote_write/mod.rs
@@ -0,0 +1,55 @@
+use snafu::prelude::*;
+
+use crate::sinks::util::buffer::metrics::MetricNormalize;
+
+mod config;
+mod request_builder;
+mod service;
+mod sink;
+
+#[cfg(test)]
+mod tests;
+
+#[cfg(all(test, feature = "prometheus-integration-tests"))]
+mod integration_tests;
+
+pub use config::RemoteWriteConfig;
+
+/// Supported compression types for Prometheus Remote Write.
+#[configurable_component]
+#[derive(Clone, Copy, Debug, Derivative)]
+#[derivative(Default)]
+#[serde(rename_all = "lowercase")]
+pub(super) enum Compression {
+    /// Snappy.
+    #[derivative(Default)]
+    Snappy,
+
+    /// Gzip.
+    Gzip,
+
+    /// Zstandard.
+    Zstd,
+}
+
+#[derive(Debug, Snafu)]
+enum Errors {
+    #[snafu(display(r#"Prometheus remote_write sink cannot accept "set" metrics"#))]
+    SetMetricInvalid,
+    #[snafu(display("aws.region required when AWS authentication is in use"))]
+    AwsRegionRequired,
+}
+
+#[derive(Clone, Eq, Hash, PartialEq)]
+struct PartitionKey {
+    tenant_id: Option<String>,
+}
+
+#[derive(Default)]
+pub struct PrometheusMetricNormalize;
+
+impl MetricNormalize for PrometheusMetricNormalize {
+    fn normalize(&mut self, state: &mut MetricSet, metric: Metric) -> Option<Metric> {
+        state.make_absolute(metric)
+    }
+}
diff --git a/src/sinks/prometheus/remote_write/request_builder.rs b/src/sinks/prometheus/remote_write/request_builder.rs
new file mode 100644
index 0000000000000..088596303d319
--- /dev/null
+++ b/src/sinks/prometheus/remote_write/request_builder.rs
@@ -0,0 +1,198 @@
+use std::io;
+
+use bytes::Bytes;
+use http::{Request, Uri};
+use vector_core::config::telemetry;
+
+use crate::{
+    http::Auth,
+    sinks::{prelude::*, prometheus::PrometheusRemoteWriteAuth, util::PartitionInnerBuffer},
+};
+
+use super::{sink::RemoteWriteMetric, PartitionKey};
+
+pub(crate) struct RemoteWriteEncoder;
+
+impl encoding::Encoder<Vec<RemoteWriteMetric>> for RemoteWriteEncoder {
+    fn encode_input(
+        &self,
+        input: Vec<RemoteWriteMetric>,
+        writer: &mut dyn std::io::Write,
+    ) -> std::io::Result<(usize, GroupedCountByteSize)> {
+        let mut byte_size = telemetry().create_request_count_byte_size();
+
+        let mut time_series = collector::TimeSeries::new();
+        for metric in metrics {
+            byte_size.add_event(&metric.metric, metric.estimated_json_encoded_size_of());
+
+            time_series.encode_metric(
+                self.default_namespace.as_deref(),
+                &self.buckets,
+                &self.quantiles,
+                &metric,
+            );
+        }
+        let request = time_series.finish();
+
+        let mut out = BytesMut::with_capacity(request.encoded_len());
+        request.encode(&mut out).expect("Out of memory");
+        let body = out.freeze();
+
+        write_all(writer, metrics.len(), body.as_ref())?;
+
+        Ok((body.len(), byte_size))
+    }
+}
+
+pub(super) struct RemoteWriteRequest {
+    request: http::Request<Bytes>,
+    finalizers: EventFinalizers,
+    metadata: RequestMetadata,
+    tenant_id: String,
+}
+
+impl Finalizable for RemoteWriteRequest {
+    fn take_finalizers(&mut self) -> EventFinalizers {
+        std::mem::take(&mut self.finalizers)
+    }
+}
+
+impl MetaDescriptive for RemoteWriteRequest {
+    fn get_metadata(&self) -> &RequestMetadata {
+        &self.metadata
+    }
+
+    fn metadata_mut(&mut self) -> &mut RequestMetadata {
+        &mut self.metadata
+    }
+}
+
+struct RemoteWriteMetadata {
+    finalizers: EventFinalizers,
+    tenant_id: Option<String>,
+}
+
+pub struct RemoteWriteRequestBuilder {
+    pub compression: super::Compression,
+    pub encoder: RemoteWriteEncoder,
+    pub http_auth: Option<Auth>,
+    pub endpoint: String,
+}
+
+pub fn build_request(
+    method: http::Method,
+    endpoint: Uri,
+    compression: super::Compression,
+    auth: Option<Auth>,
+    body: Vec<u8>,
+    tenant_id: Option<String>,
+) -> Request<hyper::Body> {
+    let content_encoding = convert_compression_to_content_encoding(compression);
+
+    let mut builder = http::Request::builder()
+        .method(method)
+        .uri(endpoint)
+        .header("X-Prometheus-Remote-Write-Version", "0.1.0")
+        .header("Content-Encoding", content_encoding)
+        .header("Content-Type", "application/x-protobuf");
+
+    if let Some(tenant_id) = &tenant_id {
+        builder = builder.header("X-Scope-OrgID", tenant_id);
+    }
+
+    let mut request = builder.body(body.into()).unwrap();
+    if let Some(http_auth) = &auth {
+        http_auth.apply(&mut request);
+    }
+
+    // if let Some(credentials_provider) = &self.credentials_provider {
+    //     sign_request(&mut request, credentials_provider, &self.aws_region).await?;
+    // }
+
+    // let (parts, body) = request.into_parts();
+    // let request: Request<hyper::Body> = hyper::Request::from_parts(parts, body.into());
+
+    request
+}
+
+impl RequestBuilder<PartitionInnerBuffer<Vec<RemoteWriteMetric>, PartitionKey>>
+    for RemoteWriteRequestBuilder
+{
+    type Metadata = RemoteWriteMetadata;
+    type Events = Vec<RemoteWriteMetric>;
+    type Encoder = RemoteWriteEncoder;
+    type Payload = Bytes;
+    type Request = RemoteWriteRequest;
+    type Error = io::Error;
+
+    fn compression(&self) -> Compression {
+        self.compression
+    }
+
+    fn encoder(&self) -> &Self::Encoder {
+        &self.encoder
+    }
+
+    fn split_input(
+        &self,
+        input: PartitionInnerBuffer<Vec<RemoteWriteMetric>, PartitionKey>,
+    ) -> (Self::Metadata, RequestMetadataBuilder, Self::Events) {
+        let (events, key) = input.into_parts();
+        let builder = RequestMetadataBuilder::from_events(&events);
+        let metadata = RemoteWriteMetadata {
+            finalizers: input.take_finalizers(),
+            tenant_id: key,
+        };
+
+        (metadata, builder, input)
+    }
+
+    fn build_request(
+        &self,
+        metadata: Self::Metadata,
+        request_metadata: RequestMetadata,
+        payload: EncodeResult<Self::Payload>,
+    ) -> Self::Request {
+        let body = payload.into_payload();
+        let request = build_request(
+            http::Method::POST,
+            &self.endpoint,
+            self.compression,
+            self.http_auth,
+            body,
+            metadata.tenant_id,
+        );
+        RemoteWriteRequest {
+            request,
+            finalizers: metadata.finalizers,
+            tenant_id: metadata.tenant_id,
+            metadata,
+        }
+    }
+}
+
+fn compress_block(compression: super::Compression, data: Bytes) -> Vec<u8> {
+    match compression {
+        super::Compression::Snappy => snap::raw::Encoder::new()
+            .compress_vec(&data)
+            .expect("snap compression failed, please report"),
+        super::Compression::Gzip => {
+            let mut buf = Vec::new();
+            flate2::read::GzEncoder::new(data.as_ref(), flate2::Compression::default())
+                .read_to_end(&mut buf)
+                .expect("gzip compression failed, please report");
+            buf
+        }
+        super::Compression::Zstd => {
+            zstd::encode_all(data.as_ref(), 0).expect("zstd compression failed, please report")
+        }
+    }
+}
+
+const fn convert_compression_to_content_encoding(compression: super::Compression) -> &'static str {
+    match compression {
+        super::Compression::Snappy => "snappy",
+        super::Compression::Gzip => "gzip",
+        super::Compression::Zstd => "zstd",
+    }
+}
diff --git a/src/sinks/prometheus/remote_write/service.rs b/src/sinks/prometheus/remote_write/service.rs
new file mode 100644
index 0000000000000..8d15ee5a49de7
--- /dev/null
+++ b/src/sinks/prometheus/remote_write/service.rs
@@ -0,0 +1,120 @@
+use std::sync::Arc;
+
+use http::StatusCode;
+
+use crate::{http::HttpClient, sinks::prelude::*};
+
+#[derive(Clone)]
+struct RemoteWriteService {
+    pub endpoint: Uri,
+    pub aws_region: Option<Region>,
+    pub credentials_provider: Option<SharedCredentialsProvider>,
+
+    default_namespace: Option<String>,
+    client: HttpClient,
+    buckets: Vec<f64>,
+    quantiles: Vec<f64>,
+    compression: Compression,
+}
+
+impl RemoteWriteService {
+    // fn encode_events(&self, metrics: Vec<Metric>) -> Bytes {
+    //     let mut time_series = collector::TimeSeries::new();
+    //     for metric in metrics {
+    //         time_series.encode_metric(
+    //             self.default_namespace.as_deref(),
+    //             &self.buckets,
+    //             &self.quantiles,
+    //             &metric,
+    //         );
+    //     }
+    //     let request = time_series.finish();
+
+    //     let mut out = BytesMut::with_capacity(request.encoded_len());
+    //     request.encode(&mut out).expect("Out of memory");
+    //     out.freeze()
+    // }
+}
+
+struct RemoteWriteResponse {
+    byte_size: usize,
+    json_size: GroupedCountByteSize,
+    response: http::Response<Bytes>,
+}
+
+impl DriverResponse for AmqpResponse {
+    fn event_status(&self) -> EventStatus {
+        if response.status().is_success() {
+            EventStatus::Delivered
+        } else {
+            EventStatus::Errored
+        }
+    }
+
+    fn events_sent(&self) -> &GroupedCountByteSize {
+        &self.json_size
+    }
+
+    fn bytes_sent(&self) -> Option<usize> {
+        Some(self.byte_size)
+    }
+}
+
+impl Service<RemoteWriteRequest> for RemoteWriteService {
+    type Response = http::Response<Bytes>;
+    type Error = crate::Error;
+    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
+
+    fn poll_ready(&mut self, _task: &mut task::Context<'_>) -> task::Poll<Result<(), Self::Error>> {
+        task::Poll::Ready(Ok(()))
+    }
+
+    // Emission of internal events for errors and dropped events is handled upstream by the caller.
+    fn call(&mut self, request: RemoteWriteRequest) -> Self::Future {
+        // let (events, key) = buffer.into_parts();
+        // let body = self.encode_events(events);
+        // let body = compress_block(self.compression, body);
+
+        let client = self.client.clone();
+        // let request_builder = Arc::clone(&self.http_request_builder);
+
+        Box::pin(async move {
+            // let request = request_builder
+            //     .build_request(http::Method::POST, body, key.tenant_id)
+            //     .await?;
+
+            let metadata = std::mem::take(request.metadata_mut());
+            let mut request = request.request;
+
+            if let Some(credentials_provider) = &self.credentials_provider {
+                sign_request(&mut request, credentials_provider, &self.aws_region).await?;
+            }
+
+            let (parts, body) = request.into_parts();
+            let request: Request<hyper::Body> = hyper::Request::from_parts(parts, body.into());
+
+            let (protocol, endpoint) = uri::protocol_endpoint(request.uri().clone());
+
+            let response = client.send(request).await?;
+            let (parts, body) = response.into_parts();
+            let body = hyper::body::to_bytes(body).await?;
+            let byte_size = body.len();
+
+            let response = hyper::Response::from_parts(parts, body);
+
+            Ok(RemoteWriteResponse {
+                byte_size,
+                json_size: todo!(),
+                response,
+            })
+        })
+    }
+}
+
+async fn sign_request(
+    request: &mut http::Request<Bytes>,
+    credentials_provider: &SharedCredentialsProvider,
+    region: &Option<Region>,
+) -> crate::Result<()> {
+    crate::aws::sign_request("aps", request, credentials_provider, region).await
+}
diff --git a/src/sinks/prometheus/remote_write/sink.rs b/src/sinks/prometheus/remote_write/sink.rs
new file mode 100644
index 0000000000000..fcaf9f2345f3a
--- /dev/null
+++ b/src/sinks/prometheus/remote_write/sink.rs
@@ -0,0 +1,147 @@
+use vector_core::event::Metric;
+
+use crate::sinks::{prelude::*, prometheus::PrometheusRemoteWriteAuth};
+
+use super::{
+    request_builder::{RemoteWriteEncoder, RemoteWriteRequestBuilder},
+    PrometheusMetricNormalize,
+};
+
+pub(super) struct RemoteWriteMetric {
+    metric: Metric,
+    tenant_id: String,
+}
+
+impl GetEventCountTags for RemoteWriteMetric {
+    fn get_tags(&self) -> TaggedEventsSent {
+        self.metric.get_tags()
+    }
+}
+
+impl EstimatedJsonEncodedSizeOf for RemoteWriteMetric {
+    fn estimated_json_encoded_size_of(&self) -> JsonSize {
+        self.metric.estimated_json_encoded_size_of()
+    }
+}
+
+#[derive(Clone, Copy, Debug, Default)]
+pub struct PrometheusRemoteWriteDefaultBatchSettings;
+
+pub(super) struct PrometheusTenantIdPartitioner;
+
+impl Partitioner for PrometheusTenantIdPartitioner {
+    type Item = PrometheusMetric;
+    type Key = String;
+
+    fn partition(&self, item: &Self::Item) -> Self::Key {
+        // TODO do this better.
+        item.tenant_id.clone()
+    }
+}
+
+impl SinkBatchSettings for PrometheusRemoteWriteDefaultBatchSettings {
+    const MAX_EVENTS: Option<usize> = Some(1_000);
+    const MAX_BYTES: Option<usize> = None;
+    const TIMEOUT_SECS: f64 = 1.0;
+}
+
+pub(super) struct RemoteWriteSink {
+    tenant_id: Option<Template>,
+    batch_settings: BatcherSettings,
+    compression: super::Compression,
+    http_auth: Option<PrometheusRemoteWriteAuth>,
+    endpoint: String,
+}
+
+impl RemoteWriteSink {
+    fn make_remote_write_event(&self, metric: Metric) -> Option<RemoteWriteMetric> {
+        let tenant_id = self.tenant_id.as_ref().and_then(|template| {
+            template
+                .render_string(&metric)
+                .map_err(|error| {
+                    emit!(TemplateRenderingError {
+                        error,
+                        field: Some("tenant_id"),
+                        drop_event: true,
+                    })
+                })
+                .ok()
+        })?;
+
+        Some(RemoteWriteMetric { metric, tenant_id })
+    }
+
+    async fn run_inner(self: Box<Self>, input: BoxStream<'_, Event>) -> Result<(), ()> {
+        let request_builder = RemoteWriteRequestBuilder {
+            endpoint: self.endpoint.clone(),
+            compression: self.compression,
+            encoder: RemoteWriteEncoder,
+            http_auth: self.http_auth,
+        };
+
+        let service = RemoteWriteService {
+            default_namespace: self.default_namespace.clone(),
+            client,
+            buckets,
+            quantiles,
+            http_request_builder,
+            compression: self.compression,
+        };
+        let service = ServiceBuilder::new().service(service);
+
+        input
+            .filter_map(|event| future::ready(event.try_into_metric()))
+            .normalized_with_default::<PrometheusMetricNormalize>()
+            .filter_map(|event| self.make_remote_write_event(event))
+            .batched_partitioned(PrometheusTenantIdPartitioner, self.batch_settings)
+            .request_builder(request_builder)
+            .into_driver(service)
+            .protocol("http")
+            .run()
+            .await
+
+        /*
+            let buffer = PartitionBuffer::new(MetricsBuffer::new(batch.size));
+            let mut normalizer = MetricNormalizer::<PrometheusMetricNormalize>::default();
+
+            request_settings
+                .partition_sink(HttpRetryLogic, service, buffer, batch.timeout)
+                .with_flat_map(move |event: Event| {
+                    let byte_size = event.size_of();
+                    let json_size = event.estimated_json_encoded_size_of();
+
+                    stream::iter(normalizer.normalize(event.into_metric()).map(|event| {
+                        let tenant_id = tenant_id.as_ref().and_then(|template| {
+                            template
+                                .render_string(&event)
+                                .map_err(|error| {
+                                    emit!(TemplateRenderingError {
+                                        error,
+                                        field: Some("tenant_id"),
+                                        drop_event: true,
+                                    })
+                                })
+                                .ok()
+                        });
+                        let key = PartitionKey { tenant_id };
+                        Ok(EncodedEvent::new(
+                            PartitionInnerBuffer::new(event, key),
+                            byte_size,
+                            json_size,
+                        ))
+                    }))
+                })
+                .sink_map_err(
+                    |error| error!(message = "Prometheus remote_write sink error.", %error),
+                )
+        ;
+        */
+    }
+}
+
+#[async_trait]
+impl StreamSink<Event> for RemoteWriteSink {
+    async fn run(self: Box<Self>, input: BoxStream<'_, Event>) -> Result<(), ()> {
+        self.run_inner(input).await
+    }
+}
diff --git a/src/sinks/prometheus/remote_write/tests.rs b/src/sinks/prometheus/remote_write/tests.rs
new file mode 100644
index 0000000000000..d7950686b486e
--- /dev/null
+++ b/src/sinks/prometheus/remote_write/tests.rs
@@ -0,0 +1,230 @@
+use crate::sinks::prelude::*;
+
+#[cfg(test)]
+mod tests {
+    use futures::StreamExt;
+    use http::HeaderMap;
+    use indoc::indoc;
+    use prometheus_parser::proto;
+    use vector_core::metric_tags;
+
+    use super::*;
+    use crate::{
+        config::SinkContext,
+        event::{MetricKind, MetricValue},
+        sinks::util::test::build_test_server,
+        test_util::{
+            self,
+            components::{assert_sink_compliance, HTTP_SINK_TAGS},
+        },
+    };
+
+    #[test]
+    fn generate_config() {
+        crate::test_util::test_generate_config::<RemoteWriteConfig>();
+    }
+
+    macro_rules! labels {
+        ( $( $name:expr => $value:expr ),* ) => {
+            vec![ $( proto::Label {
+                name: $name.to_string(),
+                value: $value.to_string()
+            }, )* ]
+        }
+    }
+
+    #[tokio::test]
+    async fn sends_request() {
+        let outputs = send_request("", vec![create_event("gauge-2".into(), 32.0)]).await;
+        assert_eq!(outputs.len(), 1);
+        let (headers, req) = &outputs[0];
+
+        assert!(!headers.contains_key("x-scope-orgid"));
+
+        assert_eq!(req.timeseries.len(), 1);
+        assert_eq!(
+            req.timeseries[0].labels,
+            labels!("__name__" => "gauge-2", "production" => "true", "region" => "us-west-1")
+        );
+        assert_eq!(req.timeseries[0].samples.len(), 1);
+        assert_eq!(req.timeseries[0].samples[0].value, 32.0);
+        assert_eq!(req.metadata.len(), 1);
+        assert_eq!(req.metadata[0].r#type, proto::MetricType::Gauge as i32);
+        assert_eq!(req.metadata[0].metric_family_name, "gauge-2");
+    }
+
+    #[tokio::test]
+    async fn sends_authenticated_request() {
+        let outputs = send_request(
+            indoc! {r#"
+                tenant_id = "tenant-%Y"
+                [auth]
+                strategy = "basic"
+                user = "user"
+                password = "password"
+            "#},
+            vec![create_event("gauge-2".into(), 32.0)],
+        )
+        .await;
+
+        assert_eq!(outputs.len(), 1);
+        let (_headers, req) = &outputs[0];
+
+        assert_eq!(req.timeseries.len(), 1);
+        assert_eq!(
+            req.timeseries[0].labels,
+            labels!("__name__" => "gauge-2", "production" => "true", "region" => "us-west-1")
+        );
+        assert_eq!(req.timeseries[0].samples.len(), 1);
+        assert_eq!(req.timeseries[0].samples[0].value, 32.0);
+        assert_eq!(req.metadata.len(), 1);
+        assert_eq!(req.metadata[0].r#type, proto::MetricType::Gauge as i32);
+        assert_eq!(req.metadata[0].metric_family_name, "gauge-2");
+    }
+
+    #[tokio::test]
+    async fn sends_authenticated_aws_request() {
+        let outputs = send_request(
+            indoc! {r#"
+                tenant_id = "tenant-%Y"
+                [aws]
+                region = "foo"
+                [auth]
+                strategy = "aws"
+                access_key_id = "foo"
+                secret_access_key = "bar"
+            "#},
+            vec![create_event("gauge-2".into(), 32.0)],
+        )
+        .await;
+
+        assert_eq!(outputs.len(), 1);
+        let (headers, _req) = &outputs[0];
+
+        let auth = headers["authorization"]
+            .to_str()
+            .expect("Missing AWS authorization header");
+        assert!(auth.starts_with("AWS4-HMAC-SHA256"));
+    }
+
+    #[tokio::test]
+    async fn sends_x_scope_orgid_header() {
+        let outputs = send_request(
+            r#"tenant_id = "tenant""#,
+            vec![create_event("gauge-3".into(), 12.0)],
+        )
+        .await;
+
+        assert_eq!(outputs.len(), 1);
+        let (headers, _) = &outputs[0];
+        assert_eq!(headers["x-scope-orgid"], "tenant");
+    }
+
+    #[tokio::test]
+    async fn sends_templated_x_scope_orgid_header() {
+        let outputs = send_request(
+            r#"tenant_id = "tenant-%Y""#,
+            vec![create_event("gauge-3".into(), 12.0)],
+        )
+        .await;
+
+        assert_eq!(outputs.len(), 1);
+        let (headers, _) = &outputs[0];
+        let orgid = headers["x-scope-orgid"]
+            .to_str()
+            .expect("Missing x-scope-orgid header");
+        assert!(orgid.starts_with("tenant-20"));
+        assert_eq!(orgid.len(), 11);
+    }
+
+    #[tokio::test]
+    async fn retains_state_between_requests() {
+        // This sink converts all incremental events to absolute, and
+        // should accumulate their totals between batches.
+        let outputs = send_request(
+            r#"batch.max_events = 1"#,
+            vec![
+                create_inc_event("counter-1".into(), 12.0),
+                create_inc_event("counter-2".into(), 13.0),
+                create_inc_event("counter-1".into(), 14.0),
+            ],
+        )
+        .await;
+
+        assert_eq!(outputs.len(), 3);
+
+        let check_output = |index: usize, name: &str, value: f64| {
+            let (_, req) = &outputs[index];
+            assert_eq!(req.timeseries.len(), 1);
+            assert_eq!(req.timeseries[0].labels, labels!("__name__" => name));
+            assert_eq!(req.timeseries[0].samples.len(), 1);
+            assert_eq!(req.timeseries[0].samples[0].value, value);
+        };
+        check_output(0, "counter-1", 12.0);
+        check_output(1, "counter-2", 13.0);
+        check_output(2, "counter-1", 26.0);
+    }
+
+    async fn send_request(
+        config: &str,
+        events: Vec<Event>,
+    ) -> Vec<(HeaderMap, proto::WriteRequest)> {
+        assert_sink_compliance(&HTTP_SINK_TAGS, async {
+            let addr = test_util::next_addr();
+            let (rx, trigger, server) = build_test_server(addr);
+            tokio::spawn(server);
+
+            let config = format!("endpoint = \"http://{}/write\"\n{}", addr, config);
+            let config: RemoteWriteConfig = toml::from_str(&config).unwrap();
+            let cx = SinkContext::default();
+
+            let (sink, _) = config.build(cx).await.unwrap();
+            sink.run_events(events).await.unwrap();
+
+            drop(trigger);
+
+            rx.map(|(parts, body)| {
+                assert_eq!(parts.method, "POST");
+                assert_eq!(parts.uri.path(), "/write");
+                let headers = parts.headers;
+                assert_eq!(headers["x-prometheus-remote-write-version"], "0.1.0");
+                assert_eq!(headers["content-encoding"], "snappy");
+                assert_eq!(headers["content-type"], "application/x-protobuf");
+
+                if config.auth.is_some() {
+                    assert!(headers.contains_key("authorization"));
+                }
+
+                let decoded = snap::raw::Decoder::new()
+                    .decompress_vec(&body)
+                    .expect("Invalid snappy compressed data");
+                let request =
+                    proto::WriteRequest::decode(Bytes::from(decoded)).expect("Invalid protobuf");
+                (headers, request)
+            })
+            .collect::<Vec<_>>()
+            .await
+        })
+        .await
+    }
+
+    pub(super) fn create_event(name: String, value: f64) -> Event {
+        Metric::new(name, MetricKind::Absolute, MetricValue::Gauge { value })
+            .with_tags(Some(metric_tags!(
+                "region" => "us-west-1",
+                "production" => "true",
+            )))
+            .with_timestamp(Some(chrono::Utc::now()))
+            .into()
+    }
+
+    fn create_inc_event(name: String, value: f64) -> Event {
+        Metric::new(
+            name,
+            MetricKind::Incremental,
+            MetricValue::Counter { value },
+        )
+        .with_timestamp(Some(chrono::Utc::now()))
+        .into()
+    }
+}

From 06a7ae3823514fc0b2c70bc76611eb4c18f8e677 Mon Sep 17 00:00:00 2001
From: Stephen Wakely <fungus.humungus@gmail.com>
Date: Mon, 18 Sep 2023 15:45:47 +0100
Subject: [PATCH 02/32] WIP

Signed-off-by: Stephen Wakely <fungus.humungus@gmail.com>
---
 src/sinks/prometheus/remote_write/config.rs   | 49 +++++++++----------
 src/sinks/prometheus/remote_write/mod.rs      |  3 +-
 .../remote_write/request_builder.rs           | 14 ++++--
 src/sinks/prometheus/remote_write/service.rs  | 18 ++++---
 src/sinks/prometheus/remote_write/sink.rs     |  4 +-
 5 files changed, 49 insertions(+), 39 deletions(-)

diff --git a/src/sinks/prometheus/remote_write/config.rs b/src/sinks/prometheus/remote_write/config.rs
index 0a3fae3a8b48e..cb3e27e8049e7 100644
--- a/src/sinks/prometheus/remote_write/config.rs
+++ b/src/sinks/prometheus/remote_write/config.rs
@@ -1,10 +1,14 @@
+#[cfg(feature = "aws-core")]
+use aws_credential_types::provider::SharedCredentialsProvider;
+#[cfg(feature = "aws-core")]
+use aws_types::region::Region;
 use http::Uri;
 use snafu::prelude::*;
 
 use crate::{
     aws::RegionOrEndpoint,
     http::{Auth, HttpClient},
-    sinks::{prelude::*, prometheus::PrometheusRemoteWriteAuth},
+    sinks::{prelude::*, prometheus::PrometheusRemoteWriteAuth, UriParseSnafu},
 };
 
 use super::{
@@ -115,22 +119,17 @@ impl SinkConfig for RemoteWriteConfig {
         let client = HttpClient::new(tls_settings, cx.proxy())?;
         let tenant_id = self.tenant_id.clone();
 
-        let (http_auth, credentials_provider, aws_region) = match &self.auth {
-            Some(PrometheusRemoteWriteAuth::Basic { user, password }) => (
-                Some(Auth::Basic {
-                    user: user.clone(),
-                    password: password.clone().into(),
-                }),
-                None,
-                None,
-            ),
-            Some(PrometheusRemoteWriteAuth::Bearer { token }) => (
-                Some(Auth::Bearer {
-                    token: token.clone(),
-                }),
-                None,
-                None,
-            ),
+        let auth = match &self.auth {
+            Some(PrometheusRemoteWriteAuth::Basic { user, password }) => Some(Auth::Basic {
+                user: user.clone(),
+                password: password.clone().into(),
+            }),
+
+            Some(PrometheusRemoteWriteAuth::Bearer { token }) => Some(Auth::Bearer {
+                token: token.clone(),
+            }),
+
+            #[cfg(feature = "aws_core")]
             Some(PrometheusRemoteWriteAuth::Aws(aws_auth)) => {
                 let region = self
                     .aws
@@ -139,13 +138,13 @@ impl SinkConfig for RemoteWriteConfig {
                     .ok_or(Errors::AwsRegionRequired)?
                     .ok_or(Errors::AwsRegionRequired)?;
 
-                (
-                    None,
-                    Some(aws_auth.credentials_provider(region.clone()).await?),
-                    Some(region),
-                )
+                Some(Auth::Aws {
+                    credentials_provider: aws_auth.credentials_provider(region.clone()).await?,
+                    region,
+                })
             }
-            None => (None, None, None),
+
+            None => None,
         };
 
         // let http_request_builder = Arc::new(RemoteWriteRequestBuilder {
@@ -160,7 +159,7 @@ impl SinkConfig for RemoteWriteConfig {
             client.clone(),
             endpoint,
             self.compression.clone(),
-            http_auth.clone(),
+            auth.clone(),
         )
         .boxed();
 
@@ -169,7 +168,7 @@ impl SinkConfig for RemoteWriteConfig {
             compression: self.compression.clone(),
             batch_settings: self.batch.validate()?.into_batcher_settings()?,
             endpoint: endpoint.clone(),
-            http_auth,
+            auth,
         };
 
         Ok((VectorSink::from_event_streamsink(sink), healthcheck))
diff --git a/src/sinks/prometheus/remote_write/mod.rs b/src/sinks/prometheus/remote_write/mod.rs
index 5dd531ce2ce80..cdc0148811894 100644
--- a/src/sinks/prometheus/remote_write/mod.rs
+++ b/src/sinks/prometheus/remote_write/mod.rs
@@ -1,6 +1,6 @@
 use snafu::prelude::*;
 
-use crate::sinks::util::buffer::metrics::MetricNormalize;
+use crate::sinks::{prelude::*, util::buffer::metrics::MetricNormalize};
 
 mod config;
 mod request_builder;
@@ -36,6 +36,7 @@ pub(super) enum Compression {
 enum Errors {
     #[snafu(display(r#"Prometheus remote_write sink cannot accept "set" metrics"#))]
     SetMetricInvalid,
+    #[cfg(feature = "aws-core")]
     #[snafu(display("aws.region required when AWS authentication is in use"))]
     AwsRegionRequired,
 }
diff --git a/src/sinks/prometheus/remote_write/request_builder.rs b/src/sinks/prometheus/remote_write/request_builder.rs
index 088596303d319..091ead9896f78 100644
--- a/src/sinks/prometheus/remote_write/request_builder.rs
+++ b/src/sinks/prometheus/remote_write/request_builder.rs
@@ -1,12 +1,16 @@
 use std::io;
 
-use bytes::Bytes;
+use bytes::{Bytes, BytesMut};
 use http::{Request, Uri};
 use vector_core::config::telemetry;
 
 use crate::{
     http::Auth,
-    sinks::{prelude::*, prometheus::PrometheusRemoteWriteAuth, util::PartitionInnerBuffer},
+    sinks::{
+        prelude::*,
+        prometheus::{collector, PrometheusRemoteWriteAuth},
+        util::PartitionInnerBuffer,
+    },
 };
 
 use super::{sink::RemoteWriteMetric, PartitionKey};
@@ -22,7 +26,7 @@ impl encoding::Encoder<Vec<RemoteWriteMetric>> for RemoteWriteEncoder {
         let mut byte_size = telemetry().create_request_count_byte_size();
 
         let mut time_series = collector::TimeSeries::new();
-        for metric in metrics {
+        for metric in input {
             byte_size.add_event(&metric.metric, metric.estimated_json_encoded_size_of());
 
             time_series.encode_metric(
@@ -38,14 +42,14 @@ impl encoding::Encoder<Vec<RemoteWriteMetric>> for RemoteWriteEncoder {
         request.encode(&mut out).expect("Out of memory");
         let body = out.freeze();
 
-        write_all(writer, metrics.len(), body.as_ref())?;
+        write_all(writer, input.len(), body.as_ref())?;
 
         Ok((body.len(), byte_size))
     }
 }
 
 pub(super) struct RemoteWriteRequest {
-    request: http::Request<Bytes>,
+    pub(super) request: http::Request<Bytes>,
     finalizers: EventFinalizers,
     metadata: RequestMetadata,
     tenant_id: String,
diff --git a/src/sinks/prometheus/remote_write/service.rs b/src/sinks/prometheus/remote_write/service.rs
index 8d15ee5a49de7..28a6667ea97a6 100644
--- a/src/sinks/prometheus/remote_write/service.rs
+++ b/src/sinks/prometheus/remote_write/service.rs
@@ -1,11 +1,17 @@
-use std::sync::Arc;
+use std::{
+    sync::Arc,
+    task::{Context, Poll},
+};
 
+use aws_sdk_cloudwatch::Region;
+use bytes::Bytes;
 use http::StatusCode;
 
+use super::request_builder::RemoteWriteRequest;
 use crate::{http::HttpClient, sinks::prelude::*};
 
 #[derive(Clone)]
-struct RemoteWriteService {
+pub(super) struct RemoteWriteService {
     pub endpoint: Uri,
     pub aws_region: Option<Region>,
     pub credentials_provider: Option<SharedCredentialsProvider>,
@@ -42,9 +48,9 @@ struct RemoteWriteResponse {
     response: http::Response<Bytes>,
 }
 
-impl DriverResponse for AmqpResponse {
+impl DriverResponse for RemoteWriteResponse {
     fn event_status(&self) -> EventStatus {
-        if response.status().is_success() {
+        if self.response.status().is_success() {
             EventStatus::Delivered
         } else {
             EventStatus::Errored
@@ -65,8 +71,8 @@ impl Service<RemoteWriteRequest> for RemoteWriteService {
     type Error = crate::Error;
     type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
 
-    fn poll_ready(&mut self, _task: &mut task::Context<'_>) -> task::Poll<Result<(), Self::Error>> {
-        task::Poll::Ready(Ok(()))
+    fn poll_ready(&mut self, _task: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        Poll::Ready(Ok(()))
     }
 
     // Emission of internal events for errors and dropped events is handled upstream by the caller.
diff --git a/src/sinks/prometheus/remote_write/sink.rs b/src/sinks/prometheus/remote_write/sink.rs
index fcaf9f2345f3a..6f3e06d7a01ca 100644
--- a/src/sinks/prometheus/remote_write/sink.rs
+++ b/src/sinks/prometheus/remote_write/sink.rs
@@ -8,7 +8,7 @@ use super::{
 };
 
 pub(super) struct RemoteWriteMetric {
-    metric: Metric,
+    pub(super) metric: Metric,
     tenant_id: String,
 }
 
@@ -30,7 +30,7 @@ pub struct PrometheusRemoteWriteDefaultBatchSettings;
 pub(super) struct PrometheusTenantIdPartitioner;
 
 impl Partitioner for PrometheusTenantIdPartitioner {
-    type Item = PrometheusMetric;
+    type Item = RemoteWriteMetric;
     type Key = String;
 
     fn partition(&self, item: &Self::Item) -> Self::Key {

From a3ef0f31d0a0a196330c34f1dc46107861b0c4b7 Mon Sep 17 00:00:00 2001
From: Stephen Wakely <fungus.humungus@gmail.com>
Date: Mon, 25 Sep 2023 14:42:23 +0100
Subject: [PATCH 03/32] Made tests pass

Signed-off-by: Stephen Wakely <fungus.humungus@gmail.com>
---
 src/sinks/prometheus/remote_write/config.rs   |  93 +++--
 .../remote_write/integration_tests.rs         |   2 -
 src/sinks/prometheus/remote_write/mod.rs      |   8 +-
 .../remote_write/request_builder.rs           | 120 ++----
 src/sinks/prometheus/remote_write/service.rs  | 189 ++++++---
 src/sinks/prometheus/remote_write/sink.rs     | 172 ++++----
 src/sinks/prometheus/remote_write/tests.rs    | 381 +++++++++---------
 src/sinks/util/buffer/metrics/normalize.rs    |   2 +-
 src/sources/prometheus/remote_write.rs        | 214 +++++-----
 9 files changed, 589 insertions(+), 592 deletions(-)

diff --git a/src/sinks/prometheus/remote_write/config.rs b/src/sinks/prometheus/remote_write/config.rs
index cb3e27e8049e7..edf40ba267a61 100644
--- a/src/sinks/prometheus/remote_write/config.rs
+++ b/src/sinks/prometheus/remote_write/config.rs
@@ -1,22 +1,23 @@
-#[cfg(feature = "aws-core")]
-use aws_credential_types::provider::SharedCredentialsProvider;
-#[cfg(feature = "aws-core")]
-use aws_types::region::Region;
+// #[cfg(feature = "aws-core")]
+// use aws_credential_types::provider::SharedCredentialsProvider;
+// #[cfg(feature = "aws-core")]
+// use aws_types::region::Region;
 use http::Uri;
 use snafu::prelude::*;
 
 use crate::{
-    aws::RegionOrEndpoint,
-    http::{Auth, HttpClient},
-    sinks::{prelude::*, prometheus::PrometheusRemoteWriteAuth, UriParseSnafu},
+    http::HttpClient,
+    sinks::{prelude::*, prometheus::PrometheusRemoteWriteAuth, util::auth::Auth, UriParseSnafu},
 };
 
 use super::{
-    request_builder::build_request,
+    service::{build_request, RemoteWriteRetryLogic, RemoteWriteService},
     sink::{PrometheusRemoteWriteDefaultBatchSettings, RemoteWriteSink},
-    Errors,
 };
 
+#[cfg(feature = "aws-core")]
+use super::Errors;
+
 /// Configuration for the `prometheus_remote_write` sink.
 #[configurable_component(sink(
     "prometheus_remote_write",
@@ -81,9 +82,10 @@ pub struct RemoteWriteConfig {
     #[configurable(derived)]
     pub auth: Option<PrometheusRemoteWriteAuth>,
 
+    #[cfg(feature = "aws-config")]
     #[configurable(derived)]
     #[configurable(metadata(docs::advanced))]
-    pub aws: Option<RegionOrEndpoint>,
+    pub aws: Option<crate::aws::RegionOrEndpoint>,
 
     #[configurable(derived)]
     #[serde(
@@ -111,25 +113,26 @@ impl SinkConfig for RemoteWriteConfig {
     async fn build(&self, cx: SinkContext) -> crate::Result<(VectorSink, Healthcheck)> {
         let endpoint = self.endpoint.parse::<Uri>().context(UriParseSnafu)?;
         let tls_settings = TlsSettings::from_options(&self.tls)?;
-        let batch = self.batch.into_batch_settings()?;
         let request_settings = self.request.unwrap_with(&TowerRequestConfig::default());
         let buckets = self.buckets.clone();
         let quantiles = self.quantiles.clone();
+        let default_namespace = self.default_namespace.clone();
 
         let client = HttpClient::new(tls_settings, cx.proxy())?;
-        let tenant_id = self.tenant_id.clone();
 
         let auth = match &self.auth {
-            Some(PrometheusRemoteWriteAuth::Basic { user, password }) => Some(Auth::Basic {
-                user: user.clone(),
-                password: password.clone().into(),
-            }),
-
-            Some(PrometheusRemoteWriteAuth::Bearer { token }) => Some(Auth::Bearer {
-                token: token.clone(),
-            }),
-
-            #[cfg(feature = "aws_core")]
+            Some(PrometheusRemoteWriteAuth::Basic { user, password }) => {
+                Some(Auth::Basic(crate::http::Auth::Basic {
+                    user: user.clone(),
+                    password: password.clone().into(),
+                }))
+            }
+            Some(PrometheusRemoteWriteAuth::Bearer { token }) => {
+                Some(Auth::Basic(crate::http::Auth::Bearer {
+                    token: token.clone(),
+                }))
+            }
+            #[cfg(feature = "aws-core")]
             Some(PrometheusRemoteWriteAuth::Aws(aws_auth)) => {
                 let region = self
                     .aws
@@ -137,38 +140,40 @@ impl SinkConfig for RemoteWriteConfig {
                     .map(|config| config.region())
                     .ok_or(Errors::AwsRegionRequired)?
                     .ok_or(Errors::AwsRegionRequired)?;
-
                 Some(Auth::Aws {
                     credentials_provider: aws_auth.credentials_provider(region.clone()).await?,
                     region,
                 })
             }
-
             None => None,
         };
 
-        // let http_request_builder = Arc::new(RemoteWriteRequestBuilder {
-        //     endpoint: endpoint.clone(),
-        //     aws_region,
-        //     credentials_provider,
-        //     http_auth,
-        //     compression: self.compression,
-        // });
-
         let healthcheck = healthcheck(
             client.clone(),
-            endpoint,
-            self.compression.clone(),
+            endpoint.clone(),
+            self.compression,
             auth.clone(),
         )
         .boxed();
 
+        let service = RemoteWriteService {
+            endpoint,
+            client,
+            auth: auth.clone(),
+            compression: self.compression,
+        };
+        let service = ServiceBuilder::new()
+            .settings(request_settings, RemoteWriteRetryLogic)
+            .service(service);
+
         let sink = RemoteWriteSink {
             tenant_id: self.tenant_id.clone(),
-            compression: self.compression.clone(),
+            compression: self.compression,
             batch_settings: self.batch.validate()?.into_batcher_settings()?,
-            endpoint: endpoint.clone(),
-            auth,
+            buckets,
+            quantiles,
+            default_namespace,
+            service,
         };
 
         Ok((VectorSink::from_event_streamsink(sink), healthcheck))
@@ -186,17 +191,9 @@ async fn healthcheck(
     auth: Option<Auth>,
 ) -> crate::Result<()> {
     let body = bytes::Bytes::new();
-    // let request = http_request_builder.do_the_thing(http::Method::GET, body.into(), None);
-    let request = build_request(
-        http::Method::GET,
-        endpoint,
-        compression,
-        auth,
-        body.into(),
-        None,
-    );
-    // TODO Sign the request
-    let response = client.send(request).await?;
+    let request =
+        build_request(http::Method::GET, &endpoint, compression, body, None, auth).await?;
+    let response = client.send(request.map(hyper::Body::from)).await?;
 
     match response.status() {
         http::StatusCode::OK => Ok(()),
diff --git a/src/sinks/prometheus/remote_write/integration_tests.rs b/src/sinks/prometheus/remote_write/integration_tests.rs
index c3429bef6db60..07f106053fdad 100644
--- a/src/sinks/prometheus/remote_write/integration_tests.rs
+++ b/src/sinks/prometheus/remote_write/integration_tests.rs
@@ -1,5 +1,3 @@
-use crate::sinks::prelude::*;
-
 use std::{collections::HashMap, ops::Range};
 
 use serde_json::Value;
diff --git a/src/sinks/prometheus/remote_write/mod.rs b/src/sinks/prometheus/remote_write/mod.rs
index cdc0148811894..137e7e309a9cb 100644
--- a/src/sinks/prometheus/remote_write/mod.rs
+++ b/src/sinks/prometheus/remote_write/mod.rs
@@ -1,6 +1,10 @@
 use snafu::prelude::*;
+use vector_core::event::Metric;
 
-use crate::sinks::{prelude::*, util::buffer::metrics::MetricNormalize};
+use crate::sinks::{
+    prelude::*,
+    util::buffer::metrics::{MetricNormalize, MetricSet},
+};
 
 mod config;
 mod request_builder;
@@ -20,7 +24,7 @@ pub use config::RemoteWriteConfig;
 #[derive(Clone, Copy, Debug, Derivative)]
 #[derivative(Default)]
 #[serde(rename_all = "lowercase")]
-pub(super) enum Compression {
+pub enum Compression {
     /// Snappy.
     #[derivative(Default)]
     Snappy,
diff --git a/src/sinks/prometheus/remote_write/request_builder.rs b/src/sinks/prometheus/remote_write/request_builder.rs
index 091ead9896f78..f13ffea2a454b 100644
--- a/src/sinks/prometheus/remote_write/request_builder.rs
+++ b/src/sinks/prometheus/remote_write/request_builder.rs
@@ -1,31 +1,32 @@
-use std::io;
+use std::io::{self, Read};
 
 use bytes::{Bytes, BytesMut};
-use http::{Request, Uri};
+use prost::Message;
 use vector_core::config::telemetry;
 
-use crate::{
-    http::Auth,
-    sinks::{
-        prelude::*,
-        prometheus::{collector, PrometheusRemoteWriteAuth},
-        util::PartitionInnerBuffer,
-    },
+use crate::sinks::{
+    prelude::*,
+    prometheus::{collector, collector::MetricCollector as _},
 };
 
 use super::{sink::RemoteWriteMetric, PartitionKey};
 
-pub(crate) struct RemoteWriteEncoder;
+pub(crate) struct RemoteWriteEncoder {
+    pub(super) default_namespace: Option<String>,
+    pub(super) buckets: Vec<f64>,
+    pub(super) quantiles: Vec<f64>,
+}
 
 impl encoding::Encoder<Vec<RemoteWriteMetric>> for RemoteWriteEncoder {
     fn encode_input(
         &self,
         input: Vec<RemoteWriteMetric>,
-        writer: &mut dyn std::io::Write,
-    ) -> std::io::Result<(usize, GroupedCountByteSize)> {
+        writer: &mut dyn io::Write,
+    ) -> io::Result<(usize, GroupedCountByteSize)> {
         let mut byte_size = telemetry().create_request_count_byte_size();
 
         let mut time_series = collector::TimeSeries::new();
+        let len = input.len();
         for metric in input {
             byte_size.add_event(&metric.metric, metric.estimated_json_encoded_size_of());
 
@@ -33,7 +34,7 @@ impl encoding::Encoder<Vec<RemoteWriteMetric>> for RemoteWriteEncoder {
                 self.default_namespace.as_deref(),
                 &self.buckets,
                 &self.quantiles,
-                &metric,
+                &metric.metric,
             );
         }
         let request = time_series.finish();
@@ -42,17 +43,18 @@ impl encoding::Encoder<Vec<RemoteWriteMetric>> for RemoteWriteEncoder {
         request.encode(&mut out).expect("Out of memory");
         let body = out.freeze();
 
-        write_all(writer, input.len(), body.as_ref())?;
+        write_all(writer, len, body.as_ref())?;
 
         Ok((body.len(), byte_size))
     }
 }
 
+#[derive(Clone)]
 pub(super) struct RemoteWriteRequest {
-    pub(super) request: http::Request<Bytes>,
+    pub(super) request: Bytes,
+    pub(super) tenant_id: Option<String>,
     finalizers: EventFinalizers,
     metadata: RequestMetadata,
-    tenant_id: String,
 }
 
 impl Finalizable for RemoteWriteRequest {
@@ -71,57 +73,17 @@ impl MetaDescriptive for RemoteWriteRequest {
     }
 }
 
-struct RemoteWriteMetadata {
+pub(super) struct RemoteWriteMetadata {
     finalizers: EventFinalizers,
     tenant_id: Option<String>,
 }
 
-pub struct RemoteWriteRequestBuilder {
-    pub compression: super::Compression,
-    pub encoder: RemoteWriteEncoder,
-    pub http_auth: Option<Auth>,
-    pub endpoint: String,
+pub(super) struct RemoteWriteRequestBuilder {
+    pub(super) compression: super::Compression,
+    pub(super) encoder: RemoteWriteEncoder,
 }
 
-pub fn build_request(
-    method: http::Method,
-    endpoint: Uri,
-    compression: super::Compression,
-    auth: Option<Auth>,
-    body: Vec<u8>,
-    tenant_id: Option<String>,
-) -> Request<hyper::Body> {
-    let content_encoding = convert_compression_to_content_encoding(compression);
-
-    let mut builder = http::Request::builder()
-        .method(method)
-        .uri(endpoint)
-        .header("X-Prometheus-Remote-Write-Version", "0.1.0")
-        .header("Content-Encoding", content_encoding)
-        .header("Content-Type", "application/x-protobuf");
-
-    if let Some(tenant_id) = &tenant_id {
-        builder = builder.header("X-Scope-OrgID", tenant_id);
-    }
-
-    let mut request = builder.body(body.into()).unwrap();
-    if let Some(http_auth) = &auth {
-        http_auth.apply(&mut request);
-    }
-
-    // if let Some(credentials_provider) = &self.credentials_provider {
-    //     sign_request(&mut request, credentials_provider, &self.aws_region).await?;
-    // }
-
-    // let (parts, body) = request.into_parts();
-    // let request: Request<hyper::Body> = hyper::Request::from_parts(parts, body.into());
-
-    request
-}
-
-impl RequestBuilder<PartitionInnerBuffer<Vec<RemoteWriteMetric>, PartitionKey>>
-    for RemoteWriteRequestBuilder
-{
+impl RequestBuilder<(PartitionKey, Vec<RemoteWriteMetric>)> for RemoteWriteRequestBuilder {
     type Metadata = RemoteWriteMetadata;
     type Events = Vec<RemoteWriteMetric>;
     type Encoder = RemoteWriteEncoder;
@@ -130,7 +92,9 @@ impl RequestBuilder<PartitionInnerBuffer<Vec<RemoteWriteMetric>, PartitionKey>>
     type Error = io::Error;
 
     fn compression(&self) -> Compression {
-        self.compression
+        // Since we also support Snappy which isn't handled by the main compression,
+        // we need to do this ourselves.
+        Compression::None
     }
 
     fn encoder(&self) -> &Self::Encoder {
@@ -139,16 +103,16 @@ impl RequestBuilder<PartitionInnerBuffer<Vec<RemoteWriteMetric>, PartitionKey>>
 
     fn split_input(
         &self,
-        input: PartitionInnerBuffer<Vec<RemoteWriteMetric>, PartitionKey>,
+        input: (PartitionKey, Vec<RemoteWriteMetric>),
     ) -> (Self::Metadata, RequestMetadataBuilder, Self::Events) {
-        let (events, key) = input.into_parts();
+        let (key, mut events) = input;
         let builder = RequestMetadataBuilder::from_events(&events);
         let metadata = RemoteWriteMetadata {
-            finalizers: input.take_finalizers(),
-            tenant_id: key,
+            finalizers: events.take_finalizers(),
+            tenant_id: key.tenant_id,
         };
 
-        (metadata, builder, input)
+        (metadata, builder, events)
     }
 
     fn build_request(
@@ -158,19 +122,13 @@ impl RequestBuilder<PartitionInnerBuffer<Vec<RemoteWriteMetric>, PartitionKey>>
         payload: EncodeResult<Self::Payload>,
     ) -> Self::Request {
         let body = payload.into_payload();
-        let request = build_request(
-            http::Method::POST,
-            &self.endpoint,
-            self.compression,
-            self.http_auth,
-            body,
-            metadata.tenant_id,
-        );
+        let body = compress_block(self.compression, body);
+
         RemoteWriteRequest {
-            request,
+            request: Bytes::from(body),
             finalizers: metadata.finalizers,
             tenant_id: metadata.tenant_id,
-            metadata,
+            metadata: request_metadata,
         }
     }
 }
@@ -192,11 +150,3 @@ fn compress_block(compression: super::Compression, data: Bytes) -> Vec<u8> {
         }
     }
 }
-
-const fn convert_compression_to_content_encoding(compression: super::Compression) -> &'static str {
-    match compression {
-        super::Compression::Snappy => "snappy",
-        super::Compression::Gzip => "gzip",
-        super::Compression::Zstd => "zstd",
-    }
-}
diff --git a/src/sinks/prometheus/remote_write/service.rs b/src/sinks/prometheus/remote_write/service.rs
index 28a6667ea97a6..1c060f6a42d8a 100644
--- a/src/sinks/prometheus/remote_write/service.rs
+++ b/src/sinks/prometheus/remote_write/service.rs
@@ -1,49 +1,62 @@
-use std::{
-    sync::Arc,
-    task::{Context, Poll},
-};
+use std::task::{Context, Poll};
+
+#[cfg(feature = "aws-core")]
+use aws_credential_types::provider::SharedCredentialsProvider;
+#[cfg(feature = "aws-core")]
+use aws_types::region::Region;
 
-use aws_sdk_cloudwatch::Region;
 use bytes::Bytes;
-use http::StatusCode;
+use http::{StatusCode, Uri};
 
 use super::request_builder::RemoteWriteRequest;
-use crate::{http::HttpClient, sinks::prelude::*};
+use crate::{
+    http::{HttpClient, HttpError},
+    internal_events::EndpointBytesSent,
+    sinks::{prelude::*, util::auth::Auth},
+};
 
-#[derive(Clone)]
-pub(super) struct RemoteWriteService {
-    pub endpoint: Uri,
-    pub aws_region: Option<Region>,
-    pub credentials_provider: Option<SharedCredentialsProvider>,
-
-    default_namespace: Option<String>,
-    client: HttpClient,
-    buckets: Vec<f64>,
-    quantiles: Vec<f64>,
-    compression: Compression,
+#[derive(Debug, Default, Clone)]
+pub(super) struct RemoteWriteRetryLogic;
+
+impl RetryLogic for RemoteWriteRetryLogic {
+    type Error = HttpError;
+    type Response = RemoteWriteResponse;
+
+    fn is_retriable_error(&self, _error: &Self::Error) -> bool {
+        true
+    }
+
+    fn should_retry_response(&self, response: &Self::Response) -> RetryAction {
+        let status = response.response.status();
+
+        match status {
+            StatusCode::TOO_MANY_REQUESTS => RetryAction::Retry("too many requests".into()),
+            StatusCode::NOT_IMPLEMENTED => {
+                RetryAction::DontRetry("endpoint not implemented".into())
+            }
+            _ if status.is_server_error() => RetryAction::Retry(
+                format!(
+                    "{}: {}",
+                    status,
+                    String::from_utf8_lossy(response.response.body())
+                )
+                .into(),
+            ),
+            _ if status.is_success() => RetryAction::Successful,
+            _ => RetryAction::DontRetry(format!("response status: {}", status).into()),
+        }
+    }
 }
 
-impl RemoteWriteService {
-    // fn encode_events(&self, metrics: Vec<Metric>) -> Bytes {
-    //     let mut time_series = collector::TimeSeries::new();
-    //     for metric in metrics {
-    //         time_series.encode_metric(
-    //             self.default_namespace.as_deref(),
-    //             &self.buckets,
-    //             &self.quantiles,
-    //             &metric,
-    //         );
-    //     }
-    //     let request = time_series.finish();
-
-    //     let mut out = BytesMut::with_capacity(request.encoded_len());
-    //     request.encode(&mut out).expect("Out of memory");
-    //     out.freeze()
-    // }
+#[derive(Clone)]
+pub(super) struct RemoteWriteService {
+    pub(super) endpoint: Uri,
+    pub(super) auth: Option<Auth>,
+    pub(super) client: HttpClient,
+    pub(super) compression: super::Compression,
 }
 
-struct RemoteWriteResponse {
-    byte_size: usize,
+pub(super) struct RemoteWriteResponse {
     json_size: GroupedCountByteSize,
     response: http::Response<Bytes>,
 }
@@ -60,14 +73,10 @@ impl DriverResponse for RemoteWriteResponse {
     fn events_sent(&self) -> &GroupedCountByteSize {
         &self.json_size
     }
-
-    fn bytes_sent(&self) -> Option<usize> {
-        Some(self.byte_size)
-    }
 }
 
 impl Service<RemoteWriteRequest> for RemoteWriteService {
-    type Response = http::Response<Bytes>;
+    type Response = RemoteWriteResponse;
     type Error = crate::Error;
     type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
 
@@ -76,47 +85,54 @@ impl Service<RemoteWriteRequest> for RemoteWriteService {
     }
 
     // Emission of internal events for errors and dropped events is handled upstream by the caller.
-    fn call(&mut self, request: RemoteWriteRequest) -> Self::Future {
-        // let (events, key) = buffer.into_parts();
-        // let body = self.encode_events(events);
-        // let body = compress_block(self.compression, body);
-
+    fn call(&mut self, mut request: RemoteWriteRequest) -> Self::Future {
         let client = self.client.clone();
-        // let request_builder = Arc::clone(&self.http_request_builder);
+        let metadata = std::mem::take(request.metadata_mut());
+        let events_byte_size = metadata.into_events_estimated_json_encoded_byte_size();
+        let endpoint = self.endpoint.clone();
+        let auth = self.auth.clone();
+        let compression = self.compression;
 
         Box::pin(async move {
-            // let request = request_builder
-            //     .build_request(http::Method::POST, body, key.tenant_id)
-            //     .await?;
-
-            let metadata = std::mem::take(request.metadata_mut());
-            let mut request = request.request;
-
-            if let Some(credentials_provider) = &self.credentials_provider {
-                sign_request(&mut request, credentials_provider, &self.aws_region).await?;
-            }
+            let request = build_request(
+                http::Method::POST,
+                &endpoint,
+                compression,
+                request.request,
+                request.tenant_id.as_ref(),
+                auth,
+            )
+            .await?;
 
             let (parts, body) = request.into_parts();
-            let request: Request<hyper::Body> = hyper::Request::from_parts(parts, body.into());
+            let request: hyper::Request<hyper::Body> =
+                hyper::Request::from_parts(parts, body.into());
 
-            let (protocol, endpoint) = uri::protocol_endpoint(request.uri().clone());
-
-            let response = client.send(request).await?;
+            let response = client.send(request.map(hyper::Body::from)).await?;
             let (parts, body) = response.into_parts();
             let body = hyper::body::to_bytes(body).await?;
             let byte_size = body.len();
 
             let response = hyper::Response::from_parts(parts, body);
 
+            if response.status().is_success() {
+                // We can't rely on the framework to emit this because we need to specify the additional `endpoint` tag.
+                emit!(EndpointBytesSent {
+                    byte_size,
+                    protocol: "http",
+                    endpoint: &endpoint.to_string(),
+                });
+            }
+
             Ok(RemoteWriteResponse {
-                byte_size,
-                json_size: todo!(),
+                json_size: events_byte_size,
                 response,
             })
         })
     }
 }
 
+#[cfg(feature = "aws-core")]
 async fn sign_request(
     request: &mut http::Request<Bytes>,
     credentials_provider: &SharedCredentialsProvider,
@@ -124,3 +140,48 @@ async fn sign_request(
 ) -> crate::Result<()> {
     crate::aws::sign_request("aps", request, credentials_provider, region).await
 }
+
+pub(super) async fn build_request(
+    method: http::Method,
+    endpoint: &Uri,
+    compression: super::Compression,
+    body: Bytes,
+    tenant_id: Option<&String>,
+    auth: Option<Auth>,
+) -> crate::Result<http::Request<Bytes>> {
+    let content_encoding = convert_compression_to_content_encoding(compression);
+
+    let mut builder = http::Request::builder()
+        .method(method)
+        .uri(endpoint)
+        .header("X-Prometheus-Remote-Write-Version", "0.1.0")
+        .header("Content-Encoding", content_encoding)
+        .header("Content-Type", "application/x-protobuf");
+
+    if let Some(tenant_id) = tenant_id {
+        builder = builder.header("X-Scope-OrgID", tenant_id);
+    }
+
+    let mut request = builder.body(body)?;
+
+    if let Some(auth) = auth {
+        match auth {
+            Auth::Basic(http_auth) => http_auth.apply(&mut request),
+            #[cfg(feature = "aws-core")]
+            Auth::Aws {
+                credentials_provider: provider,
+                region,
+            } => sign_request(&mut request, &provider, &Some(region.clone())).await?,
+        }
+    }
+
+    Ok(request)
+}
+
+const fn convert_compression_to_content_encoding(compression: super::Compression) -> &'static str {
+    match compression {
+        super::Compression::Snappy => "snappy",
+        super::Compression::Gzip => "gzip",
+        super::Compression::Zstd => "zstd",
+    }
+}
diff --git a/src/sinks/prometheus/remote_write/sink.rs b/src/sinks/prometheus/remote_write/sink.rs
index 6f3e06d7a01ca..75603ce0990cb 100644
--- a/src/sinks/prometheus/remote_write/sink.rs
+++ b/src/sinks/prometheus/remote_write/sink.rs
@@ -1,15 +1,24 @@
+use std::fmt;
+
+use vector_common::byte_size_of::ByteSizeOf;
 use vector_core::event::Metric;
 
-use crate::sinks::{prelude::*, prometheus::PrometheusRemoteWriteAuth};
+use crate::sinks::prelude::*;
 
 use super::{
-    request_builder::{RemoteWriteEncoder, RemoteWriteRequestBuilder},
-    PrometheusMetricNormalize,
+    request_builder::{RemoteWriteEncoder, RemoteWriteRequest, RemoteWriteRequestBuilder},
+    PartitionKey, PrometheusMetricNormalize,
 };
 
 pub(super) struct RemoteWriteMetric {
     pub(super) metric: Metric,
-    tenant_id: String,
+    tenant_id: Option<String>,
+}
+
+impl Finalizable for RemoteWriteMetric {
+    fn take_finalizers(&mut self) -> EventFinalizers {
+        self.metric.take_finalizers()
+    }
 }
 
 impl GetEventCountTags for RemoteWriteMetric {
@@ -24,6 +33,12 @@ impl EstimatedJsonEncodedSizeOf for RemoteWriteMetric {
     }
 }
 
+impl ByteSizeOf for RemoteWriteMetric {
+    fn allocated_bytes(&self) -> usize {
+        self.metric.allocated_bytes()
+    }
+}
+
 #[derive(Clone, Copy, Debug, Default)]
 pub struct PrometheusRemoteWriteDefaultBatchSettings;
 
@@ -31,11 +46,12 @@ pub(super) struct PrometheusTenantIdPartitioner;
 
 impl Partitioner for PrometheusTenantIdPartitioner {
     type Item = RemoteWriteMetric;
-    type Key = String;
+    type Key = PartitionKey;
 
     fn partition(&self, item: &Self::Item) -> Self::Key {
-        // TODO do this better.
-        item.tenant_id.clone()
+        PartitionKey {
+            tenant_id: item.tenant_id.clone(),
+        }
     }
 }
 
@@ -45,103 +61,89 @@ impl SinkBatchSettings for PrometheusRemoteWriteDefaultBatchSettings {
     const TIMEOUT_SECS: f64 = 1.0;
 }
 
-pub(super) struct RemoteWriteSink {
-    tenant_id: Option<Template>,
-    batch_settings: BatcherSettings,
-    compression: super::Compression,
-    http_auth: Option<PrometheusRemoteWriteAuth>,
-    endpoint: String,
+pub(super) struct RemoteWriteSink<S> {
+    pub(super) tenant_id: Option<Template>,
+    pub(super) batch_settings: BatcherSettings,
+    pub(super) compression: super::Compression,
+    pub(super) default_namespace: Option<String>,
+    pub(super) buckets: Vec<f64>,
+    pub(super) quantiles: Vec<f64>,
+    pub(super) service: S,
 }
 
-impl RemoteWriteSink {
-    fn make_remote_write_event(&self, metric: Metric) -> Option<RemoteWriteMetric> {
-        let tenant_id = self.tenant_id.as_ref().and_then(|template| {
-            template
-                .render_string(&metric)
-                .map_err(|error| {
-                    emit!(TemplateRenderingError {
-                        error,
-                        field: Some("tenant_id"),
-                        drop_event: true,
-                    })
-                })
-                .ok()
-        })?;
-
-        Some(RemoteWriteMetric { metric, tenant_id })
-    }
-
+impl<S> RemoteWriteSink<S>
+where
+    S: Service<RemoteWriteRequest> + Send + 'static,
+    S::Future: Send + 'static,
+    S::Response: DriverResponse + Send + 'static,
+    S::Error: fmt::Debug + Into<crate::Error> + Send,
+{
     async fn run_inner(self: Box<Self>, input: BoxStream<'_, Event>) -> Result<(), ()> {
         let request_builder = RemoteWriteRequestBuilder {
-            endpoint: self.endpoint.clone(),
             compression: self.compression,
-            encoder: RemoteWriteEncoder,
-            http_auth: self.http_auth,
+            encoder: RemoteWriteEncoder {
+                default_namespace: self.default_namespace.clone(),
+                buckets: self.buckets.clone(),
+                quantiles: self.quantiles.clone(),
+            },
         };
 
-        let service = RemoteWriteService {
-            default_namespace: self.default_namespace.clone(),
-            client,
-            buckets,
-            quantiles,
-            http_request_builder,
-            compression: self.compression,
-        };
-        let service = ServiceBuilder::new().service(service);
+        let batch_settings = self.batch_settings;
+        let tenant_id = self.tenant_id.clone();
+        let service = self.service;
 
         input
             .filter_map(|event| future::ready(event.try_into_metric()))
             .normalized_with_default::<PrometheusMetricNormalize>()
-            .filter_map(|event| self.make_remote_write_event(event))
-            .batched_partitioned(PrometheusTenantIdPartitioner, self.batch_settings)
-            .request_builder(request_builder)
+            .filter_map(move |event| {
+                future::ready(make_remote_write_event(tenant_id.as_ref(), event))
+            })
+            .batched_partitioned(PrometheusTenantIdPartitioner, batch_settings)
+            .request_builder(None, request_builder)
+            .filter_map(|request| async move {
+                match request {
+                    Err(e) => {
+                        error!("Failed to build Remote Write request: {:?}.", e);
+                        None
+                    }
+                    Ok(req) => Some(req),
+                }
+            })
             .into_driver(service)
-            .protocol("http")
             .run()
             .await
-
-        /*
-            let buffer = PartitionBuffer::new(MetricsBuffer::new(batch.size));
-            let mut normalizer = MetricNormalizer::<PrometheusMetricNormalize>::default();
-
-            request_settings
-                .partition_sink(HttpRetryLogic, service, buffer, batch.timeout)
-                .with_flat_map(move |event: Event| {
-                    let byte_size = event.size_of();
-                    let json_size = event.estimated_json_encoded_size_of();
-
-                    stream::iter(normalizer.normalize(event.into_metric()).map(|event| {
-                        let tenant_id = tenant_id.as_ref().and_then(|template| {
-                            template
-                                .render_string(&event)
-                                .map_err(|error| {
-                                    emit!(TemplateRenderingError {
-                                        error,
-                                        field: Some("tenant_id"),
-                                        drop_event: true,
-                                    })
-                                })
-                                .ok()
-                        });
-                        let key = PartitionKey { tenant_id };
-                        Ok(EncodedEvent::new(
-                            PartitionInnerBuffer::new(event, key),
-                            byte_size,
-                            json_size,
-                        ))
-                    }))
-                })
-                .sink_map_err(
-                    |error| error!(message = "Prometheus remote_write sink error.", %error),
-                )
-        ;
-        */
     }
 }
 
 #[async_trait]
-impl StreamSink<Event> for RemoteWriteSink {
+impl<S> StreamSink<Event> for RemoteWriteSink<S>
+where
+    S: Service<RemoteWriteRequest> + Send + 'static,
+    S::Future: Send + 'static,
+    S::Response: DriverResponse + Send + 'static,
+    S::Error: fmt::Debug + Into<crate::Error> + Send,
+{
     async fn run(self: Box<Self>, input: BoxStream<'_, Event>) -> Result<(), ()> {
         self.run_inner(input).await
     }
 }
+
+fn make_remote_write_event(
+    tenant_id: Option<&Template>,
+    metric: Metric,
+) -> Option<RemoteWriteMetric> {
+    let tenant_id = tenant_id.and_then(|template| {
+        template
+            .render_string(&metric)
+            .map_err(|error| {
+                emit!(TemplateRenderingError {
+                    error,
+                    field: Some("tenant_id"),
+                    drop_event: true,
+                })
+            })
+            .ok()
+    });
+
+    Some(RemoteWriteMetric { metric, tenant_id })
+}
diff --git a/src/sinks/prometheus/remote_write/tests.rs b/src/sinks/prometheus/remote_write/tests.rs
index d7950686b486e..da88851802a52 100644
--- a/src/sinks/prometheus/remote_write/tests.rs
+++ b/src/sinks/prometheus/remote_write/tests.rs
@@ -1,30 +1,28 @@
-use crate::sinks::prelude::*;
-
-#[cfg(test)]
-mod tests {
-    use futures::StreamExt;
-    use http::HeaderMap;
-    use indoc::indoc;
-    use prometheus_parser::proto;
-    use vector_core::metric_tags;
-
-    use super::*;
-    use crate::{
-        config::SinkContext,
-        event::{MetricKind, MetricValue},
-        sinks::util::test::build_test_server,
-        test_util::{
-            self,
-            components::{assert_sink_compliance, HTTP_SINK_TAGS},
-        },
-    };
-
-    #[test]
-    fn generate_config() {
-        crate::test_util::test_generate_config::<RemoteWriteConfig>();
-    }
+use bytes::Bytes;
+use futures::StreamExt;
+use http::HeaderMap;
+use indoc::indoc;
+use prometheus_parser::proto;
+use prost::Message;
+use vector_core::metric_tags;
+
+use super::*;
+use crate::{
+    config::SinkContext,
+    event::{MetricKind, MetricValue},
+    sinks::{prometheus::remote_write::RemoteWriteConfig, util::test::build_test_server},
+    test_util::{
+        self,
+        components::{assert_sink_compliance, HTTP_SINK_TAGS},
+    },
+};
+
+#[test]
+fn generate_config() {
+    crate::test_util::test_generate_config::<RemoteWriteConfig>();
+}
 
-    macro_rules! labels {
+macro_rules! labels {
         ( $( $name:expr => $value:expr ),* ) => {
             vec![ $( proto::Label {
                 name: $name.to_string(),
@@ -33,59 +31,60 @@ mod tests {
         }
     }
 
-    #[tokio::test]
-    async fn sends_request() {
-        let outputs = send_request("", vec![create_event("gauge-2".into(), 32.0)]).await;
-        assert_eq!(outputs.len(), 1);
-        let (headers, req) = &outputs[0];
-
-        assert!(!headers.contains_key("x-scope-orgid"));
-
-        assert_eq!(req.timeseries.len(), 1);
-        assert_eq!(
-            req.timeseries[0].labels,
-            labels!("__name__" => "gauge-2", "production" => "true", "region" => "us-west-1")
-        );
-        assert_eq!(req.timeseries[0].samples.len(), 1);
-        assert_eq!(req.timeseries[0].samples[0].value, 32.0);
-        assert_eq!(req.metadata.len(), 1);
-        assert_eq!(req.metadata[0].r#type, proto::MetricType::Gauge as i32);
-        assert_eq!(req.metadata[0].metric_family_name, "gauge-2");
-    }
+#[tokio::test]
+async fn sends_request() {
+    let outputs = send_request("", vec![create_event("gauge-2".into(), 32.0)]).await;
+    assert_eq!(outputs.len(), 1);
+    let (headers, req) = &outputs[0];
+
+    assert!(!headers.contains_key("x-scope-orgid"));
+
+    assert_eq!(req.timeseries.len(), 1);
+    assert_eq!(
+        req.timeseries[0].labels,
+        labels!("__name__" => "gauge-2", "production" => "true", "region" => "us-west-1")
+    );
+    assert_eq!(req.timeseries[0].samples.len(), 1);
+    assert_eq!(req.timeseries[0].samples[0].value, 32.0);
+    assert_eq!(req.metadata.len(), 1);
+    assert_eq!(req.metadata[0].r#type, proto::MetricType::Gauge as i32);
+    assert_eq!(req.metadata[0].metric_family_name, "gauge-2");
+}
 
-    #[tokio::test]
-    async fn sends_authenticated_request() {
-        let outputs = send_request(
-            indoc! {r#"
+#[tokio::test]
+async fn sends_authenticated_request() {
+    let outputs = send_request(
+        indoc! {r#"
                 tenant_id = "tenant-%Y"
                 [auth]
                 strategy = "basic"
                 user = "user"
                 password = "password"
             "#},
-            vec![create_event("gauge-2".into(), 32.0)],
-        )
-        .await;
-
-        assert_eq!(outputs.len(), 1);
-        let (_headers, req) = &outputs[0];
-
-        assert_eq!(req.timeseries.len(), 1);
-        assert_eq!(
-            req.timeseries[0].labels,
-            labels!("__name__" => "gauge-2", "production" => "true", "region" => "us-west-1")
-        );
-        assert_eq!(req.timeseries[0].samples.len(), 1);
-        assert_eq!(req.timeseries[0].samples[0].value, 32.0);
-        assert_eq!(req.metadata.len(), 1);
-        assert_eq!(req.metadata[0].r#type, proto::MetricType::Gauge as i32);
-        assert_eq!(req.metadata[0].metric_family_name, "gauge-2");
-    }
+        vec![create_event("gauge-2".into(), 32.0)],
+    )
+    .await;
+
+    assert_eq!(outputs.len(), 1);
+    let (_headers, req) = &outputs[0];
+
+    assert_eq!(req.timeseries.len(), 1);
+    assert_eq!(
+        req.timeseries[0].labels,
+        labels!("__name__" => "gauge-2", "production" => "true", "region" => "us-west-1")
+    );
+    assert_eq!(req.timeseries[0].samples.len(), 1);
+    assert_eq!(req.timeseries[0].samples[0].value, 32.0);
+    assert_eq!(req.metadata.len(), 1);
+    assert_eq!(req.metadata[0].r#type, proto::MetricType::Gauge as i32);
+    assert_eq!(req.metadata[0].metric_family_name, "gauge-2");
+}
 
-    #[tokio::test]
-    async fn sends_authenticated_aws_request() {
-        let outputs = send_request(
-            indoc! {r#"
+#[cfg(feature = "aws-config")]
+#[tokio::test]
+async fn sends_authenticated_aws_request() {
+    let outputs = send_request(
+        indoc! {r#"
                 tenant_id = "tenant-%Y"
                 [aws]
                 region = "foo"
@@ -94,137 +93,133 @@ mod tests {
                 access_key_id = "foo"
                 secret_access_key = "bar"
             "#},
-            vec![create_event("gauge-2".into(), 32.0)],
-        )
-        .await;
+        vec![create_event("gauge-2".into(), 32.0)],
+    )
+    .await;
 
-        assert_eq!(outputs.len(), 1);
-        let (headers, _req) = &outputs[0];
+    assert_eq!(outputs.len(), 1);
+    let (headers, _req) = &outputs[0];
 
-        let auth = headers["authorization"]
-            .to_str()
-            .expect("Missing AWS authorization header");
-        assert!(auth.starts_with("AWS4-HMAC-SHA256"));
-    }
+    let auth = headers["authorization"]
+        .to_str()
+        .expect("Missing AWS authorization header");
+    assert!(auth.starts_with("AWS4-HMAC-SHA256"));
+}
 
-    #[tokio::test]
-    async fn sends_x_scope_orgid_header() {
-        let outputs = send_request(
-            r#"tenant_id = "tenant""#,
-            vec![create_event("gauge-3".into(), 12.0)],
-        )
-        .await;
-
-        assert_eq!(outputs.len(), 1);
-        let (headers, _) = &outputs[0];
-        assert_eq!(headers["x-scope-orgid"], "tenant");
-    }
+#[tokio::test]
+async fn sends_x_scope_orgid_header() {
+    let outputs = send_request(
+        r#"tenant_id = "tenant""#,
+        vec![create_event("gauge-3".into(), 12.0)],
+    )
+    .await;
+
+    assert_eq!(outputs.len(), 1);
+    let (headers, _) = &outputs[0];
+    assert_eq!(headers["x-scope-orgid"], "tenant");
+}
 
-    #[tokio::test]
-    async fn sends_templated_x_scope_orgid_header() {
-        let outputs = send_request(
-            r#"tenant_id = "tenant-%Y""#,
-            vec![create_event("gauge-3".into(), 12.0)],
-        )
-        .await;
-
-        assert_eq!(outputs.len(), 1);
-        let (headers, _) = &outputs[0];
-        let orgid = headers["x-scope-orgid"]
-            .to_str()
-            .expect("Missing x-scope-orgid header");
-        assert!(orgid.starts_with("tenant-20"));
-        assert_eq!(orgid.len(), 11);
-    }
+#[tokio::test]
+async fn sends_templated_x_scope_orgid_header() {
+    let outputs = send_request(
+        r#"tenant_id = "tenant-%Y""#,
+        vec![create_event("gauge-3".into(), 12.0)],
+    )
+    .await;
+
+    assert_eq!(outputs.len(), 1);
+    let (headers, _) = &outputs[0];
+    let orgid = headers["x-scope-orgid"]
+        .to_str()
+        .expect("Missing x-scope-orgid header");
+    assert!(orgid.starts_with("tenant-20"));
+    assert_eq!(orgid.len(), 11);
+}
 
-    #[tokio::test]
-    async fn retains_state_between_requests() {
-        // This sink converts all incremental events to absolute, and
-        // should accumulate their totals between batches.
-        let outputs = send_request(
-            r#"batch.max_events = 1"#,
-            vec![
-                create_inc_event("counter-1".into(), 12.0),
-                create_inc_event("counter-2".into(), 13.0),
-                create_inc_event("counter-1".into(), 14.0),
-            ],
-        )
-        .await;
-
-        assert_eq!(outputs.len(), 3);
-
-        let check_output = |index: usize, name: &str, value: f64| {
-            let (_, req) = &outputs[index];
-            assert_eq!(req.timeseries.len(), 1);
-            assert_eq!(req.timeseries[0].labels, labels!("__name__" => name));
-            assert_eq!(req.timeseries[0].samples.len(), 1);
-            assert_eq!(req.timeseries[0].samples[0].value, value);
-        };
-        check_output(0, "counter-1", 12.0);
-        check_output(1, "counter-2", 13.0);
-        check_output(2, "counter-1", 26.0);
-    }
+#[tokio::test]
+async fn retains_state_between_requests() {
+    // This sink converts all incremental events to absolute, and
+    // should accumulate their totals between batches.
+    let outputs = send_request(
+        r#"batch.max_events = 1"#,
+        vec![
+            create_inc_event("counter-1".into(), 12.0),
+            create_inc_event("counter-2".into(), 13.0),
+            create_inc_event("counter-1".into(), 14.0),
+        ],
+    )
+    .await;
+
+    assert_eq!(outputs.len(), 3);
+
+    let check_output = |index: usize, name: &str, value: f64| {
+        let (_, req) = &outputs[index];
+        assert_eq!(req.timeseries.len(), 1);
+        assert_eq!(req.timeseries[0].labels, labels!("__name__" => name));
+        assert_eq!(req.timeseries[0].samples.len(), 1);
+        assert_eq!(req.timeseries[0].samples[0].value, value);
+    };
+    check_output(0, "counter-1", 12.0);
+    check_output(1, "counter-2", 13.0);
+    check_output(2, "counter-1", 26.0);
+}
 
-    async fn send_request(
-        config: &str,
-        events: Vec<Event>,
-    ) -> Vec<(HeaderMap, proto::WriteRequest)> {
-        assert_sink_compliance(&HTTP_SINK_TAGS, async {
-            let addr = test_util::next_addr();
-            let (rx, trigger, server) = build_test_server(addr);
-            tokio::spawn(server);
-
-            let config = format!("endpoint = \"http://{}/write\"\n{}", addr, config);
-            let config: RemoteWriteConfig = toml::from_str(&config).unwrap();
-            let cx = SinkContext::default();
-
-            let (sink, _) = config.build(cx).await.unwrap();
-            sink.run_events(events).await.unwrap();
-
-            drop(trigger);
-
-            rx.map(|(parts, body)| {
-                assert_eq!(parts.method, "POST");
-                assert_eq!(parts.uri.path(), "/write");
-                let headers = parts.headers;
-                assert_eq!(headers["x-prometheus-remote-write-version"], "0.1.0");
-                assert_eq!(headers["content-encoding"], "snappy");
-                assert_eq!(headers["content-type"], "application/x-protobuf");
-
-                if config.auth.is_some() {
-                    assert!(headers.contains_key("authorization"));
-                }
-
-                let decoded = snap::raw::Decoder::new()
-                    .decompress_vec(&body)
-                    .expect("Invalid snappy compressed data");
-                let request =
-                    proto::WriteRequest::decode(Bytes::from(decoded)).expect("Invalid protobuf");
-                (headers, request)
-            })
-            .collect::<Vec<_>>()
-            .await
+async fn send_request(config: &str, events: Vec<Event>) -> Vec<(HeaderMap, proto::WriteRequest)> {
+    assert_sink_compliance(&HTTP_SINK_TAGS, async {
+        let addr = test_util::next_addr();
+        let (rx, trigger, server) = build_test_server(addr);
+        tokio::spawn(server);
+
+        let config = format!("endpoint = \"http://{}/write\"\n{}", addr, config);
+        let config: RemoteWriteConfig = toml::from_str(&config).unwrap();
+        let cx = SinkContext::default();
+
+        let (sink, _) = config.build(cx).await.unwrap();
+        sink.run_events(events).await.unwrap();
+
+        drop(trigger);
+
+        rx.map(|(parts, body)| {
+            assert_eq!(parts.method, "POST");
+            assert_eq!(parts.uri.path(), "/write");
+            let headers = parts.headers;
+            assert_eq!(headers["x-prometheus-remote-write-version"], "0.1.0");
+            assert_eq!(headers["content-encoding"], "snappy");
+            assert_eq!(headers["content-type"], "application/x-protobuf");
+
+            if config.auth.is_some() {
+                assert!(headers.contains_key("authorization"));
+            }
+
+            let decoded = snap::raw::Decoder::new()
+                .decompress_vec(&body)
+                .expect("Invalid snappy compressed data");
+            let request =
+                proto::WriteRequest::decode(Bytes::from(decoded)).expect("Invalid protobuf");
+            (headers, request)
         })
+        .collect::<Vec<_>>()
         .await
-    }
-
-    pub(super) fn create_event(name: String, value: f64) -> Event {
-        Metric::new(name, MetricKind::Absolute, MetricValue::Gauge { value })
-            .with_tags(Some(metric_tags!(
-                "region" => "us-west-1",
-                "production" => "true",
-            )))
-            .with_timestamp(Some(chrono::Utc::now()))
-            .into()
-    }
+    })
+    .await
+}
 
-    fn create_inc_event(name: String, value: f64) -> Event {
-        Metric::new(
-            name,
-            MetricKind::Incremental,
-            MetricValue::Counter { value },
-        )
+pub(super) fn create_event(name: String, value: f64) -> Event {
+    Metric::new(name, MetricKind::Absolute, MetricValue::Gauge { value })
+        .with_tags(Some(metric_tags!(
+            "region" => "us-west-1",
+            "production" => "true",
+        )))
         .with_timestamp(Some(chrono::Utc::now()))
         .into()
-    }
+}
+
+fn create_inc_event(name: String, value: f64) -> Event {
+    Metric::new(
+        name,
+        MetricKind::Incremental,
+        MetricValue::Counter { value },
+    )
+    .with_timestamp(Some(chrono::Utc::now()))
+    .into()
 }
diff --git a/src/sinks/util/buffer/metrics/normalize.rs b/src/sinks/util/buffer/metrics/normalize.rs
index afd39f8c31c73..1ada2126098e6 100644
--- a/src/sinks/util/buffer/metrics/normalize.rs
+++ b/src/sinks/util/buffer/metrics/normalize.rs
@@ -85,7 +85,7 @@ type MetricEntry = (MetricData, EventMetadata);
 ///
 /// This is primarily a wrapper around `HashMap` with convenience methods to make it easier to perform
 /// normalization-specific operations.
-#[derive(Clone, Default)]
+#[derive(Clone, Default, Debug)]
 pub struct MetricSet(HashMap<MetricSeries, MetricEntry>);
 
 impl MetricSet {
diff --git a/src/sources/prometheus/remote_write.rs b/src/sources/prometheus/remote_write.rs
index 6fe08a32c7b51..9a2e4356975b6 100644
--- a/src/sources/prometheus/remote_write.rs
+++ b/src/sources/prometheus/remote_write.rs
@@ -155,11 +155,7 @@ mod test {
     use crate::{
         config::{SinkConfig, SinkContext},
         sinks::prometheus::remote_write::RemoteWriteConfig,
-        test_util::{
-            self,
-            components::{assert_source_compliance, HTTP_PUSH_SOURCE_TAGS},
-            wait_for_tcp,
-        },
+        test_util::{self, wait_for_tcp},
         tls::MaybeTlsSettings,
         SourceSender,
     };
@@ -180,54 +176,51 @@ mod test {
     }
 
     async fn receives_metrics(tls: Option<TlsEnableableConfig>) {
-        assert_source_compliance(&HTTP_PUSH_SOURCE_TAGS, async {
-            let address = test_util::next_addr();
-            let (tx, rx) = SourceSender::new_test_finalize(EventStatus::Delivered);
-
-            let proto = MaybeTlsSettings::from_config(&tls, true)
-                .unwrap()
-                .http_protocol_name();
-            let source = PrometheusRemoteWriteConfig {
-                address,
-                auth: None,
-                tls: tls.clone(),
-                acknowledgements: SourceAcknowledgementsConfig::default(),
-            };
-            let source = source
-                .build(SourceContext::new_test(tx, None))
-                .await
-                .unwrap();
-            tokio::spawn(source);
-            wait_for_tcp(address).await;
-
-            let sink = RemoteWriteConfig {
-                endpoint: format!("{}://localhost:{}/", proto, address.port()),
-                tls: tls.map(|tls| tls.options),
-                ..Default::default()
-            };
-            let (sink, _) = sink
-                .build(SinkContext::default())
-                .await
-                .expect("Error building config.");
-
-            let events = make_events();
-            let events_copy = events.clone();
-            let mut output = test_util::spawn_collect_ready(
-                async move {
-                    sink.run_events(events_copy).await.unwrap();
-                },
-                rx,
-                1,
-            )
-            .await;
-
-            // The MetricBuffer used by the sink may reorder the metrics, so
-            // put them back into order before comparing.
-            output.sort_unstable_by_key(|event| event.as_metric().name().to_owned());
+        let address = test_util::next_addr();
+        let (tx, rx) = SourceSender::new_test_finalize(EventStatus::Delivered);
 
-            vector_common::assert_event_data_eq!(events, output);
-        })
+        let proto = MaybeTlsSettings::from_config(&tls, true)
+            .unwrap()
+            .http_protocol_name();
+        let source = PrometheusRemoteWriteConfig {
+            address,
+            auth: None,
+            tls: tls.clone(),
+            acknowledgements: SourceAcknowledgementsConfig::default(),
+        };
+        let source = source
+            .build(SourceContext::new_test(tx, None))
+            .await
+            .unwrap();
+        tokio::spawn(source);
+        wait_for_tcp(address).await;
+
+        let sink = RemoteWriteConfig {
+            endpoint: format!("{}://localhost:{}/", proto, address.port()),
+            tls: tls.map(|tls| tls.options),
+            ..Default::default()
+        };
+        let (sink, _) = sink
+            .build(SinkContext::default())
+            .await
+            .expect("Error building config.");
+
+        let events = make_events();
+        let events_copy = events.clone();
+        let mut output = test_util::spawn_collect_ready(
+            async move {
+                sink.run_events(events_copy).await.unwrap();
+            },
+            rx,
+            1,
+        )
         .await;
+
+        // The MetricBuffer used by the sink may reorder the metrics, so
+        // put them back into order before comparing.
+        output.sort_unstable_by_key(|event| event.as_metric().name().to_owned());
+
+        vector_common::assert_event_data_eq!(events, output);
     }
 
     fn make_events() -> Vec<Event> {
@@ -278,69 +271,66 @@ mod test {
     /// we accept the metric, but take the last label in the list.
     #[tokio::test]
     async fn receives_metrics_duplicate_labels() {
-        assert_source_compliance(&HTTP_PUSH_SOURCE_TAGS, async {
-            let address = test_util::next_addr();
-            let (tx, rx) = SourceSender::new_test_finalize(EventStatus::Delivered);
-
-            let source = PrometheusRemoteWriteConfig {
-                address,
-                auth: None,
-                tls: None,
-                acknowledgements: SourceAcknowledgementsConfig::default(),
-            };
-            let source = source
-                .build(SourceContext::new_test(tx, None))
-                .await
-                .unwrap();
-            tokio::spawn(source);
-            wait_for_tcp(address).await;
-
-            let sink = RemoteWriteConfig {
-                endpoint: format!("http://localhost:{}/", address.port()),
-                ..Default::default()
-            };
-            let (sink, _) = sink
-                .build(SinkContext::default())
-                .await
-                .expect("Error building config.");
-
-            let timestamp = Utc::now().trunc_subsecs(3);
-
-            let events = vec![Metric::new(
-                "gauge_2",
-                MetricKind::Absolute,
-                MetricValue::Gauge { value: 41.0 },
-            )
-            .with_timestamp(Some(timestamp))
-            .with_tags(Some(metric_tags! {
-                "code" => "200".to_string(),
-                "code" => "success".to_string(),
-            }))
-            .into()];
-
-            let expected = vec![Metric::new(
-                "gauge_2",
-                MetricKind::Absolute,
-                MetricValue::Gauge { value: 41.0 },
-            )
-            .with_timestamp(Some(timestamp))
-            .with_tags(Some(metric_tags! {
-                "code" => "success".to_string(),
-            }))
-            .into()];
-
-            let output = test_util::spawn_collect_ready(
-                async move {
-                    sink.run_events(events).await.unwrap();
-                },
-                rx,
-                1,
-            )
-            .await;
+        let address = test_util::next_addr();
+        let (tx, rx) = SourceSender::new_test_finalize(EventStatus::Delivered);
 
-            vector_common::assert_event_data_eq!(expected, output);
-        })
+        let source = PrometheusRemoteWriteConfig {
+            address,
+            auth: None,
+            tls: None,
+            acknowledgements: SourceAcknowledgementsConfig::default(),
+        };
+        let source = source
+            .build(SourceContext::new_test(tx, None))
+            .await
+            .unwrap();
+        tokio::spawn(source);
+        wait_for_tcp(address).await;
+
+        let sink = RemoteWriteConfig {
+            endpoint: format!("http://localhost:{}/", address.port()),
+            ..Default::default()
+        };
+        let (sink, _) = sink
+            .build(SinkContext::default())
+            .await
+            .expect("Error building config.");
+
+        let timestamp = Utc::now().trunc_subsecs(3);
+
+        let events = vec![Metric::new(
+            "gauge_2",
+            MetricKind::Absolute,
+            MetricValue::Gauge { value: 41.0 },
+        )
+        .with_timestamp(Some(timestamp))
+        .with_tags(Some(metric_tags! {
+            "code" => "200".to_string(),
+            "code" => "success".to_string(),
+        }))
+        .into()];
+
+        let expected = vec![Metric::new(
+            "gauge_2",
+            MetricKind::Absolute,
+            MetricValue::Gauge { value: 41.0 },
+        )
+        .with_timestamp(Some(timestamp))
+        .with_tags(Some(metric_tags! {
+            "code" => "success".to_string(),
+        }))
+        .into()];
+
+        let output = test_util::spawn_collect_ready(
+            async move {
+                sink.run_events(events).await.unwrap();
+            },
+            rx,
+            1,
+        )
         .await;
+
+        vector_common::assert_event_data_eq!(expected, output);
     }
 }
 

From 31d7fb2ec2ed0b5dfe85433071f4166bdff03195 Mon Sep 17 00:00:00 2001
From: Stephen Wakely <fungus.humungus@gmail.com>
Date: Tue, 26 Sep 2023 14:32:29 +0100
Subject: [PATCH 04/32] Use shared compression

Signed-off-by: Stephen Wakely <fungus.humungus@gmail.com>
---
 src/sinks/prometheus/remote_write/mod.rs      | 10 +++
 .../remote_write/request_builder.rs           | 29 +-----
 src/sinks/prometheus/remote_write/service.rs  |  4 +-
 src/sinks/util/buffer/compression.rs          | 27 +++++-
 src/sinks/util/buffer/mod.rs                  |  9 ++
 src/sinks/util/compressor.rs                  | 11 ++-
 src/sinks/util/mod.rs                         |  1 +
 src/sinks/util/snappy.rs                      | 88 +++++++++++++++++++
 8 files changed, 147 insertions(+), 32 deletions(-)
 create mode 100644 src/sinks/util/snappy.rs

diff --git a/src/sinks/prometheus/remote_write/mod.rs b/src/sinks/prometheus/remote_write/mod.rs
index 137e7e309a9cb..c34751b313a32 100644
--- a/src/sinks/prometheus/remote_write/mod.rs
+++ b/src/sinks/prometheus/remote_write/mod.rs
@@ -36,6 +36,16 @@ pub enum Compression {
     Zstd,
 }
 
+impl From<Compression> for crate::sinks::util::Compression {
+    fn from(value: Compression) -> Self {
+        match value {
+            Compression::Snappy => crate::sinks::util::Compression::Snappy,
+            Compression::Gzip => crate::sinks::util::Compression::Gzip(Default::default()),
+            Compression::Zstd => crate::sinks::util::Compression::Zstd(Default::default()),
+        }
+    }
+}
+
 #[derive(Debug, Snafu)]
 enum Errors {
     #[snafu(display(r#"Prometheus remote_write sink cannot accept "set" metrics"#))]
diff --git a/src/sinks/prometheus/remote_write/request_builder.rs b/src/sinks/prometheus/remote_write/request_builder.rs
index f13ffea2a454b..19504df1b0b09 100644
--- a/src/sinks/prometheus/remote_write/request_builder.rs
+++ b/src/sinks/prometheus/remote_write/request_builder.rs
@@ -1,4 +1,4 @@
-use std::io::{self, Read};
+use std::io;
 
 use bytes::{Bytes, BytesMut};
 use prost::Message;
@@ -92,9 +92,7 @@ impl RequestBuilder<(PartitionKey, Vec<RemoteWriteMetric>)> for RemoteWriteReque
     type Error = io::Error;
 
     fn compression(&self) -> Compression {
-        // Since we also support Snappy which isn't handled by the main compression,
-        // we need to do this ourselves.
-        Compression::None
+        self.compression.into()
     }
 
     fn encoder(&self) -> &Self::Encoder {
@@ -121,32 +119,11 @@ impl RequestBuilder<(PartitionKey, Vec<RemoteWriteMetric>)> for RemoteWriteReque
         request_metadata: RequestMetadata,
         payload: EncodeResult<Self::Payload>,
     ) -> Self::Request {
-        let body = payload.into_payload();
-        let body = compress_block(self.compression, body);
-
         RemoteWriteRequest {
-            request: Bytes::from(body),
+            request: payload.into_payload(),
             finalizers: metadata.finalizers,
             tenant_id: metadata.tenant_id,
             metadata: request_metadata,
         }
     }
 }
-
-fn compress_block(compression: super::Compression, data: Bytes) -> Vec<u8> {
-    match compression {
-        super::Compression::Snappy => snap::raw::Encoder::new()
-            .compress_vec(&data)
-            .expect("snap compression failed, please report"),
-        super::Compression::Gzip => {
-            let mut buf = Vec::new();
-            flate2::read::GzEncoder::new(data.as_ref(), flate2::Compression::default())
-                .read_to_end(&mut buf)
-                .expect("gzip compression failed, please report");
-            buf
-        }
-        super::Compression::Zstd => {
-            zstd::encode_all(data.as_ref(), 0).expect("zstd compression failed, please report")
-        }
-    }
-}
diff --git a/src/sinks/prometheus/remote_write/service.rs b/src/sinks/prometheus/remote_write/service.rs
index 1c060f6a42d8a..9718fb0f66d6f 100644
--- a/src/sinks/prometheus/remote_write/service.rs
+++ b/src/sinks/prometheus/remote_write/service.rs
@@ -88,7 +88,7 @@ impl Service<RemoteWriteRequest> for RemoteWriteService {
     fn call(&mut self, mut request: RemoteWriteRequest) -> Self::Future {
         let client = self.client.clone();
         let metadata = std::mem::take(request.metadata_mut());
-        let events_byte_size = metadata.into_events_estimated_json_encoded_byte_size();
+        let json_size = metadata.into_events_estimated_json_encoded_byte_size();
         let endpoint = self.endpoint.clone();
         let auth = self.auth.clone();
         let compression = self.compression;
@@ -125,7 +125,7 @@ impl Service<RemoteWriteRequest> for RemoteWriteService {
             }
 
             Ok(RemoteWriteResponse {
-                json_size: events_byte_size,
+                json_size,
                 response,
             })
         })
diff --git a/src/sinks/util/buffer/compression.rs b/src/sinks/util/buffer/compression.rs
index 11796b046104a..47607f1bb4428 100644
--- a/src/sinks/util/buffer/compression.rs
+++ b/src/sinks/util/buffer/compression.rs
@@ -37,6 +37,11 @@ pub enum Compression {
     ///
     /// [zstd]: https://facebook.github.io/zstd/
     Zstd(CompressionLevel),
+
+    /// [Snappy][snappy] compression.
+    ///
+    /// [snappy]: https://github.com/google/snappy/blob/main/docs/README.md
+    Snappy,
 }
 
 impl Compression {
@@ -70,6 +75,7 @@ impl Compression {
             Self::Gzip(_) => Some("gzip"),
             Self::Zlib(_) => Some("deflate"),
             Self::Zstd(_) => Some("zstd"),
+            Self::Snappy => Some("snappy"),
         }
     }
 
@@ -78,6 +84,7 @@ impl Compression {
             Self::Gzip(_) => Some("gzip"),
             Self::Zlib(_) => Some("deflate"),
             Self::Zstd(_) => Some("zstd"),
+            Self::Snappy => Some("snappy"),
             _ => None,
         }
     }
@@ -88,6 +95,7 @@ impl Compression {
             Self::Gzip(_) => "log.gz",
             Self::Zlib(_) => "log.zz",
             Self::Zstd(_) => "log.zst",
+            Self::Snappy => "log.snappy",
         }
     }
 
@@ -97,12 +105,13 @@ impl Compression {
             Compression::Gzip(_) => 9,
             Compression::Zlib(_) => 9,
             Compression::Zstd(_) => 21,
+            Compression::Snappy => 0,
         }
     }
 
     pub const fn compression_level(self) -> CompressionLevel {
         match self {
-            Self::None => CompressionLevel::None,
+            Self::None | Self::Snappy => CompressionLevel::None,
             Self::Gzip(level) | Self::Zlib(level) | Self::Zstd(level) => level,
         }
     }
@@ -117,6 +126,7 @@ impl fmt::Display for Compression {
             Compression::Zstd(ref level) => {
                 write!(f, "zstd({})", ZstdCompressionLevel::from(*level))
             }
+            Compression::Snappy => write!(f, "snappy"),
         }
     }
 }
@@ -144,6 +154,7 @@ impl<'de> de::Deserialize<'de> for Compression {
                     "gzip" => Ok(Compression::gzip_default()),
                     "zlib" => Ok(Compression::zlib_default()),
                     "zstd" => Ok(Compression::zstd_default()),
+                    "snappy" => Ok(Compression::Snappy),
                     _ => Err(de::Error::invalid_value(
                         de::Unexpected::Str(s),
                         &r#""none" or "gzip" or "zlib" or "zstd""#,
@@ -187,9 +198,13 @@ impl<'de> de::Deserialize<'de> for Compression {
                     "gzip" => Ok(Compression::Gzip(level.unwrap_or_default())),
                     "zlib" => Ok(Compression::Zlib(level.unwrap_or_default())),
                     "zstd" => Ok(Compression::Zstd(level.unwrap_or_default())),
+                    "snappy" => match level {
+                        Some(_) => Err(de::Error::unknown_field("level", &[])),
+                        None => Ok(Compression::Snappy),
+                    },
                     algorithm => Err(de::Error::unknown_variant(
                         algorithm,
-                        &["none", "gzip", "zlib", "zstd"],
+                        &["none", "gzip", "zlib", "zstd", "snappy"],
                     )),
                 }?;
 
@@ -252,6 +267,7 @@ impl ser::Serialize for Compression {
                     serializer.serialize_str("zstd")
                 }
             }
+            Compression::Snappy => serializer.serialize_str("snappy"),
         }
     }
 }
@@ -509,6 +525,7 @@ mod test {
             (r#""none""#, Compression::None),
             (r#""gzip""#, Compression::Gzip(CompressionLevel::default())),
             (r#""zlib""#, Compression::Zlib(CompressionLevel::default())),
+            (r#""snappy""#, Compression::Snappy),
             (r#"{"algorithm": "none"}"#, Compression::None),
             (
                 r#"{"algorithm": "gzip"}"#,
@@ -551,7 +568,7 @@ mod test {
             ),
             (
                 r#"{"algorithm": "b42"}"#,
-                r#"unknown variant `b42`, expected one of `none`, `gzip`, `zlib`, `zstd` at line 1 column 20"#,
+                r#"unknown variant `b42`, expected one of `none`, `gzip`, `zlib`, `zstd`, `snappy` at line 1 column 20"#,
             ),
             (
                 r#"{"algorithm": "none", "level": "default"}"#,
@@ -581,6 +598,10 @@ mod test {
                 r#"{"algorithm": "zstd", "level": 22}"#,
                 r#"invalid value `22`, expected value in range [0, 21] at line 1 column 34"#,
             ),
+            (
+                r#"{"algorithm": "snappy", "level": 3}"#,
+                r#"unknown field `level`, there are no fields at line 1 column 35"#,
+            ),
         ];
         for (source, result) in fixtures_invalid.iter() {
             let deserialized: Result<Compression, _> = serde_json::from_str(source);
diff --git a/src/sinks/util/buffer/mod.rs b/src/sinks/util/buffer/mod.rs
index 296e14b02f9c8..f0ab3dd6af8fb 100644
--- a/src/sinks/util/buffer/mod.rs
+++ b/src/sinks/util/buffer/mod.rs
@@ -5,6 +5,7 @@ use flate2::write::{GzEncoder, ZlibEncoder};
 
 use super::{
     batch::{err_event_too_large, Batch, BatchSize, PushResult},
+    snappy::SnappyEncoder,
     zstd::ZstdEncoder,
 };
 
@@ -32,6 +33,7 @@ pub enum InnerBuffer {
     Gzip(GzEncoder<bytes::buf::Writer<BytesMut>>),
     Zlib(ZlibEncoder<bytes::buf::Writer<BytesMut>>),
     Zstd(ZstdEncoder<bytes::buf::Writer<BytesMut>>),
+    Snappy(SnappyEncoder<bytes::buf::Writer<BytesMut>>),
 }
 
 impl Buffer {
@@ -62,6 +64,7 @@ impl Buffer {
                     ZstdEncoder::new(writer, level.into())
                         .expect("Zstd encoder should not fail on init."),
                 ),
+                Compression::Snappy => InnerBuffer::Snappy(SnappyEncoder::new(writer)),
             }
         })
     }
@@ -81,6 +84,7 @@ impl Buffer {
             InnerBuffer::Zstd(inner) => {
                 inner.write_all(input).unwrap();
             }
+            InnerBuffer::Snappy(inner) => inner.write_all(input).unwrap(),
         }
     }
 
@@ -92,6 +96,7 @@ impl Buffer {
                 InnerBuffer::Gzip(inner) => inner.get_ref().get_ref().is_empty(),
                 InnerBuffer::Zlib(inner) => inner.get_ref().get_ref().is_empty(),
                 InnerBuffer::Zstd(inner) => inner.get_ref().get_ref().is_empty(),
+                InnerBuffer::Snappy(inner) => inner.is_empty(),
             })
             .unwrap_or(true)
     }
@@ -142,6 +147,10 @@ impl Batch for Buffer {
                 .finish()
                 .expect("This can't fail because the inner writer is a Vec")
                 .into_inner(),
+            Some(InnerBuffer::Snappy(inner)) => inner
+                .finish()
+                .expect("This can't fail because the inner writer is a Vec")
+                .into_inner(),
             None => BytesMut::new(),
         }
     }
diff --git a/src/sinks/util/compressor.rs b/src/sinks/util/compressor.rs
index b918893a39bd2..f6d3c01ccca87 100644
--- a/src/sinks/util/compressor.rs
+++ b/src/sinks/util/compressor.rs
@@ -3,13 +3,14 @@ use std::io;
 use bytes::{BufMut, BytesMut};
 use flate2::write::{GzEncoder, ZlibEncoder};
 
-use super::{zstd::ZstdEncoder, Compression};
+use super::{snappy::SnappyEncoder, zstd::ZstdEncoder, Compression};
 
 enum Writer {
     Plain(bytes::buf::Writer<BytesMut>),
     Gzip(GzEncoder<bytes::buf::Writer<BytesMut>>),
     Zlib(ZlibEncoder<bytes::buf::Writer<BytesMut>>),
     Zstd(ZstdEncoder<bytes::buf::Writer<BytesMut>>),
+    Snappy(SnappyEncoder<bytes::buf::Writer<BytesMut>>),
 }
 
 impl Writer {
@@ -19,6 +20,7 @@ impl Writer {
             Writer::Gzip(inner) => inner.get_ref().get_ref(),
             Writer::Zlib(inner) => inner.get_ref().get_ref(),
             Writer::Zstd(inner) => inner.get_ref().get_ref(),
+            Writer::Snappy(inner) => inner.get_ref().get_ref(),
         }
     }
 }
@@ -35,6 +37,7 @@ impl From<Compression> for Writer {
                     .expect("Zstd encoder should not fail on init.");
                 Writer::Zstd(encoder)
             }
+            Compression::Snappy => Writer::Snappy(SnappyEncoder::new(writer)),
         }
     }
 }
@@ -47,6 +50,7 @@ impl io::Write for Writer {
             Writer::Gzip(writer) => writer.write(buf),
             Writer::Zlib(writer) => writer.write(buf),
             Writer::Zstd(writer) => writer.write(buf),
+            Writer::Snappy(writer) => writer.write(buf),
         }
     }
 
@@ -56,6 +60,7 @@ impl io::Write for Writer {
             Writer::Gzip(writer) => writer.flush(),
             Writer::Zlib(writer) => writer.flush(),
             Writer::Zstd(writer) => writer.flush(),
+            Writer::Snappy(writer) => writer.flush(),
         }
     }
 }
@@ -98,6 +103,7 @@ impl Compressor {
             Writer::Gzip(writer) => writer.finish()?,
             Writer::Zlib(writer) => writer.finish()?,
             Writer::Zstd(writer) => writer.finish()?,
+            Writer::Snappy(writer) => writer.finish()?,
         }
         .into_inner();
 
@@ -125,6 +131,9 @@ impl Compressor {
             Writer::Zstd(writer) => writer
                 .finish()
                 .expect("zstd writer should not fail to finish"),
+            Writer::Snappy(writer) => writer
+                .finish()
+                .expect("snappy writer should not fail to finish"),
         }
         .into_inner()
     }
diff --git a/src/sinks/util/mod.rs b/src/sinks/util/mod.rs
index 1fcae4842cd7f..f5c09b97f380f 100644
--- a/src/sinks/util/mod.rs
+++ b/src/sinks/util/mod.rs
@@ -16,6 +16,7 @@ pub mod request_builder;
 pub mod retries;
 pub mod service;
 pub mod sink;
+pub mod snappy;
 pub mod socket_bytes_sink;
 pub mod statistic;
 pub mod tcp;
diff --git a/src/sinks/util/snappy.rs b/src/sinks/util/snappy.rs
new file mode 100644
index 0000000000000..69830bed42546
--- /dev/null
+++ b/src/sinks/util/snappy.rs
@@ -0,0 +1,88 @@
+//! An encoder for [Snappy] compression.
+//! Whilst there does exist a [Writer] implementation for Snappy, this compresses
+//! using the [Snappy frame format][frame], which is not quite what we want. So
+//! instead this encoder buffers the data in a [`Vec`] until the end. The `raw`
+//! compressor is then used to compress the data and writes it to the provided
+//! writer.
+//!
+//! [Snappy]: https://github.com/google/snappy/blob/main/docs/README.md
+//! [Writer]: https://docs.rs/snap/latest/snap/write/struct.FrameEncoder.html
+//! [frame]: https://github.com/google/snappy/blob/master/framing_format.txt
+
+use std::io;
+
+use snap::raw::Encoder;
+
+pub struct SnappyEncoder<W: io::Write> {
+    writer: W,
+    buffer: Vec<u8>,
+}
+
+impl<W: io::Write> SnappyEncoder<W> {
+    pub fn new(writer: W) -> Self {
+        Self {
+            writer,
+            buffer: Vec::new(),
+        }
+    }
+
+    pub fn finish(mut self) -> io::Result<W> {
+        let mut encoder = Encoder::new();
+        let compressed = encoder.compress_vec(&self.buffer)?;
+
+        self.writer.write_all(&compressed)?;
+
+        Ok(self.writer)
+    }
+
+    pub fn get_ref(&self) -> &W {
+        &self.writer
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.buffer.is_empty()
+    }
+}
+
+impl<W: io::Write> io::Write for SnappyEncoder<W> {
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        self.buffer.extend_from_slice(buf);
+        Ok(buf.len())
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        Ok(())
+    }
+}
+
+impl<W: io::Write + std::fmt::Debug> std::fmt::Debug for SnappyEncoder<W> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("SnappyEncoder")
+            .field("inner", &self.get_ref())
+            .finish()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::io::Write;
+
+    use bytes::{BufMut, BytesMut};
+
+    use super::*;
+
+    #[test]
+    fn is_empty() {
+        let writer = BytesMut::with_capacity(64).writer();
+        let mut encoder = SnappyEncoder::new(writer);
+
+        encoder.write_all(b"I am a potato").unwrap();
+
+        // Because we are buffering the results until the end, the writer will be
+        // empty, but our buffer won't be. The `is_empty` function is provided to
+        // allow us to determine if data has been written to the encoder without having
+        // to check the writer.
+        assert!(encoder.get_ref().get_ref().is_empty());
+        assert!(!encoder.is_empty());
+    }
+}

From 4be6d00d205def94ce8e7d22d5c38944bcb7c298 Mon Sep 17 00:00:00 2001
From: Stephen Wakely <fungus.humungus@gmail.com>
Date: Tue, 26 Sep 2023 15:45:20 +0100
Subject: [PATCH 05/32] Clippy

Signed-off-by: Stephen Wakely <fungus.humungus@gmail.com>
---
 src/sinks/util/snappy.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/sinks/util/snappy.rs b/src/sinks/util/snappy.rs
index 69830bed42546..a46edd41ebf97 100644
--- a/src/sinks/util/snappy.rs
+++ b/src/sinks/util/snappy.rs
@@ -35,7 +35,7 @@ impl<W: io::Write> SnappyEncoder<W> {
         Ok(self.writer)
     }
 
-    pub fn get_ref(&self) -> &W {
+    pub const fn get_ref(&self) -> &W {
         &self.writer
     }
 

From 59707170055ce0a68ad880965033683eba91f6aa Mon Sep 17 00:00:00 2001
From: Stephen Wakely <fungus.humungus@gmail.com>
Date: Tue, 26 Sep 2023 17:08:14 +0100
Subject: [PATCH 06/32] Some comments

Signed-off-by: Stephen Wakely <fungus.humungus@gmail.com>
---
 src/sinks/prometheus/remote_write/config.rs | 4 ----
 src/sinks/prometheus/remote_write/mod.rs    | 7 +++++++
 2 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/src/sinks/prometheus/remote_write/config.rs b/src/sinks/prometheus/remote_write/config.rs
index edf40ba267a61..6965c56ea5515 100644
--- a/src/sinks/prometheus/remote_write/config.rs
+++ b/src/sinks/prometheus/remote_write/config.rs
@@ -1,7 +1,3 @@
-// #[cfg(feature = "aws-core")]
-// use aws_credential_types::provider::SharedCredentialsProvider;
-// #[cfg(feature = "aws-core")]
-// use aws_types::region::Region;
 use http::Uri;
 use snafu::prelude::*;
 
diff --git a/src/sinks/prometheus/remote_write/mod.rs b/src/sinks/prometheus/remote_write/mod.rs
index c34751b313a32..32659aefef71a 100644
--- a/src/sinks/prometheus/remote_write/mod.rs
+++ b/src/sinks/prometheus/remote_write/mod.rs
@@ -1,3 +1,10 @@
+//! The Prometheus Remote Write [`vector_core::sink::VectorSink`].
+//! Contains the [`vector_core::sink::VectorSink`] instance that is responsible
+//! for taking a stream of [`vector_core::event::Event`] and forwarding
+//! them to a server via the [Prometheus Remote Write protocol][remote_write].
+//!
+//! [remote_write]: https://prometheus.io/docs/concepts/remote_write_spec/
+
 use snafu::prelude::*;
 use vector_core::event::Metric;
 

From 844a2ee58ae1893878c9c1f5a18c7361c5baa04c Mon Sep 17 00:00:00 2001
From: Stephen Wakely <fungus.humungus@gmail.com>
Date: Tue, 26 Sep 2023 17:51:57 +0100
Subject: [PATCH 07/32] Use HttpResponse

Signed-off-by: Stephen Wakely <fungus.humungus@gmail.com>
---
 src/sinks/prometheus/remote_write/config.rs  | 13 +++-
 src/sinks/prometheus/remote_write/service.rs | 79 ++++----------------
 src/sinks/util/http.rs                       |  4 +-
 3 files changed, 26 insertions(+), 70 deletions(-)

diff --git a/src/sinks/prometheus/remote_write/config.rs b/src/sinks/prometheus/remote_write/config.rs
index 6965c56ea5515..922adc0962783 100644
--- a/src/sinks/prometheus/remote_write/config.rs
+++ b/src/sinks/prometheus/remote_write/config.rs
@@ -3,11 +3,16 @@ use snafu::prelude::*;
 
 use crate::{
     http::HttpClient,
-    sinks::{prelude::*, prometheus::PrometheusRemoteWriteAuth, util::auth::Auth, UriParseSnafu},
+    sinks::{
+        prelude::*,
+        prometheus::PrometheusRemoteWriteAuth,
+        util::{auth::Auth, http::http_response_retry_logic},
+        UriParseSnafu,
+    },
 };
 
 use super::{
-    service::{build_request, RemoteWriteRetryLogic, RemoteWriteService},
+    service::{build_request, RemoteWriteService},
     sink::{PrometheusRemoteWriteDefaultBatchSettings, RemoteWriteSink},
 };
 
@@ -155,11 +160,11 @@ impl SinkConfig for RemoteWriteConfig {
         let service = RemoteWriteService {
             endpoint,
             client,
-            auth: auth.clone(),
+            auth,
             compression: self.compression,
         };
         let service = ServiceBuilder::new()
-            .settings(request_settings, RemoteWriteRetryLogic)
+            .settings(request_settings, http_response_retry_logic())
             .service(service);
 
         let sink = RemoteWriteSink {
diff --git a/src/sinks/prometheus/remote_write/service.rs b/src/sinks/prometheus/remote_write/service.rs
index 9718fb0f66d6f..2410059f39c82 100644
--- a/src/sinks/prometheus/remote_write/service.rs
+++ b/src/sinks/prometheus/remote_write/service.rs
@@ -6,48 +6,18 @@ use aws_credential_types::provider::SharedCredentialsProvider;
 use aws_types::region::Region;
 
 use bytes::Bytes;
-use http::{StatusCode, Uri};
+use http::Uri;
 
 use super::request_builder::RemoteWriteRequest;
 use crate::{
-    http::{HttpClient, HttpError},
+    http::HttpClient,
     internal_events::EndpointBytesSent,
-    sinks::{prelude::*, util::auth::Auth},
+    sinks::{
+        prelude::*,
+        util::{auth::Auth, http::HttpResponse},
+    },
 };
 
-#[derive(Debug, Default, Clone)]
-pub(super) struct RemoteWriteRetryLogic;
-
-impl RetryLogic for RemoteWriteRetryLogic {
-    type Error = HttpError;
-    type Response = RemoteWriteResponse;
-
-    fn is_retriable_error(&self, _error: &Self::Error) -> bool {
-        true
-    }
-
-    fn should_retry_response(&self, response: &Self::Response) -> RetryAction {
-        let status = response.response.status();
-
-        match status {
-            StatusCode::TOO_MANY_REQUESTS => RetryAction::Retry("too many requests".into()),
-            StatusCode::NOT_IMPLEMENTED => {
-                RetryAction::DontRetry("endpoint not implemented".into())
-            }
-            _ if status.is_server_error() => RetryAction::Retry(
-                format!(
-                    "{}: {}",
-                    status,
-                    String::from_utf8_lossy(response.response.body())
-                )
-                .into(),
-            ),
-            _ if status.is_success() => RetryAction::Successful,
-            _ => RetryAction::DontRetry(format!("response status: {}", status).into()),
-        }
-    }
-}
-
 #[derive(Clone)]
 pub(super) struct RemoteWriteService {
     pub(super) endpoint: Uri,
@@ -56,27 +26,8 @@ pub(super) struct RemoteWriteService {
     pub(super) compression: super::Compression,
 }
 
-pub(super) struct RemoteWriteResponse {
-    json_size: GroupedCountByteSize,
-    response: http::Response<Bytes>,
-}
-
-impl DriverResponse for RemoteWriteResponse {
-    fn event_status(&self) -> EventStatus {
-        if self.response.status().is_success() {
-            EventStatus::Delivered
-        } else {
-            EventStatus::Errored
-        }
-    }
-
-    fn events_sent(&self) -> &GroupedCountByteSize {
-        &self.json_size
-    }
-}
-
 impl Service<RemoteWriteRequest> for RemoteWriteService {
-    type Response = RemoteWriteResponse;
+    type Response = HttpResponse;
     type Error = crate::Error;
     type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
 
@@ -111,22 +62,22 @@ impl Service<RemoteWriteRequest> for RemoteWriteService {
             let response = client.send(request.map(hyper::Body::from)).await?;
             let (parts, body) = response.into_parts();
             let body = hyper::body::to_bytes(body).await?;
-            let byte_size = body.len();
-
-            let response = hyper::Response::from_parts(parts, body);
+            let raw_byte_size = body.len();
+            let http_response = hyper::Response::from_parts(parts, body);
 
-            if response.status().is_success() {
+            if http_response.status().is_success() {
                 // We can't rely on the framework to emit this because we need to specify the additional `endpoint` tag.
                 emit!(EndpointBytesSent {
-                    byte_size,
+                    byte_size: raw_byte_size,
                     protocol: "http",
                     endpoint: &endpoint.to_string(),
                 });
             }
 
-            Ok(RemoteWriteResponse {
-                json_size,
-                response,
+            Ok(HttpResponse {
+                events_byte_size: json_size,
+                http_response,
+                raw_byte_size,
             })
         })
     }
diff --git a/src/sinks/util/http.rs b/src/sinks/util/http.rs
index 435a28c9f685f..e88b9b9a77a27 100644
--- a/src/sinks/util/http.rs
+++ b/src/sinks/util/http.rs
@@ -654,8 +654,8 @@ impl ByteSizeOf for HttpRequest {
 /// Response type for use in the `Service` implementation of HTTP stream sinks.
 pub struct HttpResponse {
     pub http_response: Response<Bytes>,
-    events_byte_size: GroupedCountByteSize,
-    raw_byte_size: usize,
+    pub events_byte_size: GroupedCountByteSize,
+    pub raw_byte_size: usize,
 }
 
 impl DriverResponse for HttpResponse {

From e185761392d7ac5f5e642218baabac9699f62073 Mon Sep 17 00:00:00 2001
From: Stephen Wakely <fungus.humungus@gmail.com>
Date: Wed, 27 Sep 2023 10:03:48 +0100
Subject: [PATCH 08/32] Update request builder default

Signed-off-by: Stephen Wakely <fungus.humungus@gmail.com>
---
 src/sinks/prometheus/remote_write/sink.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/sinks/prometheus/remote_write/sink.rs b/src/sinks/prometheus/remote_write/sink.rs
index 75603ce0990cb..15d309d4cf6a9 100644
--- a/src/sinks/prometheus/remote_write/sink.rs
+++ b/src/sinks/prometheus/remote_write/sink.rs
@@ -99,7 +99,7 @@ where
                 future::ready(make_remote_write_event(tenant_id.as_ref(), event))
             })
             .batched_partitioned(PrometheusTenantIdPartitioner, batch_settings)
-            .request_builder(None, request_builder)
+            .request_builder(default_request_builder_concurrency_limit(), request_builder)
             .filter_map(|request| async move {
                 match request {
                     Err(e) => {

From ef4846925947d05fb73695ced8c3b88ccf513224 Mon Sep 17 00:00:00 2001
From: Stephen Wakely <fungus.humungus@gmail.com>
Date: Fri, 6 Oct 2023 12:04:53 +0100
Subject: [PATCH 09/32] Update PartitionBatcher to use BatchConfig

Signed-off-by: Stephen Wakely <fungus.humungus@gmail.com>
---
 .../src/stream/partitioned_batcher.rs         | 332 +++++++++---------
 src/sinks/aws_cloudwatch_logs/sink.rs         |   5 +-
 src/sinks/aws_kinesis/sink.rs                 |   4 +-
 src/sinks/azure_common/sink.rs                |   7 +-
 src/sinks/clickhouse/sink.rs                  |   4 +-
 src/sinks/datadog/logs/sink.rs                |   6 +-
 src/sinks/datadog/metrics/sink.rs             |   6 +-
 src/sinks/datadog/traces/sink.rs              |   7 +-
 src/sinks/gcs_common/sink.rs                  |   5 +-
 src/sinks/loki/sink.rs                        |   6 +-
 src/sinks/opendal_common.rs                   |   7 +-
 src/sinks/prometheus/remote_write/sink.rs     |  18 +-
 src/sinks/prometheus/remote_write/tests.rs    |   2 +
 src/sinks/s3_common/sink.rs                   |   7 +-
 src/sinks/splunk_hec/logs/sink.rs             |   3 +-
 src/sinks/splunk_hec/metrics/sink.rs          |   6 +-
 src/sinks/util/builder.rs                     |  10 +-
 17 files changed, 245 insertions(+), 190 deletions(-)

diff --git a/lib/vector-core/src/stream/partitioned_batcher.rs b/lib/vector-core/src/stream/partitioned_batcher.rs
index aa2f7e6b58bcd..165c3ad3d35ff 100644
--- a/lib/vector-core/src/stream/partitioned_batcher.rs
+++ b/lib/vector-core/src/stream/partitioned_batcher.rs
@@ -1,8 +1,6 @@
 use std::{
-    cmp,
     collections::HashMap,
     hash::{BuildHasherDefault, Hash},
-    mem,
     num::NonZeroUsize,
     pin::Pin,
     task::{ready, Context, Poll},
@@ -25,6 +23,8 @@ use crate::{
     ByteSizeOf,
 };
 
+use super::batcher::BatchConfig;
+
 /// A `KeyedTimer` based on `DelayQueue`.
 pub struct ExpirationQueue<K> {
     /// The timeout to give each new key entry
@@ -100,114 +100,6 @@ where
     }
 }
 
-/// A batch for use by `Batcher`
-///
-/// This structure is a private implementation detail that simplifies the
-/// implementation of `Batcher`. It is the actual store of items that come
-/// through the stream manipulated by `Batcher` plus limit information to signal
-/// when the `Batch` is full.
-struct Batch<I> {
-    /// The total number of `I` bytes stored, does not any overhead in this
-    /// structure.
-    allocated_bytes: usize,
-    /// The maximum number of elements allowed in this structure.
-    element_limit: usize,
-    /// The maximum number of allocated bytes (not including overhead) allowed
-    /// in this structure.
-    allocation_limit: usize,
-    /// The store of `I` elements.
-    elements: Vec<I>,
-}
-
-impl<I> ByteSizeOf for Batch<I> {
-    fn allocated_bytes(&self) -> usize {
-        self.allocated_bytes
-    }
-}
-
-impl<I> Batch<I>
-where
-    I: ByteSizeOf,
-{
-    /// Create a new Batch instance
-    ///
-    /// Creates a new batch instance with specific element and allocation limits. The element limit
-    /// is a maximum cap on the number of `I` instances. The allocation limit is a soft-max on the
-    /// number of allocated bytes stored in this batch, not taking into account overhead from this
-    /// structure itself.
-    ///
-    /// If `allocation_limit` is smaller than the size of `I` as reported by `std::mem::size_of`,
-    /// then the allocation limit will be raised such that the batch can hold a single instance of
-    /// `I`.  Likewise, `element_limit` will be raised such that it is always at least 1, ensuring
-    /// that a new batch can be pushed into.
-    fn new(element_limit: usize, allocation_limit: usize) -> Self {
-        // SAFETY: `element_limit` is always non-zero because `BatcherSettings` can only be
-        // constructed with `NonZeroUsize` versions of allocation limit/item limit.  `Batch` is also
-        // only constructable via `Batcher`.
-
-        // TODO: This may need to be reworked, because it's subtly wrong as-is.
-        // ByteSizeOf::size_of() always returns the size of the type itself, plus any "allocated
-        // bytes".  Thus, there are times when an item will be bigger than simply the size of the
-        // type itself (aka mem::size_of::<I>()) and thus than type of item would never fit in a
-        // batch where the `allocation_limit` is at or lower than the size of that item.
-        //
-        // We're counteracting this here by ensuring that the element limit is always at least 1.
-        let allocation_limit = cmp::max(allocation_limit, mem::size_of::<I>());
-        Self {
-            allocated_bytes: 0,
-            element_limit,
-            allocation_limit,
-            elements: Vec::with_capacity(128),
-        }
-    }
-
-    /// Unconditionally insert an element into the batch
-    ///
-    /// This function is similar to `push` except that the caller does not need
-    /// to call `has_space` prior to calling this and it will never
-    /// panic. Intended to be used only when insertion must not fail.
-    fn with(mut self, value: I) -> Self {
-        self.allocated_bytes += value.size_of();
-        self.elements.push(value);
-        self
-    }
-
-    /// Decompose the batch
-    ///
-    /// Called by the user when they want to get at the internal store of
-    /// items. Returns a tuple, the first element being the allocated size of
-    /// stored items and the second the store of items.
-    fn into_inner(self) -> Vec<I> {
-        self.elements
-    }
-
-    /// Whether the batch has space for a new item
-    ///
-    /// This function returns true of there is space both in terms of item count
-    /// and byte count for the given item, false otherwise.
-    fn has_space(&self, value: &I) -> bool {
-        let too_many_elements = self.elements.len() + 1 > self.element_limit;
-        let too_many_bytes = self.allocated_bytes + value.size_of() > self.allocation_limit;
-        !(too_many_elements || too_many_bytes)
-    }
-
-    /// Push an element into the batch
-    ///
-    /// This function pushes an element into the batch. Callers must be sure to
-    /// call `has_space` prior to calling this function and receive a positive
-    /// result.
-    ///
-    /// # Panics
-    ///
-    /// This function will panic if there is not sufficient space in the batch
-    /// for a new element to be inserted.
-    fn push(&mut self, value: I) {
-        assert!(self.has_space(&value));
-        self.allocated_bytes += value.size_of();
-        self.elements.push(value);
-    }
-}
-
 /// Controls the behavior of the batcher in terms of batch size and flush interval.
 ///
 /// This is a temporary solution for pushing in a fixed settings structure so we don't have to worry
@@ -287,18 +179,16 @@ impl BatcherSettings {
 }
 
 #[pin_project]
-pub struct PartitionedBatcher<St, Prt, KT>
+pub struct PartitionedBatcher<St, Prt, KT, C>
 where
     Prt: Partitioner,
 {
-    /// The total number of bytes a single batch in this struct is allowed to
-    /// hold.
-    batch_allocation_limit: usize,
-    /// The maximum number of items that are allowed per-batch
-    batch_item_limit: usize,
+    /// A closure that retrievs a new `BatchConfig` when needed to batch a
+    /// new partition.
+    state: Box<dyn Fn() -> C + Send>,
     /// The store of live batches. Note that the key here is an option type,
     /// on account of the interface of `Prt`.
-    batches: HashMap<Prt::Key, Batch<Prt::Item>, BuildHasherDefault<XxHash64>>,
+    batches: HashMap<Prt::Key, C, BuildHasherDefault<XxHash64>>,
     /// The store of 'closed' batches. When this is not empty it will be
     /// preferentially flushed prior to consuming any new items from the
     /// underlying stream.
@@ -312,44 +202,44 @@ where
     stream: Fuse<St>,
 }
 
-impl<St, Prt> PartitionedBatcher<St, Prt, ExpirationQueue<Prt::Key>>
+impl<St, Prt, C> PartitionedBatcher<St, Prt, ExpirationQueue<Prt::Key>, C>
 where
     St: Stream<Item = Prt::Item>,
     Prt: Partitioner + Unpin,
     Prt::Key: Eq + Hash + Clone,
     Prt::Item: ByteSizeOf,
+    C: BatchConfig<Prt::Item>,
 {
-    pub fn new(stream: St, partitioner: Prt, settings: BatcherSettings) -> Self {
+    pub fn new(stream: St, partitioner: Prt, settings: Box<dyn Fn() -> C + Send>) -> Self {
+        let timeout = settings().timeout();
         Self {
-            batch_allocation_limit: settings.size_limit,
-            batch_item_limit: settings.item_limit,
+            state: settings,
             batches: HashMap::default(),
             closed_batches: Vec::default(),
-            timer: ExpirationQueue::new(settings.timeout),
+            timer: ExpirationQueue::new(timeout),
             partitioner,
             stream: stream.fuse(),
         }
     }
 }
 
-impl<St, Prt, KT> PartitionedBatcher<St, Prt, KT>
+#[cfg(test)]
+impl<St, Prt, KT, C> PartitionedBatcher<St, Prt, KT, C>
 where
     St: Stream<Item = Prt::Item>,
     Prt: Partitioner + Unpin,
     Prt::Key: Eq + Hash + Clone,
     Prt::Item: ByteSizeOf,
+    C: BatchConfig<Prt::Item>,
 {
     pub fn with_timer(
         stream: St,
         partitioner: Prt,
         timer: KT,
-        batch_item_limit: NonZeroUsize,
-        batch_allocation_limit: Option<NonZeroUsize>,
+        settings: Box<dyn Fn() -> C + Send>,
     ) -> Self {
         Self {
-            batch_allocation_limit: batch_allocation_limit
-                .map_or(usize::max_value(), NonZeroUsize::get),
-            batch_item_limit: batch_item_limit.get(),
+            state: settings,
             batches: HashMap::default(),
             closed_batches: Vec::default(),
             timer,
@@ -359,13 +249,14 @@ where
     }
 }
 
-impl<St, Prt, KT> Stream for PartitionedBatcher<St, Prt, KT>
+impl<St, Prt, KT, C> Stream for PartitionedBatcher<St, Prt, KT, C>
 where
     St: Stream<Item = Prt::Item>,
     Prt: Partitioner + Unpin,
     Prt::Key: Eq + Hash + Clone,
     Prt::Item: ByteSizeOf,
     KT: KeyedTimer<Prt::Key>,
+    C: BatchConfig<Prt::Item, Batch = Vec<Prt::Item>>,
 {
     type Item = (Prt::Key, Vec<Prt::Item>);
 
@@ -385,13 +276,11 @@ where
                     // here but still be usable later if more entries are added.
                     Poll::Pending | Poll::Ready(None) => return Poll::Pending,
                     Poll::Ready(Some(item_key)) => {
-                        let batch = this
+                        let mut batch = this
                             .batches
                             .remove(&item_key)
                             .expect("batch should exist if it is set to expire");
-                        this.closed_batches.push((item_key, batch.into_inner()));
-
-                        continue;
+                        this.closed_batches.push((item_key, batch.take_batch()));
                     }
                 },
                 Poll::Ready(None) => {
@@ -405,7 +294,7 @@ where
                         this.closed_batches.extend(
                             this.batches
                                 .drain()
-                                .map(|(key, batch)| (key, batch.into_inner())),
+                                .map(|(key, mut batch)| (key, batch.take_batch())),
                         );
                         continue;
                     }
@@ -413,33 +302,37 @@ where
                 }
                 Poll::Ready(Some(item)) => {
                     let item_key = this.partitioner.partition(&item);
-                    let item_limit: usize = *this.batch_item_limit;
-                    let alloc_limit: usize = *this.batch_allocation_limit;
-
-                    if let Some(batch) = this.batches.get_mut(&item_key) {
-                        if batch.has_space(&item) {
-                            // When there's space in the partition batch just
-                            // push the item in and loop back around.
-                            batch.push(item);
-                        } else {
-                            let new_batch = Batch::new(item_limit, alloc_limit).with(item);
-                            let batch = mem::replace(batch, new_batch);
 
-                            // The batch for this partition key was set to
-                            // expire, but now it's overflowed and must be
-                            // pushed out, so now we reset the batch timeout.
-                            this.timer.insert(item_key.clone());
-
-                            this.closed_batches.push((item_key, batch.into_inner()));
+                    // Get the batch for this partition, or create a new one.
+                    let batch = match this.batches.get_mut(&item_key) {
+                        Some(batch) => batch,
+                        None => {
+                            let batch = (this.state)();
+                            this.batches.insert(item_key.clone(), batch);
+                            this.batches
+                                .get_mut(&item_key)
+                                .expect("batch has just been inserted so should exist")
                         }
-                    } else {
-                        // We have no batch yet for this partition key, so
-                        // create one and create the expiration entries as well.
-                        // This allows the batch to expire before filling up,
-                        // and vice versa.
-                        let batch = Batch::new(item_limit, alloc_limit).with(item);
-                        this.batches.insert(item_key.clone(), batch);
-                        this.timer.insert(item_key);
+                    };
+
+                    let (fits, metadata) = batch.item_fits_in_batch(&item);
+                    if !fits {
+                        // This batch is too full to accept a new item, so we move the contents of
+                        // the batch into `closed_batches` to be push out of this stream on the
+                        // next iteration.
+                        this.closed_batches
+                            .push((item_key.clone(), batch.take_batch()));
+                    }
+
+                    // Insert the item into the batch.
+                    batch.push(item, metadata);
+                    this.timer.insert(item_key.clone());
+                    if batch.is_batch_full() {
+                        // If the insertion means the batch is now full, we clear out the batch and
+                        // remove it from the list.
+                        this.closed_batches
+                            .push((item_key.clone(), batch.take_batch()));
+                        this.batches.remove(&item_key);
                     }
                 }
             }
@@ -464,7 +357,10 @@ mod test {
 
     use crate::{
         partition::Partitioner,
-        stream::partitioned_batcher::{ExpirationQueue, PartitionedBatcher},
+        stream::{
+            partitioned_batcher::{ExpirationQueue, PartitionedBatcher},
+            BatcherSettings,
+        },
         time::KeyedTimer,
     };
 
@@ -571,11 +467,12 @@ mod test {
             // key.
             let mut stream = stream::iter(stream.into_iter());
             let stream_size_hint = stream.size_hint();
-
             let item_limit = NonZeroUsize::new(item_limit as usize).unwrap();
             let allocation_limit = NonZeroUsize::new(allocation_limit as usize).unwrap();
+            let batch_settings = BatcherSettings::new(Duration::from_secs(1), allocation_limit, item_limit);
+
             let batcher = PartitionedBatcher::with_timer(&mut stream, partitioner, timer,
-                                              item_limit, Some(allocation_limit));
+                                              Box::new(move || batch_settings.into_byte_size_config()));
             let batcher_size_hint = batcher.size_hint();
 
             assert_eq!(stream_size_hint, batcher_size_hint);
@@ -597,9 +494,9 @@ mod test {
             let mut stream = stream::iter(stream.into_iter());
             let item_limit = NonZeroUsize::new(item_limit as usize).unwrap();
             let allocation_limit = NonZeroUsize::new(allocation_limit as usize).unwrap();
+            let batch_settings = BatcherSettings::new(Duration::from_secs(1), allocation_limit, item_limit);
             let mut batcher = PartitionedBatcher::with_timer(&mut stream, partitioner,
-                                                  timer, item_limit,
-                                                  Some(allocation_limit));
+                                                  timer, Box::new(move || batch_settings.into_byte_size_config()));
             let mut batcher = Pin::new(&mut batcher);
 
             loop {
@@ -668,9 +565,9 @@ mod test {
             let mut stream = stream::iter(stream.into_iter());
             let item_limit = NonZeroUsize::new(item_limit as usize).unwrap();
             let allocation_limit = NonZeroUsize::new(allocation_limit as usize).unwrap();
+            let batch_settings = BatcherSettings::new(Duration::from_secs(1), allocation_limit, item_limit);
             let mut batcher = PartitionedBatcher::with_timer(&mut stream, partitioner,
-                                                  timer, item_limit,
-                                                  Some(allocation_limit));
+                                                  timer, Box::new(move || batch_settings.clone().into_byte_size_config()));
             let mut batcher = Pin::new(&mut batcher);
 
             loop {
@@ -683,6 +580,10 @@ mod test {
                         let expected_partition = partitions
                             .get_mut(&key)
                             .expect("impossible situation");
+
+                        dbg!(&actual_batch);
+                        dbg!(&expected_partition);
+
                         for item in actual_batch {
                             assert_eq!(item, expected_partition.pop().unwrap());
                         }
@@ -695,6 +596,99 @@ mod test {
         }
     }
 
+    #[tokio::test]
+    async fn zorkwonk() {
+        // let stream = (0..20).collect::<Vec<_>>();
+        // let stream = vec![
+        //     7037641852729981347,
+        //     8644263311095074276,
+        //     4920834891500558744,
+        //     15150872024001141310,
+        //     12803461249171200340,
+        // ];
+        let stream = vec![
+            7037641852729981347,
+            8644263311095074276,
+            4920834891500558744,
+            15150872024001141310,
+            12803461249171200340,
+        ];
+        let item_limit = 1;
+        let allocation_limit = 8;
+        let partitioner = TestPartitioner {
+            key_space: NonZeroU8::new(23).unwrap(),
+        };
+        let mut count = 0;
+        let timer = TestTimer::new(
+            std::iter::from_fn(move || {
+                if count > 254 {
+                    None
+                } else {
+                    count += 1;
+                    if count % 5 == 0 {
+                        Some(Poll::Ready(Some(count)))
+                    } else {
+                        Some(Poll::Ready(None))
+                    }
+                }
+            })
+            .collect(),
+        );
+
+        // Asserts that for every received batch received the elements in
+        // the batch are not reordered within a batch. No claim is made on
+        // when batches themselves will issue, batch sizes etc.
+        let noop_waker = futures::task::noop_waker();
+        let mut cx = Context::from_waker(&noop_waker);
+
+        let mut partitions = separate_partitions(stream.clone(), &partitioner);
+        dbg!(&partitions);
+
+        let mut stream = stream::iter(stream.into_iter());
+        let item_limit = NonZeroUsize::new(item_limit as usize).unwrap();
+        let allocation_limit = NonZeroUsize::new(allocation_limit as usize).unwrap();
+        let batch_settings =
+            BatcherSettings::new(Duration::from_secs(1), allocation_limit, item_limit);
+
+        let mut batcher = PartitionedBatcher::with_timer(
+            &mut stream,
+            partitioner,
+            timer,
+            Box::new(move || batch_settings.clone().into_byte_size_config()),
+        );
+        // let mut batcher = PartitionedBatcher::new(
+        //     &mut stream,
+        //     partitioner,
+        //     // timer,
+        //     Box::new(move || batch_settings.clone().into_byte_size_config()),
+        // );
+        let mut batcher = Pin::new(&mut batcher);
+
+        loop {
+            match batcher.as_mut().poll_next(&mut cx) {
+                Poll::Pending => {}
+                Poll::Ready(None) => {
+                    break;
+                }
+                Poll::Ready(Some((key, actual_batch))) => {
+                    let expected_partition =
+                        partitions.get_mut(&key).expect("impossible situation");
+
+                    dbg!(&actual_batch);
+                    dbg!(&expected_partition);
+
+                    for item in actual_batch {
+                        assert_eq!(item, expected_partition.pop().unwrap());
+                    }
+                }
+            }
+        }
+        for v in partitions.values() {
+            assert!(v.is_empty());
+        }
+        // panic!();
+    }
+
     proptest! {
         #[test]
         fn batch_does_not_lose_items(stream: Vec<u64>,
@@ -711,9 +705,11 @@ mod test {
             let mut stream = stream::iter(stream.into_iter());
             let item_limit = NonZeroUsize::new(item_limit as usize).unwrap();
             let allocation_limit = NonZeroUsize::new(allocation_limit as usize).unwrap();
+            let batch_settings = BatcherSettings::new(Duration::from_secs(1), allocation_limit, item_limit);
             let mut batcher = PartitionedBatcher::with_timer(&mut stream, partitioner,
-                                                  timer, item_limit,
-                                                  Some(allocation_limit));
+                                                  timer, Box::new(move || batch_settings.clone().into_byte_size_config()));
+            // let mut batcher = PartitionedBatcher::new(&mut stream, partitioner,
+            //                                        Box::new(move || batch_settings.clone().into_byte_size_config()));
             let mut batcher = Pin::new(&mut batcher);
 
             let mut observed_items = 0;
diff --git a/src/sinks/aws_cloudwatch_logs/sink.rs b/src/sinks/aws_cloudwatch_logs/sink.rs
index 3c320ad6236e7..80672b3b55253 100644
--- a/src/sinks/aws_cloudwatch_logs/sink.rs
+++ b/src/sinks/aws_cloudwatch_logs/sink.rs
@@ -49,7 +49,10 @@ where
                 let age_range = start..end;
                 future::ready(age_range.contains(&req.timestamp))
             })
-            .batched_partitioned(CloudwatchPartitioner, batcher_settings)
+            .batched_partitioned(
+                CloudwatchPartitioner,
+                Box::new(move || batcher_settings.clone().into_byte_size_config()),
+            )
             .map(|(key, events)| {
                 let metadata = RequestMetadata::from_batch(
                     events.iter().map(|req| req.get_metadata().clone()),
diff --git a/src/sinks/aws_kinesis/sink.rs b/src/sinks/aws_kinesis/sink.rs
index 61a1e539418aa..eb94a159e19c2 100644
--- a/src/sinks/aws_kinesis/sink.rs
+++ b/src/sinks/aws_kinesis/sink.rs
@@ -42,6 +42,8 @@ where
     R: Record + Send + Sync + Unpin + Clone + 'static,
 {
     async fn run_inner(self: Box<Self>, input: BoxStream<'_, Event>) -> Result<(), ()> {
+        let batch_settings = self.batch_settings.clone();
+
         input
             .filter_map(|event| {
                 // Panic: This sink only accepts Logs, so this should never panic
@@ -67,7 +69,7 @@ where
                 KinesisPartitioner {
                     _phantom: PhantomData,
                 },
-                self.batch_settings,
+                Box::new(move || batch_settings.clone().into_byte_size_config()),
             )
             .map(|(key, events)| {
                 let metadata = RequestMetadata::from_batch(
diff --git a/src/sinks/azure_common/sink.rs b/src/sinks/azure_common/sink.rs
index d46910f1e4d98..a97b669c1134f 100644
--- a/src/sinks/azure_common/sink.rs
+++ b/src/sinks/azure_common/sink.rs
@@ -37,12 +37,15 @@ where
 {
     async fn run_inner(self: Box<Self>, input: BoxStream<'_, Event>) -> Result<(), ()> {
         let partitioner = self.partitioner;
-        let settings = self.batcher_settings;
+        let settings = self.batcher_settings.clone();
 
         let request_builder = self.request_builder;
 
         input
-            .batched_partitioned(partitioner, settings)
+            .batched_partitioned(
+                partitioner,
+                Box::new(move || settings.clone().into_byte_size_config()),
+            )
             .filter_map(|(key, batch)| async move {
                 // We don't need to emit an error here if the event is dropped since this will occur if the template
                 // couldn't be rendered during the partitioning. A `TemplateRenderingError` is already emitted when
diff --git a/src/sinks/clickhouse/sink.rs b/src/sinks/clickhouse/sink.rs
index c951101c8ce21..949b2748fc1fa 100644
--- a/src/sinks/clickhouse/sink.rs
+++ b/src/sinks/clickhouse/sink.rs
@@ -42,10 +42,12 @@ impl ClickhouseSink {
     }
 
     async fn run_inner(self: Box<Self>, input: BoxStream<'_, Event>) -> Result<(), ()> {
+        let batch_settings = self.batch_settings.clone();
+
         input
             .batched_partitioned(
                 KeyPartitioner::new(self.database, self.table),
-                self.batch_settings,
+                Box::new(move || batch_settings.clone().into_byte_size_config()),
             )
             .filter_map(|(key, batch)| async move { key.map(move |k| (k, batch)) })
             .request_builder(
diff --git a/src/sinks/datadog/logs/sink.rs b/src/sinks/datadog/logs/sink.rs
index 363fa55383d2b..6055866e3e85e 100644
--- a/src/sinks/datadog/logs/sink.rs
+++ b/src/sinks/datadog/logs/sink.rs
@@ -261,8 +261,12 @@ where
         let default_api_key = Arc::clone(&self.default_api_key);
 
         let partitioner = EventPartitioner;
+        let batch_settings = self.batch_settings.clone();
 
-        let input = input.batched_partitioned(partitioner, self.batch_settings);
+        let input = input.batched_partitioned(
+            partitioner,
+            Box::new(move || batch_settings.clone().into_byte_size_config()),
+        );
         input
             .request_builder(
                 default_request_builder_concurrency_limit(),
diff --git a/src/sinks/datadog/metrics/sink.rs b/src/sinks/datadog/metrics/sink.rs
index 5ceefc3c487d2..8c304e5891ecc 100644
--- a/src/sinks/datadog/metrics/sink.rs
+++ b/src/sinks/datadog/metrics/sink.rs
@@ -85,6 +85,7 @@ where
 
     async fn run_inner(self: Box<Self>, input: BoxStream<'_, Event>) -> Result<(), ()> {
         let mut splitter: MetricSplitter<AggregatedSummarySplitter> = MetricSplitter::default();
+        let batch_settings = self.batch_settings.clone();
 
         input
             // Convert `Event` to `Metric` so we don't have to deal with constant conversions.
@@ -99,7 +100,10 @@ where
             .normalized_with_default::<DatadogMetricsNormalizer>()
             // We batch metrics by their endpoint: series endpoint for counters, gauge, and sets vs sketch endpoint for
             // distributions, aggregated histograms, and sketches.
-            .batched_partitioned(DatadogMetricsTypePartitioner, self.batch_settings)
+            .batched_partitioned(
+                DatadogMetricsTypePartitioner,
+                Box::new(move || batch_settings.clone().into_byte_size_config()),
+            )
             // Aggregate counters with identical timestamps, otherwise identical counters (same
             // series and same timestamp, when rounded to whole seconds) will be dropped in a
             // last-write-wins situation when they hit the DD metrics intake.
diff --git a/src/sinks/datadog/traces/sink.rs b/src/sinks/datadog/traces/sink.rs
index 4cf502a7eebae..c28b053089297 100644
--- a/src/sinks/datadog/traces/sink.rs
+++ b/src/sinks/datadog/traces/sink.rs
@@ -108,8 +108,13 @@ where
     }
 
     async fn run_inner(self: Box<Self>, input: BoxStream<'_, Event>) -> Result<(), ()> {
+        let batch_settings = self.batch_settings.clone();
+
         input
-            .batched_partitioned(EventPartitioner, self.batch_settings)
+            .batched_partitioned(
+                EventPartitioner,
+                Box::new(move || batch_settings.clone().into_byte_size_config()),
+            )
             .incremental_request_builder(self.request_builder)
             .flat_map(stream::iter)
             .filter_map(|request| async move {
diff --git a/src/sinks/gcs_common/sink.rs b/src/sinks/gcs_common/sink.rs
index 830dbd4ca6301..7955d064052e0 100644
--- a/src/sinks/gcs_common/sink.rs
+++ b/src/sinks/gcs_common/sink.rs
@@ -45,7 +45,10 @@ where
         let request_builder = self.request_builder;
 
         input
-            .batched_partitioned(partitioner, settings)
+            .batched_partitioned(
+                partitioner,
+                Box::new(move || settings.clone().into_byte_size_config()),
+            )
             .filter_map(|(key, batch)| async move {
                 // A `TemplateRenderingError` will have been emitted by `KeyPartitioner` if the key here is `None`,
                 // thus no further `EventsDropped` event needs emitting at this stage.
diff --git a/src/sinks/loki/sink.rs b/src/sinks/loki/sink.rs
index 89ae9857ee746..dbaf4d5a29753 100644
--- a/src/sinks/loki/sink.rs
+++ b/src/sinks/loki/sink.rs
@@ -451,12 +451,16 @@ impl LokiSink {
                 NonZeroUsize::new(1).expect("static")
             }
         };
+        let batch_settings = self.batch_settings.clone();
 
         input
             .map(|event| encoder.encode_event(event))
             .filter_map(|event| async { event })
             .map(|record| filter.filter_record(record))
-            .batched_partitioned(RecordPartitioner, self.batch_settings)
+            .batched_partitioned(
+                RecordPartitioner,
+                Box::new(move || batch_settings.clone().into_byte_size_config()),
+            )
             .filter_map(|(partition, batch)| async {
                 if let Some(partition) = partition {
                     let mut count: usize = 0;
diff --git a/src/sinks/opendal_common.rs b/src/sinks/opendal_common.rs
index 3ff05f214f1c7..ba94b05e145d0 100644
--- a/src/sinks/opendal_common.rs
+++ b/src/sinks/opendal_common.rs
@@ -73,12 +73,15 @@ where
 {
     async fn run_inner(self: Box<Self>, input: BoxStream<'_, Event>) -> Result<(), ()> {
         let partitioner = self.partitioner;
-        let settings = self.batcher_settings;
+        let settings = self.batcher_settings.clone();
 
         let request_builder = self.request_builder;
 
         input
-            .batched_partitioned(partitioner, settings)
+            .batched_partitioned(
+                partitioner,
+                Box::new(move || settings.clone().into_byte_size_config()),
+            )
             .filter_map(|(key, batch)| async move {
                 // We don't need to emit an error here if the event is dropped since this will occur if the template
                 // couldn't be rendered during the partitioning. A `TemplateRenderingError` is already emitted when
diff --git a/src/sinks/prometheus/remote_write/sink.rs b/src/sinks/prometheus/remote_write/sink.rs
index 15d309d4cf6a9..15075f7737521 100644
--- a/src/sinks/prometheus/remote_write/sink.rs
+++ b/src/sinks/prometheus/remote_write/sink.rs
@@ -98,7 +98,23 @@ where
             .filter_map(move |event| {
                 future::ready(make_remote_write_event(tenant_id.as_ref(), event))
             })
-            .batched_partitioned(PrometheusTenantIdPartitioner, batch_settings)
+            .batched_partitioned(
+                PrometheusTenantIdPartitioner,
+                Box::new(move || batch_settings.clone().into_byte_size_config()),
+            )
+            // .batched(self.batch_settings.into_reducer_config(
+            //     |data: &Metric| data.size_of(),
+            //     |event_collection: &mut EventCollection, mut item: Metric| {
+            //         event_collection
+            //             .finalizers
+            //             .merge(item.metadata_mut().take_finalizers());
+            //         event_collection.events_byte_size += item.size_of();
+            //         event_collection
+            //             .events_json_byte_size
+            //             .add_event(&item, item.estimated_json_encoded_size_of());
+            //         event_collection.events.insert_update(item);
+            //     },
+            // ))
             .request_builder(default_request_builder_concurrency_limit(), request_builder)
             .filter_map(|request| async move {
                 match request {
diff --git a/src/sinks/prometheus/remote_write/tests.rs b/src/sinks/prometheus/remote_write/tests.rs
index da88851802a52..58d6430d38667 100644
--- a/src/sinks/prometheus/remote_write/tests.rs
+++ b/src/sinks/prometheus/remote_write/tests.rs
@@ -150,6 +150,8 @@ async fn retains_state_between_requests() {
     )
     .await;
 
+    dbg!(&outputs);
+
     assert_eq!(outputs.len(), 3);
 
     let check_output = |index: usize, name: &str, value: f64| {
diff --git a/src/sinks/s3_common/sink.rs b/src/sinks/s3_common/sink.rs
index c26f6f93e8a0f..57bf735fbb687 100644
--- a/src/sinks/s3_common/sink.rs
+++ b/src/sinks/s3_common/sink.rs
@@ -39,12 +39,15 @@ where
 {
     async fn run_inner(self: Box<Self>, input: BoxStream<'_, Event>) -> Result<(), ()> {
         let partitioner = self.partitioner;
-        let settings = self.batcher_settings;
+        let settings = self.batcher_settings.clone();
 
         let request_builder = self.request_builder;
 
         input
-            .batched_partitioned(partitioner, settings)
+            .batched_partitioned(
+                partitioner,
+                Box::new(move || settings.clone().into_byte_size_config()),
+            )
             .filter_map(|(key, batch)| async move { key.map(move |k| (k, batch)) })
             .request_builder(default_request_builder_concurrency_limit(), request_builder)
             .filter_map(|request| async move {
diff --git a/src/sinks/splunk_hec/logs/sink.rs b/src/sinks/splunk_hec/logs/sink.rs
index 555a3d9b95076..9ba19c5ffcad8 100644
--- a/src/sinks/splunk_hec/logs/sink.rs
+++ b/src/sinks/splunk_hec/logs/sink.rs
@@ -61,6 +61,7 @@ where
             timestamp_key: self.timestamp_key.clone(),
             endpoint_target: self.endpoint_target,
         };
+        let batch_settings = self.batch_settings.clone();
 
         input
             .map(move |event| process_log(event, &data))
@@ -77,7 +78,7 @@ where
                 } else {
                     EventPartitioner::new(None, None, None, None)
                 },
-                self.batch_settings,
+                Box::new(move || batch_settings.clone().into_byte_size_config()),
             )
             .request_builder(
                 default_request_builder_concurrency_limit(),
diff --git a/src/sinks/splunk_hec/metrics/sink.rs b/src/sinks/splunk_hec/metrics/sink.rs
index ef42d21b60370..f5e07c18d73fa 100644
--- a/src/sinks/splunk_hec/metrics/sink.rs
+++ b/src/sinks/splunk_hec/metrics/sink.rs
@@ -39,6 +39,7 @@ where
         let index = self.index.as_ref();
         let host_key = self.host_key.as_ref();
         let default_namespace = self.default_namespace.as_deref();
+        let batch_settings = self.batch_settings.clone();
 
         input
             .map(|event| (event.size_of(), event.into_metric()))
@@ -53,7 +54,10 @@ where
                     default_namespace,
                 ))
             })
-            .batched_partitioned(EventPartitioner, self.batch_settings)
+            .batched_partitioned(
+                EventPartitioner,
+                Box::new(move || batch_settings.clone().into_byte_size_config()),
+            )
             .request_builder(
                 default_request_builder_concurrency_limit(),
                 self.request_builder,
diff --git a/src/sinks/util/builder.rs b/src/sinks/util/builder.rs
index c51bf405dc4d1..19aeb3c74dc49 100644
--- a/src/sinks/util/builder.rs
+++ b/src/sinks/util/builder.rs
@@ -17,8 +17,7 @@ use vector_core::{
     partition::Partitioner,
     stream::{
         batcher::{config::BatchConfig, Batcher},
-        BatcherSettings, ConcurrentMap, Driver, DriverResponse, ExpirationQueue,
-        PartitionedBatcher,
+        ConcurrentMap, Driver, DriverResponse, ExpirationQueue, PartitionedBatcher,
     },
     ByteSizeOf,
 };
@@ -45,16 +44,17 @@ pub trait SinkBuilderExt: Stream {
     /// The stream will yield batches of events, with their partition key, when either a batch fills
     /// up or times out. [`Partitioner`] operates on a per-event basis, and has access to the event
     /// itself, and so can access any and all fields of an event.
-    fn batched_partitioned<P>(
+    fn batched_partitioned<P, C>(
         self,
         partitioner: P,
-        settings: BatcherSettings,
-    ) -> PartitionedBatcher<Self, P, ExpirationQueue<P::Key>>
+        settings: Box<dyn Fn() -> C + Send>,
+    ) -> PartitionedBatcher<Self, P, ExpirationQueue<P::Key>, C>
     where
         Self: Stream<Item = P::Item> + Sized,
         P: Partitioner + Unpin,
         P::Key: Eq + Hash + Clone,
         P::Item: ByteSizeOf,
+        C: BatchConfig<P::Item>,
     {
         PartitionedBatcher::new(self, partitioner, settings)
     }

From b439288d3bf7b5256e73bc1560b248d21a211f4e Mon Sep 17 00:00:00 2001
From: Stephen Wakely <fungus.humungus@gmail.com>
Date: Fri, 6 Oct 2023 12:04:53 +0100
Subject: [PATCH 10/32] Update PartitionBatcher to use BatchConfig

Signed-off-by: Stephen Wakely <fungus.humungus@gmail.com>
---
 .../src/stream/partitioned_batcher.rs         | 332 +++++++++---------
 src/sinks/aws_cloudwatch_logs/sink.rs         |   5 +-
 src/sinks/aws_kinesis/sink.rs                 |   4 +-
 src/sinks/azure_common/sink.rs                |   7 +-
 src/sinks/clickhouse/sink.rs                  |   4 +-
 src/sinks/datadog/logs/sink.rs                |   6 +-
 src/sinks/datadog/metrics/sink.rs             |   6 +-
 src/sinks/datadog/traces/sink.rs              |   7 +-
 src/sinks/gcs_common/sink.rs                  |   5 +-
 src/sinks/loki/sink.rs                        |   6 +-
 src/sinks/opendal_common.rs                   |   7 +-
 src/sinks/s3_common/sink.rs                   |   7 +-
 src/sinks/splunk_hec/logs/sink.rs             |   3 +-
 src/sinks/splunk_hec/metrics/sink.rs          |   6 +-
 src/sinks/util/builder.rs                     |  10 +-
 15 files changed, 226 insertions(+), 189 deletions(-)

diff --git a/lib/vector-core/src/stream/partitioned_batcher.rs b/lib/vector-core/src/stream/partitioned_batcher.rs
index aa2f7e6b58bcd..165c3ad3d35ff 100644
--- a/lib/vector-core/src/stream/partitioned_batcher.rs
+++ b/lib/vector-core/src/stream/partitioned_batcher.rs
@@ -1,8 +1,6 @@
 use std::{
-    cmp,
     collections::HashMap,
     hash::{BuildHasherDefault, Hash},
-    mem,
     num::NonZeroUsize,
     pin::Pin,
     task::{ready, Context, Poll},
@@ -25,6 +23,8 @@ use crate::{
     ByteSizeOf,
 };
 
+use super::batcher::BatchConfig;
+
 /// A `KeyedTimer` based on `DelayQueue`.
 pub struct ExpirationQueue<K> {
     /// The timeout to give each new key entry
@@ -100,114 +100,6 @@ where
     }
 }
 
-/// A batch for use by `Batcher`
-///
-/// This structure is a private implementation detail that simplifies the
-/// implementation of `Batcher`. It is the actual store of items that come
-/// through the stream manipulated by `Batcher` plus limit information to signal
-/// when the `Batch` is full.
-struct Batch<I> {
-    /// The total number of `I` bytes stored, does not any overhead in this
-    /// structure.
-    allocated_bytes: usize,
-    /// The maximum number of elements allowed in this structure.
-    element_limit: usize,
-    /// The maximum number of allocated bytes (not including overhead) allowed
-    /// in this structure.
-    allocation_limit: usize,
-    /// The store of `I` elements.
-    elements: Vec<I>,
-}
-
-impl<I> ByteSizeOf for Batch<I> {
-    fn allocated_bytes(&self) -> usize {
-        self.allocated_bytes
-    }
-}
-
-impl<I> Batch<I>
-where
-    I: ByteSizeOf,
-{
-    /// Create a new Batch instance
-    ///
-    /// Creates a new batch instance with specific element and allocation limits. The element limit
-    /// is a maximum cap on the number of `I` instances. The allocation limit is a soft-max on the
-    /// number of allocated bytes stored in this batch, not taking into account overhead from this
-    /// structure itself.
-    ///
-    /// If `allocation_limit` is smaller than the size of `I` as reported by `std::mem::size_of`,
-    /// then the allocation limit will be raised such that the batch can hold a single instance of
-    /// `I`.  Likewise, `element_limit` will be raised such that it is always at least 1, ensuring
-    /// that a new batch can be pushed into.
-    fn new(element_limit: usize, allocation_limit: usize) -> Self {
-        // SAFETY: `element_limit` is always non-zero because `BatcherSettings` can only be
-        // constructed with `NonZeroUsize` versions of allocation limit/item limit.  `Batch` is also
-        // only constructable via `Batcher`.
-
-        // TODO: This may need to be reworked, because it's subtly wrong as-is.
-        // ByteSizeOf::size_of() always returns the size of the type itself, plus any "allocated
-        // bytes".  Thus, there are times when an item will be bigger than simply the size of the
-        // type itself (aka mem::size_of::<I>()) and thus than type of item would never fit in a
-        // batch where the `allocation_limit` is at or lower than the size of that item.
-        //
-        // We're counteracting this here by ensuring that the element limit is always at least 1.
-        let allocation_limit = cmp::max(allocation_limit, mem::size_of::<I>());
-        Self {
-            allocated_bytes: 0,
-            element_limit,
-            allocation_limit,
-            elements: Vec::with_capacity(128),
-        }
-    }
-
-    /// Unconditionally insert an element into the batch
-    ///
-    /// This function is similar to `push` except that the caller does not need
-    /// to call `has_space` prior to calling this and it will never
-    /// panic. Intended to be used only when insertion must not fail.
-    fn with(mut self, value: I) -> Self {
-        self.allocated_bytes += value.size_of();
-        self.elements.push(value);
-        self
-    }
-
-    /// Decompose the batch
-    ///
-    /// Called by the user when they want to get at the internal store of
-    /// items. Returns a tuple, the first element being the allocated size of
-    /// stored items and the second the store of items.
-    fn into_inner(self) -> Vec<I> {
-        self.elements
-    }
-
-    /// Whether the batch has space for a new item
-    ///
-    /// This function returns true of there is space both in terms of item count
-    /// and byte count for the given item, false otherwise.
-    fn has_space(&self, value: &I) -> bool {
-        let too_many_elements = self.elements.len() + 1 > self.element_limit;
-        let too_many_bytes = self.allocated_bytes + value.size_of() > self.allocation_limit;
-        !(too_many_elements || too_many_bytes)
-    }
-
-    /// Push an element into the batch
-    ///
-    /// This function pushes an element into the batch. Callers must be sure to
-    /// call `has_space` prior to calling this function and receive a positive
-    /// result.
-    ///
-    /// # Panics
-    ///
-    /// This function will panic if there is not sufficient space in the batch
-    /// for a new element to be inserted.
-    fn push(&mut self, value: I) {
-        assert!(self.has_space(&value));
-        self.allocated_bytes += value.size_of();
-        self.elements.push(value);
-    }
-}
-
 /// Controls the behavior of the batcher in terms of batch size and flush interval.
 ///
 /// This is a temporary solution for pushing in a fixed settings structure so we don't have to worry
@@ -287,18 +179,16 @@ impl BatcherSettings {
 }
 
 #[pin_project]
-pub struct PartitionedBatcher<St, Prt, KT>
+pub struct PartitionedBatcher<St, Prt, KT, C>
 where
     Prt: Partitioner,
 {
-    /// The total number of bytes a single batch in this struct is allowed to
-    /// hold.
-    batch_allocation_limit: usize,
-    /// The maximum number of items that are allowed per-batch
-    batch_item_limit: usize,
+    /// A closure that retrievs a new `BatchConfig` when needed to batch a
+    /// new partition.
+    state: Box<dyn Fn() -> C + Send>,
     /// The store of live batches. Note that the key here is an option type,
     /// on account of the interface of `Prt`.
-    batches: HashMap<Prt::Key, Batch<Prt::Item>, BuildHasherDefault<XxHash64>>,
+    batches: HashMap<Prt::Key, C, BuildHasherDefault<XxHash64>>,
     /// The store of 'closed' batches. When this is not empty it will be
     /// preferentially flushed prior to consuming any new items from the
     /// underlying stream.
@@ -312,44 +202,44 @@ where
     stream: Fuse<St>,
 }
 
-impl<St, Prt> PartitionedBatcher<St, Prt, ExpirationQueue<Prt::Key>>
+impl<St, Prt, C> PartitionedBatcher<St, Prt, ExpirationQueue<Prt::Key>, C>
 where
     St: Stream<Item = Prt::Item>,
     Prt: Partitioner + Unpin,
     Prt::Key: Eq + Hash + Clone,
     Prt::Item: ByteSizeOf,
+    C: BatchConfig<Prt::Item>,
 {
-    pub fn new(stream: St, partitioner: Prt, settings: BatcherSettings) -> Self {
+    pub fn new(stream: St, partitioner: Prt, settings: Box<dyn Fn() -> C + Send>) -> Self {
+        let timeout = settings().timeout();
         Self {
-            batch_allocation_limit: settings.size_limit,
-            batch_item_limit: settings.item_limit,
+            state: settings,
             batches: HashMap::default(),
             closed_batches: Vec::default(),
-            timer: ExpirationQueue::new(settings.timeout),
+            timer: ExpirationQueue::new(timeout),
             partitioner,
             stream: stream.fuse(),
         }
     }
 }
 
-impl<St, Prt, KT> PartitionedBatcher<St, Prt, KT>
+#[cfg(test)]
+impl<St, Prt, KT, C> PartitionedBatcher<St, Prt, KT, C>
 where
     St: Stream<Item = Prt::Item>,
     Prt: Partitioner + Unpin,
     Prt::Key: Eq + Hash + Clone,
     Prt::Item: ByteSizeOf,
+    C: BatchConfig<Prt::Item>,
 {
     pub fn with_timer(
         stream: St,
         partitioner: Prt,
         timer: KT,
-        batch_item_limit: NonZeroUsize,
-        batch_allocation_limit: Option<NonZeroUsize>,
+        settings: Box<dyn Fn() -> C + Send>,
     ) -> Self {
         Self {
-            batch_allocation_limit: batch_allocation_limit
-                .map_or(usize::max_value(), NonZeroUsize::get),
-            batch_item_limit: batch_item_limit.get(),
+            state: settings,
             batches: HashMap::default(),
             closed_batches: Vec::default(),
             timer,
@@ -359,13 +249,14 @@ where
     }
 }
 
-impl<St, Prt, KT> Stream for PartitionedBatcher<St, Prt, KT>
+impl<St, Prt, KT, C> Stream for PartitionedBatcher<St, Prt, KT, C>
 where
     St: Stream<Item = Prt::Item>,
     Prt: Partitioner + Unpin,
     Prt::Key: Eq + Hash + Clone,
     Prt::Item: ByteSizeOf,
     KT: KeyedTimer<Prt::Key>,
+    C: BatchConfig<Prt::Item, Batch = Vec<Prt::Item>>,
 {
     type Item = (Prt::Key, Vec<Prt::Item>);
 
@@ -385,13 +276,11 @@ where
                     // here but still be usable later if more entries are added.
                     Poll::Pending | Poll::Ready(None) => return Poll::Pending,
                     Poll::Ready(Some(item_key)) => {
-                        let batch = this
+                        let mut batch = this
                             .batches
                             .remove(&item_key)
                             .expect("batch should exist if it is set to expire");
-                        this.closed_batches.push((item_key, batch.into_inner()));
-
-                        continue;
+                        this.closed_batches.push((item_key, batch.take_batch()));
                     }
                 },
                 Poll::Ready(None) => {
@@ -405,7 +294,7 @@ where
                         this.closed_batches.extend(
                             this.batches
                                 .drain()
-                                .map(|(key, batch)| (key, batch.into_inner())),
+                                .map(|(key, mut batch)| (key, batch.take_batch())),
                         );
                         continue;
                     }
@@ -413,33 +302,37 @@ where
                 }
                 Poll::Ready(Some(item)) => {
                     let item_key = this.partitioner.partition(&item);
-                    let item_limit: usize = *this.batch_item_limit;
-                    let alloc_limit: usize = *this.batch_allocation_limit;
-
-                    if let Some(batch) = this.batches.get_mut(&item_key) {
-                        if batch.has_space(&item) {
-                            // When there's space in the partition batch just
-                            // push the item in and loop back around.
-                            batch.push(item);
-                        } else {
-                            let new_batch = Batch::new(item_limit, alloc_limit).with(item);
-                            let batch = mem::replace(batch, new_batch);
 
-                            // The batch for this partition key was set to
-                            // expire, but now it's overflowed and must be
-                            // pushed out, so now we reset the batch timeout.
-                            this.timer.insert(item_key.clone());
-
-                            this.closed_batches.push((item_key, batch.into_inner()));
+                    // Get the batch for this partition, or create a new one.
+                    let batch = match this.batches.get_mut(&item_key) {
+                        Some(batch) => batch,
+                        None => {
+                            let batch = (this.state)();
+                            this.batches.insert(item_key.clone(), batch);
+                            this.batches
+                                .get_mut(&item_key)
+                                .expect("batch has just been inserted so should exist")
                         }
-                    } else {
-                        // We have no batch yet for this partition key, so
-                        // create one and create the expiration entries as well.
-                        // This allows the batch to expire before filling up,
-                        // and vice versa.
-                        let batch = Batch::new(item_limit, alloc_limit).with(item);
-                        this.batches.insert(item_key.clone(), batch);
-                        this.timer.insert(item_key);
+                    };
+
+                    let (fits, metadata) = batch.item_fits_in_batch(&item);
+                    if !fits {
+                        // This batch is too full to accept a new item, so we move the contents of
+                        // the batch into `closed_batches` to be push out of this stream on the
+                        // next iteration.
+                        this.closed_batches
+                            .push((item_key.clone(), batch.take_batch()));
+                    }
+
+                    // Insert the item into the batch.
+                    batch.push(item, metadata);
+                    this.timer.insert(item_key.clone());
+                    if batch.is_batch_full() {
+                        // If the insertion means the batch is now full, we clear out the batch and
+                        // remove it from the list.
+                        this.closed_batches
+                            .push((item_key.clone(), batch.take_batch()));
+                        this.batches.remove(&item_key);
                     }
                 }
             }
@@ -464,7 +357,10 @@ mod test {
 
     use crate::{
         partition::Partitioner,
-        stream::partitioned_batcher::{ExpirationQueue, PartitionedBatcher},
+        stream::{
+            partitioned_batcher::{ExpirationQueue, PartitionedBatcher},
+            BatcherSettings,
+        },
         time::KeyedTimer,
     };
 
@@ -571,11 +467,12 @@ mod test {
             // key.
             let mut stream = stream::iter(stream.into_iter());
             let stream_size_hint = stream.size_hint();
-
             let item_limit = NonZeroUsize::new(item_limit as usize).unwrap();
             let allocation_limit = NonZeroUsize::new(allocation_limit as usize).unwrap();
+            let batch_settings = BatcherSettings::new(Duration::from_secs(1), allocation_limit, item_limit);
+
             let batcher = PartitionedBatcher::with_timer(&mut stream, partitioner, timer,
-                                              item_limit, Some(allocation_limit));
+                                              Box::new(move || batch_settings.into_byte_size_config()));
             let batcher_size_hint = batcher.size_hint();
 
             assert_eq!(stream_size_hint, batcher_size_hint);
@@ -597,9 +494,9 @@ mod test {
             let mut stream = stream::iter(stream.into_iter());
             let item_limit = NonZeroUsize::new(item_limit as usize).unwrap();
             let allocation_limit = NonZeroUsize::new(allocation_limit as usize).unwrap();
+            let batch_settings = BatcherSettings::new(Duration::from_secs(1), allocation_limit, item_limit);
             let mut batcher = PartitionedBatcher::with_timer(&mut stream, partitioner,
-                                                  timer, item_limit,
-                                                  Some(allocation_limit));
+                                                  timer, Box::new(move || batch_settings.into_byte_size_config()));
             let mut batcher = Pin::new(&mut batcher);
 
             loop {
@@ -668,9 +565,9 @@ mod test {
             let mut stream = stream::iter(stream.into_iter());
             let item_limit = NonZeroUsize::new(item_limit as usize).unwrap();
             let allocation_limit = NonZeroUsize::new(allocation_limit as usize).unwrap();
+            let batch_settings = BatcherSettings::new(Duration::from_secs(1), allocation_limit, item_limit);
             let mut batcher = PartitionedBatcher::with_timer(&mut stream, partitioner,
-                                                  timer, item_limit,
-                                                  Some(allocation_limit));
+                                                  timer, Box::new(move || batch_settings.clone().into_byte_size_config()));
             let mut batcher = Pin::new(&mut batcher);
 
             loop {
@@ -683,6 +580,10 @@ mod test {
                         let expected_partition = partitions
                             .get_mut(&key)
                             .expect("impossible situation");
+
+                        dbg!(&actual_batch);
+                        dbg!(&expected_partition);
+
                         for item in actual_batch {
                             assert_eq!(item, expected_partition.pop().unwrap());
                         }
@@ -695,6 +596,99 @@ mod test {
         }
     }
 
+    #[tokio::test]
+    async fn zorkwonk() {
+        // let stream = (0..20).collect::<Vec<_>>();
+        // let stream = vec![
+        //     7037641852729981347,
+        //     8644263311095074276,
+        //     4920834891500558744,
+        //     15150872024001141310,
+        //     12803461249171200340,
+        // ];
+        let stream = vec![
+            7037641852729981347,
+            8644263311095074276,
+            4920834891500558744,
+            15150872024001141310,
+            12803461249171200340,
+        ];
+        let item_limit = 1;
+        let allocation_limit = 8;
+        let partitioner = TestPartitioner {
+            key_space: NonZeroU8::new(23).unwrap(),
+        };
+        let mut count = 0;
+        let timer = TestTimer::new(
+            std::iter::from_fn(move || {
+                if count > 254 {
+                    None
+                } else {
+                    count += 1;
+                    if count % 5 == 0 {
+                        Some(Poll::Ready(Some(count)))
+                    } else {
+                        Some(Poll::Ready(None))
+                    }
+                }
+            })
+            .collect(),
+        );
+
+        // Asserts that for every received batch received the elements in
+        // the batch are not reordered within a batch. No claim is made on
+        // when batches themselves will issue, batch sizes etc.
+        let noop_waker = futures::task::noop_waker();
+        let mut cx = Context::from_waker(&noop_waker);
+
+        let mut partitions = separate_partitions(stream.clone(), &partitioner);
+        dbg!(&partitions);
+
+        let mut stream = stream::iter(stream.into_iter());
+        let item_limit = NonZeroUsize::new(item_limit as usize).unwrap();
+        let allocation_limit = NonZeroUsize::new(allocation_limit as usize).unwrap();
+        let batch_settings =
+            BatcherSettings::new(Duration::from_secs(1), allocation_limit, item_limit);
+
+        let mut batcher = PartitionedBatcher::with_timer(
+            &mut stream,
+            partitioner,
+            timer,
+            Box::new(move || batch_settings.clone().into_byte_size_config()),
+        );
+        // let mut batcher = PartitionedBatcher::new(
+        //     &mut stream,
+        //     partitioner,
+        //     // timer,
+        //     Box::new(move || batch_settings.clone().into_byte_size_config()),
+        // );
+        let mut batcher = Pin::new(&mut batcher);
+
+        loop {
+            match batcher.as_mut().poll_next(&mut cx) {
+                Poll::Pending => {}
+                Poll::Ready(None) => {
+                    break;
+                }
+                Poll::Ready(Some((key, actual_batch))) => {
+                    let expected_partition =
+                        partitions.get_mut(&key).expect("impossible situation");
+
+                    dbg!(&actual_batch);
+                    dbg!(&expected_partition);
+
+                    for item in actual_batch {
+                        assert_eq!(item, expected_partition.pop().unwrap());
+                    }
+                }
+            }
+        }
+        for v in partitions.values() {
+            assert!(v.is_empty());
+        }
+        // panic!();
+    }
+
     proptest! {
         #[test]
         fn batch_does_not_lose_items(stream: Vec<u64>,
@@ -711,9 +705,11 @@ mod test {
             let mut stream = stream::iter(stream.into_iter());
             let item_limit = NonZeroUsize::new(item_limit as usize).unwrap();
             let allocation_limit = NonZeroUsize::new(allocation_limit as usize).unwrap();
+            let batch_settings = BatcherSettings::new(Duration::from_secs(1), allocation_limit, item_limit);
             let mut batcher = PartitionedBatcher::with_timer(&mut stream, partitioner,
-                                                  timer, item_limit,
-                                                  Some(allocation_limit));
+                                                  timer, Box::new(move || batch_settings.clone().into_byte_size_config()));
+            // let mut batcher = PartitionedBatcher::new(&mut stream, partitioner,
+            //                                        Box::new(move || batch_settings.clone().into_byte_size_config()));
             let mut batcher = Pin::new(&mut batcher);
 
             let mut observed_items = 0;
diff --git a/src/sinks/aws_cloudwatch_logs/sink.rs b/src/sinks/aws_cloudwatch_logs/sink.rs
index 3c320ad6236e7..80672b3b55253 100644
--- a/src/sinks/aws_cloudwatch_logs/sink.rs
+++ b/src/sinks/aws_cloudwatch_logs/sink.rs
@@ -49,7 +49,10 @@ where
                 let age_range = start..end;
                 future::ready(age_range.contains(&req.timestamp))
             })
-            .batched_partitioned(CloudwatchPartitioner, batcher_settings)
+            .batched_partitioned(
+                CloudwatchPartitioner,
+                Box::new(move || batcher_settings.clone().into_byte_size_config()),
+            )
             .map(|(key, events)| {
                 let metadata = RequestMetadata::from_batch(
                     events.iter().map(|req| req.get_metadata().clone()),
diff --git a/src/sinks/aws_kinesis/sink.rs b/src/sinks/aws_kinesis/sink.rs
index 61a1e539418aa..eb94a159e19c2 100644
--- a/src/sinks/aws_kinesis/sink.rs
+++ b/src/sinks/aws_kinesis/sink.rs
@@ -42,6 +42,8 @@ where
     R: Record + Send + Sync + Unpin + Clone + 'static,
 {
     async fn run_inner(self: Box<Self>, input: BoxStream<'_, Event>) -> Result<(), ()> {
+        let batch_settings = self.batch_settings.clone();
+
         input
             .filter_map(|event| {
                 // Panic: This sink only accepts Logs, so this should never panic
@@ -67,7 +69,7 @@ where
                 KinesisPartitioner {
                     _phantom: PhantomData,
                 },
-                self.batch_settings,
+                Box::new(move || batch_settings.clone().into_byte_size_config()),
             )
             .map(|(key, events)| {
                 let metadata = RequestMetadata::from_batch(
diff --git a/src/sinks/azure_common/sink.rs b/src/sinks/azure_common/sink.rs
index d46910f1e4d98..a97b669c1134f 100644
--- a/src/sinks/azure_common/sink.rs
+++ b/src/sinks/azure_common/sink.rs
@@ -37,12 +37,15 @@ where
 {
     async fn run_inner(self: Box<Self>, input: BoxStream<'_, Event>) -> Result<(), ()> {
         let partitioner = self.partitioner;
-        let settings = self.batcher_settings;
+        let settings = self.batcher_settings.clone();
 
         let request_builder = self.request_builder;
 
         input
-            .batched_partitioned(partitioner, settings)
+            .batched_partitioned(
+                partitioner,
+                Box::new(move || settings.clone().into_byte_size_config()),
+            )
             .filter_map(|(key, batch)| async move {
                 // We don't need to emit an error here if the event is dropped since this will occur if the template
                 // couldn't be rendered during the partitioning. A `TemplateRenderingError` is already emitted when
diff --git a/src/sinks/clickhouse/sink.rs b/src/sinks/clickhouse/sink.rs
index c951101c8ce21..949b2748fc1fa 100644
--- a/src/sinks/clickhouse/sink.rs
+++ b/src/sinks/clickhouse/sink.rs
@@ -42,10 +42,12 @@ impl ClickhouseSink {
     }
 
     async fn run_inner(self: Box<Self>, input: BoxStream<'_, Event>) -> Result<(), ()> {
+        let batch_settings = self.batch_settings.clone();
+
         input
             .batched_partitioned(
                 KeyPartitioner::new(self.database, self.table),
-                self.batch_settings,
+                Box::new(move || batch_settings.clone().into_byte_size_config()),
             )
             .filter_map(|(key, batch)| async move { key.map(move |k| (k, batch)) })
             .request_builder(
diff --git a/src/sinks/datadog/logs/sink.rs b/src/sinks/datadog/logs/sink.rs
index 363fa55383d2b..6055866e3e85e 100644
--- a/src/sinks/datadog/logs/sink.rs
+++ b/src/sinks/datadog/logs/sink.rs
@@ -261,8 +261,12 @@ where
         let default_api_key = Arc::clone(&self.default_api_key);
 
         let partitioner = EventPartitioner;
+        let batch_settings = self.batch_settings.clone();
 
-        let input = input.batched_partitioned(partitioner, self.batch_settings);
+        let input = input.batched_partitioned(
+            partitioner,
+            Box::new(move || batch_settings.clone().into_byte_size_config()),
+        );
         input
             .request_builder(
                 default_request_builder_concurrency_limit(),
diff --git a/src/sinks/datadog/metrics/sink.rs b/src/sinks/datadog/metrics/sink.rs
index fccfe040fdd95..e7da4f2036015 100644
--- a/src/sinks/datadog/metrics/sink.rs
+++ b/src/sinks/datadog/metrics/sink.rs
@@ -84,6 +84,7 @@ where
 
     async fn run_inner(self: Box<Self>, input: BoxStream<'_, Event>) -> Result<(), ()> {
         let mut splitter: MetricSplitter<AggregatedSummarySplitter> = MetricSplitter::default();
+        let batch_settings = self.batch_settings.clone();
 
         input
             // Convert `Event` to `Metric` so we don't have to deal with constant conversions.
@@ -98,7 +99,10 @@ where
             .normalized_with_default::<DatadogMetricsNormalizer>()
             // We batch metrics by their endpoint: series endpoint for counters, gauge, and sets vs sketch endpoint for
             // distributions, aggregated histograms, and sketches.
-            .batched_partitioned(DatadogMetricsTypePartitioner, self.batch_settings)
+            .batched_partitioned(
+                DatadogMetricsTypePartitioner,
+                Box::new(move || batch_settings.clone().into_byte_size_config()),
+            )
             // Aggregate counters with identical timestamps, otherwise identical counters (same
             // series and same timestamp, when rounded to whole seconds) will be dropped in a
             // last-write-wins situation when they hit the DD metrics intake.
diff --git a/src/sinks/datadog/traces/sink.rs b/src/sinks/datadog/traces/sink.rs
index 4cf502a7eebae..c28b053089297 100644
--- a/src/sinks/datadog/traces/sink.rs
+++ b/src/sinks/datadog/traces/sink.rs
@@ -108,8 +108,13 @@ where
     }
 
     async fn run_inner(self: Box<Self>, input: BoxStream<'_, Event>) -> Result<(), ()> {
+        let batch_settings = self.batch_settings.clone();
+
         input
-            .batched_partitioned(EventPartitioner, self.batch_settings)
+            .batched_partitioned(
+                EventPartitioner,
+                Box::new(move || batch_settings.clone().into_byte_size_config()),
+            )
             .incremental_request_builder(self.request_builder)
             .flat_map(stream::iter)
             .filter_map(|request| async move {
diff --git a/src/sinks/gcs_common/sink.rs b/src/sinks/gcs_common/sink.rs
index 830dbd4ca6301..7955d064052e0 100644
--- a/src/sinks/gcs_common/sink.rs
+++ b/src/sinks/gcs_common/sink.rs
@@ -45,7 +45,10 @@ where
         let request_builder = self.request_builder;
 
         input
-            .batched_partitioned(partitioner, settings)
+            .batched_partitioned(
+                partitioner,
+                Box::new(move || settings.clone().into_byte_size_config()),
+            )
             .filter_map(|(key, batch)| async move {
                 // A `TemplateRenderingError` will have been emitted by `KeyPartitioner` if the key here is `None`,
                 // thus no further `EventsDropped` event needs emitting at this stage.
diff --git a/src/sinks/loki/sink.rs b/src/sinks/loki/sink.rs
index 9d9046788d04b..4619be1370dd6 100644
--- a/src/sinks/loki/sink.rs
+++ b/src/sinks/loki/sink.rs
@@ -453,12 +453,16 @@ impl LokiSink {
                 NonZeroUsize::new(1).expect("static")
             }
         };
+        let batch_settings = self.batch_settings.clone();
 
         input
             .map(|event| encoder.encode_event(event))
             .filter_map(|event| async { event })
             .map(|record| filter.filter_record(record))
-            .batched_partitioned(RecordPartitioner, self.batch_settings)
+            .batched_partitioned(
+                RecordPartitioner,
+                Box::new(move || batch_settings.clone().into_byte_size_config()),
+            )
             .filter_map(|(partition, batch)| async {
                 if let Some(partition) = partition {
                     let mut count: usize = 0;
diff --git a/src/sinks/opendal_common.rs b/src/sinks/opendal_common.rs
index 3ff05f214f1c7..ba94b05e145d0 100644
--- a/src/sinks/opendal_common.rs
+++ b/src/sinks/opendal_common.rs
@@ -73,12 +73,15 @@ where
 {
     async fn run_inner(self: Box<Self>, input: BoxStream<'_, Event>) -> Result<(), ()> {
         let partitioner = self.partitioner;
-        let settings = self.batcher_settings;
+        let settings = self.batcher_settings.clone();
 
         let request_builder = self.request_builder;
 
         input
-            .batched_partitioned(partitioner, settings)
+            .batched_partitioned(
+                partitioner,
+                Box::new(move || settings.clone().into_byte_size_config()),
+            )
             .filter_map(|(key, batch)| async move {
                 // We don't need to emit an error here if the event is dropped since this will occur if the template
                 // couldn't be rendered during the partitioning. A `TemplateRenderingError` is already emitted when
diff --git a/src/sinks/s3_common/sink.rs b/src/sinks/s3_common/sink.rs
index c26f6f93e8a0f..57bf735fbb687 100644
--- a/src/sinks/s3_common/sink.rs
+++ b/src/sinks/s3_common/sink.rs
@@ -39,12 +39,15 @@ where
 {
     async fn run_inner(self: Box<Self>, input: BoxStream<'_, Event>) -> Result<(), ()> {
         let partitioner = self.partitioner;
-        let settings = self.batcher_settings;
+        let settings = self.batcher_settings.clone();
 
         let request_builder = self.request_builder;
 
         input
-            .batched_partitioned(partitioner, settings)
+            .batched_partitioned(
+                partitioner,
+                Box::new(move || settings.clone().into_byte_size_config()),
+            )
             .filter_map(|(key, batch)| async move { key.map(move |k| (k, batch)) })
             .request_builder(default_request_builder_concurrency_limit(), request_builder)
             .filter_map(|request| async move {
diff --git a/src/sinks/splunk_hec/logs/sink.rs b/src/sinks/splunk_hec/logs/sink.rs
index 555a3d9b95076..9ba19c5ffcad8 100644
--- a/src/sinks/splunk_hec/logs/sink.rs
+++ b/src/sinks/splunk_hec/logs/sink.rs
@@ -61,6 +61,7 @@ where
             timestamp_key: self.timestamp_key.clone(),
             endpoint_target: self.endpoint_target,
         };
+        let batch_settings = self.batch_settings.clone();
 
         input
             .map(move |event| process_log(event, &data))
@@ -77,7 +78,7 @@ where
                 } else {
                     EventPartitioner::new(None, None, None, None)
                 },
-                self.batch_settings,
+                Box::new(move || batch_settings.clone().into_byte_size_config()),
             )
             .request_builder(
                 default_request_builder_concurrency_limit(),
diff --git a/src/sinks/splunk_hec/metrics/sink.rs b/src/sinks/splunk_hec/metrics/sink.rs
index ef42d21b60370..f5e07c18d73fa 100644
--- a/src/sinks/splunk_hec/metrics/sink.rs
+++ b/src/sinks/splunk_hec/metrics/sink.rs
@@ -39,6 +39,7 @@ where
         let index = self.index.as_ref();
         let host_key = self.host_key.as_ref();
         let default_namespace = self.default_namespace.as_deref();
+        let batch_settings = self.batch_settings.clone();
 
         input
             .map(|event| (event.size_of(), event.into_metric()))
@@ -53,7 +54,10 @@ where
                     default_namespace,
                 ))
             })
-            .batched_partitioned(EventPartitioner, self.batch_settings)
+            .batched_partitioned(
+                EventPartitioner,
+                Box::new(move || batch_settings.clone().into_byte_size_config()),
+            )
             .request_builder(
                 default_request_builder_concurrency_limit(),
                 self.request_builder,
diff --git a/src/sinks/util/builder.rs b/src/sinks/util/builder.rs
index 617697ea84ec1..30025bbdf9bf1 100644
--- a/src/sinks/util/builder.rs
+++ b/src/sinks/util/builder.rs
@@ -17,8 +17,7 @@ use vector_core::{
     partition::Partitioner,
     stream::{
         batcher::{config::BatchConfig, Batcher},
-        BatcherSettings, ConcurrentMap, Driver, DriverResponse, ExpirationQueue,
-        PartitionedBatcher,
+        ConcurrentMap, Driver, DriverResponse, ExpirationQueue, PartitionedBatcher,
     },
     ByteSizeOf,
 };
@@ -45,16 +44,17 @@ pub trait SinkBuilderExt: Stream {
     /// The stream will yield batches of events, with their partition key, when either a batch fills
     /// up or times out. [`Partitioner`] operates on a per-event basis, and has access to the event
     /// itself, and so can access any and all fields of an event.
-    fn batched_partitioned<P>(
+    fn batched_partitioned<P, C>(
         self,
         partitioner: P,
-        settings: BatcherSettings,
-    ) -> PartitionedBatcher<Self, P, ExpirationQueue<P::Key>>
+        settings: Box<dyn Fn() -> C + Send>,
+    ) -> PartitionedBatcher<Self, P, ExpirationQueue<P::Key>, C>
     where
         Self: Stream<Item = P::Item> + Sized,
         P: Partitioner + Unpin,
         P::Key: Eq + Hash + Clone,
         P::Item: ByteSizeOf,
+        C: BatchConfig<P::Item>,
     {
         PartitionedBatcher::new(self, partitioner, settings)
     }

From cc81a6838eec43b15faa357f06a546abfea09512 Mon Sep 17 00:00:00 2001
From: Stephen Wakely <fungus.humungus@gmail.com>
Date: Fri, 6 Oct 2023 12:27:46 +0100
Subject: [PATCH 11/32] Remove zorkwonk

Signed-off-by: Stephen Wakely <fungus.humungus@gmail.com>
---
 .../src/stream/partitioned_batcher.rs         | 93 -------------------
 1 file changed, 93 deletions(-)

diff --git a/lib/vector-core/src/stream/partitioned_batcher.rs b/lib/vector-core/src/stream/partitioned_batcher.rs
index 165c3ad3d35ff..400d74aa19188 100644
--- a/lib/vector-core/src/stream/partitioned_batcher.rs
+++ b/lib/vector-core/src/stream/partitioned_batcher.rs
@@ -596,99 +596,6 @@ mod test {
         }
     }
 
-    #[tokio::test]
-    async fn zorkwonk() {
-        // let stream = (0..20).collect::<Vec<_>>();
-        // let stream = vec![
-        //     7037641852729981347,
-        //     8644263311095074276,
-        //     4920834891500558744,
-        //     15150872024001141310,
-        //     12803461249171200340,
-        // ];
-        let stream = vec![
-            7037641852729981347,
-            8644263311095074276,
-            4920834891500558744,
-            15150872024001141310,
-            12803461249171200340,
-        ];
-        let item_limit = 1;
-        let allocation_limit = 8;
-        let partitioner = TestPartitioner {
-            key_space: NonZeroU8::new(23).unwrap(),
-        };
-        let mut count = 0;
-        let timer = TestTimer::new(
-            std::iter::from_fn(move || {
-                if count > 254 {
-                    None
-                } else {
-                    count += 1;
-                    if count % 5 == 0 {
-                        Some(Poll::Ready(Some(count)))
-                    } else {
-                        Some(Poll::Ready(None))
-                    }
-                }
-            })
-            .collect(),
-        );
-
-        // Asserts that for every received batch received the elements in
-        // the batch are not reordered within a batch. No claim is made on
-        // when batches themselves will issue, batch sizes etc.
-        let noop_waker = futures::task::noop_waker();
-        let mut cx = Context::from_waker(&noop_waker);
-
-        let mut partitions = separate_partitions(stream.clone(), &partitioner);
-        dbg!(&partitions);
-
-        let mut stream = stream::iter(stream.into_iter());
-        let item_limit = NonZeroUsize::new(item_limit as usize).unwrap();
-        let allocation_limit = NonZeroUsize::new(allocation_limit as usize).unwrap();
-        let batch_settings =
-            BatcherSettings::new(Duration::from_secs(1), allocation_limit, item_limit);
-
-        let mut batcher = PartitionedBatcher::with_timer(
-            &mut stream,
-            partitioner,
-            timer,
-            Box::new(move || batch_settings.clone().into_byte_size_config()),
-        );
-        // let mut batcher = PartitionedBatcher::new(
-        //     &mut stream,
-        //     partitioner,
-        //     // timer,
-        //     Box::new(move || batch_settings.clone().into_byte_size_config()),
-        // );
-        let mut batcher = Pin::new(&mut batcher);
-
-        loop {
-            match batcher.as_mut().poll_next(&mut cx) {
-                Poll::Pending => {}
-                Poll::Ready(None) => {
-                    break;
-                }
-                Poll::Ready(Some((key, actual_batch))) => {
-                    let expected_partition =
-                        partitions.get_mut(&key).expect("impossible situation");
-
-                    dbg!(&actual_batch);
-                    dbg!(&expected_partition);
-
-                    for item in actual_batch {
-                        assert_eq!(item, expected_partition.pop().unwrap());
-                    }
-                }
-            }
-        }
-        for v in partitions.values() {
-            assert!(v.is_empty());
-        }
-        // panic!();
-    }
-
     proptest! {
         #[test]
         fn batch_does_not_lose_items(stream: Vec<u64>,

From 830272510eb37815f6479a68eac08cf78114e79e Mon Sep 17 00:00:00 2001
From: Stephen Wakely <fungus.humungus@gmail.com>
Date: Fri, 6 Oct 2023 12:41:07 +0100
Subject: [PATCH 12/32] Make into fns as fns instead

Signed-off-by: Stephen Wakely <fungus.humungus@gmail.com>
---
 lib/vector-core/src/stream/batcher/mod.rs      |  9 +++------
 .../src/stream/partitioned_batcher.rs          | 18 ++++++++----------
 src/sinks/appsignal/sink.rs                    |  2 +-
 src/sinks/aws_cloudwatch_logs/sink.rs          |  2 +-
 src/sinks/aws_kinesis/sink.rs                  |  2 +-
 src/sinks/azure_common/sink.rs                 |  2 +-
 src/sinks/azure_monitor_logs/sink.rs           |  2 +-
 src/sinks/clickhouse/sink.rs                   |  2 +-
 src/sinks/databend/sink.rs                     |  2 +-
 src/sinks/datadog/logs/sink.rs                 |  2 +-
 src/sinks/datadog/metrics/sink.rs              |  2 +-
 src/sinks/datadog/traces/sink.rs               |  2 +-
 src/sinks/elasticsearch/sink.rs                |  2 +-
 src/sinks/gcp/stackdriver/logs/sink.rs         |  5 +----
 src/sinks/gcs_common/sink.rs                   |  2 +-
 src/sinks/greptimedb/sink.rs                   |  2 +-
 src/sinks/honeycomb/sink.rs                    |  5 +----
 src/sinks/http/sink.rs                         |  2 +-
 src/sinks/loki/sink.rs                         |  2 +-
 src/sinks/new_relic/sink.rs                    |  2 +-
 src/sinks/opendal_common.rs                    |  2 +-
 src/sinks/redis/sink.rs                        |  2 +-
 src/sinks/s3_common/sink.rs                    |  2 +-
 src/sinks/splunk_hec/logs/sink.rs              |  2 +-
 src/sinks/splunk_hec/metrics/sink.rs           |  2 +-
 src/sinks/statsd/sink.rs                       |  2 +-
 26 files changed, 35 insertions(+), 46 deletions(-)

diff --git a/lib/vector-core/src/stream/batcher/mod.rs b/lib/vector-core/src/stream/batcher/mod.rs
index a7153a5a8665d..bfb20e95e4160 100644
--- a/lib/vector-core/src/stream/batcher/mod.rs
+++ b/lib/vector-core/src/stream/batcher/mod.rs
@@ -129,10 +129,7 @@ mod test {
             NonZeroUsize::new(10000).unwrap(),
             NonZeroUsize::new(2).unwrap(),
         );
-        let batcher = Batcher::new(
-            stream,
-            settings.into_item_size_config(|x: &u32| *x as usize),
-        );
+        let batcher = Batcher::new(stream, settings.as_item_size_config(|x: &u32| *x as usize));
         let batches: Vec<_> = batcher.collect().await;
         assert_eq!(batches, vec![vec![1, 2], vec![3],]);
     }
@@ -146,7 +143,7 @@ mod test {
                 NonZeroUsize::new(5).unwrap(),
                 NonZeroUsize::new(100).unwrap(),
             )
-            .into_item_size_config(|x: &u32| *x as usize),
+            .as_item_size_config(|x: &u32| *x as usize),
         );
         let batches: Vec<_> = batcher.collect().await;
         assert_eq!(
@@ -176,7 +173,7 @@ mod test {
                 NonZeroUsize::new(5).unwrap(),
                 NonZeroUsize::new(100).unwrap(),
             )
-            .into_item_size_config(|x: &u32| *x as usize),
+            .as_item_size_config(|x: &u32| *x as usize),
         );
 
         tokio::pin!(batcher);
diff --git a/lib/vector-core/src/stream/partitioned_batcher.rs b/lib/vector-core/src/stream/partitioned_batcher.rs
index 400d74aa19188..8573e86f707e0 100644
--- a/lib/vector-core/src/stream/partitioned_batcher.rs
+++ b/lib/vector-core/src/stream/partitioned_batcher.rs
@@ -129,15 +129,15 @@ impl BatcherSettings {
 
     /// A batcher config using the `ByteSizeOf` trait to determine batch sizes.
     /// The output is a  `Vec<T>`.
-    pub fn into_byte_size_config<T: ByteSizeOf>(
-        self,
+    pub fn as_byte_size_config<T: ByteSizeOf>(
+        &self,
     ) -> BatchConfigParts<SizeLimit<ByteSizeOfItemSize>, Vec<T>> {
-        self.into_item_size_config(ByteSizeOfItemSize)
+        self.as_item_size_config(ByteSizeOfItemSize)
     }
 
     /// A batcher config using the `ItemBatchSize` trait to determine batch sizes.
     /// The output is a `Vec<T>`.
-    pub fn into_item_size_config<T, I>(self, item_size: I) -> BatchConfigParts<SizeLimit<I>, Vec<T>>
+    pub fn as_item_size_config<T, I>(&self, item_size: I) -> BatchConfigParts<SizeLimit<I>, Vec<T>>
     where
         I: ItemBatchSize<T>,
     {
@@ -472,7 +472,7 @@ mod test {
             let batch_settings = BatcherSettings::new(Duration::from_secs(1), allocation_limit, item_limit);
 
             let batcher = PartitionedBatcher::with_timer(&mut stream, partitioner, timer,
-                                              Box::new(move || batch_settings.into_byte_size_config()));
+                                              Box::new(move || batch_settings.as_byte_size_config()));
             let batcher_size_hint = batcher.size_hint();
 
             assert_eq!(stream_size_hint, batcher_size_hint);
@@ -496,7 +496,7 @@ mod test {
             let allocation_limit = NonZeroUsize::new(allocation_limit as usize).unwrap();
             let batch_settings = BatcherSettings::new(Duration::from_secs(1), allocation_limit, item_limit);
             let mut batcher = PartitionedBatcher::with_timer(&mut stream, partitioner,
-                                                  timer, Box::new(move || batch_settings.into_byte_size_config()));
+                                                  timer, Box::new(move || batch_settings.as_byte_size_config()));
             let mut batcher = Pin::new(&mut batcher);
 
             loop {
@@ -567,7 +567,7 @@ mod test {
             let allocation_limit = NonZeroUsize::new(allocation_limit as usize).unwrap();
             let batch_settings = BatcherSettings::new(Duration::from_secs(1), allocation_limit, item_limit);
             let mut batcher = PartitionedBatcher::with_timer(&mut stream, partitioner,
-                                                  timer, Box::new(move || batch_settings.clone().into_byte_size_config()));
+                                                  timer, Box::new(move || batch_settings.as_byte_size_config()));
             let mut batcher = Pin::new(&mut batcher);
 
             loop {
@@ -614,9 +614,7 @@ mod test {
             let allocation_limit = NonZeroUsize::new(allocation_limit as usize).unwrap();
             let batch_settings = BatcherSettings::new(Duration::from_secs(1), allocation_limit, item_limit);
             let mut batcher = PartitionedBatcher::with_timer(&mut stream, partitioner,
-                                                  timer, Box::new(move || batch_settings.clone().into_byte_size_config()));
-            // let mut batcher = PartitionedBatcher::new(&mut stream, partitioner,
-            //                                        Box::new(move || batch_settings.clone().into_byte_size_config()));
+                                                  timer, Box::new(move || batch_settings.as_byte_size_config()));
             let mut batcher = Pin::new(&mut batcher);
 
             let mut observed_items = 0;
diff --git a/src/sinks/appsignal/sink.rs b/src/sinks/appsignal/sink.rs
index 1ad0e3e0c8779..e066b6dd94a30 100644
--- a/src/sinks/appsignal/sink.rs
+++ b/src/sinks/appsignal/sink.rs
@@ -34,7 +34,7 @@ where
                     Some(event)
                 })
             })
-            .batched(self.batch_settings.into_byte_size_config())
+            .batched(self.batch_settings.as_byte_size_config())
             .request_builder(
                 default_request_builder_concurrency_limit(),
                 AppsignalRequestBuilder {
diff --git a/src/sinks/aws_cloudwatch_logs/sink.rs b/src/sinks/aws_cloudwatch_logs/sink.rs
index 80672b3b55253..b23d752a58977 100644
--- a/src/sinks/aws_cloudwatch_logs/sink.rs
+++ b/src/sinks/aws_cloudwatch_logs/sink.rs
@@ -51,7 +51,7 @@ where
             })
             .batched_partitioned(
                 CloudwatchPartitioner,
-                Box::new(move || batcher_settings.clone().into_byte_size_config()),
+                Box::new(move || batcher_settings.as_byte_size_config()),
             )
             .map(|(key, events)| {
                 let metadata = RequestMetadata::from_batch(
diff --git a/src/sinks/aws_kinesis/sink.rs b/src/sinks/aws_kinesis/sink.rs
index eb94a159e19c2..7e72234ab1743 100644
--- a/src/sinks/aws_kinesis/sink.rs
+++ b/src/sinks/aws_kinesis/sink.rs
@@ -69,7 +69,7 @@ where
                 KinesisPartitioner {
                     _phantom: PhantomData,
                 },
-                Box::new(move || batch_settings.clone().into_byte_size_config()),
+                Box::new(move || batch_settings.as_byte_size_config()),
             )
             .map(|(key, events)| {
                 let metadata = RequestMetadata::from_batch(
diff --git a/src/sinks/azure_common/sink.rs b/src/sinks/azure_common/sink.rs
index a97b669c1134f..4f8c8c11c7fc2 100644
--- a/src/sinks/azure_common/sink.rs
+++ b/src/sinks/azure_common/sink.rs
@@ -44,7 +44,7 @@ where
         input
             .batched_partitioned(
                 partitioner,
-                Box::new(move || settings.clone().into_byte_size_config()),
+                Box::new(move || settings.as_byte_size_config()),
             )
             .filter_map(|(key, batch)| async move {
                 // We don't need to emit an error here if the event is dropped since this will occur if the template
diff --git a/src/sinks/azure_monitor_logs/sink.rs b/src/sinks/azure_monitor_logs/sink.rs
index ce06e04bb83b8..20db422f01420 100644
--- a/src/sinks/azure_monitor_logs/sink.rs
+++ b/src/sinks/azure_monitor_logs/sink.rs
@@ -39,7 +39,7 @@ where
 
     async fn run_inner(self: Box<Self>, input: BoxStream<'_, Event>) -> Result<(), ()> {
         input
-            .batched(self.batch_settings.into_byte_size_config())
+            .batched(self.batch_settings.as_byte_size_config())
             .request_builder(
                 default_request_builder_concurrency_limit(),
                 AzureMonitorLogsRequestBuilder {
diff --git a/src/sinks/clickhouse/sink.rs b/src/sinks/clickhouse/sink.rs
index 949b2748fc1fa..62af596bb1f07 100644
--- a/src/sinks/clickhouse/sink.rs
+++ b/src/sinks/clickhouse/sink.rs
@@ -47,7 +47,7 @@ impl ClickhouseSink {
         input
             .batched_partitioned(
                 KeyPartitioner::new(self.database, self.table),
-                Box::new(move || batch_settings.clone().into_byte_size_config()),
+                Box::new(move || batch_settings.as_byte_size_config()),
             )
             .filter_map(|(key, batch)| async move { key.map(move |k| (k, batch)) })
             .request_builder(
diff --git a/src/sinks/databend/sink.rs b/src/sinks/databend/sink.rs
index 87b3d1155cc9d..dcb625ce1ffe6 100644
--- a/src/sinks/databend/sink.rs
+++ b/src/sinks/databend/sink.rs
@@ -24,7 +24,7 @@ impl DatabendSink {
 
     async fn run_inner(self: Box<Self>, input: BoxStream<'_, Event>) -> Result<(), ()> {
         input
-            .batched(self.batch_settings.into_byte_size_config())
+            .batched(self.batch_settings.as_byte_size_config())
             .request_builder(
                 default_request_builder_concurrency_limit(),
                 self.request_builder,
diff --git a/src/sinks/datadog/logs/sink.rs b/src/sinks/datadog/logs/sink.rs
index 6055866e3e85e..98afac3f4ca1f 100644
--- a/src/sinks/datadog/logs/sink.rs
+++ b/src/sinks/datadog/logs/sink.rs
@@ -265,7 +265,7 @@ where
 
         let input = input.batched_partitioned(
             partitioner,
-            Box::new(move || batch_settings.clone().into_byte_size_config()),
+            Box::new(move || batch_settings.as_byte_size_config()),
         );
         input
             .request_builder(
diff --git a/src/sinks/datadog/metrics/sink.rs b/src/sinks/datadog/metrics/sink.rs
index e7da4f2036015..7788c0a55bbfa 100644
--- a/src/sinks/datadog/metrics/sink.rs
+++ b/src/sinks/datadog/metrics/sink.rs
@@ -101,7 +101,7 @@ where
             // distributions, aggregated histograms, and sketches.
             .batched_partitioned(
                 DatadogMetricsTypePartitioner,
-                Box::new(move || batch_settings.clone().into_byte_size_config()),
+                Box::new(move || batch_settings.as_byte_size_config()),
             )
             // Aggregate counters with identical timestamps, otherwise identical counters (same
             // series and same timestamp, when rounded to whole seconds) will be dropped in a
diff --git a/src/sinks/datadog/traces/sink.rs b/src/sinks/datadog/traces/sink.rs
index c28b053089297..52e7ed237c056 100644
--- a/src/sinks/datadog/traces/sink.rs
+++ b/src/sinks/datadog/traces/sink.rs
@@ -113,7 +113,7 @@ where
         input
             .batched_partitioned(
                 EventPartitioner,
-                Box::new(move || batch_settings.clone().into_byte_size_config()),
+                Box::new(move || batch_settings.as_byte_size_config()),
             )
             .incremental_request_builder(self.request_builder)
             .flat_map(stream::iter)
diff --git a/src/sinks/elasticsearch/sink.rs b/src/sinks/elasticsearch/sink.rs
index a3168d5b44809..5d3c043e233ce 100644
--- a/src/sinks/elasticsearch/sink.rs
+++ b/src/sinks/elasticsearch/sink.rs
@@ -81,7 +81,7 @@ where
             .filter_map(move |log| {
                 future::ready(process_log(log, &mode, id_key_field, &transformer))
             })
-            .batched(self.batch_settings.into_byte_size_config())
+            .batched(self.batch_settings.as_byte_size_config())
             .request_builder(
                 default_request_builder_concurrency_limit(),
                 self.request_builder,
diff --git a/src/sinks/gcp/stackdriver/logs/sink.rs b/src/sinks/gcp/stackdriver/logs/sink.rs
index 39870bca24432..fbf3e2066d420 100644
--- a/src/sinks/gcp/stackdriver/logs/sink.rs
+++ b/src/sinks/gcp/stackdriver/logs/sink.rs
@@ -36,10 +36,7 @@ where
     async fn run_inner(self: Box<Self>, input: BoxStream<'_, Event>) -> Result<(), ()> {
         input
             // Batch the input stream with size calculation based on the estimated encoded json size
-            .batched(
-                self.batch_settings
-                    .into_item_size_config(HttpJsonBatchSizer),
-            )
+            .batched(self.batch_settings.as_item_size_config(HttpJsonBatchSizer))
             // Build requests with no concurrency limit.
             .request_builder(
                 default_request_builder_concurrency_limit(),
diff --git a/src/sinks/gcs_common/sink.rs b/src/sinks/gcs_common/sink.rs
index 7955d064052e0..4e6326fd8dc2e 100644
--- a/src/sinks/gcs_common/sink.rs
+++ b/src/sinks/gcs_common/sink.rs
@@ -47,7 +47,7 @@ where
         input
             .batched_partitioned(
                 partitioner,
-                Box::new(move || settings.clone().into_byte_size_config()),
+                Box::new(move || settings.as_byte_size_config()),
             )
             .filter_map(|(key, batch)| async move {
                 // A `TemplateRenderingError` will have been emitted by `KeyPartitioner` if the key here is `None`,
diff --git a/src/sinks/greptimedb/sink.rs b/src/sinks/greptimedb/sink.rs
index fc8e75140708d..b4b36a10bf5db 100644
--- a/src/sinks/greptimedb/sink.rs
+++ b/src/sinks/greptimedb/sink.rs
@@ -37,7 +37,7 @@ impl GreptimeDBSink {
             .normalized_with_default::<GreptimeDBMetricNormalize>()
             .batched(
                 self.batch_settings
-                    .into_item_size_config(GreptimeDBBatchSizer),
+                    .as_item_size_config(GreptimeDBBatchSizer),
             )
             .map(GreptimeDBRequest::from_metrics)
             .into_driver(self.service)
diff --git a/src/sinks/honeycomb/sink.rs b/src/sinks/honeycomb/sink.rs
index 235dcd6eb4893..fe2bf4cdfc598 100644
--- a/src/sinks/honeycomb/sink.rs
+++ b/src/sinks/honeycomb/sink.rs
@@ -36,10 +36,7 @@ where
     async fn run_inner(self: Box<Self>, input: BoxStream<'_, Event>) -> Result<(), ()> {
         input
             // Batch the input stream with size calculation based on the estimated encoded json size
-            .batched(
-                self.batch_settings
-                    .into_item_size_config(HttpJsonBatchSizer),
-            )
+            .batched(self.batch_settings.as_item_size_config(HttpJsonBatchSizer))
             // Build requests with default concurrency limit.
             .request_builder(
                 default_request_builder_concurrency_limit(),
diff --git a/src/sinks/http/sink.rs b/src/sinks/http/sink.rs
index 8427201b31662..575e558e2d013 100644
--- a/src/sinks/http/sink.rs
+++ b/src/sinks/http/sink.rs
@@ -33,7 +33,7 @@ where
     async fn run_inner(self: Box<Self>, input: BoxStream<'_, Event>) -> Result<(), ()> {
         input
             // Batch the input stream with size calculation based on the configured codec
-            .batched(self.batch_settings.into_item_size_config(HttpBatchSizer {
+            .batched(self.batch_settings.as_item_size_config(HttpBatchSizer {
                 encoder: self.request_builder.encoder.encoder.clone(),
             }))
             // Build requests with default concurrency limit.
diff --git a/src/sinks/loki/sink.rs b/src/sinks/loki/sink.rs
index 4619be1370dd6..8563e2b45286a 100644
--- a/src/sinks/loki/sink.rs
+++ b/src/sinks/loki/sink.rs
@@ -461,7 +461,7 @@ impl LokiSink {
             .map(|record| filter.filter_record(record))
             .batched_partitioned(
                 RecordPartitioner,
-                Box::new(move || batch_settings.clone().into_byte_size_config()),
+                Box::new(move || batch_settings.as_byte_size_config()),
             )
             .filter_map(|(partition, batch)| async {
                 if let Some(partition) = partition {
diff --git a/src/sinks/new_relic/sink.rs b/src/sinks/new_relic/sink.rs
index fd93c90bbd61b..e8976e9131026 100644
--- a/src/sinks/new_relic/sink.rs
+++ b/src/sinks/new_relic/sink.rs
@@ -123,7 +123,7 @@ where
         let protocol = get_http_scheme_from_uri(&self.credentials.get_uri());
 
         input
-            .batched(self.batcher_settings.into_byte_size_config())
+            .batched(self.batcher_settings.as_byte_size_config())
             .request_builder(default_request_builder_concurrency_limit(), request_builder)
             .filter_map(
                 |request: Result<NewRelicApiRequest, NewRelicSinkError>| async move {
diff --git a/src/sinks/opendal_common.rs b/src/sinks/opendal_common.rs
index ba94b05e145d0..e63d570ae08d8 100644
--- a/src/sinks/opendal_common.rs
+++ b/src/sinks/opendal_common.rs
@@ -80,7 +80,7 @@ where
         input
             .batched_partitioned(
                 partitioner,
-                Box::new(move || settings.clone().into_byte_size_config()),
+                Box::new(move || settings.as_byte_size_config()),
             )
             .filter_map(|(key, batch)| async move {
                 // We don't need to emit an error here if the event is dropped since this will occur if the template
diff --git a/src/sinks/redis/sink.rs b/src/sinks/redis/sink.rs
index a9dcad68682ff..128ce05adddef 100644
--- a/src/sinks/redis/sink.rs
+++ b/src/sinks/redis/sink.rs
@@ -86,7 +86,7 @@ impl RedisSink {
 
         let mut encoder = self.encoder.clone();
         let transformer = self.transformer.clone();
-        let batcher_settings = self.batcher_settings.into_byte_size_config();
+        let batcher_settings = self.batcher_settings.as_byte_size_config();
 
         input
             .filter_map(|event| future::ready(self.make_redis_event(event)))
diff --git a/src/sinks/s3_common/sink.rs b/src/sinks/s3_common/sink.rs
index 57bf735fbb687..7378854a17dd3 100644
--- a/src/sinks/s3_common/sink.rs
+++ b/src/sinks/s3_common/sink.rs
@@ -46,7 +46,7 @@ where
         input
             .batched_partitioned(
                 partitioner,
-                Box::new(move || settings.clone().into_byte_size_config()),
+                Box::new(move || settings.as_byte_size_config()),
             )
             .filter_map(|(key, batch)| async move { key.map(move |k| (k, batch)) })
             .request_builder(default_request_builder_concurrency_limit(), request_builder)
diff --git a/src/sinks/splunk_hec/logs/sink.rs b/src/sinks/splunk_hec/logs/sink.rs
index 9ba19c5ffcad8..b4e488ac7c1da 100644
--- a/src/sinks/splunk_hec/logs/sink.rs
+++ b/src/sinks/splunk_hec/logs/sink.rs
@@ -78,7 +78,7 @@ where
                 } else {
                     EventPartitioner::new(None, None, None, None)
                 },
-                Box::new(move || batch_settings.clone().into_byte_size_config()),
+                Box::new(move || batch_settings.as_byte_size_config()),
             )
             .request_builder(
                 default_request_builder_concurrency_limit(),
diff --git a/src/sinks/splunk_hec/metrics/sink.rs b/src/sinks/splunk_hec/metrics/sink.rs
index f5e07c18d73fa..dbc847ec816b9 100644
--- a/src/sinks/splunk_hec/metrics/sink.rs
+++ b/src/sinks/splunk_hec/metrics/sink.rs
@@ -56,7 +56,7 @@ where
             })
             .batched_partitioned(
                 EventPartitioner,
-                Box::new(move || batch_settings.clone().into_byte_size_config()),
+                Box::new(move || batch_settings.as_byte_size_config()),
             )
             .request_builder(
                 default_request_builder_concurrency_limit(),
diff --git a/src/sinks/statsd/sink.rs b/src/sinks/statsd/sink.rs
index a9f0bdf44a9e2..f5c77e91b4f86 100644
--- a/src/sinks/statsd/sink.rs
+++ b/src/sinks/statsd/sink.rs
@@ -58,7 +58,7 @@ where
             // other metric types in type-specific ways i.e. incremental gauge updates use a
             // different syntax, etc.
             .normalized_with_default::<StatsdNormalizer>()
-            .batched(self.batch_settings.into_item_size_config(StatsdBatchSizer))
+            .batched(self.batch_settings.as_item_size_config(StatsdBatchSizer))
             // We build our requests "incrementally", which means that for a single batch of
             // metrics, we might generate N requests to represent all of the metrics in the batch.
             //

From e6b052d9efeb08e7e3e7bdb6de05b7981954ba2e Mon Sep 17 00:00:00 2001
From: Stephen Wakely <fungus.humungus@gmail.com>
Date: Fri, 6 Oct 2023 12:43:25 +0100
Subject: [PATCH 13/32] Spelling

Signed-off-by: Stephen Wakely <fungus.humungus@gmail.com>
---
 lib/vector-core/src/stream/partitioned_batcher.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lib/vector-core/src/stream/partitioned_batcher.rs b/lib/vector-core/src/stream/partitioned_batcher.rs
index 8573e86f707e0..1497ff9b53aa0 100644
--- a/lib/vector-core/src/stream/partitioned_batcher.rs
+++ b/lib/vector-core/src/stream/partitioned_batcher.rs
@@ -183,7 +183,7 @@ pub struct PartitionedBatcher<St, Prt, KT, C>
 where
     Prt: Partitioner,
 {
-    /// A closure that retrievs a new `BatchConfig` when needed to batch a
+    /// A closure that retrieves a new [`BatchConfig`] when needed to batch a
     /// new partition.
     state: Box<dyn Fn() -> C + Send>,
     /// The store of live batches. Note that the key here is an option type,

From 179f35173393cd7b40e7d9b3e57af27b2753d895 Mon Sep 17 00:00:00 2001
From: Stephen Wakely <fungus.humungus@gmail.com>
Date: Fri, 6 Oct 2023 13:09:02 +0100
Subject: [PATCH 14/32] Don't box the closure

Signed-off-by: Stephen Wakely <fungus.humungus@gmail.com>
---
 .../src/stream/partitioned_batcher.rs         | 22 +++++++++----------
 src/sinks/aws_cloudwatch_logs/sink.rs         |  7 +++---
 src/sinks/aws_kinesis/sink.rs                 |  2 +-
 src/sinks/azure_common/sink.rs                |  5 +----
 src/sinks/clickhouse/sink.rs                  |  7 +++---
 src/sinks/datadog/logs/sink.rs                |  5 +----
 src/sinks/datadog/metrics/sink.rs             |  7 +++---
 src/sinks/datadog/traces/sink.rs              |  5 +----
 src/sinks/gcs_common/sink.rs                  |  5 +----
 src/sinks/loki/sink.rs                        |  5 +----
 src/sinks/opendal_common.rs                   |  5 +----
 src/sinks/s3_common/sink.rs                   |  5 +----
 src/sinks/splunk_hec/logs/sink.rs             |  2 +-
 src/sinks/splunk_hec/metrics/sink.rs          |  5 +----
 src/sinks/util/builder.rs                     |  7 +++---
 15 files changed, 33 insertions(+), 61 deletions(-)

diff --git a/lib/vector-core/src/stream/partitioned_batcher.rs b/lib/vector-core/src/stream/partitioned_batcher.rs
index 1497ff9b53aa0..68bb989ed371b 100644
--- a/lib/vector-core/src/stream/partitioned_batcher.rs
+++ b/lib/vector-core/src/stream/partitioned_batcher.rs
@@ -179,13 +179,13 @@ impl BatcherSettings {
 }
 
 #[pin_project]
-pub struct PartitionedBatcher<St, Prt, KT, C>
+pub struct PartitionedBatcher<St, Prt, KT, C, F>
 where
     Prt: Partitioner,
 {
     /// A closure that retrieves a new [`BatchConfig`] when needed to batch a
     /// new partition.
-    state: Box<dyn Fn() -> C + Send>,
+    state: F,
     /// The store of live batches. Note that the key here is an option type,
     /// on account of the interface of `Prt`.
     batches: HashMap<Prt::Key, C, BuildHasherDefault<XxHash64>>,
@@ -202,15 +202,16 @@ where
     stream: Fuse<St>,
 }
 
-impl<St, Prt, C> PartitionedBatcher<St, Prt, ExpirationQueue<Prt::Key>, C>
+impl<St, Prt, C, F> PartitionedBatcher<St, Prt, ExpirationQueue<Prt::Key>, C, F>
 where
     St: Stream<Item = Prt::Item>,
     Prt: Partitioner + Unpin,
     Prt::Key: Eq + Hash + Clone,
     Prt::Item: ByteSizeOf,
     C: BatchConfig<Prt::Item>,
+    F: Fn() -> C + Send,
 {
-    pub fn new(stream: St, partitioner: Prt, settings: Box<dyn Fn() -> C + Send>) -> Self {
+    pub fn new(stream: St, partitioner: Prt, settings: F) -> Self {
         let timeout = settings().timeout();
         Self {
             state: settings,
@@ -224,20 +225,16 @@ where
 }
 
 #[cfg(test)]
-impl<St, Prt, KT, C> PartitionedBatcher<St, Prt, KT, C>
+impl<St, Prt, KT, C, F> PartitionedBatcher<St, Prt, KT, C, F>
 where
     St: Stream<Item = Prt::Item>,
     Prt: Partitioner + Unpin,
     Prt::Key: Eq + Hash + Clone,
     Prt::Item: ByteSizeOf,
     C: BatchConfig<Prt::Item>,
+    F: Fn() -> C + Send,
 {
-    pub fn with_timer(
-        stream: St,
-        partitioner: Prt,
-        timer: KT,
-        settings: Box<dyn Fn() -> C + Send>,
-    ) -> Self {
+    pub fn with_timer(stream: St, partitioner: Prt, timer: KT, settings: F) -> Self {
         Self {
             state: settings,
             batches: HashMap::default(),
@@ -249,7 +246,7 @@ where
     }
 }
 
-impl<St, Prt, KT, C> Stream for PartitionedBatcher<St, Prt, KT, C>
+impl<St, Prt, KT, C, F> Stream for PartitionedBatcher<St, Prt, KT, C, F>
 where
     St: Stream<Item = Prt::Item>,
     Prt: Partitioner + Unpin,
@@ -257,6 +254,7 @@ where
     Prt::Item: ByteSizeOf,
     KT: KeyedTimer<Prt::Key>,
     C: BatchConfig<Prt::Item, Batch = Vec<Prt::Item>>,
+    F: Fn() -> C + Send,
 {
     type Item = (Prt::Key, Vec<Prt::Item>);
 
diff --git a/src/sinks/aws_cloudwatch_logs/sink.rs b/src/sinks/aws_cloudwatch_logs/sink.rs
index b23d752a58977..5d71f8d0b605e 100644
--- a/src/sinks/aws_cloudwatch_logs/sink.rs
+++ b/src/sinks/aws_cloudwatch_logs/sink.rs
@@ -49,10 +49,9 @@ where
                 let age_range = start..end;
                 future::ready(age_range.contains(&req.timestamp))
             })
-            .batched_partitioned(
-                CloudwatchPartitioner,
-                Box::new(move || batcher_settings.as_byte_size_config()),
-            )
+            .batched_partitioned(CloudwatchPartitioner, || {
+                batcher_settings.as_byte_size_config()
+            })
             .map(|(key, events)| {
                 let metadata = RequestMetadata::from_batch(
                     events.iter().map(|req| req.get_metadata().clone()),
diff --git a/src/sinks/aws_kinesis/sink.rs b/src/sinks/aws_kinesis/sink.rs
index 7e72234ab1743..60e056d2811ad 100644
--- a/src/sinks/aws_kinesis/sink.rs
+++ b/src/sinks/aws_kinesis/sink.rs
@@ -69,7 +69,7 @@ where
                 KinesisPartitioner {
                     _phantom: PhantomData,
                 },
-                Box::new(move || batch_settings.as_byte_size_config()),
+                || batch_settings.as_byte_size_config(),
             )
             .map(|(key, events)| {
                 let metadata = RequestMetadata::from_batch(
diff --git a/src/sinks/azure_common/sink.rs b/src/sinks/azure_common/sink.rs
index 4f8c8c11c7fc2..a954217a15f33 100644
--- a/src/sinks/azure_common/sink.rs
+++ b/src/sinks/azure_common/sink.rs
@@ -42,10 +42,7 @@ where
         let request_builder = self.request_builder;
 
         input
-            .batched_partitioned(
-                partitioner,
-                Box::new(move || settings.as_byte_size_config()),
-            )
+            .batched_partitioned(partitioner, || settings.as_byte_size_config())
             .filter_map(|(key, batch)| async move {
                 // We don't need to emit an error here if the event is dropped since this will occur if the template
                 // couldn't be rendered during the partitioning. A `TemplateRenderingError` is already emitted when
diff --git a/src/sinks/clickhouse/sink.rs b/src/sinks/clickhouse/sink.rs
index 62af596bb1f07..987bc5c1e34e3 100644
--- a/src/sinks/clickhouse/sink.rs
+++ b/src/sinks/clickhouse/sink.rs
@@ -45,10 +45,9 @@ impl ClickhouseSink {
         let batch_settings = self.batch_settings.clone();
 
         input
-            .batched_partitioned(
-                KeyPartitioner::new(self.database, self.table),
-                Box::new(move || batch_settings.as_byte_size_config()),
-            )
+            .batched_partitioned(KeyPartitioner::new(self.database, self.table), || {
+                batch_settings.as_byte_size_config()
+            })
             .filter_map(|(key, batch)| async move { key.map(move |k| (k, batch)) })
             .request_builder(
                 default_request_builder_concurrency_limit(),
diff --git a/src/sinks/datadog/logs/sink.rs b/src/sinks/datadog/logs/sink.rs
index 98afac3f4ca1f..2abad6286bd47 100644
--- a/src/sinks/datadog/logs/sink.rs
+++ b/src/sinks/datadog/logs/sink.rs
@@ -263,10 +263,7 @@ where
         let partitioner = EventPartitioner;
         let batch_settings = self.batch_settings.clone();
 
-        let input = input.batched_partitioned(
-            partitioner,
-            Box::new(move || batch_settings.as_byte_size_config()),
-        );
+        let input = input.batched_partitioned(partitioner, || batch_settings.as_byte_size_config());
         input
             .request_builder(
                 default_request_builder_concurrency_limit(),
diff --git a/src/sinks/datadog/metrics/sink.rs b/src/sinks/datadog/metrics/sink.rs
index 7788c0a55bbfa..23ea4e2a77227 100644
--- a/src/sinks/datadog/metrics/sink.rs
+++ b/src/sinks/datadog/metrics/sink.rs
@@ -99,10 +99,9 @@ where
             .normalized_with_default::<DatadogMetricsNormalizer>()
             // We batch metrics by their endpoint: series endpoint for counters, gauge, and sets vs sketch endpoint for
             // distributions, aggregated histograms, and sketches.
-            .batched_partitioned(
-                DatadogMetricsTypePartitioner,
-                Box::new(move || batch_settings.as_byte_size_config()),
-            )
+            .batched_partitioned(DatadogMetricsTypePartitioner, || {
+                batch_settings.as_byte_size_config()
+            })
             // Aggregate counters with identical timestamps, otherwise identical counters (same
             // series and same timestamp, when rounded to whole seconds) will be dropped in a
             // last-write-wins situation when they hit the DD metrics intake.
diff --git a/src/sinks/datadog/traces/sink.rs b/src/sinks/datadog/traces/sink.rs
index 52e7ed237c056..dd09b57a64d86 100644
--- a/src/sinks/datadog/traces/sink.rs
+++ b/src/sinks/datadog/traces/sink.rs
@@ -111,10 +111,7 @@ where
         let batch_settings = self.batch_settings.clone();
 
         input
-            .batched_partitioned(
-                EventPartitioner,
-                Box::new(move || batch_settings.as_byte_size_config()),
-            )
+            .batched_partitioned(EventPartitioner, || batch_settings.as_byte_size_config())
             .incremental_request_builder(self.request_builder)
             .flat_map(stream::iter)
             .filter_map(|request| async move {
diff --git a/src/sinks/gcs_common/sink.rs b/src/sinks/gcs_common/sink.rs
index 4e6326fd8dc2e..1f1971a15152e 100644
--- a/src/sinks/gcs_common/sink.rs
+++ b/src/sinks/gcs_common/sink.rs
@@ -45,10 +45,7 @@ where
         let request_builder = self.request_builder;
 
         input
-            .batched_partitioned(
-                partitioner,
-                Box::new(move || settings.as_byte_size_config()),
-            )
+            .batched_partitioned(partitioner, || settings.as_byte_size_config())
             .filter_map(|(key, batch)| async move {
                 // A `TemplateRenderingError` will have been emitted by `KeyPartitioner` if the key here is `None`,
                 // thus no further `EventsDropped` event needs emitting at this stage.
diff --git a/src/sinks/loki/sink.rs b/src/sinks/loki/sink.rs
index 8563e2b45286a..1c588bcb4d78c 100644
--- a/src/sinks/loki/sink.rs
+++ b/src/sinks/loki/sink.rs
@@ -459,10 +459,7 @@ impl LokiSink {
             .map(|event| encoder.encode_event(event))
             .filter_map(|event| async { event })
             .map(|record| filter.filter_record(record))
-            .batched_partitioned(
-                RecordPartitioner,
-                Box::new(move || batch_settings.as_byte_size_config()),
-            )
+            .batched_partitioned(RecordPartitioner, || batch_settings.as_byte_size_config())
             .filter_map(|(partition, batch)| async {
                 if let Some(partition) = partition {
                     let mut count: usize = 0;
diff --git a/src/sinks/opendal_common.rs b/src/sinks/opendal_common.rs
index e63d570ae08d8..702ace4227da7 100644
--- a/src/sinks/opendal_common.rs
+++ b/src/sinks/opendal_common.rs
@@ -78,10 +78,7 @@ where
         let request_builder = self.request_builder;
 
         input
-            .batched_partitioned(
-                partitioner,
-                Box::new(move || settings.as_byte_size_config()),
-            )
+            .batched_partitioned(partitioner, || settings.as_byte_size_config())
             .filter_map(|(key, batch)| async move {
                 // We don't need to emit an error here if the event is dropped since this will occur if the template
                 // couldn't be rendered during the partitioning. A `TemplateRenderingError` is already emitted when
diff --git a/src/sinks/s3_common/sink.rs b/src/sinks/s3_common/sink.rs
index 7378854a17dd3..6d957daa4785b 100644
--- a/src/sinks/s3_common/sink.rs
+++ b/src/sinks/s3_common/sink.rs
@@ -44,10 +44,7 @@ where
         let request_builder = self.request_builder;
 
         input
-            .batched_partitioned(
-                partitioner,
-                Box::new(move || settings.as_byte_size_config()),
-            )
+            .batched_partitioned(partitioner, || settings.as_byte_size_config())
             .filter_map(|(key, batch)| async move { key.map(move |k| (k, batch)) })
             .request_builder(default_request_builder_concurrency_limit(), request_builder)
             .filter_map(|request| async move {
diff --git a/src/sinks/splunk_hec/logs/sink.rs b/src/sinks/splunk_hec/logs/sink.rs
index b4e488ac7c1da..a9a3b1d23e76c 100644
--- a/src/sinks/splunk_hec/logs/sink.rs
+++ b/src/sinks/splunk_hec/logs/sink.rs
@@ -78,7 +78,7 @@ where
                 } else {
                     EventPartitioner::new(None, None, None, None)
                 },
-                Box::new(move || batch_settings.as_byte_size_config()),
+                || batch_settings.as_byte_size_config(),
             )
             .request_builder(
                 default_request_builder_concurrency_limit(),
diff --git a/src/sinks/splunk_hec/metrics/sink.rs b/src/sinks/splunk_hec/metrics/sink.rs
index dbc847ec816b9..d8a4e23b57d77 100644
--- a/src/sinks/splunk_hec/metrics/sink.rs
+++ b/src/sinks/splunk_hec/metrics/sink.rs
@@ -54,10 +54,7 @@ where
                     default_namespace,
                 ))
             })
-            .batched_partitioned(
-                EventPartitioner,
-                Box::new(move || batch_settings.as_byte_size_config()),
-            )
+            .batched_partitioned(EventPartitioner, || batch_settings.as_byte_size_config())
             .request_builder(
                 default_request_builder_concurrency_limit(),
                 self.request_builder,
diff --git a/src/sinks/util/builder.rs b/src/sinks/util/builder.rs
index 30025bbdf9bf1..a25f0badb6f70 100644
--- a/src/sinks/util/builder.rs
+++ b/src/sinks/util/builder.rs
@@ -44,17 +44,18 @@ pub trait SinkBuilderExt: Stream {
     /// The stream will yield batches of events, with their partition key, when either a batch fills
     /// up or times out. [`Partitioner`] operates on a per-event basis, and has access to the event
     /// itself, and so can access any and all fields of an event.
-    fn batched_partitioned<P, C>(
+    fn batched_partitioned<P, C, F>(
         self,
         partitioner: P,
-        settings: Box<dyn Fn() -> C + Send>,
-    ) -> PartitionedBatcher<Self, P, ExpirationQueue<P::Key>, C>
+        settings: F,
+    ) -> PartitionedBatcher<Self, P, ExpirationQueue<P::Key>, C, F>
     where
         Self: Stream<Item = P::Item> + Sized,
         P: Partitioner + Unpin,
         P::Key: Eq + Hash + Clone,
         P::Item: ByteSizeOf,
         C: BatchConfig<P::Item>,
+        F: Fn() -> C + Send,
     {
         PartitionedBatcher::new(self, partitioner, settings)
     }

From f8967b7bae3dc68dc41fe18ec49c3edf100183cc Mon Sep 17 00:00:00 2001
From: Stephen Wakely <fungus.humungus@gmail.com>
Date: Fri, 6 Oct 2023 13:26:16 +0100
Subject: [PATCH 15/32] Clippy

Signed-off-by: Stephen Wakely <fungus.humungus@gmail.com>
---
 .../src/stream/partitioned_batcher.rs         | 20 ++++++++-----------
 src/sinks/aws_kinesis/sink.rs                 |  2 +-
 src/sinks/azure_common/sink.rs                |  2 +-
 src/sinks/clickhouse/sink.rs                  |  2 +-
 src/sinks/datadog/logs/sink.rs                |  2 +-
 src/sinks/datadog/metrics/sink.rs             |  2 +-
 src/sinks/datadog/traces/sink.rs              |  2 +-
 src/sinks/loki/sink.rs                        |  2 +-
 src/sinks/opendal_common.rs                   |  2 +-
 src/sinks/s3_common/sink.rs                   |  3 +--
 src/sinks/splunk_hec/logs/sink.rs             |  2 +-
 src/sinks/splunk_hec/metrics/sink.rs          |  2 +-
 12 files changed, 19 insertions(+), 24 deletions(-)

diff --git a/lib/vector-core/src/stream/partitioned_batcher.rs b/lib/vector-core/src/stream/partitioned_batcher.rs
index 68bb989ed371b..e3899a691d31d 100644
--- a/lib/vector-core/src/stream/partitioned_batcher.rs
+++ b/lib/vector-core/src/stream/partitioned_batcher.rs
@@ -302,15 +302,14 @@ where
                     let item_key = this.partitioner.partition(&item);
 
                     // Get the batch for this partition, or create a new one.
-                    let batch = match this.batches.get_mut(&item_key) {
-                        Some(batch) => batch,
-                        None => {
-                            let batch = (this.state)();
-                            this.batches.insert(item_key.clone(), batch);
-                            this.batches
-                                .get_mut(&item_key)
-                                .expect("batch has just been inserted so should exist")
-                        }
+                    let batch = if let Some(batch) = this.batches.get_mut(&item_key) {
+                        batch
+                    } else {
+                        let batch = (this.state)();
+                        this.batches.insert(item_key.clone(), batch);
+                        this.batches
+                            .get_mut(&item_key)
+                            .expect("batch has just been inserted so should exist")
                     };
 
                     let (fits, metadata) = batch.item_fits_in_batch(&item);
@@ -579,9 +578,6 @@ mod test {
                             .get_mut(&key)
                             .expect("impossible situation");
 
-                        dbg!(&actual_batch);
-                        dbg!(&expected_partition);
-
                         for item in actual_batch {
                             assert_eq!(item, expected_partition.pop().unwrap());
                         }
diff --git a/src/sinks/aws_kinesis/sink.rs b/src/sinks/aws_kinesis/sink.rs
index 60e056d2811ad..6dfb0be08d19e 100644
--- a/src/sinks/aws_kinesis/sink.rs
+++ b/src/sinks/aws_kinesis/sink.rs
@@ -42,7 +42,7 @@ where
     R: Record + Send + Sync + Unpin + Clone + 'static,
 {
     async fn run_inner(self: Box<Self>, input: BoxStream<'_, Event>) -> Result<(), ()> {
-        let batch_settings = self.batch_settings.clone();
+        let batch_settings = self.batch_settings;
 
         input
             .filter_map(|event| {
diff --git a/src/sinks/azure_common/sink.rs b/src/sinks/azure_common/sink.rs
index a954217a15f33..4844474232cd0 100644
--- a/src/sinks/azure_common/sink.rs
+++ b/src/sinks/azure_common/sink.rs
@@ -37,7 +37,7 @@ where
 {
     async fn run_inner(self: Box<Self>, input: BoxStream<'_, Event>) -> Result<(), ()> {
         let partitioner = self.partitioner;
-        let settings = self.batcher_settings.clone();
+        let settings = self.batcher_settings;
 
         let request_builder = self.request_builder;
 
diff --git a/src/sinks/clickhouse/sink.rs b/src/sinks/clickhouse/sink.rs
index 987bc5c1e34e3..be9c28a0e83b4 100644
--- a/src/sinks/clickhouse/sink.rs
+++ b/src/sinks/clickhouse/sink.rs
@@ -42,7 +42,7 @@ impl ClickhouseSink {
     }
 
     async fn run_inner(self: Box<Self>, input: BoxStream<'_, Event>) -> Result<(), ()> {
-        let batch_settings = self.batch_settings.clone();
+        let batch_settings = self.batch_settings;
 
         input
             .batched_partitioned(KeyPartitioner::new(self.database, self.table), || {
diff --git a/src/sinks/datadog/logs/sink.rs b/src/sinks/datadog/logs/sink.rs
index 2abad6286bd47..7b1245edb6389 100644
--- a/src/sinks/datadog/logs/sink.rs
+++ b/src/sinks/datadog/logs/sink.rs
@@ -261,7 +261,7 @@ where
         let default_api_key = Arc::clone(&self.default_api_key);
 
         let partitioner = EventPartitioner;
-        let batch_settings = self.batch_settings.clone();
+        let batch_settings = self.batch_settings;
 
         let input = input.batched_partitioned(partitioner, || batch_settings.as_byte_size_config());
         input
diff --git a/src/sinks/datadog/metrics/sink.rs b/src/sinks/datadog/metrics/sink.rs
index 23ea4e2a77227..0b50549acefe0 100644
--- a/src/sinks/datadog/metrics/sink.rs
+++ b/src/sinks/datadog/metrics/sink.rs
@@ -84,7 +84,7 @@ where
 
     async fn run_inner(self: Box<Self>, input: BoxStream<'_, Event>) -> Result<(), ()> {
         let mut splitter: MetricSplitter<AggregatedSummarySplitter> = MetricSplitter::default();
-        let batch_settings = self.batch_settings.clone();
+        let batch_settings = self.batch_settings;
 
         input
             // Convert `Event` to `Metric` so we don't have to deal with constant conversions.
diff --git a/src/sinks/datadog/traces/sink.rs b/src/sinks/datadog/traces/sink.rs
index dd09b57a64d86..384d7ef635022 100644
--- a/src/sinks/datadog/traces/sink.rs
+++ b/src/sinks/datadog/traces/sink.rs
@@ -108,7 +108,7 @@ where
     }
 
     async fn run_inner(self: Box<Self>, input: BoxStream<'_, Event>) -> Result<(), ()> {
-        let batch_settings = self.batch_settings.clone();
+        let batch_settings = self.batch_settings;
 
         input
             .batched_partitioned(EventPartitioner, || batch_settings.as_byte_size_config())
diff --git a/src/sinks/loki/sink.rs b/src/sinks/loki/sink.rs
index 1c588bcb4d78c..1b0d45bb3092b 100644
--- a/src/sinks/loki/sink.rs
+++ b/src/sinks/loki/sink.rs
@@ -453,7 +453,7 @@ impl LokiSink {
                 NonZeroUsize::new(1).expect("static")
             }
         };
-        let batch_settings = self.batch_settings.clone();
+        let batch_settings = self.batch_settings;
 
         input
             .map(|event| encoder.encode_event(event))
diff --git a/src/sinks/opendal_common.rs b/src/sinks/opendal_common.rs
index 702ace4227da7..34b4962771835 100644
--- a/src/sinks/opendal_common.rs
+++ b/src/sinks/opendal_common.rs
@@ -73,7 +73,7 @@ where
 {
     async fn run_inner(self: Box<Self>, input: BoxStream<'_, Event>) -> Result<(), ()> {
         let partitioner = self.partitioner;
-        let settings = self.batcher_settings.clone();
+        let settings = self.batcher_settings;
 
         let request_builder = self.request_builder;
 
diff --git a/src/sinks/s3_common/sink.rs b/src/sinks/s3_common/sink.rs
index 6d957daa4785b..fd6733a484a75 100644
--- a/src/sinks/s3_common/sink.rs
+++ b/src/sinks/s3_common/sink.rs
@@ -39,8 +39,7 @@ where
 {
     async fn run_inner(self: Box<Self>, input: BoxStream<'_, Event>) -> Result<(), ()> {
         let partitioner = self.partitioner;
-        let settings = self.batcher_settings.clone();
-
+        let settings = self.batcher_settings;
         let request_builder = self.request_builder;
 
         input
diff --git a/src/sinks/splunk_hec/logs/sink.rs b/src/sinks/splunk_hec/logs/sink.rs
index a9a3b1d23e76c..7fed3bf3310e0 100644
--- a/src/sinks/splunk_hec/logs/sink.rs
+++ b/src/sinks/splunk_hec/logs/sink.rs
@@ -61,7 +61,7 @@ where
             timestamp_key: self.timestamp_key.clone(),
             endpoint_target: self.endpoint_target,
         };
-        let batch_settings = self.batch_settings.clone();
+        let batch_settings = self.batch_settings;
 
         input
             .map(move |event| process_log(event, &data))
diff --git a/src/sinks/splunk_hec/metrics/sink.rs b/src/sinks/splunk_hec/metrics/sink.rs
index d8a4e23b57d77..6f6ad5d74b9dd 100644
--- a/src/sinks/splunk_hec/metrics/sink.rs
+++ b/src/sinks/splunk_hec/metrics/sink.rs
@@ -39,7 +39,7 @@ where
         let index = self.index.as_ref();
         let host_key = self.host_key.as_ref();
         let default_namespace = self.default_namespace.as_deref();
-        let batch_settings = self.batch_settings.clone();
+        let batch_settings = self.batch_settings;
 
         input
             .map(|event| (event.size_of(), event.into_metric()))

From 896d76faef19c38553e4bbaae75424da34f1415d Mon Sep 17 00:00:00 2001
From: Stephen Wakely <fungus.humungus@gmail.com>
Date: Fri, 6 Oct 2023 16:05:26 +0100
Subject: [PATCH 16/32] Only insert timeout if we add the batch to the list

Signed-off-by: Stephen Wakely <fungus.humungus@gmail.com>
---
 lib/vector-core/src/stream/partitioned_batcher.rs | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/lib/vector-core/src/stream/partitioned_batcher.rs b/lib/vector-core/src/stream/partitioned_batcher.rs
index e3899a691d31d..ab82a34bda404 100644
--- a/lib/vector-core/src/stream/partitioned_batcher.rs
+++ b/lib/vector-core/src/stream/partitioned_batcher.rs
@@ -323,13 +323,14 @@ where
 
                     // Insert the item into the batch.
                     batch.push(item, metadata);
-                    this.timer.insert(item_key.clone());
                     if batch.is_batch_full() {
                         // If the insertion means the batch is now full, we clear out the batch and
                         // remove it from the list.
                         this.closed_batches
                             .push((item_key.clone(), batch.take_batch()));
                         this.batches.remove(&item_key);
+                    } else {
+                        this.timer.insert(item_key.clone());
                     }
                 }
             }

From 817bc4623156f03453bf8f48e5109de96955d6a9 Mon Sep 17 00:00:00 2001
From: Stephen Wakely <fungus.humungus@gmail.com>
Date: Fri, 6 Oct 2023 16:28:50 +0100
Subject: [PATCH 17/32] Allow the timer to remove an item

Signed-off-by: Stephen Wakely <fungus.humungus@gmail.com>
---
 lib/vector-core/src/stream/partitioned_batcher.rs | 11 ++++++++++-
 lib/vector-core/src/time.rs                       |  3 +++
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/lib/vector-core/src/stream/partitioned_batcher.rs b/lib/vector-core/src/stream/partitioned_batcher.rs
index ab82a34bda404..e6afbf4ab79fc 100644
--- a/lib/vector-core/src/stream/partitioned_batcher.rs
+++ b/lib/vector-core/src/stream/partitioned_batcher.rs
@@ -86,6 +86,10 @@ where
         }
     }
 
+    fn remove(&mut self, item_key: &K) {
+        self.expiration_map.remove(item_key);
+    }
+
     fn poll_expired(&mut self, cx: &mut Context) -> Poll<Option<K>> {
         match ready!(self.expirations.poll_expired(cx)) {
             // No expirations yet.
@@ -329,8 +333,9 @@ where
                         this.closed_batches
                             .push((item_key.clone(), batch.take_batch()));
                         this.batches.remove(&item_key);
+                        this.timer.remove(&item_key);
                     } else {
-                        this.timer.insert(item_key.clone());
+                        this.timer.insert(item_key);
                     }
                 }
             }
@@ -392,6 +397,10 @@ mod test {
             self.valid_keys.insert(item_key);
         }
 
+        fn remove(&mut self, item_key: &u8) {
+            self.valid_keys.remove(item_key);
+        }
+
         fn poll_expired(&mut self, _cx: &mut Context) -> Poll<Option<u8>> {
             match self.responses.pop() {
                 Some(Poll::Pending) => unreachable!(),
diff --git a/lib/vector-core/src/time.rs b/lib/vector-core/src/time.rs
index 49481842cfcdd..e4e4e4975ebe4 100644
--- a/lib/vector-core/src/time.rs
+++ b/lib/vector-core/src/time.rs
@@ -21,6 +21,9 @@ pub trait KeyedTimer<K> {
     /// If the given key already exists in the timer, the underlying subtimer is reset.
     fn insert(&mut self, item_key: K);
 
+    /// Removes a subtimer from the list.
+    fn remove(&mut self, item_key: &K);
+
     /// Attempts to pull out the next expired subtimer in the queue.
     ///
     /// The key of the subtimer is returned if it has expired, otherwise, returns `None` if the

From eab7397beb9752ddbaf8146fb0ec217c2dba56df Mon Sep 17 00:00:00 2001
From: Stephen Wakely <fungus.humungus@gmail.com>
Date: Wed, 11 Oct 2023 17:21:25 +0100
Subject: [PATCH 18/32] Allow partitions to be types other than Vec

Signed-off-by: Stephen Wakely <fungus.humungus@gmail.com>
---
 lib/vector-core/src/stream/mod.rs             |   2 +-
 .../src/stream/partitioned_batcher.rs         | 256 ++++++------------
 .../remote_write/request_builder.rs           |   2 +-
 src/sinks/prometheus/remote_write/sink.rs     |   7 +-
 src/sinks/prometheus/remote_write/tests.rs    |   2 -
 src/sinks/util/builder.rs                     |   5 +-
 6 files changed, 86 insertions(+), 188 deletions(-)

diff --git a/lib/vector-core/src/stream/mod.rs b/lib/vector-core/src/stream/mod.rs
index e44a825901ec3..64e4edc3247f9 100644
--- a/lib/vector-core/src/stream/mod.rs
+++ b/lib/vector-core/src/stream/mod.rs
@@ -8,4 +8,4 @@ mod partitioned_batcher;
 pub use concurrent_map::ConcurrentMap;
 pub use driver::{Driver, DriverResponse};
 use futures_unordered_count::FuturesUnorderedCount;
-pub use partitioned_batcher::{BatchOutput, BatcherSettings, ExpirationQueue, PartitionedBatcher};
+pub use partitioned_batcher::{BatcherSettings, ExpirationQueue, PartitionedBatcher};
diff --git a/lib/vector-core/src/stream/partitioned_batcher.rs b/lib/vector-core/src/stream/partitioned_batcher.rs
index fd0c197224906..60914302f7ef6 100644
--- a/lib/vector-core/src/stream/partitioned_batcher.rs
+++ b/lib/vector-core/src/stream/partitioned_batcher.rs
@@ -183,7 +183,7 @@ impl BatcherSettings {
 }
 
 #[pin_project]
-pub struct PartitionedBatcher<St, Prt, KT, C, F>
+pub struct PartitionedBatcher<St, Prt, KT, C, F, B>
 where
     Prt: Partitioner,
 {
@@ -196,7 +196,7 @@ where
     /// The store of 'closed' batches. When this is not empty it will be
     /// preferentially flushed prior to consuming any new items from the
     /// underlying stream.
-    closed_batches: Vec<(Prt::Key, Vec<Prt::Item>)>,
+    closed_batches: Vec<(Prt::Key, B)>,
     /// The queue of pending batch expirations
     timer: KT,
     /// The partitioner for this `Batcher`
@@ -206,7 +206,7 @@ where
     stream: Fuse<St>,
 }
 
-impl<St, Prt, C, F> PartitionedBatcher<St, Prt, ExpirationQueue<Prt::Key>, C, F>
+impl<St, Prt, C, F, B> PartitionedBatcher<St, Prt, ExpirationQueue<Prt::Key>, C, F, B>
 where
     St: Stream<Item = Prt::Item>,
     Prt: Partitioner + Unpin,
@@ -229,7 +229,7 @@ where
 }
 
 #[cfg(test)]
-impl<St, Prt, KT, C, F> PartitionedBatcher<St, Prt, KT, C, F>
+impl<St, Prt, KT, C, F, B> PartitionedBatcher<St, Prt, KT, C, F, B>
 where
     St: Stream<Item = Prt::Item>,
     Prt: Partitioner + Unpin,
@@ -250,11 +250,7 @@ where
     }
 }
 
-pub trait BatchOutput<T> {}
-
-impl<T> BatchOutput<T> for Vec<T> {}
-
-impl<St, Prt, KT, C, F, B> Stream for PartitionedBatcher<St, Prt, KT, C, F>
+impl<St, Prt, KT, C, F, B> Stream for PartitionedBatcher<St, Prt, KT, C, F, B>
 where
     St: Stream<Item = Prt::Item>,
     Prt: Partitioner + Unpin,
@@ -262,9 +258,7 @@ where
     Prt::Item: ByteSizeOf,
     KT: KeyedTimer<Prt::Key>,
     C: BatchConfig<Prt::Item, Batch = B>,
-    // C: BatchConfig<Prt::Item>,
     F: Fn() -> C + Send,
-    B: BatchOutput<Prt::Item>,
 {
     type Item = (Prt::Key, B);
 
@@ -272,82 +266,80 @@ where
         self.stream.size_hint()
     }
 
-    #[allow(warnings)]
     fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
-        todo!()
-        // let mut this = self.project();
-        // loop {
-        //     if !this.closed_batches.is_empty() {
-        //         return Poll::Ready(this.closed_batches.pop());
-        //     }
-        // match this.stream.as_mut().poll_next(cx) {
-        //     Poll::Pending => match this.timer.poll_expired(cx) {
-        //         // Unlike normal streams, `DelayQueue` can return `None`
-        //         // here but still be usable later if more entries are added.
-        //         Poll::Pending | Poll::Ready(None) => return Poll::Pending,
-        //         Poll::Ready(Some(item_key)) => {
-        //             let mut batch = this
-        //                 .batches
-        //                 .remove(&item_key)
-        //                 .expect("batch should exist if it is set to expire");
-        //             this.closed_batches.push((item_key, batch.take_batch()));
-        //         }
-        //     },
-        //     Poll::Ready(None) => {
-        //         // Now that the underlying stream is closed, we need to
-        //         // clear out our batches, including all expiration
-        //         // entries. If we had any batches to hand over, we have to
-        //         // continue looping so the caller can drain them all before
-        //         // we finish.
-        //         if !this.batches.is_empty() {
-        //             this.timer.clear();
-        //             this.closed_batches.extend(
-        //                 this.batches
-        //                     .drain()
-        //                     .map(|(key, mut batch)| (key, batch.take_batch())),
-        //             );
-        //             continue;
-        //         }
-        //         return Poll::Ready(None);
-        //     }
-        //     Poll::Ready(Some(item)) => {
-        //         let item_key = this.partitioner.partition(&item);
-
-        //         // Get the batch for this partition, or create a new one.
-        //         let batch = if let Some(batch) = this.batches.get_mut(&item_key) {
-        //             batch
-        //         } else {
-        //             let batch = (this.state)();
-        //             this.batches.insert(item_key.clone(), batch);
-        //             this.batches
-        //                 .get_mut(&item_key)
-        //                 .expect("batch has just been inserted so should exist")
-        //         };
-
-        //         let (fits, metadata) = batch.item_fits_in_batch(&item);
-        //         if !fits {
-        //             // This batch is too full to accept a new item, so we move the contents of
-        //             // the batch into `closed_batches` to be push out of this stream on the
-        //             // next iteration.
-        //             this.closed_batches
-        //                 .push((item_key.clone(), batch.take_batch()));
-        //         }
-
-        //         // Insert the item into the batch.
-        //         batch.push(item, metadata);
-        //         if batch.is_batch_full() {
-        //             // If the insertion means the batch is now full, we clear out the batch and
-        //             // remove it from the list.
-        //             this.closed_batches
-        //                 .push((item_key.clone(), batch.take_batch()));
-        //             this.batches.remove(&item_key);
-        //             this.timer.remove(&item_key);
-        //         } else {
-        //             this.timer.insert(item_key);
-        //         }
-        //     }
-        // }
-        // }
+        let mut this = self.project();
+        loop {
+            if !this.closed_batches.is_empty() {
+                return Poll::Ready(this.closed_batches.pop());
+            }
+            match this.stream.as_mut().poll_next(cx) {
+                Poll::Pending => match this.timer.poll_expired(cx) {
+                    // Unlike normal streams, `DelayQueue` can return `None`
+                    // here but still be usable later if more entries are added.
+                    Poll::Pending | Poll::Ready(None) => return Poll::Pending,
+                    Poll::Ready(Some(item_key)) => {
+                        let mut batch = this
+                            .batches
+                            .remove(&item_key)
+                            .expect("batch should exist if it is set to expire");
+                        this.closed_batches.push((item_key, batch.take_batch()));
+                    }
+                },
+                Poll::Ready(None) => {
+                    // Now that the underlying stream is closed, we need to
+                    // clear out our batches, including all expiration
+                    // entries. If we had any batches to hand over, we have to
+                    // continue looping so the caller can drain them all before
+                    // we finish.
+                    if !this.batches.is_empty() {
+                        this.timer.clear();
+                        this.closed_batches.extend(
+                            this.batches
+                                .drain()
+                                .map(|(key, mut batch)| (key, batch.take_batch())),
+                        );
+                        continue;
+                    }
+                    return Poll::Ready(None);
+                }
+                Poll::Ready(Some(item)) => {
+                    let item_key = this.partitioner.partition(&item);
+
+                    // Get the batch for this partition, or create a new one.
+                    let batch = if let Some(batch) = this.batches.get_mut(&item_key) {
+                        batch
+                    } else {
+                        let batch = (this.state)();
+                        this.batches.insert(item_key.clone(), batch);
+                        this.batches
+                            .get_mut(&item_key)
+                            .expect("batch has just been inserted so should exist")
+                    };
+
+                    let (fits, metadata) = batch.item_fits_in_batch(&item);
+                    if !fits {
+                        // This batch is too full to accept a new item, so we move the contents of
+                        // the batch into `closed_batches` to be push out of this stream on the
+                        // next iteration.
+                        this.closed_batches
+                            .push((item_key.clone(), batch.take_batch()));
+                    }
+
+                    // Insert the item into the batch.
+                    batch.push(item, metadata);
+                    if batch.is_batch_full() {
+                        // If the insertion means the batch is now full, we clear out the batch and
+                        // remove it from the list.
+                        this.closed_batches
+                            .push((item_key.clone(), batch.take_batch()));
+                        this.batches.remove(&item_key);
+                        this.timer.remove(&item_key);
+                    } else {
+                        this.timer.insert(item_key);
+                    }
+                }
+            }
+        }
     }
 }
 
@@ -608,98 +600,6 @@ mod test {
         }
     }
 
-    #[tokio::test]
-    async fn zorkwonk() {
-        // let stream = (0..20).collect::<Vec<_>>();
-        // let stream = vec![
-        //     7037641852729981347,
-        //     8644263311095074276,
-        //     4920834891500558744,
-        //     15150872024001141310,
-        //     12803461249171200340,
-        // ];
-        let stream = vec![
-            7037641852729981347,
-            8644263311095074276,
-            4920834891500558744,
-            15150872024001141310,
-            12803461249171200340,
-        ];
-        let item_limit = 1;
-        let allocation_limit = 8;
-        let partitioner = TestPartitioner {
-            key_space: NonZeroU8::new(23).unwrap(),
-        };
-        let mut count = 0;
-        let timer = TestTimer::new(
-            std::iter::from_fn(move || {
-                if count > 254 {
-                    None
-                } else {
-                    count += 1;
-                    if count % 5 == 0 {
-                        Some(Poll::Ready(Some(count)))
-                    } else {
-                        Some(Poll::Ready(None))
-                    }
-                }
-            })
-            .collect(),
-        );
-
-        // Asserts that for every received batch received the elements in
-        // the batch are not reordered within a batch. No claim is made on
-        // when batches themselves will issue, batch sizes etc.
-        let noop_waker = futures::task::noop_waker();
-        let mut cx = Context::from_waker(&noop_waker);
-
-        let mut partitions = separate_partitions(stream.clone(), &partitioner);
-
-        let mut stream = stream::iter(stream.into_iter());
-        let item_limit = NonZeroUsize::new(item_limit as usize).unwrap();
-        let allocation_limit = NonZeroUsize::new(allocation_limit as usize).unwrap();
-        let batch_settings =
-            BatcherSettings::new(Duration::from_secs(1), allocation_limit, item_limit);
-
-        let mut batcher = PartitionedBatcher::with_timer(
-            &mut stream,
-            partitioner,
-            timer,
-            Box::new(move || batch_settings.clone().as_byte_size_config()),
-        );
-        // let mut batcher = PartitionedBatcher::new(
-        //     &mut stream,
-        //     partitioner,
-        //     // timer,
-        //     Box::new(move || batch_settings.clone().into_byte_size_config()),
-        // );
-        let mut batcher = Pin::new(&mut batcher);
-
-        loop {
-            match batcher.as_mut().poll_next(&mut cx) {
-                Poll::Pending => {}
-                Poll::Ready(None) => {
-                    break;
-                }
-                Poll::Ready(Some((key, actual_batch))) => {
-                    let expected_partition =
-                        partitions.get_mut(&key).expect("impossible situation");
-
-                    dbg!(&actual_batch);
-                    dbg!(&expected_partition);
-
-                    for item in actual_batch {
-                        assert_eq!(item, expected_partition.pop().unwrap());
-                    }
-                }
-            }
-        }
-        for v in partitions.values() {
-            assert!(v.is_empty());
-        }
-        // panic!();
-    }
-
     proptest! {
         #[test]
         fn batch_does_not_lose_items(stream: Vec<u64>,
diff --git a/src/sinks/prometheus/remote_write/request_builder.rs b/src/sinks/prometheus/remote_write/request_builder.rs
index 8f6d37514b76c..8161dc2138535 100644
--- a/src/sinks/prometheus/remote_write/request_builder.rs
+++ b/src/sinks/prometheus/remote_write/request_builder.rs
@@ -103,7 +103,7 @@ impl RequestBuilder<(PartitionKey, EventCollection)> for RemoteWriteRequestBuild
         &self,
         input: (PartitionKey, EventCollection),
     ) -> (Self::Metadata, RequestMetadataBuilder, Self::Events) {
-        let (key, mut events) = input;
+        let (key, events) = input;
         let finalizers = events.finalizers;
         let metrics = events.events.into_metrics();
         let builder = RequestMetadataBuilder::from_events(&metrics);
diff --git a/src/sinks/prometheus/remote_write/sink.rs b/src/sinks/prometheus/remote_write/sink.rs
index 56ac22fa86ff3..5dc511c1b87af 100644
--- a/src/sinks/prometheus/remote_write/sink.rs
+++ b/src/sinks/prometheus/remote_write/sink.rs
@@ -87,8 +87,6 @@ impl Default for EventCollection {
     }
 }
 
-impl<T> vector_core::stream::BatchOutput<T> for Vec<T> {}
-
 pub(super) struct RemoteWriteSink<S> {
     pub(super) tenant_id: Option<Template>,
     pub(super) batch_settings: BatcherSettings,
@@ -126,7 +124,7 @@ where
             .filter_map(move |event| {
                 future::ready(make_remote_write_event(tenant_id.as_ref(), event))
             })
-            .batched_partitioned(PrometheusTenantIdPartitioner, || fun_name(batch_settings))
+            .batched_partitioned(PrometheusTenantIdPartitioner, || reducer(batch_settings))
             .request_builder(default_request_builder_concurrency_limit(), request_builder)
             .filter_map(|request| async move {
                 match request {
@@ -143,7 +141,8 @@ where
     }
 }
 
-fn fun_name(
+/// Create a reducer from the batch settings.
+fn reducer(
     batch_settings: BatcherSettings,
 ) -> BatchConfigParts<
     SizeLimit<ByteSizeOfItemSize>,
diff --git a/src/sinks/prometheus/remote_write/tests.rs b/src/sinks/prometheus/remote_write/tests.rs
index 58d6430d38667..da88851802a52 100644
--- a/src/sinks/prometheus/remote_write/tests.rs
+++ b/src/sinks/prometheus/remote_write/tests.rs
@@ -150,8 +150,6 @@ async fn retains_state_between_requests() {
     )
     .await;
 
-    dbg!(&outputs);
-
     assert_eq!(outputs.len(), 3);
 
     let check_output = |index: usize, name: &str, value: f64| {
diff --git a/src/sinks/util/builder.rs b/src/sinks/util/builder.rs
index a25f0badb6f70..a36eb1e79493d 100644
--- a/src/sinks/util/builder.rs
+++ b/src/sinks/util/builder.rs
@@ -44,11 +44,11 @@ pub trait SinkBuilderExt: Stream {
     /// The stream will yield batches of events, with their partition key, when either a batch fills
     /// up or times out. [`Partitioner`] operates on a per-event basis, and has access to the event
     /// itself, and so can access any and all fields of an event.
-    fn batched_partitioned<P, C, F>(
+    fn batched_partitioned<P, C, F, B>(
         self,
         partitioner: P,
         settings: F,
-    ) -> PartitionedBatcher<Self, P, ExpirationQueue<P::Key>, C, F>
+    ) -> PartitionedBatcher<Self, P, ExpirationQueue<P::Key>, C, F, B>
     where
         Self: Stream<Item = P::Item> + Sized,
         P: Partitioner + Unpin,
@@ -56,6 +56,7 @@ pub trait SinkBuilderExt: Stream {
         P::Item: ByteSizeOf,
         C: BatchConfig<P::Item>,
         F: Fn() -> C + Send,
+        // B: BatchOutput<P::Item>,
     {
         PartitionedBatcher::new(self, partitioner, settings)
     }

From 2e3ea52f4a666992751342a0fdbf27f211ddaffc Mon Sep 17 00:00:00 2001
From: Stephen Wakely <fungus.humungus@gmail.com>
Date: Thu, 12 Oct 2023 09:48:34 +0100
Subject: [PATCH 19/32] Clippy

Signed-off-by: Stephen Wakely <fungus.humungus@gmail.com>
---
 src/sinks/azure_common/sink.rs | 2 +-
 src/sinks/opendal_common.rs    | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/sinks/azure_common/sink.rs b/src/sinks/azure_common/sink.rs
index a954217a15f33..4844474232cd0 100644
--- a/src/sinks/azure_common/sink.rs
+++ b/src/sinks/azure_common/sink.rs
@@ -37,7 +37,7 @@ where
 {
     async fn run_inner(self: Box<Self>, input: BoxStream<'_, Event>) -> Result<(), ()> {
         let partitioner = self.partitioner;
-        let settings = self.batcher_settings.clone();
+        let settings = self.batcher_settings;
 
         let request_builder = self.request_builder;
 
diff --git a/src/sinks/opendal_common.rs b/src/sinks/opendal_common.rs
index 702ace4227da7..34b4962771835 100644
--- a/src/sinks/opendal_common.rs
+++ b/src/sinks/opendal_common.rs
@@ -73,7 +73,7 @@ where
 {
     async fn run_inner(self: Box<Self>, input: BoxStream<'_, Event>) -> Result<(), ()> {
         let partitioner = self.partitioner;
-        let settings = self.batcher_settings.clone();
+        let settings = self.batcher_settings;
 
         let request_builder = self.request_builder;
 

From aa54d4ec345e41173947d9ec01b666092bcbe1f2 Mon Sep 17 00:00:00 2001
From: Stephen Wakely <fungus.humungus@gmail.com>
Date: Thu, 12 Oct 2023 11:01:51 +0100
Subject: [PATCH 20/32] Added test for aggregation

Signed-off-by: Stephen Wakely <fungus.humungus@gmail.com>
---
 src/sinks/prometheus/remote_write/tests.rs | 26 ++++++++++++++++++++++
 1 file changed, 26 insertions(+)

diff --git a/src/sinks/prometheus/remote_write/tests.rs b/src/sinks/prometheus/remote_write/tests.rs
index da88851802a52..cffcb3aef502d 100644
--- a/src/sinks/prometheus/remote_write/tests.rs
+++ b/src/sinks/prometheus/remote_write/tests.rs
@@ -164,6 +164,32 @@ async fn retains_state_between_requests() {
     check_output(2, "counter-1", 26.0);
 }
 
+#[tokio::test]
+async fn aggregates_batches() {
+    let outputs = send_request(
+        r#"batch.max_events = 2"#,
+        vec![
+            create_inc_event("counter-1".into(), 12.0),
+            create_inc_event("counter-2".into(), 13.0),
+            create_inc_event("counter-1".into(), 14.0),
+            create_inc_event("counter-2".into(), 14.0),
+        ],
+    )
+    .await;
+
+    assert_eq!(outputs.len(), 1);
+
+    let (_, req) = &outputs[0];
+    assert_eq!(req.timeseries.len(), 2);
+    assert_eq!(req.timeseries[0].labels, labels!("__name__" => "counter-1"));
+    assert_eq!(req.timeseries[0].samples.len(), 1);
+    assert_eq!(req.timeseries[0].samples[0].value, 26.0);
+
+    assert_eq!(req.timeseries[1].labels, labels!("__name__" => "counter-2"));
+    assert_eq!(req.timeseries[1].samples.len(), 1);
+    assert_eq!(req.timeseries[1].samples[0].value, 27.0);
+}
+
 async fn send_request(config: &str, events: Vec<Event>) -> Vec<(HeaderMap, proto::WriteRequest)> {
     assert_sink_compliance(&HTTP_SINK_TAGS, async {
         let addr = test_util::next_addr();

From 3af06d5c84805dc32f5bf02a79ed51cc3c4d7609 Mon Sep 17 00:00:00 2001
From: Stephen Wakely <fungus.humungus@gmail.com>
Date: Fri, 13 Oct 2023 10:42:36 +0100
Subject: [PATCH 21/32] Allow a custom object to be used for the reducer

Signed-off-by: Stephen Wakely <fungus.humungus@gmail.com>
---
 .../src/stream/partitioned_batcher.rs         | 18 +++----
 src/sinks/prometheus/remote_write/sink.rs     | 53 +++++++++----------
 src/sinks/prometheus/remote_write/tests.rs    |  4 +-
 src/sinks/vector/sink.rs                      |  8 +--
 4 files changed, 39 insertions(+), 44 deletions(-)

diff --git a/lib/vector-core/src/stream/partitioned_batcher.rs b/lib/vector-core/src/stream/partitioned_batcher.rs
index 60914302f7ef6..8afa5e7b0c829 100644
--- a/lib/vector-core/src/stream/partitioned_batcher.rs
+++ b/lib/vector-core/src/stream/partitioned_batcher.rs
@@ -16,14 +16,13 @@ use crate::{
     partition::Partitioner,
     stream::batcher::{
         config::BatchConfigParts,
-        data::BatchReduce,
         limiter::{ByteSizeOfItemSize, ItemBatchSize, SizeLimit},
     },
     time::KeyedTimer,
     ByteSizeOf,
 };
 
-use super::batcher::BatchConfig;
+use super::batcher::{data::BatchData, BatchConfig};
 
 /// A `KeyedTimer` based on `DelayQueue`.
 pub struct ExpirationQueue<K> {
@@ -158,16 +157,15 @@ impl BatcherSettings {
     }
 
     /// A batcher config using the `ItemBatchSize` trait to determine batch sizes.
-    /// The output is built with the supplied reducer function.
-    pub fn into_reducer_config<I, T, F, S>(
-        self,
+    /// The output is built with the supplied object implementing [`BatchData`].
+    pub fn as_reducer_config<I, T, B>(
+        &self,
         item_size: I,
-        reducer: F,
-    ) -> BatchConfigParts<SizeLimit<I>, BatchReduce<F, S>>
+        reducer: B,
+    ) -> BatchConfigParts<SizeLimit<I>, B>
     where
         I: ItemBatchSize<T>,
-        F: FnMut(&mut S, T) + Send,
-        S: Default,
+        B: BatchData<T>,
     {
         BatchConfigParts {
             batch_limiter: SizeLimit {
@@ -176,7 +174,7 @@ impl BatcherSettings {
                 current_size: 0,
                 item_size_calculator: item_size,
             },
-            batch_data: BatchReduce::new(reducer),
+            batch_data: reducer,
             timeout: self.timeout,
         }
     }
diff --git a/src/sinks/prometheus/remote_write/sink.rs b/src/sinks/prometheus/remote_write/sink.rs
index 5dc511c1b87af..a6116667e8edd 100644
--- a/src/sinks/prometheus/remote_write/sink.rs
+++ b/src/sinks/prometheus/remote_write/sink.rs
@@ -3,11 +3,7 @@ use std::fmt;
 use vector_common::byte_size_of::ByteSizeOf;
 use vector_core::{
     event::Metric,
-    stream::batcher::{
-        config::BatchConfigParts,
-        data::BatchReduce,
-        limiter::{ByteSizeOfItemSize, SizeLimit},
-    },
+    stream::batcher::{data::BatchData, limiter::ByteSizeOfItemSize},
 };
 
 use crate::sinks::{prelude::*, util::buffer::metrics::MetricSet};
@@ -87,6 +83,27 @@ impl Default for EventCollection {
     }
 }
 
+impl BatchData<RemoteWriteMetric> for EventCollection {
+    type Batch = Self;
+
+    fn len(&self) -> usize {
+        self.events.len()
+    }
+
+    fn take_batch(&mut self) -> Self::Batch {
+        std::mem::take(self)
+    }
+
+    fn push_item(&mut self, mut item: RemoteWriteMetric) {
+        self.finalizers
+            .merge(item.metric.metadata_mut().take_finalizers());
+        self.events_byte_size += item.size_of();
+        self.events_json_byte_size
+            .add_event(&item.metric, item.estimated_json_encoded_size_of());
+        self.events.insert_update(item.metric);
+    }
+}
+
 pub(super) struct RemoteWriteSink<S> {
     pub(super) tenant_id: Option<Template>,
     pub(super) batch_settings: BatcherSettings,
@@ -124,7 +141,9 @@ where
             .filter_map(move |event| {
                 future::ready(make_remote_write_event(tenant_id.as_ref(), event))
             })
-            .batched_partitioned(PrometheusTenantIdPartitioner, || reducer(batch_settings))
+            .batched_partitioned(PrometheusTenantIdPartitioner, || {
+                batch_settings.as_reducer_config(ByteSizeOfItemSize, EventCollection::default())
+            })
             .request_builder(default_request_builder_concurrency_limit(), request_builder)
             .filter_map(|request| async move {
                 match request {
@@ -141,28 +160,6 @@ where
     }
 }
 
-/// Create a reducer from the batch settings.
-fn reducer(
-    batch_settings: BatcherSettings,
-) -> BatchConfigParts<
-    SizeLimit<ByteSizeOfItemSize>,
-    BatchReduce<impl FnMut(&mut EventCollection, RemoteWriteMetric) + Send, EventCollection>,
-> {
-    batch_settings.into_reducer_config(
-        ByteSizeOfItemSize,
-        |event_collection: &mut EventCollection, mut item: RemoteWriteMetric| {
-            event_collection
-                .finalizers
-                .merge(item.metric.metadata_mut().take_finalizers());
-            event_collection.events_byte_size += item.size_of();
-            event_collection
-                .events_json_byte_size
-                .add_event(&item.metric, item.estimated_json_encoded_size_of());
-            event_collection.events.insert_update(item.metric);
-        },
-    )
-}
-
 #[async_trait]
 impl<S> StreamSink<Event> for RemoteWriteSink<S>
 where
diff --git a/src/sinks/prometheus/remote_write/tests.rs b/src/sinks/prometheus/remote_write/tests.rs
index cffcb3aef502d..36ae905c4198a 100644
--- a/src/sinks/prometheus/remote_write/tests.rs
+++ b/src/sinks/prometheus/remote_write/tests.rs
@@ -167,11 +167,11 @@ async fn retains_state_between_requests() {
 #[tokio::test]
 async fn aggregates_batches() {
     let outputs = send_request(
-        r#"batch.max_events = 2"#,
+        r#"batch.max_events = 3"#,
         vec![
             create_inc_event("counter-1".into(), 12.0),
-            create_inc_event("counter-2".into(), 13.0),
             create_inc_event("counter-1".into(), 14.0),
+            create_inc_event("counter-2".into(), 13.0),
             create_inc_event("counter-2".into(), 14.0),
         ],
     )
diff --git a/src/sinks/vector/sink.rs b/src/sinks/vector/sink.rs
index fe77de6837883..a2fae74fc4a94 100644
--- a/src/sinks/vector/sink.rs
+++ b/src/sinks/vector/sink.rs
@@ -7,7 +7,7 @@ use tower::Service;
 use vector_common::request_metadata::GroupedCountByteSize;
 use vector_core::{
     config::telemetry,
-    stream::{BatcherSettings, DriverResponse},
+    stream::{batcher::data::BatchReduce, BatcherSettings, DriverResponse},
     ByteSizeOf, EstimatedJsonEncodedSizeOf,
 };
 
@@ -71,14 +71,14 @@ where
                     wrapper: EventWrapper::from(event),
                 }
             })
-            .batched(self.batch_settings.into_reducer_config(
+            .batched(self.batch_settings.as_reducer_config(
                 |data: &EventData| data.wrapper.encoded_len(),
-                |event_collection: &mut EventCollection, item: EventData| {
+                BatchReduce::new(|event_collection: &mut EventCollection, item: EventData| {
                     event_collection.finalizers.merge(item.finalizers);
                     event_collection.events.push(item.wrapper);
                     event_collection.events_byte_size += item.byte_size;
                     event_collection.events_json_byte_size += item.json_byte_size;
-                },
+                }),
             ))
             .map(|event_collection| {
                 let builder = RequestMetadataBuilder::new(

From ea600fac2e4a4eb13d6856553277c7b5a5faf6c8 Mon Sep 17 00:00:00 2001
From: Stephen Wakely <fungus.humungus@gmail.com>
Date: Mon, 16 Oct 2023 10:14:11 +0100
Subject: [PATCH 22/32] Make aggregating optional

Signed-off-by: Stephen Wakely <fungus.humungus@gmail.com>
---
 src/sinks/prometheus/remote_write/config.rs | 22 ++++++++-
 src/sinks/prometheus/remote_write/sink.rs   | 55 ++++++++++++++++++---
 src/sinks/util/buffer/metrics/normalize.rs  | 10 ++--
 src/sinks/util/builder.rs                   |  1 -
 4 files changed, 73 insertions(+), 15 deletions(-)

diff --git a/src/sinks/prometheus/remote_write/config.rs b/src/sinks/prometheus/remote_write/config.rs
index 922adc0962783..b75edda87377b 100644
--- a/src/sinks/prometheus/remote_write/config.rs
+++ b/src/sinks/prometheus/remote_write/config.rs
@@ -19,6 +19,19 @@ use super::{
 #[cfg(feature = "aws-core")]
 use super::Errors;
 
+/// The batch config for remote write.
+#[configurable_component]
+#[derive(Clone, Copy, Debug, Default)]
+pub struct RemoteWriteBatchConfig {
+    #[configurable(derived)]
+    #[serde(flatten)]
+    pub batch_settings: BatchConfig<PrometheusRemoteWriteDefaultBatchSettings>,
+
+    /// Set to true if incremental metrics within the batch should be aggregated.
+    #[serde(default = "crate::serde::default_true")]
+    pub aggregate: bool,
+}
+
 /// Configuration for the `prometheus_remote_write` sink.
 #[configurable_component(sink(
     "prometheus_remote_write",
@@ -61,7 +74,7 @@ pub struct RemoteWriteConfig {
 
     #[configurable(derived)]
     #[serde(default)]
-    pub batch: BatchConfig<PrometheusRemoteWriteDefaultBatchSettings>,
+    pub batch: RemoteWriteBatchConfig,
 
     #[configurable(derived)]
     #[serde(default)]
@@ -170,7 +183,12 @@ impl SinkConfig for RemoteWriteConfig {
         let sink = RemoteWriteSink {
             tenant_id: self.tenant_id.clone(),
             compression: self.compression,
-            batch_settings: self.batch.validate()?.into_batcher_settings()?,
+            aggregate: self.batch.aggregate,
+            batch_settings: self
+                .batch
+                .batch_settings
+                .validate()?
+                .into_batcher_settings()?,
             buckets,
             quantiles,
             default_namespace,
diff --git a/src/sinks/prometheus/remote_write/sink.rs b/src/sinks/prometheus/remote_write/sink.rs
index a6116667e8edd..2386497cd3a4f 100644
--- a/src/sinks/prometheus/remote_write/sink.rs
+++ b/src/sinks/prometheus/remote_write/sink.rs
@@ -64,23 +64,60 @@ impl SinkBatchSettings for PrometheusRemoteWriteDefaultBatchSettings {
     const TIMEOUT_SECS: f64 = 1.0;
 }
 
-#[derive(Clone)]
+pub(super) enum BatchedMetrics {
+    Aggregated(MetricSet),
+    All(Vec<Metric>),
+}
+
+impl BatchedMetrics {
+    pub(super) fn into_metrics(self) -> Vec<Metric> {
+        match self {
+            BatchedMetrics::Aggregated(metrics) => metrics.into_metrics(),
+            BatchedMetrics::All(metrics) => metrics,
+        }
+    }
+
+    pub(super) fn insert_update(&mut self, metric: Metric) {
+        match self {
+            BatchedMetrics::Aggregated(metrics) => metrics.insert_update(metric),
+            BatchedMetrics::All(metrics) => metrics.push(metric),
+        }
+    }
+
+    pub(super) fn len(&self) -> usize {
+        match self {
+            BatchedMetrics::Aggregated(metrics) => metrics.len(),
+            BatchedMetrics::All(metrics) => metrics.len(),
+        }
+    }
+}
+
 pub(super) struct EventCollection {
     pub(super) finalizers: EventFinalizers,
-    pub(super) events: MetricSet,
+    pub(super) events: BatchedMetrics,
     pub(super) events_byte_size: usize,
     pub(super) events_json_byte_size: GroupedCountByteSize,
 }
 
-impl Default for EventCollection {
-    fn default() -> Self {
+impl EventCollection {
+    /// Creates a new event collection that will either aggregate the incremental metrics
+    /// or store all the metrics, depending on the value of the `aggregate` parameter.
+    fn new(aggregate: bool) -> Self {
         Self {
             finalizers: Default::default(),
-            events: Default::default(),
+            events: if aggregate {
+                BatchedMetrics::Aggregated(Default::default())
+            } else {
+                BatchedMetrics::All(Default::default())
+            },
             events_byte_size: Default::default(),
             events_json_byte_size: telemetry().create_request_count_byte_size(),
         }
     }
+
+    fn is_aggregated(&self) -> bool {
+        matches!(self.events, BatchedMetrics::Aggregated(_))
+    }
 }
 
 impl BatchData<RemoteWriteMetric> for EventCollection {
@@ -91,7 +128,9 @@ impl BatchData<RemoteWriteMetric> for EventCollection {
     }
 
     fn take_batch(&mut self) -> Self::Batch {
-        std::mem::take(self)
+        let mut new = Self::new(self.is_aggregated());
+        std::mem::swap(self, &mut new);
+        new
     }
 
     fn push_item(&mut self, mut item: RemoteWriteMetric) {
@@ -107,6 +146,7 @@ impl BatchData<RemoteWriteMetric> for EventCollection {
 pub(super) struct RemoteWriteSink<S> {
     pub(super) tenant_id: Option<Template>,
     pub(super) batch_settings: BatcherSettings,
+    pub(super) aggregate: bool,
     pub(super) compression: super::Compression,
     pub(super) default_namespace: Option<String>,
     pub(super) buckets: Vec<f64>,
@@ -142,7 +182,8 @@ where
                 future::ready(make_remote_write_event(tenant_id.as_ref(), event))
             })
             .batched_partitioned(PrometheusTenantIdPartitioner, || {
-                batch_settings.as_reducer_config(ByteSizeOfItemSize, EventCollection::default())
+                batch_settings
+                    .as_reducer_config(ByteSizeOfItemSize, EventCollection::new(self.aggregate))
             })
             .request_builder(default_request_builder_concurrency_limit(), request_builder)
             .filter_map(|request| async move {
diff --git a/src/sinks/util/buffer/metrics/normalize.rs b/src/sinks/util/buffer/metrics/normalize.rs
index 1ada2126098e6..16812fcd26333 100644
--- a/src/sinks/util/buffer/metrics/normalize.rs
+++ b/src/sinks/util/buffer/metrics/normalize.rs
@@ -1,5 +1,4 @@
-use std::collections::HashMap;
-
+use indexmap::IndexMap;
 use vector_core::event::{
     metric::{MetricData, MetricSeries},
     EventMetadata, Metric, MetricKind,
@@ -83,10 +82,11 @@ type MetricEntry = (MetricData, EventMetadata);
 
 /// Metric storage for use with normalization.
 ///
-/// This is primarily a wrapper around `HashMap` with convenience methods to make it easier to perform
+/// This is primarily a wrapper around [`IndexMap`] (to ensure insertion order
+/// is maintained) with convenience methods to make it easier to perform
 /// normalization-specific operations.
 #[derive(Clone, Default, Debug)]
-pub struct MetricSet(HashMap<MetricSeries, MetricEntry>);
+pub struct MetricSet(IndexMap<MetricSeries, MetricEntry>);
 
 impl MetricSet {
     /// Creates an empty `MetricSet` with the specified capacity.
@@ -94,7 +94,7 @@ impl MetricSet {
     /// The metric set will be able to hold at least `capacity` elements without reallocating. If `capacity` is 0, the
     /// metric set will not allocate.
     pub fn with_capacity(capacity: usize) -> Self {
-        Self(HashMap::with_capacity(capacity))
+        Self(IndexMap::with_capacity(capacity))
     }
 
     /// Returns the number of elements in the set.
diff --git a/src/sinks/util/builder.rs b/src/sinks/util/builder.rs
index a36eb1e79493d..d30c8d0a7bd87 100644
--- a/src/sinks/util/builder.rs
+++ b/src/sinks/util/builder.rs
@@ -56,7 +56,6 @@ pub trait SinkBuilderExt: Stream {
         P::Item: ByteSizeOf,
         C: BatchConfig<P::Item>,
         F: Fn() -> C + Send,
-        // B: BatchOutput<P::Item>,
     {
         PartitionedBatcher::new(self, partitioner, settings)
     }

From 36ef82c00ece4659bc5660cb1b1964c661f5c207 Mon Sep 17 00:00:00 2001
From: Stephen Wakely <fungus.humungus@gmail.com>
Date: Mon, 16 Oct 2023 13:51:42 +0100
Subject: [PATCH 23/32] Adde test for non aggregation

Signed-off-by: Stephen Wakely <fungus.humungus@gmail.com>
---
 src/sinks/prometheus/remote_write/sink.rs  |  2 +-
 src/sinks/prometheus/remote_write/tests.rs | 39 ++++++++++++++++++++++
 2 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/src/sinks/prometheus/remote_write/sink.rs b/src/sinks/prometheus/remote_write/sink.rs
index 2386497cd3a4f..34f6838aa09e1 100644
--- a/src/sinks/prometheus/remote_write/sink.rs
+++ b/src/sinks/prometheus/remote_write/sink.rs
@@ -115,7 +115,7 @@ impl EventCollection {
         }
     }
 
-    fn is_aggregated(&self) -> bool {
+    const fn is_aggregated(&self) -> bool {
         matches!(self.events, BatchedMetrics::Aggregated(_))
     }
 }
diff --git a/src/sinks/prometheus/remote_write/tests.rs b/src/sinks/prometheus/remote_write/tests.rs
index 36ae905c4198a..190d3723bce6d 100644
--- a/src/sinks/prometheus/remote_write/tests.rs
+++ b/src/sinks/prometheus/remote_write/tests.rs
@@ -190,6 +190,45 @@ async fn aggregates_batches() {
     assert_eq!(req.timeseries[1].samples[0].value, 27.0);
 }
 
+#[tokio::test]
+async fn doesnt_aggregate_batches() {
+    let outputs = send_request(
+        indoc! {
+            r#"
+            batch.max_events = 3
+            batch.aggregate = false
+            "#
+        },
+        vec![
+            create_inc_event("counter-1".into(), 12.0),
+            create_inc_event("counter-1".into(), 14.0),
+            create_inc_event("counter-2".into(), 13.0),
+            create_inc_event("counter-2".into(), 14.0),
+        ],
+    )
+    .await;
+
+    assert_eq!(outputs.len(), 2);
+
+    // The first three metrics are in the first batch.
+    let (_, req) = &outputs[0];
+    assert_eq!(req.timeseries.len(), 2);
+    assert_eq!(req.timeseries[0].labels, labels!("__name__" => "counter-1"));
+    assert_eq!(req.timeseries[0].samples.len(), 2);
+    assert_eq!(req.timeseries[0].samples[0].value, 12.0);
+    assert_eq!(req.timeseries[0].samples[1].value, 26.0);
+
+    assert_eq!(req.timeseries[1].labels, labels!("__name__" => "counter-2"));
+    assert_eq!(req.timeseries[1].samples.len(), 1);
+    assert_eq!(req.timeseries[1].samples[0].value, 13.0);
+
+    // The last metric is in the last batch.
+    let (_, req) = &outputs[1];
+    assert_eq!(req.timeseries[0].labels, labels!("__name__" => "counter-2"));
+    assert_eq!(req.timeseries[0].samples.len(), 1);
+    assert_eq!(req.timeseries[0].samples[0].value, 27.0);
+}
+
 async fn send_request(config: &str, events: Vec<Event>) -> Vec<(HeaderMap, proto::WriteRequest)> {
     assert_sink_compliance(&HTTP_SINK_TAGS, async {
         let addr = test_util::next_addr();

From 7dec12de31805e7818e5ca4788bfb505085da1bd Mon Sep 17 00:00:00 2001
From: Stephen Wakely <fungus.humungus@gmail.com>
Date: Mon, 16 Oct 2023 15:25:17 +0100
Subject: [PATCH 24/32] Clippy

Signed-off-by: Stephen Wakely <fungus.humungus@gmail.com>
---
 src/sinks/gcp/stackdriver/metrics/sink.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/sinks/gcp/stackdriver/metrics/sink.rs b/src/sinks/gcp/stackdriver/metrics/sink.rs
index 538a666bdd91d..91ba22a640824 100644
--- a/src/sinks/gcp/stackdriver/metrics/sink.rs
+++ b/src/sinks/gcp/stackdriver/metrics/sink.rs
@@ -66,7 +66,7 @@ where
                 })
             })
             .normalized_with_default::<StackdriverMetricsNormalize>()
-            .batched(self.batch_settings.into_byte_size_config())
+            .batched(self.batch_settings.as_byte_size_config())
             .request_builder(
                 default_request_builder_concurrency_limit(),
                 self.request_builder,

From 3f24424bfa27aba1110c2b2a64db1146e50e0d5b Mon Sep 17 00:00:00 2001
From: Stephen Wakely <fungus.humungus@gmail.com>
Date: Tue, 17 Oct 2023 11:29:20 +0100
Subject: [PATCH 25/32] Component docs

Signed-off-by: Stephen Wakely <fungus.humungus@gmail.com>
---
 .../components/sinks/base/prometheus_remote_write.cue      | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/website/cue/reference/components/sinks/base/prometheus_remote_write.cue b/website/cue/reference/components/sinks/base/prometheus_remote_write.cue
index 4f825fdac5752..e7ac30201089a 100644
--- a/website/cue/reference/components/sinks/base/prometheus_remote_write.cue
+++ b/website/cue/reference/components/sinks/base/prometheus_remote_write.cue
@@ -190,9 +190,14 @@ base: components: sinks: prometheus_remote_write: configuration: {
 		}
 	}
 	batch: {
-		description: "Event batching behavior."
+		description: "The batch config for remote write."
 		required:    false
 		type: object: options: {
+			aggregate: {
+				description: "Set to true if incremental metrics within the batch should be aggregated."
+				required:    false
+				type: bool: default: false
+			}
 			max_bytes: {
 				description: """
 					The maximum size of a batch that is processed by a sink.

From faf8356495bead3ab934b7c18a553af82ecb9fa3 Mon Sep 17 00:00:00 2001
From: Stephen Wakely <fungus.humungus@gmail.com>
Date: Wed, 18 Oct 2023 14:27:28 +0100
Subject: [PATCH 26/32] Feedback from Kyle and Doug

Signed-off-by: Stephen Wakely <fungus.humungus@gmail.com>
---
 src/sinks/prometheus/remote_write/config.rs   | 15 ++++++--
 src/sinks/prometheus/remote_write/mod.rs      |  2 -
 src/sinks/prometheus/remote_write/service.rs  | 38 +++++++++++--------
 src/sinks/prometheus/remote_write/sink.rs     | 22 +++++------
 .../sinks/base/prometheus_remote_write.cue    |  2 +-
 5 files changed, 47 insertions(+), 32 deletions(-)

diff --git a/src/sinks/prometheus/remote_write/config.rs b/src/sinks/prometheus/remote_write/config.rs
index b75edda87377b..b40db5cdf2b32 100644
--- a/src/sinks/prometheus/remote_write/config.rs
+++ b/src/sinks/prometheus/remote_write/config.rs
@@ -21,17 +21,26 @@ use super::Errors;
 
 /// The batch config for remote write.
 #[configurable_component]
-#[derive(Clone, Copy, Debug, Default)]
+#[derive(Clone, Copy, Debug)]
 pub struct RemoteWriteBatchConfig {
     #[configurable(derived)]
     #[serde(flatten)]
     pub batch_settings: BatchConfig<PrometheusRemoteWriteDefaultBatchSettings>,
 
-    /// Set to true if incremental metrics within the batch should be aggregated.
+    /// Whether or not to aggregate metrics within a batch.
     #[serde(default = "crate::serde::default_true")]
     pub aggregate: bool,
 }
 
+impl Default for RemoteWriteBatchConfig {
+    fn default() -> Self {
+        Self {
+            batch_settings: Default::default(),
+            aggregate: true,
+        }
+    }
+}
+
 /// Configuration for the `prometheus_remote_write` sink.
 #[configurable_component(sink(
     "prometheus_remote_write",
@@ -212,7 +221,7 @@ async fn healthcheck(
     let body = bytes::Bytes::new();
     let request =
         build_request(http::Method::GET, &endpoint, compression, body, None, auth).await?;
-    let response = client.send(request.map(hyper::Body::from)).await?;
+    let response = client.send(request).await?;
 
     match response.status() {
         http::StatusCode::OK => Ok(()),
diff --git a/src/sinks/prometheus/remote_write/mod.rs b/src/sinks/prometheus/remote_write/mod.rs
index 32659aefef71a..3c1ce7c90e11d 100644
--- a/src/sinks/prometheus/remote_write/mod.rs
+++ b/src/sinks/prometheus/remote_write/mod.rs
@@ -55,8 +55,6 @@ impl From<Compression> for crate::sinks::util::Compression {
 
 #[derive(Debug, Snafu)]
 enum Errors {
-    #[snafu(display(r#"Prometheus remote_write sink cannot accept "set" metrics"#))]
-    SetMetricInvalid,
     #[cfg(feature = "aws-core")]
     #[snafu(display("aws.region required when AWS authentication is in use"))]
     AwsRegionRequired,
diff --git a/src/sinks/prometheus/remote_write/service.rs b/src/sinks/prometheus/remote_write/service.rs
index 2410059f39c82..d4360fd466f6e 100644
--- a/src/sinks/prometheus/remote_write/service.rs
+++ b/src/sinks/prometheus/remote_write/service.rs
@@ -18,6 +18,17 @@ use crate::{
     },
 };
 
+/// Constants for header strings.
+mod headers {
+    pub(super) const X_PROMETHEUS_REMOTE_WRITE_VERSION: &str = "X-Prometheus-Remote-Write-Version";
+    pub(super) const CONTENT_ENCODING: &str = "Content-Encoding";
+    pub(super) const CONTENT_TYPE: &str = "Content-Type";
+    pub(super) const X_SCOPE_ORGID: &str = "X-Scope-OrgID";
+
+    pub(super) const VERSION: &str = "0.1.0";
+    pub(super) const APPLICATION_X_PROTOBUF: &str = "application/x-protobuf";
+}
+
 #[derive(Clone)]
 pub(super) struct RemoteWriteService {
     pub(super) endpoint: Uri,
@@ -38,14 +49,16 @@ impl Service<RemoteWriteRequest> for RemoteWriteService {
     // Emission of internal events for errors and dropped events is handled upstream by the caller.
     fn call(&mut self, mut request: RemoteWriteRequest) -> Self::Future {
         let client = self.client.clone();
-        let metadata = std::mem::take(request.metadata_mut());
-        let json_size = metadata.into_events_estimated_json_encoded_byte_size();
         let endpoint = self.endpoint.clone();
         let auth = self.auth.clone();
         let compression = self.compression;
 
         Box::pin(async move {
-            let request = build_request(
+            let metadata = std::mem::take(request.metadata_mut());
+            let json_size = metadata.into_events_estimated_json_encoded_byte_size();
+            let raw_byte_size = request.request.len();
+
+            let http_request = build_request(
                 http::Method::POST,
                 &endpoint,
                 compression,
@@ -55,14 +68,9 @@ impl Service<RemoteWriteRequest> for RemoteWriteService {
             )
             .await?;
 
-            let (parts, body) = request.into_parts();
-            let request: hyper::Request<hyper::Body> =
-                hyper::Request::from_parts(parts, body.into());
-
-            let response = client.send(request.map(hyper::Body::from)).await?;
+            let response = client.send(http_request).await?;
             let (parts, body) = response.into_parts();
             let body = hyper::body::to_bytes(body).await?;
-            let raw_byte_size = body.len();
             let http_response = hyper::Response::from_parts(parts, body);
 
             if http_response.status().is_success() {
@@ -99,18 +107,18 @@ pub(super) async fn build_request(
     body: Bytes,
     tenant_id: Option<&String>,
     auth: Option<Auth>,
-) -> crate::Result<http::Request<Bytes>> {
+) -> crate::Result<http::Request<hyper::Body>> {
     let content_encoding = convert_compression_to_content_encoding(compression);
 
     let mut builder = http::Request::builder()
         .method(method)
         .uri(endpoint)
-        .header("X-Prometheus-Remote-Write-Version", "0.1.0")
-        .header("Content-Encoding", content_encoding)
-        .header("Content-Type", "application/x-protobuf");
+        .header(headers::X_PROMETHEUS_REMOTE_WRITE_VERSION, headers::VERSION)
+        .header(headers::CONTENT_ENCODING, content_encoding)
+        .header(headers::CONTENT_TYPE, headers::APPLICATION_X_PROTOBUF);
 
     if let Some(tenant_id) = tenant_id {
-        builder = builder.header("X-Scope-OrgID", tenant_id);
+        builder = builder.header(headers::X_SCOPE_ORGID, tenant_id);
     }
 
     let mut request = builder.body(body)?;
@@ -126,7 +134,7 @@ pub(super) async fn build_request(
         }
     }
 
-    Ok(request)
+    Ok(request.map(hyper::Body::from))
 }
 
 const fn convert_compression_to_content_encoding(compression: super::Compression) -> &'static str {
diff --git a/src/sinks/prometheus/remote_write/sink.rs b/src/sinks/prometheus/remote_write/sink.rs
index 34f6838aa09e1..a642a60b29a99 100644
--- a/src/sinks/prometheus/remote_write/sink.rs
+++ b/src/sinks/prometheus/remote_write/sink.rs
@@ -45,6 +45,12 @@ impl ByteSizeOf for RemoteWriteMetric {
 #[derive(Clone, Copy, Debug, Default)]
 pub struct PrometheusRemoteWriteDefaultBatchSettings;
 
+impl SinkBatchSettings for PrometheusRemoteWriteDefaultBatchSettings {
+    const MAX_EVENTS: Option<usize> = Some(1_000);
+    const MAX_BYTES: Option<usize> = None;
+    const TIMEOUT_SECS: f64 = 1.0;
+}
+
 pub(super) struct PrometheusTenantIdPartitioner;
 
 impl Partitioner for PrometheusTenantIdPartitioner {
@@ -58,36 +64,30 @@ impl Partitioner for PrometheusTenantIdPartitioner {
     }
 }
 
-impl SinkBatchSettings for PrometheusRemoteWriteDefaultBatchSettings {
-    const MAX_EVENTS: Option<usize> = Some(1_000);
-    const MAX_BYTES: Option<usize> = None;
-    const TIMEOUT_SECS: f64 = 1.0;
-}
-
 pub(super) enum BatchedMetrics {
     Aggregated(MetricSet),
-    All(Vec<Metric>),
+    Unaggregated(Vec<Metric>),
 }
 
 impl BatchedMetrics {
     pub(super) fn into_metrics(self) -> Vec<Metric> {
         match self {
             BatchedMetrics::Aggregated(metrics) => metrics.into_metrics(),
-            BatchedMetrics::All(metrics) => metrics,
+            BatchedMetrics::Unaggregated(metrics) => metrics,
         }
     }
 
     pub(super) fn insert_update(&mut self, metric: Metric) {
         match self {
             BatchedMetrics::Aggregated(metrics) => metrics.insert_update(metric),
-            BatchedMetrics::All(metrics) => metrics.push(metric),
+            BatchedMetrics::Unaggregated(metrics) => metrics.push(metric),
         }
     }
 
     pub(super) fn len(&self) -> usize {
         match self {
             BatchedMetrics::Aggregated(metrics) => metrics.len(),
-            BatchedMetrics::All(metrics) => metrics.len(),
+            BatchedMetrics::Unaggregated(metrics) => metrics.len(),
         }
     }
 }
@@ -108,7 +108,7 @@ impl EventCollection {
             events: if aggregate {
                 BatchedMetrics::Aggregated(Default::default())
             } else {
-                BatchedMetrics::All(Default::default())
+                BatchedMetrics::Unaggregated(Default::default())
             },
             events_byte_size: Default::default(),
             events_json_byte_size: telemetry().create_request_count_byte_size(),
diff --git a/website/cue/reference/components/sinks/base/prometheus_remote_write.cue b/website/cue/reference/components/sinks/base/prometheus_remote_write.cue
index e7ac30201089a..e6273898f319f 100644
--- a/website/cue/reference/components/sinks/base/prometheus_remote_write.cue
+++ b/website/cue/reference/components/sinks/base/prometheus_remote_write.cue
@@ -196,7 +196,7 @@ base: components: sinks: prometheus_remote_write: configuration: {
 			aggregate: {
 				description: "Set to true if incremental metrics within the batch should be aggregated."
 				required:    false
-				type: bool: default: false
+				type: bool: default: true
 			}
 			max_bytes: {
 				description: """

From c4d3ee304ac0f5462175d31983a647f11c78c070 Mon Sep 17 00:00:00 2001
From: Stephen Wakely <fungus.humungus@gmail.com>
Date: Wed, 18 Oct 2023 14:57:20 +0100
Subject: [PATCH 27/32] Use generic compression options

Signed-off-by: Stephen Wakely <fungus.humungus@gmail.com>
---
 src/sinks/prometheus/remote_write/config.rs   |  4 +--
 src/sinks/prometheus/remote_write/mod.rs      | 27 ----------------
 .../remote_write/request_builder.rs           |  4 +--
 src/sinks/prometheus/remote_write/service.rs  | 17 +++-------
 .../sinks/base/prometheus_remote_write.cue    | 31 ++++++++++++++-----
 5 files changed, 33 insertions(+), 50 deletions(-)

diff --git a/src/sinks/prometheus/remote_write/config.rs b/src/sinks/prometheus/remote_write/config.rs
index b40db5cdf2b32..b8171bfaa4b51 100644
--- a/src/sinks/prometheus/remote_write/config.rs
+++ b/src/sinks/prometheus/remote_write/config.rs
@@ -121,7 +121,7 @@ pub struct RemoteWriteConfig {
     #[configurable(derived)]
     #[configurable(metadata(docs::advanced))]
     #[serde(default)]
-    pub compression: super::Compression,
+    pub compression: Compression,
 }
 
 impl_generate_config_from_default!(RemoteWriteConfig);
@@ -215,7 +215,7 @@ impl SinkConfig for RemoteWriteConfig {
 async fn healthcheck(
     client: HttpClient,
     endpoint: Uri,
-    compression: super::Compression,
+    compression: Compression,
     auth: Option<Auth>,
 ) -> crate::Result<()> {
     let body = bytes::Bytes::new();
diff --git a/src/sinks/prometheus/remote_write/mod.rs b/src/sinks/prometheus/remote_write/mod.rs
index 3c1ce7c90e11d..3442c88de812f 100644
--- a/src/sinks/prometheus/remote_write/mod.rs
+++ b/src/sinks/prometheus/remote_write/mod.rs
@@ -26,33 +26,6 @@ mod integration_tests;
 
 pub use config::RemoteWriteConfig;
 
-/// Supported compression types for Prometheus Remote Write.
-#[configurable_component]
-#[derive(Clone, Copy, Debug, Derivative)]
-#[derivative(Default)]
-#[serde(rename_all = "lowercase")]
-pub enum Compression {
-    /// Snappy.
-    #[derivative(Default)]
-    Snappy,
-
-    /// Gzip.
-    Gzip,
-
-    /// Zstandard.
-    Zstd,
-}
-
-impl From<Compression> for crate::sinks::util::Compression {
-    fn from(value: Compression) -> Self {
-        match value {
-            Compression::Snappy => crate::sinks::util::Compression::Snappy,
-            Compression::Gzip => crate::sinks::util::Compression::Gzip(Default::default()),
-            Compression::Zstd => crate::sinks::util::Compression::Zstd(Default::default()),
-        }
-    }
-}
-
 #[derive(Debug, Snafu)]
 enum Errors {
     #[cfg(feature = "aws-core")]
diff --git a/src/sinks/prometheus/remote_write/request_builder.rs b/src/sinks/prometheus/remote_write/request_builder.rs
index 8161dc2138535..a2e156c42cfb8 100644
--- a/src/sinks/prometheus/remote_write/request_builder.rs
+++ b/src/sinks/prometheus/remote_write/request_builder.rs
@@ -79,7 +79,7 @@ pub(super) struct RemoteWriteMetadata {
 }
 
 pub(super) struct RemoteWriteRequestBuilder {
-    pub(super) compression: super::Compression,
+    pub(super) compression: Compression,
     pub(super) encoder: RemoteWriteEncoder,
 }
 
@@ -92,7 +92,7 @@ impl RequestBuilder<(PartitionKey, EventCollection)> for RemoteWriteRequestBuild
     type Error = io::Error;
 
     fn compression(&self) -> Compression {
-        self.compression.into()
+        self.compression
     }
 
     fn encoder(&self) -> &Self::Encoder {
diff --git a/src/sinks/prometheus/remote_write/service.rs b/src/sinks/prometheus/remote_write/service.rs
index d4360fd466f6e..fb459d7c4bd8e 100644
--- a/src/sinks/prometheus/remote_write/service.rs
+++ b/src/sinks/prometheus/remote_write/service.rs
@@ -103,20 +103,21 @@ async fn sign_request(
 pub(super) async fn build_request(
     method: http::Method,
     endpoint: &Uri,
-    compression: super::Compression,
+    compression: Compression,
     body: Bytes,
     tenant_id: Option<&String>,
     auth: Option<Auth>,
 ) -> crate::Result<http::Request<hyper::Body>> {
-    let content_encoding = convert_compression_to_content_encoding(compression);
-
     let mut builder = http::Request::builder()
         .method(method)
         .uri(endpoint)
         .header(headers::X_PROMETHEUS_REMOTE_WRITE_VERSION, headers::VERSION)
-        .header(headers::CONTENT_ENCODING, content_encoding)
         .header(headers::CONTENT_TYPE, headers::APPLICATION_X_PROTOBUF);
 
+    if let Some(content_encoding) = compression.content_encoding() {
+        builder = builder.header(headers::CONTENT_ENCODING, content_encoding);
+    }
+
     if let Some(tenant_id) = tenant_id {
         builder = builder.header(headers::X_SCOPE_ORGID, tenant_id);
     }
@@ -136,11 +137,3 @@ pub(super) async fn build_request(
 
     Ok(request.map(hyper::Body::from))
 }
-
-const fn convert_compression_to_content_encoding(compression: super::Compression) -> &'static str {
-    match compression {
-        super::Compression::Snappy => "snappy",
-        super::Compression::Gzip => "gzip",
-        super::Compression::Zstd => "zstd",
-    }
-}
diff --git a/website/cue/reference/components/sinks/base/prometheus_remote_write.cue b/website/cue/reference/components/sinks/base/prometheus_remote_write.cue
index e6273898f319f..2749ac28b6a35 100644
--- a/website/cue/reference/components/sinks/base/prometheus_remote_write.cue
+++ b/website/cue/reference/components/sinks/base/prometheus_remote_write.cue
@@ -194,7 +194,7 @@ base: components: sinks: prometheus_remote_write: configuration: {
 		required:    false
 		type: object: options: {
 			aggregate: {
-				description: "Set to true if incremental metrics within the batch should be aggregated."
+				description: "Whether or not to aggregate metrics within a batch."
 				required:    false
 				type: bool: default: true
 			}
@@ -239,14 +239,31 @@ base: components: sinks: prometheus_remote_write: configuration: {
 		}
 	}
 	compression: {
-		description: "Supported compression types for Prometheus Remote Write."
-		required:    false
+		description: """
+			Compression configuration.
+
+			All compression algorithms use the default compression level unless otherwise specified.
+			"""
+		required: false
 		type: string: {
-			default: "snappy"
+			default: "none"
 			enum: {
-				gzip:   "Gzip."
-				snappy: "Snappy."
-				zstd:   "Zstandard."
+				gzip: """
+					[Gzip][gzip] compression.
+
+					[gzip]: https://www.gzip.org/
+					"""
+				none: "No compression."
+				zlib: """
+					[Zlib][zlib] compression.
+
+					[zlib]: https://zlib.net/
+					"""
+				zstd: """
+					[Zstandard][zstd] compression.
+
+					[zstd]: https://facebook.github.io/zstd/
+					"""
 			}
 		}
 	}

From 3f54b76ac7f279ebb2015b23a80d1d1a01871705 Mon Sep 17 00:00:00 2001
From: Stephen Wakely <fungus.humungus@gmail.com>
Date: Thu, 19 Oct 2023 10:38:03 +0100
Subject: [PATCH 28/32] Default compression to Snappy

Signed-off-by: Stephen Wakely <fungus.humungus@gmail.com>
---
 src/sinks/prometheus/remote_write/config.rs | 23 ++++++++++-----------
 1 file changed, 11 insertions(+), 12 deletions(-)

diff --git a/src/sinks/prometheus/remote_write/config.rs b/src/sinks/prometheus/remote_write/config.rs
index b8171bfaa4b51..8f9f065642bdf 100644
--- a/src/sinks/prometheus/remote_write/config.rs
+++ b/src/sinks/prometheus/remote_write/config.rs
@@ -21,7 +21,8 @@ use super::Errors;
 
 /// The batch config for remote write.
 #[configurable_component]
-#[derive(Clone, Copy, Debug)]
+#[derive(Clone, Copy, Debug, Derivative)]
+#[derivative(Default)]
 pub struct RemoteWriteBatchConfig {
     #[configurable(derived)]
     #[serde(flatten)]
@@ -29,24 +30,17 @@ pub struct RemoteWriteBatchConfig {
 
     /// Whether or not to aggregate metrics within a batch.
     #[serde(default = "crate::serde::default_true")]
+    #[derivative(Default(value = "true"))]
     pub aggregate: bool,
 }
 
-impl Default for RemoteWriteBatchConfig {
-    fn default() -> Self {
-        Self {
-            batch_settings: Default::default(),
-            aggregate: true,
-        }
-    }
-}
-
 /// Configuration for the `prometheus_remote_write` sink.
 #[configurable_component(sink(
     "prometheus_remote_write",
     "Deliver metric data to a Prometheus remote write endpoint."
 ))]
-#[derive(Clone, Debug, Default)]
+#[derive(Clone, Debug, Derivative)]
+#[derivative(Default)]
 #[serde(deny_unknown_fields)]
 pub struct RemoteWriteConfig {
     /// The endpoint to send data to.
@@ -120,10 +114,15 @@ pub struct RemoteWriteConfig {
 
     #[configurable(derived)]
     #[configurable(metadata(docs::advanced))]
-    #[serde(default)]
+    #[serde(default = "default_compression")]
+    #[derivative(Default(value = "default_compression()"))]
     pub compression: Compression,
 }
 
+const fn default_compression() -> Compression {
+    Compression::Snappy
+}
+
 impl_generate_config_from_default!(RemoteWriteConfig);
 
 #[async_trait::async_trait]

From 17d46f2b027cba989fbd85d114ab348dad95aa88 Mon Sep 17 00:00:00 2001
From: Stephen Wakely <fungus.humungus@gmail.com>
Date: Thu, 19 Oct 2023 12:00:14 +0100
Subject: [PATCH 29/32] Update docs

Signed-off-by: Stephen Wakely <fungus.humungus@gmail.com>
---
 .../reference/components/sinks/base/prometheus_remote_write.cue | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/website/cue/reference/components/sinks/base/prometheus_remote_write.cue b/website/cue/reference/components/sinks/base/prometheus_remote_write.cue
index 2749ac28b6a35..ac45926ec9b45 100644
--- a/website/cue/reference/components/sinks/base/prometheus_remote_write.cue
+++ b/website/cue/reference/components/sinks/base/prometheus_remote_write.cue
@@ -246,7 +246,7 @@ base: components: sinks: prometheus_remote_write: configuration: {
 			"""
 		required: false
 		type: string: {
-			default: "none"
+			default: "snappy"
 			enum: {
 				gzip: """
 					[Gzip][gzip] compression.

From 411eb343353ab7aa5c54020389cd31645c829b3a Mon Sep 17 00:00:00 2001
From: Stephen Wakely <fungus.humungus@gmail.com>
Date: Fri, 20 Oct 2023 10:00:14 +0100
Subject: [PATCH 30/32] Add snappy to the compression docs

Signed-off-by: Stephen Wakely <fungus.humungus@gmail.com>
---
 src/sinks/util/buffer/compression.rs                       | 7 +++++++
 website/cue/reference/components/sinks/base/appsignal.cue  | 5 +++++
 .../components/sinks/base/aws_cloudwatch_logs.cue          | 5 +++++
 .../components/sinks/base/aws_cloudwatch_metrics.cue       | 5 +++++
 .../components/sinks/base/aws_kinesis_firehose.cue         | 5 +++++
 .../components/sinks/base/aws_kinesis_streams.cue          | 5 +++++
 website/cue/reference/components/sinks/base/aws_s3.cue     | 5 +++++
 website/cue/reference/components/sinks/base/axiom.cue      | 5 +++++
 website/cue/reference/components/sinks/base/azure_blob.cue | 5 +++++
 website/cue/reference/components/sinks/base/clickhouse.cue | 5 +++++
 .../cue/reference/components/sinks/base/datadog_logs.cue   | 5 +++++
 .../cue/reference/components/sinks/base/datadog_traces.cue | 5 +++++
 .../cue/reference/components/sinks/base/elasticsearch.cue  | 5 +++++
 .../reference/components/sinks/base/gcp_cloud_storage.cue  | 5 +++++
 website/cue/reference/components/sinks/base/http.cue       | 5 +++++
 website/cue/reference/components/sinks/base/humio_logs.cue | 5 +++++
 .../cue/reference/components/sinks/base/humio_metrics.cue  | 5 +++++
 website/cue/reference/components/sinks/base/new_relic.cue  | 5 +++++
 .../components/sinks/base/prometheus_remote_write.cue      | 5 +++++
 .../reference/components/sinks/base/splunk_hec_logs.cue    | 5 +++++
 .../reference/components/sinks/base/splunk_hec_metrics.cue | 5 +++++
 website/cue/reference/components/sinks/base/webhdfs.cue    | 5 +++++
 22 files changed, 112 insertions(+)

diff --git a/src/sinks/util/buffer/compression.rs b/src/sinks/util/buffer/compression.rs
index 47607f1bb4428..943232ad25be5 100644
--- a/src/sinks/util/buffer/compression.rs
+++ b/src/sinks/util/buffer/compression.rs
@@ -330,11 +330,18 @@ impl Configurable for Compression {
             "[zstd]: https://facebook.github.io/zstd/",
         );
 
+        let snappy_string_subschema = generate_string_schema(
+            "Snappy",
+            Some("[Snappy][snappy] compression."),
+            "[snappy]: https://github.com/google/snappy/blob/main/docs/README.md",
+        );
+
         let mut all_string_oneof_subschema = generate_one_of_schema(&[
             none_string_subschema,
             gzip_string_subschema,
             zlib_string_subschema,
             zstd_string_subschema,
+            snappy_string_subschema,
         ]);
         apply_base_metadata(&mut all_string_oneof_subschema, string_metadata);
 
diff --git a/website/cue/reference/components/sinks/base/appsignal.cue b/website/cue/reference/components/sinks/base/appsignal.cue
index 2eebeb475dc99..824204cc4d6bd 100644
--- a/website/cue/reference/components/sinks/base/appsignal.cue
+++ b/website/cue/reference/components/sinks/base/appsignal.cue
@@ -78,6 +78,11 @@ base: components: sinks: appsignal: configuration: {
 					[gzip]: https://www.gzip.org/
 					"""
 				none: "No compression."
+				snappy: """
+					[Snappy][snappy] compression.
+
+					[snappy]: https://github.com/google/snappy/blob/main/docs/README.md
+					"""
 				zlib: """
 					[Zlib][zlib] compression.
 
diff --git a/website/cue/reference/components/sinks/base/aws_cloudwatch_logs.cue b/website/cue/reference/components/sinks/base/aws_cloudwatch_logs.cue
index 68e8b1614b1e4..8af01b34e91d1 100644
--- a/website/cue/reference/components/sinks/base/aws_cloudwatch_logs.cue
+++ b/website/cue/reference/components/sinks/base/aws_cloudwatch_logs.cue
@@ -180,6 +180,11 @@ base: components: sinks: aws_cloudwatch_logs: configuration: {
 					[gzip]: https://www.gzip.org/
 					"""
 				none: "No compression."
+				snappy: """
+					[Snappy][snappy] compression.
+
+					[snappy]: https://github.com/google/snappy/blob/main/docs/README.md
+					"""
 				zlib: """
 					[Zlib][zlib] compression.
 
diff --git a/website/cue/reference/components/sinks/base/aws_cloudwatch_metrics.cue b/website/cue/reference/components/sinks/base/aws_cloudwatch_metrics.cue
index daa94b379ffda..16e90e50e5193 100644
--- a/website/cue/reference/components/sinks/base/aws_cloudwatch_metrics.cue
+++ b/website/cue/reference/components/sinks/base/aws_cloudwatch_metrics.cue
@@ -177,6 +177,11 @@ base: components: sinks: aws_cloudwatch_metrics: configuration: {
 					[gzip]: https://www.gzip.org/
 					"""
 				none: "No compression."
+				snappy: """
+					[Snappy][snappy] compression.
+
+					[snappy]: https://github.com/google/snappy/blob/main/docs/README.md
+					"""
 				zlib: """
 					[Zlib][zlib] compression.
 
diff --git a/website/cue/reference/components/sinks/base/aws_kinesis_firehose.cue b/website/cue/reference/components/sinks/base/aws_kinesis_firehose.cue
index f1a647e602fd5..870af3559cbac 100644
--- a/website/cue/reference/components/sinks/base/aws_kinesis_firehose.cue
+++ b/website/cue/reference/components/sinks/base/aws_kinesis_firehose.cue
@@ -180,6 +180,11 @@ base: components: sinks: aws_kinesis_firehose: configuration: {
 					[gzip]: https://www.gzip.org/
 					"""
 				none: "No compression."
+				snappy: """
+					[Snappy][snappy] compression.
+
+					[snappy]: https://github.com/google/snappy/blob/main/docs/README.md
+					"""
 				zlib: """
 					[Zlib][zlib] compression.
 
diff --git a/website/cue/reference/components/sinks/base/aws_kinesis_streams.cue b/website/cue/reference/components/sinks/base/aws_kinesis_streams.cue
index f723f9e8ad5b2..32c63f7478b31 100644
--- a/website/cue/reference/components/sinks/base/aws_kinesis_streams.cue
+++ b/website/cue/reference/components/sinks/base/aws_kinesis_streams.cue
@@ -180,6 +180,11 @@ base: components: sinks: aws_kinesis_streams: configuration: {
 					[gzip]: https://www.gzip.org/
 					"""
 				none: "No compression."
+				snappy: """
+					[Snappy][snappy] compression.
+
+					[snappy]: https://github.com/google/snappy/blob/main/docs/README.md
+					"""
 				zlib: """
 					[Zlib][zlib] compression.
 
diff --git a/website/cue/reference/components/sinks/base/aws_s3.cue b/website/cue/reference/components/sinks/base/aws_s3.cue
index 16b10bc9018bc..131fa69ce057c 100644
--- a/website/cue/reference/components/sinks/base/aws_s3.cue
+++ b/website/cue/reference/components/sinks/base/aws_s3.cue
@@ -264,6 +264,11 @@ base: components: sinks: aws_s3: configuration: {
 					[gzip]: https://www.gzip.org/
 					"""
 				none: "No compression."
+				snappy: """
+					[Snappy][snappy] compression.
+
+					[snappy]: https://github.com/google/snappy/blob/main/docs/README.md
+					"""
 				zlib: """
 					[Zlib][zlib] compression.
 
diff --git a/website/cue/reference/components/sinks/base/axiom.cue b/website/cue/reference/components/sinks/base/axiom.cue
index a20276b2ee26b..2de88265fc8f8 100644
--- a/website/cue/reference/components/sinks/base/axiom.cue
+++ b/website/cue/reference/components/sinks/base/axiom.cue
@@ -43,6 +43,11 @@ base: components: sinks: axiom: configuration: {
 					[gzip]: https://www.gzip.org/
 					"""
 				none: "No compression."
+				snappy: """
+					[Snappy][snappy] compression.
+
+					[snappy]: https://github.com/google/snappy/blob/main/docs/README.md
+					"""
 				zlib: """
 					[Zlib][zlib] compression.
 
diff --git a/website/cue/reference/components/sinks/base/azure_blob.cue b/website/cue/reference/components/sinks/base/azure_blob.cue
index a56680098c0fc..a0760192996bf 100644
--- a/website/cue/reference/components/sinks/base/azure_blob.cue
+++ b/website/cue/reference/components/sinks/base/azure_blob.cue
@@ -127,6 +127,11 @@ base: components: sinks: azure_blob: configuration: {
 					[gzip]: https://www.gzip.org/
 					"""
 				none: "No compression."
+				snappy: """
+					[Snappy][snappy] compression.
+
+					[snappy]: https://github.com/google/snappy/blob/main/docs/README.md
+					"""
 				zlib: """
 					[Zlib][zlib] compression.
 
diff --git a/website/cue/reference/components/sinks/base/clickhouse.cue b/website/cue/reference/components/sinks/base/clickhouse.cue
index 6be7c94fa2a6c..03c3c13bba14a 100644
--- a/website/cue/reference/components/sinks/base/clickhouse.cue
+++ b/website/cue/reference/components/sinks/base/clickhouse.cue
@@ -122,6 +122,11 @@ base: components: sinks: clickhouse: configuration: {
 					[gzip]: https://www.gzip.org/
 					"""
 				none: "No compression."
+				snappy: """
+					[Snappy][snappy] compression.
+
+					[snappy]: https://github.com/google/snappy/blob/main/docs/README.md
+					"""
 				zlib: """
 					[Zlib][zlib] compression.
 
diff --git a/website/cue/reference/components/sinks/base/datadog_logs.cue b/website/cue/reference/components/sinks/base/datadog_logs.cue
index 7de6e0a300989..da23c24d449c6 100644
--- a/website/cue/reference/components/sinks/base/datadog_logs.cue
+++ b/website/cue/reference/components/sinks/base/datadog_logs.cue
@@ -76,6 +76,11 @@ base: components: sinks: datadog_logs: configuration: {
 				[gzip]: https://www.gzip.org/
 				"""
 			none: "No compression."
+			snappy: """
+				[Snappy][snappy] compression.
+
+				[snappy]: https://github.com/google/snappy/blob/main/docs/README.md
+				"""
 			zlib: """
 				[Zlib][zlib] compression.
 
diff --git a/website/cue/reference/components/sinks/base/datadog_traces.cue b/website/cue/reference/components/sinks/base/datadog_traces.cue
index edf6ec6d65bc7..1564e8f720f2d 100644
--- a/website/cue/reference/components/sinks/base/datadog_traces.cue
+++ b/website/cue/reference/components/sinks/base/datadog_traces.cue
@@ -76,6 +76,11 @@ base: components: sinks: datadog_traces: configuration: {
 				[gzip]: https://www.gzip.org/
 				"""
 			none: "No compression."
+			snappy: """
+				[Snappy][snappy] compression.
+
+				[snappy]: https://github.com/google/snappy/blob/main/docs/README.md
+				"""
 			zlib: """
 				[Zlib][zlib] compression.
 
diff --git a/website/cue/reference/components/sinks/base/elasticsearch.cue b/website/cue/reference/components/sinks/base/elasticsearch.cue
index 98b9dbdeff3a6..a5c4fb951f76b 100644
--- a/website/cue/reference/components/sinks/base/elasticsearch.cue
+++ b/website/cue/reference/components/sinks/base/elasticsearch.cue
@@ -279,6 +279,11 @@ base: components: sinks: elasticsearch: configuration: {
 					[gzip]: https://www.gzip.org/
 					"""
 				none: "No compression."
+				snappy: """
+					[Snappy][snappy] compression.
+
+					[snappy]: https://github.com/google/snappy/blob/main/docs/README.md
+					"""
 				zlib: """
 					[Zlib][zlib] compression.
 
diff --git a/website/cue/reference/components/sinks/base/gcp_cloud_storage.cue b/website/cue/reference/components/sinks/base/gcp_cloud_storage.cue
index 47f6bf3b03df7..5caae5b7b0b52 100644
--- a/website/cue/reference/components/sinks/base/gcp_cloud_storage.cue
+++ b/website/cue/reference/components/sinks/base/gcp_cloud_storage.cue
@@ -151,6 +151,11 @@ base: components: sinks: gcp_cloud_storage: configuration: {
 					[gzip]: https://www.gzip.org/
 					"""
 				none: "No compression."
+				snappy: """
+					[Snappy][snappy] compression.
+
+					[snappy]: https://github.com/google/snappy/blob/main/docs/README.md
+					"""
 				zlib: """
 					[Zlib][zlib] compression.
 
diff --git a/website/cue/reference/components/sinks/base/http.cue b/website/cue/reference/components/sinks/base/http.cue
index 58fcd9c005c81..19069c666f090 100644
--- a/website/cue/reference/components/sinks/base/http.cue
+++ b/website/cue/reference/components/sinks/base/http.cue
@@ -122,6 +122,11 @@ base: components: sinks: http: configuration: {
 					[gzip]: https://www.gzip.org/
 					"""
 				none: "No compression."
+				snappy: """
+					[Snappy][snappy] compression.
+
+					[snappy]: https://github.com/google/snappy/blob/main/docs/README.md
+					"""
 				zlib: """
 					[Zlib][zlib] compression.
 
diff --git a/website/cue/reference/components/sinks/base/humio_logs.cue b/website/cue/reference/components/sinks/base/humio_logs.cue
index 2892e356c6784..8416b61927fd5 100644
--- a/website/cue/reference/components/sinks/base/humio_logs.cue
+++ b/website/cue/reference/components/sinks/base/humio_logs.cue
@@ -75,6 +75,11 @@ base: components: sinks: humio_logs: configuration: {
 					[gzip]: https://www.gzip.org/
 					"""
 				none: "No compression."
+				snappy: """
+					[Snappy][snappy] compression.
+
+					[snappy]: https://github.com/google/snappy/blob/main/docs/README.md
+					"""
 				zlib: """
 					[Zlib][zlib] compression.
 
diff --git a/website/cue/reference/components/sinks/base/humio_metrics.cue b/website/cue/reference/components/sinks/base/humio_metrics.cue
index 2d39e4dca0b90..c45509ce09b10 100644
--- a/website/cue/reference/components/sinks/base/humio_metrics.cue
+++ b/website/cue/reference/components/sinks/base/humio_metrics.cue
@@ -75,6 +75,11 @@ base: components: sinks: humio_metrics: configuration: {
 					[gzip]: https://www.gzip.org/
 					"""
 				none: "No compression."
+				snappy: """
+					[Snappy][snappy] compression.
+
+					[snappy]: https://github.com/google/snappy/blob/main/docs/README.md
+					"""
 				zlib: """
 					[Zlib][zlib] compression.
 
diff --git a/website/cue/reference/components/sinks/base/new_relic.cue b/website/cue/reference/components/sinks/base/new_relic.cue
index 1929466c31c80..ac289488930d2 100644
--- a/website/cue/reference/components/sinks/base/new_relic.cue
+++ b/website/cue/reference/components/sinks/base/new_relic.cue
@@ -92,6 +92,11 @@ base: components: sinks: new_relic: configuration: {
 					[gzip]: https://www.gzip.org/
 					"""
 				none: "No compression."
+				snappy: """
+					[Snappy][snappy] compression.
+
+					[snappy]: https://github.com/google/snappy/blob/main/docs/README.md
+					"""
 				zlib: """
 					[Zlib][zlib] compression.
 
diff --git a/website/cue/reference/components/sinks/base/prometheus_remote_write.cue b/website/cue/reference/components/sinks/base/prometheus_remote_write.cue
index ac45926ec9b45..8858c0f99ed74 100644
--- a/website/cue/reference/components/sinks/base/prometheus_remote_write.cue
+++ b/website/cue/reference/components/sinks/base/prometheus_remote_write.cue
@@ -254,6 +254,11 @@ base: components: sinks: prometheus_remote_write: configuration: {
 					[gzip]: https://www.gzip.org/
 					"""
 				none: "No compression."
+				snappy: """
+					[Snappy][snappy] compression.
+
+					[snappy]: https://github.com/google/snappy/blob/main/docs/README.md
+					"""
 				zlib: """
 					[Zlib][zlib] compression.
 
diff --git a/website/cue/reference/components/sinks/base/splunk_hec_logs.cue b/website/cue/reference/components/sinks/base/splunk_hec_logs.cue
index e79d4cdb46261..fb7cc67ed8b22 100644
--- a/website/cue/reference/components/sinks/base/splunk_hec_logs.cue
+++ b/website/cue/reference/components/sinks/base/splunk_hec_logs.cue
@@ -116,6 +116,11 @@ base: components: sinks: splunk_hec_logs: configuration: {
 					[gzip]: https://www.gzip.org/
 					"""
 				none: "No compression."
+				snappy: """
+					[Snappy][snappy] compression.
+
+					[snappy]: https://github.com/google/snappy/blob/main/docs/README.md
+					"""
 				zlib: """
 					[Zlib][zlib] compression.
 
diff --git a/website/cue/reference/components/sinks/base/splunk_hec_metrics.cue b/website/cue/reference/components/sinks/base/splunk_hec_metrics.cue
index 96c8517576f39..820454b8e867d 100644
--- a/website/cue/reference/components/sinks/base/splunk_hec_metrics.cue
+++ b/website/cue/reference/components/sinks/base/splunk_hec_metrics.cue
@@ -102,6 +102,11 @@ base: components: sinks: splunk_hec_metrics: configuration: {
 					[gzip]: https://www.gzip.org/
 					"""
 				none: "No compression."
+				snappy: """
+					[Snappy][snappy] compression.
+
+					[snappy]: https://github.com/google/snappy/blob/main/docs/README.md
+					"""
 				zlib: """
 					[Zlib][zlib] compression.
 
diff --git a/website/cue/reference/components/sinks/base/webhdfs.cue b/website/cue/reference/components/sinks/base/webhdfs.cue
index 0dd49a0d683b8..9339e23be45e4 100644
--- a/website/cue/reference/components/sinks/base/webhdfs.cue
+++ b/website/cue/reference/components/sinks/base/webhdfs.cue
@@ -75,6 +75,11 @@ base: components: sinks: webhdfs: configuration: {
 					[gzip]: https://www.gzip.org/
 					"""
 				none: "No compression."
+				snappy: """
+					[Snappy][snappy] compression.
+
+					[snappy]: https://github.com/google/snappy/blob/main/docs/README.md
+					"""
 				zlib: """
 					[Zlib][zlib] compression.
 

From 7e5146269e0140bfcbe670e61c803e8cfa7df0f5 Mon Sep 17 00:00:00 2001
From: Stephen Wakely <fungus.humungus@gmail.com>
Date: Fri, 20 Oct 2023 10:14:07 +0100
Subject: [PATCH 31/32] Remove proptest file

Signed-off-by: Stephen Wakely <fungus.humungus@gmail.com>
---
 .../proptest-regressions/stream/partitioned_batcher.txt    | 7 -------
 1 file changed, 7 deletions(-)
 delete mode 100644 lib/vector-core/proptest-regressions/stream/partitioned_batcher.txt

diff --git a/lib/vector-core/proptest-regressions/stream/partitioned_batcher.txt b/lib/vector-core/proptest-regressions/stream/partitioned_batcher.txt
deleted file mode 100644
index 507ce908df049..0000000000000
--- a/lib/vector-core/proptest-regressions/stream/partitioned_batcher.txt
+++ /dev/null
@@ -1,7 +0,0 @@
-# Seeds for failure cases proptest has generated in the past. It is
-# automatically read and these particular cases re-run before any
-# novel cases are generated.
-#
-# It is recommended to check this file in to source control so that
-# everyone who runs the test benefits from these saved cases.
-cc 129cf5357923c46bfb0534871fb03baa08aa96e741bdbd855864263c22e303d3 # shrinks to stream = [0], item_limit = 1, allocation_limit = 8, partitioner = TestPartitioner { key_space: 1 }, timer = TestTimer { responses: [], valid_keys: {} }

From e014b2941514f405f441b1a83c300da194aa8446 Mon Sep 17 00:00:00 2001
From: Stephen Wakely <fungus.humungus@gmail.com>
Date: Tue, 24 Oct 2023 13:56:33 +0100
Subject: [PATCH 32/32] Snappy is no longer optional

Signed-off-by: Stephen Wakely <fungus.humungus@gmail.com>
---
 Cargo.toml | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 3d48ecf208b66..99686fcfd88ff 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -316,7 +316,7 @@ seahash = { version = "4.1.0", default-features = false }
 semver = { version = "1.0.20", default-features = false, features = ["serde", "std"], optional = true }
 smallvec = { version = "1", default-features = false, features = ["union", "serde"] }
 snafu = { version = "0.7.5", default-features = false, features = ["futures"] }
-snap = { version = "1.1.0", default-features = false, optional = true }
+snap = { version = "1.1.0", default-features = false }
 socket2 = { version = "0.5.5", default-features = false }
 stream-cancel = { version = "0.8.1", default-features = false }
 strip-ansi-escapes = { version = "0.2.0", default-features = false }
@@ -560,9 +560,9 @@ sources-splunk_hec = ["dep:roaring"]
 sources-statsd = ["sources-utils-net", "tokio-util/net"]
 sources-stdin = ["tokio-util/io"]
 sources-syslog = ["codecs-syslog", "sources-utils-net", "tokio-util/net"]
-sources-utils-http = ["dep:snap", "sources-utils-http-auth", "sources-utils-http-encoding", "sources-utils-http-error", "sources-utils-http-prelude"]
+sources-utils-http = ["sources-utils-http-auth", "sources-utils-http-encoding", "sources-utils-http-error", "sources-utils-http-prelude"]
 sources-utils-http-auth = ["sources-utils-http-error"]
-sources-utils-http-encoding = ["dep:snap", "sources-utils-http-error"]
+sources-utils-http-encoding = ["sources-utils-http-error"]
 sources-utils-http-error = []
 sources-utils-http-prelude = ["sources-utils-http", "sources-utils-http-auth", "sources-utils-http-encoding", "sources-utils-http-error"]
 sources-utils-http-query = []
@@ -715,7 +715,7 @@ sinks-nats = ["dep:async-nats", "dep:nkeys"]
 sinks-new_relic_logs = ["sinks-http"]
 sinks-new_relic = []
 sinks-papertrail = ["dep:syslog"]
-sinks-prometheus = ["dep:base64", "dep:prometheus-parser", "dep:snap"]
+sinks-prometheus = ["dep:base64", "dep:prometheus-parser"]
 sinks-pulsar = ["dep:apache-avro", "dep:pulsar", "dep:lru"]
 sinks-redis = ["dep:redis"]
 sinks-sematext = ["sinks-elasticsearch", "sinks-influxdb"]