Factor out ReusablePool

PiperOrigin-RevId: 488783477

mediapipe/gpu/BUILD

@@ -607,6 +607,7 @@ cc_library(
         ":gpu_buffer",
         ":gpu_shared_data_header",
         ":multi_pool",
+        ":reusable_pool",
         "//mediapipe/framework:calculator_context",
         "//mediapipe/framework:calculator_node",
         "//mediapipe/framework/port:logging",
@@ -615,6 +616,16 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "reusable_pool",
+    hdrs = ["reusable_pool.h"],
+    deps = [
+        ":multi_pool",
+        "@com_google_absl//absl/functional:any_invocable",
+        "@com_google_absl//absl/synchronization",
+    ],
+)
+
 cc_library(
     name = "multi_pool",
     hdrs = ["multi_pool.h"],

mediapipe/gpu/gl_texture_buffer.h

@@ -71,6 +71,11 @@ class GlTextureBuffer
   // Create a texture with a copy of the data in image_frame.
   static std::unique_ptr<GlTextureBuffer> Create(const ImageFrame& image_frame);
 
+  static std::unique_ptr<GlTextureBuffer> Create(
+      const internal::GpuBufferSpec& spec) {
+    return Create(spec.width, spec.height, spec.format);
+  }
+
   // Wraps an existing texture, but does not take ownership of it.
   // deletion_callback is invoked when the GlTextureBuffer is released, so
   // the caller knows that the texture is no longer in use.

mediapipe/gpu/gl_texture_buffer_pool.cc

@@ -16,79 +16,4 @@
 
 #include "absl/synchronization/mutex.h"
 
-namespace mediapipe {
+namespace mediapipe {}  // namespace mediapipe
-
-GlTextureBufferPool::GlTextureBufferPool(int width, int height,
-                                         GpuBufferFormat format, int keep_count)
-    : width_(width),
-      height_(height),
-      format_(format),
-      keep_count_(keep_count) {}
-
-GlTextureBufferSharedPtr GlTextureBufferPool::GetBuffer() {
-  std::unique_ptr<GlTextureBuffer> buffer;
-  bool reuse = false;
-
-  {
-    absl::MutexLock lock(&mutex_);
-    if (available_.empty()) {
-      buffer = GlTextureBuffer::Create(width_, height_, format_);
-      if (!buffer) return nullptr;
-    } else {
-      buffer = std::move(available_.back());
-      available_.pop_back();
-      reuse = true;
-    }
-
-    ++in_use_count_;
-  }
-
-  // This needs to wait on consumer sync points, therefore it should not be
-  // done while holding the mutex.
-  if (reuse) {
-    buffer->Reuse();
-  }
-
-  // Return a shared_ptr with a custom deleter that adds the buffer back
-  // to our available list.
-  std::weak_ptr<GlTextureBufferPool> weak_pool(shared_from_this());
-  return std::shared_ptr<GlTextureBuffer>(
-      buffer.release(), [weak_pool](GlTextureBuffer* buf) {
-        auto pool = weak_pool.lock();
-        if (pool) {
-          pool->Return(absl::WrapUnique(buf));
-        } else {
-          delete buf;
-        }
-      });
-}
-
-std::pair<int, int> GlTextureBufferPool::GetInUseAndAvailableCounts() {
-  absl::MutexLock lock(&mutex_);
-  return {in_use_count_, available_.size()};
-}
-
-void GlTextureBufferPool::Return(std::unique_ptr<GlTextureBuffer> buf) {
-  std::vector<std::unique_ptr<GlTextureBuffer>> trimmed;
-  {
-    absl::MutexLock lock(&mutex_);
-    --in_use_count_;
-    available_.emplace_back(std::move(buf));
-    TrimAvailable(&trimmed);
-  }
-  // The trimmed buffers will be released without holding the lock.
-}
-
-void GlTextureBufferPool::TrimAvailable(
-    std::vector<std::unique_ptr<GlTextureBuffer>>* trimmed) {
-  int keep = std::max(keep_count_ - in_use_count_, 0);
-  if (available_.size() > keep) {
-    auto trim_it = std::next(available_.begin(), keep);
-    if (trimmed) {
-      std::move(trim_it, available_.end(), std::back_inserter(*trimmed));
-    }
-    available_.erase(trim_it, available_.end());
-  }
-}
-
-}  // namespace mediapipe

mediapipe/gpu/gl_texture_buffer_pool.h

@@ -24,11 +24,11 @@
 #include "absl/synchronization/mutex.h"
 #include "mediapipe/gpu/gl_texture_buffer.h"
 #include "mediapipe/gpu/multi_pool.h"
+#include "mediapipe/gpu/reusable_pool.h"
 
 namespace mediapipe {
 
-class GlTextureBufferPool
+class GlTextureBufferPool : public ReusablePool<GlTextureBuffer> {
-    : public std::enable_shared_from_this<GlTextureBufferPool> {
  public:
   // Creates a pool. This pool will manage buffers of the specified dimensions,
   // and will keep keep_count buffers around for reuse.
@@ -37,52 +37,32 @@ class GlTextureBufferPool
   static std::shared_ptr<GlTextureBufferPool> Create(int width, int height,
                                                      GpuBufferFormat format,
                                                      int keep_count) {
-    return std::shared_ptr<GlTextureBufferPool>(
+    return Create({width, height, format}, {.keep_count = keep_count});
-        new GlTextureBufferPool(width, height, format, keep_count));
   }
 
   static std::shared_ptr<GlTextureBufferPool> Create(
       const internal::GpuBufferSpec& spec, const MultiPoolOptions& options) {
-    return Create(spec.width, spec.height, spec.format, options.keep_count);
+    return std::shared_ptr<GlTextureBufferPool>(
+        new GlTextureBufferPool(spec, options));
   }
 
-  // Obtains a buffers. May either be reused or created anew.
+  int width() const { return spec_.width; }
-  // A GlContext must be current when this is called.
+  int height() const { return spec_.height; }
-  GlTextureBufferSharedPtr GetBuffer();
+  GpuBufferFormat format() const { return spec_.format; }
-
-  int width() const { return width_; }
-  int height() const { return height_; }
-  GpuBufferFormat format() const { return format_; }
-
-  // This method is meant for testing.
-  std::pair<int, int> GetInUseAndAvailableCounts();
 
   static GlTextureBufferSharedPtr CreateBufferWithoutPool(
       const internal::GpuBufferSpec& spec) {
-    return GlTextureBuffer::Create(spec.width, spec.height, spec.format);
+    return GlTextureBuffer::Create(spec);
   }
 
- private:
+ protected:
-  GlTextureBufferPool(int width, int height, GpuBufferFormat format,
+  GlTextureBufferPool(const internal::GpuBufferSpec& spec,
-                      int keep_count);
+                      const MultiPoolOptions& options)
+      : ReusablePool<GlTextureBuffer>(
+            [this] { return GlTextureBuffer::Create(spec_); }, options),
+        spec_(spec) {}
 
-  // Return a buffer to the pool.
+  const internal::GpuBufferSpec spec_;
-  void Return(std::unique_ptr<GlTextureBuffer> buf);
-
-  // If the total number of buffers is greater than keep_count, destroys any
-  // surplus buffers that are no longer in use.
-  void TrimAvailable(std::vector<std::unique_ptr<GlTextureBuffer>>* trimmed)
-      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
-
-  const int width_;
-  const int height_;
-  const GpuBufferFormat format_;
-  const int keep_count_;
-
-  absl::Mutex mutex_;
-  int in_use_count_ ABSL_GUARDED_BY(mutex_) = 0;
-  std::vector<std::unique_ptr<GlTextureBuffer>> available_
-      ABSL_GUARDED_BY(mutex_);
 };
 
 }  // namespace mediapipe

mediapipe/gpu/reusable_pool.h

@@ -0,0 +1,145 @@
+// Copyright 2019 The MediaPipe Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Consider this file an implementation detail. None of this is part of the
+// public API.
+
+#ifndef MEDIAPIPE_GPU_REUSABLE_POOL_H_
+#define MEDIAPIPE_GPU_REUSABLE_POOL_H_
+
+#include <utility>
+#include <vector>
+
+#include "absl/functional/any_invocable.h"
+#include "absl/synchronization/mutex.h"
+#include "mediapipe/gpu/multi_pool.h"
+
+namespace mediapipe {
+
+template <class Item>
+class ReusablePool : public std::enable_shared_from_this<ReusablePool<Item>> {
+ public:
+  using ItemFactory = absl::AnyInvocable<std::unique_ptr<Item>() const>;
+
+  // Creates a pool. This pool will manage buffers of the specified dimensions,
+  // and will keep keep_count buffers around for reuse.
+  // We enforce creation as a shared_ptr so that we can use a weak reference in
+  // the buffers' deleters.
+  static std::shared_ptr<ReusablePool<Item>> Create(
+      ItemFactory item_factory, const MultiPoolOptions& options) {
+    return std::shared_ptr<ReusablePool<Item>>(
+        new ReusablePool<Item>(std::move(item_factory), options));
+  }
+
+  // Obtains a buffer. May either be reused or created anew.
+  // A GlContext must be current when this is called.
+  std::shared_ptr<Item> GetBuffer();
+
+  // This method is meant for testing.
+  std::pair<int, int> GetInUseAndAvailableCounts();
+
+ protected:
+  ReusablePool(ItemFactory item_factory, const MultiPoolOptions& options)
+      : item_factory_(std::move(item_factory)),
+        keep_count_(options.keep_count) {}
+
+ private:
+  // Return a buffer to the pool.
+  void Return(std::unique_ptr<Item> buf);
+
+  // If the total number of buffers is greater than keep_count, destroys any
+  // surplus buffers that are no longer in use.
+  void TrimAvailable(std::vector<std::unique_ptr<Item>>* trimmed)
+      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
+
+  const ItemFactory item_factory_;
+  const int keep_count_;
+
+  absl::Mutex mutex_;
+  int in_use_count_ ABSL_GUARDED_BY(mutex_) = 0;
+  std::vector<std::unique_ptr<Item>> available_ ABSL_GUARDED_BY(mutex_);
+};
+
+template <class Item>
+inline std::shared_ptr<Item> ReusablePool<Item>::GetBuffer() {
+  std::unique_ptr<Item> buffer;
+  bool reuse = false;
+
+  {
+    absl::MutexLock lock(&mutex_);
+    if (available_.empty()) {
+      buffer = item_factory_();
+      if (!buffer) return nullptr;
+    } else {
+      buffer = std::move(available_.back());
+      available_.pop_back();
+      reuse = true;
+    }
+
+    ++in_use_count_;
+  }
+
+  // This needs to wait on consumer sync points, therefore it should not be
+  // done while holding the mutex.
+  if (reuse) {
+    buffer->Reuse();
+  }
+
+  // Return a shared_ptr with a custom deleter that adds the buffer back
+  // to our available list.
+  std::weak_ptr<ReusablePool<Item>> weak_pool(this->shared_from_this());
+  return std::shared_ptr<Item>(buffer.release(), [weak_pool](Item* buf) {
+    auto pool = weak_pool.lock();
+    if (pool) {
+      pool->Return(absl::WrapUnique(buf));
+    } else {
+      delete buf;
+    }
+  });
+}
+
+template <class Item>
+inline std::pair<int, int> ReusablePool<Item>::GetInUseAndAvailableCounts() {
+  absl::MutexLock lock(&mutex_);
+  return {in_use_count_, available_.size()};
+}
+
+template <class Item>
+void ReusablePool<Item>::Return(std::unique_ptr<Item> buf) {
+  std::vector<std::unique_ptr<Item>> trimmed;
+  {
+    absl::MutexLock lock(&mutex_);
+    --in_use_count_;
+    available_.emplace_back(std::move(buf));
+    TrimAvailable(&trimmed);
+  }
+  // The trimmed buffers will be released without holding the lock.
+}
+
+template <class Item>
+void ReusablePool<Item>::TrimAvailable(
+    std::vector<std::unique_ptr<Item>>* trimmed) {
+  int keep = std::max(keep_count_ - in_use_count_, 0);
+  if (available_.size() > keep) {
+    auto trim_it = std::next(available_.begin(), keep);
+    if (trimmed) {
+      std::move(trim_it, available_.end(), std::back_inserter(*trimmed));
+    }
+    available_.erase(trim_it, available_.end());
+  }
+}
+
+}  // namespace mediapipe
+
+#endif  // MEDIAPIPE_GPU_REUSABLE_POOL_H_