pool.hh

Go to the documentation of this file.

#pragma once
 
#include <functional>
#include <kj/async.h>
#include <limits>
#include <list>
#include <memory>
#include <cassert>
#include <optional>
#include <vector>
 
#include "lix/libutil/async.hh"
#include "lix/libutil/result.hh"
#include "lix/libutil/sync.hh"
#include "lix/libutil/ref.hh"
#include "lix/libutil/types.hh"
 
namespace nix {

template <class R>
class Pool
{
public:

    typedef std::function<ref<R>()> Factory;

    typedef std::function<bool(const ref<R> &)> Validator;
 
private:
 
    Factory factory;
    Validator validator;
 
    struct State
    {
        size_t inUse = 0;
        size_t max;
        std::vector<ref<R>> idle;
        std::list<kj::Own<kj::CrossThreadPromiseFulfiller<void>>> waiters;
 
        void notify()
        {
            for (auto & waiter : waiters) {
                waiter->fulfill();
            }
            waiters.clear();
        }
    };
 
    Sync<State> state;
 
public:
 
    Pool(size_t max = std::numeric_limits<size_t>::max(),
        const Factory & factory = []() { return make_ref<R>(); },
        const Validator & validator = [](ref<R> r) { return true; })
        : factory(factory)
        , validator(validator)
    {
        auto state_(state.lock());
        state_->max = max;
    }
 
    void incCapacity()
    {
        auto state_(state.lock());
        state_->max++;
        /* we could wakeup here, but this is only used when we're
         * about to nest Pool usages, and we want to save the slot for
         * the nested use if we can
         */
    }
 
    void decCapacity()
    {
        auto state_(state.lock());
        state_->max--;
    }
 
    ~Pool()
    {
        auto state_(state.lock());
        assert(!state_->inUse);
        state_->max = 0;
        state_->idle.clear();
    }
 
    class Handle
    {
    private:
        Pool & pool;
        std::shared_ptr<R> r;
        bool bad = false;
 
        friend Pool;
 
        Handle(Pool & pool, std::shared_ptr<R> r) : pool(pool), r(r) { }
 
    public:
        Handle(Handle && h) : pool(h.pool), r(h.r) { h.r.reset(); }
 
        Handle(const Handle & l) = delete;
 
        ~Handle()
        {
            if (!r) return;
            {
                auto state_(pool.state.lock());
                if (!bad)
                    state_->idle.push_back(ref<R>::unsafeFromPtr(r));
                assert(state_->inUse);
                state_->inUse--;
                state_->notify();
            }
        }
 
        R * operator -> () { return &*r; }
        R & operator * () { return *r; }
 
        void markBad() { bad = true; }
    };
 
private:
    void getFailed()
    {
        auto state_(state.lock());
        state_->inUse--;
        state_->notify();
    }
 
    // lock lifetimes must always be short, and *NEVER* cross a yield point.
    // we ensure this by using explicit continuations instead of coroutines.
    kj::Promise<Result<std::optional<Handle>>> tryGet()
    try {
        auto state_(state.lock());
 
        /* If we're over the maximum number of instance, we need
           to wait until a slot becomes available. */
        if (state_->idle.empty() && state_->inUse >= state_->max) {
            auto pfp = kj::newPromiseAndCrossThreadFulfiller<void>();
            state_->waiters.push_back(std::move(pfp.fulfiller));
            return pfp.promise.then([this] { return tryGet(); });
        }
 
        while (!state_->idle.empty()) {
            auto p = state_->idle.back();
            state_->idle.pop_back();
            if (validator(p)) {
                state_->inUse++;
                return {Handle(*this, p)};
            }
        }
 
        state_->inUse++;
        return {std::nullopt};
    } catch (...) {
        return {result::current_exception()};
    }
 
public:
    kj::Promise<Result<Handle>> get()
    try {
        if (auto existing = LIX_TRY_AWAIT(tryGet())) {
            co_return std::move(*existing);
        }
 
        /* We need to create a new instance. Because that might take a
           while, we don't hold the lock in the meantime. */
        try {
            Handle h(*this, factory());
            co_return h;
        } catch (...) {
            getFailed();
            throw;
        }
    } catch (...) {
        co_return result::current_exception();
    }
 
    size_t count()
    {
        auto state_(state.lock());
        return state_->idle.size() + state_->inUse;
    }
 
    size_t capacity()
    {
        return state.lock()->max;
    }
};
 
}