Extend cluster v2 to output the published distribution state.

This is the baseline cluster state + per bucket space states.

5 files changed, 86 insertions, 4 deletions

diff --git a/clustercontroller-utils/src/main/java/com/yahoo/vespa/clustercontroller/utils/staterestapi/response/DistributionState.java b/clustercontroller-utils/src/main/java/com/yahoo/vespa/clustercontroller/utils/staterestapi/response/DistributionState.java
new file mode 100644
index 00000000000..53a4a07fa76
--- /dev/null
+++ b/clustercontroller-utils/src/main/java/com/yahoo/vespa/clustercontroller/utils/staterestapi/response/DistributionState.java
@@ -0,0 +1,26 @@
+// Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.clustercontroller.utils.staterestapi.response;
+
+import java.util.Map;
+
+/**
+ * A distribution state representing the baseline cluster state and per bucket space states.
+ */
+public class DistributionState {
+    private final String baselineState;
+    private final Map<String, String> bucketSpaceStates;
+
+    public DistributionState(String baselineState,
+                             Map<String, String> bucketSpaceStates) {
+        this.baselineState = baselineState;
+        this.bucketSpaceStates = bucketSpaceStates;
+    }
+
+    public String getBaselineState() {
+        return baselineState;
+    }
+
+    public Map<String, String> getBucketSpaceStates() {
+        return bucketSpaceStates;
+    }
+}
diff --git a/clustercontroller-utils/src/main/java/com/yahoo/vespa/clustercontroller/utils/staterestapi/response/DistributionStates.java b/clustercontroller-utils/src/main/java/com/yahoo/vespa/clustercontroller/utils/staterestapi/response/DistributionStates.java
new file mode 100644
index 00000000000..e7de760bb64
--- /dev/null
+++ b/clustercontroller-utils/src/main/java/com/yahoo/vespa/clustercontroller/utils/staterestapi/response/DistributionStates.java
@@ -0,0 +1,10 @@
+// Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.clustercontroller.utils.staterestapi.response;
+
+/**
+ * Interface to get the published distribution state.
+ */
+public interface DistributionStates {
+
+    DistributionState getPublishedState();
+}
diff --git a/clustercontroller-utils/src/main/java/com/yahoo/vespa/clustercontroller/utils/staterestapi/response/UnitResponse.java b/clustercontroller-utils/src/main/java/com/yahoo/vespa/clustercontroller/utils/staterestapi/response/UnitResponse.java
index 583820f5535..61ed6662493 100644
--- a/clustercontroller-utils/src/main/java/com/yahoo/vespa/clustercontroller/utils/staterestapi/response/UnitResponse.java
+++ b/clustercontroller-utils/src/main/java/com/yahoo/vespa/clustercontroller/utils/staterestapi/response/UnitResponse.java
@@ -9,5 +9,6 @@ public interface UnitResponse {
     CurrentUnitState getCurrentState();
     Map<String, SubUnitList> getSubUnits();
     UnitMetrics getMetrics();
+    DistributionStates getDistributionStates();
 
 }
diff --git a/clustercontroller-utils/src/main/java/com/yahoo/vespa/clustercontroller/utils/staterestapi/server/JsonWriter.java b/clustercontroller-utils/src/main/java/com/yahoo/vespa/clustercontroller/utils/staterestapi/server/JsonWriter.java
index 11b0dc33610..a5a0e6cbe3d 100644
--- a/clustercontroller-utils/src/main/java/com/yahoo/vespa/clustercontroller/utils/staterestapi/server/JsonWriter.java
+++ b/clustercontroller-utils/src/main/java/com/yahoo/vespa/clustercontroller/utils/staterestapi/server/JsonWriter.java
@@ -2,6 +2,7 @@
 package com.yahoo.vespa.clustercontroller.utils.staterestapi.server;
 
 import com.yahoo.vespa.clustercontroller.utils.staterestapi.response.*;
+import org.codehaus.jettison.json.JSONArray;
 import org.codehaus.jettison.json.JSONException;
 import org.codehaus.jettison.json.JSONObject;
 
@@ -29,13 +30,25 @@ public class JsonWriter {
 
     public void fillInJson(UnitResponse data, JSONObject json) throws Exception {
         UnitAttributes attributes = data.getAttributes();
-        if (attributes != null) fillInJson(attributes, json);
+        if (attributes != null) {
+            fillInJson(attributes, json);
+        }
         CurrentUnitState stateData = data.getCurrentState();
-        if (stateData != null) fillInJson(stateData, json);
+        if (stateData != null) {
+            fillInJson(stateData, json);
+        }
         UnitMetrics metrics = data.getMetrics();
-        if (metrics != null) fillInJson(metrics, json);
+        if (metrics != null) {
+            fillInJson(metrics, json);
+        }
         Map<String, SubUnitList> subUnits = data.getSubUnits();
-        if (subUnits != null) fillInJson(subUnits, json);
+        if (subUnits != null) {
+            fillInJson(subUnits, json);
+        }
+        DistributionStates distributionStates = data.getDistributionStates();
+        if (distributionStates != null) {
+            fillInJson(distributionStates, json);
+        }
     }
 
     public void fillInJson(CurrentUnitState stateData, JSONObject json) throws Exception {
@@ -85,6 +98,26 @@ public class JsonWriter {
         }
     }
 
+    private static void fillInJson(DistributionStates states, JSONObject json) throws Exception {
+        JSONObject statesJson = new JSONObject();
+        statesJson.put("published", distributionStateToJson(states.getPublishedState()));
+        json.put("distribution-states", statesJson);
+    }
+
+    private static JSONObject distributionStateToJson(DistributionState state) throws Exception {
+        JSONObject result = new JSONObject();
+        result.put("baseline", state.getBaselineState());
+        JSONArray bucketSpacesJson = new JSONArray();
+        result.put("bucket-spaces", bucketSpacesJson);
+        for (Map.Entry<String, String> entry : state.getBucketSpaceStates().entrySet()) {
+            JSONObject bucketSpaceJson = new JSONObject();
+            bucketSpaceJson.put("name", entry.getKey());
+            bucketSpaceJson.put("state", entry.getValue());
+            bucketSpacesJson.put(bucketSpaceJson);
+        }
+        return result;
+    }
+
     public JSONObject createErrorJson(String description) {
         JSONObject o = new JSONObject();
         try{
diff --git a/clustercontroller-utils/src/test/java/com/yahoo/vespa/clustercontroller/utils/staterestapi/DummyStateApi.java b/clustercontroller-utils/src/test/java/com/yahoo/vespa/clustercontroller/utils/staterestapi/DummyStateApi.java
index a54653ddd13..3c120f031d3 100644
--- a/clustercontroller-utils/src/test/java/com/yahoo/vespa/clustercontroller/utils/staterestapi/DummyStateApi.java
+++ b/clustercontroller-utils/src/test/java/com/yahoo/vespa/clustercontroller/utils/staterestapi/DummyStateApi.java
@@ -71,6 +71,10 @@ public class DummyStateApi implements StateRestAPI {
                 }
                 return result;
             }
+            @Override
+            public DistributionStates getDistributionStates() {
+                return null;
+            }
         };
     }
     private UnitResponse getClusterState(final DummyBackend.Cluster cluster, final int recursive) {
@@ -91,6 +95,10 @@ public class DummyStateApi implements StateRestAPI {
                 }
                 return result;
             }
+            @Override
+            public DistributionStates getDistributionStates() {
+                return null;
+            }
         };
     }
     private UnitResponse getNodeState(final DummyBackend.Node node) {
@@ -135,6 +143,10 @@ public class DummyStateApi implements StateRestAPI {
                     }
                 };
             }
+            @Override
+            public DistributionStates getDistributionStates() {
+                return null;
+            }
         };
 
     }