From fb8577140e2d499b8c3c382ce35588cf8314a2b8 Mon Sep 17 00:00:00 2001
From: brunobat <brunobat@gmail.com>
Date: Tue, 17 Oct 2023 17:38:28 +0100
Subject: [PATCH] OTel flush timeout

---
 .../OpenTelemetryDestroyerTest.java           | 48 +++++++++++++++++++
 .../opentelemetry/OpenTelemetryDestroyer.java | 19 +++++++-
 .../config/runtime/OTelRuntimeConfig.java     |  8 ++++
 3 files changed, 73 insertions(+), 2 deletions(-)
 create mode 100644 extensions/opentelemetry/deployment/src/test/java/io/quarkus/opentelemetry/deployment/OpenTelemetryDestroyerTest.java

diff --git a/extensions/opentelemetry/deployment/src/test/java/io/quarkus/opentelemetry/deployment/OpenTelemetryDestroyerTest.java b/extensions/opentelemetry/deployment/src/test/java/io/quarkus/opentelemetry/deployment/OpenTelemetryDestroyerTest.java
new file mode 100644
index 0000000000000..d0fe61b90b23d
--- /dev/null
+++ b/extensions/opentelemetry/deployment/src/test/java/io/quarkus/opentelemetry/deployment/OpenTelemetryDestroyerTest.java
@@ -0,0 +1,48 @@
+package io.quarkus.opentelemetry.deployment;
+
+import static org.hamcrest.Matchers.is;
+
+import jakarta.ws.rs.GET;
+import jakarta.ws.rs.Path;
+
+import org.jboss.shrinkwrap.api.asset.StringAsset;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+import io.quarkus.opentelemetry.OpenTelemetryDestroyer;
+import io.quarkus.opentelemetry.deployment.common.TestSpanExporter;
+import io.quarkus.opentelemetry.deployment.common.TestSpanExporterProvider;
+import io.quarkus.test.QuarkusDevModeTest;
+import io.restassured.RestAssured;
+
+public class OpenTelemetryDestroyerTest {
+
+    @RegisterExtension
+    final static QuarkusDevModeTest TEST = new QuarkusDevModeTest()
+            .withApplicationRoot((jar) -> jar
+                    .addClasses(TestSpanExporter.class,
+                            TestSpanExporterProvider.class,
+                            HelloResource.class)
+                    .addAsResource(new StringAsset(TestSpanExporterProvider.class.getCanonicalName()),
+                            "META-INF/services/io.opentelemetry.sdk.autoconfigure.spi.traces.ConfigurableSpanExporterProvider")
+                    .add(new StringAsset(
+                            "quarkus.otel.traces.exporter=test-span-exporter\n" +
+                                    "quarkus.otel.experimental.shutdown-wait-time=PT60S\n"),
+                            "application.properties"));
+
+    @Test
+    void getShutdownWaitTime() {
+        RestAssured.when()
+                .get("/hello").then()
+                .statusCode(200)
+                .body(is("PT1M"));
+    }
+
+    @Path("/hello")
+    public static class HelloResource {
+        @GET
+        public String hello() {
+            return OpenTelemetryDestroyer.getShutdownWaitTime().toString();
+        }
+    }
+}
diff --git a/extensions/opentelemetry/runtime/src/main/java/io/quarkus/opentelemetry/OpenTelemetryDestroyer.java b/extensions/opentelemetry/runtime/src/main/java/io/quarkus/opentelemetry/OpenTelemetryDestroyer.java
index 67bacb2fc9a0b..b76b36bccfa1d 100644
--- a/extensions/opentelemetry/runtime/src/main/java/io/quarkus/opentelemetry/OpenTelemetryDestroyer.java
+++ b/extensions/opentelemetry/runtime/src/main/java/io/quarkus/opentelemetry/OpenTelemetryDestroyer.java
@@ -1,21 +1,36 @@
 package io.quarkus.opentelemetry;
 
+import static java.util.concurrent.TimeUnit.MILLISECONDS;
+
+import java.time.Duration;
 import java.util.Map;
 
 import jakarta.enterprise.context.spi.CreationalContext;
 
+import org.eclipse.microprofile.config.ConfigProvider;
+
 import io.opentelemetry.api.OpenTelemetry;
 import io.opentelemetry.sdk.OpenTelemetrySdk;
 import io.quarkus.arc.BeanDestroyer;
+import io.smallrye.config.SmallRyeConfig;
 
 public class OpenTelemetryDestroyer implements BeanDestroyer<OpenTelemetry> {
     @Override
     public void destroy(OpenTelemetry openTelemetry, CreationalContext<OpenTelemetry> creationalContext,
             Map<String, Object> params) {
         if (openTelemetry instanceof OpenTelemetrySdk) {
+            // between flush and shutdown we will wait shutdown-wait-time, at the most.
+            var waitTime = getShutdownWaitTime().dividedBy(2);
             var openTelemetrySdk = ((OpenTelemetrySdk) openTelemetry);
-            openTelemetrySdk.getSdkTracerProvider().forceFlush();
-            openTelemetrySdk.getSdkTracerProvider().shutdown();
+            openTelemetrySdk.getSdkTracerProvider().forceFlush().join(waitTime.toMillis(), MILLISECONDS);
+            openTelemetrySdk.getSdkTracerProvider().shutdown().join(waitTime.toMillis(), MILLISECONDS);
         }
     }
+
+    public static Duration getShutdownWaitTime() {
+        var config = ConfigProvider.getConfig().unwrap(SmallRyeConfig.class);
+        var waitTime = config.getOptionalValue("quarkus.otel.experimental.shutdown-wait-time", Duration.class)
+                .orElse(Duration.ofSeconds(1));
+        return waitTime;
+    }
 }
diff --git a/extensions/opentelemetry/runtime/src/main/java/io/quarkus/opentelemetry/runtime/config/runtime/OTelRuntimeConfig.java b/extensions/opentelemetry/runtime/src/main/java/io/quarkus/opentelemetry/runtime/config/runtime/OTelRuntimeConfig.java
index b7eb750b2e36f..d38e1fc83fe7a 100644
--- a/extensions/opentelemetry/runtime/src/main/java/io/quarkus/opentelemetry/runtime/config/runtime/OTelRuntimeConfig.java
+++ b/extensions/opentelemetry/runtime/src/main/java/io/quarkus/opentelemetry/runtime/config/runtime/OTelRuntimeConfig.java
@@ -1,5 +1,6 @@
 package io.quarkus.opentelemetry.runtime.config.runtime;
 
+import java.time.Duration;
 import java.util.List;
 import java.util.Optional;
 
@@ -63,4 +64,11 @@ public interface OTelRuntimeConfig {
      */
     @WithName("experimental.resource.disabled-keys")
     Optional<List<String>> experimentalResourceDisabledKeys();
+
+    /**
+     * The maximum amount of time Quarkus will wait for the OpenTelemetry SDK to flush unsent spans and shutdown.
+     */
+    @WithName("experimental.shutdown-wait-time")
+    @WithDefault("1s")
+    Duration experimentalShutdownWaitTime();
 }