From ee5a52c90723992af9b705168a0fe9d88b88ded0 Mon Sep 17 00:00:00 2001
From: Stephen Wakely <fungus.humungus@gmail.com>
Date: Fri, 2 Jun 2023 17:45:32 +0100
Subject: [PATCH] fix(loki sink): use json size of unencoded event (#17572)

Currently the `loki` sink is encoding the event and then taking the
estimated json size of that encoded event. This is wrong. All sinks
should take the estimated json size of the unencoded event.

There is an open question around whether we should be taking the json
size before or after the `only_fields` and `except_fields` are applied.
I'm currently trying to get an answer to that. Currently everything is
before.

The `loki` sink works a little bit different to the other stream based
sinks. Most sinks pass the event unencoded to the request builder. It is
at this point that the json size of the metrics are calculated. However,
`loki` encodes the event and passes the encoded value to the request
builder.

This PR changes it so it also passes the json size to the request
builder so it can use that value to calculate the metrics.

Signed-off-by: Stephen Wakely <fungus.humungus@gmail.com>
---
 src/sinks/loki/event.rs | 18 ++----------------
 src/sinks/loki/sink.rs  |  4 +++-
 2 files changed, 5 insertions(+), 17 deletions(-)

diff --git a/src/sinks/loki/event.rs b/src/sinks/loki/event.rs
index 76389e2c64fb9d..1aede42e359547 100644
--- a/src/sinks/loki/event.rs
+++ b/src/sinks/loki/event.rs
@@ -139,21 +139,6 @@ impl ByteSizeOf for LokiEvent {
     }
 }
 
-/// This implementation approximates the `Serialize` implementation below, without any allocations.
-impl EstimatedJsonEncodedSizeOf for LokiEvent {
-    fn estimated_json_encoded_size_of(&self) -> JsonSize {
-        static BRACKETS_SIZE: JsonSize = JsonSize::new(2);
-        static COLON_SIZE: JsonSize = JsonSize::new(1);
-        static QUOTES_SIZE: JsonSize = JsonSize::new(2);
-
-        BRACKETS_SIZE
-            + QUOTES_SIZE
-            + self.timestamp.estimated_json_encoded_size_of()
-            + COLON_SIZE
-            + self.event.estimated_json_encoded_size_of()
-    }
-}
-
 impl Serialize for LokiEvent {
     fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
     where
@@ -172,6 +157,7 @@ pub struct LokiRecord {
     pub partition: PartitionKey,
     pub labels: Labels,
     pub event: LokiEvent,
+    pub json_byte_size: JsonSize,
     pub finalizers: EventFinalizers,
 }
 
@@ -187,7 +173,7 @@ impl ByteSizeOf for LokiRecord {
 
 impl EstimatedJsonEncodedSizeOf for LokiRecord {
     fn estimated_json_encoded_size_of(&self) -> JsonSize {
-        self.event.estimated_json_encoded_size_of()
+        self.json_byte_size
     }
 }
 
diff --git a/src/sinks/loki/sink.rs b/src/sinks/loki/sink.rs
index 766526bb12bff5..b5871d4e876e45 100644
--- a/src/sinks/loki/sink.rs
+++ b/src/sinks/loki/sink.rs
@@ -12,7 +12,7 @@ use vector_core::{
     partition::Partitioner,
     sink::StreamSink,
     stream::BatcherSettings,
-    ByteSizeOf,
+    ByteSizeOf, EstimatedJsonEncodedSizeOf,
 };
 
 use super::{
@@ -268,6 +268,7 @@ impl EventEncoder {
     pub(super) fn encode_event(&mut self, mut event: Event) -> Option<LokiRecord> {
         let tenant_id = self.key_partitioner.partition(&event);
         let finalizers = event.take_finalizers();
+        let json_byte_size = event.estimated_json_encoded_size_of();
         let mut labels = self.build_labels(&event);
         self.remove_label_fields(&mut event);
 
@@ -302,6 +303,7 @@ impl EventEncoder {
             },
             partition,
             finalizers,
+            json_byte_size,
         })
     }
 }