out_kinesis_streams: Add simple_aggregation operation

Signed-off-by: Shelby Hagman <[email protected]>

Author: ShelbyZ

plugins/out_kinesis_streams/kinesis.c

@@ -308,10 +308,10 @@ static int cb_kinesis_init(struct flb_output_instance *ins,
     return -1;
 }
 
-static struct flush *new_flush_buffer(const char *tag, int tag_len)
+static struct flush *new_flush_buffer(struct flb_kinesis *ctx, const char *tag, int tag_len)
 {
     struct flush *buf;
-
+    int ret;
 
     buf = flb_calloc(1, sizeof(struct flush));
     if (!buf) {
@@ -338,6 +338,18 @@ static struct flush *new_flush_buffer(const char *tag, int tag_len)
     buf->tag = tag;
     buf->tag_len = tag_len;
 
+    /* Initialize aggregation buffer if simple_aggregation is enabled */
+    buf->agg_buf_initialized = FLB_FALSE;
+    if (ctx->simple_aggregation) {
+        ret = flb_aws_aggregation_init(&buf->agg_buf, MAX_EVENT_SIZE);
+        if (ret < 0) {
+            flb_plg_error(ctx->ins, "Failed to initialize aggregation buffer");
+            kinesis_flush_destroy(buf);
+            return NULL;
+        }
+        buf->agg_buf_initialized = FLB_TRUE;
+    }
+
     return buf;
 }
 
@@ -353,7 +365,7 @@ static void cb_kinesis_flush(struct flb_event_chunk *event_chunk,
     (void) i_ins;
     (void) config;
 
-    buf = new_flush_buffer(event_chunk->tag, flb_sds_len(event_chunk->tag));
+    buf = new_flush_buffer(ctx, event_chunk->tag, flb_sds_len(event_chunk->tag));
     if (!buf) {
         flb_plg_error(ctx->ins, "Failed to construct flush buffer");
         FLB_OUTPUT_RETURN(FLB_RETRY);
@@ -503,6 +515,13 @@ static struct flb_config_map config_map[] = {
      "$HOME/.aws/ directory."
     },
 
+    {
+     FLB_CONFIG_MAP_BOOL, "simple_aggregation", "false",
+     0, FLB_TRUE, offsetof(struct flb_kinesis, simple_aggregation),
+     "Enable simple aggregation to combine multiple records into single API calls. "
+     "This reduces the number of requests and can improve throughput."
+    },
+
     /* EOF */
     {0}
 };

plugins/out_kinesis_streams/kinesis.h

@@ -26,6 +26,7 @@
 #include <fluent-bit/flb_http_client.h>
 #include <fluent-bit/flb_aws_util.h>
 #include <fluent-bit/flb_signv4.h>
+#include <fluent-bit/aws/flb_aws_aggregation.h>
 
 #define DEFAULT_TIME_KEY_FORMAT "%Y-%m-%dT%H:%M:%S"
 
@@ -56,6 +57,10 @@ struct flush {
     char *event_buf;
     size_t event_buf_size;
 
+    /* aggregation buffer for simple_aggregation mode */
+    struct flb_aws_agg_buffer agg_buf;
+    int agg_buf_initialized;
+
     int records_sent;
     int records_processed;
 
@@ -92,6 +97,7 @@ struct flb_kinesis {
     const char *log_key;
     const char *external_id;
     int retry_requests;
+    int simple_aggregation;
     char *sts_endpoint;
     int custom_endpoint;
     uint16_t port;

plugins/out_kinesis_streams/kinesis_api.c

@@ -36,6 +36,7 @@
 #include <fluent-bit/flb_http_client.h>
 #include <fluent-bit/flb_utils.h>
 #include <fluent-bit/flb_base64.h>
+#include <fluent-bit/aws/flb_aws_aggregation.h>
 
 #include <monkey/mk_core.h>
 #include <msgpack.h>
@@ -51,6 +52,9 @@
 
 #define ERR_CODE_EXCEEDED_THROUGHPUT "ProvisionedThroughputExceededException"
 
+/* Forward declarations */
+static int send_log_events(struct flb_kinesis *ctx, struct flush *buf);
+
 static struct flb_aws_header put_records_target_header = {
     .key = "X-Amz-Target",
     .key_len = 12,
@@ -202,6 +206,29 @@ static int end_put_payload(struct flb_kinesis *ctx, struct flush *buf,
 }
 
 
+/*
+ * Process event with simple aggregation (Kinesis Streams version)
+ * Uses shared aggregation implementation
+ */
+static int process_event_simple_aggregation(struct flb_kinesis *ctx, struct flush *buf,
+                                            const msgpack_object *obj, struct flb_time *tms,
+                                            struct flb_config *config)
+{
+    return flb_aws_aggregation_process_event(&buf->agg_buf,
+                                            buf->tmp_buf,
+                                            buf->tmp_buf_size,
+                                            &buf->tmp_buf_offset,
+                                            obj,
+                                            tms,
+                                            config,
+                                            ctx->ins,
+                                            ctx->stream_name,
+                                            ctx->log_key,
+                                            ctx->time_key,
+                                            ctx->time_key_format,
+                                            MAX_EVENT_SIZE);
+}
+
 /*
  * Processes the msgpack object
  * -1 = failure, record not added
@@ -391,6 +418,70 @@ static void reset_flush_buf(struct flb_kinesis *ctx, struct flush *buf) {
     buf->data_size += strlen(ctx->stream_name);
 }
 
+/* Finalize and send aggregated record (Kinesis Streams version - no final newline) */
+static int send_aggregated_record(struct flb_kinesis *ctx, struct flush *buf) {
+    int ret;
+    size_t agg_size;
+    size_t b64_len;
+    struct kinesis_event *event;
+
+    /* Finalize without final newline (Kinesis Streams doesn't need it) */
+    ret = flb_aws_aggregation_finalize(&buf->agg_buf, 0, &agg_size);
+    if (ret < 0) {
+        return 0;
+    }
+
+    /* Base64 encode the aggregated record */
+    size_t size = (agg_size * 1.5) + 4;
+    if (buf->event_buf == NULL || buf->event_buf_size < size) {
+        flb_free(buf->event_buf);
+        buf->event_buf = flb_malloc(size);
+        buf->event_buf_size = size;
+        if (buf->event_buf == NULL) {
+            flb_errno();
+            return -1;
+        }
+    }
+
+    ret = flb_base64_encode((unsigned char *) buf->event_buf, size, &b64_len,
+                            (unsigned char *) buf->agg_buf.agg_buf, agg_size);
+    if (ret != 0) {
+        flb_errno();
+        return -1;
+    }
+    agg_size = b64_len;
+
+    /* Copy to tmp_buf */
+    if (buf->tmp_buf_size < agg_size) {
+        flb_plg_error(ctx->ins, "Aggregated record too large for buffer");
+        flb_aws_aggregation_reset(&buf->agg_buf);
+        return 0;
+    }
+
+    memcpy(buf->tmp_buf, buf->event_buf, agg_size);
+
+    /* Create event record */
+    event = &buf->events[0];
+    event->json = buf->tmp_buf;
+    event->len = agg_size;
+    event->timestamp.tv_sec = 0;
+    event->timestamp.tv_nsec = 0;
+    buf->event_index = 1;
+
+    /* Calculate data_size for the payload */
+    buf->data_size = PUT_RECORDS_HEADER_LEN + PUT_RECORDS_FOOTER_LEN;
+    buf->data_size += strlen(ctx->stream_name);
+    buf->data_size += agg_size + PUT_RECORDS_PER_RECORD_LEN;
+
+    /* Send the aggregated record */
+    ret = send_log_events(ctx, buf);
+    
+    /* Reset aggregation buffer */
+    flb_aws_aggregation_reset(&buf->agg_buf);
+    
+    return ret;
+}
+
 /* constructs a put payload, and then sends */
 static int send_log_events(struct flb_kinesis *ctx, struct flush *buf) {
     int ret;
@@ -476,10 +567,47 @@ static int add_event(struct flb_kinesis *ctx, struct flush *buf,
     size_t event_bytes = 0;
 
     if (buf->event_index == 0) {
-        /* init */
         reset_flush_buf(ctx, buf);
     }
 
+    /* Use simple aggregation if enabled */
+    if (ctx->simple_aggregation) {
+retry_add_event_agg:
+        retry_add = FLB_FALSE;
+        ret = process_event_simple_aggregation(ctx, buf, obj, tms, config);
+        if (ret < 0) {
+            return -1;
+        }
+        else if (ret == 1) {
+            /* Buffer full - send aggregated record and retry */
+            ret = send_aggregated_record(ctx, buf);
+            if (ret == 0 && buf->agg_buf.agg_buf_offset == 0) {
+                /* No data was aggregated - discard this record to avoid infinite loop */
+                flb_plg_warn(ctx->ins, "Discarding unprocessable record (aggregation buffer empty), %s",
+                             ctx->stream_name);
+                reset_flush_buf(ctx, buf);
+                return 0;
+            }
+            reset_flush_buf(ctx, buf);
+            if (ret < 0) {
+                return -1;
+            }
+            retry_add = FLB_TRUE;
+        }
+        else if (ret == 2) {
+            flb_plg_warn(ctx->ins, "Discarding large or unprocessable record, %s",
+                         ctx->stream_name);
+            return 0;
+        }
+
+        if (retry_add == FLB_TRUE) {
+            goto retry_add_event_agg;
+        }
+
+        return 0;
+    }
+
+    /* Normal processing without aggregation */
 retry_add_event:
     retry_add = FLB_FALSE;
     ret = process_event(ctx, buf, obj, tms, config);
@@ -488,16 +616,13 @@ static int add_event(struct flb_kinesis *ctx, struct flush *buf,
     }
     else if (ret == 1) {
         if (buf->event_index <= 0) {
-            /* somehow the record was larger than our entire request buffer */
             flb_plg_warn(ctx->ins, "Discarding massive log record, %s",
                          ctx->stream_name);
-            return 0; /* discard this record and return to caller */
+            return 0;
         }
-        /* send logs and then retry the add */
         retry_add = FLB_TRUE;
         goto send;
     } else if (ret == 2) {
-        /* discard this record and return to caller */
         flb_plg_warn(ctx->ins, "Discarding large or unprocessable record, %s",
                      ctx->stream_name);
         return 0;
@@ -508,17 +633,14 @@ static int add_event(struct flb_kinesis *ctx, struct flush *buf,
 
     if ((buf->data_size + event_bytes) > PUT_RECORDS_PAYLOAD_SIZE) {
         if (buf->event_index <= 0) {
-            /* somehow the record was larger than our entire request buffer */
             flb_plg_warn(ctx->ins, "[size=%zu] Discarding massive log record, %s",
                          event_bytes, ctx->stream_name);
-            return 0; /* discard this record and return to caller */
+            return 0;
         }
-        /* do not send this event */
         retry_add = FLB_TRUE;
         goto send;
     }
 
-    /* send is not needed yet, return to caller */
     buf->data_size += event_bytes;
     buf->event_index++;
 
@@ -633,7 +755,12 @@ int process_and_send_to_kinesis(struct flb_kinesis *ctx, struct flush *buf,
     flb_log_event_decoder_destroy(&log_decoder);
 
     /* send any remaining events */
-    ret = send_log_events(ctx, buf);
+    if (ctx->simple_aggregation) {
+        ret = send_aggregated_record(ctx, buf);
+    }
+    else {
+        ret = send_log_events(ctx, buf);
+    }
     reset_flush_buf(ctx, buf);
     if (ret < 0) {
         return -1;
@@ -981,6 +1108,9 @@ int put_records(struct flb_kinesis *ctx, struct flush *buf,
 void kinesis_flush_destroy(struct flush *buf)
 {
     if (buf) {
+        if (buf->agg_buf_initialized) {
+            flb_aws_aggregation_destroy(&buf->agg_buf);
+        }
         flb_free(buf->tmp_buf);
         flb_free(buf->out_buf);
         flb_free(buf->events);