Delegate responsibility/timing of cleaning up old components to caller

5 files changed, 39 insertions, 16 deletions

diff --git a/container-core/src/main/java/com/yahoo/container/core/config/HandlersConfigurerDi.java b/container-core/src/main/java/com/yahoo/container/core/config/HandlersConfigurerDi.java
index be4bc556dde..7148ed6fb45 100644
--- a/container-core/src/main/java/com/yahoo/container/core/config/HandlersConfigurerDi.java
+++ b/container-core/src/main/java/com/yahoo/container/core/config/HandlersConfigurerDi.java
@@ -71,7 +71,8 @@ public class HandlersConfigurerDi {
 
         this.vespaContainer = vespaContainer;
         container = new Container(subscriberFactory, configId, deconstructor, osgiWrapper);
-        waitForNextComponentGeneration(discInjector, true);
+        Runnable cleanupTask = waitForNextComponentGeneration(discInjector, true);
+        cleanupTask.run();
     }
 
     private static class ContainerAndDiOsgi extends OsgiImpl implements OsgiWrapper {
@@ -106,11 +107,15 @@ public class HandlersConfigurerDi {
 
     /**
      * Wait for new config to arrive and produce the new graph
+     * @return Task for deconstructing previous component graph and bundles
      */
-    public void waitForNextComponentGeneration(Injector discInjector, boolean isInitializing) {
-        currentGraph = container.waitForNextComponentGeneration(currentGraph,
-                                                      createFallbackInjector(vespaContainer, discInjector),
-                                                      isInitializing);
+    public Runnable waitForNextComponentGeneration(Injector discInjector, boolean isInitializing) {
+        Container.ComponentGraphResult result = container.waitForNextComponentGeneration(
+                this.currentGraph,
+                createFallbackInjector(vespaContainer, discInjector),
+                isInitializing);
+        this.currentGraph = result.newGraph();
+        return result.oldComponentsCleanupTask();
     }
 
     @SuppressWarnings("deprecation")
diff --git a/container-core/src/main/java/com/yahoo/container/core/config/testutil/HandlersConfigurerTestWrapper.java b/container-core/src/main/java/com/yahoo/container/core/config/testutil/HandlersConfigurerTestWrapper.java
index 5f5ee99c19f..c96f90bc0b3 100644
--- a/container-core/src/main/java/com/yahoo/container/core/config/testutil/HandlersConfigurerTestWrapper.java
+++ b/container-core/src/main/java/com/yahoo/container/core/config/testutil/HandlersConfigurerTestWrapper.java
@@ -120,7 +120,8 @@ public class HandlersConfigurerTestWrapper {
 
     public void reloadConfig() {
         configurer.reloadConfig(++lastGeneration);
-        configurer.waitForNextComponentGeneration(guiceInjector(), false);
+        Runnable cleanupTask = configurer.waitForNextComponentGeneration(guiceInjector(), false);
+        cleanupTask.run();
     }
 
     public void shutdown() {
diff --git a/container-core/src/main/java/com/yahoo/container/di/Container.java b/container-core/src/main/java/com/yahoo/container/di/Container.java
index 844d1dc1151..2726d0f69d4 100644
--- a/container-core/src/main/java/com/yahoo/container/di/Container.java
+++ b/container-core/src/main/java/com/yahoo/container/di/Container.java
@@ -71,14 +71,14 @@ public class Container {
         });
     }
 
-    public ComponentGraph waitForNextComponentGeneration(ComponentGraph oldGraph, Injector fallbackInjector, boolean isInitializing) {
+    public ComponentGraphResult waitForNextComponentGeneration(ComponentGraph oldGraph, Injector fallbackInjector, boolean isInitializing) {
         try {
             Collection<Bundle> obsoleteBundles = new HashSet<>();
             ComponentGraph newGraph = waitForNewConfigGenAndCreateGraph(oldGraph, fallbackInjector, isInitializing, obsoleteBundles);
             newGraph.reuseNodes(oldGraph);
             constructComponents(newGraph);
-            deconstructObsoleteComponents(oldGraph, newGraph, obsoleteBundles);
-            return newGraph;
+            Runnable cleanupTask = createPreviousGraphDeconstructionTask(oldGraph, newGraph, obsoleteBundles);
+            return new ComponentGraphResult(newGraph, cleanupTask);
         } catch (Throwable t) {
             invalidateGeneration(oldGraph.generation(), t);
             throw t;
@@ -159,9 +159,9 @@ public class Container {
         });
     }
 
-    private void deconstructObsoleteComponents(ComponentGraph oldGraph,
-                                               ComponentGraph newGraph,
-                                               Collection<Bundle> obsoleteBundles) {
+    private Runnable createPreviousGraphDeconstructionTask(ComponentGraph oldGraph,
+                                                           ComponentGraph newGraph,
+                                                           Collection<Bundle> obsoleteBundles) {
         Map<Object, ?> newComponents = new IdentityHashMap<>(newGraph.size());
         for (Object component : newGraph.allConstructedComponentsAndProviders())
             newComponents.put(component, null);
@@ -171,7 +171,7 @@ public class Container {
             if ( ! newComponents.containsKey(component))
                 obsoleteComponents.add(component);
 
-        componentDeconstructor.deconstruct(obsoleteComponents, obsoleteBundles);
+        return () -> componentDeconstructor.deconstruct(obsoleteComponents, obsoleteBundles);
     }
 
     private Set<Bundle> installApplicationBundles(Map<ConfigKey<? extends ConfigInstance>, ConfigInstance> configsIncludingBootstrapConfigs) {
@@ -281,4 +281,17 @@ public class Container {
         return BundleInstantiationSpecification.getFromStrings(config.id(), config.classId(), config.bundle());
     }
 
+    public static class ComponentGraphResult {
+        private final ComponentGraph newGraph;
+        private final Runnable oldComponentsCleanupTask;
+
+        public ComponentGraphResult(ComponentGraph newGraph, Runnable oldComponentsCleanupTask) {
+            this.newGraph = newGraph;
+            this.oldComponentsCleanupTask = oldComponentsCleanupTask;
+        }
+
+        public ComponentGraph newGraph() { return newGraph; }
+        public Runnable oldComponentsCleanupTask() { return oldComponentsCleanupTask; }
+    }
+
 }
diff --git a/container-core/src/test/java/com/yahoo/container/di/ContainerTest.java b/container-core/src/test/java/com/yahoo/container/di/ContainerTest.java
index 4ea1321cd70..020aa1e1ae4 100644
--- a/container-core/src/test/java/com/yahoo/container/di/ContainerTest.java
+++ b/container-core/src/test/java/com/yahoo/container/di/ContainerTest.java
@@ -333,11 +333,13 @@ public class ContainerTest extends ContainerTestBase {
     }
 
     ComponentGraph getNewComponentGraph(Container container, ComponentGraph oldGraph) {
-        return container.waitForNextComponentGeneration(oldGraph, Guice.createInjector(), true);
+        Container.ComponentGraphResult result = container.waitForNextComponentGeneration(oldGraph, Guice.createInjector(), true);
+        result.oldComponentsCleanupTask().run();
+        return result.newGraph();
     }
 
     ComponentGraph getNewComponentGraph(Container container) {
-        return container.waitForNextComponentGeneration(new ComponentGraph(), Guice.createInjector(), true);
+        return container.waitForNextComponentGeneration(new ComponentGraph(), Guice.createInjector(), true).newGraph();
     }
 
     private ComponentTakingConfig createComponentTakingConfig(ComponentGraph componentGraph) {
diff --git a/container-core/src/test/java/com/yahoo/container/di/ContainerTestBase.java b/container-core/src/test/java/com/yahoo/container/di/ContainerTestBase.java
index 9565f26fb1e..4f4ffe3732a 100644
--- a/container-core/src/test/java/com/yahoo/container/di/ContainerTestBase.java
+++ b/container-core/src/test/java/com/yahoo/container/di/ContainerTestBase.java
@@ -65,7 +65,9 @@ public class ContainerTestBase {
                             throw new UnsupportedOperationException("getBundle not supported.");
                         }
                     });
-            componentGraph = container.waitForNextComponentGeneration(componentGraph, Guice.createInjector(), true);
+            Container.ComponentGraphResult result = container.waitForNextComponentGeneration(this.componentGraph, Guice.createInjector(), true);
+            result.oldComponentsCleanupTask().run();
+            this.componentGraph = result.newGraph();
         } catch (Exception e) {
             throw new RuntimeException(e);
         }