Refactor TransportKafka to manage responses by thread

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;
 }

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_;
 };

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.