Add JVM shutdown hook to flush telemetry

trask · trask · commit e052ccac55d9 · 2020-04-06T10:05:17.000-07:00
diff --git a/agent/agent-tooling/src/main/java/com/microsoft/applicationinsights/agent/internal/BeforeAgentInstaller.java b/agent/agent-tooling/src/main/java/com/microsoft/applicationinsights/agent/internal/BeforeAgentInstaller.java
@@ -28,6 +28,10 @@
 import java.util.Map;
 import java.util.Properties;
 import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 
 import com.google.common.base.Strings;
 import com.microsoft.applicationinsights.TelemetryClient;
@@ -60,6 +64,7 @@
 import io.opentelemetry.auto.bootstrap.instrumentation.aiappid.AiAppId;
 import io.opentelemetry.auto.config.Config;
 import io.opentelemetry.auto.config.ConfigOverride;
+import io.opentelemetry.sdk.OpenTelemetrySdk;
 import org.apache.http.HttpHost;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.slf4j.Logger;
@@ -143,11 +148,26 @@ private static void start(Instrumentation instrumentation, File agentJarFile) th
         if (fixedRateSampling != null && fixedRateSampling.percentage != null) {
             Global.setFixedRateSamplingPercentage(fixedRateSampling.percentage);
         }
-        TelemetryClient telemetryClient = new TelemetryClient();
+        final TelemetryClient telemetryClient = new TelemetryClient();
         Global.setTelemetryClient(telemetryClient);
         AiAppId.setSupplier(new AppIdSupplier());
         // 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, TimeUnit.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);
+                }
+            }
+        });
     }
 
     @Nullable
diff --git a/agent/agent-tooling/src/main/java/io/opentelemetry/auto/tooling/TracerInstaller.java b/agent/agent-tooling/src/main/java/io/opentelemetry/auto/tooling/TracerInstaller.java
@@ -30,6 +30,7 @@ public static void installAgentTracer() {
                             .setSampler(new FixedRateSampler(fixedRateSamplingPercentage))
                             .build());
         }
+        // if changing the span processor to something async, flush it in the shutdown hook before flushing TelemetryClient
         OpenTelemetrySdk.getTracerProvider()
                 .addSpanProcessor(SimpleSpansProcessor.newBuilder(new Exporter(telemetryClient)).build());
     }
diff --git a/core/src/main/java/com/microsoft/applicationinsights/TelemetryClient.java b/core/src/main/java/com/microsoft/applicationinsights/TelemetryClient.java
@@ -23,6 +23,7 @@
 
 import java.util.Date;
 import java.util.Map;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicLong;
 
 import com.google.common.base.Strings;
@@ -455,6 +456,10 @@ public void flush() {
         getChannel().flush();
     }
 
+    public void shutdown(long timeout, TimeUnit timeUnit) throws InterruptedException {
+        getChannel().shutdown(timeout, timeUnit);
+    }
+
     /**
      * Gets the channel used by the client.
      */
diff --git a/core/src/main/java/com/microsoft/applicationinsights/channel/TelemetryChannel.java b/core/src/main/java/com/microsoft/applicationinsights/channel/TelemetryChannel.java
@@ -53,7 +53,7 @@ public interface TelemetryChannel {
      * @param timeout Time to try and stop
      * @param timeUnit The units of the 'timeout' parameter
      */
-    void stop(long timeout, TimeUnit timeUnit);
+    void shutdown(long timeout, TimeUnit timeUnit) throws InterruptedException;
 
     /**
      * Flushes the data that the channel might have internally.
diff --git a/core/src/main/java/com/microsoft/applicationinsights/channel/concrete/TelemetryChannelBase.java b/core/src/main/java/com/microsoft/applicationinsights/channel/concrete/TelemetryChannelBase.java
@@ -70,7 +70,6 @@ public abstract class TelemetryChannelBase<T> implements TelemetryChannel {
     private TransmitterFactory transmitterFactory;
     private AtomicLong itemsSent = new AtomicLong(0);
 
-    protected boolean stopped = false;
     protected boolean isInitailized = false;
 
     protected TelemetriesTransmitter<T> telemetriesTransmitter;
@@ -260,30 +259,9 @@ public void setDeveloperMode(boolean developerMode) {
         }
     }
 
-    /**
-     * Stops on going work
-     */
     @Override
-    public synchronized void stop(long timeout, TimeUnit timeUnit) {
-        try {
-            if (stopped) {
-                return;
-            }
-
-            telemetriesTransmitter.stop(timeout, timeUnit);
-            stopped = true;
-        } catch (ThreadDeath td) {
-            throw td;
-        } catch (Throwable t) {
-            try {
-                logger.error("Exception generated while stopping telemetry transmitter");
-                logger.trace("Exception generated while stopping telemetry transmitter", t);
-            } catch (ThreadDeath td) {
-                throw td;
-            } catch (Throwable t2) {
-                // chomp
-            }
-        }
+    public synchronized void shutdown(long timeout, TimeUnit timeUnit) throws InterruptedException {
+        telemetriesTransmitter.shutdown(timeout, timeUnit);
     }
 
     /**
diff --git a/core/src/main/java/com/microsoft/applicationinsights/internal/channel/TelemetriesTransmitter.java b/core/src/main/java/com/microsoft/applicationinsights/internal/channel/TelemetriesTransmitter.java
@@ -49,5 +49,5 @@ public interface TelemetriesFetcher<T> {
 
     boolean sendNow(Collection<T> telemetries);
 
-    void stop(long timeout, TimeUnit timeUnit);
+    void shutdown(long timeout, TimeUnit timeUnit) throws InterruptedException;
 }
diff --git a/core/src/main/java/com/microsoft/applicationinsights/internal/channel/TransmissionDispatcher.java b/core/src/main/java/com/microsoft/applicationinsights/internal/channel/TransmissionDispatcher.java
@@ -33,6 +33,6 @@
 public interface TransmissionDispatcher {
     void dispatch(Transmission transmission);
 
-    void stop(long timeout, TimeUnit timeUnit);
+    void shutdown(long timeout, TimeUnit timeUnit) throws InterruptedException;
 }
 
diff --git a/core/src/main/java/com/microsoft/applicationinsights/internal/channel/TransmissionOutputAsync.java b/core/src/main/java/com/microsoft/applicationinsights/internal/channel/TransmissionOutputAsync.java
@@ -36,6 +36,6 @@
 public interface TransmissionOutputAsync {
     boolean sendAsync(Transmission transmission);
 
-    void stop(long timeout, TimeUnit timeUnit);
+    void shutdown(long timeout, TimeUnit timeUnit) throws InterruptedException;
 }
 
diff --git a/core/src/main/java/com/microsoft/applicationinsights/internal/channel/TransmissionOutputSync.java b/core/src/main/java/com/microsoft/applicationinsights/internal/channel/TransmissionOutputSync.java
@@ -21,8 +21,6 @@
 
 package com.microsoft.applicationinsights.internal.channel;
 
-import java.util.concurrent.TimeUnit;
-
 import com.microsoft.applicationinsights.internal.channel.common.Transmission;
 
 /**
@@ -35,7 +33,5 @@
  */
 public interface TransmissionOutputSync {
     boolean sendSync(Transmission transmission);
-
-    void stop(long timeout, TimeUnit timeUnit);
 }
 
diff --git a/core/src/main/java/com/microsoft/applicationinsights/internal/channel/TransmissionsLoader.java b/core/src/main/java/com/microsoft/applicationinsights/internal/channel/TransmissionsLoader.java
@@ -21,13 +21,11 @@
 
 package com.microsoft.applicationinsights.internal.channel;
 
-import com.microsoft.applicationinsights.internal.shutdown.Stoppable;
-
-import java.util.concurrent.TimeUnit;
-
 /**
  * Created by gupele on 12/22/2014.
  */
-public interface TransmissionsLoader extends Stoppable {
+public interface TransmissionsLoader {
     boolean load(boolean waitForThreadsToStart);
+
+    void shutdown();
 }
diff --git a/core/src/main/java/com/microsoft/applicationinsights/internal/channel/common/ActiveTransmissionFileSystemOutput.java b/core/src/main/java/com/microsoft/applicationinsights/internal/channel/common/ActiveTransmissionFileSystemOutput.java
@@ -90,8 +90,8 @@ public void run() {
     }
 
     @Override
-    public void stop(long timeout, TimeUnit timeUnit) {
-        actualOutput.stop(timeout, timeUnit);
-        ThreadPoolUtils.stop(threadPool, timeout, timeUnit);
+    public void shutdown(long timeout, TimeUnit timeUnit) throws InterruptedException {
+        threadPool.shutdown();
+        threadPool.awaitTermination(timeout, timeUnit);
     }
 }
diff --git a/core/src/main/java/com/microsoft/applicationinsights/internal/channel/common/ActiveTransmissionLoader.java b/core/src/main/java/com/microsoft/applicationinsights/internal/channel/common/ActiveTransmissionLoader.java
@@ -22,7 +22,6 @@
 package com.microsoft.applicationinsights.internal.channel.common;
 
 import java.util.concurrent.CountDownLatch;
-import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 import com.google.common.base.Preconditions;
@@ -157,22 +156,9 @@ public synchronized boolean load(boolean waitForThreadsToStart) {
     }
 
     @Override
-    public void stop(long timeout, TimeUnit timeUnit) {
+    public void shutdown() {
         done.set(true);
         interruptAllThreads();
-        joinAllThreads();
-    }
-
-    private void joinAllThreads() {
-        for (Thread thread : threads) {
-            try {
-                thread.join();
-            } catch (InterruptedException e) {
-                logger.error("Interrupted during join of active transmission loader, exception: {}", e.toString());
-                Thread.currentThread().interrupt();
-                break;
-            }
-        }
     }
 
     private void interruptAllThreads() {
diff --git a/core/src/main/java/com/microsoft/applicationinsights/internal/channel/common/ActiveTransmissionNetworkOutput.java b/core/src/main/java/com/microsoft/applicationinsights/internal/channel/common/ActiveTransmissionNetworkOutput.java
@@ -96,9 +96,9 @@ public void run() {
     }
 
     @Override
-    public void stop(long timeout, TimeUnit timeUnit) {
-        actualOutput.stop(timeout, timeUnit);
-        ThreadPoolUtils.stop(outputThreads, timeout, timeUnit);
+    public void shutdown(long timeout, TimeUnit timeUnit) throws InterruptedException {
+        outputThreads.shutdown();
+        outputThreads.awaitTermination(timeout, timeUnit);
     }
 
     public int getNumberOfMaxThreads() {
diff --git a/core/src/main/java/com/microsoft/applicationinsights/internal/channel/common/NonBlockingDispatcher.java b/core/src/main/java/com/microsoft/applicationinsights/internal/channel/common/NonBlockingDispatcher.java
@@ -24,6 +24,7 @@
 import java.util.concurrent.TimeUnit;
 
 import com.google.common.base.Preconditions;
+import com.google.common.base.Stopwatch;
 import com.microsoft.applicationinsights.internal.channel.TransmissionDispatcher;
 import com.microsoft.applicationinsights.internal.channel.TransmissionOutputAsync;
 
@@ -59,9 +60,13 @@ public void dispatch(Transmission transmission) {
     }
 
     @Override
-    public void stop(long timeout, TimeUnit timeUnit) {
+    public void shutdown(long timeout, TimeUnit timeUnit) throws InterruptedException {
+        Stopwatch stopwatch = Stopwatch.createStarted();
         for (TransmissionOutputAsync output : transmissionOutputs) {
-            output.stop(timeout, timeUnit);
+            long remaining = timeout - stopwatch.elapsed(timeUnit);
+            if (remaining > 0) {
+                output.shutdown(remaining, timeUnit);
+            }
         }
     }
 }
diff --git a/core/src/main/java/com/microsoft/applicationinsights/internal/channel/common/TransmissionFileSystemOutput.java b/core/src/main/java/com/microsoft/applicationinsights/internal/channel/common/TransmissionFileSystemOutput.java
@@ -41,7 +41,6 @@
 import java.util.Collection;
 import java.util.Collections;
 import java.util.Comparator;
-import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicLong;
 
 import com.microsoft.applicationinsights.internal.channel.TransmissionOutputSync;
@@ -164,10 +163,6 @@ public boolean sendSync(Transmission transmission) {
         return true;
     }
 
-    @Override
-    public void stop(long timeout, TimeUnit timeUnit) {
-    }
-
     public Transmission fetchOldestFile() {
         try {
             Optional<File> oldestFile = fetchOldestFromCache();
diff --git a/core/src/main/java/com/microsoft/applicationinsights/internal/channel/common/TransmissionNetworkOutput.java b/core/src/main/java/com/microsoft/applicationinsights/internal/channel/common/TransmissionNetworkOutput.java
@@ -42,7 +42,6 @@
 import java.io.IOException;
 import java.net.SocketException;
 import java.net.UnknownHostException;
-import java.util.concurrent.TimeUnit;
 
 /**
  * The class is responsible for the actual sending of
@@ -127,24 +126,6 @@ public void setTransmissionDispatcher(TransmissionDispatcher transmissionDispatc
         this.transmissionDispatcher = transmissionDispatcher;
     }
 
-    /**
-     * Stops all threads from sending data.
-     *
-     * @param timeout
-     *            The timeout to wait, which is not relevant here.
-     * @param timeUnit
-     *            The time unit, which is not relevant in this method.
-     */
-    @Override
-    public synchronized void stop(long timeout, TimeUnit timeUnit) {
-        if (stopped) {
-            return;
-        }
-
-        httpClient.close();
-        stopped = true;
-    }
-
     /**
      * Tries to send a
      * {@link com.microsoft.applicationinsights.internal.channel.common.Transmission}
diff --git a/core/src/main/java/com/microsoft/applicationinsights/internal/channel/common/TransmitterImpl.java b/core/src/main/java/com/microsoft/applicationinsights/internal/channel/common/TransmitterImpl.java
@@ -28,6 +28,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 
+import com.google.common.base.Stopwatch;
 import com.microsoft.applicationinsights.internal.channel.TelemetriesTransmitter;
 import com.microsoft.applicationinsights.internal.channel.TelemetrySerializer;
 import com.microsoft.applicationinsights.internal.channel.TransmissionDispatcher;
@@ -245,9 +246,14 @@ public void run() {
     }
 
     @Override
-    public void stop(long timeout, TimeUnit timeUnit) {
-        transmissionsLoader.stop(timeout, timeUnit);
-        ThreadPoolUtils.stop(threadPool, timeout, timeUnit);
-        transmissionDispatcher.stop(timeout, timeUnit);
+    public void shutdown(long timeout, TimeUnit timeUnit) throws InterruptedException {
+        Stopwatch stopwatch = Stopwatch.createStarted();
+        transmissionsLoader.shutdown();
+        threadPool.shutdown();
+        threadPool.awaitTermination(timeout, timeUnit);
+        long remaining = timeout - stopwatch.elapsed(timeUnit);
+        if (remaining > 0) {
+            transmissionDispatcher.shutdown(remaining, timeUnit);
+        }
     }
 }
diff --git a/core/src/test/java/com/microsoft/applicationinsights/TelemetryClientTests.java b/core/src/test/java/com/microsoft/applicationinsights/TelemetryClientTests.java
@@ -141,7 +141,7 @@ public void send(Telemetry item) {
             }
 
             @Override
-            public void stop(long timeout, TimeUnit timeUnit) {
+            public void shutdown(long timeout, TimeUnit timeUnit) {
 
             }
 
diff --git a/core/src/test/java/com/microsoft/applicationinsights/TestFramework/StubTelemetryChannel.java b/core/src/test/java/com/microsoft/applicationinsights/TestFramework/StubTelemetryChannel.java
@@ -23,7 +23,7 @@ public void send(Telemetry item) {
     }
 
     @Override
-    public void stop(long timeout, TimeUnit timeUnit) {
+    public void shutdown(long timeout, TimeUnit timeUnit) {
 
     }
 
diff --git a/core/src/test/java/com/microsoft/applicationinsights/internal/channel/common/ActiveTransmissionLoaderTest.java b/core/src/test/java/com/microsoft/applicationinsights/internal/channel/common/ActiveTransmissionLoaderTest.java
@@ -129,7 +129,7 @@ private void testFilesOnDiskAreLoaded(int amount, boolean putFilesFirst) throws
 
         } finally {
             if (tested != null) {
-                tested.stop(1L, TimeUnit.SECONDS);
+                tested.shutdown();
             }
 
             if (folder != null && folder.exists()) {
diff --git a/core/src/test/java/com/microsoft/applicationinsights/internal/channel/common/ActiveTransmissionNetworkOutputTest.java b/core/src/test/java/com/microsoft/applicationinsights/internal/channel/common/ActiveTransmissionNetworkOutputTest.java
diff --git a/core/src/test/java/com/microsoft/applicationinsights/internal/channel/common/TelemetryBufferTest.java b/core/src/test/java/com/microsoft/applicationinsights/internal/channel/common/TelemetryBufferTest.java
diff --git a/core/src/test/java/com/microsoft/applicationinsights/internal/channel/common/TransmitterImplTest.java b/core/src/test/java/com/microsoft/applicationinsights/internal/channel/common/TransmitterImplTest.java

Original file line number	Diff line number	Diff line change
`@@ -30,6 +30,7 @@ public static void installAgentTracer() {`
`30`	`30`	`.setSampler(new FixedRateSampler(fixedRateSamplingPercentage))`
`31`	`31`	`.build());`
`32`	`32`	`}`
	`33`	`+ // if changing the span processor to something async, flush it in the shutdown hook before flushing TelemetryClient`
`33`	`34`	`OpenTelemetrySdk.getTracerProvider()`
`34`	`35`	`.addSpanProcessor(SimpleSpansProcessor.newBuilder(new Exporter(telemetryClient)).build());`
`35`	`36`	`}`
Original file line number	Diff line number	Diff line change
`@@ -49,5 +49,5 @@ public interface TelemetriesFetcher<T> {`
`49`	`49`
`50`	`50`	`boolean sendNow(Collection<T> telemetries);`
`51`	`51`
`52`		`- void stop(long timeout, TimeUnit timeUnit);`
	`52`	`+ void shutdown(long timeout, TimeUnit timeUnit) throws InterruptedException;`
`53`	`53`	`}`
Original file line number	Diff line number	Diff line change
`@@ -33,6 +33,6 @@`
`33`	`33`	`public interface TransmissionDispatcher {`
`34`	`34`	`void dispatch(Transmission transmission);`
`35`	`35`
`36`		`- void stop(long timeout, TimeUnit timeUnit);`
	`36`	`+ void shutdown(long timeout, TimeUnit timeUnit) throws InterruptedException;`
`37`	`37`	`}`
`38`	`38`
Original file line number	Diff line number	Diff line change
`@@ -36,6 +36,6 @@`
`36`	`36`	`public interface TransmissionOutputAsync {`
`37`	`37`	`boolean sendAsync(Transmission transmission);`
`38`	`38`
`39`		`- void stop(long timeout, TimeUnit timeUnit);`
	`39`	`+ void shutdown(long timeout, TimeUnit timeUnit) throws InterruptedException;`
`40`	`40`	`}`
`41`	`41`