Move flush hook to CvPixelBufferPoolWrapper constructor

This unifies the implementation of GpuBufferMultiPool::GetBufferFromSimplePool.

PiperOrigin-RevId: 488782173

mediapipe/gpu/cv_pixel_buffer_pool_wrapper.cc

@@ -23,18 +23,18 @@
 
 namespace mediapipe {
 
-CvPixelBufferPoolWrapper::CvPixelBufferPoolWrapper(int width, int height,
-                                                   GpuBufferFormat format,
-                                                   CFTimeInterval maxAge) {
+CvPixelBufferPoolWrapper::CvPixelBufferPoolWrapper(
+    int width, int height, GpuBufferFormat format, CFTimeInterval maxAge,
+    std::function<void(void)> flush_texture_caches) {
   OSType cv_format = CVPixelFormatForGpuBufferFormat(format);
   CHECK_NE(cv_format, -1) << "unsupported pixel format";
   pool_ = MakeCFHolderAdopting(
       /* keep count is 0 because the age param keeps buffers around anyway */
       CreateCVPixelBufferPool(width, height, cv_format, 0, maxAge));
+  flush_texture_caches_ = std::move(flush_texture_caches);
 }
 
-CFHolder<CVPixelBufferRef> CvPixelBufferPoolWrapper::GetBuffer(
-    std::function<void(void)> flush) {
+CFHolder<CVPixelBufferRef> CvPixelBufferPoolWrapper::GetBuffer() {
   CVPixelBufferRef buffer;
   int threshold = 1;
   NSMutableDictionary* auxAttributes =
@@ -47,12 +47,12 @@ CFHolder<CVPixelBufferRef> CvPixelBufferPoolWrapper::GetBuffer(
         kCFAllocatorDefault, *pool_, (__bridge CFDictionaryRef)auxAttributes,
         &buffer);
     if (err != kCVReturnWouldExceedAllocationThreshold) break;
-    if (flush && !tried_flushing) {
+    if (flush_texture_caches_ && !tried_flushing) {
       // Call the flush function to potentially release old holds on buffers
       // and try again to create a pixel buffer.
       // This is used to flush CV texture caches, which may retain buffers until
       // flushed.
-      flush();
+      flush_texture_caches_();
       tried_flushing = true;
     } else {
       ++threshold;

mediapipe/gpu/cv_pixel_buffer_pool_wrapper.h

@@ -32,8 +32,9 @@ namespace mediapipe {
 class CvPixelBufferPoolWrapper {
  public:
   CvPixelBufferPoolWrapper(int width, int height, GpuBufferFormat format,
-                           CFTimeInterval maxAge);
-  CFHolder<CVPixelBufferRef> GetBuffer(std::function<void(void)> flush);
+                           CFTimeInterval maxAge,
+                           std::function<void(void)> flush_texture_caches);
+  CFHolder<CVPixelBufferRef> GetBuffer();
 
   int GetBufferCount() const { return count_; }
   std::string GetDebugString() const;
@@ -46,6 +47,7 @@ class CvPixelBufferPoolWrapper {
  private:
   CFHolder<CVPixelBufferPoolRef> pool_;
   int count_ = 0;
+  std::function<void(void)> flush_texture_caches_;
 };
 
 }  // namespace mediapipe

mediapipe/gpu/gpu_buffer_multi_pool.cc

@@ -48,12 +48,8 @@ static constexpr int kRequestCountScrubInterval = 50;
 std::shared_ptr<GpuBufferMultiPool::SimplePool>
 GpuBufferMultiPool::MakeSimplePool(const GpuBufferMultiPool::BufferSpec& spec) {
   return std::make_shared<CvPixelBufferPoolWrapper>(
-      spec.width, spec.height, spec.format, kMaxInactiveBufferAge);
-}
-
-GpuBuffer GpuBufferMultiPool::GetBufferFromSimplePool(
-    BufferSpec spec, GpuBufferMultiPool::SimplePool& pool) {
-  return GpuBuffer(pool.GetBuffer(flush_platform_caches_));
+      spec.width, spec.height, spec.format, kMaxInactiveBufferAge,
+      flush_platform_caches_);
 }
 
 #else
@@ -64,11 +60,6 @@ GpuBufferMultiPool::MakeSimplePool(const BufferSpec& spec) {
                                      kKeepCount);
 }
 
-GpuBuffer GpuBufferMultiPool::GetBufferFromSimplePool(
-    BufferSpec spec, GpuBufferMultiPool::SimplePool& pool) {
-  return GpuBuffer(pool.GetBuffer());
-}
-
 #endif  // MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
 
 std::shared_ptr<GpuBufferMultiPool::SimplePool> GpuBufferMultiPool::RequestPool(
@@ -102,6 +93,11 @@ GpuBuffer GpuBufferMultiPool::GetBuffer(int width, int height,
   }
 }
 
+GpuBuffer GpuBufferMultiPool::GetBufferFromSimplePool(
+    BufferSpec spec, GpuBufferMultiPool::SimplePool& pool) {
+  return GpuBuffer(pool.GetBuffer());
+}
+
 GpuBuffer GpuBufferMultiPool::GetBufferWithoutPool(const BufferSpec& spec) {
   return GpuBuffer(SimplePool::CreateBufferWithoutPool(spec.width, spec.height,
                                                        spec.format));