diff --git a/api/all/src/main/java/io/opentelemetry/api/trace/ArrayBasedTraceStateBuilder.java b/api/all/src/main/java/io/opentelemetry/api/trace/ArrayBasedTraceStateBuilder.java index 09c3d1e0303..c72daddb614 100644 --- a/api/all/src/main/java/io/opentelemetry/api/trace/ArrayBasedTraceStateBuilder.java +++ b/api/all/src/main/java/io/opentelemetry/api/trace/ArrayBasedTraceStateBuilder.java @@ -92,8 +92,13 @@ public TraceStateBuilder remove(String key) { } for (int i = 0; i < reversedEntries.size(); i += 2) { if (reversedEntries.get(i).equals(key)) { - reversedEntries.set(i + 1, null); - numEntries--; + // Only account for the removal if the entry is still present. A repeated remove of an + // already-removed key must be a no-op, mirroring the guard in put(); otherwise numEntries + // is decremented twice and build() drops unrelated entries or emits a null-valued entry. + if (reversedEntries.get(i + 1) != null) { + reversedEntries.set(i + 1, null); + numEntries--; + } return this; } } diff --git a/api/all/src/test/java/io/opentelemetry/api/trace/TraceStateTest.java b/api/all/src/test/java/io/opentelemetry/api/trace/TraceStateTest.java index 22a04b4176b..a4924fb66a3 100644 --- a/api/all/src/test/java/io/opentelemetry/api/trace/TraceStateTest.java +++ b/api/all/src/test/java/io/opentelemetry/api/trace/TraceStateTest.java @@ -303,6 +303,40 @@ void removeNotPresent() { .isEqualTo(multiValueTraceState); } + @Test + void removeAlreadyRemovedKeyKeepsOtherEntries() { + TraceState state = + TraceState.builder() + .put("a", "1") + .put("b", "2") + .remove("a") + .remove("a") // removing an already-removed key must be a no-op + .build(); + assertThat(state.get("b")).isEqualTo("2"); + assertThat(state.size()).isEqualTo(1); + } + + @Test + void removeAlreadyRemovedKeyDoesNotProduceNullValue() { + TraceState state = TraceState.builder().put("a", "1").remove("a").remove("a").build(); + assertThat(state.isEmpty()).isTrue(); + assertThat(state.get("a")).isNull(); + } + + @Test + void removeAlreadyRemovedKeyWithMultipleEntriesDoesNotThrow() { + assertThatCode( + () -> + TraceState.builder() + .put("a", "1") + .put("b", "2") + .put("c", "3") + .remove("a") + .remove("a") + .build()) + .doesNotThrowAnyException(); + } + @Test void addAndRemoveEntry() { assertThat(