Transfer ownership from a unique_ptr to a ThreadLocalPtr

Summary:
This `ThreadLocalPtr::reset` overload will accept a `unique_ptr`. It's actually not totally exception safe, simply because `ElementWrapper::set` is not exception safe. The best I can say is that my additional code is exactly as safe as the underlying implemenation.

liketolivedangerously

Reviewed By: ericniebler

Differential Revision: D3271563

fbshipit-source-id: 774bcf31924b1ed4e29a6cb1c0a36ad710ab6034

folly/ThreadLocal.h

@@ -193,6 +193,34 @@ class ThreadLocalPtr {
     return get() != nullptr;
   }
 
+  /**
+   * reset() that transfers ownership from a smart pointer
+   */
+  template <
+      typename SourceT,
+      typename Deleter,
+      typename = typename std::enable_if<
+          std::is_convertible<SourceT*, T*>::value>::type>
+  void reset(std::unique_ptr<SourceT, Deleter> source) {
+    auto deleter = [delegate = source.get_deleter()](
+        T * ptr, TLPDestructionMode) {
+      delegate(ptr);
+    };
+    reset(source.release(), deleter);
+  }
+
+  /**
+   * reset() that transfers ownership from a smart pointer with the default
+   * deleter
+   */
+  template <
+      typename SourceT,
+      typename = typename std::enable_if<
+          std::is_convertible<SourceT*, T*>::value>::type>
+  void reset(std::unique_ptr<SourceT> source) {
+    reset(source.release());
+  }
+
   /**
    * reset() with a custom deleter:
    * deleter(T* ptr, TLPDestructionMode mode)

folly/test/ThreadLocalTest.cpp

@@ -36,6 +36,7 @@
 #include <gtest/gtest.h>
 
 #include <folly/Baton.h>
+#include <folly/Memory.h>
 #include <folly/experimental/io/FsUtil.h>
 
 using namespace folly;
@@ -48,7 +49,7 @@ struct Widget {
   }
 
   static void customDeleter(Widget* w, TLPDestructionMode mode) {
-    totalVal_ += (mode == TLPDestructionMode::ALL_THREADS) * 1000;
+    totalVal_ += (mode == TLPDestructionMode::ALL_THREADS) ? 1000 : 1;
     delete w;
   }
 };
@@ -72,6 +73,37 @@ TEST(ThreadLocalPtr, CustomDeleter1) {
         w.reset(new Widget(), Widget::customDeleter);
         w.get()->val_ += 10;
       }).join();
+    EXPECT_EQ(11, Widget::totalVal_);
+  }
+  EXPECT_EQ(11, Widget::totalVal_);
+}
+
+TEST(ThreadLocalPtr, CustomDeleterOwnershipTransfer) {
+  Widget::totalVal_ = 0;
+  {
+    ThreadLocalPtr<Widget> w;
+    auto deleter = [](Widget* ptr) {
+      Widget::customDeleter(ptr, TLPDestructionMode::THIS_THREAD);
+    };
+    std::unique_ptr<Widget, typeof(deleter)> source(new Widget(), deleter);
+    std::thread([&w, &source]() {
+      w.reset(std::move(source));
+      w.get()->val_ += 10;
+    }).join();
+    EXPECT_EQ(11, Widget::totalVal_);
+  }
+  EXPECT_EQ(11, Widget::totalVal_);
+}
+
+TEST(ThreadLocalPtr, DefaultDeleterOwnershipTransfer) {
+  Widget::totalVal_ = 0;
+  {
+    ThreadLocalPtr<Widget> w;
+    auto source = folly::make_unique<Widget>();
+    std::thread([&w, &source]() {
+      w.reset(std::move(source));
+      w.get()->val_ += 10;
+    }).join();
     EXPECT_EQ(10, Widget::totalVal_);
   }
   EXPECT_EQ(10, Widget::totalVal_);