Integrate encoding::Encoder with loki sink (#12558)

Signed-off-by: Pablo Sichert <[email protected]>

Author: pablosichert

lib/codecs/src/encoding/format/logfmt.rs

@@ -0,0 +1,77 @@
+use bytes::BytesMut;
+use serde::{Deserialize, Serialize};
+use tokio_util::codec::Encoder;
+use vector_common::encode_logfmt;
+use vector_core::{config::DataType, event::Event, schema};
+
+/// Config used to build a `LogfmtSerializer`.
+#[derive(Debug, Clone, Default, Deserialize, Serialize)]
+pub struct LogfmtSerializerConfig;
+
+impl LogfmtSerializerConfig {
+    /// Creates a new `LogfmtSerializerConfig`.
+    pub const fn new() -> Self {
+        Self
+    }
+
+    /// Build the `LogfmtSerializer` from this configuration.
+    pub const fn build(&self) -> LogfmtSerializer {
+        LogfmtSerializer
+    }
+
+    /// The data type of events that are accepted by `JsonSerializer`.
+    pub fn input_type(&self) -> DataType {
+        DataType::Log
+    }
+
+    /// The schema required by the serializer.
+    pub fn schema_requirement(&self) -> schema::Requirement {
+        // While technically we support `Value` variants that can't be losslessly serialized to
+        // logfmt, we don't want to enforce that limitation to users yet.
+        schema::Requirement::empty()
+    }
+}
+
+/// Serializer that converts an `Event` to bytes using the logfmt format.
+#[derive(Debug, Clone)]
+pub struct LogfmtSerializer;
+
+impl LogfmtSerializer {
+    /// Creates a new `LogfmtSerializer`.
+    pub const fn new() -> Self {
+        Self
+    }
+}
+
+impl Encoder<Event> for LogfmtSerializer {
+    type Error = vector_core::Error;
+
+    fn encode(&mut self, event: Event, buffer: &mut BytesMut) -> Result<(), Self::Error> {
+        let log = event.as_log();
+        let string = encode_logfmt::to_string(log.as_map())?;
+        buffer.extend_from_slice(string.as_bytes());
+
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use bytes::BytesMut;
+    use vector_common::btreemap;
+    use vector_core::event::Value;
+
+    #[test]
+    fn serialize_logfmt() {
+        let event = Event::from(btreemap! {
+            "foo" => Value::from("bar")
+        });
+        let mut serializer = LogfmtSerializer::new();
+        let mut bytes = BytesMut::new();
+
+        serializer.encode(event, &mut bytes).unwrap();
+
+        assert_eq!(bytes.freeze(), "foo=bar");
+    }
+}

lib/codecs/src/encoding/format/mod.rs

@@ -4,6 +4,7 @@
 #![deny(missing_docs)]
 
 mod json;
+mod logfmt;
 mod native;
 mod native_json;
 mod raw_message;
@@ -13,6 +14,7 @@ use std::fmt::Debug;
 
 use dyn_clone::DynClone;
 pub use json::{JsonSerializer, JsonSerializerConfig};
+pub use logfmt::{LogfmtSerializer, LogfmtSerializerConfig};
 pub use native::{NativeSerializer, NativeSerializerConfig};
 pub use native_json::{NativeJsonSerializer, NativeJsonSerializerConfig};
 pub use raw_message::{RawMessageSerializer, RawMessageSerializerConfig};

lib/codecs/src/encoding/mod.rs

@@ -8,9 +8,9 @@ use std::fmt::Debug;
 
 use bytes::BytesMut;
 pub use format::{
-    JsonSerializer, JsonSerializerConfig, NativeJsonSerializer, NativeJsonSerializerConfig,
-    NativeSerializer, NativeSerializerConfig, RawMessageSerializer, RawMessageSerializerConfig,
-    TextSerializer, TextSerializerConfig,
+    JsonSerializer, JsonSerializerConfig, LogfmtSerializer, LogfmtSerializerConfig,
+    NativeJsonSerializer, NativeJsonSerializerConfig, NativeSerializer, NativeSerializerConfig,
+    RawMessageSerializer, RawMessageSerializerConfig, TextSerializer, TextSerializerConfig,
 };
 pub use framing::{
     BoxedFramer, BoxedFramingError, BytesEncoder, BytesEncoderConfig, CharacterDelimitedEncoder,
@@ -183,6 +183,8 @@ impl tokio_util::codec::Encoder<()> for Framer {
 pub enum SerializerConfig {
     /// Configures the `JsonSerializer`.
     Json,
+    /// Configures the `LogfmtSerializer`.
+    Logfmt,
     /// Configures the `NativeSerializer`.
     Native,
     /// Configures the `NativeJsonSerializer`.
@@ -199,6 +201,12 @@ impl From<JsonSerializerConfig> for SerializerConfig {
     }
 }
 
+impl From<LogfmtSerializerConfig> for SerializerConfig {
+    fn from(_: LogfmtSerializerConfig) -> Self {
+        Self::Logfmt
+    }
+}
+
 impl From<NativeSerializerConfig> for SerializerConfig {
     fn from(_: NativeSerializerConfig) -> Self {
         Self::Native
@@ -228,6 +236,7 @@ impl SerializerConfig {
     pub const fn build(&self) -> Serializer {
         match self {
             SerializerConfig::Json => Serializer::Json(JsonSerializerConfig.build()),
+            SerializerConfig::Logfmt => Serializer::Logfmt(LogfmtSerializerConfig.build()),
             SerializerConfig::Native => Serializer::Native(NativeSerializerConfig.build()),
             SerializerConfig::NativeJson => {
                 Serializer::NativeJson(NativeJsonSerializerConfig.build())
@@ -243,6 +252,7 @@ impl SerializerConfig {
     pub fn input_type(&self) -> DataType {
         match self {
             SerializerConfig::Json => JsonSerializerConfig.input_type(),
+            SerializerConfig::Logfmt => LogfmtSerializerConfig.input_type(),
             SerializerConfig::Native => NativeSerializerConfig.input_type(),
             SerializerConfig::NativeJson => NativeJsonSerializerConfig.input_type(),
             SerializerConfig::RawMessage => RawMessageSerializerConfig.input_type(),
@@ -254,6 +264,7 @@ impl SerializerConfig {
     pub fn schema_requirement(&self) -> schema::Requirement {
         match self {
             SerializerConfig::Json => JsonSerializerConfig.schema_requirement(),
+            SerializerConfig::Logfmt => LogfmtSerializerConfig.schema_requirement(),
             SerializerConfig::Native => NativeSerializerConfig.schema_requirement(),
             SerializerConfig::NativeJson => NativeJsonSerializerConfig.schema_requirement(),
             SerializerConfig::RawMessage => RawMessageSerializerConfig.schema_requirement(),
@@ -267,6 +278,8 @@ impl SerializerConfig {
 pub enum Serializer {
     /// Uses a `JsonSerializer` for serialization.
     Json(JsonSerializer),
+    /// Uses a `LogfmtSerializer` for serialization.
+    Logfmt(LogfmtSerializer),
     /// Uses a `NativeSerializer` for serialization.
     Native(NativeSerializer),
     /// Uses a `NativeJsonSerializer` for serialization.
@@ -283,6 +296,12 @@ impl From<JsonSerializer> for Serializer {
     }
 }
 
+impl From<LogfmtSerializer> for Serializer {
+    fn from(serializer: LogfmtSerializer) -> Self {
+        Self::Logfmt(serializer)
+    }
+}
+
 impl From<NativeSerializer> for Serializer {
     fn from(serializer: NativeSerializer) -> Self {
         Self::Native(serializer)
@@ -313,6 +332,7 @@ impl tokio_util::codec::Encoder<Event> for Serializer {
     fn encode(&mut self, event: Event, buffer: &mut BytesMut) -> Result<(), Self::Error> {
         match self {
             Serializer::Json(serializer) => serializer.encode(event, buffer),
+            Serializer::Logfmt(serializer) => serializer.encode(event, buffer),
             Serializer::Native(serializer) => serializer.encode(event, buffer),
             Serializer::NativeJson(serializer) => serializer.encode(event, buffer),
             Serializer::RawMessage(serializer) => serializer.encode(event, buffer),

lib/codecs/src/lib.rs

@@ -19,7 +19,8 @@ pub use decoding::{SyslogDeserializer, SyslogDeserializerConfig};
 pub use encoding::{
     BytesEncoder, BytesEncoderConfig, CharacterDelimitedEncoder, CharacterDelimitedEncoderConfig,
     JsonSerializer, JsonSerializerConfig, LengthDelimitedEncoder, LengthDelimitedEncoderConfig,
-    NativeJsonSerializer, NativeJsonSerializerConfig, NativeSerializer, NativeSerializerConfig,
-    NewlineDelimitedEncoder, NewlineDelimitedEncoderConfig, RawMessageSerializer,
-    RawMessageSerializerConfig, TextSerializer, TextSerializerConfig,
+    LogfmtSerializer, LogfmtSerializerConfig, NativeJsonSerializer, NativeJsonSerializerConfig,
+    NativeSerializer, NativeSerializerConfig, NewlineDelimitedEncoder,
+    NewlineDelimitedEncoderConfig, RawMessageSerializer, RawMessageSerializerConfig,
+    TextSerializer, TextSerializerConfig,
 };

src/codecs/encoder.rs

@@ -116,6 +116,7 @@ impl Encoder<Framer> {
             (Serializer::Native(_), _) => "application/octet-stream",
             (
                 Serializer::Json(_)
+                | Serializer::Logfmt(_)
                 | Serializer::NativeJson(_)
                 | Serializer::RawMessage(_)
                 | Serializer::Text(_),

src/sinks/aws_s3/config.rs

@@ -158,9 +158,13 @@ impl S3SinkConfig {
                 // TODO: We probably want to use something like octet framing here.
                 return Err("Native encoding is not implemented for this sink yet".into());
             }
-            (None, Serializer::NativeJson(_) | Serializer::RawMessage(_) | Serializer::Text(_)) => {
-                NewlineDelimitedEncoder::new().into()
-            }
+            (
+                None,
+                Serializer::Logfmt(_)
+                | Serializer::NativeJson(_)
+                | Serializer::RawMessage(_)
+                | Serializer::Text(_),
+            ) => NewlineDelimitedEncoder::new().into(),
         };
         let encoder = Encoder::<Framer>::new(framer, serializer);
 

src/sinks/azure_blob/config.rs

@@ -143,9 +143,13 @@ impl AzureBlobSinkConfig {
                 // TODO: We probably want to use something like octet framing here.
                 return Err("Native encoding is not implemented for this sink yet".into());
             }
-            (None, Serializer::NativeJson(_) | Serializer::RawMessage(_) | Serializer::Text(_)) => {
-                NewlineDelimitedEncoder::new().into()
-            }
+            (
+                None,
+                Serializer::Logfmt(_)
+                | Serializer::NativeJson(_)
+                | Serializer::RawMessage(_)
+                | Serializer::Text(_),
+            ) => NewlineDelimitedEncoder::new().into(),
         };
         let encoder = Encoder::<Framer>::new(framer, serializer);
 

src/sinks/gcp/cloud_storage.rs

@@ -286,9 +286,13 @@ impl RequestSettings {
             (Some(framer), _) => framer,
             (None, Serializer::Json(_)) => CharacterDelimitedEncoder::new(b',').into(),
             (None, Serializer::Native(_)) => LengthDelimitedEncoder::new().into(),
-            (None, Serializer::NativeJson(_) | Serializer::RawMessage(_) | Serializer::Text(_)) => {
-                NewlineDelimitedEncoder::new().into()
-            }
+            (
+                None,
+                Serializer::Logfmt(_)
+                | Serializer::NativeJson(_)
+                | Serializer::RawMessage(_)
+                | Serializer::Text(_),
+            ) => NewlineDelimitedEncoder::new().into(),
         };
         let encoder = Encoder::<Framer>::new(framer, serializer);
         let acl = config

src/sinks/loki/config.rs

@@ -1,9 +1,12 @@
 use std::collections::HashMap;
 
+use codecs::encoding::SerializerConfig;
+use codecs::{JsonSerializerConfig, LogfmtSerializerConfig, TextSerializerConfig};
 use futures::future::FutureExt;
 use serde::{Deserialize, Serialize};
 
 use super::{healthcheck::healthcheck, sink::LokiSink};
+use crate::sinks::util::encoding::{EncodingConfigAdapter, EncodingConfigMigrator};
 use crate::sinks::util::Compression;
 use crate::{
     config::{AcknowledgementsConfig, GenerateConfig, Input, SinkConfig, SinkContext},
@@ -18,15 +21,28 @@ use crate::{
     tls::{TlsConfig, TlsSettings},
 };
 
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct EncodingMigrator;
+
+impl EncodingConfigMigrator for EncodingMigrator {
+    type Codec = Encoding;
+
+    fn migrate(codec: &Self::Codec) -> SerializerConfig {
+        match codec {
+            Encoding::Json => JsonSerializerConfig::new().into(),
+            Encoding::Text => TextSerializerConfig::new().into(),
+            Encoding::Logfmt => LogfmtSerializerConfig::new().into(),
+        }
+    }
+}
+
 #[derive(Clone, Debug, Deserialize, Serialize)]
 #[serde(deny_unknown_fields)]
 pub struct LokiConfig {
     pub endpoint: UriSerde,
-    pub encoding: EncodingConfig<Encoding>,
-
+    pub encoding: EncodingConfigAdapter<EncodingConfig<Encoding>, EncodingMigrator>,
     pub tenant_id: Option<Template>,
     pub labels: HashMap<Template, Template>,
-
     #[serde(default = "crate::serde::default_false")]
     pub remove_label_fields: bool,
     #[serde(default = "crate::serde::default_true")]
@@ -35,17 +51,12 @@ pub struct LokiConfig {
     pub compression: Compression,
     #[serde(default)]
     pub out_of_order_action: OutOfOrderAction,
-
     pub auth: Option<Auth>,
-
     #[serde(default)]
     pub request: TowerRequestConfig,
-
     #[serde(default)]
     pub batch: BatchConfig<LokiDefaultBatchSettings>,
-
     pub tls: Option<TlsConfig>,
-
     #[serde(
         default,
         deserialize_with = "crate::serde::bool_or_struct",
@@ -85,7 +96,7 @@ impl GenerateConfig for LokiConfig {
     fn generate_config() -> toml::Value {
         toml::from_str(
             r#"endpoint = "http://localhost:3100"
-            encoding = "json"
+            encoding.codec = "json"
             labels = {}"#,
         )
         .unwrap()
@@ -132,7 +143,7 @@ impl SinkConfig for LokiConfig {
     }
 
     fn input(&self) -> Input {
-        Input::log()
+        Input::new(self.encoding.config().input_type())
     }
 
     fn sink_type(&self) -> &'static str {

src/sinks/loki/event.rs

@@ -1,5 +1,6 @@
 use std::{collections::HashMap, io};
 
+use bytes::Bytes;
 use serde::{ser::SerializeSeq, Serialize};
 use vector_core::{
     event::{EventFinalizers, Finalizable},
@@ -52,7 +53,7 @@ impl From<Vec<LokiRecord>> for LokiBatch {
 #[derive(Clone, Debug)]
 pub struct LokiEvent {
     pub timestamp: i64,
-    pub event: String,
+    pub event: Bytes,
 }
 
 impl ByteSizeOf for LokiEvent {
@@ -68,7 +69,8 @@ impl Serialize for LokiEvent {
     {
         let mut seq = serializer.serialize_seq(Some(2))?;
         seq.serialize_element(&self.timestamp.to_string())?;
-        seq.serialize_element(&self.event)?;
+        let event = String::from_utf8_lossy(&self.event);
+        seq.serialize_element(&event)?;
         seq.end()
     }
 }