Fixes ExpediaGroup#222: A memory leak in DashboardData. (ExpediaGroup…

…#224)

(cherry picked from commit d8d256f)

components/proxy/src/main/java/com/hotels/styx/admin/dashboard/DashboardData.java

@@ -25,12 +25,13 @@
 import com.google.common.eventbus.EventBus;
 import com.google.common.eventbus.Subscribe;
 import com.hotels.styx.Version;
-import com.hotels.styx.api.extension.OriginsSnapshot;
-import com.hotels.styx.api.extension.OriginsChangeListener;
 import com.hotels.styx.api.MetricRegistry;
+import com.hotels.styx.api.extension.OriginsChangeListener;
+import com.hotels.styx.api.extension.OriginsSnapshot;
 import com.hotels.styx.api.extension.service.BackendService;
 import com.hotels.styx.api.extension.service.spi.Registry;
 
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
@@ -89,6 +90,10 @@ public long publishTime() {
         return System.currentTimeMillis();
     }
 
+    void unregister() {
+        this.downstream.unregister();
+    }
+
     /**
      * Styx-related data.
      */
@@ -112,7 +117,7 @@ public String version() {
         }
 
         @JsonProperty("uptime")
-        public String uptime() {
+        String uptime() {
             return uptimeGauge == null ? null : uptimeGauge.getValue();
         }
 
@@ -130,7 +135,7 @@ public final class Downstream implements Registry.ChangeListener<BackendService>
         private final Supplier<Map<String, Integer>> responsesSupplier;
 
         private Downstream() {
-            updateBackendsFromRegistry();
+            this.backends = updateBackendsFromRegistry();
 
             this.responsesSupplier = new ResponseCodeSupplier(metrics, COUNTER, "origins.response.status", false);
             backendServicesRegistry.addListener(this);
@@ -166,14 +171,22 @@ Backend backend(String backendId) {
 
         @Override
         public void onChange(Registry.Changes<BackendService> changes) {
-            updateBackendsFromRegistry();
+            this.backends = updateBackendsFromRegistry();
         }
 
-        private void updateBackendsFromRegistry() {
-            this.backends = stream(backendServicesRegistry.get().spliterator(), false)
+        private List<Backend> updateBackendsFromRegistry() {
+            unregister();
+
+            return stream(backendServicesRegistry.get().spliterator(), false)
                     .map(Backend::new)
                     .collect(toList());
         }
+
+        void unregister() {
+            if (backends != null) {
+                backends.forEach(Backend::unregister);
+            }
+        }
     }
 
     /**
@@ -183,6 +196,8 @@ public final class Backend {
         private final String id;
         private final String name;
         private final List<Origin> origin;
+        private List<Origin> registeredOrigins;
+
         private final Supplier<Map<String, Integer>> responsesSupplier;
         private final Requests requests;
         private final List<String> status;
@@ -194,8 +209,12 @@ private Backend(BackendService application) {
             this.requests = new Requests("origins." + application.id());
 
             this.origin = application.origins().stream().map(Origin::new).collect(toList());
+            this.registeredOrigins = new ArrayList<>();
 
-            this.origin.forEach(eventBus::register);
+            this.origin.forEach(origin -> {
+                eventBus.register(origin);
+                registeredOrigins.add(origin);
+            });
 
             /* IMPORTANT NOTE: We are using guava transforms here instead of java 8 stream-map-collect because
               the guava transforms are backed by the original objects and reflect changes in them. */
@@ -206,6 +225,13 @@ private Backend(BackendService application) {
             this.responsesSupplier = new ResponseCodeSupplier(metrics, METER, prefix, true);
         }
 
+        void unregister() {
+            registeredOrigins.forEach(origin -> {
+                eventBus.unregister(origin);
+            });
+            registeredOrigins = new ArrayList<>();
+        }
+
         @JsonProperty("id")
         public String id() {
             return id;

components/proxy/src/main/java/com/hotels/styx/admin/dashboard/DashboardDataSupplier.java

@@ -59,6 +59,10 @@ public void onChange(Registry.Changes<BackendService> changes) {
     }
 
     private DashboardData updateDashboardData(Registry<BackendService> backendServices) {
+        if (this.data != null) {
+            this.data.unregister();
+        }
+
         return new DashboardData(environment.metricRegistry(), backendServices, jvmRouteName, buildInfo, environment.eventBus());
     }
 

components/proxy/src/test/java/com/hotels/styx/admin/dashboard/DashboardDataTest.java

@@ -48,7 +48,10 @@
 import static org.hamcrest.Matchers.closeTo;
 import static org.hamcrest.Matchers.containsInAnyOrder;
 import static org.hamcrest.Matchers.is;
+import static org.mockito.Matchers.any;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 public class DashboardDataTest {
@@ -292,6 +295,24 @@ public void providesOriginConnectionPoolData() {
         assertThat(connectionsPool.pending(), is(345));
     }
 
+    @Test
+    public void unsubscribesFromEventBus() {
+        EventBus eventBus = mock(EventBus.class);
+        MemoryBackedRegistry<BackendService> backendServicesRegistry = new MemoryBackedRegistry<>();
+        backendServicesRegistry.add(application("app", origin("app-01", "localhost", 9090)));
+        backendServicesRegistry.add(application("test", origin("test-01", "localhost", 9090)));
+
+        DashboardData dashbaord = new DashboardData(metricRegistry, backendServicesRegistry, "styx-prod1-presentation-01", new Version("releaseTag"), eventBus);
+
+        // Twice for each backend. One during backend construction, another from BackendServicesRegistry listener callback.
+        verify(eventBus, times(4)).register(any(DashboardData.Origin.class));
+
+        dashbaord.unregister();
+
+        verify(eventBus, times(4)).unregister(any(DashboardData.Origin.class));
+    }
+
+
     // makes the generics explicit for the compiler (to avoid having a cast every time () -> value is used).
     private <T> Gauge<T> gauge(T value) {
         return () -> value;