out_loki: add compress option for gzip

Signed-off-by: Kazuhiro Suzuki <[email protected]>
fluent · Sep 21, 2023 · cf81573 · cf81573
1 parent e5f12f9
commit cf81573
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 4 deletions.
diff --git a/plugins/out_loki/loki.c b/plugins/out_loki/loki.c
@@ -28,6 +28,7 @@
 #include <fluent-bit/record_accessor/flb_ra_parser.h>
 #include <fluent-bit/flb_mp.h>
 #include <fluent-bit/flb_log_event_decoder.h>
+#include <fluent-bit/flb_gzip.h>
 
 #include <ctype.h>
 #include <sys/stat.h>
@@ -905,6 +906,7 @@ static struct flb_loki *loki_config_create(struct flb_output_instance *ins,
     int io_flags = 0;
     struct flb_loki *ctx;
     struct flb_upstream *upstream;
+    char *compress;
 
     /* Create context */
     ctx = flb_calloc(1, sizeof(struct flb_loki));
@@ -951,6 +953,15 @@ static struct flb_loki *loki_config_create(struct flb_output_instance *ins,
         }
     }
 
+    /* Compress (gzip) */
+    compress = (char *) flb_output_get_property("compress", ins);
+    ctx->compress_gzip = FLB_FALSE;
+    if (compress) {
+        if (strcasecmp(compress, "gzip") == 0) {
+            ctx->compress_gzip = FLB_TRUE;
+        }
+    }
+
     /* Line Format */
     if (strcasecmp(ctx->line_format, "json") == 0) {
         ctx->out_line_format = FLB_LOKI_FMT_JSON;
@@ -1487,6 +1498,9 @@ static void cb_loki_flush(struct flb_event_chunk *event_chunk,
     int out_ret = FLB_OK;
     size_t b_sent;
     flb_sds_t payload = NULL;
+    flb_sds_t out_buf = NULL;
+    size_t out_size;
+    int compressed = FLB_FALSE;
     struct flb_loki *ctx = out_context;
     struct flb_connection *u_conn;
     struct flb_http_client *c;
@@ -1520,31 +1534,54 @@ static void cb_loki_flush(struct flb_event_chunk *event_chunk,
                                    flb_sds_len(event_chunk->tag),
                                    event_chunk->data, event_chunk->size,
                                    &dynamic_tenant_id->value);
+
+
     if (!payload) {
         flb_plg_error(ctx->ins, "cannot compose request payload");
 
         FLB_OUTPUT_RETURN(FLB_RETRY);
     }
 
+    /* Map buffer */
+    out_buf = payload;
+    out_size = flb_sds_len(payload);
+
+    if (ctx->compress_gzip == FLB_TRUE) {
+        ret = flb_gzip_compress((void *) payload, flb_sds_len(payload), (void **) &out_buf, &out_size);
+        if (ret == -1) {
+            flb_plg_error(ctx->ins,
+                          "cannot gzip payload, disabling compression");
+        } else {
+            compressed = FLB_TRUE;
+            /* payload is not longer needed */
+            flb_sds_destroy(payload);
+        }
+    }
+
     /* Lookup an available connection context */
     u_conn = flb_upstream_conn_get(ctx->u);
     if (!u_conn) {
         flb_plg_error(ctx->ins, "no upstream connections available");
 
-        flb_sds_destroy(payload);
+        if (compressed) {
+          flb_free(out_buf);
+        }
+        else {
+          flb_sds_destroy(out_buf);
+        }
 
         FLB_OUTPUT_RETURN(FLB_RETRY);
     }
 
     /* Create HTTP client context */
     c = flb_http_client(u_conn, FLB_HTTP_POST, FLB_LOKI_URI,
-                        payload, flb_sds_len(payload),
+                        out_buf, out_size,
                         ctx->tcp_host, ctx->tcp_port,
                         NULL, 0);
     if (!c) {
         flb_plg_error(ctx->ins, "cannot create HTTP client context");
 
-        flb_sds_destroy(payload);
+        flb_sds_destroy(out_buf);
         flb_upstream_conn_release(u_conn);
 
         FLB_OUTPUT_RETURN(FLB_RETRY);
@@ -1568,6 +1605,12 @@ static void cb_loki_flush(struct flb_event_chunk *event_chunk,
                         FLB_LOKI_CT, sizeof(FLB_LOKI_CT) - 1,
                         FLB_LOKI_CT_JSON, sizeof(FLB_LOKI_CT_JSON) - 1);
 
+    if (compressed == FLB_TRUE) {
+        flb_http_add_header(c,
+                            FLB_LOKI_HEADER_CE, sizeof(FLB_LOKI_HEADER_CE) - 1,
+                            FLB_LOKI_CE_GZIP, sizeof(FLB_LOKI_CE_GZIP) - 1);
+    }
+
     /* Add X-Scope-OrgID header */
     if (dynamic_tenant_id->value != NULL) {
         flb_http_add_header(c,
@@ -1583,7 +1626,12 @@ static void cb_loki_flush(struct flb_event_chunk *event_chunk,
 
     /* Send HTTP request */
     ret = flb_http_do(c, &b_sent);
-    flb_sds_destroy(payload);
+    if (compressed) {
+      flb_free(out_buf);
+    }
+    else {
+      flb_sds_destroy(out_buf);
+    }
 
     /* Validate HTTP client return status */
     if (ret == 0) {
@@ -1760,6 +1808,12 @@ static struct flb_config_map config_map[] = {
      "Set bearer token auth"
     },
 
+    {
+     FLB_CONFIG_MAP_STR, "compress", NULL,
+     0, FLB_FALSE, 0,
+     "Set payload compression in network transfer. Option available is 'gzip'"
+    },
+
     /* EOF */
     {0}
 };

diff --git a/plugins/out_loki/loki.h b/plugins/out_loki/loki.h
@@ -32,6 +32,8 @@
 #define FLB_LOKI_HOST            "127.0.0.1"
 #define FLB_LOKI_PORT            3100
 #define FLB_LOKI_HEADER_SCOPE    "X-Scope-OrgID"
+#define FLB_LOKI_HEADER_CE       "Content-Encoding"
+#define FLB_LOKI_CE_GZIP         "gzip"
 
 #define FLB_LOKI_KV_STR    0     /* sds string */
 #define FLB_LOKI_KV_RA     1     /* record accessor */
@@ -58,6 +60,7 @@ struct flb_loki {
     flb_sds_t line_format;
     flb_sds_t tenant_id;
     flb_sds_t tenant_id_key_config;
+    int compress_gzip;
 
     /* HTTP Auth */
     flb_sds_t http_user;