out_azure_logs_ingestion: split payloads over the 1 MiB request limit

plugins/out_azure_logs_ingestion/azure_logs_ingestion.c

@@ -242,72 +242,28 @@ flb_sds_t get_az_li_token(struct flb_az_li *ctx)
     return token_return;
 }
 
-static void cb_azure_logs_ingestion_flush(struct flb_event_chunk *event_chunk,
-                           struct flb_output_flush *out_flush,
-                           struct flb_input_instance *i_ins,
-                           void *out_context,
-                           struct flb_config *config)
+/* Sends one payload to the data collection endpoint */
+static int az_li_send(struct flb_az_li *ctx, flb_sds_t token,
+                      void *payload, size_t payload_size, int is_compressed)
 {
     int ret;
     int flush_status;
     size_t b_sent;
-    size_t json_payload_size;
-    void* final_payload;
-    size_t final_payload_size;
-    flb_sds_t token;
     struct flb_connection *u_conn;
     struct flb_http_client *c = NULL;
-    int is_compressed = FLB_FALSE;
-    flb_sds_t json_payload = NULL;
-    struct flb_az_li *ctx = out_context;
-    (void) i_ins;
-    (void) config;
 
     /* Get upstream connection */
     u_conn = flb_upstream_conn_get(ctx->u_dce);
     if (!u_conn) {
-        FLB_OUTPUT_RETURN(FLB_RETRY);
-    }
-
-    /* Convert binary logs into a JSON payload */
-    ret = az_li_format(event_chunk->data, event_chunk->size,
-                       &json_payload, &json_payload_size, ctx,
-                       config);
-    if (ret == -1) {
-        flb_upstream_conn_release(u_conn);
-        FLB_OUTPUT_RETURN(FLB_ERROR);
-    }
-
-    /* Get OAuth2 token */
-    token = get_az_li_token(ctx);
-    if (!token) {
-        flush_status = FLB_RETRY;
-        goto cleanup;
-    }
-
-    /* Map buffer */
-    final_payload = json_payload;
-    final_payload_size = json_payload_size;
-    if (ctx->compress_enabled == FLB_TRUE) {
-        ret = flb_gzip_compress((void *) json_payload, json_payload_size,
-                                &final_payload, &final_payload_size);
-        if (ret == -1) {
-            flb_plg_error(ctx->ins,
-                          "cannot gzip payload, disabling compression");
-        }
-        else {
-            is_compressed = FLB_TRUE;
-            flb_plg_debug(ctx->ins, "enabled payload gzip compression");
-            /* JSON buffer will be cleared at cleanup: */
-        }
+        return FLB_RETRY;
     }
 
     /* Compose HTTP Client request */
     c = flb_http_client(u_conn, FLB_HTTP_POST, ctx->dce_u_url,
-                        final_payload, final_payload_size, NULL, 0, NULL, 0);
+                        payload, payload_size, NULL, 0, NULL, 0);
 
     if (!c) {
-        flb_plg_warn(ctx->ins, "retrying payload bytes=%lu", final_payload_size);
+        flb_plg_warn(ctx->ins, "retrying payload bytes=%lu", payload_size);
         flush_status = FLB_RETRY;
         goto cleanup;
     }
@@ -343,34 +299,149 @@ static void cb_azure_logs_ingestion_flush(struct flb_event_chunk *event_chunk,
             else {
                 flb_plg_warn(ctx->ins, "http_status=%i", c->resp.status);
             }
-            flb_plg_debug(ctx->ins, "retrying payload bytes=%lu", final_payload_size);
+            flb_plg_debug(ctx->ins, "retrying payload bytes=%lu", payload_size);
             flush_status = FLB_RETRY;
             goto cleanup;
         }
     }
 
 cleanup:
-    /* cleanup */
-    if (json_payload) {
+    if (c) {
+        flb_http_client_destroy(c);
+    }
+    flb_upstream_conn_release(u_conn);
+
+    return flush_status;
+}
+
+/*
+ * Formats records [first, last) of the chunk and sends them. The Logs
+ * Ingestion API rejects payloads above FLB_AZ_LI_MAX_PAYLOAD_SIZE, so a
+ * payload over the limit is split in two halves that are sent separately.
+ * A single record is sent as it is, whatever its size.
+ */
+static int az_li_flush_records(struct flb_az_li *ctx, struct flb_config *config,
+                               flb_sds_t token, const char *data,
+                               size_t *offsets, size_t first, size_t last)
+{
+    int ret;
+    size_t middle;
+    size_t json_payload_size;
+    size_t final_payload_size;
+    void *final_payload;
+    int is_compressed = FLB_FALSE;
+    flb_sds_t json_payload = NULL;
+
+    /* Convert binary logs into a JSON payload */
+    ret = az_li_format(data + offsets[first], offsets[last] - offsets[first],
+                       &json_payload, &json_payload_size, ctx, config);
+    if (ret == -1) {
+        return FLB_ERROR;
+    }
+
+    /* Map buffer */
+    final_payload = json_payload;
+    final_payload_size = json_payload_size;
+    if (ctx->compress_enabled == FLB_TRUE) {
+        ret = flb_gzip_compress((void *) json_payload, json_payload_size,
+                                &final_payload, &final_payload_size);
+        if (ret == -1) {
+            flb_plg_error(ctx->ins,
+                          "cannot gzip payload, disabling compression");
+            final_payload = json_payload;
+            final_payload_size = json_payload_size;
+        }
+        else {
+            is_compressed = FLB_TRUE;
+            flb_plg_debug(ctx->ins, "enabled payload gzip compression");
+        }
+    }
+
+    if (final_payload_size > FLB_AZ_LI_MAX_PAYLOAD_SIZE && last - first > 1) {
+        if (is_compressed == FLB_TRUE) {
+            flb_free(final_payload);
+        }
         flb_sds_destroy(json_payload);
+
+        middle = first + (last - first) / 2;
+        flb_plg_debug(ctx->ins, "payload of %zu records is %zu bytes, splitting it",
+                      last - first, final_payload_size);
+
+        /* Stop at the first failure, the chunk is retried as a whole */
+        ret = az_li_flush_records(ctx, config, token, data, offsets, first, middle);
+        if (ret != FLB_OK) {
+            return ret;
+        }
+        return az_li_flush_records(ctx, config, token, data, offsets, middle, last);
     }
 
-    /* release compressed payload */
+    ret = az_li_send(ctx, token, final_payload, final_payload_size, is_compressed);
+
     if (is_compressed == FLB_TRUE) {
         flb_free(final_payload);
     }
+    flb_sds_destroy(json_payload);
 
-    if (c) {
-        flb_http_client_destroy(c);
+    return ret;
+}
+
+static void cb_azure_logs_ingestion_flush(struct flb_event_chunk *event_chunk,
+                           struct flb_output_flush *out_flush,
+                           struct flb_input_instance *i_ins,
+                           void *out_context,
+                           struct flb_config *config)
+{
+    int ret;
+    int flush_status;
+    size_t count = 0;
+    size_t records;
+    size_t *offsets;
+    flb_sds_t token;
+    struct flb_log_event_decoder log_decoder;
+    struct flb_log_event log_event;
+    struct flb_az_li *ctx = out_context;
+    (void) i_ins;
+    (void) out_flush;
+
+    /* Record the end offset of every record so the chunk can be split */
+    records = flb_mp_count_log_records(event_chunk->data, event_chunk->size);
+    offsets = flb_calloc(records + 1, sizeof(size_t));
+    if (!offsets) {
+        flb_errno();
+        FLB_OUTPUT_RETURN(FLB_RETRY);
     }
-    if (u_conn) {
-        flb_upstream_conn_release(u_conn);
+
+    ret = flb_log_event_decoder_init(&log_decoder, (char *) event_chunk->data,
+                                     event_chunk->size);
+    if (ret != FLB_EVENT_DECODER_SUCCESS) {
+        flb_free(offsets);
+        FLB_OUTPUT_RETURN(FLB_ERROR);
+    }
+    while (count < records &&
+           flb_log_event_decoder_next(&log_decoder, &log_event) ==
+           FLB_EVENT_DECODER_SUCCESS) {
+        offsets[++count] = log_decoder.offset;
+    }
+    flb_log_event_decoder_destroy(&log_decoder);
+
+    if (count == 0) {
+        flb_free(offsets);
+        FLB_OUTPUT_RETURN(FLB_OK);
     }
 
-    /* destory token at last after HTTP call has finished */
-    if (token) {
-        flb_sds_destroy(token);
+    /* Get OAuth2 token */
+    token = get_az_li_token(ctx);
+    if (!token) {
+        flb_free(offsets);
+        FLB_OUTPUT_RETURN(FLB_RETRY);
     }
+
+    flush_status = az_li_flush_records(ctx, config, token, event_chunk->data,
+                                       offsets, 0, count);
+
+    flb_free(offsets);
+    /* destroy token at last after HTTP calls have finished */
+    flb_sds_destroy(token);
     FLB_OUTPUT_RETURN(flush_status);
 }
 

plugins/out_azure_logs_ingestion/azure_logs_ingestion.h

@@ -33,6 +33,8 @@
 #define FLB_AZ_LI_TLS_MODE          FLB_IO_TLS
 /* refresh token every 60 minutes */
 #define FLB_AZ_LI_TOKEN_TIMEOUT 3600
+/* maximum request body accepted by the Logs Ingestion API (1 MiB) */
+#define FLB_AZ_LI_MAX_PAYLOAD_SIZE  1048576
 
 #include <fluent-bit/flb_info.h>
 #include <fluent-bit/flb_output.h>

tests/integration/scenarios/out_azure_logs_ingestion/config/out_azure_logs_ingestion_large_chunk.yaml

@@ -0,0 +1,29 @@
+service:
+  flush: 1
+  log_level: info
+  http_server: on
+  http_port: ${FLUENT_BIT_HTTP_MONITORING_PORT}
+
+pipeline:
+  inputs:
+    - name: dummy
+      tag: azure_logs_ingestion
+      dummy: '{"message":"${LARGE_MESSAGE}"}'
+      samples: 1
+      copies: 500
+
+  outputs:
+    - name: azure_logs_ingestion
+      match: azure_logs_ingestion
+      tenant_id: suite-tenant
+      client_id: suite-client
+      client_secret: suite-secret
+      auth_url: http://127.0.0.1:${TEST_SUITE_OAUTH_PORT}/oauth/token
+      dce_url: https://localhost:${TEST_SUITE_HTTP_PORT}
+      dcr_id: dcr-suite
+      table_name: suite_CL
+      compress: off
+      tls.verify: on
+      tls.verify_hostname: on
+      tls.vhost: localhost
+      tls.ca_file: ${CERTIFICATE_TEST}

tests/integration/scenarios/out_azure_logs_ingestion/tests/test_out_azure_logs_ingestion_001.py

@@ -15,7 +15,7 @@
 
 
 class Service:
-    def __init__(self, config_file):
+    def __init__(self, config_file, extra_env=None):
         self.config_file = os.path.abspath(os.path.join(os.path.dirname(__file__), "../config", config_file))
         cert_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../in_splunk/certificate"))
         self.tls_crt_file = os.path.join(cert_dir, "certificate.pem")
@@ -28,6 +28,7 @@ def __init__(self, config_file):
             extra_env={
                 "CERTIFICATE_TEST": self.tls_crt_file,
                 "PRIVATE_KEY_TEST": self.tls_key_file,
+                **(extra_env or {}),
             },
             pre_start=self._start_receiver,
             post_stop=self._stop_receiver,
@@ -154,3 +155,45 @@ def test_out_azure_logs_ingestion_legacy_oauth2_and_payload_format():
     assert payload[0]["source"] == "dummy"
     assert payload[0]["level"] == "info"
     assert isinstance(payload[0]["@timestamp"], (int, float))
+
+
+def test_out_azure_logs_ingestion_splits_payloads_over_1mib():
+    # 500 records of about 3 KB serialize to about 1.5 MiB, above the 1 MiB
+    # limit of the Logs Ingestion API, so the chunk must go out in several
+    # requests that are each under the limit.
+    message = "x" * 3000
+    service = Service(
+        "out_azure_logs_ingestion_large_chunk.yaml",
+        extra_env={"LARGE_MESSAGE": message},
+    )
+    service.start()
+    configure_http_response(status_code=200, body={"status": "received"})
+    configure_oauth_token_response(
+        status_code=200,
+        body={"access_token": "oauth-access-token", "token_type": "Bearer", "expires_in": 300},
+    )
+
+    data_path = "/dataCollectionRules/dcr-suite/streams/Custom-suite_CL"
+
+    def data_requests():
+        return [request for request in data_storage["requests"] if request["path"] == data_path]
+
+    def all_records_received():
+        requests_seen = data_requests()
+        if sum(len(request["json"]) for request in requests_seen) >= 500:
+            return requests_seen
+        return None
+
+    requests_seen = service.service.wait_for_condition(
+        all_records_received,
+        timeout=20,
+        interval=0.5,
+        description="500 records split over several requests",
+    )
+    service.stop()
+
+    assert len(requests_seen) >= 2
+    for request in requests_seen:
+        assert len(request["raw_data"].encode("utf-8")) <= 1024 * 1024
+        assert all(record["message"] == message for record in request["json"])
+    assert sum(len(request["json"]) for request in requests_seen) == 500