PinnableSlice move assignment

Summary:
Allow `std::move(pinnable_slice)`.
Closes https://github.com/facebook/rocksdb/pull/2997

Differential Revision: D6036782

Pulled By: yiwu-arbug

fbshipit-source-id: 583fb0419a97e437ff530f4305822341cd3381fa

CMakeLists.txt

@@ -779,6 +779,7 @@ if(WITH_TESTS)
         options/options_test.cc
         table/block_based_filter_block_test.cc
         table/block_test.cc
+        table/cleanable_test.cc
         table/cuckoo_table_builder_test.cc
         table/cuckoo_table_reader_test.cc
         table/full_filter_block_test.cc
@@ -801,6 +802,7 @@ if(WITH_TESTS)
         util/hash_test.cc
         util/heap_test.cc
         util/rate_limiter_test.cc
+        util/slice_test.cc
         util/slice_transform_test.cc
         util/timer_queue_test.cc
         util/thread_list_test.cc

Makefile

@@ -494,6 +494,7 @@ TESTS = \
 	repair_test \
 	env_timed_test \
 	write_prepared_transaction_test \
+	slice_test \
 
 PARALLEL_TEST = \
 	backupable_db_test \
@@ -1477,6 +1478,9 @@ range_del_aggregator_test: db/range_del_aggregator_test.o db/db_test_util.o $(LI
 blob_db_test: utilities/blob_db/blob_db_test.o $(LIBOBJECTS) $(TESTHARNESS)
 	$(AM_LINK)
 
+slice_test: util/slice_test.o $(LIBOBJECTS) $(TESTHARNESS)
+	$(AM_LINK)
+
 #-------------------------------------------------
 # make install related stuff
 INSTALL_PATH ?= /usr/local

TARGETS

@@ -463,6 +463,7 @@ ROCKS_TESTS = [['arena_test', 'util/arena_test.cc', 'serial'],
  ['repair_test', 'db/repair_test.cc', 'serial'],
  ['sim_cache_test', 'utilities/simulator_cache/sim_cache_test.cc', 'serial'],
  ['skiplist_test', 'memtable/skiplist_test.cc', 'serial'],
+ ['slice_test', 'util/slice_test.cc', 'serial'],
  ['slice_transform_test', 'util/slice_transform_test.cc', 'serial'],
  ['spatial_db_test', 'utilities/spatialdb/spatial_db_test.cc', 'serial'],
  ['sst_dump_test', 'tools/sst_dump_test.cc', 'serial'],

include/rocksdb/cleanable.h

@@ -25,6 +25,15 @@ class Cleanable {
  public:
   Cleanable();
   ~Cleanable();
+
+  // No copy constructor and copy assignment allowed.
+  Cleanable(Cleanable&) = delete;
+  Cleanable& operator=(Cleanable&) = delete;
+
+  // Move consturctor and move assignment is allowed.
+  Cleanable(Cleanable&&);
+  Cleanable& operator=(Cleanable&&);
+
   // Clients are allowed to register function/arg1/arg2 triples that
   // will be invoked when this iterator is destroyed.
   //

include/rocksdb/slice.h

@@ -129,6 +129,31 @@ class PinnableSlice : public Slice, public Cleanable {
   PinnableSlice() { buf_ = &self_space_; }
   explicit PinnableSlice(std::string* buf) { buf_ = buf; }
 
+  // No copy constructor and copy assignment allowed.
+  PinnableSlice(PinnableSlice&) = delete;
+  PinnableSlice& operator=(PinnableSlice&) = delete;
+
+  PinnableSlice(PinnableSlice&& other) { *this = std::move(other); }
+
+  PinnableSlice& operator=(PinnableSlice&& other) {
+    if (this != &other) {
+      // cleanup itself.
+      Reset();
+
+      Slice::operator=(other);
+      Cleanable::operator=(std::move(other));
+      pinned_ = other.pinned_;
+      if (!pinned_ && other.buf_ == &other.self_space_) {
+        self_space_ = std::move(other.self_space_);
+        buf_ = &self_space_;
+        data_ = buf_->data();
+      } else {
+        buf_ = other.buf_;
+      }
+    }
+    return *this;
+  }
+
   inline void PinSlice(const Slice& s, CleanupFunction f, void* arg1,
                        void* arg2) {
     assert(!pinned_);

src.mk

@@ -312,6 +312,7 @@ MAIN_SOURCES =                                                    \
   options/options_test.cc                                               \
   table/block_based_filter_block_test.cc                                \
   table/block_test.cc                                                   \
+  table/cleanable_test.cc                                               \
   table/cuckoo_table_builder_test.cc                                    \
   table/cuckoo_table_reader_test.cc                                     \
   table/full_filter_block_test.cc                                       \
@@ -336,6 +337,7 @@ MAIN_SOURCES =                                                    \
   util/filelock_test.cc                                                 \
   util/log_write_bench.cc                                               \
   util/rate_limiter_test.cc                                             \
+  util/slice_test.cc                                                    \
   util/slice_transform_test.cc                                          \
   util/timer_queue_test.cc                                              \
   util/thread_list_test.cc                                              \

table/iterator.cc

@@ -21,6 +21,19 @@ Cleanable::Cleanable() {
 
 Cleanable::~Cleanable() { DoCleanup(); }
 
+Cleanable::Cleanable(Cleanable&& other) {
+  *this = std::move(other);
+}
+
+Cleanable& Cleanable::operator=(Cleanable&& other) {
+  if (this != &other) {
+    cleanup_ = other.cleanup_;
+    other.cleanup_.function = nullptr;
+    other.cleanup_.next = nullptr;
+  }
+  return *this;
+}
+
 // If the entire linked list was on heap we could have simply add attach one
 // link list to another. However the head is an embeded object to avoid the cost
 // of creating objects for most of the use cases when the Cleanable has only one

util/slice_test.cc

@@ -0,0 +1,70 @@
+//  Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
+//  This source code is licensed under both the GPLv2 (found in the
+//  COPYING file in the root directory) and Apache 2.0 License
+//  (found in the LICENSE.Apache file in the root directory).
+
+#include "port/stack_trace.h"
+#include "rocksdb/slice.h"
+#include "util/testharness.h"
+
+namespace rocksdb {
+
+class SliceTest : public testing::Test {};
+
+namespace {
+void BumpCounter(void* arg1, void* arg2) {
+  (*reinterpret_cast<int*>(arg1))++;
+}
+}  // anonymous namespace
+
+TEST_F(SliceTest, PinnableSliceMoveConstruct) {
+  for (int i = 0; i < 3; i++) {
+    int orig_cleanup = 0;
+    int moved_cleanup = 0;
+    PinnableSlice* s1 = nullptr;
+    std::string external_storage;
+    switch (i) {
+      case 0:
+        s1 = new PinnableSlice();
+        *(s1->GetSelf()) = "foo";
+        s1->PinSelf();
+        s1->RegisterCleanup(BumpCounter, &moved_cleanup, nullptr);
+        break;
+      case 1:
+        s1 = new PinnableSlice(&external_storage);
+        *(s1->GetSelf()) = "foo";
+        s1->PinSelf();
+        s1->RegisterCleanup(BumpCounter, &moved_cleanup, nullptr);
+        break;
+      case 2:
+        s1 = new PinnableSlice();
+        s1->PinSlice("foo", BumpCounter, &moved_cleanup, nullptr);
+        break;
+    }
+    ASSERT_EQ("foo", s1->ToString());
+    PinnableSlice* s2 = new PinnableSlice();
+    s2->PinSelf("bar");
+    ASSERT_EQ("bar", s2->ToString());
+    s2->RegisterCleanup(BumpCounter, &orig_cleanup, nullptr);
+    *s2 = std::move(*s1);
+    ASSERT_EQ("foo", s2->ToString());
+    ASSERT_EQ(1, orig_cleanup);
+    ASSERT_EQ(0, moved_cleanup);
+    delete s1;
+    // ASAN will check if it will access storage of s1, which is deleted.
+    ASSERT_EQ("foo", s2->ToString());
+    ASSERT_EQ(1, orig_cleanup);
+    ASSERT_EQ(0, moved_cleanup);
+    delete s2;
+    ASSERT_EQ(1, orig_cleanup);
+    ASSERT_EQ(1, moved_cleanup);
+  }
+}
+
+}  // namespace rocksdb
+
+int main(int argc, char** argv) {
+  rocksdb::port::InstallStackTraceHandler();
+  ::testing::InitGoogleTest(&argc, argv);
+  return RUN_ALL_TESTS();
+}