ZOOKEEPER-3301: Enfore the quota limit

zookeeper-docs/src/main/resources/markdown/zookeeperQuotas.md

@@ -46,6 +46,9 @@ and `-b` (for bytes).
 The ZooKeeper quota is stored in ZooKeeper itself in /zookeeper/quota. To disable other people from
 changing the quota's set the ACL for /zookeeper/quota such that only admins are able to read and write to it.
 
+- The quota supports the soft and hard quota.the soft quota just logs the warning info when exceeding the quota, but the hard quota
+also throws a QuotaExceededException. When setting soft and hard quota on the same path, the hard quota has the priority.
+
 <a name="Listing+Quotas"></a>
 
 ### Listing Quotas

zookeeper-server/src/main/java/org/apache/zookeeper/KeeperException.java

@@ -146,6 +146,8 @@ public static KeeperException create(Code code) {
                 return new ReconfigDisabledException();
             case REQUESTTIMEOUT:
                 return new RequestTimeoutException();
+            case QuotaExceeded:
+                return new QuotaExceededException();
             case OK:
             default:
                 throw new IllegalArgumentException("Invalid exception code");
@@ -397,7 +399,9 @@ public static enum Code implements CodeDeprecated {
         /** Request not completed within max allowed time.*/
         REQUESTTIMEOUT (-122),
         /** Attempts to perform a reconfiguration operation when reconfiguration feature is disabled. */
-        RECONFIGDISABLED(-123);
+        RECONFIGDISABLED(-123),
+        /** Exceeded the quota that setted on the path.*/
+        QuotaExceeded(-125);
 
         private static final Map<Integer,Code> lookup
             = new HashMap<Integer,Code>();
@@ -484,6 +488,8 @@ static String getCodeMessage(Code code) {
                 return "No such watcher";
             case RECONFIGDISABLED:
                 return "Reconfig is disabled";
+            case QuotaExceeded:
+                return "Quota has exceeded";
             default:
                 return "Unknown error " + code;
         }
@@ -856,4 +862,17 @@ public RequestTimeoutException() {
             super(Code.REQUESTTIMEOUT);
         }
     }
+
+    /**
+     * @see Code#QuotaExceeded
+     */
+    @InterfaceAudience.Public
+    public static class QuotaExceededException extends KeeperException {
+        public QuotaExceededException() {
+            super(Code.QuotaExceeded);
+        }
+        public QuotaExceededException(String path) {
+            super(Code.QuotaExceeded, path);
+        }
+    }
 }

zookeeper-server/src/main/java/org/apache/zookeeper/Quotas.java

@@ -48,9 +48,19 @@ public class Quotas {
      * return the quota path associated with this
      * prefix
      * @param path the actual path in zookeeper.
-     * @return the limit quota path
+     * @return the quota path
      */
     public static String quotaPath(String path) {
+        return quotaZookeeper + path;
+    }
+
+    /**
+     * return the limit quota path associated with this
+     * prefix
+     * @param path the actual path in zookeeper.
+     * @return the limit quota path
+     */
+    public static String limitPath(String path) {
         return quotaZookeeper + path +
         "/" + limitNode;
     }
@@ -65,4 +75,14 @@ public static String statPath(String path) {
         return quotaZookeeper + path + "/" +
         statNode;
     }
+
+    /**
+     * return the real path associated with this
+     * quotaPath.
+     * @param quotaPath the quotaPath which's started with /zookeeper/quota
+     * @return the real path associated with this quotaPath.
+     */
+    public static String trimQuotaPath(String quotaPath) {
+        return quotaPath.substring(quotaZookeeper.length());
+    }
 }

zookeeper-server/src/main/java/org/apache/zookeeper/StatsTrack.java

@@ -18,14 +18,25 @@
 
 package org.apache.zookeeper;
 
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+import org.apache.commons.lang.StringUtils;
+
 /**
  * a class that represents the stats associated with quotas
  */
 public class StatsTrack {
-    private int count;
-    private long bytes;
-    private String countStr = "count";
-    private String byteStr = "bytes";
+
+    private static final String countStr = "count";
+    private static final String countHardLimitStr = "countHardLimit";
+
+    private static final String byteStr = "bytes";
+    private static final String byteHardLimitStr = "byteHardLimit";
+
+    private final Map<String, Long> stats = new HashMap<>();
 
     /**
      * a default constructor for
@@ -34,22 +45,23 @@ public class StatsTrack {
     public StatsTrack() {
         this(null);
     }
+
     /**
-     * the stat string should be of the form count=int,bytes=long
-     * if stats is called with null the count and bytes are initialized
-     * to -1.
-     * @param stats the stat string to be intialized with
+     * the stat string should be of the form key1str=long,key2str=long,..
+     * where either , or ; are valid separators
+     * uninitialized values are returned as -1
+     * @param stats the stat string to be initialized with
      */
     public StatsTrack(String stats) {
-        if (stats == null) {
-            stats = "count=-1,bytes=-1";
+        this.stats.clear();
+        if (stats == null || stats.length() == 0) {
+            return;
         }
-        String[] split = stats.split(",");
-        if (split.length != 2) {
-            throw new IllegalArgumentException("invalid string " + stats);
+        String[] keyValuePairs = stats.split("[,;]+");
+        for (String keyValuePair : keyValuePairs) {
+            String[] kv = keyValuePair.split("=");
+            this.stats.put(kv[0], Long.parseLong(StringUtils.isEmpty(kv[1]) ? "-1" : kv[1]));
         }
-        count = Integer.parseInt(split[0].split("=")[1]);
-        bytes = Long.parseLong(split[1].split("=")[1]);
     }
 
 
@@ -59,7 +71,7 @@ public StatsTrack(String stats) {
      * @return the count as part of this string
      */
     public int getCount() {
-        return this.count;
+        return (int) getValue(countStr);
     }
 
     /**
@@ -69,7 +81,25 @@ public int getCount() {
      *            the count to set with
      */
     public void setCount(int count) {
-        this.count = count;
+        setValue(countStr, count);
+    }
+
+    /**
+     * get the count of nodes allowed as part of quota (hard limit)
+     *
+     * @return the count as part of this string
+     */
+    public int getCountHardLimit() {
+        return (int) getValue(countHardLimitStr);
+    }
+
+    /**
+     * set the count hard limit
+     *
+     * @param count the count limit to set
+     */
+    public void setCountHardLimit(int count) {
+        setValue(countHardLimitStr, count);
     }
 
     /**
@@ -78,24 +108,97 @@ public void setCount(int count) {
      * @return the bytes as part of this string
      */
     public long getBytes() {
-        return this.bytes;
+        return getValue(byteStr);
     }
 
     /**
-     * set teh bytes for this stat tracker.
+     * set the bytes for this stat tracker.
      *
      * @param bytes
      *            the bytes to set with
      */
     public void setBytes(long bytes) {
-        this.bytes = bytes;
+        setValue(byteStr, bytes);
+    }
+
+    /**
+     * get the count of bytes allowed as part of quota (hard limit)
+     *
+     * @return the bytes as part of this string
+     */
+    public long getByteHardLimit() {
+        return getValue(byteHardLimitStr);
+    }
+
+    /**
+     * set the byte hard limit
+     *
+     * @param bytes the byte limit to set
+     */
+    public void setByteHardLimit(long bytes) {
+        setValue(byteHardLimitStr, bytes);
+    }
+
+    /**
+     * get helper to lookup a given key
+     *
+     * @param key the key to lookup
+     * @return key's value or -1 if it doesn't exist
+     */
+    private long getValue(String key) {
+        Long val = this.stats.get(key);
+        return val == null ? -1 : val.longValue();
+    }
+
+    /**
+     * set helper to set the value for the specified key
+     *
+     * @param key   the key to set
+     * @param value the value to set
+     */
+    private void setValue(String key, long value) {
+        this.stats.put(key, value);
     }
 
-    @Override
     /*
      * returns the string that maps to this stat tracking.
+     *
+     * Builds a string of the form
+     * "count=4,bytes=5=;countHardLimit=10;byteHardLimit=10"
+     *
+     * This string is slightly hacky to preserve compatibility with 3.4.3 and
+     * older parser. In particular, count must be first, bytes must be second,
+     * all new fields must use a separator that is not a "," (so, ";"), and the
+     * seemingly spurious "=" after the bytes field is essential to allowing
+     * it to be parseable by the old parsing code.
      */
+    @Override
     public String toString() {
-        return countStr + "=" + count + "," + byteStr + "=" + bytes;
+        StringBuilder buf = new StringBuilder();
+        ArrayList<String> keys = new ArrayList<>(stats.keySet());
+
+        // Special handling for count=,byte= to enforce them coming first
+        // for backwards compatibility
+        keys.remove(countStr);
+        keys.remove(byteStr);
+        buf.append(countStr);
+        buf.append("=");
+        buf.append(getCount());
+        buf.append(",");
+        buf.append(byteStr);
+        buf.append("=");
+        buf.append(getBytes());
+        if (!keys.isEmpty()) {
+            // Add extra = to trick old parsing code so it will ignore new flags
+            buf.append("=");
+            Collections.sort(keys);
+            for(String key : keys) {
+                buf.append(";");
+                buf.append(key);
+                buf.append("=");
+                buf.append(stats.get(key));
+            }
+        }
+        return buf.toString();
     }
 }

zookeeper-server/src/main/java/org/apache/zookeeper/ZooKeeperMain.java

@@ -295,8 +295,7 @@ protected void connectToZK(String newHost) throws InterruptedException, IOExcept
         zk = new ZooKeeperAdmin(host, Integer.parseInt(cl.getOption("timeout")), new MyWatcher(), readOnly);
     }
 
-    public static void main(String args[]) throws CliException, IOException, InterruptedException
-    {
+    public static void main(String args[]) throws CliException, IOException, InterruptedException, KeeperException {
         ZooKeeperMain main = new ZooKeeperMain(args);
         main.run();
     }
@@ -311,7 +310,7 @@ public ZooKeeperMain(ZooKeeper zk) {
       this.zk = zk;
     }
 
-    void run() throws CliException, IOException, InterruptedException {
+    void run() throws CliException, IOException, InterruptedException, KeeperException {
         if (cl.getCommand() == null) {
             System.out.println("Welcome to ZooKeeper!");
 
@@ -372,7 +371,7 @@ void run() throws CliException, IOException, InterruptedException {
         System.exit(exitCode);
     }
 
-    public void executeLine(String line) throws CliException, InterruptedException, IOException {
+    public void executeLine(String line) throws CliException, InterruptedException, IOException, KeeperException {
       if (!line.equals("")) {
         cl.parseCommand(line);
         addToHistory(commandCount,line);
@@ -425,8 +424,8 @@ public static boolean delQuota(ZooKeeper zk, String path,
             boolean bytes, boolean numNodes)
         throws KeeperException, IOException, InterruptedException
     {
-        String parentPath = Quotas.quotaZookeeper + path;
-        String quotaPath = Quotas.quotaZookeeper + path + "/" + Quotas.limitNode;
+        String parentPath = Quotas.quotaPath(path);
+        String quotaPath = Quotas.limitPath(path);
         if (zk.exists(quotaPath, false) == null) {
             System.out.println("Quota does not exist for " + path);
             return true;
@@ -518,7 +517,7 @@ public static boolean createQuota(ZooKeeper zk, String path,
         // are not the children then this path
         // is an ancestor of some path that
         // already has quota
-        String realPath = Quotas.quotaZookeeper + path;
+        String realPath = Quotas.quotaPath(path);
         try {
             List<String> children = zk.getChildren(realPath, false);
             for (String child: children) {
@@ -589,7 +588,7 @@ public static boolean createQuota(ZooKeeper zk, String path,
         return true;
     }
 
-    protected boolean processCmd(MyCommandOptions co) throws CliException, IOException, InterruptedException {
+    protected boolean processCmd(MyCommandOptions co) throws CliException, IOException, InterruptedException, KeeperException {
         boolean watch = false;
         try {
             watch = processZKCmd(co);
@@ -601,7 +600,7 @@ protected boolean processCmd(MyCommandOptions co) throws CliException, IOExcepti
         return watch;
     }
 
-    protected boolean processZKCmd(MyCommandOptions co) throws CliException, IOException, InterruptedException {
+    protected boolean processZKCmd(MyCommandOptions co) throws CliException, IOException, InterruptedException, KeeperException {
         String[] args = co.getArgArray();
         String cmd = co.getCommand();
         if (args.length < 1) {

zookeeper-server/src/main/java/org/apache/zookeeper/cli/CliCommand.java

@@ -17,8 +17,11 @@
  */
 package org.apache.zookeeper.cli;
 
+import java.io.IOException;
 import java.io.PrintStream;
 import java.util.Map;
+
+import org.apache.zookeeper.KeeperException;
 import org.apache.zookeeper.ZooKeeper;
 
 /**

zookeeper-server/src/main/java/org/apache/zookeeper/cli/CliWrapperException.java

@@ -53,6 +53,8 @@ private static String getMessage(Throwable cause) {
                 return "No quorum of new config is connected and " +
                         "up-to-date with the leader of last commmitted config - try invoking reconfiguration after " +
                         "new servers are connected and synced";
+            } else if (keeperException instanceof KeeperException.QuotaExceededException) {
+                return "Quota has exceeded : " + keeperException.getPath();
             }
         }
         return cause.getMessage();