ZOOKEEPER-3354: Improve efficiency of DeleteAllCommand

Author: Brian Nixon <nixon@fb.com>

Reviewers: andor@apache.org, szepet95@gmail.com, eolivelli@apache.org

Closes #899 from enixon/delete-all

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

@@ -24,6 +24,9 @@ import java.util.Collections;
 import java.util.List;
 import java.util.Queue;
 
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.Semaphore;
+import org.apache.zookeeper.AsyncCallback.MultiCallback;
 import org.apache.zookeeper.AsyncCallback.StringCallback;
 import org.apache.zookeeper.AsyncCallback.VoidCallback;
 import org.apache.zookeeper.KeeperException.Code;
@@ -44,7 +47,7 @@ public class ZKUtil {
      *
      * @throws IllegalArgumentException if an invalid path is specified
      */
-    public static void deleteRecursive(ZooKeeper zk, final String pathRoot)
+    public static boolean deleteRecursive(ZooKeeper zk, final String pathRoot, final int batchSize)
         throws InterruptedException, KeeperException
     {
         PathUtils.validatePath(pathRoot);
@@ -52,12 +55,53 @@ public class ZKUtil {
         List<String> tree = listSubTreeBFS(zk, pathRoot);
         LOG.debug("Deleting " + tree);
         LOG.debug("Deleting " + tree.size() + " subnodes ");
-        for (int i = tree.size() - 1; i >= 0 ; --i) {
-            //Delete the leaves first and eventually get rid of the root
-            zk.delete(tree.get(i), -1); //Delete all versions of the node with -1.
+
+        return deleteInBatch(zk, tree, batchSize);
+    }
+
+    private static class BatchedDeleteCbContext {
+        public Semaphore sem;
+        public AtomicBoolean success;
+
+        public BatchedDeleteCbContext(int rateLimit) {
+            sem = new Semaphore(rateLimit);
+            success = new AtomicBoolean(true);
         }
     }
 
+    private static boolean deleteInBatch(ZooKeeper zk, List<String> tree, int batchSize)
+        throws InterruptedException
+    {
+        int rateLimit = 10;
+        List<Op> ops = new ArrayList<>();
+        BatchedDeleteCbContext context = new BatchedDeleteCbContext(rateLimit);
+        MultiCallback cb = (rc, path, ctx, opResults) -> {
+            ((BatchedDeleteCbContext)ctx).sem.release();
+            if (rc != Code.OK.intValue()) {
+                ((BatchedDeleteCbContext)ctx).success.set(false);
+            }
+        };
+
+        // Delete the leaves first and eventually get rid of the root
+        for (int i = tree.size() - 1; i >= 0 ; --i) {
+            // Create Op to delete all versions of the node with -1.
+            ops.add(Op.delete(tree.get(i), -1));
+
+            if (ops.size() == batchSize || i == 0) {
+                if (!context.success.get()) {
+                    // fail fast
+                    break;
+                }
+                context.sem.acquire();
+                zk.multi(ops, cb, context);
+                ops = new ArrayList<>();
+            }
+        }
+
+        // wait for all callbacks to finish
+        context.sem.acquire(rateLimit);
+        return context.success.get();
+    }
 
     /**
      * Recursively delete the node with the given path. (async version).

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

@@ -18,6 +18,7 @@
 package org.apache.zookeeper.cli;
 
 import org.apache.commons.cli.CommandLine;
+import org.apache.commons.cli.Option;
 import org.apache.commons.cli.Options;
 import org.apache.commons.cli.ParseException;
 import org.apache.commons.cli.Parser;
@@ -32,19 +33,23 @@ public class DeleteAllCommand extends CliCommand {
 
     private static Options options = new Options();
     private String[] args;
+    private CommandLine cl;
+
+    static {
+        options.addOption(new Option("b", true, "batch size"));
+    }
 
     public DeleteAllCommand() {
         this("deleteall");
     }
 
     public DeleteAllCommand(String cmdStr) {
-        super(cmdStr, "path");
+        super(cmdStr, "path [-b batch size]");
     }
     
     @Override
     public CliCommand parse(String[] cmdArgs) throws CliParseException {
         Parser parser = new PosixParser();
-        CommandLine cl;
         try {
             cl = parser.parse(options, cmdArgs);
         } catch (ParseException ex) {
@@ -61,10 +66,19 @@ public class DeleteAllCommand extends CliCommand {
     @Override
     public boolean exec() throws CliException {
         printDeprecatedWarning();
-        
+        int batchSize;
+        try {
+            batchSize = cl.hasOption("b") ? Integer.parseInt(cl.getOptionValue("b")) : 1000;
+        } catch (NumberFormatException e) {
+            throw new MalformedCommandException("-b argument must be an int value");
+        }
+
         String path = args[1];
         try {
-            ZKUtil.deleteRecursive(zk, path);
+            boolean success = ZKUtil.deleteRecursive(zk, path, batchSize);
+            if (!success) {
+                err.println("Failed to delete some node(s) in the subtree!");
+            }
         } catch (IllegalArgumentException ex) {
             throw new MalformedPathException(ex.getMessage());
         } catch (KeeperException|InterruptedException ex) {

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

@@ -23,6 +23,7 @@ import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.PrintStream;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 import java.util.concurrent.atomic.AtomicInteger;
 
@@ -34,6 +35,8 @@ import org.apache.zookeeper.client.HostProvider;
 import org.apache.zookeeper.client.StaticHostProvider;
 import org.apache.zookeeper.client.ZKClientConfig;
 import org.apache.zookeeper.common.StringUtils;
+import org.apache.zookeeper.data.ACL;
+import org.apache.zookeeper.data.Id;
 import org.apache.zookeeper.data.Stat;
 import org.apache.zookeeper.test.ClientBase;
 import org.junit.Assert;
@@ -49,7 +52,93 @@ public class ZooKeeperTest extends ClientBase {
     private static final String LINE_SEPARATOR = System.getProperty("line.separator", "\n");
 
     @Test
-    public void testDeleteRecursive() throws IOException, InterruptedException, CliException, KeeperException {
+    public void testDeleteRecursive()
+        throws IOException, InterruptedException, KeeperException
+    {
+        final ZooKeeper zk = createClient();
+        setupDataTree(zk);
+
+        Assert.assertTrue(ZKUtil.deleteRecursive(zk, "/a/c", 1000));
+        List<String> children = zk.getChildren("/a", false);
+        Assert.assertEquals("1 children - c should be deleted ", 1, children.size());
+        Assert.assertTrue(children.contains("b"));
+
+        Assert.assertTrue(ZKUtil.deleteRecursive(zk, "/a", 1000));
+        Assert.assertNull(zk.exists("/a", null));
+    }
+
+    @Test
+    public void testDeleteRecursiveFail()
+            throws IOException, InterruptedException, KeeperException
+    {
+        final ZooKeeper zk = createClient();
+        setupDataTree(zk);
+
+        ACL deleteProtection = new ACL(ZooDefs.Perms.DELETE,
+                new Id("digest", "user:tl+z3z0vO6PfPfEENfLF96E6pM0="/* password is test */));
+        List<ACL> acls = Arrays.asList(
+                new ACL(ZooDefs.Perms.READ, Ids.ANYONE_ID_UNSAFE),
+                deleteProtection
+        );
+
+        // poison the well
+        zk.create("/a/c/0/surprise", "".getBytes(), Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
+        Assert.assertEquals(1, zk.getACL("/a/c/0", new Stat()).size());
+        zk.setACL("/a/c/0", acls, -1);
+        Assert.assertEquals(2, zk.getACL("/a/c/0", new Stat()).size());
+
+        Assert.assertFalse(ZKUtil.deleteRecursive(zk, "/a/c", 1000));
+        List<String> children = zk.getChildren("/a", false);
+        Assert.assertEquals("2 children - c should fail to be deleted ", 2, children.size());
+        Assert.assertTrue(children.contains("b"));
+
+        Assert.assertTrue(ZKUtil.deleteRecursive(zk, "/a/b", 1000));
+        children = zk.getChildren("/a", false);
+        Assert.assertEquals("1 children - b should be deleted ", 1, children.size());
+
+        // acquire immunity to poison
+        zk.addAuthInfo(deleteProtection.getId().getScheme(), "user:test".getBytes());
+
+        Assert.assertTrue(ZKUtil.deleteRecursive(zk, "/a", 1000));
+        Assert.assertNull(zk.exists("/a", null));
+    }
+
+    private void setupDataTree(ZooKeeper zk) throws KeeperException, InterruptedException {
+        // making sure setdata works on /
+        zk.setData("/", "some".getBytes(), -1);
+        zk.create("/a", "some".getBytes(), Ids.OPEN_ACL_UNSAFE,
+                CreateMode.PERSISTENT);
+
+        zk.create("/a/b", "some".getBytes(), Ids.OPEN_ACL_UNSAFE,
+                CreateMode.PERSISTENT);
+
+        zk.create("/a/b/v", "some".getBytes(), Ids.OPEN_ACL_UNSAFE,
+                CreateMode.PERSISTENT);
+
+        for (int i = 1000; i < 3000; ++i) {
+            zk.create("/a/b/v/" + i, "some".getBytes(), Ids.OPEN_ACL_UNSAFE,
+                    CreateMode.PERSISTENT);
+        }
+
+        zk.create("/a/c", "some".getBytes(), Ids.OPEN_ACL_UNSAFE,
+                CreateMode.PERSISTENT);
+
+        zk.create("/a/c/v", "some".getBytes(), Ids.OPEN_ACL_UNSAFE,
+                CreateMode.PERSISTENT);
+
+        for (int i = 0; i < 500; ++i) {
+            zk.create("/a/c/" + i, "some".getBytes(), Ids.OPEN_ACL_UNSAFE,
+                    CreateMode.PERSISTENT);
+        }
+        List<String> children = zk.getChildren("/a", false);
+
+        Assert.assertEquals("2 children - b & c should be present ", 2, children.size());
+        Assert.assertTrue(children.contains("b"));
+        Assert.assertTrue(children.contains("c"));
+    }
+
+    @Test
+    public void testDeleteRecursiveCli() throws IOException, InterruptedException, CliException, KeeperException {
         final ZooKeeper zk = createClient();
         // making sure setdata works on /
         zk.setData("/", "some".getBytes(), -1);