Also set the distributor wanted state when safe-setting the storage node state

This is done as part of the SAFE REST API call to set the node state of a
storage node to ensure atomicity of the state change, reduce the number of
state changes, and minimize the time to complete the state changes.

The right way to think about the safe-set is then: In order to safely set a
storage node to (e.g.) maintenance, the distributor will also have to be set to
down. And so on for the various permutations of state transitions.

4 files changed, 247 insertions, 12 deletions

diff --git a/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/ClusterInfo.java b/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/ClusterInfo.java
index 6f921eef25e..ff92fc6867f 100644
--- a/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/ClusterInfo.java
+++ b/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/ClusterInfo.java
@@ -50,6 +50,8 @@ public class ClusterInfo {
     /** Returns the configured nodes of this as a read-only map indexed on node index (distribution key) */
     Map<Integer, ConfiguredNode> getConfiguredNodes() { return Collections.unmodifiableMap(nodes); }
 
+    boolean hasConfiguredNode(int index) { return nodes.containsKey(index); }
+
     /** Sets the nodes which belongs to this cluster */
     void setNodes(Collection<ConfiguredNode> newNodes, ContentCluster owner, Distribution distribution) {
         // Remove info for removed nodes
diff --git a/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/ContentCluster.java b/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/ContentCluster.java
index ad7da755b98..0ff59c26c13 100644
--- a/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/ContentCluster.java
+++ b/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/ContentCluster.java
@@ -206,7 +206,6 @@ public class ContentCluster {
     }
 
     public boolean hasConfiguredNode(int index) {
-        return clusterInfo.getConfiguredNodes().containsKey(index);
+        return clusterInfo.hasConfiguredNode(index);
     }
-
 }
diff --git a/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/restapiv2/requests/SetNodeStateRequest.java b/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/restapiv2/requests/SetNodeStateRequest.java
index cf41707db75..79441bc4bb4 100644
--- a/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/restapiv2/requests/SetNodeStateRequest.java
+++ b/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/restapiv2/requests/SetNodeStateRequest.java
@@ -5,16 +5,18 @@ import com.yahoo.log.LogLevel;
 import com.yahoo.vdslib.state.ClusterState;
 import com.yahoo.vdslib.state.Node;
 import com.yahoo.vdslib.state.NodeState;
+import com.yahoo.vdslib.state.NodeType;
 import com.yahoo.vdslib.state.State;
+import com.yahoo.vespa.clustercontroller.core.ContentCluster;
 import com.yahoo.vespa.clustercontroller.core.NodeInfo;
 import com.yahoo.vespa.clustercontroller.core.NodeStateChangeChecker;
 import com.yahoo.vespa.clustercontroller.core.RemoteClusterControllerTask;
-import com.yahoo.vespa.clustercontroller.core.ContentCluster;
 import com.yahoo.vespa.clustercontroller.core.listeners.NodeStateOrHostInfoChangeHandler;
 import com.yahoo.vespa.clustercontroller.core.restapiv2.Id;
 import com.yahoo.vespa.clustercontroller.core.restapiv2.MissingIdException;
 import com.yahoo.vespa.clustercontroller.core.restapiv2.Request;
-import com.yahoo.vespa.clustercontroller.utils.staterestapi.errors.*;
+import com.yahoo.vespa.clustercontroller.utils.staterestapi.errors.InvalidContentException;
+import com.yahoo.vespa.clustercontroller.utils.staterestapi.errors.StateRestApiException;
 import com.yahoo.vespa.clustercontroller.utils.staterestapi.requests.SetUnitStateRequest;
 import com.yahoo.vespa.clustercontroller.utils.staterestapi.response.SetResponse;
 import com.yahoo.vespa.clustercontroller.utils.staterestapi.response.UnitState;
@@ -77,7 +79,7 @@ public class SetNodeStateRequest extends Request<SetResponse> {
             Node node,
             NodeStateOrHostInfoChangeHandler stateListener,
             ClusterState currentClusterState) throws StateRestApiException {
-        if ( ! cluster.getConfiguredNodes().containsKey(node.getIndex())) {
+        if ( ! cluster.hasConfiguredNode(node.getIndex())) {
             throw new MissingIdException(cluster.getName(), node);
         }
         NodeInfo nodeInfo = cluster.getNodeInfo(node);
@@ -95,15 +97,92 @@ public class SetNodeStateRequest extends Request<SetResponse> {
                 " wanted-state=" + wantedState +
                 " new-wanted-state=" + newWantedState +
                 " change-check=" + result);
+
+        boolean success = ensureWantedState(
+                result,
+                newWantedState,
+                condition,
+                nodeInfo,
+                cluster,
+                stateListener);
+
+        // If the state was successfully set, just return an "ok" message back.
+        String reason = success ? "ok" : result.getReason();
+        return new SetResponse(reason, success);
+    }
+
+    /**
+     * Returns true if the current/old wanted state already matches the requested
+     * wanted state, or the requested state has been accepted as the new wanted state.
+     */
+    private static boolean ensureWantedState(NodeStateChangeChecker.Result result,
+                                             NodeState newWantedState,
+                                             SetUnitStateRequest.Condition condition,
+                                             NodeInfo nodeInfo,
+                                             ContentCluster cluster,
+                                             NodeStateOrHostInfoChangeHandler stateListener) {
         if (result.settingWantedStateIsAllowed()) {
-            nodeInfo.setWantedState(newWantedState);
-            stateListener.handleNewWantedNodeState(nodeInfo, newWantedState);
+            setNewWantedState(nodeInfo, newWantedState, stateListener);
         }
 
-        // wasModified is true if the new/current State equals the wanted state in the request.
-        boolean wasModified = result.settingWantedStateIsAllowed() || result.wantedStateAlreadySet();
-        // If the state was successfully set, just return an "ok" message back.
-        String reason = wasModified ? "ok" : result.getReason();
-        return new SetResponse(reason, wasModified);
+        // True if the wanted state was or has just been set to newWantedState
+        boolean success = result.settingWantedStateIsAllowed() || result.wantedStateAlreadySet();
+
+        if (success && condition == SetUnitStateRequest.Condition.SAFE && nodeInfo.isStorage()) {
+            // In safe-mode, setting the storage node must be accompanied by changing the state
+            // of the distributor. E.g. setting the storage node to maintenance may cause
+            // feeding issues unless distributor is also set down.
+
+            ensureDistributorWantedState(cluster, nodeInfo.getNodeIndex(), newWantedState, stateListener);
+        }
+
+        return success;
+    }
+
+    /**
+     * Set the wanted state on the distributor to something appropriate given the storage is being
+     * set to (or is equal to) newStorageWantedState.
+     */
+    private static void ensureDistributorWantedState(ContentCluster cluster,
+                                                     int index,
+                                                     NodeState newStorageWantedState,
+                                                     NodeStateOrHostInfoChangeHandler stateListener) {
+        Node distributorNode = new Node(NodeType.DISTRIBUTOR, index);
+        NodeInfo nodeInfo = cluster.getNodeInfo(distributorNode);
+        if (nodeInfo == null) {
+            throw new IllegalStateException("Missing distributor at index " +
+                    distributorNode.getIndex());
+        }
+
+        State newState;
+        switch (newStorageWantedState.getState()) {
+            case MAINTENANCE:
+                newState = State.DOWN;
+                break;
+            case RETIRED:
+                newState = State.UP;
+                break;
+            default:
+                newState = newStorageWantedState.getState();
+                if (!newState.validWantedNodeState(distributorNode.getType())) {
+                    throw new IllegalStateException("Distributor cannot be set to wanted state " +
+                            newState);
+                }
+        }
+
+        NodeState newWantedState = new NodeState(distributorNode.getType(), newState);
+        newWantedState.setDescription(newStorageWantedState.getDescription());
+
+        NodeState currentWantedState = nodeInfo.getUserWantedState();
+        if (newWantedState.getState() != currentWantedState.getState()) {
+            setNewWantedState(nodeInfo, newWantedState, stateListener);
+        }
+    }
+
+    private static void setNewWantedState(NodeInfo nodeInfo,
+                                          NodeState newWantedState,
+                                          NodeStateOrHostInfoChangeHandler stateListener) {
+        nodeInfo.setWantedState(newWantedState);
+        stateListener.handleNewWantedNodeState(nodeInfo, newWantedState);
     }
 }
diff --git a/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/restapiv2/requests/SetNodeStateRequestTest.java b/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/restapiv2/requests/SetNodeStateRequestTest.java
new file mode 100644
index 00000000000..9239b8774b0
--- /dev/null
+++ b/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/restapiv2/requests/SetNodeStateRequestTest.java
@@ -0,0 +1,155 @@
+// Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+package com.yahoo.vespa.clustercontroller.core.restapiv2.requests;
+
+import com.yahoo.vdslib.state.ClusterState;
+import com.yahoo.vdslib.state.Node;
+import com.yahoo.vdslib.state.NodeState;
+import com.yahoo.vdslib.state.NodeType;
+import com.yahoo.vdslib.state.State;
+import com.yahoo.vespa.clustercontroller.core.ContentCluster;
+import com.yahoo.vespa.clustercontroller.core.NodeInfo;
+import com.yahoo.vespa.clustercontroller.core.NodeStateChangeChecker;
+import com.yahoo.vespa.clustercontroller.core.listeners.NodeStateOrHostInfoChangeHandler;
+import com.yahoo.vespa.clustercontroller.utils.staterestapi.errors.StateRestApiException;
+import com.yahoo.vespa.clustercontroller.utils.staterestapi.requests.SetUnitStateRequest;
+import com.yahoo.vespa.clustercontroller.utils.staterestapi.response.SetResponse;
+import com.yahoo.vespa.clustercontroller.utils.staterestapi.response.UnitState;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+
+import static org.mockito.Matchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+public class SetNodeStateRequestTest {
+    public static final String REASON = "operator";
+    ContentCluster cluster = mock(ContentCluster.class);
+    SetUnitStateRequest.Condition condition = SetUnitStateRequest.Condition.SAFE;
+    Map<String, UnitState> newStates = new HashMap<>();
+    UnitState unitState = mock(UnitState.class);
+    private final int NODE_INDEX = 2;
+    Node storageNode = new Node(NodeType.STORAGE, NODE_INDEX);
+    NodeStateOrHostInfoChangeHandler stateListener = mock(NodeStateOrHostInfoChangeHandler.class);
+    ClusterState currentClusterState = mock(ClusterState.class);
+
+    @Before
+    public void setUp() {
+        newStates.put("user", unitState);
+    }
+
+    @Test
+    public void testUpToMaintenance() throws StateRestApiException {
+        testSetStateRequest(
+                "maintenance",
+                State.UP, State.UP,
+                NodeStateChangeChecker.Result.allowSettingOfWantedState(),
+                Optional.of(State.MAINTENANCE), Optional.of(State.DOWN));
+    }
+
+    @Test
+    public void testUpToDown() throws StateRestApiException {
+        testSetStateRequest(
+                "down",
+                State.UP, State.UP,
+                NodeStateChangeChecker.Result.allowSettingOfWantedState(),
+                Optional.of(State.DOWN), Optional.of(State.DOWN));
+    }
+
+    @Test
+    public void testMaintenanceToUp() throws StateRestApiException {
+        testSetStateRequest(
+                "up",
+                State.MAINTENANCE, State.DOWN,
+                NodeStateChangeChecker.Result.allowSettingOfWantedState(),
+                Optional.of(State.UP), Optional.of(State.UP));
+    }
+
+    @Test
+    public void testDownToUp() throws StateRestApiException {
+        testSetStateRequest(
+                "up",
+                State.DOWN, State.DOWN,
+                NodeStateChangeChecker.Result.allowSettingOfWantedState(),
+                Optional.of(State.UP), Optional.of(State.UP));
+    }
+
+    @Test
+    public void testOnlyStorageInMaintenaceToMaintenance() throws StateRestApiException {
+        testSetStateRequest(
+                "maintenance",
+                State.MAINTENANCE, State.UP,
+                NodeStateChangeChecker.Result.createAlreadySet(),
+                Optional.empty(), Optional.of(State.DOWN));
+    }
+
+    @Test
+    public void testNoOpMaintenaceToMaintenance() throws StateRestApiException {
+        testSetStateRequest(
+                "maintenance",
+                State.MAINTENANCE, State.DOWN,
+                NodeStateChangeChecker.Result.createAlreadySet(),
+                Optional.empty(), Optional.empty());
+    }
+
+    private void testSetStateRequest(
+            String wantedStateString,
+            State storageWantedState,
+            State distributorWantedState,
+            NodeStateChangeChecker.Result result,
+            Optional<State> expectedNewStorageWantedState,
+            Optional<State> expectedNewDistributorWantedState) throws StateRestApiException {
+        when(cluster.hasConfiguredNode(NODE_INDEX)).thenReturn(true);
+
+        NodeInfo storageNodeInfo = mock(NodeInfo.class);
+        when(cluster.getNodeInfo(storageNode)).thenReturn(storageNodeInfo);
+        NodeState storageNodeState = new NodeState(NodeType.STORAGE, storageWantedState);
+        when(storageNodeInfo.getUserWantedState()).thenReturn(storageNodeState);
+
+        when(unitState.getId()).thenReturn(wantedStateString);
+        when(unitState.getReason()).thenReturn(REASON);
+
+        when(cluster.calculateEffectOfNewState(any(), any(), any(), any(), any())).thenReturn(result);
+
+        when(storageNodeInfo.isStorage()).thenReturn(storageNode.getType() == NodeType.STORAGE);
+        when(storageNodeInfo.getNodeIndex()).thenReturn(storageNode.getIndex());
+
+        NodeInfo distributorNodeInfo = mock(NodeInfo.class);
+        Node distributorNode = new Node(NodeType.DISTRIBUTOR, NODE_INDEX);
+        when(cluster.getNodeInfo(distributorNode)).thenReturn(distributorNodeInfo);
+
+        NodeState distributorNodeState = new NodeState(distributorNode.getType(), distributorWantedState);
+        when(distributorNodeInfo.getUserWantedState()).thenReturn(distributorNodeState);
+
+        setWantedState();
+
+        if (expectedNewStorageWantedState.isPresent()) {
+            NodeState expectedNewStorageNodeState =
+                    new NodeState(NodeType.STORAGE, expectedNewStorageWantedState.get());
+            verify(storageNodeInfo).setWantedState(expectedNewStorageNodeState);
+            verify(stateListener).handleNewWantedNodeState(storageNodeInfo, expectedNewStorageNodeState);
+        }
+
+        if (expectedNewDistributorWantedState.isPresent()) {
+            NodeState expectedNewDistributorNodeState =
+                    new NodeState(NodeType.DISTRIBUTOR, expectedNewDistributorWantedState.get());
+            verify(distributorNodeInfo).setWantedState(expectedNewDistributorNodeState);
+            verify(stateListener).handleNewWantedNodeState(distributorNodeInfo, expectedNewDistributorNodeState);
+        }
+    }
+
+    private SetResponse setWantedState() throws StateRestApiException {
+        return SetNodeStateRequest.setWantedState(
+                cluster,
+                condition,
+                newStates,
+                storageNode,
+                stateListener,
+                currentClusterState);
+    }
+}
\ No newline at end of file