shared_ptr that removes itself from a container that owns it, is there a better way?

Question

What I want to do is basically queue a bunch to task objects to a container, where the task can remove itself from the queue. But I also don't want the object to be destroyed when it removes itself, so it can continue to finish whatever the work is doing.

So, a safe way to do this is either call RemoveSelf() when the work is done, or use a keepAlive reference then continue to do the work. I've verified that this does indeed work, while the DoWorkUnsafe will always crash after a few iterations.

I'm not particularly happy with the solution, because I have to either remember to call RemoveSelf() at the end of work being done, or remember to use a keepAlive, otherwise it will cause undefined behavior.

Another problem is that if someone decides to iterate through the ownerList and do work, it would invalidate the iterator as they iterate, which is also unsafe.

Alternatively, I know I can instead put the task onto a separate "cleanup" queue and destroy finished tasks separately. But this method seemed neater to me, but with too many caveats.

Is there a better pattern to handle something like this?

#include <memory>
#include <unordered_set>

class SelfDestruct : public std::enable_shared_from_this<SelfDestruct> {
public:
    SelfDestruct(std::unordered_set<std::shared_ptr<SelfDestruct>> &ownerSet)
        : _ownerSet(ownerSet){}

    void DoWorkUnsafe() {
        RemoveSelf();
        DoWork();
    }

    void DoWorkSafe() {
        DoWork();
        RemoveSelf();
    }

    void DoWorkAlsoSafe() {
        auto keepAlive = RemoveSelf();
        DoWork();
    }


    std::shared_ptr<SelfDestruct> RemoveSelf() { 
        auto keepAlive = shared_from_this();
        _ownerSet.erase(keepAlive);
        return keepAlive;
    };

private:

    void DoWork() {
        for (auto i = 0; i < 100; ++i)
            _dummy.push_back(i);
    }

    std::unordered_set<std::shared_ptr<SelfDestruct>> &_ownerSet;
    std::vector<int> _dummy;
};

TEST_CASE("Self destruct should not cause undefined behavior") {

    std::unordered_set<std::shared_ptr<SelfDestruct>> ownerSet;

    for (auto i = 0; i < 100; ++i)
        ownerSet.emplace(std::make_shared<SelfDestruct>(ownerSet));

    while (!ownerSet.empty()) {
        (*ownerSet.begin())->DoWorkSafe();
    }
}

Answer 1

There is a good design principle that says each class should have exactly one purpose. A "task object" should exist to perform that task. When you start adding additional responsibilities, you tend to end up with a mess. Messes can include having to remember to call a certain method after completing the primary purpose, or having to remember to use a hacky workaround to keep the object alive. Messes are often a sign of inadequate thought put into the design. Being unhappy with a mess speaks well of your potential for good design.

Let us backtrack and look at the real problem. There are task objects stored in a container. The container decides when to invoke each task. The task must be removed from the container before the next task is invoked (so that it is not invoked again). It looks to me like the responsibility for removing elements from the container should fall to the container.

So we'll re-envision your class without that "SelfDestruct" mess. Your task objects exist to perform a task. They are probably polymorphic, hence the need for a container of pointers to task objects rather than a container of task objects. The task objects don't care how they are managed; that is work for someone else.

class Task {
public:
    Task() {}
    // Other constructors, the destructor, assignment operators, etc. go here

    void DoWork() {
        // Stuff is done here.
        // The work might involve adding tasks to the queue.
    }
};

Now focus on the container. The container (more precisely, the container's owner) is responsible for adding and removing elements. So do that. You seem to prefer removing the element before invoking it. That seems like a good idea to me, but don't try to pawn off the removal on the task. Instead use a helper function, keeping this logic at the abstraction level of the container's owner.

// Extract the first element of `ownerSet`. That is, remove it and return it.
// ASSUMES: `ownerSet` is not empty
std::shared_ptr<Task> extract(std::unordered_set<std::shared_ptr<Task>>& ownerSet)
{
    auto begin = ownerSet.begin();
    std::shared_ptr<Task> first{*begin};
    ownerSet.erase(begin);
    return first;
}

TEST_CASE("Removal from the container should not cause undefined behavior") {

    std::unordered_set<std::shared_ptr<Task>> ownerSet;

    for (int i = 0; i < 100; ++i)
        ownerSet.emplace(std::make_shared<Task>());

    while (!ownerSet.empty()) {
        // The temporary returned by extract() will live until the semicolon,
        // so it will (barely) outlive the call to DoWork().
        extract(ownerSet)->DoWork();
        // This is equivalent to:
        //auto todo{extract(ownerSet)};
        //todo->DoWork();
    }

}

From one perspective, this is an almost trivial change from your approach, as all I did was shift a responsibility from the task object to the owner of the container. Yet with this shift, the mess disappears. The same steps are performed, but they make sense and are almost forced when moved to a more appropriate context. Clean design tends to lead to clean implementation.