Also use AnnotatedClusterState for derived bucket space states in ClusterStateBundle.

9 files changed, 71 insertions, 53 deletions

diff --git a/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/AnnotatedClusterState.java b/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/AnnotatedClusterState.java
index 9bf36cca947..32cbcf33d8e 100644
--- a/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/AnnotatedClusterState.java
+++ b/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/AnnotatedClusterState.java
@@ -9,7 +9,7 @@ import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
 
-public class AnnotatedClusterState {
+public class AnnotatedClusterState implements Cloneable {
 
     private final ClusterState clusterState;
     private final Map<Node, NodeStateReason> nodeStateReasons;
@@ -48,6 +48,16 @@ public class AnnotatedClusterState {
         return clusterStateReason;
     }
 
+    public AnnotatedClusterState clone() {
+        return cloneWithClusterState(clusterState.clone());
+    }
+
+    public AnnotatedClusterState cloneWithClusterState(ClusterState newClusterState) {
+        return new AnnotatedClusterState(newClusterState,
+                getClusterStateReason(),
+                getNodeStateReasons());
+    }
+
     @Override
     public String toString() {
         return clusterState.toString();
diff --git a/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/ClusterStateBundle.java b/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/ClusterStateBundle.java
index 09273ee5656..9291c8277da 100644
--- a/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/ClusterStateBundle.java
+++ b/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/ClusterStateBundle.java
@@ -23,7 +23,7 @@ import java.util.stream.Collectors;
 public class ClusterStateBundle {
 
     private final AnnotatedClusterState baselineState;
-    private final Map<String, ClusterState> derivedBucketSpaceStates;
+    private final Map<String, AnnotatedClusterState> derivedBucketSpaceStates;
 
     public static class Builder {
         private final AnnotatedClusterState baselineState;
@@ -53,15 +53,15 @@ public class ClusterStateBundle {
             if (stateDeriver == null || bucketSpaces == null || bucketSpaces.isEmpty()) {
                 return ClusterStateBundle.ofBaselineOnly(baselineState);
             }
-            Map<String, ClusterState> derived = bucketSpaces.stream()
+            Map<String, AnnotatedClusterState> derived = bucketSpaces.stream()
                     .collect(Collectors.toMap(
                             Function.identity(),
-                            s -> stateDeriver.derivedFrom(baselineState.getClusterState(), s)));
+                            s -> stateDeriver.derivedFrom(baselineState, s)));
             return new ClusterStateBundle(baselineState, derived);
         }
     }
 
-    private ClusterStateBundle(AnnotatedClusterState baselineState, Map<String, ClusterState> derivedBucketSpaceStates) {
+    private ClusterStateBundle(AnnotatedClusterState baselineState, Map<String, AnnotatedClusterState> derivedBucketSpaceStates) {
         this.baselineState = baselineState;
         this.derivedBucketSpaceStates = Collections.unmodifiableMap(derivedBucketSpaceStates);
     }
@@ -70,7 +70,7 @@ public class ClusterStateBundle {
         return new Builder(baselineState);
     }
 
-    public static ClusterStateBundle of(AnnotatedClusterState baselineState, Map<String, ClusterState> derivedBucketSpaceStates) {
+    public static ClusterStateBundle of(AnnotatedClusterState baselineState, Map<String, AnnotatedClusterState> derivedBucketSpaceStates) {
         return new ClusterStateBundle(baselineState, derivedBucketSpaceStates);
     }
 
@@ -86,17 +86,16 @@ public class ClusterStateBundle {
         return baselineState.getClusterState();
     }
 
-    public Map<String, ClusterState> getDerivedBucketSpaceStates() {
+    public Map<String, AnnotatedClusterState> getDerivedBucketSpaceStates() {
         return derivedBucketSpaceStates;
     }
 
     public ClusterStateBundle cloneWithMapper(Function<ClusterState, ClusterState> mapper) {
-        AnnotatedClusterState clonedBaseline = new AnnotatedClusterState(
-                mapper.apply(baselineState.getClusterState().clone()),
-                baselineState.getClusterStateReason(),
-                baselineState.getNodeStateReasons());
-        Map<String, ClusterState> clonedDerived = derivedBucketSpaceStates.entrySet().stream()
-                .collect(Collectors.toMap(e -> e.getKey(), e -> mapper.apply(e.getValue().clone())));
+        AnnotatedClusterState clonedBaseline = baselineState.cloneWithClusterState(
+                mapper.apply(baselineState.getClusterState().clone()));
+        Map<String, AnnotatedClusterState> clonedDerived = derivedBucketSpaceStates.entrySet().stream()
+                .collect(Collectors.toMap(e -> e.getKey(), e -> e.getValue().cloneWithClusterState(
+                        mapper.apply(e.getValue().getClusterState().clone()))));
         return new ClusterStateBundle(clonedBaseline, clonedDerived);
     }
 
@@ -114,7 +113,7 @@ public class ClusterStateBundle {
         // FIXME we currently treat mismatching bucket space sets as unchanged to avoid breaking some tests
         return derivedBucketSpaceStates.entrySet().stream()
                 .allMatch(entry -> other.derivedBucketSpaceStates.getOrDefault(entry.getKey(), entry.getValue())
-                        .similarToIgnoringInitProgress(entry.getValue()));
+                        .getClusterState().similarToIgnoringInitProgress(entry.getValue().getClusterState()));
     }
 
     public int getVersion() {
@@ -126,7 +125,7 @@ public class ClusterStateBundle {
         if (derivedBucketSpaceStates.isEmpty()) {
             return String.format("ClusterStateBundle('%s')", baselineState);
         }
-        Map<String, ClusterState> orderedStates = new TreeMap<>(derivedBucketSpaceStates);
+        Map<String, AnnotatedClusterState> orderedStates = new TreeMap<>(derivedBucketSpaceStates);
         return String.format("ClusterStateBundle('%s', %s)", baselineState, orderedStates.entrySet().stream()
                 .map(e -> String.format("%s '%s'", e.getKey(), e.getValue()))
                 .collect(Collectors.joining(", ")));
diff --git a/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/ClusterStateDeriver.java b/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/ClusterStateDeriver.java
index 9e42e1da649..aea6db7c2f6 100644
--- a/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/ClusterStateDeriver.java
+++ b/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/ClusterStateDeriver.java
@@ -14,5 +14,5 @@ public interface ClusterStateDeriver {
      * @param bucketSpace The name of the bucket space for which the state should be derived
      * @return A cluster state instance representing the derived state, or <em>state</em> unchanged.
      */
-    ClusterState derivedFrom(ClusterState state, String bucketSpace);
+    AnnotatedClusterState derivedFrom(AnnotatedClusterState state, String bucketSpace);
 }
diff --git a/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/MaintenanceWhenPendingGlobalMerges.java b/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/MaintenanceWhenPendingGlobalMerges.java
index 2f4ca09b584..14dc36b1d7d 100644
--- a/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/MaintenanceWhenPendingGlobalMerges.java
+++ b/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/MaintenanceWhenPendingGlobalMerges.java
@@ -28,17 +28,16 @@ public class MaintenanceWhenPendingGlobalMerges implements ClusterStateDeriver {
     }
 
     @Override
-    public ClusterState derivedFrom(ClusterState baselineState, String bucketSpace) {
-        ClusterState derivedState = baselineState.clone();
+    public AnnotatedClusterState derivedFrom(AnnotatedClusterState baselineState, String bucketSpace) {
+        AnnotatedClusterState derivedState = baselineState.clone();
         if (!bucketSpace.equals(bucketSpaceToDerive)) {
             return derivedState;
         }
-        Set<Integer> incompleteNodeIndices = nodesWithMergesNotDone(baselineState);
+        Set<Integer> incompleteNodeIndices = nodesWithMergesNotDone(baselineState.getClusterState());
         if (incompleteNodeIndices.isEmpty()) {
             return derivedState; // Nothing to do
         }
-        incompleteNodeIndices.forEach(nodeIndex -> derivedState.setNodeState(Node.ofStorage(nodeIndex),
-                new NodeState(NodeType.STORAGE, State.MAINTENANCE)));
+        incompleteNodeIndices.forEach(nodeIndex -> setNodeInMaintenance(derivedState, nodeIndex));
         return derivedState;
     }
 
@@ -56,6 +55,11 @@ public class MaintenanceWhenPendingGlobalMerges implements ClusterStateDeriver {
         return incompleteNodes;
     }
 
+    private void setNodeInMaintenance(AnnotatedClusterState derivedState, int nodeIndex) {
+        derivedState.getClusterState().setNodeState(Node.ofStorage(nodeIndex),
+                new NodeState(NodeType.STORAGE, State.MAINTENANCE));
+    }
+
     private boolean contentNodeIsAvailable(ClusterState state, int nodeIndex) {
         return state.getNodeState(Node.ofStorage(nodeIndex)).getState().oneOf("uir");
     }
diff --git a/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/rpc/SlimeClusterStateBundleCodec.java b/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/rpc/SlimeClusterStateBundleCodec.java
index c37bd8313a9..f6034cb34ff 100644
--- a/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/rpc/SlimeClusterStateBundleCodec.java
+++ b/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/rpc/SlimeClusterStateBundleCodec.java
@@ -48,9 +48,9 @@ public class SlimeClusterStateBundleCodec implements ClusterStateBundleCodec {
         ClusterState baseline = ClusterState.stateFromString(states.field("baseline").asString());
 
         Inspector spaces = states.field("spaces");
-        Map<String, ClusterState> derivedStates = new HashMap<>();
+        Map<String, AnnotatedClusterState> derivedStates = new HashMap<>();
         spaces.traverse(((ObjectTraverser)(key, value) -> {
-            derivedStates.put(key, ClusterState.stateFromString(value.asString()));
+            derivedStates.put(key, AnnotatedClusterState.withoutAnnotations(ClusterState.stateFromString(value.asString())));
         }));
 
         return ClusterStateBundle.of(AnnotatedClusterState.withoutAnnotations(baseline), derivedStates);
diff --git a/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/ClusterStateBundleTest.java b/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/ClusterStateBundleTest.java
index 04a0451bc6d..afb98b9d3d3 100644
--- a/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/ClusterStateBundleTest.java
+++ b/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/ClusterStateBundleTest.java
@@ -4,8 +4,6 @@ package com.yahoo.vespa.clustercontroller.core;
 import com.yahoo.vdslib.state.*;
 import org.junit.Test;
 
-import java.text.ParseException;
-
 import static org.hamcrest.CoreMatchers.equalTo;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.Assert.assertFalse;
@@ -14,11 +12,11 @@ import static org.junit.Assert.assertTrue;
 public class ClusterStateBundleTest {
 
     private static ClusterState stateOf(String state) {
-        try {
-            return new ClusterState(state);
-        } catch (ParseException e) {
-            throw new RuntimeException(e);
-        }
+        return ClusterState.stateFromString(state);
+    }
+
+    private static AnnotatedClusterState annotatedStateOf(String state) {
+        return AnnotatedClusterState.withoutAnnotations(stateOf(state));
     }
 
     private static ClusterStateBundle createTestBundle(boolean modifyDefaultSpace) {
@@ -26,11 +24,13 @@ public class ClusterStateBundleTest {
                 .builder(AnnotatedClusterState.withoutAnnotations(stateOf("distributor:2 storage:2")))
                 .bucketSpaces("default", "global", "narnia")
                 .stateDeriver((state, space) -> {
-                    ClusterState derived = state.clone();
+                    AnnotatedClusterState derived = state.clone();
                     if (space.equals("default") && modifyDefaultSpace) {
-                        derived.setNodeState(Node.ofStorage(0), new NodeState(NodeType.STORAGE, State.DOWN));
+                        derived.getClusterState()
+                                .setNodeState(Node.ofStorage(0), new NodeState(NodeType.STORAGE, State.DOWN));
                     } else if (space.equals("narnia")) {
-                        derived.setNodeState(Node.ofDistributor(0), new NodeState(NodeType.DISTRIBUTOR, State.DOWN));
+                        derived.getClusterState()
+                                .setNodeState(Node.ofDistributor(0), new NodeState(NodeType.DISTRIBUTOR, State.DOWN));
                     }
                     return derived;
                 })
@@ -46,9 +46,9 @@ public class ClusterStateBundleTest {
         ClusterStateBundle bundle = createTestBundle();
         assertThat(bundle.getBaselineClusterState(), equalTo(stateOf("distributor:2 storage:2")));
         assertThat(bundle.getDerivedBucketSpaceStates().size(), equalTo(3));
-        assertThat(bundle.getDerivedBucketSpaceStates().get("default"), equalTo(stateOf("distributor:2 storage:2 .0.s:d")));
-        assertThat(bundle.getDerivedBucketSpaceStates().get("global"), equalTo(stateOf("distributor:2 storage:2")));
-        assertThat(bundle.getDerivedBucketSpaceStates().get("narnia"), equalTo(stateOf("distributor:2 .0.s:d storage:2")));
+        assertThat(bundle.getDerivedBucketSpaceStates().get("default"), equalTo(annotatedStateOf("distributor:2 storage:2 .0.s:d")));
+        assertThat(bundle.getDerivedBucketSpaceStates().get("global"), equalTo(annotatedStateOf("distributor:2 storage:2")));
+        assertThat(bundle.getDerivedBucketSpaceStates().get("narnia"), equalTo(annotatedStateOf("distributor:2 .0.s:d storage:2")));
     }
 
     @Test
@@ -56,9 +56,9 @@ public class ClusterStateBundleTest {
         ClusterStateBundle bundle = createTestBundle().clonedWithVersionSet(123);
         assertThat(bundle.getBaselineClusterState(), equalTo(stateOf("version:123 distributor:2 storage:2")));
         assertThat(bundle.getDerivedBucketSpaceStates().size(), equalTo(3));
-        assertThat(bundle.getDerivedBucketSpaceStates().get("default"), equalTo(stateOf("version:123 distributor:2 storage:2 .0.s:d")));
-        assertThat(bundle.getDerivedBucketSpaceStates().get("global"), equalTo(stateOf("version:123 distributor:2 storage:2")));
-        assertThat(bundle.getDerivedBucketSpaceStates().get("narnia"), equalTo(stateOf("version:123 distributor:2 .0.s:d storage:2")));
+        assertThat(bundle.getDerivedBucketSpaceStates().get("default"), equalTo(annotatedStateOf("version:123 distributor:2 storage:2 .0.s:d")));
+        assertThat(bundle.getDerivedBucketSpaceStates().get("global"), equalTo(annotatedStateOf("version:123 distributor:2 storage:2")));
+        assertThat(bundle.getDerivedBucketSpaceStates().get("narnia"), equalTo(annotatedStateOf("version:123 distributor:2 .0.s:d storage:2")));
     }
 
     @Test
@@ -83,7 +83,7 @@ public class ClusterStateBundleTest {
     @Test
     public void toString_without_bucket_space_states_prints_only_baseline_state() {
         ClusterStateBundle bundle = ClusterStateBundle.ofBaselineOnly(
-                AnnotatedClusterState.withoutAnnotations(stateOf("distributor:2 storage:2")));
+                annotatedStateOf("distributor:2 storage:2"));
         assertThat(bundle.toString(), equalTo("ClusterStateBundle('distributor:2 storage:2')"));
     }
 
diff --git a/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/ClusterStateBundleUtil.java b/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/ClusterStateBundleUtil.java
index e097874682a..00c2194205d 100644
--- a/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/ClusterStateBundleUtil.java
+++ b/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/ClusterStateBundleUtil.java
@@ -14,7 +14,8 @@ public class ClusterStateBundleUtil {
 
     public static ClusterStateBundle makeBundle(String baselineState, StateMapping... bucketSpaceStates) {
         return ClusterStateBundle.of(AnnotatedClusterState.withoutAnnotations(ClusterState.stateFromString(baselineState)),
-                Stream.of(bucketSpaceStates).collect(Collectors.toMap(sm -> sm.bucketSpace, sm -> sm.state)));
+                Stream.of(bucketSpaceStates).collect(Collectors.toMap(sm -> sm.bucketSpace,
+                        sm -> AnnotatedClusterState.withoutAnnotations(sm.state))));
     }
 
 }
diff --git a/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/MaintenanceWhenPendingGlobalMergesTest.java b/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/MaintenanceWhenPendingGlobalMergesTest.java
index 0e2b9557cb9..9e17127d22b 100644
--- a/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/MaintenanceWhenPendingGlobalMergesTest.java
+++ b/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/MaintenanceWhenPendingGlobalMergesTest.java
@@ -27,12 +27,16 @@ public class MaintenanceWhenPendingGlobalMergesTest {
         return FixedBucketSpaces.globalSpace();
     }
 
+    private static AnnotatedClusterState stateFromString(String stateStr) {
+        return AnnotatedClusterState.withoutAnnotations(ClusterState.stateFromString(stateStr));
+    }
+
     @Test
     public void no_nodes_set_to_maintenance_in_global_bucket_space_state() {
         Fixture f = new Fixture();
         when(f.mockPendingChecker.mayHaveMergesPending(eq(globalSpace()), anyInt())).thenReturn(true); // False returned by default otherwise
-        ClusterState derived = f.deriver.derivedFrom(ClusterState.stateFromString("distributor:2 storage:2"), globalSpace());
-        assertThat(derived, equalTo(ClusterState.stateFromString("distributor:2 storage:2")));
+        AnnotatedClusterState derived = f.deriver.derivedFrom(stateFromString("distributor:2 storage:2"), globalSpace());
+        assertThat(derived, equalTo(stateFromString("distributor:2 storage:2")));
     }
 
     @Test
@@ -40,32 +44,32 @@ public class MaintenanceWhenPendingGlobalMergesTest {
         Fixture f = new Fixture();
         when(f.mockPendingChecker.mayHaveMergesPending(globalSpace(), 1)).thenReturn(true);
         when(f.mockPendingChecker.mayHaveMergesPending(globalSpace(), 3)).thenReturn(true);
-        ClusterState derived = f.deriver.derivedFrom(ClusterState.stateFromString("distributor:5 storage:5"), defaultSpace());
-        assertThat(derived, equalTo(ClusterState.stateFromString("distributor:5 storage:5 .1.s:m .3.s:m")));
+        AnnotatedClusterState derived = f.deriver.derivedFrom(stateFromString("distributor:5 storage:5"), defaultSpace());
+        assertThat(derived, equalTo(stateFromString("distributor:5 storage:5 .1.s:m .3.s:m")));
     }
 
     @Test
     public void no_nodes_set_to_maintenance_when_no_merges_pending() {
         Fixture f = new Fixture();
-        ClusterState derived = f.deriver.derivedFrom(ClusterState.stateFromString("distributor:5 storage:5"), defaultSpace());
-        assertThat(derived, equalTo(ClusterState.stateFromString("distributor:5 storage:5")));
+        AnnotatedClusterState derived = f.deriver.derivedFrom(stateFromString("distributor:5 storage:5"), defaultSpace());
+        assertThat(derived, equalTo(stateFromString("distributor:5 storage:5")));
     }
 
     @Test
     public void default_space_merges_do_not_count_towards_maintenance() {
         Fixture f = new Fixture();
         when(f.mockPendingChecker.mayHaveMergesPending(eq(defaultSpace()), anyInt())).thenReturn(true);
-        ClusterState derived = f.deriver.derivedFrom(ClusterState.stateFromString("distributor:2 storage:2"), defaultSpace());
-        assertThat(derived, equalTo(ClusterState.stateFromString("distributor:2 storage:2")));
+        AnnotatedClusterState derived = f.deriver.derivedFrom(stateFromString("distributor:2 storage:2"), defaultSpace());
+        assertThat(derived, equalTo(stateFromString("distributor:2 storage:2")));
     }
 
     @Test
     public void nodes_only_set_to_maintenance_when_marked_up_init_or_retiring() {
         Fixture f = new Fixture();
         when(f.mockPendingChecker.mayHaveMergesPending(eq(globalSpace()), anyInt())).thenReturn(true);
-        ClusterState derived = f.deriver.derivedFrom(ClusterState.stateFromString("distributor:5 storage:5 .1.s:m .2.s:r .3.s:i .4.s:d"), defaultSpace());
+        AnnotatedClusterState derived = f.deriver.derivedFrom(stateFromString("distributor:5 storage:5 .1.s:m .2.s:r .3.s:i .4.s:d"), defaultSpace());
         // TODO reconsider role of retired here... It should not have merges pending towards it in the general case, but may be out of sync
-        assertThat(derived, equalTo(ClusterState.stateFromString("distributor:5 storage:5 .0.s:m .1.s:m .2.s:m .3.s:m .4.s:d")));
+        assertThat(derived, equalTo(stateFromString("distributor:5 storage:5 .0.s:m .1.s:m .2.s:m .3.s:m .4.s:d")));
     }
 
 }
diff --git a/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/StateVersionTrackerTest.java b/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/StateVersionTrackerTest.java
index 0859ee0e409..c92f1badae8 100644
--- a/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/StateVersionTrackerTest.java
+++ b/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/StateVersionTrackerTest.java
@@ -246,9 +246,9 @@ public class StateVersionTrackerTest {
                 .builder(AnnotatedClusterState.withoutAnnotations(stateOf("distributor:1 storage:1")))
                 .bucketSpaces("default")
                 .stateDeriver((state, space) -> {
-                    ClusterState derived = state.clone();
+                    AnnotatedClusterState derived = state.clone();
                     if (alteredDefaultState) {
-                        derived.setNodeState(Node.ofStorage(0), new NodeState(NodeType.STORAGE, State.DOWN));
+                        derived.getClusterState().setNodeState(Node.ofStorage(0), new NodeState(NodeType.STORAGE, State.DOWN));
                     }
                     return derived;
                 })