1 files changed, 17 insertions, 11 deletions

diff --git a/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/database/DatabaseHandler.java b/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/database/DatabaseHandler.java
index 0c6c773a9bc..01b8ed48c80 100644
--- a/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/database/DatabaseHandler.java
+++ b/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/database/DatabaseHandler.java
@@ -10,8 +10,8 @@ import com.yahoo.vespa.clustercontroller.core.FleetControllerContext;
 import com.yahoo.vespa.clustercontroller.core.FleetController;
 import com.yahoo.vespa.clustercontroller.core.NodeInfo;
 import com.yahoo.vespa.clustercontroller.core.Timer;
-import com.yahoo.vespa.clustercontroller.core.listeners.NodeAddedOrRemovedListener;
-import com.yahoo.vespa.clustercontroller.core.listeners.NodeStateOrHostInfoChangeHandler;
+import com.yahoo.vespa.clustercontroller.core.listeners.SlobrokListener;
+import com.yahoo.vespa.clustercontroller.core.listeners.NodeListener;
 import org.apache.zookeeper.KeeperException;
 
 import java.io.PrintWriter;
@@ -32,8 +32,8 @@ public class DatabaseHandler {
     public interface DatabaseContext {
         ContentCluster getCluster();
         FleetController getFleetController();
-        NodeAddedOrRemovedListener getNodeAddedOrRemovedListener();
-        NodeStateOrHostInfoChangeHandler getNodeStateUpdateListener();
+        SlobrokListener getNodeAddedOrRemovedListener();
+        NodeListener getNodeStateUpdateListener();
     }
 
     private static class Data {
@@ -191,9 +191,8 @@ public class DatabaseHandler {
                 // being called, but after receiving a database loss event.
                 clearSessionMetaData(false);
                 fleetControllerContext.log(logger, Level.INFO, "Setting up new ZooKeeper session at " + zooKeeperAddress);
-                DatabaseFactory.Params params = new DatabaseFactory.Params()
-                        .cluster(cluster)
-                        .nodeIndex(fleetControllerContext.id().index())
+                DatabaseFactory.Params params = new DatabaseFactory
+                        .Params()
                         .databaseAddress(zooKeeperAddress)
                         .databaseSessionTimeout(zooKeeperSessionTimeout)
                         .databaseListener(dbListener);
@@ -426,10 +425,10 @@ public class DatabaseHandler {
         }
     }
 
-    public void saveWantedStates(DatabaseContext databaseContext) {
+    public boolean saveWantedStates(DatabaseContext databaseContext) {
         fleetControllerContext.log(logger, Level.FINE, () -> "Checking whether wanted states have changed compared to zookeeper version.");
         Map<Node, NodeState> wantedStates = new TreeMap<>();
-        for (NodeInfo info : databaseContext.getCluster().getNodeInfo()) {
+        for (NodeInfo info : databaseContext.getCluster().getNodeInfos()) {
             if (!info.getUserWantedState().equals(new NodeState(info.getNode().getType(), State.UP))) {
                 wantedStates.put(info.getNode(), info.getUserWantedState());
             }
@@ -444,6 +443,9 @@ public class DatabaseHandler {
             fleetControllerContext.log(logger, Level.FINE, () -> "Scheduling new wanted states to be stored into zookeeper.");
             pendingStore.wantedStates = wantedStates;
             doNextZooKeeperTask(databaseContext);
+            return true;
+        } else {
+            return false;
         }
     }
 
@@ -466,7 +468,11 @@ public class DatabaseHandler {
         boolean altered = false;
         for (Node node : wantedStates.keySet()) {
             NodeInfo nodeInfo = databaseContext.getCluster().getNodeInfo(node);
-            if (nodeInfo == null) continue; // ignore wanted state of nodes which doesn't exist
+            if (nodeInfo == null) {
+                databaseContext.getNodeStateUpdateListener().handleRemovedNode(node);
+                altered = true;
+                continue;
+            }
             NodeState wantedState = wantedStates.get(node);
             if ( ! nodeInfo.getUserWantedState().equals(wantedState)) {
                 nodeInfo.setWantedState(wantedState);
@@ -477,7 +483,7 @@ public class DatabaseHandler {
         }
 
         // Remove wanted state from any node having a wanted state set that is no longer valid
-        for (NodeInfo info : databaseContext.getCluster().getNodeInfo()) {
+        for (NodeInfo info : databaseContext.getCluster().getNodeInfos()) {
             NodeState wantedState = wantedStates.get(info.getNode());
             if (wantedState == null && !info.getUserWantedState().equals(new NodeState(info.getNode().getType(), State.UP))) {
                 info.setWantedState(null);