From 16a1bb6acf3e51563a68e21de08ed87e11502c52 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20Komor?= <mkomor@graftcode.com>
Date: Tue, 13 Oct 2026 14:52:37 +0200
Subject: [PATCH] Refactor TransportKafka to manage responses by thread

- Replaced the single lastResponse_ vector with a map to store responses by thread ID, allowing for thread-safe access and improved response handling.
- Introduced a shared mutex to synchronize access to the response map.
- Updated SendCommand and ReadResponse methods to utilize the new response management system, ensuring proper response retrieval and storage.
---
 kafka/KafkaPlugin/TransportKafka.cpp          | 25 +++++++++++++++----
 kafka/KafkaPlugin/TransportKafka.h            |  6 ++++-
 .../ServiceBusPlugin/ServiceBusClient.cpp     | 10 +++++++-
 3 files changed, 34 insertions(+), 7 deletions(-)

diff --git a/kafka/KafkaPlugin/TransportKafka.cpp b/kafka/KafkaPlugin/TransportKafka.cpp
index a25129c..0a2156a 100644
--- a/kafka/KafkaPlugin/TransportKafka.cpp
+++ b/kafka/KafkaPlugin/TransportKafka.cpp
@@ -4,6 +4,7 @@
 #include <nlohmann/json.hpp>
 
 #include <cstring>
+#include <utility>
 
 #if defined(_WIN32)
 #define KAFKA_PLUGIN_EXPORT extern "C" __declspec(dllexport)
@@ -65,16 +66,30 @@ TransportKafka::TransportKafka(const char*, unsigned short, const char* configJs
 int TransportKafka::Initialize(byte, byte, byte) { return 0; }
 
 int TransportKafka::SendCommand(byte* messageByteArray, int32_t messageByteArrayLen) {
-    lastResponse_ = client_->call(messageByteArray, static_cast<std::size_t>(messageByteArrayLen));
-    return static_cast<int>(lastResponse_.size());
+    std::vector<unsigned char> response =
+        client_->call(messageByteArray, static_cast<std::size_t>(messageByteArrayLen));
+    const int responseSize = static_cast<int>(response.size());
+    std::unique_lock<std::shared_mutex> lock(responseByThreadMutex_);
+    responseByThread_[std::this_thread::get_id()] = std::move(response);
+    return responseSize;
 }
 
 int TransportKafka::ReadResponse(byte* responseByteArray, int32_t responseByteArrayLen) {
-    if (responseByteArrayLen < static_cast<int32_t>(lastResponse_.size())) {
+    std::vector<unsigned char> response;
+    {
+        std::unique_lock<std::shared_mutex> lock(responseByThreadMutex_);
+        const auto it = responseByThread_.find(std::this_thread::get_id());
+        if (it == responseByThread_.end()) {
+            return -1;
+        }
+        response = std::move(it->second);
+        responseByThread_.erase(it);
+    }
+    if (responseByteArrayLen < static_cast<int32_t>(response.size())) {
         return -1;
     }
-    if (!lastResponse_.empty()) {
-        std::memcpy(responseByteArray, lastResponse_.data(), lastResponse_.size());
+    if (!response.empty()) {
+        std::memcpy(responseByteArray, response.data(), response.size());
     }
     return 0;
 }
diff --git a/kafka/KafkaPlugin/TransportKafka.h b/kafka/KafkaPlugin/TransportKafka.h
index e8dcbdc..6988b5b 100644
--- a/kafka/KafkaPlugin/TransportKafka.h
+++ b/kafka/KafkaPlugin/TransportKafka.h
@@ -1,7 +1,10 @@
 #pragma once
 #include "ITransport.h"
 #include "KafkaClient.h"
+#include <map>
 #include <memory>
+#include <shared_mutex>
+#include <thread>
 #include <vector>
 
 class TransportKafka : public Hypertube::Native::Interfaces::ITransport {
@@ -13,5 +16,6 @@ class TransportKafka : public Hypertube::Native::Interfaces::ITransport {
 
 private:
     std::unique_ptr<KafkaClient> client_;
-    std::vector<unsigned char> lastResponse_;
+    std::map<std::thread::id, std::vector<unsigned char>> responseByThread_;
+    std::shared_mutex responseByThreadMutex_;
 };
diff --git a/servicebus/ServiceBusPlugin/ServiceBusClient.cpp b/servicebus/ServiceBusPlugin/ServiceBusClient.cpp
index a708660..0bde8c0 100644
--- a/servicebus/ServiceBusPlugin/ServiceBusClient.cpp
+++ b/servicebus/ServiceBusPlugin/ServiceBusClient.cpp
@@ -13,11 +13,13 @@
 
 #include <nlohmann/json.hpp>
 
+#include <atomic>
 #include <chrono>
 #include <cstring>
 #include <filesystem>
 #include <fstream>
 #include <limits>
+#include <mutex>
 #include <optional>
 #include <sstream>
 #include <string>
@@ -40,11 +42,15 @@ namespace
 	// specific session. Matches the value used by the official .NET/Java SDKs.
 	constexpr const char* SessionFilterName = "com.microsoft:session-filter";
 
+	std::recursive_mutex callMutex_;
+	std::atomic<std::uint64_t> correlationCounter_{ 0 };
+
 	std::string BuildCorrelationId()
 	{
 		const auto now = std::chrono::steady_clock::now().time_since_epoch().count();
+		const auto sequence = correlationCounter_.fetch_add(1, std::memory_order_relaxed);
 		std::ostringstream oss;
-		oss << "graftcode-" << std::this_thread::get_id() << "-" << now;
+		oss << "graftcode-" << std::this_thread::get_id() << "-" << now << "-" << sequence;
 		return oss.str();
 	}
 
@@ -252,6 +258,7 @@ bool ServiceBusClient::PublishOneWay(
 	const byte* payload,
 	std::size_t payloadSize)
 {
+	std::lock_guard<std::recursive_mutex> callLock(callMutex_);
 	if (config.topic.empty()) {
 		return false;
 	}
@@ -331,6 +338,7 @@ bool ServiceBusClient::SendRpc(
 	std::size_t payloadSize,
 	std::vector<byte>& responsePayload)
 {
+	std::lock_guard<std::recursive_mutex> callLock(callMutex_);
 	responsePayload.clear();
 
 	// One-way (topic) mode: publish and return immediately without a reply.