Address review comment

Signed-off-by: Gagan Juneja <gjjuneja@amazon.com>
opensearch-project · Oct 4, 2023 · f13dc76 · f13dc76
1 parent e8b0339
commit f13dc76
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 6 deletions.
diff --git a/...src/main/java/org/opensearch/telemetry/tracing/channels/TraceableTcpTransportChannel.java b/...src/main/java/org/opensearch/telemetry/tracing/channels/TraceableTcpTransportChannel.java
@@ -16,7 +16,6 @@
 import org.opensearch.telemetry.tracing.SpanScope;
 import org.opensearch.telemetry.tracing.Tracer;
 import org.opensearch.transport.BaseTcpTransportChannel;
-import org.opensearch.transport.TcpChannel;
 import org.opensearch.transport.TcpTransportChannel;
 import org.opensearch.transport.TransportChannel;
 
@@ -50,12 +49,11 @@ public TraceableTcpTransportChannel(TcpTransportChannel delegate, Span span, Tra
      * @param delegate delegate
      * @param span     span
      * @param tracer tracer
-     * @param tcpChannel tcpChannel
      * @return transport channel
      */
-    public static TransportChannel create(TcpTransportChannel delegate, final Span span, final Tracer tracer, final TcpChannel tcpChannel) {
+    public static TransportChannel create(TcpTransportChannel delegate, final Span span, final Tracer tracer) {
         if (FeatureFlags.isEnabled(FeatureFlags.TELEMETRY) == true) {
-            tcpChannel.addCloseListener(new ActionListener<Void>() {
+            delegate.getChannel().addCloseListener(new ActionListener<Void>() {
                 @Override
                 public void onResponse(Void unused) {
                     onFailure(null);

diff --git a/server/src/main/java/org/opensearch/transport/InboundHandler.java b/server/src/main/java/org/opensearch/transport/InboundHandler.java
@@ -211,7 +211,7 @@ private <T extends TransportRequest> void handleRequest(TcpChannel channel, Head
                     header.isHandshake(),
                     message.takeBreakerReleaseControl()
                 );
-                TransportChannel traceableTransportChannel = TraceableTcpTransportChannel.create(transportChannel, span, tracer, channel);
+                TransportChannel traceableTransportChannel = TraceableTcpTransportChannel.create(transportChannel, span, tracer);
                 try {
                     handshaker.handleHandshake(traceableTransportChannel, requestId, stream);
                 } catch (Exception e) {
@@ -241,7 +241,7 @@ private <T extends TransportRequest> void handleRequest(TcpChannel channel, Head
                     header.isHandshake(),
                     message.takeBreakerReleaseControl()
                 );
-                TransportChannel traceableTransportChannel = TraceableTcpTransportChannel.create(transportChannel, span, tracer, channel);
+                TransportChannel traceableTransportChannel = TraceableTcpTransportChannel.create(transportChannel, span, tracer);
                 try {
                     messageListener.onRequestReceived(requestId, action);
                     if (message.isShortCircuit()) {