Add tests

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

@@ -27,14 +27,6 @@ public class PaginationNextPage {
     private long minCzxid;
     private int minCzxidOffset;
 
-    public PaginationNextPage() {
-    }
-
-    public PaginationNextPage(long minCzxid, int minCzxidOffset) {
-        this.minCzxid = minCzxid;
-        this.minCzxidOffset = minCzxidOffset;
-    }
-
     public long getMinCzxid() {
         return minCzxid;
     }

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

@@ -183,8 +183,8 @@ public class DataTree {
     // Constants used to calculate response packet length for the paginated children list.
     // packetLength = (childNameLength + PAGINATION_PACKET_CHILD_EXTRA_BYTES) * numChildren
     //                + PAGINATION_PACKET_BASE_BYTES
-    private static final int PAGINATION_PACKET_BASE_BYTES;
-    private static final int PAGINATION_PACKET_CHILD_EXTRA_BYTES;
+    public static final int PAGINATION_PACKET_BASE_BYTES;
+    public static final int PAGINATION_PACKET_CHILD_EXTRA_BYTES;
 
     static {
         PAGINATION_PACKET_BASE_BYTES = getPaginationPacketLength(Collections.emptyList());
@@ -868,7 +868,7 @@ public List<String> getPaginatedChildren(String path, Stat stat, Watcher watcher
                     n.copyStat(stat);
                 }
                 allChildren = n.getChildren();
-                if (isPacketLengthBelowMaxBuffer(computeChildrenPacketLength(allChildren))) {
+                if (canPacketFitInMaxBuffer(computeChildrenPacketLength(allChildren))) {
                     isBelowMaxBuffer = true;
                     if (watcher != null) {
                         childWatches.addWatch(path, watcher);
@@ -916,7 +916,7 @@ public List<String> getPaginatedChildren(String path, Stat stat, Watcher watcher
                  index++) {
                 String child = targetChildren.get(index).getPath();
                 packetLength += child.length() + PAGINATION_PACKET_CHILD_EXTRA_BYTES;
-                if (!isPacketLengthBelowMaxBuffer(packetLength)) {
+                if (!canPacketFitInMaxBuffer(packetLength)) {
                     // Stop adding more children to ensure packet is below max buffer
                     break;
                 }
@@ -939,8 +939,8 @@ public List<String> getPaginatedChildren(String path, Stat stat, Watcher watcher
         return paginatedChildren;
     }
 
-    private boolean isPacketLengthBelowMaxBuffer(int packetLength) {
-        return packetLength < BinaryInputArchive.maxBuffer;
+    private boolean canPacketFitInMaxBuffer(int packetLength) {
+        return packetLength <= BinaryInputArchive.maxBuffer;
     }
 
     private void buildChildrenPathWithStat(DataNode n, String path, Stat stat, long minCzxId,

zookeeper-server/src/test/java/org/apache/zookeeper/server/DataTreeTest.java

@@ -310,7 +310,7 @@ public void run() {
     }
 
     @Test(timeout = 60000)
-    public void getChildrenPaginated() throws NodeExistsException, NoNodeException {
+    public void testGetChildrenPaginated() throws NodeExistsException, NoNodeException {
         final String rootPath = "/children";
         final int firstCzxId = 1000;
         final int countNodes = 10;
@@ -373,7 +373,7 @@ public void getChildrenPaginated() throws NodeExistsException, NoNodeException {
     }
 
     @Test(timeout = 60000)
-    public void getChildrenPaginatedWithOffset() throws NodeExistsException, NoNodeException {
+    public void testGetChildrenPaginatedWithOffset() throws NodeExistsException, NoNodeException {
         final String rootPath = "/children";
         final int childrenCzxId = 1000;
         final int countNodes = 9;
@@ -455,18 +455,25 @@ public void getChildrenPaginatedWithOffset() throws NodeExistsException, NoNodeE
     }
 
     @Test(timeout = 60000)
-    public void getChildrenPaginatedEmpty() throws NodeExistsException, NoNodeException {
+    public void testGetChildrenPaginatedEmpty() throws NodeExistsException, NoNodeException {
         final String rootPath = "/children";
 
         //  Create the parent node
         DataTree dt = new DataTree();
         dt.createNode(rootPath, new byte[0], null, 0, dt.getNode("/").stat.getCversion() + 1, 1, 1);
 
-        //  Asking from a negative would give me all children, and set the watch
+        // Asking from a negative would give me all children, and set the watch
+        // This goes to the pagination branch.
         int curWatchCount = dt.getWatchCount();
         List<String> result = dt.getPaginatedChildren(rootPath, null, new DummyWatcher(), 100, -1, 0, null);
         assertTrue("The result should be empty", result.isEmpty());
         assertEquals("The watch should have been set", curWatchCount + 1, dt.getWatchCount());
+
+        // Specify max int to fetch all children - trying to return all all without sorting, and set the watch
+        curWatchCount = dt.getWatchCount();
+        result = dt.getPaginatedChildren(rootPath, null, new DummyWatcher(), Integer.MAX_VALUE, 0, 0, null);
+        assertTrue("The result should be empty", result.isEmpty());
+        assertEquals("The watch should have been set", curWatchCount + 1, dt.getWatchCount());
     }
 
     private class DummyWatcher implements Watcher {

zookeeper-server/src/test/java/org/apache/zookeeper/test/GetChildrenPaginatedTest.java

@@ -28,16 +28,21 @@
 import java.util.Set;
 import java.util.UUID;
 import org.apache.jute.BinaryInputArchive;
+import org.apache.jute.BinaryOutputArchive;
 import org.apache.zookeeper.CreateMode;
 import org.apache.zookeeper.KeeperException;
 import org.apache.zookeeper.OpResult;
+import org.apache.zookeeper.PaginationNextPage;
 import org.apache.zookeeper.RemoteIterator;
 import org.apache.zookeeper.Transaction;
 import org.apache.zookeeper.WatchedEvent;
 import org.apache.zookeeper.Watcher;
+import org.apache.zookeeper.ZooDefs;
 import org.apache.zookeeper.ZooDefs.Ids;
 import org.apache.zookeeper.ZooKeeper;
 import org.apache.zookeeper.data.Stat;
+import org.apache.zookeeper.server.DataTree;
+import org.apache.zookeeper.server.quorum.BufferStats;
 import org.junit.Assert;
 import org.junit.Test;
 
@@ -366,9 +371,7 @@ public void testGetAllChildrenPaginatedOnePage() throws KeeperException, Interru
     }
 
     /*
-     * Tests if all children's packet exceeds jute.maxbuffer, it can still successfully fetch them.
-     * The packet length computation formula is also tested through this test. Otherwise, it'll
-     * fail to return all children with the paginated API.
+     * Tests if all children's packet exceeds jute.maxbuffer, the paginated API can still successfully fetch them.
      */
     @Test(timeout = 60000)
     public void testGetAllChildrenPaginatedMultiPages() throws InterruptedException, KeeperException {
@@ -406,4 +409,125 @@ public void testGetAllChildrenPaginatedMultiPages() throws InterruptedException,
         Assert.assertEquals(numChildren, actualChildren.size());
         Assert.assertEquals(expectedChildren, new HashSet<>(actualChildren));
     }
+
+    /*
+     * Tests the packet length formula in DataTree as expected. The packet length calculated
+     * is equal to the actual buffer size of the client response.
+     */
+    @Test(timeout = 60000)
+    public void testPacketLengthFormula() throws KeeperException, InterruptedException {
+        final String basePath = "/testPagination-" + UUID.randomUUID().toString();
+        zk.create(basePath, new byte[0], Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
+
+        Set<String> expectedChildren = new HashSet<>();
+        int numChildren = 50 + random.nextInt(50);
+
+        for (int i = 0; i < numChildren; i++) {
+            String child = UUID.randomUUID().toString();
+            String childPath = basePath + "/" + child;
+            zk.create(childPath, new byte[0], Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
+            expectedChildren.add(child);
+        }
+
+        // Fetch max number of children that can fit into one page. The client response stat should record the same
+        // buffer size as the packet length calculated by the formula in DataTree.
+        PaginationNextPage nextPage = new PaginationNextPage();
+
+        // First page
+        List<String> firstPage = zk.getChildren(basePath, null, numChildren, 0, 0, nextPage);
+        BufferStats clientResponseStats = serverFactory.getZooKeeperServer().serverStats().getClientResponseStats();
+        int expectedPacketLength =
+                (DataTree.PAGINATION_PACKET_CHILD_EXTRA_BYTES + UUID.randomUUID().toString().length())
+                        * firstPage.size() + DataTree.PAGINATION_PACKET_BASE_BYTES;
+        int actualBufferSize = clientResponseStats.getLastBufferSize();
+
+        Assert.assertEquals("The page should be the last page.",
+                ZooDefs.GetChildrenPaginated.lastPageMinCzxid, nextPage.getMinCzxid());
+        Assert.assertEquals("The client response buffer size should be equal to the packet length calculated by "
+                + "pagination formula in DataTree.", expectedPacketLength, actualBufferSize);
+        Assert.assertTrue(actualBufferSize <= BinaryInputArchive.maxBuffer);
+
+        Assert.assertEquals(expectedChildren, new HashSet<>(firstPage));
+    }
+
+    /*
+     * Tests below logic:
+     * 1. the packet length formula is as expected. The calculated length is equal to the actual
+     * client response's buffer size.
+     * 2. logic of adding children on server side is correct. It does not miss the last child.
+     * 3. watch is only added in the last page.
+     */
+    @Test(timeout = 60000)
+    public void testMaxNumChildrenPageWatch() throws KeeperException, InterruptedException {
+        // Get the max number of UUID children that can return in a page.
+        int numChildren = (BinaryInputArchive.maxBuffer - DataTree.PAGINATION_PACKET_BASE_BYTES)
+                / (UUID.randomUUID().toString().length() + DataTree.PAGINATION_PACKET_CHILD_EXTRA_BYTES);
+        // The packet will exceed 1 MB.
+        numChildren++;
+
+        final String basePath = "/testPagination-" + UUID.randomUUID().toString();
+
+        zk.create(basePath, new byte[0], Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
+
+        Set<String> expectedChildren = new HashSet<>();
+
+        for (int i = 0; i < numChildren; i += 1000) {
+            Transaction transaction = zk.transaction();
+            for (int j = i; j < i + 1000 && j < numChildren; j++) {
+                String child = UUID.randomUUID().toString();
+                String childPath = basePath + "/" + child;
+                transaction.create(childPath, new byte[0], Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
+                expectedChildren.add(child);
+            }
+            transaction.commit();
+        }
+
+        // Fetch max number of children that can fit into one page. The client response stat should record the same
+        // buffer size as the packet length calculated by the formula in DataTree.
+        PaginationNextPage nextPage = new PaginationNextPage();
+        FireOnlyOnceWatcher fireOnlyOnceWatcher = new FireOnlyOnceWatcher();
+
+        // First page
+        List<String> firstPage = zk.getChildren(basePath, fireOnlyOnceWatcher, numChildren, 0, 0, nextPage);
+        BufferStats clientResponseStats = serverFactory.getZooKeeperServer().serverStats().getClientResponseStats();
+        int expectedPacketLength =
+                (DataTree.PAGINATION_PACKET_CHILD_EXTRA_BYTES + UUID.randomUUID().toString().length())
+                        * firstPage.size() + DataTree.PAGINATION_PACKET_BASE_BYTES;
+        int actualBufferSize = clientResponseStats.getLastBufferSize();
+
+        Assert.assertNotEquals("The page should not be the last page.",
+                ZooDefs.GetChildrenPaginated.lastPageMinCzxid, nextPage.getMinCzxid());
+        Assert.assertEquals("The client response buffer size should be equal to the packet length calculated by "
+                + "pagination formula in DataTree.", expectedPacketLength, actualBufferSize);
+        Assert.assertTrue(actualBufferSize <= BinaryInputArchive.maxBuffer);
+
+        // Verify the index on server is correct: not adding the last child in the first page.
+        Assert.assertEquals(numChildren - 1, firstPage.size());
+
+        // Modify the first child of each page.
+        // This should not trigger additional watches or create duplicates in the set of children returned
+        zk.setData(basePath + "/" + firstPage.get(0), new byte[3], -1);
+
+        synchronized(fireOnlyOnceWatcher) {
+            Assert.assertEquals("Watch should not have fired yet", 0, fireOnlyOnceWatcher.watchFiredCount);
+        }
+
+        // Second page
+        List<String> secondPage = zk.getChildren(basePath, fireOnlyOnceWatcher, numChildren, nextPage.getMinCzxid(),
+                nextPage.getMinCzxidOffset(), nextPage);
+
+        // Verify the logic of adding children to page in DataTree is correct.
+        Assert.assertEquals(ZooDefs.GetChildrenPaginated.lastPageMinCzxid, nextPage.getMinCzxid());
+        Assert.assertEquals("Should expect only 1 child left in the second page", 1, secondPage.size());
+
+        zk.setData(basePath + "/" + secondPage.get(0), new byte[3], -1);
+
+        synchronized(fireOnlyOnceWatcher) {
+            Assert.assertEquals("Watch has fired", 0, fireOnlyOnceWatcher.watchFiredCount);
+        }
+
+        Set<String> actualSet = new HashSet<>(firstPage);
+        actualSet.addAll(secondPage);
+        Assert.assertEquals(expectedChildren, actualSet);
+    }
 }