From cb703c9d03a01cc6eea1f2fe2f167aa1e11dd011 Mon Sep 17 00:00:00 2001
From: Jim Paton <jpaton2@fb.com>
Date: Wed, 21 Aug 2013 18:27:48 -0700
Subject: [PATCH] Allow WriteBatch::Handler to abort iteration

Summary:
Sometimes you don't need to iterate through the whole WriteBatch. This diff makes the Handler member functions return a bool that indicates whether to abort or not. If they return true, the iteration stops.

One thing I just thought of is that this will break backwards-compability. Maybe it would be better to add a virtual member function WriteBatch::Handler::ShouldAbort() that returns false by default. Comments requested.

I still have to add a new unit test for the abort code, but let's finalize the API first.

Test Plan: make -j32 check

Reviewers: dhruba, haobo, vamsi, emayanke

Reviewed By: dhruba

CC: leveldb

Differential Revision: https://reviews.facebook.net/D12339
---
 db/write_batch.cc             |  6 +++-
 db/write_batch_test.cc        | 65 +++++++++++++++++++++++++++++------
 include/leveldb/write_batch.h |  4 +++
 3 files changed, 63 insertions(+), 12 deletions(-)

diff --git a/db/write_batch.cc b/db/write_batch.cc
index 317599542..6537f3721 100644
--- a/db/write_batch.cc
+++ b/db/write_batch.cc
@@ -48,6 +48,10 @@ void WriteBatch::Handler::LogData(const Slice& blob) {
   // them.
 }
 
+bool WriteBatch::Handler::Continue() {
+  return true;
+}
+
 void WriteBatch::Clear() {
   rep_.clear();
   rep_.resize(kHeader);
@@ -66,7 +70,7 @@ Status WriteBatch::Iterate(Handler* handler) const {
   input.remove_prefix(kHeader);
   Slice key, value, blob;
   int found = 0;
-  while (!input.empty()) {
+  while (!input.empty() && handler->Continue()) {
     char tag = input[0];
     input.remove_prefix(1);
     switch (tag) {
diff --git a/db/write_batch_test.cc b/db/write_batch_test.cc
index ce6104c7f..c20c9c0c0 100644
--- a/db/write_batch_test.cc
+++ b/db/write_batch_test.cc
@@ -134,6 +134,24 @@ TEST(WriteBatchTest, Append) {
   ASSERT_EQ(4, b1.Count());
 }
 
+namespace {
+  struct TestHandler : public WriteBatch::Handler {
+    std::string seen;
+    virtual void Put(const Slice& key, const Slice& value) {
+      seen += "Put(" + key.ToString() + ", " + value.ToString() + ")";
+    }
+    virtual void Merge(const Slice& key, const Slice& value) {
+      seen += "Merge(" + key.ToString() + ", " + value.ToString() + ")";
+    }
+    virtual void LogData(const Slice& blob) {
+      seen += "LogData(" + blob.ToString() + ")";
+    }
+    virtual void Delete(const Slice& key) {
+      seen += "Delete(" + key.ToString() + ")";
+    }
+  };
+}
+
 TEST(WriteBatchTest, Blob) {
   WriteBatch batch;
   batch.Put(Slice("k1"), Slice("v1"));
@@ -151,30 +169,55 @@ TEST(WriteBatchTest, Blob) {
             "Put(k3, v3)@2",
             PrintContents(&batch));
 
-  struct Handler : public WriteBatch::Handler {
-    std::string seen;
+  TestHandler handler;
+  batch.Iterate(&handler);
+  ASSERT_EQ(
+            "Put(k1, v1)"
+            "Put(k2, v2)"
+            "Put(k3, v3)"
+            "LogData(blob1)"
+            "Delete(k2)"
+            "LogData(blob2)"
+            "Merge(foo, bar)",
+            handler.seen);
+}
+
+TEST(WriteBatchTest, Continue) {
+  WriteBatch batch;
+
+  struct Handler : public TestHandler {
+    int num_seen = 0;
     virtual void Put(const Slice& key, const Slice& value) {
-      seen += "Put(" + key.ToString() + ", " + value.ToString() + ")";
+      ++num_seen;
+      TestHandler::Put(key, value);
     }
     virtual void Merge(const Slice& key, const Slice& value) {
-      seen += "Merge(" + key.ToString() + ", " + value.ToString() + ")";
+      ++num_seen;
+      TestHandler::Merge(key, value);
     }
     virtual void LogData(const Slice& blob) {
-      seen += "LogData(" + blob.ToString() + ")";
+      ++num_seen;
+      TestHandler::LogData(blob);
     }
     virtual void Delete(const Slice& key) {
-      seen += "Delete(" + key.ToString() + ")";
+      ++num_seen;
+      TestHandler::Delete(key);
+    }
+    virtual bool Continue() override {
+      return num_seen < 3;
     }
   } handler;
+
+  batch.Put(Slice("k1"), Slice("v1"));
+  batch.PutLogData(Slice("blob1"));
+  batch.Delete(Slice("k1"));
+  batch.PutLogData(Slice("blob2"));
+  batch.Merge(Slice("foo"), Slice("bar"));
   batch.Iterate(&handler);
   ASSERT_EQ(
             "Put(k1, v1)"
-            "Put(k2, v2)"
-            "Put(k3, v3)"
             "LogData(blob1)"
-            "Delete(k2)"
-            "LogData(blob2)"
-            "Merge(foo, bar)",
+            "Delete(k1)",
             handler.seen);
 }
 
diff --git a/include/leveldb/write_batch.h b/include/leveldb/write_batch.h
index b2cfc5348..4564b01fb 100644
--- a/include/leveldb/write_batch.h
+++ b/include/leveldb/write_batch.h
@@ -69,6 +69,10 @@ class WriteBatch {
     // The default implementation of LogData does nothing.
     virtual void LogData(const Slice& blob);
     virtual void Delete(const Slice& key) = 0;
+    // Continue is called by WriteBatch::Iterate. If it returns false,
+    // iteration is halted. Otherwise, it continues iterating. The default
+    // implementation always returns true.
+    virtual bool Continue();
   };
   Status Iterate(Handler* handler) const;