Also use AnnotatedClusterState for derived bucket space states in ClusterStateBundle.

5 files changed, 36 insertions, 23 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);