add EnableMasterFromThis

Summary:
when T derives from EnableSharedFromThis<T> and T is placed in a MasterPtr<T> then users of T have access to MasterPtr<T> functionality.

`EnableSharedFromThis<T>::masterLockFromThis()`
`EnableSharedFromThis<T>::masterRefFromThis()`

Reviewed By: andriigrynenko

Differential Revision: D19583025

fbshipit-source-id: ee12e9de30fd844b5be36c39d7e1ade830e27bbb

folly/experimental/MasterPtr.h

@@ -17,6 +17,7 @@
 #pragma once
 
 #include <memory>
+#include <mutex>
 
 #include <folly/synchronization/Baton.h>
 
@@ -24,9 +25,76 @@
 
 namespace folly {
 
+template <typename T>
+class MasterPtr;
+
 template <typename T>
 class MasterPtrRef;
 
+/**
+ * EnableMasterFromThis provides an object with appropriate access to the
+ * functionality of the MasterPtr holding this.
+ */
+template <typename T>
+class EnableMasterFromThis {
+  template <class O>
+  static void set(
+      EnableMasterFromThis<O>* that,
+      const std::shared_ptr<std::shared_ptr<O>>& outerPtrShared) {
+    that->outerPtrWeak_ = outerPtrShared;
+    that->lockedPtrWeak_ = *outerPtrShared;
+  }
+  template <class O>
+  static auto set(O*, const std::shared_ptr<std::shared_ptr<T>>&) ->
+      typename std::enable_if<
+          !std::is_base_of<EnableMasterFromThis<T>, O>::value>::type {}
+
+ public:
+  // Gets a non-owning reference to the pointer. MasterPtr::join() does *NOT*
+  // wait for outstanding MasterPtrRef objects to be released.
+  MasterPtrRef<T> masterRefFromThis() {
+    return MasterPtrRef<T>(outerPtrWeak_);
+  }
+
+  // Gets a non-owning reference to the pointer. MasterPtr::join() does *NOT*
+  // wait for outstanding MasterPtrRef objects to be released.
+  MasterPtrRef<const T> masterRefFromThis() const {
+    return MasterPtrRef<const T>(outerPtrWeak_);
+  }
+
+  // Attempts to lock a pointer to this. Returns null if pointer is not set or
+  // if join() was called (even if the call to join() hasn't returned yet).
+  std::shared_ptr<T> masterLockFromThis() {
+    if (auto outerPtr = outerPtrWeak_.lock()) {
+      return *outerPtr;
+    }
+    return nullptr;
+  }
+  // Attempts to lock a pointer to this. Returns null if pointer is not set or
+  // if join() was called (even if the call to join() hasn't returned yet).
+  std::shared_ptr<T const> masterLockFromThis() const {
+    if (auto outerPtr = outerPtrWeak_.lock()) {
+      return *outerPtr;
+    }
+    return nullptr;
+  }
+
+  // returns the cached weak_ptr<T>
+  std::weak_ptr<T> masterWeakFromThis() noexcept {
+    return lockedPtrWeak_;
+  }
+  // returns the cached weak_ptr<T>
+  std::weak_ptr<T const> masterWeakFromThis() const noexcept {
+    return lockedPtrWeak_;
+  }
+
+ private:
+  friend class MasterPtr<T>;
+
+  std::weak_ptr<std::shared_ptr<T>> outerPtrWeak_;
+  std::weak_ptr<T> lockedPtrWeak_;
+};
+
 /**
  * MasterPtr should be used to achieve deterministic destruction of objects with
  * shared ownership.
@@ -39,19 +107,26 @@ class MasterPtrRef;
 template <typename T>
 class MasterPtr {
  public:
-  MasterPtr();
+  MasterPtr() = delete;
   MasterPtr(std::unique_ptr<T> ptr) {
     set(std::move(ptr));
   }
   ~MasterPtr() {
-    if (innerPtr_) {
+    if (*this) {
       LOG(FATAL) << "MasterPtr has to be joined explicitly.";
     }
   }
 
+  explicit operator bool() const {
+    return !!innerPtr_;
+  }
+
   // Attempts to lock a pointer. Returns null if pointer is not set or if join()
   // was called (even if the call to join() hasn't returned yet).
   std::shared_ptr<T> lock() const {
+    if (!*this) {
+      return nullptr;
+    }
     if (auto outerPtr = outerPtrWeak_.lock()) {
       return *outerPtr;
     }
@@ -62,7 +137,7 @@ class MasterPtr {
   // destroys the object in the current thread.
   // Can not be called concurrently with set().
   void join() {
-    if (!innerPtr_) {
+    if (!*this) {
       return;
     }
 
@@ -74,7 +149,7 @@ class MasterPtr {
   // Sets the pointer. Can not be called concurrently with lock() or join() or
   // ref().
   void set(std::unique_ptr<T> ptr) {
-    if (innerPtr_) {
+    if (*this) {
       LOG(FATAL) << "MasterPtr has to be joined before being set.";
     }
 
@@ -89,6 +164,7 @@ class MasterPtr {
     outerPtrShared_ =
         std::make_shared<std::shared_ptr<T>>(std::move(innerPtrShared));
     outerPtrWeak_ = outerPtrShared_;
+    EnableMasterFromThis<T>::set(innerPtr_.get(), outerPtrShared_);
   }
 
   // Gets a non-owning reference to the pointer. join() does *NOT* wait for
@@ -99,12 +175,17 @@ class MasterPtr {
 
  private:
   friend class MasterPtrRef<T>;
+
   folly::Baton<> joinBaton_;
   std::shared_ptr<std::shared_ptr<T>> outerPtrShared_;
   std::weak_ptr<std::shared_ptr<T>> outerPtrWeak_;
   std::unique_ptr<T> innerPtr_;
 };
 
+/**
+ * MasterPtrRef is a non-owning reference to the pointer. MasterPtr::join() does
+ * *NOT* wait for outstanding MasterPtrRef objects to be released.
+ */
 template <typename T>
 class MasterPtrRef {
  public:
@@ -118,6 +199,7 @@ class MasterPtrRef {
   }
 
  private:
+  friend class EnableMasterFromThis<T>;
   friend class MasterPtr<T>;
   /* implicit */ MasterPtrRef(std::weak_ptr<std::shared_ptr<T>> outerPtrWeak)
       : outerPtrWeak_(std::move(outerPtrWeak)) {}

folly/experimental/test/MasterPtrTest.cpp

@@ -64,3 +64,60 @@ TEST(MasterPtrTest, Basic) {
       joinFuture.wait_for(std::chrono::milliseconds{100}),
       std::future_status::ready);
 }
+
+struct Mastered : folly::EnableMasterFromThis<Mastered> {
+  std::shared_ptr<Mastered> get_shared() {
+    return masterLockFromThis();
+  }
+};
+
+TEST(MasterPtrTest, EnableMasterFromThis) {
+  auto ptr = std::make_unique<Mastered>();
+  auto rawPtr = ptr.get();
+
+  auto masterPtr = folly::MasterPtr<Mastered>{std::move(ptr)};
+  auto masterPtrRef = masterPtr.ref();
+
+  auto lockedPtr1 = masterPtr.lock();
+  auto lockedPtr2 = masterPtrRef.lock();
+  EXPECT_EQ(lockedPtr1.get(), rawPtr);
+  EXPECT_EQ(lockedPtr2.get(), rawPtr);
+
+  EXPECT_EQ(lockedPtr1.use_count(), 3);
+  EXPECT_EQ(lockedPtr2.use_count(), 3);
+
+  auto lockedPtr3 = lockedPtr1->get_shared();
+  EXPECT_EQ(lockedPtr3.use_count(), 4);
+  EXPECT_EQ(lockedPtr3.get(), rawPtr);
+
+  auto joinFuture = std::async(std::launch::async, [&] { masterPtr.join(); });
+
+  auto lockFailFuture = std::async(std::launch::async, [&] {
+    while (masterPtr.lock()) {
+      std::this_thread::yield();
+    }
+  });
+
+  EXPECT_EQ(
+      lockFailFuture.wait_for(std::chrono::milliseconds{100}),
+      std::future_status::ready);
+
+  EXPECT_EQ(lockedPtr1.use_count(), 3);
+  EXPECT_EQ(lockedPtr2.use_count(), 3);
+  EXPECT_EQ(lockedPtr3.use_count(), 3);
+
+  EXPECT_EQ(masterPtr.lock().get(), nullptr);
+  EXPECT_EQ(masterPtrRef.lock().get(), nullptr);
+
+  EXPECT_EQ(
+      joinFuture.wait_for(std::chrono::milliseconds{100}),
+      std::future_status::timeout);
+
+  lockedPtr1.reset();
+  lockedPtr2.reset();
+  lockedPtr3.reset();
+
+  EXPECT_EQ(
+      joinFuture.wait_for(std::chrono::milliseconds{100}),
+      std::future_status::ready);
+}