Set traceId for topic sessions

alex268 · Jan 13, 2025 · 9c98639 · 9c98639
1 parent ccb2a17
commit 9c98639
Show file tree

Hide file tree

Showing 6 changed files with 87 additions and 78 deletions.
diff --git a/topic/src/main/java/tech/ydb/topic/TopicRpc.java b/topic/src/main/java/tech/ydb/topic/TopicRpc.java
@@ -75,9 +75,13 @@ CompletableFuture<Result<YdbTopic.DescribeConsumerResult>> describeConsumer(
     CompletableFuture<Status> updateOffsetsInTransaction(YdbTopic.UpdateOffsetsInTransactionRequest request,
                                                          GrpcRequestSettings settings);
 
-    GrpcReadWriteStream<YdbTopic.StreamWriteMessage.FromServer, YdbTopic.StreamWriteMessage.FromClient> writeSession();
+    GrpcReadWriteStream<YdbTopic.StreamWriteMessage.FromServer, YdbTopic.StreamWriteMessage.FromClient> writeSession(
+            String traceId
+    );
 
-    GrpcReadWriteStream<YdbTopic.StreamReadMessage.FromServer, YdbTopic.StreamReadMessage.FromClient> readSession();
+    GrpcReadWriteStream<YdbTopic.StreamReadMessage.FromServer, YdbTopic.StreamReadMessage.FromClient> readSession(
+            String traceId
+    );
 
     ScheduledExecutorService getScheduler();
 }
diff --git a/topic/src/main/java/tech/ydb/topic/impl/GrpcTopicRpc.java b/topic/src/main/java/tech/ydb/topic/impl/GrpcTopicRpc.java
@@ -91,21 +91,21 @@ public CompletableFuture<Status> updateOffsetsInTransaction(YdbTopic.UpdateOffse
     }
 
     @Override
-    public GrpcReadWriteStream<
-        YdbTopic.StreamWriteMessage.FromServer,
-        YdbTopic.StreamWriteMessage.FromClient
-        > writeSession() {
-        return transport.readWriteStreamCall(TopicServiceGrpc.getStreamWriteMethod(),
-                GrpcRequestSettings.newBuilder().build());
+    public GrpcReadWriteStream<YdbTopic.StreamWriteMessage.FromServer, YdbTopic.StreamWriteMessage.FromClient>
+            writeSession(String streamId) {
+        return transport.readWriteStreamCall(
+                TopicServiceGrpc.getStreamWriteMethod(),
+                GrpcRequestSettings.newBuilder().withTraceId(streamId).build()
+        );
     }
 
     @Override
-    public GrpcReadWriteStream<
-            YdbTopic.StreamReadMessage.FromServer,
-            YdbTopic.StreamReadMessage.FromClient
-            > readSession() {
-        return transport.readWriteStreamCall(TopicServiceGrpc.getStreamReadMethod(),
-                GrpcRequestSettings.newBuilder().build());
+    public GrpcReadWriteStream<YdbTopic.StreamReadMessage.FromServer, YdbTopic.StreamReadMessage.FromClient>
+            readSession(String streamId) {
+        return transport.readWriteStreamCall(
+                TopicServiceGrpc.getStreamReadMethod(),
+                GrpcRequestSettings.newBuilder().withTraceId(streamId).build()
+        );
     }
 
     @Override

diff --git a/topic/src/main/java/tech/ydb/topic/read/impl/ReadSession.java b/topic/src/main/java/tech/ydb/topic/read/impl/ReadSession.java
@@ -15,8 +15,11 @@
 public abstract class ReadSession extends SessionBase<FromServer, FromClient> {
     private static final Logger logger = LoggerFactory.getLogger(ReadSession.class);
 
-    public ReadSession(TopicRpc rpc) {
-        super(rpc.readSession());
+    protected final String streamId;
+
+    public ReadSession(TopicRpc rpc, String streamId) {
+        super(rpc.readSession(streamId));
+        this.streamId = streamId;
     }
 
     @Override