Finish span in HttpResponseWrapper when closing, if the response was not parsed

Author: ygree

dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/main/java/datadog/trace/instrumentation/openai_java/HttpResponseWrapper.java

@@ -7,6 +7,7 @@
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import java.io.InputStream;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.BiConsumer;
 import org.jetbrains.annotations.NotNull;
 import org.slf4j.Logger;
@@ -43,6 +44,7 @@ public static <T> CompletableFuture<HttpResponseFor<T>> wrapFuture(
   private final HttpResponseFor<T> delegate;
   private final AgentSpan span;
   private final BiConsumer<AgentSpan, T> decorate;
+  private final AtomicBoolean finished = new AtomicBoolean(false);
 
   private HttpResponseWrapper(
       HttpResponseFor<T> delegate, AgentSpan span, BiConsumer<AgentSpan, T> decorate) {
@@ -58,6 +60,7 @@ public T parse() {
       parsed = delegate.parse();
     } catch (Throwable err) {
       DECORATE.finishSpan(span, err);
+      finished.set(true);
       throw err;
     }
     try {
@@ -66,6 +69,7 @@ public T parse() {
       log.debug("Span decorator failed", t);
     } finally {
       DECORATE.finishSpan(span, null);
+      finished.set(true);
     }
     return parsed;
   }
@@ -89,6 +93,9 @@ public InputStream body() {
 
   @Override
   public void close() {
+    if (finished.compareAndSet(false, true)) {
+      DECORATE.finishSpan(span, null);
+    }
     delegate.close();
   }
 }