From f13dc761b0e487a1dcc2f561d68cf7e498d26a05 Mon Sep 17 00:00:00 2001
From: Gagan Juneja <gjjuneja@amazon.com>
Date: Thu, 5 Oct 2023 03:40:05 +0530
Subject: [PATCH] Address review comment

Signed-off-by: Gagan Juneja <gjjuneja@amazon.com>
---
 .../tracing/channels/TraceableTcpTransportChannel.java      | 6 ++----
 .../main/java/org/opensearch/transport/InboundHandler.java  | 4 ++--
 2 files changed, 4 insertions(+), 6 deletions(-)

diff --git a/server/src/main/java/org/opensearch/telemetry/tracing/channels/TraceableTcpTransportChannel.java b/server/src/main/java/org/opensearch/telemetry/tracing/channels/TraceableTcpTransportChannel.java
index 5ecea2fef1351..333e06eb037cb 100644
--- a/server/src/main/java/org/opensearch/telemetry/tracing/channels/TraceableTcpTransportChannel.java
+++ b/server/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
index f6e90aed860b3..c14a53e799319 100644
--- 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()) {