smart intrusive reference counting

6 files changed, 393 insertions, 0 deletions

diff --git a/vespalib/CMakeLists.txt b/vespalib/CMakeLists.txt
index 8175e75875a..b689d587bee 100644
--- a/vespalib/CMakeLists.txt
+++ b/vespalib/CMakeLists.txt
@@ -140,6 +140,7 @@ vespa_define_module(
     src/tests/process
     src/tests/programoptions
     src/tests/random
+    src/tests/ref_counted
     src/tests/referencecounter
     src/tests/regex
     src/tests/rendezvous
diff --git a/vespalib/src/tests/ref_counted/CMakeLists.txt b/vespalib/src/tests/ref_counted/CMakeLists.txt
new file mode 100644
index 00000000000..74258dc5e67
--- /dev/null
+++ b/vespalib/src/tests/ref_counted/CMakeLists.txt
@@ -0,0 +1,9 @@
+# Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+vespa_add_executable(vespalib_ref_counted_test_app TEST
+    SOURCES
+    ref_counted_test.cpp
+    DEPENDS
+    vespalib
+    GTest::GTest
+)
+vespa_add_test(NAME vespalib_ref_counted_test_app COMMAND vespalib_ref_counted_test_app)
diff --git a/vespalib/src/tests/ref_counted/ref_counted_test.cpp b/vespalib/src/tests/ref_counted/ref_counted_test.cpp
new file mode 100644
index 00000000000..751d5b7c506
--- /dev/null
+++ b/vespalib/src/tests/ref_counted/ref_counted_test.cpp
@@ -0,0 +1,218 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include <vespa/vespalib/util/ref_counted.h>
+#include <vespa/vespalib/util/gate.h>
+#include <vespa/vespalib/util/thread.h>
+#include <vespa/vespalib/gtest/gtest.h>
+
+using namespace vespalib;
+
+struct Base : enable_ref_counted {
+    static std::atomic<int> ctor_cnt;
+    static std::atomic<int> dtor_cnt;
+    int val;
+    Base(int val_in) : val(val_in) {
+        ctor_cnt.fetch_add(1, std::memory_order_relaxed);
+    }
+    ~Base() {
+        dtor_cnt.fetch_add(1, std::memory_order_relaxed);
+    }
+};
+std::atomic<int> Base::ctor_cnt = 0;
+std::atomic<int> Base::dtor_cnt = 0;
+
+struct Leaf : Base {
+    static std::atomic<int> ctor_cnt;
+    static std::atomic<int> dtor_cnt;
+    Leaf(int val_in) : Base(val_in) {
+        ctor_cnt.fetch_add(1, std::memory_order_relaxed);
+    }
+    ~Leaf() {
+        dtor_cnt.fetch_add(1, std::memory_order_relaxed);
+    }
+};
+std::atomic<int> Leaf::ctor_cnt = 0;
+std::atomic<int> Leaf::dtor_cnt = 0;
+
+// check that the expected number of objects have been created and
+// destroyed while this object was in scope.
+struct CheckObjects {
+    int expect_base;
+    int expect_leaf;
+    int old_base_ctor;
+    int old_base_dtor;
+    int old_leaf_ctor;
+    int old_leaf_dtor;
+    CheckObjects(int expect_base_in = 0, int expect_leaf_in = 0)
+      : expect_base(expect_base_in),
+        expect_leaf(expect_leaf_in)
+    {
+        old_base_ctor = Base::ctor_cnt.load(std::memory_order_relaxed);
+        old_base_dtor = Base::dtor_cnt.load(std::memory_order_relaxed);
+        old_leaf_ctor = Leaf::ctor_cnt.load(std::memory_order_relaxed);
+        old_leaf_dtor = Leaf::dtor_cnt.load(std::memory_order_relaxed);
+    }
+    ~CheckObjects() {
+        int base_ctor_diff = Base::ctor_cnt.load(std::memory_order_relaxed) - old_base_ctor;
+        int base_dtor_diff = Base::dtor_cnt.load(std::memory_order_relaxed) - old_base_dtor;
+        int leaf_ctor_diff = Leaf::ctor_cnt.load(std::memory_order_relaxed) - old_leaf_ctor;
+        int leaf_dtor_diff = Leaf::dtor_cnt.load(std::memory_order_relaxed) - old_leaf_dtor;
+        EXPECT_EQ(base_ctor_diff, expect_base);
+        EXPECT_EQ(base_dtor_diff, expect_base);
+        EXPECT_EQ(leaf_ctor_diff, expect_leaf);
+        EXPECT_EQ(leaf_dtor_diff, expect_leaf);
+    }
+};
+
+TEST(RefCountedTest, create_empty_ref_counted) {
+    CheckObjects check;
+    ref_counted<Base> empty;
+    EXPECT_FALSE(empty);
+}
+
+TEST(RefCountedTest, make_ref_counted) {
+    CheckObjects check(2, 1);
+    auto ref1 = make_ref_counted<Base>(10);
+    static_assert(std::same_as<decltype(ref1),ref_counted<Base>>);
+    EXPECT_TRUE(ref1);
+    EXPECT_EQ((*ref1).val, 10);
+    EXPECT_EQ(ref1->val, 10);
+    auto ref2 = make_ref_counted<Leaf>(20);
+    static_assert(std::same_as<decltype(ref2),ref_counted<Leaf>>);
+    EXPECT_TRUE(ref2);
+    EXPECT_EQ((*ref2).val, 20);
+    EXPECT_EQ(ref2->val, 20);
+}
+
+TEST(RefCountedTest, ref_counted_from) {
+    CheckObjects check(1, 1);
+    auto ref = make_ref_counted<Leaf>(10);
+    Leaf &leaf = *ref;
+    Base &base = leaf;
+    EXPECT_EQ(ref->count_refs(), 1);
+    auto from_leaf = ref_counted_from(leaf);
+    static_assert(std::same_as<decltype(from_leaf),ref_counted<Leaf>>);
+    auto from_base = ref_counted_from(base);
+    static_assert(std::same_as<decltype(from_base),ref_counted<Base>>);
+    EXPECT_EQ(ref->count_refs(), 3);
+    EXPECT_EQ(from_base->val, 10);
+}
+
+TEST(RefCountedTest, use_internal_api) {
+    CheckObjects check(1);
+    Base *raw = new Base(20);
+    EXPECT_EQ(raw->count_refs(), 1);
+    ref_counted<Base> ref = ref_counted<Base>::internal_attach(raw);
+    EXPECT_EQ(ref->count_refs(), 1);
+    EXPECT_EQ(ref.internal_detach(), raw);
+    EXPECT_EQ(raw->count_refs(), 1);
+    raw->internal_addref();
+    EXPECT_EQ(raw->count_refs(), 2);
+    raw->internal_subref();
+    EXPECT_EQ(raw->count_refs(), 1);
+    raw->internal_subref();
+}
+
+TEST(RefCountedTest, move_ref_counted) {
+    for (bool has_src: {true, false}) {
+        for (bool has_dst: {true, false}) {
+            for (bool same: {true, false}) {
+                if (same) {
+                    CheckObjects check(has_src + has_dst);
+                    ref_counted<Base> src = has_src ? make_ref_counted<Base>(10) : ref_counted<Base>();
+                    ref_counted<Base> dst = has_dst ? make_ref_counted<Base>(20) : ref_counted<Base>();
+                    dst = std::move(src);
+                    EXPECT_EQ(dst, has_src);
+                    EXPECT_FALSE(src);
+                    if (has_src) {
+                        EXPECT_EQ(dst->val, 10);
+                        EXPECT_EQ(dst->count_refs(), 1);
+                    }
+                } else {
+                    CheckObjects check(has_src + has_dst, has_src + has_dst);
+                    ref_counted<Leaf> src = has_src ? make_ref_counted<Leaf>(10) : ref_counted<Leaf>();
+                    ref_counted<Base> dst = has_dst ? make_ref_counted<Leaf>(20) : ref_counted<Leaf>();
+                    dst = std::move(src);
+                    EXPECT_EQ(dst, has_src);
+                    EXPECT_FALSE(src);
+                    if (has_src) {
+                        EXPECT_EQ(dst->val, 10);
+                        EXPECT_EQ(dst->count_refs(), 1);
+                    }
+                }
+            }
+        }
+    }
+}
+
+TEST(RefCountedTest, copy_ref_counted) {
+    for (bool has_src: {true, false}) {
+        for (bool has_dst: {true, false}) {
+            for (bool same: {true, false}) {
+                if (same) {
+                    CheckObjects check(2);
+                    ref_counted<Base> empty;
+                    ref_counted<Base> obj1 = make_ref_counted<Base>(10);
+                    ref_counted<Base> obj2 = make_ref_counted<Base>(20);
+                    ref_counted<Base> src = has_src ? obj1 : empty;
+                    ref_counted<Base> dst = has_dst ? obj2 : empty;
+                    dst = src;
+                    EXPECT_EQ(dst, has_src);
+                    EXPECT_EQ(src, has_src);
+                    if (has_src) {
+                        EXPECT_EQ(dst->val, 10);
+                        EXPECT_EQ(dst->count_refs(), 3);
+                    }
+                } else {
+                    CheckObjects check(2, 2);
+                    ref_counted<Leaf> empty;
+                    ref_counted<Leaf> obj1 = make_ref_counted<Leaf>(10);
+                    ref_counted<Leaf> obj2 = make_ref_counted<Leaf>(20);
+                    ref_counted<Leaf> src = has_src ? obj1 : empty;
+                    ref_counted<Base> dst = has_dst ? obj2 : empty;
+                    dst = src;
+                    EXPECT_EQ(dst, has_src);
+                    EXPECT_EQ(src, has_src);
+                    if (has_src) {
+                        EXPECT_EQ(dst->val, 10);
+                        EXPECT_EQ(dst->count_refs(), 3);
+                    }
+                }
+            }
+        }
+    }
+}
+
+TEST(RefCountedTest, compile_errors_when_uncommented) {
+    struct Foo {};
+    [[maybe_unused]] Foo foo;
+    // ref_counted<Foo> empty;
+    // auto ref1 = make_ref_counted<Foo>();
+    // auto ref2 = ref_counted_from(foo);
+}
+
+TEST(RefCountedTest, with_threads) {
+    CheckObjects check(2,1);
+    ThreadPool pool;
+    Gate gate;
+    {
+        auto a = make_ref_counted<Base>(10);
+        auto b = make_ref_counted<Leaf>(20);
+        for (int i = 0; i < 8; ++i) {
+            pool.start([&gate,a,b]()
+                       {
+                           gate.await();
+                           for (int j = 0; j < 100000; ++j) {
+                               auto c = a;
+                               auto d = b;
+                               EXPECT_EQ(c->val, 10);
+                               EXPECT_EQ(d->val, 20);
+                           }
+                       });
+        }
+    }
+    gate.countDown();
+    pool.join();
+}
+
+GTEST_MAIN_RUN_ALL_TESTS()
diff --git a/vespalib/src/vespa/vespalib/util/CMakeLists.txt b/vespalib/src/vespa/vespalib/util/CMakeLists.txt
index c8536fc68c1..663b3b65638 100644
--- a/vespalib/src/vespa/vespalib/util/CMakeLists.txt
+++ b/vespalib/src/vespa/vespalib/util/CMakeLists.txt
@@ -63,6 +63,7 @@ vespa_add_library(vespalib_vespalib_util OBJECT
     programoptions.cpp
     random.cpp
     rcuvector.cpp
+    ref_counted.cpp
     regexp.cpp
     require.cpp
     resource_limits.cpp
diff --git a/vespalib/src/vespa/vespalib/util/ref_counted.cpp b/vespalib/src/vespa/vespalib/util/ref_counted.cpp
new file mode 100644
index 00000000000..a6928082c58
--- /dev/null
+++ b/vespalib/src/vespa/vespalib/util/ref_counted.cpp
@@ -0,0 +1,47 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include "ref_counted.h"
+#include <cassert>
+
+namespace vespalib {
+
+void
+enable_ref_counted::internal_addref() const noexcept
+{
+    // relaxed because:
+    // the thread obtaining the new reference already has a reference
+    auto prev = _refs.fetch_add(1, std::memory_order_relaxed);
+    assert(prev > 0);
+}
+
+void
+enable_ref_counted::internal_subref() const noexcept
+{
+    // release because:
+    // our changes to the object must be visible to the deleter
+    auto prev = _refs.fetch_sub(1, std::memory_order_release);
+    assert(prev > 0);
+    if (prev == 1) {
+        // acquire because:
+        // we need to see all object changes before deleting it
+        std::atomic_thread_fence(std::memory_order_acquire);
+        delete this;
+    }
+}
+
+int32_t
+enable_ref_counted::count_refs() const noexcept {
+    auto result = _refs.load(std::memory_order_relaxed);
+    assert(result > 0);
+    return result;
+}
+
+enable_ref_counted::~enable_ref_counted() noexcept
+{
+    // protect against early/double delete and memory overwrites
+    assert(_refs.load(std::memory_order_relaxed) == 0);
+    assert(_guard == MAGIC);
+    _guard = 0;
+}
+
+}
diff --git a/vespalib/src/vespa/vespalib/util/ref_counted.h b/vespalib/src/vespa/vespalib/util/ref_counted.h
new file mode 100644
index 00000000000..6465a564891
--- /dev/null
+++ b/vespalib/src/vespa/vespalib/util/ref_counted.h
@@ -0,0 +1,117 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#pragma once
+
+#include <atomic>
+#include <concepts>
+#include <utility>
+
+// This file contains code that implements intrusive reference
+// counting with smart handles (ref_counted<T> points to T, T inherits
+// enable_ref_counted).
+
+// Functions with names starting with 'internal' are not intended for
+// direct use, but are available for completeness (enables incremental
+// re-write of code from bald to smart pointers).
+
+namespace vespalib {
+
+// This is the actual reference count. It cannot be moved or copied,
+// but it can be changed even when const. Classes that want to be
+// counted need to inherit from this. The destructor is virtual to
+// make sure the right one is called. This also enables the subref
+// function to directly delete the shared object.
+class enable_ref_counted
+{
+    static constexpr uint32_t MAGIC = 0xcc56a933;
+private:
+    uint32_t _guard;
+    mutable std::atomic<int32_t> _refs;
+protected:
+    enable_ref_counted() noexcept : _guard(MAGIC), _refs(1) {}
+public:
+    virtual ~enable_ref_counted() noexcept;
+    enable_ref_counted(enable_ref_counted &&) = delete;
+    enable_ref_counted(const enable_ref_counted &) = delete;
+    enable_ref_counted &operator=(enable_ref_counted &&) = delete;
+    enable_ref_counted &operator=(const enable_ref_counted &) = delete;
+    void internal_addref() const noexcept;
+    void internal_subref() const noexcept;
+    int32_t count_refs() const noexcept;
+};
+
+// This is the handle to a shared object. The handle itself is not
+// thread safe.
+template <typename T>
+class ref_counted
+{
+    // carefully placed here to give the best error messages
+    static_assert(std::derived_from<T,enable_ref_counted>);
+    template <typename X> friend class ref_counted;
+private:
+    T *_ptr;
+    ref_counted(T *ptr) noexcept : _ptr(ptr) {}
+    void maybe_subref() noexcept {
+        if (_ptr) {
+            _ptr->internal_subref();
+        }
+    }
+    static T *maybe_addref(T *ptr) noexcept {
+        if (ptr) {
+            ptr->internal_addref();
+        }
+        return ptr;
+    }
+public:
+    ref_counted() noexcept : _ptr(nullptr) {}
+    ref_counted(ref_counted &&rhs) noexcept : _ptr(std::exchange(rhs._ptr, nullptr)) {}
+    ref_counted(const ref_counted &rhs) noexcept : _ptr(maybe_addref(rhs._ptr)) {}
+    ref_counted &operator=(ref_counted &&rhs) noexcept {
+        maybe_subref();
+        _ptr = std::exchange(rhs._ptr, nullptr);
+        return *this;
+    }
+    ref_counted &operator=(const ref_counted &rhs) noexcept {
+        maybe_subref();
+        _ptr = maybe_addref(rhs._ptr);
+        return *this;
+    }
+    template <typename X> ref_counted(ref_counted<X> &&rhs) noexcept : _ptr(std::exchange(rhs._ptr, nullptr)) {}
+    template <typename X> ref_counted(const ref_counted<X> &rhs) noexcept : _ptr(maybe_addref(rhs._ptr)) {}
+    template <typename X> ref_counted &operator=(ref_counted<X> &&rhs) noexcept {
+        maybe_subref();
+        _ptr = std::exchange(rhs._ptr, nullptr);
+        return *this;
+    }
+    template <typename X> ref_counted &operator=(const ref_counted<X> &rhs) noexcept {
+        maybe_subref();
+        _ptr = maybe_addref(rhs._ptr);
+        return *this;
+    }
+    T *operator->() const noexcept { return _ptr; }
+    T &operator*() const noexcept { return *_ptr; }
+    operator bool() const noexcept { return (_ptr != nullptr); }
+    ~ref_counted() noexcept { maybe_subref(); }
+    // NB: will not call subref
+    T *internal_detach() noexcept { return std::exchange(_ptr, nullptr); }
+    // NB: will not call addref
+    static ref_counted internal_attach(T *ptr) noexcept { return ref_counted(ptr); }
+};
+
+// similar to make_shared;
+// create a reference counted object and return a handle to it
+template <typename T, typename ... Args>
+ref_counted<T> make_ref_counted(Args && ... args) {
+    return ref_counted<T>::internal_attach(new T(std::forward<Args>(args)...));
+}
+
+// similar to shared_from_this;
+// create a new handle to a reference counted object
+// (NB: object must still be valid)
+template <typename T>
+ref_counted<T> ref_counted_from(T &t) {
+    t.internal_addref();
+    return ref_counted<T>::internal_attach(std::addressof(t));
+}
+
+}