support PopSavePoint for WriteBatch

Summary:
Try to fix https://github.com/facebook/rocksdb/issues/1969
Closes https://github.com/facebook/rocksdb/pull/2170

Differential Revision: D4907333

Pulled By: yiwu-arbug

fbshipit-source-id: 417b420ff668e6c2fd0dad42a94c57385012edc5

HISTORY.md

@@ -1,5 +1,8 @@
 # Rocksdb Change Log
 ## Unreleased
+### Public API Change
+* Introduce WriteBatch::PopSavePoint to pop the most recent save point explicitly.
+
 ### New Features
 * DB::ResetStats() to reset internal stats.
 * Statistics::Reset() to reset user stats.

db/c.cc

@@ -1362,6 +1362,10 @@ void rocksdb_writebatch_rollback_to_save_point(rocksdb_writebatch_t* b,
   SaveError(errptr, b->rep.RollbackToSavePoint());
 }
 
+void rocksdb_writebatch_pop_save_point(rocksdb_writebatch_t* b, char** errptr) {
+  SaveError(errptr, b->rep.PopSavePoint());
+}
+
 rocksdb_writebatch_wi_t* rocksdb_writebatch_wi_create(size_t reserved_bytes, unsigned char overwrite_key) {
   rocksdb_writebatch_wi_t* b = new rocksdb_writebatch_wi_t;
   b->rep = new WriteBatchWithIndex(BytewiseComparator(), reserved_bytes, overwrite_key);

db/c_test.c

@@ -587,10 +587,13 @@ int main(int argc, char** argv) {
   {
     rocksdb_writebatch_t* wb = rocksdb_writebatch_create();
     rocksdb_writebatch_set_save_point(wb);
+    rocksdb_writebatch_set_save_point(wb);
     const char* k_list[2] = {"z", "ap"};
     const size_t k_sizes[2] = {1, 2};
     const char* v_list[3] = {"x", "y", "z"};
     const size_t v_sizes[3] = {1, 1, 1};
+    rocksdb_writebatch_pop_save_point(wb, &err);
+    CheckNoError(err);
     rocksdb_writebatch_putv(wb, 2, k_list, k_sizes, 3, v_list, v_sizes);
     rocksdb_writebatch_rollback_to_save_point(wb, &err);
     CheckNoError(err);

db/write_batch.cc

@@ -803,6 +803,17 @@ Status WriteBatch::RollbackToSavePoint() {
   return Status::OK();
 }
 
+Status WriteBatch::PopSavePoint() {
+  if (save_points_ == nullptr || save_points_->stack.size() == 0) {
+    return Status::NotFound();
+  }
+
+  // Pop the most recent savepoint off the stack
+  save_points_->stack.pop();
+
+  return Status::OK();
+}
+
 class MemTableInserter : public WriteBatch::Handler {
 
   SequenceNumber sequence_;

db/write_batch_test.cc

@@ -861,6 +861,19 @@ TEST_F(WriteBatchTest, SavePointTest) {
   s = batch2.RollbackToSavePoint();
   ASSERT_TRUE(s.IsNotFound());
   ASSERT_EQ("", PrintContents(&batch2));
+
+  WriteBatch batch3;
+
+  s = batch3.PopSavePoint();
+  ASSERT_TRUE(s.IsNotFound());
+  ASSERT_EQ("", PrintContents(&batch3));
+
+  batch3.SetSavePoint();
+  batch3.Delete("A");
+
+  s = batch3.PopSavePoint();
+  ASSERT_OK(s);
+  ASSERT_EQ("Delete(A)@0", PrintContents(&batch3));
 }
 
 TEST_F(WriteBatchTest, MemoryLimitTest) {

include/rocksdb/c.h

@@ -457,6 +457,8 @@ extern ROCKSDB_LIBRARY_API void rocksdb_writebatch_set_save_point(
     rocksdb_writebatch_t*);
 extern ROCKSDB_LIBRARY_API void rocksdb_writebatch_rollback_to_save_point(
     rocksdb_writebatch_t*, char** errptr);
+extern ROCKSDB_LIBRARY_API void rocksdb_writebatch_pop_save_point(
+    rocksdb_writebatch_t*, char** errptr);
 
 /* Write batch with index */
 

include/rocksdb/utilities/write_batch_with_index.h

@@ -207,6 +207,12 @@ class WriteBatchWithIndex : public WriteBatchBase {
   //         or other Status on corruption.
   Status RollbackToSavePoint() override;
 
+  // Pop the most recent save point.
+  // If there is no previous call to SetSavePoint(), Status::NotFound()
+  // will be returned.
+  // Otherwise returns Status::OK().
+  Status PopSavePoint() override;
+
   void SetMaxBytes(size_t max_bytes) override;
 
  private:

include/rocksdb/write_batch.h

@@ -166,6 +166,12 @@ class WriteBatch : public WriteBatchBase {
   // Otherwise returns Status::OK().
   Status RollbackToSavePoint() override;
 
+  // Pop the most recent save point.
+  // If there is no previous call to SetSavePoint(), Status::NotFound()
+  // will be returned.
+  // Otherwise returns Status::OK().
+  Status PopSavePoint() override;
+
   // Support for iterating over the contents of a batch.
   class Handler {
    public:

include/rocksdb/write_batch_base.h

@@ -112,6 +112,12 @@ class WriteBatchBase {
   // Clear().
   virtual Status RollbackToSavePoint() = 0;
 
+  // Pop the most recent save point.
+  // If there is no previous call to SetSavePoint(), Status::NotFound()
+  // will be returned.
+  // Otherwise returns Status::OK().
+  virtual Status PopSavePoint() = 0;
+
   // Sets the maximum size of the write batch in bytes. 0 means no limit.
   virtual void SetMaxBytes(size_t max_bytes) = 0;
 };

utilities/write_batch_with_index/write_batch_with_index.cc

@@ -832,6 +832,10 @@ Status WriteBatchWithIndex::RollbackToSavePoint() {
   return s;
 }
 
+Status WriteBatchWithIndex::PopSavePoint() {
+  return rep->write_batch.PopSavePoint();
+}
+
 void WriteBatchWithIndex::SetMaxBytes(size_t max_bytes) {
   rep->write_batch.SetMaxBytes(max_bytes);
 }