C++ std::shared_ptr implementation

Question

Took a shot at implementing std::shared_ptr, with a thread-safe refcount and weak count. Didn't do weak_ptr, I'm doing this for learning purposes, and I think from an educational standpoint there's not much to learn from adding it.

Also tried to implement the "we know where you live" optimization, with the control block and the underlying memory allocated together in makeShared, but it feels sloppy and error prone and would love advice on if it can be done more cleanly.

#include <atomic>
#include <iostream>

template <typename T>
class ControlBlock {
public:
    ControlBlock(): 
        m_ptr{nullptr},
        m_ref_count{0}, 
        m_weak_count{0}
    {}

    ControlBlock(T* ptr): 
        m_ptr{ptr},
        m_ref_count{1}, 
        m_weak_count{0}
    {}

    void add_reference() {
        m_ref_count.fetch_add(1, std::memory_order_relaxed);
    }

    void add_weak_reference() {
        m_weak_count.fetch_add(1, std::memory_order_relaxed);
    }

    void remove_reference() {
        // The release memory barrier ensures all operations on this thread are 
        // visible to other threads before a decrement happens 
        if (m_ref_count.fetch_sub(1, std::memory_order_release) == 1) {
            // the acquire memory barrier makes sure we delete the pointer 
            // and control block only after the ref count has gone to 0

            std::atomic_thread_fence(std::memory_order_acquire);
            m_ptr->~T();

            if (m_weak_count.load(std::memory_order_relaxed) == 0) {
                delete this;
            }
        }
    }

    void remove_weak_reference() {
        if (m_weak_count.fetch_sub(1, std::memory_order_release) == 1) {
            std::atomic_thread_fence(std::memory_order_acquire);
            if (m_ref_count.load(std::memory_order_relaxed) == 0) {
                delete this;
            }
        }
    }

    T* get() const noexcept {
        return m_ptr;
    }

    std::size_t use_count() const noexcept {
        return m_ref_count.load(std::memory_order_relaxed);
    }

    std::size_t weak_count() const noexcept {
        return m_weak_count.load(std::memory_order_relaxed);
    }

    template<typename... Args>
    static ControlBlock* create (Args&&... args) {
        std::size_t cb_size = sizeof(ControlBlock);
        std::size_t t_size = sizeof(T);

        std::size_t t_alignment = alignof(T);
        std::size_t remainder = cb_size % t_alignment; 
        std::size_t padding = remainder ? (t_alignment - remainder) : 0;
        std::size_t t_offset = cb_size + padding; 

        char* chunk = static_cast<char*>(::operator new(cb_size + t_size));

        T* ptr = nullptr;
        try {
            ptr = new (chunk + t_offset) T(std::forward<Args>(args)...);
        } catch (...) {
            ::operator delete (chunk);
            throw;
        }

        try {
            return new(chunk) ControlBlock{ptr};
        } catch(...) {
            ptr->~T();
            ::operator delete(chunk);
            throw;
        }
    }

private:
    T* m_ptr; 

    std::atomic<std::size_t> m_ref_count;
    std::atomic<std::size_t> m_weak_count;
};

template<typename T>
class SharedPtr {
public:
    SharedPtr() noexcept : cb{new ControlBlock<T>()} {}
    SharedPtr(T* item) : cb{new ControlBlock<T>{item}} {}

    template<typename U>
    requires(std::is_convertible_v<U*, T*>)
    SharedPtr(SharedPtr<U>&& other) noexcept 
        : SharedPtr{other.cb}
    {
        other.cb = nullptr;
    }

    SharedPtr(const SharedPtr& other) : cb{other.cb} {
        cb->add_reference();
    }

    SharedPtr (SharedPtr&& other) : cb{other.cb} {
        other.cb = nullptr;
    }

    SharedPtr& operator=(const SharedPtr& other) {
        if (this == &other) return *this;

        SharedPtr temp{other};
        swap(temp);
        return *this;
    }

    SharedPtr& operator=(SharedPtr&& other) {
        if (this == &other) return *this;

        SharedPtr temp{std::move(other)};
        swap(temp);
        return *this;
    }

    int use_count() const noexcept {
        return cb -> use_count();
    }

    void swap(SharedPtr& other) noexcept {
        std::swap(cb, other.cb);
    }

    T* reset() {
        T* result = get();
        SharedPtr new_ptr;
        swap(new_ptr);

        return result;
    }

    T* reset(T* new_val) {
        T* result = get();

        SharedPtr new_ptr{new_val};
        swap(new_ptr);

        return result;
    }

    T* get() const noexcept {
        return cb->get();
    }

    T& operator*() const {
        return *get();
    }

    T* operator->() const noexcept {
        return get();
    }

    explicit operator bool() const noexcept {
        return get() != nullptr;
    }

    ~SharedPtr() {
        if (cb) {
            cb -> remove_reference();
        }   
    }
private:
    ControlBlock<T>* cb;
    
    // Make all specializations of SharedPtr friends
    // with one another
    template<typename U>
    friend class SharedPtr;

    // Make makeShared a friend of the SharedPtr class
    // so we can use the private ControlBlock constructor
    template<typename U, typename...Args>
    friend SharedPtr<U> makeShared(Args&&... args);

    SharedPtr(ControlBlock<T>* cb) : cb{cb} {}

    template <typename U>
    SharedPtr(ControlBlock<U>* cb) :
        cb{reinterpret_cast<ControlBlock<T>*>(cb)} {}
};

template<typename T, typename... Args>
SharedPtr<T> makeShared(Args&&... args) {
    auto cb = ControlBlock<T>::create(std::forward<Args>(args)...);
    return SharedPtr<T>{cb};
}

Added some unit tests as well, based on the example behavior on cppreference, but again would love some advice on how to really stress test this, especially since the control block creation is so manual and error prone.

// Have to define the test module first, always.
#define BOOST_TEST_MODULE sharedptrtest

#ifdef BOOST_TEST_DYN_LINK
#   include <boost/test/unit_test.hpp>
#else
#   include <boost/test/included/unit_test.hpp>
#endif // BOOST_TEST_DYN_LINK

#include <boost/test/data/test_case.hpp>
#include <boost/test/data/monomorphic.hpp>

#include "SharedPtr.h"

struct Foo {
    int id{0};
    Foo (int i = 0) : id{i} {}
    ~Foo() = default;

    int val() {
        return id;
    }
};

struct B
{
    virtual ~B() = default; 
    virtual int val() const {return 0;}
};
 
struct D : B
{
    D() { }
    ~D() {}
 
    int val() const override {
        return 2;
    }
};


BOOST_AUTO_TEST_CASE(default_constructor_test)
{
    {
        SharedPtr<Foo> s1;
    }   
}

BOOST_AUTO_TEST_CASE(runtime_poly)
{
    SharedPtr<B> p = makeShared<D>();
    assert(p -> val() == 2);
}

BOOST_AUTO_TEST_CASE(copy_constructor_test)
{
    SharedPtr<Foo> s1(new Foo{10});
    BOOST_CHECK_EQUAL(s1.use_count(),  1);
    {
        SharedPtr<Foo> s2{s1}; 
        BOOST_CHECK_EQUAL(s1.use_count(),  2);
    }

    BOOST_CHECK_EQUAL(s1.use_count(),  1);   
}


BOOST_AUTO_TEST_CASE(copy_assigment_test)
{
    SharedPtr<Foo> s1(new Foo{10});
    BOOST_CHECK_EQUAL(s1.use_count(), 1);
    {
        SharedPtr<Foo> s2 = s1; 
        BOOST_CHECK_EQUAL(s1.use_count(), 2);
    }

    BOOST_CHECK_EQUAL(s1.use_count(), 1);   
}

BOOST_AUTO_TEST_CASE(move_assigment_test)
{
    SharedPtr<Foo> s1(new Foo{10});
    BOOST_CHECK_EQUAL(s1.use_count(), 1);

    SharedPtr<Foo> s2 = std::move(s1); 
    BOOST_CHECK_EQUAL(s2.use_count(), 1);   
}


BOOST_AUTO_TEST_CASE(reset_test)
{
    SharedPtr<Foo> s1(new Foo{10});
    BOOST_CHECK_EQUAL(s1->val(), 10);
    BOOST_CHECK_EQUAL(s1.use_count(), 1);

    s1.reset(new Foo{20});
    BOOST_CHECK_EQUAL(s1->val(), 20);
    BOOST_CHECK_EQUAL(s1.use_count(), 1);
    
    s1.reset();
    BOOST_CHECK_EQUAL(s1.use_count(), 0);
}


BOOST_AUTO_TEST_CASE(swap_test)
{
    SharedPtr<Foo> s1 = makeShared<Foo>(100);
    SharedPtr<Foo> s2 = makeShared<Foo>(200);

    BOOST_CHECK_EQUAL(s1->val(), 100);
    BOOST_CHECK_EQUAL(s2->val(), 200);

    BOOST_CHECK_EQUAL(s1.use_count(), 1);
    BOOST_CHECK_EQUAL(s1.use_count(), 1);

    s1.swap(s2);

    BOOST_CHECK_EQUAL(s1->val(), 200);
    BOOST_CHECK_EQUAL(s2->val(), 100);

    BOOST_CHECK_EQUAL(s1.use_count(), 1);
    BOOST_CHECK_EQUAL(s1.use_count(), 1);
}

BOOST_AUTO_TEST_CASE(get_and_deref_test)
{
    int* pInt = new int{42};
    SharedPtr shInt = makeShared<int>(42);

    BOOST_CHECK_EQUAL(*pInt, *shInt);
    BOOST_CHECK_EQUAL(*pInt, *(shInt.get()));
}

BOOST_AUTO_TEST_CASE(bool_test)
{
    SharedPtr<int> ptr;
    BOOST_CHECK(!(static_cast<bool>(ptr)));

    ptr = makeShared<int>(7);
    BOOST_CHECK((static_cast<bool>(ptr)));
}
```

Answer 1

Make ControlBlock's constructors noexcept

The constructors of ControlBlock are very trivial and never throw an exception, so mark these constructors noexcept. That then greatly simplifies create(), where you now only need one try-catch block:

try {
    T *ptr = new (chunk + t_offset) T(std::forward<Args>(args)...);
    return new(chunk) ControlBlock{ptr};
} catch (...) {
    ::operator delete (chunk);
    throw;
}

Incorrect handling of alignment restrictions

There are two obvious mistakes in handling the alignment of T in create():

1. You don't reserve any memory for the padding.
2. You calculate the offset before you even know what address chunk will point to.

You should first allocate cb_size + t_size + t_alignment bytes. Then once you have the pointer to that allocation, you can check how much you have to add to chunk + cb_size to get an address that is correctly aligned for T.

Missing deletion of T

Unless you created the ControlBlock and the T object in one allocation, then there is a delete m_ptr missing. Note that you have to be careful to not both call m_ptr->~T() and delete m_ptr in that case.

Unnecessary checks for self-assignment

You added checks for self-(move)-assignment and early out in those cases. But without these checks, your code is still correct, thanks to the fact that you have used the copy-and-swap idiom. You might do more work than necessary if you do a self-assignment, but consider that it is very unlikely that someone assigns an object to itself, and you actually slightly reduce the performance of the normal case.

The move-assignment operator can be made noexcept

Since the move-assignment operator just does a move-construction and a swap, both of which are noexcept, the move-assignment operator itself can also be made noexcept.

Assignment operators should also be templates

For the same reason you made the move constructor a template, the assignment operators should also be templates, with the same requires clause.

Missing some functionality

Didn't do weak_ptr, I'm doing this for learning purposes, and I think from an educational standpoint there's not much to learn from adding it.

I think you are wrong about that. Implementing all the member functions of std::weak_ptr correctly is perhaps not rocket science, but you don't know about all the corner cases you might have to handle until you try it. Also consider that without it, you can't actually test correct handling of weak references in your test suite.

There is also other functionality missing that is in std::shared_ptr, like support for arrays, owner_before(), the member types element_type and weak_type, other non-member functions that work with std::shared_ptr, std::atomic<std::shared_ptr>, support for custom allocators and deleters, constructors/assignment operators that take a std::unique_ptr as argument, and more.

Add unit tests for ControlBlock as well

In addition to testing SharedPtr, you should add unit tests for ControlBlock as well, just to make sure the internals of your code are also working as you, the developer, are expecting.

Answer 2

There are Numerous Bugs.

First, the Converting Move Constructor

The deleter calls the decrement member function of the control block of ControlBlock, which calls m_ptr->~T();. And the converting constructors to make a SharedPtr<T> from a SharedPtr<U> create a ControlBlock<T>. This means that ~T() will be called instead of ~U(), which will fail unless U inherits a virtual destructor from T.

However, the Standard ([util.smartptr.shared.const]) specifies that this will also work for static, not dynamic, inheritance. Additionally, a cast to a virtual base class pointer might have a different object representation from the original pointer.

Another bug, at least if you meant this to do what the Standard specifies, is that the converting constructor creates a new control block with its own reference count, instead of sharing the reference count.

And this creates a serious memory-safety bug! Let's take a look at this constructor:

    template<typename U>
    requires(std::is_convertible_v<U*, T*>)
    SharedPtr(SharedPtr<U>&& other) noexcept 
        : SharedPtr{other.cb}
    {
        other.cb = nullptr;
    }

This calls the following constructor:

    ControlBlock(T* ptr): 
        m_ptr{ptr},
        m_ref_count{1}, 
        m_weak_count{0}
    {}

Which makes a shallow copy of the pointer with its own reference count of 1.

But we only checked that other was a rvalue reference—not that it was the last reference to its memory! When it isn’t we now have two control blocks that both will try to delete the same pointer. Memory-corruption bug! Simple example:

void foo(SharedPtr<SomeClass>&&);

void bar() {
    SharedPtr<SomeCompatibleClass> p(new SomeCompatibleClass);
    // We still need p afterward, but foo consumes its input. So, pass it a copy!
    foo(std::move(SharedPtr<SomeCompatibleClass>(p)));
    // Why am I getting heap-corruption bugs from here on out?
}

What does this code do? It’s trying to call foo(p), which expects a SharedPtr<SomeClass>&&, with a SharedPtr<SomeCompatibleClass>&&. We’re moving from a copy of the SharedPtr with a reference count of 2, so as not to destroy the original p.

To resolve the overload, the compiler looks for and finds a template constructor that works out to SharedPtr<SomeClass>(SharedPtr<SomeCompatibleClass>&&). It creates a temporary SharedPtr<SomeClass> with its own control block and reference count of 1. When foo returns, the lifetime of this temporary expires. Its destructor frees the SomeClass while another shared pointer to it is still alive. The destructor of p gets called when bar returns—freeing its m_ptr a second time. Whoops!

The Thread-Safety is Also Not Correct

You frequently perform multiple loads or updates on individual data members but depend on them being consistent. For example, your decrement function has this snippet of code:

    void remove_weak_reference() {
        if (m_weak_count.fetch_sub(1, std::memory_order_release) == 1) {
            std::atomic_thread_fence(std::memory_order_acquire);
            if (m_ref_count.load(std::memory_order_relaxed) == 0) {
                delete this;
            }
        }
    }

This is not thread-safe! Both individual operations are atomic, but they’re separate operations that might return inconsistent values. For example, if the thread that decrements ref_count and the thread that decrements weak_count see the older values of the other count, neither thread will ever see both equal to 0, causing a memory leak.

What you need to do is retrieve and update both counts in the same atomic operation. You also want this operation to be a lock-free atomic compare-and-swap. On variants of x86 that support the cmpxchg16b instruction, mainstream compilers will generate it if you tell it to support a recent-enough instruction set, and if you pack the control block into a 16-byte struct or class declared with alignas(16). I would suggest, however, that you test for this at compile time, using something like (adjust this to work on your implementation):

    static_assert( std::atomic<ControlBlock<void>>::is_always_lock_free,
                   "ControlBlock does not support lock-free atomic operations with these compiler settings." );

Since it’s likely that your final implementation will need to be larger than 16 bytes, consider atomically fetching and updating just a struct or bitfield containing the two reference counts, which can be eight bytes long. On many architectures, this can use ordinary 64-bit instructions.

You absolutely want to avoid the overhead of mutex locking on every single operation!

Block Allocation Looks More Complex than it Needs to Be

The new operator already aligns memory correctly for the type it’s allocating, so I’m not sure why there’s so much low-level stuff in ControlBlock<T>::create(). If you do need low-level allocation of aligned memory, use std::aligned_alloc.

You Will Need a Different Approach to Support the Missing Features

You probably left them out for simplicity, but some of them would require significant redesigns to add later. These include:

• Each SharedPtr<T> cannot reference a ControlBlock<T> if reference counts are to be correctly shared with a SharedPtr<U>. You either want there to be a single ControlBlock class, or an abstract base class ControlBlock<Deleter>.
• Therefore, m_ptr should be moved into the shared_ptr instances. This also speeds up access to it and enables you to support pointer conversions that change the pointer’s object representation, or std::dynamic_ptr_cast.
• If you support std::get_deleter, you’ll need to store the deleter in the control block and remember its type, so you will want each ControlBlock<Deleter> to inherit from the same abstract base class. If retrieving the deleter is unnecessary, you can get by with a single ControlBlock class that stores the deleter in a std::function<void()> with [[no_unique_address]].
• Therefore, a ControlBlock cannot in the general case be small enough to have a lock-free, thread-safe implementation, and you will need to declare some other atomic type to read and update both counts in the same atomic operation.
• If you want to support std::get_deleter<Deleter, T>, it will be very helpful to have a helper function virtual void* ControlBlock::get_deleter_helper(const std::type_info&& dtype) const noexcept, called from static_cast<Deleter<T>*>(p.cb->get_deleter_helper(typeid(Deleter<T>))).