From 5f90b0352ed91fad01e7c33c0a3b6853d318f029 Mon Sep 17 00:00:00 2001
From: jack-berg <34418638+jack-berg@users.noreply.github.com>
Date: Tue, 7 Jan 2025 13:04:44 -0600
Subject: [PATCH] Fix span setStatus (#6990)

---
 .../io/opentelemetry/sdk/trace/SdkSpan.java   | 18 ++++-
 .../opentelemetry/sdk/trace/SdkSpanTest.java  | 67 ++++++++++++++++---
 2 files changed, 76 insertions(+), 9 deletions(-)

diff --git a/sdk/trace/src/main/java/io/opentelemetry/sdk/trace/SdkSpan.java b/sdk/trace/src/main/java/io/opentelemetry/sdk/trace/SdkSpan.java
index 59eec41179f..f3821c55540 100644
--- a/sdk/trace/src/main/java/io/opentelemetry/sdk/trace/SdkSpan.java
+++ b/sdk/trace/src/main/java/io/opentelemetry/sdk/trace/SdkSpan.java
@@ -434,10 +434,26 @@ public ReadWriteSpan setStatus(StatusCode statusCode, @Nullable String descripti
       if (!isModifiableByCurrentThread()) {
         logger.log(Level.FINE, "Calling setStatus() on an ended Span.");
         return this;
-      } else if (this.status.getStatusCode() == StatusCode.OK) {
+      }
+
+      // If current status is OK, ignore further attempts to change it
+      if (this.status.getStatusCode() == StatusCode.OK) {
         logger.log(Level.FINE, "Calling setStatus() on a Span that is already set to OK.");
         return this;
       }
+
+      // Ignore attempts to set status to UNSET
+      if (statusCode == StatusCode.UNSET) {
+        logger.log(Level.FINE, "Ignoring call to setStatus() with status UNSET.");
+        return this;
+      }
+
+      // Ignore description when status is not ERROR
+      if (description != null && statusCode != StatusCode.ERROR) {
+        logger.log(Level.FINE, "Ignoring setStatus() description since status is not ERROR.");
+        description = null;
+      }
+
       this.status = StatusData.create(statusCode, description);
     }
     return this;
diff --git a/sdk/trace/src/test/java/io/opentelemetry/sdk/trace/SdkSpanTest.java b/sdk/trace/src/test/java/io/opentelemetry/sdk/trace/SdkSpanTest.java
index 4b44607b1d6..180fd377b0c 100644
--- a/sdk/trace/src/test/java/io/opentelemetry/sdk/trace/SdkSpanTest.java
+++ b/sdk/trace/src/test/java/io/opentelemetry/sdk/trace/SdkSpanTest.java
@@ -63,11 +63,16 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
 import java.util.stream.IntStream;
+import java.util.stream.Stream;
 import javax.annotation.Nullable;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
 import org.mockito.Mock;
 import org.mockito.Mockito;
 import org.mockito.junit.jupiter.MockitoExtension;
@@ -1336,15 +1341,61 @@ void onStartOnEndNotRequired() {
     verify(spanProcessor, never()).onEnd(any());
   }
 
-  @Test
-  void setStatusCannotOverrideStatusOK() {
+  @ParameterizedTest
+  @MethodSource("setStatusArgs")
+  void setStatus(Consumer<Span> spanConsumer, StatusData expectedSpanData) {
     SdkSpan testSpan = createTestRootSpan();
-    testSpan.setStatus(StatusCode.OK);
-    assertThat(testSpan.toSpanData().getStatus().getStatusCode()).isEqualTo(StatusCode.OK);
-    testSpan.setStatus(StatusCode.ERROR);
-    assertThat(testSpan.toSpanData().getStatus().getStatusCode()).isEqualTo(StatusCode.OK);
-    testSpan.setStatus(StatusCode.UNSET);
-    assertThat(testSpan.toSpanData().getStatus().getStatusCode()).isEqualTo(StatusCode.OK);
+    spanConsumer.accept(testSpan);
+    assertThat(testSpan.toSpanData().getStatus()).isEqualTo(expectedSpanData);
+  }
+
+  private static Stream<Arguments> setStatusArgs() {
+    return Stream.of(
+        // Default status is UNSET
+        Arguments.of(spanConsumer(span -> {}), StatusData.unset()),
+        // Simple cases
+        Arguments.of(spanConsumer(span -> span.setStatus(StatusCode.OK)), StatusData.ok()),
+        Arguments.of(spanConsumer(span -> span.setStatus(StatusCode.ERROR)), StatusData.error()),
+        // UNSET is ignored
+        Arguments.of(
+            spanConsumer(span -> span.setStatus(StatusCode.OK).setStatus(StatusCode.UNSET)),
+            StatusData.ok()),
+        Arguments.of(
+            spanConsumer(span -> span.setStatus(StatusCode.ERROR).setStatus(StatusCode.UNSET)),
+            StatusData.error()),
+        // Description is ignored unless status is ERROR
+        Arguments.of(
+            spanConsumer(span -> span.setStatus(StatusCode.UNSET, "description")),
+            StatusData.unset()),
+        Arguments.of(
+            spanConsumer(span -> span.setStatus(StatusCode.OK, "description")), StatusData.ok()),
+        Arguments.of(
+            spanConsumer(span -> span.setStatus(StatusCode.ERROR, "description")),
+            StatusData.create(StatusCode.ERROR, "description")),
+        // ERROR is ignored if status is OK
+        Arguments.of(
+            spanConsumer(
+                span -> span.setStatus(StatusCode.OK).setStatus(StatusCode.ERROR, "description")),
+            StatusData.ok()),
+        // setStatus ignored after span is ended
+        Arguments.of(
+            spanConsumer(
+                span -> {
+                  span.end();
+                  span.setStatus(StatusCode.OK);
+                }),
+            StatusData.unset()),
+        Arguments.of(
+            spanConsumer(
+                span -> {
+                  span.end();
+                  span.setStatus(StatusCode.ERROR);
+                }),
+            StatusData.unset()));
+  }
+
+  private static Consumer<Span> spanConsumer(Consumer<Span> spanConsumer) {
+    return spanConsumer;
   }
 
   private SdkSpan createTestSpanWithAttributes(Map<AttributeKey, Object> attributes) {