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,21 +169,7 @@ TEST(WriteBatchTest, Blob) { "Put(k3, v3)@2", PrintContents(&batch)); - struct Handler : 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() + ")"; - } - } handler; + TestHandler handler; batch.Iterate(&handler); ASSERT_EQ( "Put(k1, v1)" @@ -178,6 +182,45 @@ TEST(WriteBatchTest, Blob) { handler.seen); } +TEST(WriteBatchTest, Continue) { + WriteBatch batch; + + struct Handler : public TestHandler { + int num_seen = 0; + virtual void Put(const Slice& key, const Slice& value) { + ++num_seen; + TestHandler::Put(key, value); + } + virtual void Merge(const Slice& key, const Slice& value) { + ++num_seen; + TestHandler::Merge(key, value); + } + virtual void LogData(const Slice& blob) { + ++num_seen; + TestHandler::LogData(blob); + } + virtual void Delete(const Slice& key) { + ++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)" + "LogData(blob1)" + "Delete(k1)", + handler.seen); +} + } // namespace leveldb int main(int argc, char** argv) { 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;