1 files changed, 176 insertions, 0 deletions
diff --git a/third_party/nix/src/libutil/pool.hh b/third_party/nix/src/libutil/pool.hh
new file mode 100644
index 000000000000..b5c3c4b5c43f
--- /dev/null
+++ b/third_party/nix/src/libutil/pool.hh
@@ -0,0 +1,176 @@
+#pragma once
+
+#include <cassert>
+#include <functional>
+#include <limits>
+#include <list>
+#include <memory>
+
+#include "libutil/ref.hh"
+#include "libutil/sync.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. */
+  using Validator = std::function<bool(const ref<R>&)>;
+
+ 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:
+  explicit 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>(r));
+        }
+        assert(state_->inUse);
+        state_->inUse--;
+      }
+      pool.wakeup.notify_one();
+    }
+
+    R* operator->() { return &*r; }
+    R& operator*() { return *r; }
+
+    void markBad() { bad = true; }
+  };
+
+  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--;
+      wakeup.notify_one();
+      throw;
+    }
+  }
+
+  size_t count() {
+    auto state_(state.lock());
+    return state_->idle.size() + state_->inUse;
+  }
+
+  size_t capacity() { return state.lock()->max; }
+
+  void flushBad() {
+    auto state_(state.lock());
+    std::vector<ref<R>> left;
+    for (auto& p : state_->idle) {
+      if (validator(p)) {
+        left.push_back(p);
+      }
+    }
+    std::swap(state_->idle, left);
+  }
+};
+
+}  // namespace nix