Wait for borrows to be relinquished before destroying the borrowed ob…

.bazeliskrc

@@ -1 +1 @@
-USE_BAZEL_VERSION=6.5.0
+USE_BAZEL_VERSION=6.6.0

include/stout/borrowed_ptr.h

@@ -92,6 +92,24 @@ class TypeErasedBorrowable {
     }
   }
 
+  // Transitions to 'Destructing', which prevents any new borrows from
+  // being taken, and then waits for all outstanding borrows to be
+  // relinquished.
+  //
+  // Every type deriving from 'TypeErasedBorrowable' must call this in
+  // its destructor before its members are destroyed: by the time
+  // '~TypeErasedBorrowable()' runs those members have already been
+  // destroyed, which is too late to wait for borrowers that might
+  // still be using them ('~TypeErasedBorrowable()' checks this
+  // contract and fails fast if it has not been followed).
+  void DestructingAndWait() {
+    auto state = State::Borrowing;
+    if (!tally_.Update(state, State::Destructing)) {
+      LOG(FATAL) << "Unable to transition to Destructing from state " << state;
+    }
+    WaitUntilBorrowsEquals(0);
+  }
+
  protected:
   TypeErasedBorrowable()
     : tally_(State::Borrowing) {}
@@ -107,16 +125,18 @@ class TypeErasedBorrowable {
   }
 
   virtual ~TypeErasedBorrowable() {
-    auto state = State::Borrowing;
-    if (!tally_.Update(state, State::Destructing)) {
-      LOG(FATAL) << "Unable to transition to Destructing from state " << state;
-    } else {
-      // NOTE: it's possible that we'll block forever if exceptions
-      // were thrown and destruction was not successful.
-      // if (!std::uncaught_exceptions() > 0) {
-      WaitUntilBorrowsEquals(0);
-      // }
-    }
+    // Check the contract that the destructor of the deriving type has
+    // already called 'DestructingAndWait()', i.e., that we have
+    // transitioned to 'Destructing' and that all borrows have been
+    // relinquished: the members of the deriving type have already
+    // been destroyed by the time we get here, which is too late to
+    // wait for borrowers that might still be using them.
+    auto [state, count] = tally_.Load();
+    CHECK(state == State::Destructing && count == 0)
+        << "The destructor of a type deriving from "
+           "'TypeErasedBorrowable' must call 'DestructingAndWait()' "
+           "(state: "
+        << state << ", outstanding borrows: " << count << ")";
   }
 
   enum class State : uint8_t {
@@ -193,6 +213,13 @@ class Borrowable : public TypeErasedBorrowable {
     : TypeErasedBorrowable(std::move(that)),
       t_(std::move(that.t_)) {}
 
+  ~Borrowable() {
+    // Wait for all borrows to be relinquished before 't_', the object
+    // that they refer to, gets destroyed; we need to wait here because
+    // this is where the actual data gets destroyed, which is too late.
+    DestructingAndWait();
+  }
+
   borrowed_ref<T> Borrow() {
     auto state = State::Borrowing;
     if (tally_.Increment(state)) {
@@ -255,6 +282,14 @@ class Borrowable<std::unique_ptr<T>> : public TypeErasedBorrowable {
     : TypeErasedBorrowable(std::move(that)),
       t_(std::move(that.t_)) {}
 
+  ~Borrowable() {
+    // Wait for all borrows to be relinquished before 't_', and with it
+    // the object that they refer to, gets destroyed; we need to wait here
+    // because this is where the actual data gets destroyed,
+    // which is too late.
+    DestructingAndWait();
+  }
+
   borrowed_ref<T> Borrow() {
     auto state = State::Borrowing;
     if (tally_.Increment(state)) {

include/stout/stateful-tally.h

@@ -30,6 +30,17 @@ struct StatefulTally {
     return S(value.load() >> ((sizeof(size_t) - 1) * 8));
   }
 
+  // Returns the current (state, count) pair decoded from a single
+  // atomic load, i.e., a consistent snapshot of both.
+  std::pair<S, size_t> Load() const {
+    size_t loaded = value.load();
+
+    size_t count = (loaded << 8) >> 8;
+    size_t state = loaded >> ((sizeof(size_t) - 1) * 8);
+
+    return std::make_pair(S(state), count);
+  }
+
   template <typename Predicate>
   std::pair<S, size_t> Wait(Predicate&& predicate) {
     for (AtomicBackoff b;; b.pause()) {

tests/borrowed_ptr.cc

@@ -1,6 +1,7 @@
 #include "stout/borrowed_ptr.h"
 
 #include <atomic>
+#include <memory>
 #include <string>
 #include <thread>
 #include <vector>
@@ -422,6 +423,9 @@ TEST(BorrowTest, EnableBorrowableFromThis) {
    public:
     Foo(int i)
       : i(i) {}
+    ~Foo() {
+      DestructingAndWait();
+    }
     int i = 0;
   };
 
@@ -442,6 +446,10 @@ TEST(BorrowTest, EnableBorrowableFromThisMove) {
    public:
     Foo(int i)
       : i(i) {}
+    Foo(Foo&& that) = default;
+    ~Foo() {
+      DestructingAndWait();
+    }
     int i = 0;
   };
 
@@ -464,6 +472,9 @@ TEST(BorrowTest, EnableBorrowableFromThisCopy) {
    public:
     Foo(int i)
       : i(i) {}
+    ~Foo() {
+      DestructingAndWait();
+    }
     int i = 0;
   };
 
@@ -479,3 +490,112 @@ TEST(BorrowTest, EnableBorrowableFromThisCopy) {
 
   EXPECT_EQ(borrowed->i, 42);
 }
+
+
+TEST(BorrowTest, DestructWaitsForBorrows) {
+  auto* s = new Borrowable<string>("hello world");
+
+  borrowed_ref<string> borrowed = s->Borrow();
+
+  atomic<bool> destructed(false);
+
+  thread t([&]() {
+    // Waits for 'borrowed' to be relinquished.
+    delete s;
+    destructed.store(true);
+  });
+
+  // The destructor must wait for our borrow to be relinquished
+  // before destroying the borrowed object, so this read is safe no
+  // matter how far 'delete' has gotten on the other thread.
+  EXPECT_EQ("hello world", *borrowed);
+
+  EXPECT_FALSE(destructed.load());
+
+  {
+    borrowed_ref<string> relinquished = std::move(borrowed);
+  }
+
+  t.join();
+
+  EXPECT_TRUE(destructed.load());
+}
+
+
+TEST(BorrowTest, DestructWaitsForBorrowsUniquePtr) {
+  auto* s = new Borrowable<unique_ptr<string>>(
+      std::make_unique<string>("hello world"));
+
+  borrowed_ref<string> borrowed = s->Borrow();
+
+  atomic<bool> destructed(false);
+
+  thread t([&]() {
+    // Waits for 'borrowed' to be relinquished.
+    delete s;
+    destructed.store(true);
+  });
+
+  // The destructor must wait for our borrow to be relinquished
+  // before destroying the borrowed object, so this read is safe no
+  // matter how far 'delete' has gotten on the other thread.
+  EXPECT_EQ("hello world", *borrowed);
+
+  EXPECT_FALSE(destructed.load());
+
+  {
+    borrowed_ref<string> relinquished = std::move(borrowed);
+  }
+
+  t.join();
+
+  EXPECT_TRUE(destructed.load());
+}
+
+
+TEST(BorrowTest, EnableBorrowableFromThisDestructWaitsForBorrows) {
+  class Foo : public enable_borrowable_from_this<Foo> {
+   public:
+    Foo(int i)
+      : i(i) {}
+    ~Foo() {
+      DestructingAndWait();
+    }
+    int i = 0;
+  };
+
+  auto* foo = new Foo(42);
+
+  borrowed_ptr<Foo> borrowed = foo->Borrow();
+
+  atomic<bool> destructed(false);
+
+  thread t([&]() {
+    // 'DestructingAndWait()' waits for 'borrowed' to be relinquished.
+    delete foo;
+    destructed.store(true);
+  });
+
+  // 'Foo's destructor calls 'DestructingAndWait()' before its members
+  // are destroyed, so this read is safe no matter how far 'delete'
+  // has gotten on the other thread.
+  EXPECT_EQ(borrowed->i, 42);
+
+  EXPECT_FALSE(destructed.load());
+
+  borrowed.relinquish();
+
+  t.join();
+
+  EXPECT_TRUE(destructed.load());
+}
+
+
+TEST(BorrowDeathTest, EnableBorrowableFromThisDestructingAndWaitContract) {
+  class Foo : public enable_borrowable_from_this<Foo> {};
+
+  // Destroying a type deriving from 'enable_borrowable_from_this'
+  // whose destructor did not call 'DestructingAndWait()' violates the
+  // contract checked by '~TypeErasedBorrowable()'.
+  EXPECT_DEATH({ Foo foo; }, "must call 'DestructingAndWait");
+}