diff --git a/zookeeper-server/src/main/java/org/apache/zookeeper/server/PrepRequestProcessor.java b/zookeeper-server/src/main/java/org/apache/zookeeper/server/PrepRequestProcessor.java index d263fd75c43..04089ccd45d 100644 --- a/zookeeper-server/src/main/java/org/apache/zookeeper/server/PrepRequestProcessor.java +++ b/zookeeper-server/src/main/java/org/apache/zookeeper/server/PrepRequestProcessor.java @@ -38,6 +38,7 @@ import org.apache.zookeeper.KeeperException.Code; import org.apache.zookeeper.MultiOperationRecord; import org.apache.zookeeper.Op; +import org.apache.zookeeper.StatsTrack; import org.apache.zookeeper.ZooDefs; import org.apache.zookeeper.ZooDefs.OpCode; import org.apache.zookeeper.common.ConfigException; @@ -104,6 +105,17 @@ public class PrepRequestProcessor extends ZooKeeperCriticalThread implements Req private final boolean digestEnabled; private DigestCalculator digestCalculator; + /** + * The quota-relevant changes (node count and data bytes, per quota + * prefix) of the operations that were already validated as part of the + * request currently being processed. The operations of a multi + * transaction are validated one by one before any of them is applied to + * the data tree, so the quota stat nodes do not yet account for the + * earlier operations of the same transaction. Only accessed from the + * request processing thread. + */ + private final Map pendingQuotaChanges = new HashMap<>(); + ZooKeeperServer zks; public enum DigestOpCode { @@ -361,6 +373,7 @@ protected void pRequest2Txn(int type, long zxid, Request request, Record record) if (nodeRecord.childCount > 0) { throw new KeeperException.NotEmptyException(path); } + zks.checkQuota(path, nodeRecord.data, null, OpCode.delete, pendingQuotaChanges); request.setTxn(new DeleteTxn(path)); parentRecord = parentRecord.duplicate(request.getHdr().getZxid()); parentRecord.childCount--; @@ -381,7 +394,7 @@ protected void pRequest2Txn(int type, long zxid, Request request, Record record) validatePath(path, request.sessionId); nodeRecord = getRecordForPath(path); zks.checkACL(request.cnxn, nodeRecord.acl, ZooDefs.Perms.WRITE, request.authInfo, path, null); - zks.checkQuota(path, nodeRecord.data, setDataRequest.getData(), OpCode.setData); + zks.checkQuota(path, nodeRecord.data, setDataRequest.getData(), OpCode.setData, pendingQuotaChanges); int newVersion = checkAndIncVersion(nodeRecord.stat.getVersion(), setDataRequest.getVersion(), path); request.setTxn(new SetDataTxn(path, setDataRequest.getData(), newVersion)); nodeRecord = nodeRecord.duplicate(request.getHdr().getZxid()); @@ -679,7 +692,7 @@ private void pRequest2TxnCreate(int type, Request request, Record record) throws throw new KeeperException.NoChildrenForEphemeralsException(path); } int newCversion = parentRecord.stat.getCversion() + 1; - zks.checkQuota(path, null, data, OpCode.create); + zks.checkQuota(path, null, data, OpCode.create, pendingQuotaChanges); if (type == OpCode.createContainer) { request.setTxn(new CreateContainerTxn(path, data, listACL, newCversion)); } else if (type == OpCode.createTTL) { @@ -758,6 +771,7 @@ private static int checkAndIncVersion(int currentVersion, int expectedVersion, S protected void pRequest(Request request) throws RequestProcessorException { request.setHdr(null); request.setTxn(null); + pendingQuotaChanges.clear(); if (!request.isThrottled()) { pRequestHelper(request); diff --git a/zookeeper-server/src/main/java/org/apache/zookeeper/server/ZooKeeperServer.java b/zookeeper-server/src/main/java/org/apache/zookeeper/server/ZooKeeperServer.java index 080f4f8638d..a695be406de 100644 --- a/zookeeper-server/src/main/java/org/apache/zookeeper/server/ZooKeeperServer.java +++ b/zookeeper-server/src/main/java/org/apache/zookeeper/server/ZooKeeperServer.java @@ -2111,30 +2111,64 @@ public void checkACL(ServerCnxn cnxn, List acl, int perm, List ids, Str * @param data * the data to be set, or {@code null} for none * @param type - * currently, create and setData need to check quota + * currently, create and setData need to check quota; delete + * only records the released count and bytes + * @param pendingChanges + * accumulates, per quota prefix, the count and byte changes + * made by the operations of the current request that were + * already validated; a multi transaction is validated before + * any of it is applied, so the stat nodes alone do not + * reflect the earlier operations of the transaction */ - public void checkQuota(String path, byte[] lastData, byte[] data, int type) throws KeeperException.QuotaExceededException { + public void checkQuota(String path, byte[] lastData, byte[] data, int type, + Map pendingChanges) throws KeeperException.QuotaExceededException { if (!enforceQuota) { return; } long dataBytes = (data == null) ? 0 : data.length; + long lastDataBytes = (lastData == null) ? 0 : lastData.length; ZKDatabase zkDatabase = getZKDatabase(); String lastPrefix = zkDatabase.getDataTree().getMaxPrefixWithQuota(path); if (StringUtils.isEmpty(lastPrefix)) { return; } - final String namespace = PathUtils.getTopNamespace(path); + long bytesDiff; + long countDiff; switch (type) { case OpCode.create: - checkQuota(lastPrefix, dataBytes, 1, namespace); + bytesDiff = dataBytes; + countDiff = 1; break; case OpCode.setData: - checkQuota(lastPrefix, dataBytes - (lastData == null ? 0 : lastData.length), 0, namespace); + bytesDiff = dataBytes - lastDataBytes; + countDiff = 0; + break; + case OpCode.delete: + // a delete cannot exceed a quota, but the released count and + // bytes have to be remembered so that later operations in the + // same transaction are checked against the correct usage + bytesDiff = -lastDataBytes; + countDiff = -1; break; default: throw new IllegalArgumentException("Unsupported OpCode for checkQuota: " + type); } + + if (type != OpCode.delete) { + final String namespace = PathUtils.getTopNamespace(path); + checkQuota(lastPrefix, bytesDiff, countDiff, namespace, pendingChanges); + } + + StatsTrack pending = pendingChanges.get(lastPrefix); + if (pending == null) { + pending = new StatsTrack(); + pending.setCount(0); + pending.setBytes(0); + pendingChanges.put(lastPrefix, pending); + } + pending.setCount(pending.getCount() + countDiff); + pending.setBytes(pending.getBytes() + bytesDiff); } /** @@ -2148,9 +2182,13 @@ public void checkQuota(String path, byte[] lastData, byte[] data, int type) thro * the diff to be added to the count * @param namespace * the namespace for collecting quota exceeded errors + * @param pendingChanges + * the count and byte changes made by operations that were + * already validated as part of the current request but are + * not yet reflected in the quota stat node */ - private void checkQuota(String lastPrefix, long bytesDiff, long countDiff, String namespace) - throws KeeperException.QuotaExceededException { + private void checkQuota(String lastPrefix, long bytesDiff, long countDiff, String namespace, + Map pendingChanges) throws KeeperException.QuotaExceededException { LOG.debug("checkQuota: lastPrefix={}, bytesDiff={}, countDiff={}", lastPrefix, bytesDiff, countDiff); // now check the quota we set @@ -2187,9 +2225,17 @@ private void checkQuota(String lastPrefix, long bytesDiff, long countDiff, Strin currentStats = new StatsTrack(node.data); } + long pendingCount = 0; + long pendingBytes = 0; + StatsTrack pending = pendingChanges.get(lastPrefix); + if (pending != null) { + pendingCount = pending.getCount(); + pendingBytes = pending.getBytes(); + } + //check the Count Quota if (checkCountQuota) { - long newCount = currentStats.getCount() + countDiff; + long newCount = currentStats.getCount() + pendingCount + countDiff; boolean isCountHardLimit = limitStats.getCountHardLimit() > -1; long countLimit = isCountHardLimit ? limitStats.getCountHardLimit() : limitStats.getCount(); @@ -2205,7 +2251,7 @@ private void checkQuota(String lastPrefix, long bytesDiff, long countDiff, Strin //check the Byte Quota if (checkByteQuota) { - long newBytes = currentStats.getBytes() + bytesDiff; + long newBytes = currentStats.getBytes() + pendingBytes + bytesDiff; boolean isByteHardLimit = limitStats.getByteHardLimit() > -1; long byteLimit = isByteHardLimit ? limitStats.getByteHardLimit() : limitStats.getBytes(); if (newBytes > byteLimit) { diff --git a/zookeeper-server/src/test/java/org/apache/zookeeper/test/ZooKeeperQuotaTest.java b/zookeeper-server/src/test/java/org/apache/zookeeper/test/ZooKeeperQuotaTest.java index 211f91a476f..f5b687a5736 100644 --- a/zookeeper-server/src/test/java/org/apache/zookeeper/test/ZooKeeperQuotaTest.java +++ b/zookeeper-server/src/test/java/org/apache/zookeeper/test/ZooKeeperQuotaTest.java @@ -473,6 +473,90 @@ public void testMultiCreateThenSetDataShouldFail() throws Exception { assertNull(zk.exists(subPath, null)); } + @Test + public void testMultiCreatesExceedingCountHardQuotaShouldFail() throws Exception { + final String path = "/c"; + + zk.create(path, null, Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT); + + final StatsTrack st = new StatsTrack(); + st.setCountHardLimit(3); + SetQuotaCommand.createQuota(zk, path, st); + + // each create on its own stays within the limit; only the + // transaction as a whole exceeds it + final List ops = Arrays.asList( + Op.create(path + "/1", null, Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT), + Op.create(path + "/2", null, Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT), + Op.create(path + "/3", null, Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT)); + + try { + zk.multi(ops); + fail("should fail transaction when the creates together exceed the count hard quota"); + } catch (QuotaExceededException e) { + //expected + } + + assertNull(zk.exists(path + "/1", null)); + assertNull(zk.exists(path + "/2", null)); + assertNull(zk.exists(path + "/3", null)); + } + + @Test + public void testMultiCreatesExceedingBytesHardQuotaShouldFail() throws Exception { + final String path = "/b"; + + zk.create(path, null, Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT); + + final byte[] data5b = "Hello".getBytes(StandardCharsets.UTF_8); + + final StatsTrack st = new StatsTrack(); + st.setByteHardLimit(3 * data5b.length - 1); + SetQuotaCommand.createQuota(zk, path, st); + + final List ops = Arrays.asList( + Op.create(path + "/1", data5b, Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT), + Op.create(path + "/2", data5b, Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT), + Op.create(path + "/3", data5b, Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT)); + + try { + zk.multi(ops); + fail("should fail transaction when the creates together exceed the byte hard quota"); + } catch (QuotaExceededException e) { + //expected + } + + assertNull(zk.exists(path + "/1", null)); + assertNull(zk.exists(path + "/2", null)); + assertNull(zk.exists(path + "/3", null)); + } + + @Test + public void testMultiDeleteThenCreateWithinCountHardQuotaShouldWork() throws Exception { + final String path = "/replace"; + final String oldChild = path + "/old"; + final String newChild = path + "/new"; + + zk.create(path, null, Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT); + + final StatsTrack st = new StatsTrack(); + st.setCountHardLimit(2); + SetQuotaCommand.createQuota(zk, path, st); + + zk.create(oldChild, null, Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT); + + // the node count is at the hard limit, but the transaction deletes + // as much as it creates, so it should be allowed + final List ops = Arrays.asList( + Op.delete(oldChild, -1), + Op.create(newChild, null, Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT)); + + zk.multi(ops); + + assertNull(zk.exists(oldChild, null)); + assertNotNull(zk.exists(newChild, null)); + } + @Test public void testDeleteBytesQuota() throws Exception {