Fix ChannelFetcher memory leak

Author: trask

core/src/main/java/com/microsoft/applicationinsights/TelemetryClient.java

@@ -49,17 +49,11 @@
  * General overview https://docs.microsoft.com/azure/application-insights/app-insights-api-custom-events-metrics
  */
 public class TelemetryClient {
-    private final class TelemetryClientChannelFetcher implements ChannelFetcher {
-        public TelemetryChannel fetch() {
-            return getChannel();
-        }
-    }
 
     private final TelemetryConfiguration configuration;
     private volatile TelemetryContext context;
     private TelemetryChannel channel;
 
-    private static final Object TELEMETRY_STOP_HOOK_LOCK = new Object();
     private static final Object TELEMETRY_CONTEXT_LOCK = new Object();
 
     private static AtomicLong generateCounter = new AtomicLong(0);
@@ -72,10 +66,6 @@ public TelemetryClient(TelemetryConfiguration configuration) {
             configuration = TelemetryConfiguration.getActive();
         }
 
-        synchronized (TELEMETRY_STOP_HOOK_LOCK) {
-            SDKShutdownActivity.INSTANCE.register(new TelemetryClientChannelFetcher());
-        }
-
         this.configuration = configuration;
     }
 
@@ -507,6 +497,7 @@ public void flush() {
     TelemetryChannel getChannel() {
         if (this.channel == null) {
             this.channel = configuration.getChannel();
+            SDKShutdownActivity.INSTANCE.register(this.channel);
         }
 
         return this.channel;

core/src/main/java/com/microsoft/applicationinsights/internal/shutdown/SDKShutdownActivity.java

@@ -24,6 +24,8 @@
 import java.io.Closeable;
 import java.util.List;
 import java.util.ArrayList;
+import java.util.Map;
+import java.util.WeakHashMap;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.TimeUnit;
 
@@ -51,10 +53,16 @@ public enum SDKShutdownActivity {
     private static class SDKShutdownAction implements Runnable {
         private boolean stopped = false;
 
+        private final Map<TelemetryChannel, Boolean> channels = new WeakHashMap<>();
+        @Deprecated
         private final List<ChannelFetcher> fetchers = new ArrayList<ChannelFetcher>();
         private final List<Stoppable> stoppables = new ArrayList<Stoppable>();
         private final List<Closeable> closeables = new ArrayList<Closeable>();
 
+        public synchronized void register(TelemetryChannel channel) {
+            channels.put(channel, true);
+        }
+
         public synchronized void register(ChannelFetcher fetcher) {
             fetchers.add(fetcher);
         }
@@ -126,23 +134,29 @@ private void stopInternalLogger() {
          * Make sure no exception is thrown!
          */
         private void stopChannels() {
+            for (TelemetryChannel channel : channels.keySet()) {
+                stopChannel(channel);
+            }
             for (ChannelFetcher fetcher : fetchers) {
+                stopChannel(fetcher.fetch());
+            }
+        }
+
+        private void stopChannel(TelemetryChannel channelToStop) {
+            try {
+                if (channelToStop != null) {
+                    channelToStop.stop(getPerThreadTimeout(), getPerThreadTimeUnit());
+                }
+            } catch (ThreadDeath td) {
+                throw td;
+            } catch (Throwable t) {
                 try {
-                    TelemetryChannel channelToStop = fetcher.fetch();
-                    if (channelToStop != null) {
-                        channelToStop.stop(getPerThreadTimeout(), getPerThreadTimeUnit());
-                    }
+                    InternalLogger.INSTANCE.error("Failed to stop channel: '%s'", t.toString());
+                    InternalLogger.INSTANCE.trace("Stack trace generated is %s", ExceptionUtils.getStackTrace(t));
                 } catch (ThreadDeath td) {
                     throw td;
-                } catch (Throwable t) {
-                    try {
-                        InternalLogger.INSTANCE.error("Failed to stop channel: '%s'", t.toString());
-                        InternalLogger.INSTANCE.trace("Stack trace generated is %s", ExceptionUtils.getStackTrace(t));
-                    } catch (ThreadDeath td) {
-                        throw td;
-                    } catch (Throwable t2) {
-                        // chomp
-                    }
+                } catch (Throwable t2) {
+                    // chomp
                 }
             }
         }
@@ -190,6 +204,10 @@ private void closeClosables() {
 
     private static volatile SDKShutdownAction shutdownAction;
 
+    public void register(TelemetryChannel channel) {
+        getShutdownAction().register(channel);
+    }
+
     public void register(ChannelFetcher fetcher) {
         getShutdownAction().register(fetcher);
     }

core/src/main/java/com/microsoft/applicationinsights/internal/util/ChannelFetcher.java

@@ -26,6 +26,7 @@
 /**
  * Created by gupele on 2/2/2015.
  */
+@Deprecated
 public interface ChannelFetcher {
     TelemetryChannel fetch();
 }

core/src/test/java/com/microsoft/applicationinsights/TelemetryClientMemoryTest.java

@@ -0,0 +1,62 @@
+/*
+ * ApplicationInsights-Java
+ * Copyright (c) Microsoft Corporation
+ * All rights reserved.
+ *
+ * MIT License
+ * Permission is hereby granted, free of charge, to any person obtaining a copy of this
+ * software and associated documentation files (the ""Software""), to deal in the Software
+ * without restriction, including without limitation the rights to use, copy, modify, merge,
+ * publish, distribute, sublicense, and/or sell copies of the Software, and to permit
+ * persons to whom the Software is furnished to do so, subject to the following conditions:
+ * The above copyright notice and this permission notice shall be included in all copies or
+ * substantial portions of the Software.
+ * THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED,
+ * INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR
+ * PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE
+ * FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
+ * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+ * DEALINGS IN THE SOFTWARE.
+ */
+
+package com.microsoft.applicationinsights;
+
+import java.util.LinkedList;
+import java.util.List;
+
+import com.microsoft.applicationinsights.channel.TelemetryChannel;
+import com.microsoft.applicationinsights.telemetry.Telemetry;
+import org.junit.*;
+import org.mockito.Matchers;
+import org.mockito.Mockito;
+import org.mockito.invocation.InvocationOnMock;
+import org.mockito.stubbing.Answer;
+
+import static org.mockito.Mockito.mock;
+
+public final class TelemetryClientMemoryTest {
+
+    @Test
+    public void shouldNotRunOOM() {
+        TelemetryConfiguration configuration = new TelemetryConfiguration();
+        configuration.setInstrumentationKey("00000000-0000-0000-0000-000000000000");
+        TelemetryChannel channel = mock(TelemetryChannel.class);
+        configuration.setChannel(channel);
+
+        final List<Telemetry> eventsSent = new LinkedList<Telemetry>();
+        // Setting the channel to add the sent telemetries to a collection, so they could be verified in tests.
+        Mockito.doAnswer(new Answer() {
+            @Override
+            public Object answer(InvocationOnMock invocation) throws Throwable {
+                Telemetry telemetry = ((Telemetry) invocation.getArguments()[0]);
+                eventsSent.add(telemetry);
+
+                return null;
+            }
+        }).when(channel).send(Matchers.any(Telemetry.class));
+
+        for (int i = 0; i < 100000000; i++) {
+            new TelemetryClient(configuration).getChannel();
+        }
+    }
+}