Revert "Revert "Revert "Revert "Revert "Revert "Bratseth/lockdown zk""""""

2 files changed, 94 insertions, 5 deletions

diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeRepository.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeRepository.java
index 7751abb1d4b..327246e3099 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeRepository.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeRepository.java
@@ -14,6 +14,7 @@ import com.yahoo.vespa.hosted.provision.node.filter.NodeFilter;
 import com.yahoo.vespa.hosted.provision.node.filter.NodeListFilter;
 import com.yahoo.vespa.hosted.provision.node.filter.StateFilter;
 import com.yahoo.vespa.hosted.provision.persistence.CuratorDatabaseClient;
+import com.yahoo.vespa.zookeeper.ZooKeeperServer;
 
 import java.time.Clock;
 import java.util.ArrayList;
@@ -55,6 +56,7 @@ import java.util.stream.Collectors;
 public class NodeRepository extends AbstractComponent {
 
     private final CuratorDatabaseClient zkClient;
+    private final Curator curator;
 
     /**
      * Creates a node repository form a zookeeper provider.
@@ -71,6 +73,7 @@ public class NodeRepository extends AbstractComponent {
      */
     public NodeRepository(NodeFlavors flavors, Curator curator, Clock clock) {
         this.zkClient = new CuratorDatabaseClient(flavors, curator, clock);
+        this.curator = curator;
 
         // read and write all nodes to make sure they are stored in the latest version of the serialized format
         for (Node.State state : Node.State.values())
@@ -126,7 +129,9 @@ public class NodeRepository extends AbstractComponent {
                 throw new IllegalArgumentException("Cannot add " + node.hostname() + ": A node with this name already exists");
         }
         try (Mutex lock = lockUnallocated()) {
-            return zkClient.addNodes(nodes);
+            List<Node> addedNodes = zkClient.addNodes(nodes);
+            updateAllowedHosts();
+            return addedNodes;
         }
     }
 
@@ -242,14 +247,16 @@ public class NodeRepository extends AbstractComponent {
     /**
      * Removes a node. A node must be in the failed or parked state before it can be removed.
      *
-     * @return true if the node was removed, false if it was not found
+     * @return true if the node was removed, false if it was not found in one of these states
      */
     public boolean remove(String hostname) {
         Optional<Node> nodeToRemove = getNode(hostname, Node.State.failed, Node.State.parked);
-        if ( ! nodeToRemove.isPresent()) 
-            return false;
+        if ( ! nodeToRemove.isPresent()) return false;
+
         try (Mutex lock = lock(nodeToRemove.get())) {
-            return zkClient.removeNode(nodeToRemove.get().state(), hostname);
+            boolean removed = zkClient.removeNode(nodeToRemove.get().state(), hostname);
+            updateAllowedHosts();
+            return removed;
         }
     }
 
@@ -340,5 +347,21 @@ public class NodeRepository extends AbstractComponent {
     private Mutex lock(Node node) {
         return node.allocation().isPresent() ? lock(node.allocation().get().owner()) : lockUnallocated();
     }
+    
+    private void updateAllowedHosts() {
+        StringBuilder s = new StringBuilder();
+        
+        // Add tenant hosts
+        for (Node node : getNodes(Node.Type.tenant))
+            s.append(node.hostname()).append(",");
+
+        // Add the zooKeeper servers
+        for (String hostPort : curator.connectionSpec().split(","))
+            s.append(hostPort.split(":")[0]).append(",");
+
+        if (s.length() > 0)
+            s.setLength(s.length()-1); // remove last comma
+        System.setProperty(ZooKeeperServer.ZOOKEEPER_VESPA_CLIENTS_PROPERTY, s.toString());
+    }
 
 }
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/NodeRepositoryTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/NodeRepositoryTest.java
new file mode 100644
index 00000000000..dc4744c9eaf
--- /dev/null
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/NodeRepositoryTest.java
@@ -0,0 +1,66 @@
+package com.yahoo.vespa.hosted.provision;
+
+import com.yahoo.test.ManualClock;
+import com.yahoo.vespa.config.nodes.NodeRepositoryConfig;
+import com.yahoo.vespa.curator.mock.MockCurator;
+import com.yahoo.vespa.hosted.provision.node.Configuration;
+import com.yahoo.vespa.hosted.provision.node.NodeFlavors;
+import com.yahoo.vespa.hosted.provision.testutils.FlavorConfigBuilder;
+import com.yahoo.vespa.zookeeper.ZooKeeperServer;
+import org.junit.Test;
+import java.time.Clock;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Optional;
+import java.util.Set;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+/**
+ * tests basic operation of the node repository
+ * 
+ * @author bratseth
+ */
+public class NodeRepositoryTest {
+
+    @Test
+    public void nodeRepositoryTest() {
+        NodeFlavors nodeFlavors = new NodeFlavors(createConfig());
+        Clock clock = new ManualClock();
+        MockCurator curator = new MockCurator();
+        curator.setConnectionSpec("server1:1234,server2:5678");
+        NodeRepository nodeRepository = new NodeRepository(nodeFlavors, curator, clock);
+
+        assertEquals(0, nodeRepository.getNodes(Node.Type.tenant).size());
+
+        List<Node> nodes = new ArrayList<>();
+        nodes.add(nodeRepository.createNode("id1", "host1", Optional.empty(), new Configuration(nodeFlavors.getFlavorOrThrow("default")), Node.Type.tenant));
+        nodes.add(nodeRepository.createNode("id2", "host2", Optional.empty(), new Configuration(nodeFlavors.getFlavorOrThrow("default")), Node.Type.tenant));
+        nodes.add(nodeRepository.createNode("id3", "host3", Optional.empty(), new Configuration(nodeFlavors.getFlavorOrThrow("default")), Node.Type.tenant));
+        nodeRepository.addNodes(nodes);
+
+        assertEquals(3, nodeRepository.getNodes(Node.Type.tenant).size());
+        assertEquals(asSet("host1,host2,host3,server1,server2"), asSet(System.getProperty(ZooKeeperServer.ZOOKEEPER_VESPA_CLIENTS_PROPERTY)));
+        
+        nodeRepository.move("host2", Node.State.parked);
+        assertTrue(nodeRepository.remove("host2"));
+
+        assertEquals(2, nodeRepository.getNodes(Node.Type.tenant).size());
+        assertEquals(asSet("host1,host3,server1,server2"), asSet(System.getProperty(ZooKeeperServer.ZOOKEEPER_VESPA_CLIENTS_PROPERTY)));
+    }
+    
+    private Set<String> asSet(String s) {
+        return new HashSet<>(Arrays.asList(s.split(",")));
+    }
+
+    private NodeRepositoryConfig createConfig() {
+        FlavorConfigBuilder b = new FlavorConfigBuilder();
+        b.addFlavor("default", 2., 4., 100, "BARE_METAL").cost(3);
+        b.addFlavor("small", 1., 2., 50, "BARE_METAL").cost(2);
+        return b.build();
+    }
+
+}