From ceded4535d432eb201bc470c5bb3e781861d3bca Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Simon=20Gra=CC=88tzer?= <simon@graetzer.org>
Date: Wed, 10 Oct 2018 20:55:42 -0700
Subject: [PATCH] WriteBatch::Iterate wrongly returns Status::Corruption
 (#4478)

Summary:
Wrong I overwrite `WriteBatch::Handler::Continue` to return _false_ at some point, I always get the `Status::Corruption` error.
I don't think this check is used correctly here: The counter in `found` cannot reflect all entries in the WriteBatch when we exit the loop early.
Pull Request resolved: https://github.com/facebook/rocksdb/pull/4478

Differential Revision: D10317416

Pulled By: yiwu-arbug

fbshipit-source-id: cccae3382805035f9b3239b66682b5fcbba6bb61
---
 db/write_batch.cc | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/db/write_batch.cc b/db/write_batch.cc
index 295fba22e..3fcee08e4 100644
--- a/db/write_batch.cc
+++ b/db/write_batch.cc
@@ -414,8 +414,13 @@ Status WriteBatch::Iterate(Handler* handler) const {
   char tag = 0;
   uint32_t column_family = 0;  // default
   bool last_was_try_again = false;
-  while (((s.ok() && !input.empty()) || UNLIKELY(s.IsTryAgain())) &&
-         handler->Continue()) {
+  bool handler_continue = true;
+  while (((s.ok() && !input.empty()) || UNLIKELY(s.IsTryAgain()))) {
+    handler_continue = handler->Continue();
+    if (!handler_continue) {
+      break;
+    }
+
     if (LIKELY(!s.IsTryAgain())) {
       last_was_try_again = false;
       tag = 0;
@@ -583,7 +588,7 @@ Status WriteBatch::Iterate(Handler* handler) const {
   if (!s.ok()) {
     return s;
   }
-  if (found != WriteBatchInternal::Count(this)) {
+  if (handler_continue && found != WriteBatchInternal::Count(this)) {
     return Status::Corruption("WriteBatch has wrong count");
   } else {
     return Status::OK();