Address review comment

Signed-off-by: Gagan Juneja <[email protected]>
Gaganjuneja · Oct 3, 2023 · 9f46a9a · 9f46a9a
1 parent eccc611
commit 9f46a9a
Show file tree

Hide file tree

Showing 6 changed files with 24 additions and 128 deletions.
diff --git a/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/OTelSpan.java b/plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/OTelSpan.java
@@ -57,7 +57,7 @@ public void addAttribute(String key, Boolean value) {
 
     @Override
     public void setError(Exception exception) {
-        delegateSpan.setStatus(StatusCode.ERROR, exception.getMessage());
+        delegateSpan.setStatus(StatusCode.ERROR, exception != null ? exception.getMessage() : "no description");
     }
 
     @Override

diff --git a/server/src/main/java/org/opensearch/telemetry/tracing/SpanBuilder.java b/server/src/main/java/org/opensearch/telemetry/tracing/SpanBuilder.java
@@ -135,11 +135,13 @@ private static Attributes buildSpanAttributes(String action, Transport.Connectio
      * @return context
      */
     public static SpanCreationContext from(String action, TcpChannel tcpChannel) {
-        return new SpanCreationContext(createSpanName(action, tcpChannel), buildSpanAttributes(action, tcpChannel));
+        return SpanCreationContext.server().name(createSpanName(action, tcpChannel)).attributes(buildSpanAttributes(action, tcpChannel));
     }
 
     private static String createSpanName(String action, TcpChannel tcpChannel) {
-        return action + SEPARATOR + tcpChannel.getLocalAddress().getHostString();
+        return action + SEPARATOR + (tcpChannel.getRemoteAddress() != null
+            ? tcpChannel.getRemoteAddress().getHostString()
+            : tcpChannel.getLocalAddress().getHostString());
     }
 
     private static Attributes buildSpanAttributes(String action, TcpChannel tcpChannel) {

diff --git a/server/src/main/java/org/opensearch/telemetry/tracing/channels/TraceableTcpChannel.java b/server/src/main/java/org/opensearch/telemetry/tracing/channels/TraceableTcpChannel.java
diff --git a/...er/src/main/java/org/opensearch/telemetry/tracing/channels/TraceableTransportChannel.java b/...er/src/main/java/org/opensearch/telemetry/tracing/channels/TraceableTransportChannel.java
@@ -15,7 +15,6 @@
 import org.opensearch.telemetry.tracing.Span;
 import org.opensearch.telemetry.tracing.SpanScope;
 import org.opensearch.telemetry.tracing.Tracer;
-import org.opensearch.telemetry.tracing.listener.TraceableActionListener;
 import org.opensearch.transport.TcpChannel;
 import org.opensearch.transport.TransportChannel;
 
@@ -32,8 +31,6 @@ public class TraceableTransportChannel implements TransportChannel {
 
     private final TcpChannel tcpChannel;
 
-    private final static ActionListener<Void> DUMMY_ACTION_LISTENER = ActionListener.wrap(() -> {});
-
     /**
      * Constructor.
      * @param delegate delegate
@@ -57,7 +54,20 @@ public TraceableTransportChannel(TransportChannel delegate, Span span, Tracer tr
      */
     public static TransportChannel create(TransportChannel delegate, final Span span, final Tracer tracer, final TcpChannel tcpChannel) {
         if (FeatureFlags.isEnabled(FeatureFlags.TELEMETRY) == true) {
-            tcpChannel.addCloseListener(TraceableActionListener.create(DUMMY_ACTION_LISTENER, span, tracer));
+            tcpChannel.addCloseListener(new ActionListener<Void>() {
+                @Override
+                public void onResponse(Void unused) {
+                    onFailure(null);
+                }
+
+                @Override
+                public void onFailure(Exception e) {
+                    span.addEvent("The TransportChannel was closed without sending the response");
+                    span.setError(e);
+                    span.endSpan();
+                }
+            });
+
             return new TraceableTransportChannel(delegate, span, tracer, tcpChannel);
         } else {
             return delegate;

diff --git a/server/src/main/java/org/opensearch/transport/InboundHandler.java b/server/src/main/java/org/opensearch/transport/InboundHandler.java
@@ -271,6 +271,10 @@ private <T extends TransportRequest> void handleRequest(TcpChannel channel, Head
                     sendErrorResponse(action, traceableTransportChannel, e);
                 }
             }
+        } catch (Exception e) {
+            span.setError(e);
+            span.endSpan();
+            throw e;
         }
     }
 

diff --git a/test/telemetry/src/main/java/org/opensearch/test/telemetry/tracing/MockSpan.java b/test/telemetry/src/main/java/org/opensearch/test/telemetry/tracing/MockSpan.java
@@ -166,7 +166,7 @@ public Long getEndTime() {
     }
 
     public void setError(Exception exception) {
-        putMetadata("ERROR", exception.getMessage());
+        putMetadata("ERROR", exception != null ? exception.getMessage() : null);
     }
 
     private static class IdGenerator {