in_opentelemetry: reduce protobuf decode allocation churn

plugins/in_opentelemetry/opentelemetry_logs.c

@@ -18,17 +18,83 @@
  */
 
 #include <fluent-bit/flb_input_plugin.h>
+#include <fluent-bit/flb_mem.h>
 #include <fluent-bit/flb_sds.h>
 #include <fluent-bit/flb_pack.h>
 #include <fluent-bit/flb_log_event_encoder.h>
 #include <fluent-bit/flb_time.h>
 #include <fluent-bit/flb_opentelemetry.h>
 #include <fluent-otel-proto/fluent-otel.h>
 
+#include <cfl/cfl_arena.h>
 
 #include "opentelemetry.h"
 #include "opentelemetry_utils.h"
 
+#define OTEL_PROTOBUF_ARENA_MIN_PAYLOAD_SIZE 65536
+#define OTEL_PROTOBUF_ARENA_INITIAL_CHUNK_SIZE 4096
+#define OTEL_PROTOBUF_ARENA_MAX_CHUNK_SIZE 65536
+
+static void *protobuf_arena_chunk_malloc(void *context, size_t size)
+{
+    (void) context;
+
+    return flb_malloc(size);
+}
+
+static void protobuf_arena_chunk_free(void *context, void *pointer)
+{
+    (void) context;
+
+    flb_free(pointer);
+}
+
+static void *protobuf_arena_alloc(void *allocator_data, size_t size)
+{
+    struct cfl_arena *arena;
+
+    arena = allocator_data;
+    if (size == 0) {
+        size = 1;
+    }
+
+    return cfl_arena_malloc(arena, size);
+}
+
+static void protobuf_arena_free(void *allocator_data, void *pointer)
+{
+    (void) allocator_data;
+    (void) pointer;
+}
+
+static struct cfl_arena *protobuf_arena_create(void)
+{
+    struct cfl_arena_options options;
+
+    cfl_arena_options_init(&options);
+    options.chunk_size = OTEL_PROTOBUF_ARENA_INITIAL_CHUNK_SIZE;
+    options.maximum_chunk_size = OTEL_PROTOBUF_ARENA_MAX_CHUNK_SIZE;
+    options.malloc_fn = protobuf_arena_chunk_malloc;
+    options.free_fn = protobuf_arena_chunk_free;
+
+    return cfl_arena_create_with_options(&options);
+}
+
+static Opentelemetry__Proto__Collector__Logs__V1__ExportLogsServiceRequest *
+protobuf_logs_unpack(ProtobufCAllocator *allocator, size_t size, const uint8_t *data)
+{
+    return opentelemetry__proto__collector__logs__v1__export_logs_service_request__unpack(
+        allocator, size, data);
+}
+
+static void protobuf_logs_free(
+    Opentelemetry__Proto__Collector__Logs__V1__ExportLogsServiceRequest *logs,
+    ProtobufCAllocator *allocator)
+{
+    opentelemetry__proto__collector__logs__v1__export_logs_service_request__free_unpacked(
+        logs, allocator);
+}
+
 /*
  * OTLP encoding functions to pack the log records as msgpack
  * ----------------------------------------------------------
@@ -327,9 +393,12 @@ static int binary_payload_to_msgpack(struct flb_opentelemetry *ctx,
     int log_record_index;
     char *logs_body_key;
     int scope_has_schema_url;
+    struct cfl_arena *protobuf_arena;
     struct flb_mp_map_header mh;
     struct flb_mp_map_header mh_tmp;
     struct flb_time tm;
+    ProtobufCAllocator arena_allocator;
+    ProtobufCAllocator *protobuf_allocator;
 
     msgpack_packer *mp_pck;
     msgpack_packer *mp_pck_meta;
@@ -347,9 +416,22 @@ static int binary_payload_to_msgpack(struct flb_opentelemetry *ctx,
 
     mp_pck = &encoder->body.packer;
     mp_pck_meta = &encoder->metadata.packer;
+    input_logs = NULL;
+    protobuf_arena = NULL;
+    protobuf_allocator = NULL;
+
+    if (in_size >= OTEL_PROTOBUF_ARENA_MIN_PAYLOAD_SIZE) {
+        protobuf_arena = protobuf_arena_create();
+        if (protobuf_arena != NULL) {
+            arena_allocator.alloc = protobuf_arena_alloc;
+            arena_allocator.free = protobuf_arena_free;
+            arena_allocator.allocator_data = protobuf_arena;
+            protobuf_allocator = &arena_allocator;
+        }
+    }
 
     /* unpack logs from protobuf payload */
-    input_logs = opentelemetry__proto__collector__logs__v1__export_logs_service_request__unpack(NULL, in_size, in_buf);
+    input_logs = protobuf_logs_unpack(protobuf_allocator, in_size, in_buf);
     if (input_logs == NULL) {
         flb_plg_warn(ctx->ins, "failed to unpack input logs from OpenTelemetry payload");
         ret = -1;
@@ -673,9 +755,9 @@ static int binary_payload_to_msgpack(struct flb_opentelemetry *ctx,
 
  binary_payload_to_msgpack_end:
     if (input_logs) {
-        opentelemetry__proto__collector__logs__v1__export_logs_service_request__free_unpacked(
-                                            input_logs, NULL);
+        protobuf_logs_free(input_logs, protobuf_allocator);
     }
+    cfl_arena_destroy(protobuf_arena);
 
     if (ret != 0) {
         return -1;

tests/integration/scenarios/in_opentelemetry/tests/test_in_opentelemetry_001.py

@@ -678,6 +678,49 @@ def test_opentelemetry_to_opentelemetry_basic_log():
         assert item["record_attributes"]["example_key"] == "example_value"
 
 
+def test_in_opentelemetry_large_protobuf_logs():
+    service = Service("001-fluent-bit.yaml")
+    source = ExportLogsServiceRequest()
+    source.ParseFromString(service.build_otel_payload("test_logs_001.in.json", "logs"))
+
+    payload = ExportLogsServiceRequest()
+    payload.CopyFrom(source)
+    source_records = list(source.resource_logs[0].scope_logs[0].log_records)
+
+    while payload.ByteSize() < 128 * 1024:
+        payload.resource_logs[0].scope_logs[0].log_records.extend(source_records)
+
+    expected_record_count = sum(
+        len(scope_logs.log_records)
+        for resource_logs in payload.resource_logs
+        for scope_logs in resource_logs.scope_logs
+    )
+
+    service.start()
+    try:
+        logs_before = len(data_storage["logs"])
+        response = service.send_raw_request("/v1/logs", payload.SerializeToString())
+
+        assert response.status_code == 201
+        received_record_count = service.service.wait_for_condition(
+            lambda: count if (
+                count := sum(
+                    len(scope_logs.log_records)
+                    for received in data_storage["logs"][logs_before:]
+                    for resource_logs in received.resource_logs
+                    for scope_logs in resource_logs.scope_logs
+                )
+            ) >= expected_record_count else None,
+            timeout=20,
+            interval=0.25,
+            description="all large protobuf log records",
+        )
+    finally:
+        service.stop()
+
+    assert received_record_count == expected_record_count
+
+
 # Start a Fluent Bit Pipeline with Dummy message and then it gets handle by OpenTelemetry output, the config
 # aims to populate traceId and spanId fields with the values from the Dummy message.
 #
@@ -797,10 +840,14 @@ def test_opentelemetry_to_opentelemetry_parent_child_traces():
 def test_in_opentelemetry_rejects_invalid_logs_payload():
     service = Service("001-fluent-bit.yaml")
     service.start()
-    response = service.send_raw_request("/v1/logs", b"not-a-valid-otlp-payload")
-    service.stop()
+    try:
+        small_response = service.send_raw_request("/v1/logs", b"not-a-valid-otlp-payload")
+        large_response = service.send_raw_request("/v1/logs", b"\x80" * (64 * 1024))
+    finally:
+        service.stop()
 
-    assert response.status_code >= 400
+    assert small_response.status_code >= 400
+    assert large_response.status_code >= 400
     assert len(data_storage["logs"]) == 0