ZOOKEEPER-4219: Quota checks break setData in multi transactions (backport)

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

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) {

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:

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();