From 7ef3185ecbb84567c8759350f1baa30907756c02 Mon Sep 17 00:00:00 2001
From: Camillo Lugaresi <camillol@google.com>
Date: Tue, 15 Nov 2022 16:01:56 -0800
Subject: [PATCH] Allow customizing MultiPool options

These don't need to be constants.

PiperOrigin-RevId: 488782713
---
 mediapipe/gpu/gpu_buffer_multi_pool.cc | 23 +++++------------------
 mediapipe/gpu/gpu_buffer_multi_pool.h  | 22 +++++++++++++++++++++-
 2 files changed, 26 insertions(+), 19 deletions(-)

diff --git a/mediapipe/gpu/gpu_buffer_multi_pool.cc b/mediapipe/gpu/gpu_buffer_multi_pool.cc
index d03ae06a..44f1d40d 100644
--- a/mediapipe/gpu/gpu_buffer_multi_pool.cc
+++ b/mediapipe/gpu/gpu_buffer_multi_pool.cc
@@ -23,26 +23,12 @@
 
 namespace mediapipe {
 
-// Keep this many buffers allocated for a given frame size.
-static constexpr int kKeepCount = 2;
-// The maximum size of the GpuBufferMultiPool. When the limit is reached, the
-// oldest BufferSpec will be dropped.
-static constexpr int kMaxPoolCount = 10;
-// Time in seconds after which an inactive buffer can be dropped from the pool.
-// Currently only used with CVPixelBufferPool.
-static constexpr float kMaxInactiveBufferAge = 0.25;
-// Skip allocating a buffer pool until at least this many requests have been
-// made for a given BufferSpec.
-static constexpr int kMinRequestsBeforePool = 2;
-// Do a deeper flush every this many requests.
-static constexpr int kRequestCountScrubInterval = 50;
-
 #if MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
 
 std::shared_ptr<GpuBufferMultiPool::SimplePool>
 GpuBufferMultiPool::MakeSimplePool(const GpuBufferMultiPool::BufferSpec& spec) {
   return std::make_shared<CvPixelBufferPoolWrapper>(
-      spec.width, spec.height, spec.format, kMaxInactiveBufferAge,
+      spec.width, spec.height, spec.format, options_.max_inactive_buffer_age,
       flush_platform_caches_);
 }
 
@@ -51,7 +37,7 @@ GpuBufferMultiPool::MakeSimplePool(const GpuBufferMultiPool::BufferSpec& spec) {
 std::shared_ptr<GpuBufferMultiPool::SimplePool>
 GpuBufferMultiPool::MakeSimplePool(const BufferSpec& spec) {
   return GlTextureBufferPool::Create(spec.width, spec.height, spec.format,
-                                     kKeepCount);
+                                     options_.keep_count);
 }
 
 #endif  // MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
@@ -64,11 +50,12 @@ std::shared_ptr<GpuBufferMultiPool::SimplePool> GpuBufferMultiPool::RequestPool(
     absl::MutexLock lock(&mutex_);
     pool =
         cache_.Lookup(spec, [this](const BufferSpec& spec, int request_count) {
-          return (request_count >= kMinRequestsBeforePool)
+          return (request_count >= options_.min_requests_before_pool)
                      ? MakeSimplePool(spec)
                      : nullptr;
         });
-    evicted = cache_.Evict(kMaxPoolCount, kRequestCountScrubInterval);
+    evicted = cache_.Evict(options_.max_pool_count,
+                           options_.request_count_scrub_interval);
   }
   // Evicted pools, and their buffers, will be released without holding the
   // lock.
diff --git a/mediapipe/gpu/gpu_buffer_multi_pool.h b/mediapipe/gpu/gpu_buffer_multi_pool.h
index 7feb39ad..1396bcdb 100644
--- a/mediapipe/gpu/gpu_buffer_multi_pool.h
+++ b/mediapipe/gpu/gpu_buffer_multi_pool.h
@@ -42,9 +42,28 @@ namespace mediapipe {
 struct GpuSharedData;
 class CvPixelBufferPoolWrapper;
 
+struct MultiPoolOptions {
+  // Keep this many buffers allocated for a given frame size.
+  int keep_count = 2;
+  // The maximum size of the GpuBufferMultiPool. When the limit is reached, the
+  // oldest BufferSpec will be dropped.
+  int max_pool_count = 10;
+  // Time in seconds after which an inactive buffer can be dropped from the
+  // pool. Currently only used with CVPixelBufferPool.
+  float max_inactive_buffer_age = 0.25;
+  // Skip allocating a buffer pool until at least this many requests have been
+  // made for a given BufferSpec.
+  int min_requests_before_pool = 2;
+  // Do a deeper flush every this many requests.
+  int request_count_scrub_interval = 50;
+};
+
+static constexpr MultiPoolOptions kDefaultMultiPoolOptions;
+
 class GpuBufferMultiPool {
  public:
-  GpuBufferMultiPool() {}
+  GpuBufferMultiPool(MultiPoolOptions options = kDefaultMultiPoolOptions)
+      : options_(options) {}
 
   // Obtains a buffer. May either be reused or created anew.
   GpuBuffer GetBuffer(int width, int height,
@@ -85,6 +104,7 @@ class GpuBufferMultiPool {
   // pool, in which case the caller should invoke CreateBufferWithoutPool.
   std::shared_ptr<SimplePool> RequestPool(const BufferSpec& spec);
 
+  MultiPoolOptions options_;
   absl::Mutex mutex_;
   mediapipe::ResourceCache<BufferSpec, std::shared_ptr<SimplePool>> cache_
       ABSL_GUARDED_BY(mutex_);