Address review comment

Signed-off-by: Gagan Juneja <[email protected]>
Gaganjuneja · Oct 4, 2023 · e8b0339 · e8b0339
1 parent 6b3c132
commit e8b0339
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 10 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
@@ -17,6 +17,7 @@
 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;
 
 import java.io.IOException;
@@ -35,10 +36,9 @@ public class TraceableTcpTransportChannel extends BaseTcpTransportChannel {
      * @param delegate delegate
      * @param span span
      * @param tracer tracer
-     * @param channel channel
      */
-    public TraceableTcpTransportChannel(TransportChannel delegate, Span span, Tracer tracer, TcpChannel channel) {
-        super(channel);
+    public TraceableTcpTransportChannel(TcpTransportChannel delegate, Span span, Tracer tracer) {
+        super(delegate.getChannel());
         this.delegate = delegate;
         this.span = span;
         this.tracer = tracer;
@@ -53,7 +53,7 @@ public TraceableTcpTransportChannel(TransportChannel delegate, Span span, Tracer
      * @param tcpChannel tcpChannel
      * @return transport channel
      */
-    public static TransportChannel create(TransportChannel delegate, final Span span, final Tracer tracer, final TcpChannel tcpChannel) {
+    public static TransportChannel create(TcpTransportChannel delegate, final Span span, final Tracer tracer, final TcpChannel tcpChannel) {
         if (FeatureFlags.isEnabled(FeatureFlags.TELEMETRY) == true) {
             tcpChannel.addCloseListener(new ActionListener<Void>() {
                 @Override
@@ -69,7 +69,7 @@ public void onFailure(Exception e) {
                 }
             });
 
-            return new TraceableTcpTransportChannel(delegate, span, tracer, tcpChannel);
+            return new TraceableTcpTransportChannel(delegate, span, tracer);
         } else {
             return delegate;
         }

diff --git a/server/src/main/java/org/opensearch/transport/BaseTcpTransportChannel.java b/server/src/main/java/org/opensearch/transport/BaseTcpTransportChannel.java
@@ -29,4 +29,5 @@ public BaseTcpTransportChannel(TcpChannel channel) {
     public TcpChannel getChannel() {
         return channel;
     }
+
 }
diff --git a/server/src/main/java/org/opensearch/transport/TcpTransportChannel.java b/server/src/main/java/org/opensearch/transport/TcpTransportChannel.java
@@ -50,7 +50,6 @@ public final class TcpTransportChannel extends BaseTcpTransportChannel {
 
     private final AtomicBoolean released = new AtomicBoolean();
     private final OutboundHandler outboundHandler;
-    private final TcpChannel channel;
     private final String action;
     private final long requestId;
     private final Version version;
@@ -73,7 +72,6 @@ public final class TcpTransportChannel extends BaseTcpTransportChannel {
         super(channel);
         this.version = version;
         this.features = features;
-        this.channel = channel;
         this.outboundHandler = outboundHandler;
         this.action = action;
         this.requestId = requestId;
@@ -84,7 +82,7 @@ public final class TcpTransportChannel extends BaseTcpTransportChannel {
 
     @Override
     public String getProfileName() {
-        return channel.getProfile();
+        return getChannel().getProfile();
     }
 
     @Override
@@ -94,7 +92,7 @@ public void sendResponse(TransportResponse response) throws IOException {
                 // update outbound network time with current time before sending response over network
                 ((QuerySearchResult) response).getShardSearchRequest().setOutboundNetworkTime(System.currentTimeMillis());
             }
-            outboundHandler.sendResponse(version, features, channel, requestId, action, response, compressResponse, isHandshake);
+            outboundHandler.sendResponse(version, features, getChannel(), requestId, action, response, compressResponse, isHandshake);
         } finally {
             release(false);
         }
@@ -103,7 +101,7 @@ public void sendResponse(TransportResponse response) throws IOException {
     @Override
     public void sendResponse(Exception exception) throws IOException {
         try {
-            outboundHandler.sendErrorResponse(version, features, channel, requestId, action, exception);
+            outboundHandler.sendErrorResponse(version, features, getChannel(), requestId, action, exception);
         } finally {
             release(true);
         }
-Original file line number
+Diff line change
@@ Expand Up / @@ -29,4 +29,5 @@ public BaseTcpTransportChannel(TcpChannel channel) { @@
         public TcpChannel getChannel() {
             return channel;
         }
     }