diff --git a/generic_server.cc b/generic_server.cc
index 185cb9f79b..55eedc2214 100644
--- a/generic_server.cc
+++ b/generic_server.cc
@@ -28,9 +28,10 @@ class counted_data_source_impl : public data_source_impl {
         if (_cpu_concurrency.stopped) {
             return fun();
         }
+        size_t units = _cpu_concurrency.units.count();
         _cpu_concurrency.units.return_all();
-        return fun().finally([this] () {
-            _cpu_concurrency.units.adopt(consume_units(_cpu_concurrency.semaphore, 1));
+        return fun().finally([this, units] () {
+            _cpu_concurrency.units.adopt(consume_units(_cpu_concurrency.semaphore, units));
         });
     };
 public:
@@ -56,9 +57,10 @@ class counted_data_sink_impl : public data_sink_impl {
         if (_cpu_concurrency.stopped) {
             return fun();
         }
+        size_t units = _cpu_concurrency.units.count();
         _cpu_concurrency.units.return_all();
-        return fun().finally([this] () {
-            _cpu_concurrency.units.adopt(consume_units(_cpu_concurrency.semaphore, 1));
+        return fun().finally([this, units] () {
+            _cpu_concurrency.units.adopt(consume_units(_cpu_concurrency.semaphore, units));
         });
     };
 public: