1 files changed, 54 insertions, 15 deletions

diff --git a/src/libutil/pool.hh b/src/libutil/pool.hh
index d63912e28aa9..4b865a19389e 100644
--- a/src/libutil/pool.hh
+++ b/src/libutil/pool.hh
@@ -1,8 +1,10 @@
 #pragma once
 
-#include <memory>
-#include <list>
 #include <functional>
+#include <limits>
+#include <list>
+#include <memory>
+#include <cassert>
 
 #include "sync.hh"
 #include "ref.hh"
@@ -39,37 +41,58 @@ private:
 
     struct State
     {
-        unsigned int count = 0;
-        std::list<ref<R>> idle;
+        size_t inUse = 0;
+        size_t max;
+        std::vector<ref<R>> idle;
     };
 
     Sync<State> state;
 
+    std::condition_variable_any wakeup;
+
 public:
 
-    Pool(const Factory & factory = []() { return make_ref<R>(); })
+    Pool(size_t max = std::numeric_limits<size_t>::max,
+        const Factory & factory = []() { return make_ref<R>(); })
         : factory(factory)
-    { }
+    {
+        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;
-        ref<R> r;
+        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) { abort(); }
+        Handle(Handle && h) : pool(h.pool), r(h.r) { h.r.reset(); }
 
         Handle(const Handle & l) = delete;
 
         ~Handle()
         {
-            auto state_(pool.state.lock());
-            state_->idle.push_back(r);
+            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; }
@@ -80,22 +103,38 @@ public:
     {
         {
             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);
+
             if (!state_->idle.empty()) {
                 auto p = state_->idle.back();
                 state_->idle.pop_back();
+                state_->inUse++;
                 return Handle(*this, p);
             }
-            state_->count++;
+
+            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;
         }
-        /* Note: we don't hold the lock while creating a new instance,
-           because creation might take a long time. */
-        return Handle(*this, factory());
     }
 
     unsigned int count()
     {
         auto state_(state.lock());
-        return state_->count;
+        return state_->count + state_->inUse;
     }
 };