From 0cd65e800a59349d1122d6f9c46d77e31bcc8b74 Mon Sep 17 00:00:00 2001 From: Gagan Juneja Date: Sun, 20 Aug 2023 22:28:18 +0530 Subject: [PATCH] Fix upadting the parent span Signed-off-by: Gagan Juneja --- .../telemetry/tracing/DefaultTracer.java | 6 ++++- .../opensearch/telemetry/tracing/Span.java | 6 +++++ .../telemetry/tracing/DefaultTracerTests.java | 22 +++++++++++++++++++ .../telemetry/tracing/OTelSpan.java | 5 +++++ .../telemetry/tracing/OTelSpanTests.java | 9 ++++++++ .../test/telemetry/tracing/MockSpan.java | 5 +---- 6 files changed, 48 insertions(+), 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 index 2f3a425f96703..ccfaad073884b 100644 --- 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 index 6cb1c8234f3de..ad0ccd9feec1c 100644 --- 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 index 07abd43c8dd7b..7050f85923295 100644 --- 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 index ba63df4ae47a1..1d13aa1bb2b16 100644 --- 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 index fc92ab36908e1..64556eb8104ea 100644 --- 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 index c22a395a6e17c..2dd1cd1278a77 100644 --- 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;