AAD branch: fix shutdown hook (#1705)

* Fix flushing shutdown hook

* Introduce low-level client where retry and write-to-disk will take place

* Optimizations

* More comments

* Spotbugs

* LGTM

* Comments and renames

* Return buffers to pool on exception

* Comment out work-in-progress

* Extract class

Author: trask

agent/agent-tooling/build.gradle

@@ -60,6 +60,7 @@ dependencies {
     implementation group: 'org.apache.httpcomponents', name: 'httpclient', version: versions.apacheHttpClient
 
     implementation group: 'org.checkerframework', name: 'checker-qual', version: versions.checker
+    implementation 'com.google.code.findbugs:annotations:3.0.1'
 
     implementation(project(':core')) {
         // excluding commons-logging and replacing it with jcl-over-slf4j (below)

agent/agent-tooling/gradle/dependency-locks/compileClasspath.lockfile

@@ -16,6 +16,7 @@ com.fasterxml.jackson.datatype:jackson-datatype-jsr310:2.12.2
 com.fasterxml.jackson.module:jackson-module-jaxb-annotations:2.12.2
 com.fasterxml.jackson:jackson-bom:2.12.2
 com.fasterxml.woodstox:woodstox-core:6.2.4
+com.google.code.findbugs:annotations:3.0.1
 com.google.code.findbugs:jsr305:3.0.2
 com.google.errorprone:error_prone_annotations:2.5.1
 com.google.guava:failureaccess:1.0.1
@@ -64,6 +65,7 @@ io.projectreactor:reactor-core:3.4.3
 jakarta.activation:jakarta.activation-api:1.2.1
 jakarta.xml.bind:jakarta.xml.bind-api:2.3.2
 net.bytebuddy:byte-buddy:1.10.18
+net.jcip:jcip-annotations:1.0
 org.apache.commons:commons-lang3:3.11
 org.apache.commons:commons-text:1.9
 org.apache.httpcomponents:httpclient:4.5.13

agent/agent-tooling/gradle/dependency-locks/runtimeClasspath.lockfile

@@ -24,6 +24,7 @@ com.github.stephenc.jcip:jcip-annotations:1.0-1
 com.google.auto.service:auto-service-annotations:1.0-rc7
 com.google.auto.service:auto-service:1.0-rc7
 com.google.auto:auto-common:0.10
+com.google.code.findbugs:annotations:3.0.1
 com.google.code.findbugs:jsr305:3.0.2
 com.google.code.gson:gson:2.8.2
 com.google.errorprone:error_prone_annotations:2.5.1
@@ -99,6 +100,7 @@ net.bytebuddy:byte-buddy-agent:1.10.18
 net.bytebuddy:byte-buddy:1.10.18
 net.java.dev.jna:jna-platform:5.7.0
 net.java.dev.jna:jna:5.7.0
+net.jcip:jcip-annotations:1.0
 net.minidev:accessors-smart:1.2
 net.minidev:json-smart:2.3
 org.apache.commons:commons-lang3:3.11

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

@@ -28,6 +28,7 @@
 import com.microsoft.applicationinsights.agent.bootstrap.diagnostics.SdkVersionFinder;
 import com.microsoft.applicationinsights.agent.internal.instrumentation.sdk.*;
 import com.microsoft.applicationinsights.agent.internal.wasbootstrap.MainEntryPoint;
+import com.microsoft.applicationinsights.agent.internal.wasbootstrap.OpenTelemetryConfigurer;
 import com.microsoft.applicationinsights.agent.internal.wasbootstrap.configuration.Configuration;
 import com.microsoft.applicationinsights.agent.internal.wasbootstrap.configuration.Configuration.JmxMetric;
 import com.microsoft.applicationinsights.agent.internal.wasbootstrap.configuration.Configuration.ProcessorConfig;
@@ -46,6 +47,7 @@
 import com.microsoft.applicationinsights.profiler.config.ServiceProfilerServiceConfig;
 import io.opentelemetry.instrumentation.api.aisdk.AiLazyConfiguration;
 import io.opentelemetry.javaagent.spi.ComponentInstaller;
+import io.opentelemetry.sdk.common.CompletableResultCode;
 import org.apache.http.HttpHost;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.slf4j.Logger;
@@ -169,21 +171,7 @@ private static void start(Instrumentation instrumentation) {
 
         // this is currently used by Micrometer instrumentation in addition to 2.x SDK
         BytecodeUtil.setDelegate(new BytecodeUtilImpl());
-        Runtime.getRuntime().addShutdownHook(new Thread() {
-            @Override
-            public void run() {
-                startupLogger.debug("running shutdown hook");
-                try {
-                    telemetryClient.flush();
-                    telemetryClient.shutdown(5, SECONDS);
-                    startupLogger.debug("completed shutdown hook");
-                } catch (InterruptedException e) {
-                    startupLogger.debug("interrupted while flushing telemetry during shutdown");
-                } catch (Throwable t) {
-                    startupLogger.debug(t.getMessage(), t);
-                }
-            }
-        });
+        Runtime.getRuntime().addShutdownHook(new ShutdownHook(telemetryClient));
 
         RpConfiguration rpConfiguration = MainEntryPoint.getRpConfiguration();
         if (rpConfiguration != null) {
@@ -303,4 +291,37 @@ private static ParamXmlElement newParamXml(String name, String value) {
         paramXml.setValue(value);
         return paramXml;
     }
+
+    private static class ShutdownHook extends Thread {
+        private final TelemetryClient telemetryClient;
+
+        public ShutdownHook(TelemetryClient telemetryClient) {
+            this.telemetryClient = telemetryClient;
+        }
+
+        @Override
+        public void run() {
+            startupLogger.debug("running shutdown hook");
+            CompletableResultCode otelFlush = OpenTelemetryConfigurer.flush();
+            CompletableResultCode result = new CompletableResultCode();
+            otelFlush.whenComplete(() -> {
+                    CompletableResultCode batchingClientFlush = telemetryClient.flushChannelBatcher();
+                    batchingClientFlush.whenComplete(() -> {
+                        if (otelFlush.isSuccess() && batchingClientFlush.isSuccess()) {
+                            result.succeed();
+                        } else {
+                            result.fail();
+                        }
+                    });
+            });
+            result.join(5, SECONDS);
+            if (result.isSuccess()) {
+                startupLogger.debug("flushing telemetry on shutdown completed successfully");
+            } else if (Thread.interrupted()) {
+                startupLogger.debug("interrupted while flushing telemetry on shutdown");
+            } else {
+                startupLogger.debug("flushing telemetry on shutdown has taken more than 5 seconds, shutting down anyways...");
+            }
+        }
+    }
 }

agent/agent-tooling/src/main/java/com/microsoft/applicationinsights/agent/internal/instrumentation/sdk/BytecodeUtilImpl.java

@@ -334,7 +334,7 @@ private SeverityLevel getSeverityLevel(int value) {
     @Override
     public void flush() {
         // this is not null because sdk instrumentation is not added until Global.setTelemetryClient() is called
-        checkNotNull(Global.getTelemetryClient()).flush();
+        checkNotNull(Global.getTelemetryClient()).flushChannelBatcher();
     }
 
     @Override

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

@@ -15,14 +15,24 @@
 import com.microsoft.applicationinsights.agent.internal.propagator.DelegatingPropagator;
 import com.microsoft.applicationinsights.agent.internal.sampling.DelegatingSampler;
 import com.microsoft.applicationinsights.agent.internal.sampling.Samplers;
+import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
 import io.opentelemetry.sdk.autoconfigure.spi.SdkTracerProviderConfigurer;
+import io.opentelemetry.sdk.common.CompletableResultCode;
 import io.opentelemetry.sdk.trace.SdkTracerProviderBuilder;
 import io.opentelemetry.sdk.trace.export.BatchSpanProcessor;
 import io.opentelemetry.sdk.trace.export.SpanExporter;
 
 public class OpenTelemetryConfigurer implements SdkTracerProviderConfigurer {
 
+    private static volatile BatchSpanProcessor batchSpanProcessor;
+
+    public static CompletableResultCode flush() {
+        return batchSpanProcessor.forceFlush();
+    }
+
     @Override
+    @SuppressFBWarnings(value = "ST_WRITE_TO_STATIC_FROM_INSTANCE_METHOD",
+                        justification = "this method is only called once during initialization")
     public void configure(SdkTracerProviderBuilder tracerProvider) {
         TelemetryClient telemetryClient = Global.getTelemetryClient();
         if (telemetryClient == null) {
@@ -47,37 +57,23 @@ public void configure(SdkTracerProviderBuilder tracerProvider) {
         // Reversing the order of processors before passing it to SpanProcessor
         Collections.reverse(processors);
 
-        SpanExporter exporter = new Exporter(TelemetryClient.getActive());
+        SpanExporter currExporter = new Exporter(TelemetryClient.getActive());
 
         // NOTE if changing the span processor to something async, flush it in the shutdown hook before flushing TelemetryClient
         if (!processors.isEmpty()) {
-            SpanExporter currExporter = null;
             for (ProcessorConfig processorConfig : processors) {
-
-                if (currExporter == null) {
-                    currExporter = processorConfig.type == ProcessorType.attribute ?
-                            new ExporterWithAttributeProcessor(processorConfig, exporter) :
-                            new ExporterWithSpanProcessor(processorConfig, exporter);
-
-                } else {
-                    currExporter = processorConfig.type == ProcessorType.attribute ?
-                            new ExporterWithAttributeProcessor(processorConfig, currExporter) :
-                            new ExporterWithSpanProcessor(processorConfig, currExporter);
-                }
+                currExporter = processorConfig.type == ProcessorType.attribute ?
+                        new ExporterWithAttributeProcessor(processorConfig, currExporter) :
+                        new ExporterWithSpanProcessor(processorConfig, currExporter);
             }
-
-            // using batch size 1 here because batching is done at a lower level
-            // but still using BatchSpanProcessor in order to get off of the application thread as soon as possible
-            tracerProvider.addSpanProcessor(BatchSpanProcessor.builder(currExporter)
-                    .setMaxExportBatchSize(1)
-                    .build());
-
-        } else {
-            // using batch size 1 here because batching is done at a lower level
-            // but still using BatchSpanProcessor in order to get off of the application thread as soon as possible
-            tracerProvider.addSpanProcessor(BatchSpanProcessor.builder(exporter)
-                    .setMaxExportBatchSize(1)
-                    .build());
         }
+
+        // using BatchSpanProcessor in order to get off of the application thread as soon as possible
+        // using batch size 1 because need to convert to SpanData as soon as possible to grab data for live metrics
+        // real batching is done at a lower level
+        batchSpanProcessor = BatchSpanProcessor.builder(currExporter)
+                .setMaxExportBatchSize(1)
+                .build();
+        tracerProvider.addSpanProcessor(batchSpanProcessor);
     }
 }

agent/exporter/src/main/java/com/microsoft/applicationinsights/agent/Exporter.java

@@ -126,13 +126,13 @@ public CompletableResultCode export(Collection<SpanData> spans) {
         }
 
         try {
-            List<TelemetryItem> telemetryItems = new ArrayList<>();
             for (SpanData span : spans) {
                 logger.debug("exporting span: {}", span);
-                export(span, telemetryItems);
+                export(span);
             }
-            telemetryClient.trackAsync(telemetryItems);
-            // FIXME (trask)
+            // batching, retry, throttling, and writing to disk on failure occur downstream
+            // for simplicity not reporting back success/failure from this layer
+            // only that it was successfully delivered to the next layer
             return CompletableResultCode.ofSuccess();
         } catch (Throwable t) {
             logger.error(t.getMessage(), t);
@@ -150,30 +150,30 @@ public CompletableResultCode shutdown() {
         return CompletableResultCode.ofSuccess();
     }
 
-    private void export(SpanData span, List<TelemetryItem> telemetryItems) {
+    private void export(SpanData span) {
         SpanKind kind = span.getKind();
         String instrumentationName = span.getInstrumentationLibraryInfo().getName();
         Matcher matcher = COMPONENT_PATTERN.matcher(instrumentationName);
         String stdComponent = matcher.matches() ? matcher.group(1) : null;
         if (kind == SpanKind.INTERNAL) {
             Boolean isLog = span.getAttributes().get(AI_LOG_KEY);
             if (isLog != null && isLog) {
-                exportLogSpan(span, telemetryItems);
+                exportLogSpan(span);
             } else if ("spring-scheduling".equals(stdComponent) && !span.getParentSpanContext().isValid()) {
                 // TODO (trask) need semantic convention for determining whether to map INTERNAL to request or
                 //  dependency (or need clarification to use SERVER for this)
-                exportRequest(span, telemetryItems);
+                exportRequest(span);
             } else {
-                exportRemoteDependency(span, true, telemetryItems);
+                exportRemoteDependency(span, true);
             }
         } else if (kind == SpanKind.CLIENT || kind == SpanKind.PRODUCER) {
-            exportRemoteDependency(span, false, telemetryItems);
+            exportRemoteDependency(span, false);
         } else if (kind == SpanKind.CONSUMER && !span.getParentSpanContext().isRemote()) {
             // TODO need spec clarification, but it seems polling for messages can be CONSUMER also
             //  in which case the span will not have a remote parent and should be treated as a dependency instead of a request
-            exportRemoteDependency(span, false, telemetryItems);
+            exportRemoteDependency(span, false);
         } else if (kind == SpanKind.SERVER || kind == SpanKind.CONSUMER) {
-            exportRequest(span, telemetryItems);
+            exportRequest(span);
         } else {
             throw new UnsupportedOperationException(kind.name());
         }
@@ -195,8 +195,7 @@ private static List<TelemetryExceptionDetails> minimalParse(String errorStack) {
         return Collections.singletonList(details);
     }
 
-    private void exportRemoteDependency(SpanData span, boolean inProc,
-                                        List<TelemetryItem> telemetryItems) {
+    private void exportRemoteDependency(SpanData span, boolean inProc) {
         TelemetryItem telemetry = new TelemetryItem();
         RemoteDependencyData data = new RemoteDependencyData();
         telemetryClient.initRemoteDependencyTelemetry(telemetry, data);
@@ -228,8 +227,8 @@ private void exportRemoteDependency(SpanData span, boolean inProc,
 
         float samplingPercentage = getSamplingPercentage(span.getSpanContext().getTraceState());
         telemetry.setSampleRate(samplingPercentage);
-        telemetryItems.add(telemetry);
-        exportEvents(span, samplingPercentage, telemetryItems);
+        telemetryClient.trackAsync(telemetry);
+        exportEvents(span, samplingPercentage);
     }
 
     private static float getSamplingPercentage(TraceState traceState) {
@@ -259,16 +258,16 @@ private void applySemanticConventions(Attributes attributes, RemoteDependencyDat
         }
     }
 
-    private void exportLogSpan(SpanData span, List<TelemetryItem> telemetryItems) {
+    private void exportLogSpan(SpanData span) {
         String errorStack = span.getAttributes().get(AI_LOG_ERROR_STACK_KEY);
         if (errorStack == null) {
-            trackTrace(span, telemetryItems);
+            trackTrace(span);
         } else {
-            trackTraceAsException(span, errorStack, telemetryItems);
+            trackTraceAsException(span, errorStack);
         }
     }
 
-    private void trackTrace(SpanData span, List<TelemetryItem> telemetryItems) {
+    private void trackTrace(SpanData span) {
         Attributes attributes = span.getAttributes();
         String level = attributes.get(AI_LOG_LEVEL_KEY);
         String loggerName = attributes.get(AI_LOGGER_NAME_KEY);
@@ -292,10 +291,10 @@ private void trackTrace(SpanData span, List<TelemetryItem> telemetryItems) {
 
         float samplingPercentage = getSamplingPercentage(span.getSpanContext().getTraceState());
         telemetry.setSampleRate(samplingPercentage);
-        telemetryItems.add(telemetry);
+        telemetryClient.trackAsync(telemetry);
     }
 
-    private void trackTraceAsException(SpanData span, String errorStack, List<TelemetryItem> telemetryItems) {
+    private void trackTraceAsException(SpanData span, String errorStack) {
         Attributes attributes = span.getAttributes();
         String level = attributes.get(AI_LOG_LEVEL_KEY);
         String loggerName = attributes.get(AI_LOGGER_NAME_KEY);
@@ -318,7 +317,7 @@ private void trackTraceAsException(SpanData span, String errorStack, List<Teleme
 
         float samplingPercentage = getSamplingPercentage(span.getSpanContext().getTraceState());
         telemetry.setSampleRate(samplingPercentage);
-        telemetryItems.add(telemetry);
+        telemetryClient.trackAsync(telemetry);
     }
 
     private static void setLoggerProperties(MonitorDomain data, String level, String loggerName) {
@@ -498,7 +497,7 @@ private static int getDefaultPortForDbSystem(String dbSystem) {
         }
     }
 
-    private void exportRequest(SpanData span, List<TelemetryItem> telemetryItems) {
+    private void exportRequest(SpanData span) {
         TelemetryItem telemetry = new TelemetryItem();
         RequestData data = new RequestData();
         telemetryClient.initRequestTelemetry(telemetry, data);
@@ -585,8 +584,8 @@ private void exportRequest(SpanData span, List<TelemetryItem> telemetryItems) {
 
         float samplingPercentage = getSamplingPercentage(span.getSpanContext().getTraceState());
         telemetry.setSampleRate(samplingPercentage);
-        telemetryItems.add(telemetry);
-        exportEvents(span, samplingPercentage, telemetryItems);
+        telemetryClient.trackAsync(telemetry);
+        exportEvents(span, samplingPercentage);
     }
 
     private String getTelemetryName(SpanData span) {
@@ -611,7 +610,7 @@ private static String nullAwareConcat(String str1, String str2, String separator
         return str1 + separator + str2;
     }
 
-    private void exportEvents(SpanData span, float samplingPercentage, List<TelemetryItem> telemetryItems) {
+    private void exportEvents(SpanData span, float samplingPercentage) {
         for (EventData event : span.getEvents()) {
             boolean lettuce51 =
                     span.getInstrumentationLibraryInfo().getName().equals("io.opentelemetry.javaagent.lettuce-5.1");
@@ -635,17 +634,17 @@ private void exportEvents(SpanData span, float samplingPercentage, List<Telemetr
                 // TODO map OpenTelemetry exception to Application Insights exception better
                 String stacktrace = event.getAttributes().get(SemanticAttributes.EXCEPTION_STACKTRACE);
                 if (stacktrace != null) {
-                    trackException(stacktrace, span, operationId, span.getSpanId(), samplingPercentage, telemetryItems);
+                    trackException(stacktrace, span, operationId, span.getSpanId(), samplingPercentage);
                 }
             } else {
                 telemetry.setSampleRate(samplingPercentage);
-                telemetryItems.add(telemetry);
+                telemetryClient.trackAsync(telemetry);
             }
         }
     }
 
     private void trackException(String errorStack, SpanData span, String operationId,
-                                String id, float samplingPercentage, List<TelemetryItem> telemetryItems) {
+                                String id, float samplingPercentage) {
         TelemetryItem telemetry = new TelemetryItem();
         TelemetryExceptionData data = new TelemetryExceptionData();
         telemetryClient.initExceptionTelemetry(telemetry, data);
@@ -655,7 +654,7 @@ private void trackException(String errorStack, SpanData span, String operationId
         telemetry.setTime(getFormattedTime(span.getEndEpochNanos()));
         telemetry.setSampleRate(samplingPercentage);
         data.setExceptions(minimalParse(errorStack));
-        telemetryItems.add(telemetry);
+        telemetryClient.trackAsync(telemetry);
     }
 
     private static final long NANOSECONDS_PER_DAY = DAYS.toNanos(1);

agent/instrumentation/gradle/dependency-locks/runtimeClasspath.lockfile

@@ -25,6 +25,7 @@ com.github.stephenc.jcip:jcip-annotations:1.0-1
 com.google.auto.service:auto-service-annotations:1.0-rc7
 com.google.auto.service:auto-service:1.0-rc7
 com.google.auto:auto-common:0.10
+com.google.code.findbugs:annotations:3.0.1
 com.google.code.findbugs:jsr305:3.0.2
 com.google.code.gson:gson:2.8.2
 com.google.errorprone:error_prone_annotations:2.5.1
@@ -171,6 +172,7 @@ net.bytebuddy:byte-buddy-agent:1.10.18
 net.bytebuddy:byte-buddy:1.10.18
 net.java.dev.jna:jna-platform:5.7.0
 net.java.dev.jna:jna:5.7.0
+net.jcip:jcip-annotations:1.0
 net.minidev:accessors-smart:1.2
 net.minidev:json-smart:2.3
 org.apache.commons:commons-lang3:3.11