4
\$\begingroup\$

The idea here is that we have some resource that is relatively very expensive to generate, but once we have a bunch of them in a pool, we can keep reusing them instead of generating new ones. "HTTP connections" would be an example. In my particular use-case, it's "2048-bit RSA keys" (let's pass over the obligatory what-a-terrible-idea-that-is). Generating a new key takes about 200ms, so I just need a pool that can dole out existing keys faster than 200ms.

My other design goal is that the thing needs to be thread-safe (against maybe thousands of threads at a time), and in fact it's going to be accessed from greenthreads, not std::threads, so I can't use std::mutex; I have to do everything lock-freely.

If you dole out shared read access to items, then you have to think about what happens if thread A is reading an item while thread B is trying to delete/overwrite it. I don't want to deal with that. So I made my AtomicRoundRobinPool dole out exclusive access to items. When you get_item from the pool, you get a unique_ptr to it (nobody else can see it but you). When you're done using the item, you may choose to put_item the unique_ptr back, or you may drop it on the floor; the pool class doesn't care.

If you put_item when there's no room left in the pool, the pool just kicks your item back to you. But we also want the ability to inject fresh blood into a full pool. So I made a force_put_item method, too; it always inserts the item you give it, but it might do so by evicting an item (in which case it kicks the evicted item back to you).

I also wrote a try_quick_get_item method that just checks the current round-robin slot and doesn't iterate over the entire array, but I don't think I'd end up using that method for anything in real life.

#include <atomic>
#include <cassert>
#include <memory>

template<class T, int N>
class AtomicRoundRobinPool {
    std::atomic<T*> array_[N] {};
    std::atomic<int> get_idx_ {0};
    std::atomic<int> put_idx_ {0};
    std::atomic<int> size_ {0};

    int postincrement_mod_N(std::atomic<int>& x) {
        int expected = x.load(std::memory_order_relaxed);
        while (true) {
            int desired = (expected + 1) % N;
            if (x.compare_exchange_strong(expected, desired)) {
                return expected;
            }
        }
    }

public:
    static constexpr int capacity() { return N; }

    int approximate_size() const {
        return size_.load();
    }

    std::unique_ptr<T> try_quick_get_item() noexcept {
        // Try grabbing the first slot we see. If the table is
        // very full, then we expect this usually to work.
        int idx = postincrement_mod_N(get_idx_);
        T *result = array_[idx].exchange(nullptr);
        if (result != nullptr) {
            size_.fetch_sub(1);
        }
        return std::unique_ptr<T>(result);
    }

    std::unique_ptr<T> get_item() noexcept {
        int idx = postincrement_mod_N(get_idx_);
        for (int i=0; i < N; ++i) {
            T *result = array_[idx].exchange(nullptr);
            if (result != nullptr) {
                size_.fetch_sub(1);
                return std::unique_ptr<T>(result);
            }
            idx = (idx + 1) % N;
        }
        // If we've gone around the whole array once and found no items,
        // we should give up.
        return nullptr;
    }

    std::unique_ptr<T> put_item(std::unique_ptr<T> item) noexcept {
        assert(item != nullptr);
        int idx = postincrement_mod_N(put_idx_);
        for (int i=0; i < N; ++i) {
            T *expect_null = nullptr;
            if (array_[idx].compare_exchange_strong(expect_null, item.get())) {
                item.release();
                size_.fetch_add(1);
                return nullptr;
            }
            idx = (idx + 1) % N;
        }
        // If we've gone around the whole array once and found
        // no empty slots, we should give up.
        return item;
    }

    std::unique_ptr<T> force_put_item(std::unique_ptr<T> item) noexcept {
        assert(item != nullptr);
        int idx = postincrement_mod_N(put_idx_);
        T *removed_item = array_[idx].exchange(item.release());
        if (removed_item == nullptr) {
            // We found an empty slot on our first try; excellent.
            size_.fetch_add(1);
            return nullptr;
        } else {
            // `item` is now in the pool, but at the cost of `removed_item`.
            // Put `removed_item` back in the pool somewhere, if possible.
            return put_item(std::unique_ptr<T>(removed_item));
        }
    }

    ~AtomicRoundRobinPool() {
        for (std::atomic<T*>& elt : array_) {
            delete elt.load();
        }
    }
};

And here's my example usage. Sadly, I can't think of any good way to unit-test this thing.

#include <chrono>
#include <stdio.h>
#include <string>
#include <thread>

int main() {
    AtomicRoundRobinPool<std::string, 10> pool;
    auto generate_new_item = [&]() {
        static std::atomic<int> i{0};
        std::this_thread::sleep_for(std::chrono::milliseconds(200));
        return std::make_unique<std::string>(std::to_string(++i));
    };
    auto producer = [&]() {
        for (int i=0; i < 100'000; ++i) {
            puts("Producer is generating new item");
            auto new_item = generate_new_item();
            if (pool.force_put_item(std::move(new_item))) {
                puts("Producer put the new item but removed an old one");
            } else {
                puts("Producer put the new item");
            }
        }
    };
    auto consumer = [&]() {
        for (int i=0; i < 100'000; ++i) {
            std::unique_ptr<std::string> item = pool.get_item();
            if (item != nullptr) {
                printf("Consumer got item: %s\n", item->c_str());
            } else {
                puts("Consumer must generate new item");
                item = generate_new_item();
                // Always put this new item into the pool; it's new blood!
                item = pool.force_put_item(std::move(item));
                if (item != nullptr) {
                    puts("Consumer put the new item but removed an old one");
                } else {
                    puts("Consumer put the new item");
                }
            }
            // At this point, we have either item==nullptr or
            // item== an item that we got from the pool, which means
            // it's been used at least once before.
            if (pool.approximate_size() < pool.capacity() / 2) {
                // The pool is low; put this item back.
                if (item != nullptr) {
                    if (pool.put_item(std::move(item))) {
                        puts("Consumer failed to putback the old item");
                    } else {
                        puts("Consumer putback the old item");
                    }
                }
            }
        }
    };

    std::thread ts[] = {
        std::thread(producer),
        std::thread(consumer),
        std::thread(consumer),
        std::thread(consumer),
        std::thread(consumer),
        std::thread(consumer),
    };

    for (auto&& t : ts) {
        t.join();
    }
}
```
\$\endgroup\$
  • \$\begingroup\$ Isn't it possible for std atomic to use a mutex? If you cannot use a mutex, then you might want to (static_)assert that your atomics are lock free. Or not use atomic. \$\endgroup\$ – sudo rm -rf slash Aug 11 '19 at 13:42

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.