From 3e7fc88167302ca4ae8734616bccd3f31e4fe538 Mon Sep 17 00:00:00 2001
From: Andrew Kryczka <andrew.kryczka2@gmail.com>
Date: Fri, 26 May 2023 18:15:14 -0700
Subject: [PATCH] add WriteBatch::Release() (#11482)

Summary:
Together with the existing constructor,
`explicit WriteBatch(std::string&& rep)`, this enables transferring `WriteBatch` via its `std::string` representation. Associated info like KV checksums are dropped but the caller can use `WriteBatch::VerifyChecksum()` before taking ownership if needed.

Pull Request resolved: https://github.com/facebook/rocksdb/pull/11482

Reviewed By: cbi42

Differential Revision: D46233884

Pulled By: ajkr

fbshipit-source-id: 6bc64a6e75fb7bbf61d08c09520fc3705a7b44d8
---
 HISTORY.md                    |  1 +
 db/write_batch.cc             |  6 ++++++
 db/write_batch_test.cc        | 16 ++++++++++++++++
 include/rocksdb/write_batch.h |  3 +++
 4 files changed, 26 insertions(+)

diff --git a/HISTORY.md b/HISTORY.md
index 11a7c45ac..f0742e1cd 100644
--- a/HISTORY.md
+++ b/HISTORY.md
@@ -2,6 +2,7 @@
 ## Unreleased
 ### New Features
 * Add a new option OptimisticTransactionDBOptions::shared_lock_buckets that enables sharing mutexes for validating transactions between DB instances, for better balancing memory efficiency and validation contention across DB instances. Different column families and DBs also now use different hash seeds in this validation, so that the same set of key names will not contend across DBs or column families.
+* Add `WriteBatch::Release()` that releases the batch's serialized data to the caller.
 
 ### Public API Changes
 * Add `WaitForCompact()` to wait for all flush and compactions jobs to finish. Jobs to wait include the unscheduled (queued, but not scheduled yet).
diff --git a/db/write_batch.cc b/db/write_batch.cc
index 726b14f09..47e34e6f9 100644
--- a/db/write_batch.cc
+++ b/db/write_batch.cc
@@ -293,6 +293,12 @@ size_t WriteBatch::GetProtectionBytesPerKey() const {
   return 0;
 }
 
+std::string WriteBatch::Release() {
+  std::string ret = std::move(rep_);
+  Clear();
+  return ret;
+}
+
 bool WriteBatch::HasPut() const {
   return (ComputeContentFlags() & ContentFlags::HAS_PUT) != 0;
 }
diff --git a/db/write_batch_test.cc b/db/write_batch_test.cc
index 4bd74f71e..174052644 100644
--- a/db/write_batch_test.cc
+++ b/db/write_batch_test.cc
@@ -247,6 +247,22 @@ TEST_F(WriteBatchTest, SingleDeletion) {
   ASSERT_EQ(2u, batch.Count());
 }
 
+TEST_F(WriteBatchTest, OwnershipTransfer) {
+  Random rnd(301);
+  WriteBatch put_batch;
+  ASSERT_OK(put_batch.Put(rnd.RandomString(16) /* key */,
+                          rnd.RandomString(1024) /* value */));
+
+  // (1) Verify `Release()` transfers string data ownership
+  const char* expected_data = put_batch.Data().data();
+  std::string batch_str = put_batch.Release();
+  ASSERT_EQ(expected_data, batch_str.data());
+
+  // (2) Verify constructor transfers string data ownership
+  WriteBatch move_batch(std::move(batch_str));
+  ASSERT_EQ(expected_data, move_batch.Data().data());
+}
+
 namespace {
 struct TestHandler : public WriteBatch::Handler {
   std::string seen;
diff --git a/include/rocksdb/write_batch.h b/include/rocksdb/write_batch.h
index 61ba5a739..307cd7559 100644
--- a/include/rocksdb/write_batch.h
+++ b/include/rocksdb/write_batch.h
@@ -356,6 +356,9 @@ class WriteBatch : public WriteBatchBase {
   // Retrieve the serialized version of this batch.
   const std::string& Data() const { return rep_; }
 
+  // Release the serialized data and clear this batch.
+  std::string Release();
+
   // Retrieve data size of the batch.
   size_t GetDataSize() const { return rep_.size(); }