Fix sporadic failures for smoke test NonDaemonThreads (#2289)

* Fix NonDaemonThreads sporadic failure

* Better implementation

Author: trask

agent/agent-tooling/src/main/java/com/microsoft/applicationinsights/agent/internal/init/OpenTelemetryConfigurer.java

@@ -40,6 +40,7 @@
 import com.microsoft.applicationinsights.agent.internal.telemetry.TelemetryClient;
 import io.opentelemetry.api.common.Attributes;
 import io.opentelemetry.api.common.AttributesBuilder;
+import io.opentelemetry.sdk.autoconfigure.AutoConfiguredOpenTelemetrySdkBuilder;
 import io.opentelemetry.sdk.autoconfigure.spi.AutoConfigurationCustomizer;
 import io.opentelemetry.sdk.autoconfigure.spi.AutoConfigurationCustomizerProvider;
 import io.opentelemetry.sdk.autoconfigure.spi.ConfigProperties;
@@ -48,6 +49,7 @@
 import io.opentelemetry.sdk.logs.export.BatchLogProcessor;
 import io.opentelemetry.sdk.logs.export.LogExporter;
 import io.opentelemetry.sdk.metrics.SdkMeterProviderBuilder;
+import io.opentelemetry.sdk.metrics.export.MetricReader;
 import io.opentelemetry.sdk.metrics.export.PeriodicMetricReader;
 import io.opentelemetry.sdk.trace.SdkTracerProviderBuilder;
 import io.opentelemetry.sdk.trace.data.SpanData;
@@ -59,6 +61,7 @@
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
+import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
 import javax.annotation.Nullable;
 
@@ -67,6 +70,10 @@ public class OpenTelemetryConfigurer implements AutoConfigurationCustomizerProvi
 
   @Nullable public static LoggerExporter loggerExporter;
 
+  @Nullable private static BatchLogProcessor batchLogProcessor;
+  @Nullable private static BatchSpanProcessor batchSpanProcessor;
+  @Nullable private static MetricReader metricReader;
+
   @Override
   public void customize(AutoConfigurationCustomizer autoConfiguration) {
     TelemetryClient telemetryClient = TelemetryClient.getActive();
@@ -75,6 +82,9 @@ public void customize(AutoConfigurationCustomizer autoConfiguration) {
       return;
     }
 
+    // TODO (trask) add this method to AutoConfigurationCustomizer upstream?
+    ((AutoConfiguredOpenTelemetrySdkBuilder) autoConfiguration).registerShutdownHook(false);
+
     Configuration configuration = MainEntryPoint.getConfiguration();
 
     autoConfiguration
@@ -84,6 +94,41 @@ public void customize(AutoConfigurationCustomizer autoConfiguration) {
             (builder, config) -> configureLogging(builder, telemetryClient, configuration))
         .addMeterProviderCustomizer(
             (builder, config) -> configureMetrics(builder, telemetryClient, configuration));
+
+    Runtime.getRuntime()
+        .addShutdownHook(new Thread(() -> flushAll(telemetryClient).join(10, TimeUnit.SECONDS)));
+  }
+
+  private static CompletableResultCode flushAll(TelemetryClient telemetryClient) {
+    List<CompletableResultCode> results = new ArrayList<>();
+    if (batchSpanProcessor != null) {
+      results.add(batchSpanProcessor.forceFlush());
+    }
+    if (metricReader != null) {
+      results.add(metricReader.forceFlush());
+    }
+    if (batchLogProcessor != null) {
+      results.add(batchLogProcessor.forceFlush());
+    }
+    CompletableResultCode overallResult = new CompletableResultCode();
+    CompletableResultCode initialResult = CompletableResultCode.ofAll(results);
+    initialResult.whenComplete(
+        () -> {
+          if (initialResult.isSuccess()) {
+            CompletableResultCode telemetryClientResult = telemetryClient.forceFlush();
+            telemetryClientResult.whenComplete(
+                () -> {
+                  if (telemetryClientResult.isSuccess()) {
+                    overallResult.succeed();
+                  } else {
+                    overallResult.fail();
+                  }
+                });
+          } else {
+            overallResult.fail();
+          }
+        });
+    return overallResult;
   }
 
   private static SdkTracerProviderBuilder configureTracing(
@@ -147,13 +192,12 @@ private static SdkTracerProviderBuilder configureTracing(
               telemetryClient, configuration, configuration.preview.captureHttpServer4xxAsError);
 
       // using BatchSpanProcessor in order to get off of the application thread as soon as possible
-      BatchSpanProcessor batchSpanProcessor =
+      batchSpanProcessor =
           BatchSpanProcessor.builder(spanExporter)
               .setScheduleDelay(getBatchProcessorDelay())
               .build();
 
-      tracerProvider.addSpanProcessor(
-          new TelemetryClientFlushingSpanProcessor(batchSpanProcessor, telemetryClient));
+      tracerProvider.addSpanProcessor(batchSpanProcessor);
     }
 
     return tracerProvider;
@@ -210,18 +254,15 @@ private static SdkLogEmitterProviderBuilder configureLogging(
     LogExporter logExporter = createLogExporter(telemetryClient, configuration);
 
     // using BatchLogProcessor in order to get off of the application thread as soon as possible
-    BatchLogProcessor batchLogProcessor =
+    batchLogProcessor =
         BatchLogProcessor.builder(logExporter).setScheduleDelay(getBatchProcessorDelay()).build();
 
-    TelemetryClientFlushingLogProcessor telemetryClientFlushingLogProcessor =
-        new TelemetryClientFlushingLogProcessor(batchLogProcessor, telemetryClient);
-
     // inherited attributes log processor also handles operation name, ikey and role name attributes
     // and these all need access to Span.current(), so must be run before passing off to the
     // BatchLogProcessor
     return builder.addLogProcessor(
         new InheritedAttributesLogProcessor(
-            configuration.preview.inheritedAttributes, telemetryClientFlushingLogProcessor));
+            configuration.preview.inheritedAttributes, batchLogProcessor));
   }
 
   private static LogExporter createLogExporter(
@@ -282,15 +323,12 @@ private static SdkMeterProviderBuilder configureMetrics(
       TelemetryClient telemetryClient,
       Configuration configuration) {
 
-    PeriodicMetricReader metricReader =
+    metricReader =
         PeriodicMetricReader.builder(new AzureMonitorMetricExporter(telemetryClient))
             .setInterval(Duration.ofSeconds(configuration.preview.metricIntervalSeconds))
             .build();
 
-    TelemetryClientFlushingMetricReader telemetryClientFlushingMetricReader =
-        new TelemetryClientFlushingMetricReader(metricReader, telemetryClient);
-
-    return builder.registerMetricReader(telemetryClientFlushingMetricReader);
+    return builder.registerMetricReader(metricReader);
   }
 
   private static class BackCompatHttpUrlProcessor implements SpanExporter {