ml: group: stream: Prepare to handle limitations of multiline

Signed-off-by: Hiroshi Hatake <[email protected]>
Co-authored-by: Eduardo Silva <[email protected]>

Author: cosmo0920

include/fluent-bit/multiline/flb_ml.h

@@ -52,6 +52,13 @@
 /* Default multiline buffer size: 4Kb */
 #define FLB_ML_BUF_SIZE         1024*4
 
+/* Default limit for concatenated multiline messages: 2MB */
+#define FLB_ML_BUFFER_LIMIT_DEFAULT  (1024 * 1024 * 2)
+
+/* Return codes */
+#define FLB_MULTILINE_OK         0
+#define FLB_MULTILINE_TRUNCATED  1
+
 /* Maximum number of groups per stream */
 #define FLB_ML_MAX_GROUPS       6
 
@@ -107,6 +114,9 @@ struct flb_ml_stream_group {
     msgpack_packer mp_pck;      /* temporary msgpack packer              */
     struct flb_time mp_time;    /* multiline time parsed from first line */
 
+    /* parent stream reference */
+    struct flb_ml_stream *stream;
+
     struct mk_list _head;
 };
 
@@ -275,6 +285,9 @@ struct flb_ml {
     struct flb_log_event_encoder log_event_encoder;
     struct flb_log_event_decoder log_event_decoder;
     struct flb_config *config;             /* Fluent Bit context */
+
+    /* Limit for concatenated multiline messages */
+    size_t buffer_limit;
 };
 
 struct flb_ml *flb_ml_create(struct flb_config *ctx, char *name);

include/fluent-bit/multiline/flb_ml_group.h

@@ -28,4 +28,12 @@ struct flb_ml_group *flb_ml_group_create(struct flb_ml *ml);
 void flb_ml_group_destroy(struct flb_ml_group *group);
 int flb_ml_group_add_parser(struct flb_ml *ctx, struct flb_ml_parser_ins *p);
 
+/*
+ * Append data to a multiline stream group respecting the configured
+ * buffer limit. The length of the appended data might be reduced if
+ * the limit is reached.
+ */
+int flb_ml_group_cat(struct flb_ml_stream_group *group,
+                     const char *data, size_t len);
+
 #endif

src/multiline/flb_ml_group.c

@@ -20,6 +20,7 @@
 #include <fluent-bit/flb_info.h>
 #include <fluent-bit/flb_mem.h>
 #include <fluent-bit/flb_log.h>
+#include <fluent-bit/flb_sds.h>
 #include <fluent-bit/multiline/flb_ml.h>
 #include <fluent-bit/multiline/flb_ml_parser.h>
 
@@ -84,3 +85,36 @@ void flb_ml_group_destroy(struct flb_ml_group *group)
     mk_list_del(&group->_head);
     flb_free(group);
 }
+
+int flb_ml_group_cat(struct flb_ml_stream_group *group,
+                     const char *data, size_t len)
+{
+    size_t avail;
+    size_t limit;
+    int    ret;
+    int    status = FLB_MULTILINE_OK;
+
+    limit = group->stream->ml->buffer_limit;
+    if (limit > 0) {
+        if (flb_sds_len(group->buf) >= limit) {
+            return FLB_MULTILINE_TRUNCATED;
+        }
+
+        avail = limit - flb_sds_len(group->buf);
+        if (len > avail) {
+            len = avail;
+            status = FLB_MULTILINE_TRUNCATED;
+        }
+    }
+
+    if (len == 0) {
+        return status;
+    }
+
+    ret = flb_sds_cat_safe(&group->buf, data, len);
+    if (ret == -1) {
+        return -1;
+    }
+
+    return status;
+}

src/multiline/flb_ml_stream.c

@@ -79,6 +79,9 @@ static struct flb_ml_stream_group *stream_group_create(struct flb_ml_stream *mst
     msgpack_sbuffer_init(&group->mp_sbuf);
     msgpack_packer_init(&group->mp_pck, &group->mp_sbuf, msgpack_sbuffer_write);
 
+    /* parent stream reference */
+    group->stream = mst;
+
     mk_list_add(&group->_head, &mst->groups);
 
     return group;