1 files changed, 151 insertions, 0 deletions
diff --git a/src/libutil/pool.hh b/src/libutil/pool.hh
new file mode 100644
index 000000000000..f291cd578388
--- /dev/null
+++ b/src/libutil/pool.hh
@@ -0,0 +1,151 @@
+#pragma once
+
+#include <functional>
+#include <limits>
+#include <list>
+#include <memory>
+#include <cassert>
+
+#include "sync.hh"
+#include "ref.hh"
+
+namespace nix {
+
+/* This template class implements a simple pool manager of resources
+   of some type R, such as database connections. It is used as
+   follows:
+
+     class Connection { ... };
+
+     Pool<Connection> pool;
+
+     {
+       auto conn(pool.get());
+       conn->exec("select ...");
+     }
+
+   Here, the Connection object referenced by ‘conn’ is automatically
+   returned to the pool when ‘conn’ goes out of scope.
+*/
+
+template <class R>
+class Pool
+{
+public:
+
+    /* A function that produces new instances of R on demand. */
+    typedef std::function<ref<R>()> Factory;
+
+    /* A function that checks whether an instance of R is still
+       usable. Unusable instances are removed from the pool. */
+    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;
+    };
+
+    Sync<State> state;
+
+    std::condition_variable wakeup;
+
+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;
+    }
+
+    ~Pool()
+    {
+        auto state_(state.lock());
+        assert(!state_->inUse);
+        state_->max = 0;
+        state_->idle.clear();
+    }
+
+    class Handle
+    {
+    private:
+        Pool & pool;
+        std::shared_ptr<R> r;
+
+        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());
+                state_->idle.push_back(ref<R>(r));
+                assert(state_->inUse);
+                state_->inUse--;
+            }
+            pool.wakeup.notify_one();
+        }
+
+        R * operator -> () { return &*r; }
+        R & operator * () { return *r; }
+    };
+
+    Handle get()
+    {
+        {
+            auto state_(state.lock());
+
+            /* If we're over the maximum number of instance, we need
+               to wait until a slot becomes available. */
+            while (state_->idle.empty() && state_->inUse >= state_->max)
+                state_.wait(wakeup);
+
+            while (!state_->idle.empty()) {
+                auto p = state_->idle.back();
+                state_->idle.pop_back();
+                if (validator(p)) {
+                    state_->inUse++;
+                    return Handle(*this, p);
+                }
+            }
+
+            state_->inUse++;
+        }
+
+        /* 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());
+            return h;
+        } catch (...) {
+            auto state_(state.lock());
+            state_->inUse--;
+            throw;
+        }
+    }
+
+    unsigned int count()
+    {
+        auto state_(state.lock());
+        return state_->idle.size() + state_->inUse;
+    }
+};
+
+}