SOLR-10374: Implement set-policy and remove-policy APIs

solr/core/src/java/org/apache/solr/cloud/autoscaling/AutoScalingHandler.java

@@ -106,10 +106,56 @@ public class AutoScalingHandler extends RequestHandlerBase implements Permission
         case "resume-trigger":
           handleResumeTrigger(req, rsp, op);
           break;
+        case "set-policy":
+          handleSetPolicies(req, rsp, op);
+          break;
+        case "remove-policy":
+          handleRemovePolicy(req, rsp, op);
       }
     }
   }
 
+  private void handleRemovePolicy(SolrQueryRequest req, SolrQueryResponse rsp, CommandOperation op) throws KeeperException, InterruptedException {
+    String policyName = (String) op.getCommandData();
+
+    if (policyName.trim().length() == 0) {
+      throw new SolrException(SolrException.ErrorCode.BAD_REQUEST, "The policy name cannot be empty");
+    }
+    Map<String, Object> autoScalingConf = zkReadAutoScalingConf(container.getZkController().getZkStateReader());
+    Map<String, Object> policies = (Map<String, Object>) autoScalingConf.get("policies");
+    if (policies == null || !policies.containsKey(policyName)) {
+      throw new SolrException(SolrException.ErrorCode.BAD_REQUEST, "No policy exists with name: " + policyName);
+    }
+
+    zkSetPolicies(container.getZkController().getZkStateReader(), policyName, null);
+    rsp.getValues().add("result", "success");
+  }
+
+  private void handleSetPolicies(SolrQueryRequest req, SolrQueryResponse rsp, CommandOperation op) throws KeeperException, InterruptedException {
+    String policyName = op.getStr("name");
+
+    if (policyName == null || policyName.trim().length() == 0) {
+      throw new SolrException(SolrException.ErrorCode.BAD_REQUEST, "The policy name cannot be null or empty");
+    }
+
+    Set<String> keys = op.getDataMap().keySet();
+    boolean isValid = false;
+    for (String key : keys) {
+      if (key.equals("conditions") || key.equals("preferences")) isValid = true;
+      else if(!key.equals("name")){
+        isValid = false;
+        break;
+      }
+    }
+    if (!isValid) {
+      throw new SolrException(SolrException.ErrorCode.BAD_REQUEST,
+          "No conditions or peferences are specified for the policy " + policyName);
+    }
+
+    zkSetPolicies(container.getZkController().getZkStateReader(), policyName, op.getValuesExcluding("name"));
+    rsp.getValues().add("result", "success");
+  }
+
   private void handleResumeTrigger(SolrQueryRequest req, SolrQueryResponse rsp, CommandOperation op) throws KeeperException, InterruptedException {
     String triggerName = op.getStr("name");
 
@@ -396,6 +442,30 @@ public class AutoScalingHandler extends RequestHandlerBase implements Permission
     }
   }
 
+  private void zkSetPolicies(ZkStateReader reader, String policyName, Map<String, Object> policyProperties) throws KeeperException, InterruptedException {
+    while (true) {
+      Stat stat = new Stat();
+      ZkNodeProps loaded = null;
+      byte[] data = reader.getZkClient().getData(SOLR_AUTOSCALING_CONF_PATH, null, stat, true);
+      loaded = ZkNodeProps.load(data);
+      Map<String, Object> policies = (Map<String, Object>) loaded.get("policies");
+      if (policies == null) policies = new HashMap<>(1);
+      if (policyProperties != null) {
+        policies.put(policyName, policyProperties);
+      } else {
+        policies.remove(policyName);
+      }
+      loaded = loaded.plus("policies", policies);
+      try {
+        reader.getZkClient().setData(SOLR_AUTOSCALING_CONF_PATH, Utils.toJSON(loaded), stat.getVersion(), true);
+      } catch (KeeperException.BadVersionException bve) {
+        // somebody else has changed the configuration so we must retry
+        continue;
+      }
+      break;
+    }
+  }
+
   private Map<String, Object> zkReadAutoScalingConf(ZkStateReader reader) throws KeeperException, InterruptedException {
     byte[] data = reader.getZkClient().getData(SOLR_AUTOSCALING_CONF_PATH, null, null, true);
     ZkNodeProps loaded = ZkNodeProps.load(data);

solr/core/src/resources/apispec/autoscaling.Commands.json

@@ -179,6 +179,36 @@
       "required": [
         "name"
       ]
+    },
+    "set-policy" : {
+      "type":"object",
+      "description": "The set-policy command allows you to add and update policies",
+      "properties": {
+        "name": {
+          "type": "string",
+          "description": "The name of the policy"
+        },
+        "conditions" : {
+          "type": "array",
+          "description": "Conditions of the policy"
+        },
+        "preferences": {
+          "type": "array",
+          "description": "Preferences of the policy"
+        }
+      },
+      "oneOf": [{
+        "required": ["name", "preferences"]
+      },{
+        "required": ["name", "conditions"]
+      },{
+        "required": ["name", "preferences", "conditions"]
+      }],
+      "additionalProperties": false
+    },
+    "remove-policy": {
+      "description": "The remove-policy command allows you to remove a policy",
+      "type": "string"
     }
   }
 }

solr/core/src/test/org/apache/solr/cloud/autoscaling/AutoScalingHandlerTest.java

@@ -345,6 +345,88 @@ public class AutoScalingHandlerTest extends SolrCloudTestCase {
     } catch (HttpSolrClient.RemoteSolrException e) {
       // expected
     }
+
+    // add multiple poilicies
+    String setPolicyCommand =  "{\n" +
+        "\t\"set-policy\": {\n" +
+        "\t\t\"name\" : \"default\",\n" +
+        "\t\t\"preferences\": [\n" +
+        "\t\t\t{\n" +
+        "\t\t\t\t\"minimize\": \"replicas\",\n" +
+        "\t\t\t\t\"precision\": 3\n" +
+        "\t\t\t},\n" +
+        "\t\t\t{\n" +
+        "\t\t\t\t\"maximize\": \"freedisk\",\n" +
+        "\t\t\t\t\"precision\": 100\n" +
+        "\t\t\t}\n" +
+        "\t\t]\t\t\n" +
+        "\t}, \n" +
+        "\t\"set-policy\": {\n" +
+        "\t\t\"name\" : \"policy1\",\n" +
+        "\t\t\"preferences\": [\n" +
+        "\t\t\t{\n" +
+        "\t\t\t\t\"minimize\": \"cpu\",\n" +
+        "\t\t\t\t\"precision\": 10\n" +
+        "\t\t\t}\n" +
+        "\t\t]\n" +
+        "\t}\n" +
+        "}";
+    req = new AutoScalingRequest(SolrRequest.METHOD.POST, path, setPolicyCommand);
+    response = solrClient.request(req);
+    assertEquals(response.get("result").toString(), "success");
+    data = zkClient().getData(SOLR_AUTOSCALING_CONF_PATH, null, null, true);
+    loaded = ZkNodeProps.load(data);
+    Map<String, Object> policies = (Map<String, Object>) loaded.get("policies");
+    assertNotNull(policies);
+    assertNotNull(policies.get("default"));
+    assertNotNull(policies.get("policy1"));
+
+    // update default policy
+    setPolicyCommand = "{\n" +
+        "\t\"set-policy\": {\n" +
+        "\t\t\"name\" : \"default\",\n" +
+        "\t\t\"preferences\": [\n" +
+        "\t\t\t{\n" +
+        "\t\t\t\t\"minimize\": \"replicas\",\n" +
+        "\t\t\t\t\"precision\": 3\n" +
+        "\t\t\t}\n" +
+        "\t\t]\t\t\n" +
+        "\t}\n" +
+        "}";
+    req = new AutoScalingRequest(SolrRequest.METHOD.POST, path, setPolicyCommand);
+    response = solrClient.request(req);
+    assertEquals(response.get("result").toString(), "success");
+    data = zkClient().getData(SOLR_AUTOSCALING_CONF_PATH, null, null, true);
+    loaded = ZkNodeProps.load(data);
+    policies = (Map<String, Object>) loaded.get("policies");
+    Map<String,Object> properties = (Map<String, Object>) policies.get("default");
+    List preferences = (List) properties.get("preferences");
+    assertEquals(1, preferences.size());
+
+    // policy is not valid
+    setPolicyCommand = "{\n" +
+        "\t\"set-policy\": {\n" +
+        "\t\t\"name\" : \"default\"\t\n" +
+        "\t}\n" +
+        "}";
+    req = new AutoScalingRequest(SolrRequest.METHOD.POST, path, setPolicyCommand);
+    try {
+      response = solrClient.request(req);
+      fail("Adding a policy without conditions or preferences should have failed");
+    } catch (HttpSolrClient.RemoteSolrException e) {
+      // expected
+    }
+
+    String removePolicyCommand = "{\n" +
+        "\t\"remove-policy\" : \"policy1\"\n" +
+        "}";
+    req = new AutoScalingRequest(SolrRequest.METHOD.POST, path, removePolicyCommand);
+    response = solrClient.request(req);
+    assertEquals(response.get("result").toString(), "success");
+    data = zkClient().getData(SOLR_AUTOSCALING_CONF_PATH, null, null, true);
+    loaded = ZkNodeProps.load(data);
+    policies = (Map<String, Object>) loaded.get("policies");
+    assertNull(policies.get("policy1"));
   }
 
   static class AutoScalingRequest extends SolrRequest {