chore(observability): count byte_size after transforming event

lib/vector-core/src/event/log_event.rs

@@ -293,10 +293,14 @@ impl LogEvent {
     }
 
     pub fn get_by_meaning(&self, meaning: impl AsRef<str>) -> Option<&Value> {
-        self.metadata()
-            .schema_definition()
-            .meaning_path(meaning.as_ref())
-            .and_then(|path| self.get(path))
+        if let Some(dropped) = self.metadata().dropped_field(&meaning) {
+            Some(dropped)
+        } else {
+            self.metadata()
+                .schema_definition()
+                .meaning_path(meaning.as_ref())
+                .and_then(|path| self.get(path))
+        }
     }
 
     // TODO(Jean): Once the event API uses `Lookup`, the allocation here can be removed.

lib/vector-core/src/event/metadata.rs

@@ -46,6 +46,14 @@ pub struct EventMetadata {
     /// TODO(Jean): must not skip serialization to track schemas across restarts.
     #[serde(default = "default_schema_definition", skip)]
     schema_definition: Arc<schema::Definition>,
+
+    /// A store of values that may be dropped during the encoding process but may be needed
+    /// later on. The map is indexed by meaning.
+    /// Currently this is just used for the `service`. If the service field is dropped by `only_fields`
+    /// we need to ensure it is still available later on for emitting metrics tagged by the service.
+    /// This field could almost be keyed by `&'static str`, but because it needs to be deserializable
+    /// we have to use `String`.
+    dropped_fields: BTreeMap<String, Value>,
 }
 
 fn default_metadata_value() -> Value {
@@ -123,6 +131,17 @@ impl EventMetadata {
     pub fn set_splunk_hec_token(&mut self, secret: Arc<str>) {
         self.secrets.insert(SPLUNK_HEC_TOKEN, secret);
     }
+
+    /// Adds the value to the dropped fields list, returning the
+    /// position the value will be found in the list.
+    pub fn add_dropped_field(&mut self, meaning: String, value: Value) {
+        self.dropped_fields.insert(meaning, value);
+    }
+
+    /// Fetches the dropped field by meaning.
+    pub fn dropped_field(&self, meaning: impl AsRef<str>) -> Option<&Value> {
+        self.dropped_fields.get(meaning.as_ref())
+    }
 }
 
 impl Default for EventMetadata {
@@ -134,6 +153,7 @@ impl Default for EventMetadata {
             schema_definition: default_schema_definition(),
             source_id: None,
             upstream_id: None,
+            dropped_fields: BTreeMap::new(),
         }
     }
 }

src/codecs/encoding/transformer.rs

@@ -128,11 +128,23 @@ impl Transformer {
 
     fn apply_only_fields(&self, log: &mut LogEvent) {
         if let Some(only_fields) = self.only_fields.as_ref() {
-            let old_value = std::mem::replace(log.value_mut(), Value::Object(BTreeMap::new()));
+            let mut old_value = std::mem::replace(log.value_mut(), Value::Object(BTreeMap::new()));
 
             for field in only_fields {
-                if let Some(value) = old_value.get(field) {
-                    log.insert((PathPrefix::Event, field), value.clone());
+                if let Some(value) = old_value.remove(field, true) {
+                    log.insert((PathPrefix::Event, field), value);
+                }
+            }
+
+            // We may need the service field to apply tags to emitted metrics after the log message has been pruned. If there
+            // is a service meaning, we move this value to `dropped_fields` in the metadata.
+            // If the field is still in the new log message after pruning it will have been removed from `old_value` above.
+            let service_path = log.metadata().schema_definition().meaning_path("service");
+            if let Some(service_path) = service_path {
+                let mut new_log = LogEvent::from(old_value);
+                if let Some(service) = new_log.remove(service_path) {
+                    log.metadata_mut()
+                        .add_dropped_field("service".to_string(), service);
                 }
             }
         }
@@ -141,7 +153,20 @@ impl Transformer {
     fn apply_except_fields(&self, log: &mut LogEvent) {
         if let Some(except_fields) = self.except_fields.as_ref() {
             for field in except_fields {
-                log.remove(field.as_str());
+                let value = log.remove(field.as_str());
+
+                // If we are removing the service field we need to store this in a `dropped_fields` list as we may need to
+                // refer to this later when emitting metrics.
+                if let Some(v) = value {
+                    if let Some(path) = log.metadata().schema_definition().meaning_path("service") {
+                        use lookup::path::TargetPath;
+
+                        if &path.value_path().to_string() == field {
+                            log.metadata_mut()
+                                .add_dropped_field("service".to_string(), v);
+                        }
+                    }
+                }
             }
         }
     }
@@ -213,10 +238,15 @@ pub enum TimestampFormat {
 #[cfg(test)]
 mod tests {
     use indoc::indoc;
-    use vector_core::config::log_schema;
+    use lookup::path::parse_target_path;
+    use vector_common::btreemap;
+    use vector_core::config::{log_schema, LogNamespace};
+    use vrl::value::Kind;
+
+    use crate::config::schema;
 
     use super::*;
-    use std::collections::BTreeMap;
+    use std::{collections::BTreeMap, sync::Arc};
 
     #[test]
     fn serialize() {
@@ -374,4 +404,83 @@ mod tests {
         "#});
         assert!(config.is_err())
     }
+
+    #[test]
+    fn only_fields_with_service() {
+        let transformer: Transformer = toml::from_str(r#"only_fields = ["message"]"#).unwrap();
+        let mut log = LogEvent::default();
+        {
+            log.insert("message", 1);
+            log.insert("thing.service", "carrot");
+        }
+
+        let schema = schema::Definition::new_with_default_metadata(
+            Kind::object(btreemap! {
+                "thing" => Kind::object(btreemap! {
+                    "service" => Kind::bytes(),
+                })
+            }),
+            [LogNamespace::Vector],
+        );
+
+        let schema = schema.with_meaning(parse_target_path("thing.service").unwrap(), "service");
+
+        let mut event = Event::from(log);
+
+        event
+            .metadata_mut()
+            .set_schema_definition(&Arc::new(schema));
+
+        transformer.transform(&mut event);
+        assert!(event.as_mut_log().contains("message"));
+
+        // Event no longer contains the service field.
+        assert!(!event.as_mut_log().contains("thing.service"));
+
+        // But we can still get the service by meaning.
+        assert_eq!(
+            &Value::from("carrot"),
+            event.as_log().get_by_meaning("service").unwrap()
+        );
+    }
+
+    #[test]
+    fn except_fields_with_service() {
+        let transformer: Transformer =
+            toml::from_str(r#"except_fields = ["thing.service"]"#).unwrap();
+        let mut log = LogEvent::default();
+        {
+            log.insert("message", 1);
+            log.insert("thing.service", "carrot");
+        }
+
+        let schema = schema::Definition::new_with_default_metadata(
+            Kind::object(btreemap! {
+                "thing" => Kind::object(btreemap! {
+                    "service" => Kind::bytes(),
+                })
+            }),
+            [LogNamespace::Vector],
+        );
+
+        let schema = schema.with_meaning(parse_target_path("thing.service").unwrap(), "service");
+
+        let mut event = Event::from(log);
+
+        event
+            .metadata_mut()
+            .set_schema_definition(&Arc::new(schema));
+
+        transformer.transform(&mut event);
+        assert!(event.as_mut_log().contains("message"));
+
+        // Event no longer contains the service field.
+        assert!(!event.as_mut_log().contains("thing.service"));
+
+        // But we can still get the service by meaning.
+        assert_eq!(
+            &Value::from("carrot"),
+            event.as_log().get_by_meaning("service").unwrap()
+        );
+    }
 }

src/sinks/amqp/encoder.rs

@@ -3,6 +3,7 @@ use crate::sinks::prelude::*;
 use bytes::BytesMut;
 use std::io;
 use tokio_util::codec::Encoder as _;
+use vector_core::config::telemetry;
 
 #[derive(Clone, Debug)]
 pub(super) struct AmqpEncoder {
@@ -11,9 +12,17 @@ pub(super) struct AmqpEncoder {
 }
 
 impl encoding::Encoder<Event> for AmqpEncoder {
-    fn encode_input(&self, mut input: Event, writer: &mut dyn io::Write) -> io::Result<usize> {
+    fn encode_input(
+        &self,
+        mut input: Event,
+        writer: &mut dyn io::Write,
+    ) -> io::Result<(usize, GroupedCountByteSize)> {
         let mut body = BytesMut::new();
         self.transformer.transform(&mut input);
+
+        let mut byte_size = telemetry().create_request_count_byte_size();
+        byte_size.add_event(&input, input.estimated_json_encoded_size_of());
+
         let mut encoder = self.encoder.clone();
         encoder
             .encode(input, &mut body)
@@ -22,6 +31,6 @@ impl encoding::Encoder<Event> for AmqpEncoder {
         let body = body.freeze();
         write_all(writer, 1, body.as_ref())?;
 
-        Ok(body.len())
+        Ok((body.len(), byte_size))
     }
 }

src/sinks/azure_blob/test.rs

@@ -4,7 +4,8 @@ use codecs::{
     encoding::{Framer, FramingConfig},
     NewlineDelimitedEncoder, TextSerializerConfig,
 };
-use vector_core::partition::Partitioner;
+use vector_common::request_metadata::GroupedCountByteSize;
+use vector_core::{partition::Partitioner, EstimatedJsonEncodedSizeOf};
 
 use super::config::AzureBlobSinkConfig;
 use super::request_builder::AzureBlobRequestOptions;
@@ -68,10 +69,13 @@ fn azure_blob_build_request_without_compression() {
         compression,
     };
 
+    let mut byte_size = GroupedCountByteSize::new_untagged();
+    byte_size.add_event(&log, log.estimated_json_encoded_size_of());
+
     let (metadata, request_metadata_builder, _events) =
         request_options.split_input((key, vec![log]));
 
-    let payload = EncodeResult::uncompressed(Bytes::new());
+    let payload = EncodeResult::uncompressed(Bytes::new(), byte_size);
     let request_metadata = request_metadata_builder.build(&payload);
     let request = request_options.build_request(metadata, request_metadata, payload);
 
@@ -112,10 +116,14 @@ fn azure_blob_build_request_with_compression() {
         ),
         compression,
     };
+
+    let mut byte_size = GroupedCountByteSize::new_untagged();
+    byte_size.add_event(&log, log.estimated_json_encoded_size_of());
+
     let (metadata, request_metadata_builder, _events) =
         request_options.split_input((key, vec![log]));
 
-    let payload = EncodeResult::uncompressed(Bytes::new());
+    let payload = EncodeResult::uncompressed(Bytes::new(), byte_size);
     let request_metadata = request_metadata_builder.build(&payload);
     let request = request_options.build_request(metadata, request_metadata, payload);
 
@@ -157,10 +165,13 @@ fn azure_blob_build_request_with_time_format() {
         compression,
     };
 
+    let mut byte_size = GroupedCountByteSize::new_untagged();
+    byte_size.add_event(&log, log.estimated_json_encoded_size_of());
+
     let (metadata, request_metadata_builder, _events) =
         request_options.split_input((key, vec![log]));
 
-    let payload = EncodeResult::uncompressed(Bytes::new());
+    let payload = EncodeResult::uncompressed(Bytes::new(), byte_size);
     let request_metadata = request_metadata_builder.build(&payload);
     let request = request_options.build_request(metadata, request_metadata, payload);
 
@@ -205,10 +216,13 @@ fn azure_blob_build_request_with_uuid() {
         compression,
     };
 
+    let mut byte_size = GroupedCountByteSize::new_untagged();
+    byte_size.add_event(&log, log.estimated_json_encoded_size_of());
+
     let (metadata, request_metadata_builder, _events) =
         request_options.split_input((key, vec![log]));
 
-    let payload = EncodeResult::uncompressed(Bytes::new());
+    let payload = EncodeResult::uncompressed(Bytes::new(), byte_size);
     let request_metadata = request_metadata_builder.build(&payload);
     let request = request_options.build_request(metadata, request_metadata, payload);
 

src/sinks/datadog/logs/sink.rs

@@ -7,7 +7,7 @@ use futures::stream::{BoxStream, StreamExt};
 use lookup::event_path;
 use snafu::Snafu;
 use tower::Service;
-use vector_common::request_metadata::RequestMetadata;
+use vector_common::request_metadata::{GroupedCountByteSize, RequestMetadata};
 use vector_core::{
     event::{Event, EventFinalizers, Finalizable, Value},
     partition::Partitioner,
@@ -125,7 +125,11 @@ impl JsonEncoding {
 }
 
 impl crate::sinks::util::encoding::Encoder<Vec<Event>> for JsonEncoding {
-    fn encode_input(&self, mut input: Vec<Event>, writer: &mut dyn io::Write) -> io::Result<usize> {
+    fn encode_input(
+        &self,
+        mut input: Vec<Event>,
+        writer: &mut dyn io::Write,
+    ) -> io::Result<(usize, GroupedCountByteSize)> {
         for event in input.iter_mut() {
             let log = event.as_mut_log();
             let message_path = log
@@ -219,7 +223,7 @@ impl RequestBuilder<(Option<Arc<str>>, Vec<Event>)> for LogRequestBuilder {
         // to (un)compressed size limitations.
         let mut buf = Vec::new();
         let n_events = events.len();
-        let uncompressed_size = self.encoder().encode_input(events, &mut buf)?;
+        let (uncompressed_size, byte_size) = self.encoder().encode_input(events, &mut buf)?;
         if uncompressed_size > MAX_PAYLOAD_BYTES {
             return Err(RequestBuildError::PayloadTooBig);
         }
@@ -230,9 +234,13 @@ impl RequestBuilder<(Option<Arc<str>>, Vec<Event>)> for LogRequestBuilder {
         let bytes = compressor.into_inner().freeze();
 
         if self.compression.is_compressed() {
-            Ok(EncodeResult::compressed(bytes, uncompressed_size))
+            Ok(EncodeResult::compressed(
+                bytes,
+                uncompressed_size,
+                byte_size,
+            ))
         } else {
-            Ok(EncodeResult::uncompressed(bytes))
+            Ok(EncodeResult::uncompressed(bytes, byte_size))
         }
     }