enhancement(kubernetes source): Kubernetes partial message merge

.meta/sources/kubernetes.toml.bak

@@ -38,6 +38,14 @@ of `kube-system`. Which is by default excluded, \
 unless there are any non empty `include` options.\
 """
 
+[sources.kubernetes.options.auto_partial_merge]
+type = "bool"
+common = true
+default = true
+description = """\
+Setting this to `false` will disable the automatic merging of partial events.\
+"""
+
 [sources.kubernetes.output.log.fields.container_name]
 type = "string"
 examples = ["vector"]

Cargo.toml

@@ -228,7 +228,7 @@ sources-docker = ["shiplift"]
 sources-file = ["bytesize"]
 sources-journald = []
 sources-kafka = ["owning_ref"]
-sources-kubernetes = ["sources-file", "transforms-json_parser", "transforms-regex_parser"]
+sources-kubernetes = ["sources-file", "transforms-json_parser", "transforms-regex_parser", "transforms-merge"]
 sources-logplex = ["warp", "sources-tls"]
 sources-prometheus = []
 sources-http = ["warp", "sources-tls"]

src/sources/kubernetes/file_source_builder.rs

@@ -537,6 +537,7 @@ mod tests {
             include_pod_uids: vec!["a027f09d8f18234519fa930f8fa71234".to_owned()],
             include_container_names: vec!["busybox".to_owned()],
             include_namespaces: vec!["telemetry".to_owned(), "app".to_owned()],
+            ..KubernetesConfig::default()
         };
 
         assert_eq!(

src/sources/kubernetes/mod.rs

@@ -12,6 +12,7 @@ use crate::{
     sources::Source,
     topology::config::{DataType, GlobalOptions, SourceConfig},
     transforms::{
+        merge::{Merge, TrailingNewlineNormalizer},
         regex_parser::{RegexParser, RegexParserConfig},
         Transform,
     },
@@ -39,12 +40,15 @@ enum BuildError {
     IllegalCharacterInUid { uid: String },
 }
 
-#[derive(Deserialize, Serialize, Debug, Clone, Default)]
+#[derive(Deserialize, Serialize, Debug, Clone, Derivative)]
+#[derivative(Default)]
 #[serde(deny_unknown_fields, default)]
 pub struct KubernetesConfig {
     include_container_names: Vec<String>,
     include_pod_uids: Vec<String>,
     include_namespaces: Vec<String>,
+    #[derivative(Default(value = "true"))]
+    auto_partial_merge: bool,
 }
 
 #[typetag::serde(name = "kubernetes")]
@@ -73,12 +77,23 @@ impl SourceConfig for KubernetesConfig {
         let mut parse_message = message_parser::build_message_parser()?;
 
         // Kubernetes source
-        let source = file_recv
+        let stream = file_recv
             .filter_map(move |event| transform_file.transform(event))
             .filter_map(move |event| parse_message.transform(event))
-            .filter_map(move |event| now.filter(event))
-            .map(remove_ending_newline)
-            .filter_map(move |event| transform_pod_uid.transform(event))
+            .filter_map(move |event| now.filter(event));
+
+        let stream: Box<dyn Stream<Item = Event, Error = _> + Send> = if self.auto_partial_merge {
+            let mut transform_merge_partial_events = transform_merge_partial_events();
+            Box::new(
+                stream.filter_map(move |event| transform_merge_partial_events.transform(event)),
+            )
+        } else {
+            Box::new(stream)
+        };
+
+        let stream = stream.filter_map(move |event| transform_pod_uid.transform(event));
+
+        let source = stream
             .forward(out.sink_map_err(drop))
             .map(drop)
             .join(file_source)
@@ -120,18 +135,6 @@ impl TimeFilter {
     }
 }
 
-fn remove_ending_newline(mut event: Event) -> Event {
-    if let Some(Value::Bytes(msg)) = event
-        .as_mut_log()
-        .get_mut(&event::log_schema().message_key())
-    {
-        if msg.ends_with(&['\n' as u8]) {
-            msg.truncate(msg.len() - 1);
-        }
-    }
-    event
-}
-
 fn transform_file() -> crate::Result<Box<dyn Transform>> {
     let mut config = RegexParserConfig::default();
 
@@ -210,6 +213,17 @@ fn transform_pod_uid() -> crate::Result<ApplicableTransform> {
     Ok(ApplicableTransform::Candidates(transforms))
 }
 
+fn transform_merge_partial_events() -> Merge<TrailingNewlineNormalizer> {
+    let message_key = event::log_schema().message_key();
+    Merge::new(
+        TrailingNewlineNormalizer {
+            probe_field: message_key.clone(),
+        },
+        vec![message_key.clone()],
+        vec!["pod_uid".into()],
+    )
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -271,4 +285,75 @@ mod tests {
 
         has(&event, "object_uid", "306cd636-0c6d-11ea-9079-1c1b0de4d755");
     }
+
+    #[test]
+    fn partial_events_merge() {
+        let message_key = event::log_schema().message_key();
+        let sample_pod_uid = "qwerty";
+
+        let part1 = {
+            let mut event = Event::new_empty_log();
+            event
+                .as_mut_log()
+                .insert("pod_uid", sample_pod_uid.to_owned());
+            event.as_mut_log().insert(message_key, "hello".to_owned());
+            event
+        };
+        let part2 = {
+            let mut event = Event::new_empty_log();
+            event
+                .as_mut_log()
+                .insert("pod_uid", sample_pod_uid.to_owned());
+            event
+                .as_mut_log()
+                .insert(message_key, " world!\n".to_owned());
+            event
+        };
+
+        let mut transform = transform_merge_partial_events();
+
+        assert!(transform.transform(part1).is_none());
+        let event = transform.transform(part2).unwrap();
+
+        has(&event, "pod_uid", sample_pod_uid);
+        has(&event, message_key, "hello world!");
+    }
+
+    #[test]
+    fn partial_events_merge_separate_streams() {
+        let message_key = event::log_schema().message_key();
+
+        let part1 = {
+            let mut event = Event::new_empty_log();
+            event.as_mut_log().insert("pod_uid", "qwerty1".to_owned());
+            event.as_mut_log().insert(message_key, "par".to_owned());
+            event
+        };
+        let part2 = {
+            let mut event = Event::new_empty_log();
+            event.as_mut_log().insert("pod_uid", "qwerty2".to_owned());
+            event
+                .as_mut_log()
+                .insert(message_key, "non-partial!\n".to_owned());
+            event
+        };
+        let part3 = {
+            let mut event = Event::new_empty_log();
+            event.as_mut_log().insert("pod_uid", "qwerty1".to_owned());
+            event.as_mut_log().insert(message_key, "tial!\n".to_owned());
+            event
+        };
+
+        let mut transform = transform_merge_partial_events();
+
+        assert!(transform.transform(part1).is_none());
+
+        let event = transform.transform(part2).unwrap();
+        has(&event, "pod_uid", "qwerty2");
+        has(&event, message_key, "non-partial!");
+
+        let event = transform.transform(part3).unwrap();
+        has(&event, "pod_uid", "qwerty1");
+        has(&event, message_key, "partial!");
+    }
 }

src/sources/kubernetes/test.rs

@@ -15,6 +15,7 @@ use kube::{
 };
 use serde::de::DeserializeOwned;
 use serde_json::Value;
+use std::fmt::Write;
 use uuid::Uuid;
 
 static NAMESPACE_MARKER: &'static str = "$(TEST_NAMESPACE)";
@@ -671,3 +672,40 @@ fn kube_diff_pod_uid() {
         false
     });
 }
+
+#[test]
+fn kube_partial() {
+    let namespace = format!("partial-{}", Uuid::new_v4());
+    let message = {
+        let mut s = String::new();
+        for i in 0..8 {
+            write!(s, "{}", i.to_string().repeat(8 * 1024)).unwrap();
+        }
+        s
+    };
+    assert_eq!(message.len(), 64 * 1024); // 64 kb
+    let user_namespace = user_namespace(&namespace);
+
+    let kube = Kube::new(&namespace);
+    let user = Kube::new(&user_namespace);
+
+    // Start vector
+    let vector = start_vector(&kube, user_namespace.as_str(), None);
+
+    // Start echo
+    let _echo = echo(&user, "echo", &message);
+    // Verify logs
+    wait_for(|| {
+        // If any daemon logged message, done.
+        for line in logs(&kube, &vector) {
+            if line["message"].as_str().unwrap() == message {
+                // Very long message arrived as a single (merged) message.
+                // DONE
+                return true;
+            } else {
+                debug!(namespace=%namespace,log=%line);
+            }
+        }
+        false
+    });
+}