From 1f872e6dda1ae2a47d2c795e1ce5cbd61125d592 Mon Sep 17 00:00:00 2001
From: Attila Szakacs <attila.szakacs@axoflow.com>
Date: Fri, 11 Oct 2024 14:48:20 +0200
Subject: [PATCH] grpc/otel: support dynamic headers

Signed-off-by: Attila Szakacs <attila.szakacs@axoflow.com>
---
 modules/grpc/otel/otel-dest-worker.cpp | 24 +++++++++++-------------
 modules/grpc/otel/otel-dest-worker.hpp |  1 +
 modules/grpc/otel/otel-dest.cpp        |  6 ++++++
 modules/grpc/otel/otel-dest.hpp        |  2 +-
 4 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/modules/grpc/otel/otel-dest-worker.cpp b/modules/grpc/otel/otel-dest-worker.cpp
index 28ffeb0dd4..d25d0c3a89 100644
--- a/modules/grpc/otel/otel-dest-worker.cpp
+++ b/modules/grpc/otel/otel-dest-worker.cpp
@@ -344,6 +344,12 @@ DestWorker::insert(LogMessage *msg)
       g_assert_not_reached();
     }
 
+  if (!client_context.get())
+    {
+      client_context = std::make_unique<::grpc::ClientContext>();
+      prepare_context_dynamic(*client_context, msg);
+    }
+
   if (should_initiate_flush())
     return log_threaded_dest_worker_flush(&super->super, LTF_FLUSH_NORMAL);
 
@@ -408,11 +414,8 @@ _map_grpc_status_to_log_threaded_result(const ::grpc::Status &status)
 LogThreadedResult
 DestWorker::flush_log_records()
 {
-  ::grpc::ClientContext client_context;
-  prepare_context(client_context);
-
   logs_service_response.Clear();
-  ::grpc::Status status = logs_service_stub->Export(&client_context, logs_service_request,
+  ::grpc::Status status = logs_service_stub->Export(client_context.get(), logs_service_request,
                                                     &logs_service_response);
   owner.metrics.insert_grpc_request_stats(status);
   LogThreadedResult result = _map_grpc_status_to_log_threaded_result(status);
@@ -429,11 +432,8 @@ DestWorker::flush_log_records()
 LogThreadedResult
 DestWorker::flush_metrics()
 {
-  ::grpc::ClientContext client_context;
-  prepare_context(client_context);
-
   metrics_service_response.Clear();
-  ::grpc::Status status = metrics_service_stub->Export(&client_context, metrics_service_request,
+  ::grpc::Status status = metrics_service_stub->Export(client_context.get(), metrics_service_request,
                                                        &metrics_service_response);
   owner.metrics.insert_grpc_request_stats(status);
   LogThreadedResult result = _map_grpc_status_to_log_threaded_result(status);
@@ -450,11 +450,8 @@ DestWorker::flush_metrics()
 LogThreadedResult
 DestWorker::flush_spans()
 {
-  ::grpc::ClientContext client_context;
-  prepare_context(client_context);
-
   trace_service_response.Clear();
-  ::grpc::Status status = trace_service_stub->Export(&client_context, trace_service_request,
+  ::grpc::Status status = trace_service_stub->Export(client_context.get(), trace_service_request,
                                                      &trace_service_response);
   owner.metrics.insert_grpc_request_stats(status);
   LogThreadedResult result = _map_grpc_status_to_log_threaded_result(status);
@@ -471,7 +468,7 @@ DestWorker::flush_spans()
 LogThreadedResult
 DestWorker::flush(LogThreadedFlushMode mode)
 {
-  LogThreadedResult result;
+  LogThreadedResult result = LTR_SUCCESS;
 
   if (mode == LTF_FLUSH_EXPEDITE)
     return LTR_RETRY;
@@ -498,6 +495,7 @@ DestWorker::flush(LogThreadedFlushMode mode)
     }
 
 exit:
+  client_context.reset();
   logs_service_request.Clear();
   metrics_service_request.Clear();
   trace_service_request.Clear();
diff --git a/modules/grpc/otel/otel-dest-worker.hpp b/modules/grpc/otel/otel-dest-worker.hpp
index 62f171b068..d30a95c5d8 100644
--- a/modules/grpc/otel/otel-dest-worker.hpp
+++ b/modules/grpc/otel/otel-dest-worker.hpp
@@ -82,6 +82,7 @@ class DestWorker : public syslogng::grpc::DestWorker
 
 protected:
   std::shared_ptr<::grpc::Channel> channel;
+  std::unique_ptr<::grpc::ClientContext> client_context;
   std::unique_ptr<LogsService::Stub> logs_service_stub;
   std::unique_ptr<MetricsService::Stub> metrics_service_stub;
   std::unique_ptr<TraceService::Stub> trace_service_stub;
diff --git a/modules/grpc/otel/otel-dest.cpp b/modules/grpc/otel/otel-dest.cpp
index 2e9f388c8c..3f8422c138 100644
--- a/modules/grpc/otel/otel-dest.cpp
+++ b/modules/grpc/otel/otel-dest.cpp
@@ -28,6 +28,12 @@ using namespace syslogng::grpc::otel;
 
 /* C++ Implementations */
 
+DestDriver::DestDriver(GrpcDestDriver *s) :
+  syslogng::grpc::DestDriver(s)
+{
+  this->enable_dynamic_headers();
+}
+
 const char *
 DestDriver::generate_persist_name()
 {
diff --git a/modules/grpc/otel/otel-dest.hpp b/modules/grpc/otel/otel-dest.hpp
index 4c26ae9deb..a53ac06e91 100644
--- a/modules/grpc/otel/otel-dest.hpp
+++ b/modules/grpc/otel/otel-dest.hpp
@@ -43,7 +43,7 @@ namespace otel {
 class DestDriver: public syslogng::grpc::DestDriver
 {
 public:
-  DestDriver(GrpcDestDriver *s) : syslogng::grpc::DestDriver(s) {};
+  DestDriver(GrpcDestDriver *s);
 
   const char *format_stats_key(StatsClusterKeyBuilder *kb);
   const char *generate_persist_name();