fix(aws sinks): ensure component_sent_bytes_total carries component and region labels

changelog.d/20356_aws_component_labels.fix.md

@@ -0,0 +1,3 @@
+AWS sinks (`aws_cloudwatch_logs`, `aws_s3`, `aws_kinesis_firehose`, `aws_kinesis_streams`, `aws_sns`, `aws_sqs`) now emit `component_sent_bytes_total` through the Driver instead of the transport layer, ensuring `component_id`, `component_kind`, `component_type`, and `region` labels are always present.
+
+authors: clee2691

lib/vector-common/src/internal_event/bytes_sent.rs

@@ -8,8 +8,15 @@ use super::{ByteSize, CounterName, Protocol, SharedString};
 crate::registered_event!(
     BytesSent {
         protocol: SharedString,
+        extra_labels: Vec<(SharedString, SharedString)>,
     } => {
-        bytes_sent: Counter = counter!(CounterName::ComponentSentBytesTotal, "protocol" => self.protocol.clone()),
+        bytes_sent: Counter = {
+            let mut labels: Vec<(String, String)> = vec![("protocol".to_string(), self.protocol.to_string())];
+            for (k, v) in &self.extra_labels {
+                labels.push((k.to_string(), v.to_string()));
+            }
+            counter!(CounterName::ComponentSentBytesTotal, &labels)
+        },
         protocol: SharedString = self.protocol,
     }
 
@@ -23,6 +30,7 @@ impl From<Protocol> for BytesSent {
     fn from(protocol: Protocol) -> Self {
         Self {
             protocol: protocol.0,
+            extra_labels: vec![],
         }
     }
 }

lib/vector-stream/src/driver.rs

@@ -44,6 +44,7 @@ pub struct Driver<St, Svc> {
     input: St,
     service: Svc,
     protocol: Option<SharedString>,
+    extra_labels: Vec<(SharedString, SharedString)>,
 }
 
 impl<St, Svc> Driver<St, Svc> {
@@ -52,6 +53,7 @@ impl<St, Svc> Driver<St, Svc> {
             input,
             service,
             protocol: None,
+            extra_labels: vec![],
         }
     }
 
@@ -64,6 +66,13 @@ impl<St, Svc> Driver<St, Svc> {
         self.protocol = Some(protocol.into());
         self
     }
+
+    /// Add an extra label to the `BytesSent` metric emitted by this driver.
+    #[must_use]
+    pub fn label(mut self, key: impl Into<SharedString>, value: impl Into<SharedString>) -> Self {
+        self.extra_labels.push((key.into(), value.into()));
+        self
+    }
 }
 
 impl<St, Svc> Driver<St, Svc>
@@ -92,12 +101,18 @@ where
             input,
             mut service,
             protocol,
+            extra_labels,
         } = self;
 
         let batched_input = input.ready_chunks(1024);
         pin!(batched_input);
 
-        let bytes_sent = protocol.map(|protocol| register(BytesSent { protocol }));
+        let bytes_sent = protocol.map(|protocol| {
+            register(BytesSent {
+                protocol,
+                extra_labels,
+            })
+        });
         let events_sent = RegisteredEventCache::new(());
 
         loop {

src/aws/mod.rs

@@ -183,9 +183,47 @@ pub async fn create_client<T>(
 where
     T: ClientBuilder,
 {
-    create_client_and_region::<T>(builder, auth, region, endpoint, proxy, tls_options, timeout)
-        .await
-        .map(|(client, _)| client)
+    build_client_inner::<T>(
+        builder,
+        auth,
+        region,
+        endpoint,
+        proxy,
+        tls_options,
+        timeout,
+        true,
+    )
+    .await
+    .map(|(client, _)| client)
+}
+
+/// Like [`create_client`], but suppresses transport-level `AwsBytesSent` emission.
+///
+/// Use this for sinks that report bytes through the Driver to avoid double-counting
+/// `component_sent_bytes_total`.
+pub async fn create_client_without_transport_metrics<T>(
+    builder: &T,
+    auth: &AwsAuthentication,
+    region: Option<Region>,
+    endpoint: Option<String>,
+    proxy: &ProxyConfig,
+    tls_options: Option<&TlsConfig>,
+    timeout: Option<&AwsTimeout>,
+) -> crate::Result<(T::Client, Region)>
+where
+    T: ClientBuilder,
+{
+    build_client_inner::<T>(
+        builder,
+        auth,
+        region,
+        endpoint,
+        proxy,
+        tls_options,
+        timeout,
+        false,
+    )
+    .await
 }
 
 /// Create the SDK client and resolve the region using the provided settings.
@@ -198,6 +236,33 @@ pub async fn create_client_and_region<T>(
     tls_options: Option<&TlsConfig>,
     timeout: Option<&AwsTimeout>,
 ) -> crate::Result<(T::Client, Region)>
+where
+    T: ClientBuilder,
+{
+    build_client_inner::<T>(
+        builder,
+        auth,
+        region,
+        endpoint,
+        proxy,
+        tls_options,
+        timeout,
+        true,
+    )
+    .await
+}
+
+#[allow(clippy::too_many_arguments)]
+async fn build_client_inner<T>(
+    builder: &T,
+    auth: &AwsAuthentication,
+    region: Option<Region>,
+    endpoint: Option<String>,
+    proxy: &ProxyConfig,
+    tls_options: Option<&TlsConfig>,
+    timeout: Option<&AwsTimeout>,
+    emit_bytes_sent: bool,
+) -> crate::Result<(T::Client, Region)>
 where
     T: ClientBuilder,
 {
@@ -216,7 +281,7 @@ where
     let connector = AwsHttpClient {
         http: connector,
         region: region.clone(),
-        emit_bytes_sent: true,
+        emit_bytes_sent,
     };
 
     // Build the configuration first.
@@ -487,8 +552,7 @@ where
             return HttpConnectorFuture::new(self.call_inner(req));
         }
 
-        let bytes_sent = Arc::new(AtomicUsize::new(0));
-
+        let bytes_sent = Arc::new(std::sync::atomic::AtomicUsize::new(0));
         let req = req.map(|body| {
             let bytes_sent = Arc::clone(&bytes_sent);
             body.map_preserve_contents(move |body| {

src/sinks/aws_cloudwatch_logs/config.rs

@@ -7,8 +7,12 @@ use tower::ServiceBuilder;
 use vector_lib::{codecs::JsonSerializerConfig, configurable::configurable_component, schema};
 use vrl::value::Kind;
 
+use aws_config::Region;
+
 use crate::{
-    aws::{AwsAuthentication, ClientBuilder, RegionOrEndpoint, create_client},
+    aws::{
+        AwsAuthentication, ClientBuilder, RegionOrEndpoint, create_client_without_transport_metrics,
+    },
     codecs::{Encoder, EncodingConfig},
     config::{
         AcknowledgementsConfig, DataType, GenerateConfig, Input, ProxyConfig, SinkConfig,
@@ -189,8 +193,11 @@ pub struct CloudwatchLogsSinkConfig {
 }
 
 impl CloudwatchLogsSinkConfig {
-    pub async fn create_client(&self, proxy: &ProxyConfig) -> crate::Result<CloudwatchLogsClient> {
-        create_client::<CloudwatchLogsClientBuilder>(
+    pub async fn create_client(
+        &self,
+        proxy: &ProxyConfig,
+    ) -> crate::Result<(CloudwatchLogsClient, Region)> {
+        create_client_without_transport_metrics::<CloudwatchLogsClientBuilder>(
             &CloudwatchLogsClientBuilder {},
             &self.auth,
             self.region.region(),
@@ -218,12 +225,13 @@ impl SinkConfig for CloudwatchLogsSinkConfig {
 
         let batcher_settings = self.batch.into_batcher_settings()?;
         let request_settings = self.request.tower.into_settings();
-        let client = self.create_client(cx.proxy()).await?;
+        let (client, resolved_region) = self.create_client(cx.proxy()).await?;
         let svc = ServiceBuilder::new()
             .settings(request_settings, CloudwatchRetryLogic::new())
             .service(CloudwatchLogsPartitionSvc::new(
                 self.clone(),
                 client.clone(),
+                resolved_region.to_string(),
             )?);
         let transformer = self.encoding.transformer();
         let serializer = self.encoding.build()?;
@@ -237,7 +245,7 @@ impl SinkConfig for CloudwatchLogsSinkConfig {
                 transformer,
                 encoder,
             },
-
+            region: resolved_region.to_string(),
             service: svc,
         };
         Ok((VectorSink::from_event_streamsink(sink), healthcheck))

src/sinks/aws_cloudwatch_logs/integration_tests.rs

@@ -576,7 +576,7 @@ async fn cloudwatch_healthcheck() {
         confinement: Default::default(),
     };
 
-    let client = config.create_client(&ProxyConfig::default()).await.unwrap();
+    let (client, _resolved_region) = config.create_client(&ProxyConfig::default()).await.unwrap();
     healthcheck(config, client).await.unwrap();
 }
 

src/sinks/aws_cloudwatch_logs/request.rs

@@ -22,7 +22,10 @@ use http::{HeaderValue, header::HeaderName};
 use indexmap::IndexMap;
 use tokio::sync::oneshot;
 
-use crate::sinks::aws_cloudwatch_logs::{config::Retention, service::CloudwatchError};
+use crate::sinks::{
+    aws_cloudwatch_logs::{config::Retention, service::CloudwatchError},
+    util::EncodedLength,
+};
 
 pub struct CloudwatchFuture {
     client: Client,
@@ -32,6 +35,8 @@ pub struct CloudwatchFuture {
     retention_enabled: bool,
     events: Vec<Vec<InputLogEvent>>,
     token_tx: Option<oneshot::Sender<Option<String>>>,
+    current_batch_bytes: usize,
+    accumulated_bytes_sent: usize,
 }
 
 struct Client {
@@ -55,6 +60,10 @@ enum State {
 }
 
 impl CloudwatchFuture {
+    fn batch_message_bytes(batch: &[InputLogEvent]) -> usize {
+        batch.iter().map(|e| e.encoded_length()).sum()
+    }
+
     /// Panics if events.is_empty()
     #[allow(clippy::too_many_arguments)]
     pub(super) fn new(
@@ -82,10 +91,12 @@ impl CloudwatchFuture {
             tags,
         };
 
-        let state = if let Some(token) = token {
-            State::Put(client.put_logs(Some(token), events.pop().expect("No Events to send")))
+        let (state, current_batch_bytes) = if let Some(token) = token {
+            let batch = events.pop().expect("No Events to send");
+            let bytes = Self::batch_message_bytes(&batch);
+            (State::Put(client.put_logs(Some(token), batch)), bytes)
         } else {
-            State::DescribeStream(client.describe_stream())
+            (State::DescribeStream(client.describe_stream()), 0)
         };
 
         let retention_enabled = retention.enabled;
@@ -98,6 +109,8 @@ impl CloudwatchFuture {
             create_missing_group,
             create_missing_stream,
             retention_enabled,
+            current_batch_bytes,
+            accumulated_bytes_sent: 0,
         }
     }
 }
@@ -145,6 +158,7 @@ impl Future for CloudwatchFuture {
 
                         let token = stream.upload_sequence_token;
 
+                        self.current_batch_bytes = Self::batch_message_bytes(&events);
                         info!(message = "Putting logs.", token = ?token);
                         self.state = State::Put(self.client.put_logs(token, events));
                     } else if self.create_missing_stream {
@@ -210,10 +224,21 @@ impl Future for CloudwatchFuture {
                 State::Put(fut) => {
                     let next_token = match ready!(fut.poll_unpin(cx)) {
                         Ok(resp) => resp.next_sequence_token,
-                        Err(err) => return Poll::Ready(Err(CloudwatchError::Put(err))),
+                        Err(err) => {
+                            if self.accumulated_bytes_sent > 0 {
+                                return Poll::Ready(Err(CloudwatchError::PutPartial {
+                                    error: err,
+                                    bytes_sent: self.accumulated_bytes_sent,
+                                }));
+                            }
+                            return Poll::Ready(Err(CloudwatchError::Put(err)));
+                        }
                     };
 
+                    self.accumulated_bytes_sent += self.current_batch_bytes;
+
                     if let Some(events) = self.events.pop() {
+                        self.current_batch_bytes = Self::batch_message_bytes(&events);
                         debug!(message = "Putting logs.", next_token = ?next_token);
                         self.state = State::Put(self.client.put_logs(next_token, events));
                     } else {

src/sinks/aws_cloudwatch_logs/request_builder.rs

@@ -109,8 +109,8 @@ impl CloudwatchRequestBuilder {
             return None;
         }
 
-        let bytes_len =
-            NonZeroUsize::new(message_bytes.len()).expect("payload should never be zero length");
+        let bytes_len = NonZeroUsize::new(message_bytes.len() + BATCH_SIZE_OVERHEAD)
+            .expect("payload should never be zero length");
         let metadata = builder.with_request_size(bytes_len);
 
         Some(CloudwatchRequest {

src/sinks/aws_cloudwatch_logs/retry.rs

@@ -45,7 +45,7 @@ impl<Request: Send + Sync + 'static, Response: Send + Sync + 'static> RetryLogic
     #[allow(clippy::cognitive_complexity)] // long, but just a hair over our limit
     fn is_retriable_error(&self, error: &Self::Error) -> bool {
         match error {
-            CloudwatchError::Put(err) => {
+            CloudwatchError::Put(err) | CloudwatchError::PutPartial { error: err, .. } => {
                 if let SdkError::ServiceError(inner) = err {
                     let err = inner.err();
                     if matches!(err, PutLogEventsError::ServiceUnavailableException(_)) {