Fix upadting the parent span

Signed-off-by: Gagan Juneja <[email protected]>
opensearch-project · Aug 20, 2023 · 0cd65e8 · 0cd65e8
1 parent bf10ff7
commit 0cd65e8
Show file tree

Hide file tree

Showing 6 changed files with 48 additions and 5 deletions.
diff --git a/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/DefaultTracer.java b/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/DefaultTracer.java
@@ -77,7 +77,11 @@ public SpanContext getCurrentSpan() {
     private void endSpan(Span span) {
         if (span != null) {
             span.endSpan();
-            setCurrentSpanInContext(span.getParentSpan());
+            if (span.getParentSpan() != null && !span.getParentSpan().hasEnded()) {
+                setCurrentSpanInContext(span.getParentSpan());
+            } else {
+                setCurrentSpanInContext(null);
+            }
         }
     }
 

diff --git a/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/Span.java b/libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/Span.java
@@ -90,4 +90,10 @@ public interface Span {
      */
     String getSpanId();
 
+    /**
+     * Returns whether the span is ended or not.
+     * @return Span end Status.
+     */
+    boolean hasEnded();
+
 }
diff --git a/libs/telemetry/src/test/java/org/opensearch/telemetry/tracing/DefaultTracerTests.java b/libs/telemetry/src/test/java/org/opensearch/telemetry/tracing/DefaultTracerTests.java
@@ -68,6 +68,28 @@ public void testCreateSpanWithAttributesWithParentMock() {
         verify(mockTracerContextStorage, never()).get(TracerContextStorage.CURRENT_SPAN);
     }
 
+    public void testCreateSpaWithParentEndsBeforeChild() {
+        DefaultTracer defaultTracer = new DefaultTracer(mockTracingTelemetry, mockTracerContextStorage);
+        Attributes attributes = Attributes.create().addAttribute("name", "value");
+        when(mockTracingTelemetry.createSpan("span_name", mockParentSpan, attributes)).thenReturn(mockSpan);
+        when(mockParentSpan.hasEnded()).thenReturn(true);
+        SpanScope childSpanScope = defaultTracer.startSpan("span_name", new SpanContext(mockParentSpan), attributes);
+        childSpanScope.close();
+        verify(mockTracingTelemetry).createSpan("span_name", mockParentSpan, attributes);
+        verify(mockTracerContextStorage).put(TracerContextStorage.CURRENT_SPAN, null);
+    }
+
+    public void testCreateSpaWithParentEndsAfterChild() {
+        DefaultTracer defaultTracer = new DefaultTracer(mockTracingTelemetry, mockTracerContextStorage);
+        Attributes attributes = Attributes.create().addAttribute("name", "value");
+        when(mockTracingTelemetry.createSpan("span_name", mockParentSpan, attributes)).thenReturn(mockSpan);
+        when(mockParentSpan.hasEnded()).thenReturn(false);
+        SpanScope childSpanScope = defaultTracer.startSpan("span_name", new SpanContext(mockParentSpan), attributes);
+        childSpanScope.close();
+        verify(mockTracingTelemetry).createSpan("span_name", mockParentSpan, attributes);
+        verify(mockTracerContextStorage).put(TracerContextStorage.CURRENT_SPAN, mockParentSpan);
+    }
+
     public void testCreateSpanWithAttributes() {
         TracingTelemetry tracingTelemetry = new MockTracingTelemetry();
         DefaultTracer defaultTracer = new DefaultTracer(

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
@@ -69,6 +69,11 @@ public String getSpanId() {
         return delegateSpan.getSpanContext().getSpanId();
     }
 
+    @Override
+    public boolean hasEnded() {
+        return !delegateSpan.isRecording();
+    }
+
     io.opentelemetry.api.trace.Span getDelegateSpan() {
         return delegateSpan;
     }

diff --git a/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/OTelSpanTests.java b/plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/OTelSpanTests.java
@@ -31,6 +31,15 @@ public void testEndSpanTest() {
         verify(mockSpan).end();
     }
 
+    public void testHasEndSpanTest() {
+        Span mockSpan = getMockSpan();
+        when(mockSpan.isRecording()).thenReturn(false);
+        OTelSpan oTelSpan = new OTelSpan("spanName", mockSpan, null);
+        oTelSpan.endSpan();
+        assertTrue(oTelSpan.hasEnded());
+        verify(mockSpan).end();
+    }
+
     public void testAddAttributeString() {
         Span mockSpan = getMockSpan();
         OTelSpan oTelSpan = new OTelSpan("spanName", mockSpan, null);

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
@@ -141,10 +141,7 @@ public String getSpanId() {
         return spanId;
     }
 
-    /**
-     * Returns whether the span is ended or not.
-     * @return span end status.
-     */
+    @Override
     public boolean hasEnded() {
         synchronized (lock) {
             return hasEnded;