From 1e2b5f7321fca982350f5f2fbb36d06ada21de20 Mon Sep 17 00:00:00 2001 From: Arkesh Mishra <118651144+arkmish@users.noreply.github.com> Date: Mon, 19 Oct 2026 20:13:15 +0530 Subject: [PATCH] ZOOKEEPER-4219: Quota checks break setData in multi transactions Backport of apache/zookeeper eb1569e4f (released in 3.7.0/3.8.0). checkQuota now receives the previous data from the prepared ChangeRecord instead of looking the node up in DataTree, which NPE'd for a create followed by setData in the same multi(). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../resources/markdown/zookeeperQuotas.md | 2 +- .../server/PrepRequestProcessor.java | 4 +-- .../zookeeper/server/ZooKeeperServer.java | 8 ++---- .../test/ZooKeeperQuotaEnforcementTest.java | 26 +++++++++++++++++++ 4 files changed, 31 insertions(+), 9 deletions(-) diff --git a/zookeeper-docs/src/main/resources/markdown/zookeeperQuotas.md b/zookeeper-docs/src/main/resources/markdown/zookeeperQuotas.md index de51e7c8286..1ee6d17cf0e 100644 --- a/zookeeper-docs/src/main/resources/markdown/zookeeperQuotas.md +++ b/zookeeper-docs/src/main/resources/markdown/zookeeperQuotas.md @@ -31,7 +31,7 @@ limitations under the License. ZooKeeper has both namespace and bytes quotas. You can use the ZooKeeperMain class to set up quotas. This fork backports the native quota enforcement released in Apache ZooKeeper 3.8.7 -(ZOOKEEPER-3301). Per-namespace quota metrics +(ZOOKEEPER-3301 and ZOOKEEPER-4219). Per-namespace quota metrics (ZOOKEEPER-4211) are not included. Soft limits produce warnings; hard limits reject writes with `QuotaExceededException`. Both checks require the startup JVM property `-Dzookeeper.enforceQuota=true`, which defaults to 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 93266645862..b217c863549 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 @@ -399,7 +399,7 @@ protected void pRequest2Txn(int type, long zxid, Request request, Record record, validatePath(path, request.sessionId); nodeRecord = getRecordForPath(path); zks.checkACL(request, nodeRecord.acl, ZooDefs.Perms.WRITE, path, null); - zks.checkQuota(path, setDataRequest.getData(), OpCode.setData); + zks.checkQuota(path, nodeRecord.data, setDataRequest.getData(), OpCode.setData); int newVersion = checkAndIncVersion(nodeRecord.stat.getVersion(), setDataRequest.getVersion(), path); request.setTxn(new SetDataTxn(path, setDataRequest.getData(), newVersion)); nodeRecord = nodeRecord.duplicate(request.getHdr().getZxid()); @@ -703,7 +703,7 @@ private void pRequest2TxnCreate(int type, Request request, Record record, boolea throw new KeeperException.NoChildrenForEphemeralsException(path); } int newCversion = parentRecord.stat.getCversion() + 1; - zks.checkQuota(path, data, OpCode.create); + zks.checkQuota(path, null, data, OpCode.create); if (type == OpCode.createContainer) { request.setTxn(new CreateContainerTxn(path, data, listACL, newCversion)); } else if (type == OpCode.createTTL) { 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 8eb61d317d7..00e1eb1aafe 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 @@ -2081,10 +2081,11 @@ private void checkACLWithContext(ServerAuthenticationProvider.ServerObjs serverO * Check committed quota usage before preparing a create or setData transaction. * * @param path path being modified + * @param lastData previous data from the prepared change record * @param data proposed data * @param type create or setData opcode */ - public void checkQuota(String path, byte[] data, int type) throws KeeperException.QuotaExceededException { + public void checkQuota(String path, byte[] lastData, byte[] data, int type) throws KeeperException.QuotaExceededException { if (!enforceQuota) { return; } @@ -2098,11 +2099,6 @@ public void checkQuota(String path, byte[] data, int type) throws KeeperExceptio checkQuota(lastPrefix, dataBytes, 1); break; case OpCode.setData: - DataNode node = getZKDatabase().getDataTree().getNode(path); - byte[] lastData; - synchronized (node) { - lastData = node.getData(); - } checkQuota(lastPrefix, dataBytes - (lastData == null ? 0 : lastData.length), 0); break; default: diff --git a/zookeeper-server/src/test/java/org/apache/zookeeper/test/ZooKeeperQuotaEnforcementTest.java b/zookeeper-server/src/test/java/org/apache/zookeeper/test/ZooKeeperQuotaEnforcementTest.java index f2324949382..1772d5b4567 100644 --- a/zookeeper-server/src/test/java/org/apache/zookeeper/test/ZooKeeperQuotaEnforcementTest.java +++ b/zookeeper-server/src/test/java/org/apache/zookeeper/test/ZooKeeperQuotaEnforcementTest.java @@ -24,11 +24,13 @@ import static org.junit.Assert.assertNull; import static org.junit.Assert.assertTrue; import static org.junit.Assert.fail; +import java.util.Arrays; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; import org.apache.zookeeper.CreateMode; import org.apache.zookeeper.KeeperException; +import org.apache.zookeeper.Op; import org.apache.zookeeper.Quotas; import org.apache.zookeeper.StatsTrack; import org.apache.zookeeper.ZooDefs.Ids; @@ -96,6 +98,30 @@ public void testAboveLimitUpdatesMatchReleasedBehavior() throws Exception { assertEquals(4L, usage(zk, "/limited").getBytes()); } + @Test + public void testMultiCreateThenSetDataUsesPendingData() throws Exception { + ZooKeeper zk = createClient(); + createQuota(zk, "/limited", -1, 5, new byte[0]); + zk.multi(Arrays.asList( + Op.create("/limited/child", new byte[1], Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT.toFlag()), + Op.setData("/limited/child", new byte[4], -1))); + assertArrayEquals(new byte[4], zk.getData("/limited/child", false, null)); + assertEquals(4L, usage(zk, "/limited").getBytes()); + } + + @Test + public void testRejectedMultiDoesNotApplyEarlierOperations() throws Exception { + ZooKeeper zk = createClient(); + createQuota(zk, "/limited", -1, 5, new byte[0]); + expectQuotaExceeded(() -> zk.multi(Arrays.asList( + Op.create("/limited/child", new byte[0], Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT.toFlag()), + Op.setData("/limited/child", new byte[6], -1)))); + assertNull(zk.exists("/limited/child", false)); + assertEquals(1L, usage(zk, "/limited").getCount()); + assertEquals(0L, usage(zk, "/limited").getBytes()); + zk.create("/limited/after", new byte[5], Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT); + } + @Test public void testAsyncQuotaErrorKeepsConnectionUsable() throws Exception { ZooKeeper zk = createClient();