ZOOKEEPER-3116: Make the DataTree.approximateDataSize more efficient

Cache the approximate data tree size, and update it when txns applied to data tree, this could make this query not that expensive.

Author: Fangmin Lyu <allenlyu@fb.com>

Reviewers: andor@apache.org

Closes #594 from lvfangmin/ZOOKEEPER-3116 and squashes the following commits:

ad2ce216 [Fangmin Lyu] move to synchronize when visiting node.data
2e0cbb17 [Fangmin Lyu] add test case to test cached approximate data tree size
21243f5e [Fangmin Lyu] cache the approximate data size in DataTree

src/java/main/org/apache/zookeeper/server/DataNode.java

@@ -135,11 +135,6 @@ public class DataNode implements Record {
         return Collections.unmodifiableSet(children);
     }
 
-    public synchronized long getApproximateDataSize() {
-        if(null==data) return 0;
-        return data.length;
-    }
-
     synchronized public void copyStat(Stat to) {
         to.setAversion(stat.getAversion());
         to.setCtime(stat.getCtime());

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

@@ -66,6 +66,7 @@ import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicLong;
 
 /**
  * This class maintains the tree data structure. It doesn't have any networking
@@ -90,6 +91,9 @@ public class DataTree {
 
     private final WatchManager childWatches = new WatchManager();
 
+    /** cached total size of paths and data for all DataNodes */
+    private final AtomicLong nodeDataSize = new AtomicLong(0);
+
     /** the root of zookeeper tree */
     private static final String rootZookeeper = "/";
 
@@ -199,13 +203,24 @@ public class DataTree {
         for (Map.Entry<String, DataNode> entry : nodes.entrySet()) {
             DataNode value = entry.getValue();
             synchronized (value) {
-                result += entry.getKey().length();
-                result += value.getApproximateDataSize();
+                result += getNodeSize(entry.getKey(), value.data);
             }
         }
         return result;
     }
 
+    /**
+     * Get the size of the node based on path and data length.
+     */
+    private static long getNodeSize(String path, byte[] data) {
+        return (path == null ? 0 : path.length())
+                + (data == null ? 0 : data.length);
+    }
+
+    public long cachedApproximateDataSize() {
+        return nodeDataSize.get();
+    }
+
     /**
      * This is a pointer to the root of the DataTree. It is the source of truth,
      * but we usually use the nodes hashmap to find nodes in the tree.
@@ -236,6 +251,8 @@ public class DataTree {
         nodes.put(quotaZookeeper, quotaDataNode);
 
         addConfigNode();
+
+        nodeDataSize.set(approximateDataSize());
     }
 
     /**
@@ -468,6 +485,7 @@ public class DataTree {
             Long longval = aclCache.convertAcls(acl);
             DataNode child = new DataNode(data, longval, stat);
             parent.addChild(childName);
+            nodeDataSize.addAndGet(getNodeSize(path, child.data));
             nodes.put(path, child);
             EphemeralType ephemeralType = EphemeralType.get(ephemeralOwner);
             if (ephemeralType == EphemeralType.CONTAINER) {
@@ -534,6 +552,7 @@ public class DataTree {
         nodes.remove(path);
         synchronized (node) {
             aclCache.removeUsage(node.acl);
+            nodeDataSize.addAndGet(-getNodeSize(path, node.data));
         }
         DataNode parent = nodes.get(parentName);
         if (parent == null) {
@@ -609,6 +628,7 @@ public class DataTree {
           this.updateBytes(lastPrefix, (data == null ? 0 : data.length)
               - (lastdata == null ? 0 : lastdata.length));
         }
+        nodeDataSize.addAndGet(getNodeSize(path, data) - getNodeSize(path, lastdata));
         dataWatches.triggerWatch(path, EventType.NodeDataChanged);
         return s;
     }
@@ -1183,6 +1203,7 @@ public class DataTree {
         aclCache.deserialize(ia);
         nodes.clear();
         pTrie.clear();
+        nodeDataSize.set(0);
         String path = ia.readString("path");
         while (!"/".equals(path)) {
             DataNode node = new DataNode();
@@ -1220,6 +1241,9 @@ public class DataTree {
             path = ia.readString("path");
         }
         nodes.put("/", root);
+
+        nodeDataSize.set(approximateDataSize());
+
         // we are done with deserializing the
         // the datatree
         // update the quotas - create path trie

src/java/main/org/apache/zookeeper/server/DataTreeBean.java

@@ -25,17 +25,17 @@ import org.apache.zookeeper.jmx.ZKMBeanInfo;
  */
 public class DataTreeBean implements DataTreeMXBean, ZKMBeanInfo {
     DataTree dataTree;
-    
+
     public DataTreeBean(org.apache.zookeeper.server.DataTree dataTree){
         this.dataTree = dataTree;
     }
-    
+
     public int getNodeCount() {
         return dataTree.getNodeCount();
     }
 
     public long approximateDataSize() {
-        return dataTree.approximateDataSize();
+        return dataTree.cachedApproximateDataSize();
     }
 
     public int countEphemerals() {

src/java/main/org/apache/zookeeper/server/admin/Commands.java

@@ -324,7 +324,7 @@ public class Commands {
 
             response.put("watch_count", zkdb.getDataTree().getWatchCount());
             response.put("ephemerals_count", zkdb.getDataTree().getEphemeralsCount());
-            response.put("approximate_data_size", zkdb.getDataTree().approximateDataSize());
+            response.put("approximate_data_size", zkdb.getDataTree().cachedApproximateDataSize());
 
             OSMXBean osMbean = new OSMXBean();
             response.put("open_file_descriptor_count", osMbean.getOpenFileDescriptorCount());

src/java/main/org/apache/zookeeper/server/command/MonitorCommand.java

@@ -60,7 +60,7 @@ public class MonitorCommand extends AbstractFourLetterCommand {
 
         print("watch_count", zkdb.getDataTree().getWatchCount());
         print("ephemerals_count", zkdb.getDataTree().getEphemeralsCount());
-        print("approximate_data_size", zkdb.getDataTree().approximateDataSize());
+        print("approximate_data_size", zkdb.getDataTree().cachedApproximateDataSize());
 
         OSMXBean osMbean = new OSMXBean();
         if (osMbean != null && osMbean.getUnix() == true) {

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

@@ -281,4 +281,24 @@ public class DataTreeTest extends ZKTestCase {
             "expected to have the same acl", ZooDefs.Ids.OPEN_ACL_UNSAFE,
             tree.getACL("/bug", new Stat()));
     }
+
+    @Test
+    public void testCachedApproximateDataSize() throws Exception {
+        DataTree dt = new DataTree();
+        long initialSize = dt.approximateDataSize();
+        Assert.assertEquals(dt.cachedApproximateDataSize(), dt.approximateDataSize());
+
+        // create a node
+        dt.createNode("/testApproximateDataSize", new byte[20], null, -1, 1, 1, 1);
+        dt.createNode("/testApproximateDataSize1", new byte[20], null, -1, 1, 1, 1);
+        Assert.assertEquals(dt.cachedApproximateDataSize(), dt.approximateDataSize());
+
+        // update data
+        dt.setData("/testApproximateDataSize1", new byte[32], -1, 1, 1);
+        Assert.assertEquals(dt.cachedApproximateDataSize(), dt.approximateDataSize());
+
+        // delete a node
+        dt.deleteNode("/testApproximateDataSize", -1);
+        Assert.assertEquals(dt.cachedApproximateDataSize(), dt.approximateDataSize());
+    }
 }