Enables MultiPool and ReusablePool to pass on absl::Status returns or…

…iginating from object factory methods.

PiperOrigin-RevId: 604844100

mediapipe/framework/formats/BUILD

@@ -527,6 +527,7 @@ cc_library(
         "@com_google_absl//absl/log:absl_check",
         "@com_google_absl//absl/log:absl_log",
         "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/status",
         "@com_google_absl//absl/synchronization",
     ] + select({
         "//mediapipe/gpu:disable_gpu": [],

mediapipe/framework/formats/hardware_buffer_pool.h

@@ -5,7 +5,9 @@
 #include <utility>
 
 #include "absl/log/absl_check.h"
+#include "absl/status/statusor.h"
 #include "mediapipe/framework/formats/hardware_buffer.h"
+#include "mediapipe/framework/port/status_macros.h"
 #include "mediapipe/gpu/multi_pool.h"
 #include "mediapipe/gpu/reusable_pool.h"
 
@@ -31,12 +33,10 @@ class HardwareBufferSpecPool : public ReusablePool<HardwareBuffer> {
     return std::shared_ptr<HardwareBufferSpecPool>(
         new HardwareBufferSpecPool(spec, options));
   }
-  static std::unique_ptr<HardwareBuffer> CreateBufferWithoutPool(
-      const HardwareBufferSpec& spec) {
-    auto hardware_buffer = HardwareBuffer::Create(spec);
-    // TODO Don't crash - return absl::Status to caller.
-    ABSL_CHECK_OK(hardware_buffer);
-    return std::make_unique<HardwareBuffer>(std::move(*hardware_buffer));
+  static absl::StatusOr<std::unique_ptr<HardwareBuffer>>
+  CreateBufferWithoutPool(const HardwareBufferSpec& spec) {
+    MP_ASSIGN_OR_RETURN(auto hardware_buffer, HardwareBuffer::Create(spec));
+    return std::make_unique<HardwareBuffer>(std::move(hardware_buffer));
   }
   const HardwareBufferSpec& spec() const { return spec_; }
 
@@ -62,7 +62,8 @@ class HardwareBufferPool
       : MultiPool<internal::HardwareBufferSpecPool, HardwareBufferSpec,
                   std::shared_ptr<HardwareBuffer>>(options) {}
 
-  std::shared_ptr<HardwareBuffer> GetBuffer(const HardwareBufferSpec& spec) {
+  absl::StatusOr<std::shared_ptr<HardwareBuffer>> GetBuffer(
+      const HardwareBufferSpec& spec) {
     return Get(spec);
   }
 };

mediapipe/framework/formats/hardware_buffer_pool_test.cc

@@ -3,6 +3,7 @@
 #include <cstdint>
 
 #include "mediapipe/framework/formats/hardware_buffer.h"
+#include "mediapipe/framework/port/status_matchers.h"
 #include "mediapipe/gpu/multi_pool.h"
 #include "testing/base/public/gunit.h"
 
@@ -38,12 +39,16 @@ TEST(HardwareBufferPoolTest, ShouldPoolHardwareBuffer) {
   HardwareBuffer* hardware_buffer_ptr = nullptr;
   // First request instantiates new HardwareBuffer.
   {
-    auto hardware_buffer = hardware_buffer_pool.GetBuffer(hardware_buffer_spec);
+    MP_ASSERT_OK_AND_ASSIGN(
+        auto hardware_buffer,
+        hardware_buffer_pool.GetBuffer(hardware_buffer_spec));
     hardware_buffer_ptr = hardware_buffer.get();
   }
   // Second request returns same HardwareBuffer.
   {
-    auto hardware_buffer = hardware_buffer_pool.GetBuffer(hardware_buffer_spec);
+    MP_ASSERT_OK_AND_ASSIGN(
+        auto hardware_buffer,
+        hardware_buffer_pool.GetBuffer(hardware_buffer_spec));
     EXPECT_EQ(hardware_buffer.get(), hardware_buffer_ptr);
   }
 }
@@ -54,15 +59,17 @@ TEST(HardwareBufferPoolTest, ShouldReturnNewHardwareBuffer) {
   HardwareBuffer* hardware_buffer_ptr = nullptr;
   // First request instantiates new HardwareBuffer.
   {
-    auto hardware_buffer = hardware_buffer_pool.GetBuffer(
-        GetTestHardwareBufferSpec(/*size_bytes=*/123));
+    MP_ASSERT_OK_AND_ASSIGN(auto hardware_buffer,
+                            hardware_buffer_pool.GetBuffer(
+                                GetTestHardwareBufferSpec(/*size_bytes=*/123)));
     hardware_buffer_ptr = hardware_buffer.get();
     EXPECT_NE(hardware_buffer_ptr, nullptr);
   }
   // Second request with different size returns new HardwareBuffer.
   {
-    auto hardware_buffer = hardware_buffer_pool.GetBuffer(
-        GetTestHardwareBufferSpec(/*size_bytes=*/567));
+    MP_ASSERT_OK_AND_ASSIGN(auto hardware_buffer,
+                            hardware_buffer_pool.GetBuffer(
+                                GetTestHardwareBufferSpec(/*size_bytes=*/567)));
     EXPECT_NE(hardware_buffer.get(), hardware_buffer_ptr);
   }
 }

mediapipe/framework/formats/image_multi_pool.cc

@@ -104,7 +104,9 @@ ImageMultiPool::SimplePoolGpu ImageMultiPool::MakeSimplePoolGpu(
 
 Image ImageMultiPool::GetBufferFromSimplePool(
     IBufferSpec spec, const ImageMultiPool::SimplePoolGpu& pool) {
-  return Image(pool->GetBuffer());
+  auto buffer = pool->GetBuffer();
+  ABSL_CHECK_OK(buffer);
+  return Image(*std::move(buffer));
 }
 
 #endif  // MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER

mediapipe/framework/formats/tensor.cc

@@ -498,7 +498,7 @@ void Tensor::Invalidate() {
       mtl_resources_->device = nil;
     }
 #if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
-    // Don't need to wait for the resource to be deleted bacause if will be
+    // Don't need to wait for the resource to be deleted because if will be
     // released on last reference deletion inside the OpenGL driver.
     std::swap(cleanup_gl_tex, opengl_texture2d_);
     std::swap(cleanup_gl_fb, frame_buffer_);
@@ -530,7 +530,7 @@ void Tensor::Invalidate() {
     absl::MutexLock lock(&view_mutex_);
     ReleaseAhwbStuff();
 
-    // Don't need to wait for the resource to be deleted bacause if will be
+    // Don't need to wait for the resource to be deleted because if will be
     // released on last reference deletion inside the OpenGL driver.
 #if MEDIAPIPE_OPENGL_ES_VERSION >= MEDIAPIPE_OPENGL_ES_30
     std::swap(cleanup_gl_tex, opengl_texture2d_);
@@ -675,7 +675,7 @@ Tensor::CpuWriteView Tensor::GetCpuWriteView(
 void Tensor::AllocateCpuBuffer() const {
   if (!cpu_buffer_) {
 #ifdef MEDIAPIPE_TENSOR_USE_AHWB
-    if (use_ahwb_ && AllocateAHardwareBuffer()) return;
+    if (use_ahwb_ && AllocateAHardwareBuffer().ok()) return;
 #endif  // MEDIAPIPE_TENSOR_USE_AHWB
 #if MEDIAPIPE_METAL_ENABLED
     cpu_buffer_ = AllocateVirtualMemory(bytes());

mediapipe/framework/formats/tensor.h

@@ -27,6 +27,7 @@
 
 #include "absl/container/flat_hash_map.h"
 #include "absl/log/absl_check.h"
+#include "absl/status/status.h"
 #include "absl/synchronization/mutex.h"
 #include "mediapipe/framework/formats/tensor/internal.h"
 #include "mediapipe/framework/memory_manager.h"
@@ -406,7 +407,7 @@ class Tensor {
   // If the input parameter is 'true' then wait for the writing to be finished.
   mutable FinishingFunc ahwb_written_;
   mutable std::function<void()> release_callback_;
-  bool AllocateAHardwareBuffer(int size_alignment = 0) const;
+  absl::Status AllocateAHardwareBuffer(int size_alignment = 0) const;
   void CreateEglSyncAndFd() const;
 #endif  // MEDIAPIPE_TENSOR_USE_AHWB
   // Use Ahwb for other views: OpenGL / CPU buffer.

mediapipe/framework/formats/tensor_ahwb.cc

@@ -212,7 +212,7 @@ Tensor::AHardwareBufferView Tensor::GetAHardwareBufferReadView() const {
       << "Tensor conversion between OpenGL texture and AHardwareBuffer is not "
          "supported.";
   bool transfer = ahwb_ == nullptr;
-  ABSL_CHECK(AllocateAHardwareBuffer())
+  ABSL_CHECK_OK(AllocateAHardwareBuffer())
       << "AHardwareBuffer is not supported on the target system.";
   valid_ |= kValidAHardwareBuffer;
   if (transfer) {
@@ -250,7 +250,7 @@ void Tensor::CreateEglSyncAndFd() const {
 Tensor::AHardwareBufferView Tensor::GetAHardwareBufferWriteView(
     int size_alignment) const {
   auto lock(absl::make_unique<absl::MutexLock>(&view_mutex_));
-  ABSL_CHECK(AllocateAHardwareBuffer(size_alignment))
+  ABSL_CHECK_OK(AllocateAHardwareBuffer(size_alignment))
       << "AHardwareBuffer is not supported on the target system.";
   valid_ = kValidAHardwareBuffer;
   return {ahwb_.get(),
@@ -260,7 +260,7 @@ Tensor::AHardwareBufferView Tensor::GetAHardwareBufferWriteView(
           &release_callback_,  std::move(lock)};
 }
 
-bool Tensor::AllocateAHardwareBuffer(int size_alignment) const {
+absl::Status Tensor::AllocateAHardwareBuffer(int size_alignment) const {
   // Mark current tracking key as Ahwb-use.
   if (auto it = ahwb_usage_track_.find(ahwb_tracking_key_);
       it != ahwb_usage_track_.end()) {
@@ -287,23 +287,18 @@ bool Tensor::AllocateAHardwareBuffer(int size_alignment) const {
                  HardwareBufferSpec::AHARDWAREBUFFER_USAGE_CPU_READ_OFTEN |
                  HardwareBufferSpec::AHARDWAREBUFFER_USAGE_GPU_DATA_BUFFER;
     if (hardware_buffer_pool_ == nullptr) {
-      auto new_ahwb = HardwareBuffer::Create(spec);
-      if (!new_ahwb.ok()) {
-        ABSL_LOG(ERROR) << "Allocation of NDK Hardware Buffer failed: "
-                        << new_ahwb.status();
-        return false;
-      }
-      ahwb_ = std::make_shared<HardwareBuffer>(std::move(*new_ahwb));
+      MP_ASSIGN_OR_RETURN(auto new_ahwb, HardwareBuffer::Create(spec));
+      ahwb_ = std::make_shared<HardwareBuffer>(std::move(new_ahwb));
     } else {
-      ahwb_ = hardware_buffer_pool_->GetBuffer(spec);
+      MP_ASSIGN_OR_RETURN(ahwb_, hardware_buffer_pool_->GetBuffer(spec));
     }
   }
-  return true;
+  return absl::OkStatus();
 }
 
 bool Tensor::AllocateAhwbMapToSsbo() const {
   if (__builtin_available(android 26, *)) {
-    if (AllocateAHardwareBuffer()) {
+    if (AllocateAHardwareBuffer().ok()) {
       if (MapAHardwareBufferToGlBuffer(ahwb_->GetAHardwareBuffer(), bytes())
               .ok()) {
         glBindBuffer(GL_SHADER_STORAGE_BUFFER, 0);

mediapipe/gpu/BUILD

@@ -441,6 +441,7 @@ cc_library(
         "//mediapipe/objc:CFHolder",
         "//mediapipe/objc:util",
         "@com_google_absl//absl/log:absl_check",
+        "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/synchronization",
     ],
 )
@@ -674,6 +675,7 @@ cc_library(
         "@com_google_absl//absl/log:absl_check",
         "@com_google_absl//absl/log:absl_log",
         "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
     ] + select({
         "//conditions:default": [],
         "//mediapipe:apple": [
@@ -689,16 +691,10 @@ cc_library(
     hdrs = ["gl_texture_buffer_pool.h"],
     visibility = ["//visibility:public"],
     deps = [
-        ":gl_base",
         ":gl_texture_buffer",
-        ":gpu_buffer",
-        ":gpu_shared_data_header",
         ":multi_pool",
         ":reusable_pool",
-        "//mediapipe/framework:calculator_context",
-        "//mediapipe/framework:calculator_node",
-        "//mediapipe/framework/port:logging",
-        "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/synchronization",
     ],
 )
@@ -708,15 +704,23 @@ cc_library(
     hdrs = ["reusable_pool.h"],
     deps = [
         ":multi_pool",
+        "//mediapipe/framework/port:ret_check",
+        "//mediapipe/framework/port:status",
+        "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/functional:any_invocable",
+        "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/synchronization",
     ],
 )
 
 cc_library(
     name = "multi_pool",
     hdrs = ["multi_pool.h"],
-    deps = ["//mediapipe/util:resource_cache"],
+    deps = [
+        "//mediapipe/framework/port:status",
+        "//mediapipe/util:resource_cache",
+        "@com_google_absl//absl/status:statusor",
+    ],
 )
 
 cc_library(
@@ -753,6 +757,7 @@ cc_library(
         "//mediapipe/util:resource_cache",
         "@com_google_absl//absl/hash",
         "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/synchronization",
     ] + select({
         "//conditions:default": [
@@ -868,6 +873,7 @@ cc_library(
         "@com_google_absl//absl/log:absl_check",
         "@com_google_absl//absl/log:absl_log",
         "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/synchronization",
     ] + select({
         "//conditions:default": [

mediapipe/gpu/MPPMetalHelper.cc

@@ -211,14 +211,19 @@ class MetalHelperLegacySupport {
 
 - (mediapipe::GpuBuffer)mediapipeGpuBufferWithWidth:(int)width
                                              height:(int)height {
-  return _gpuResources->gpu_buffer_pool().GetBuffer(width, height);
+  auto gpu_buffer = _gpuResources->gpu_buffer_pool().GetBuffer(width, height);
+  ABSL_CHECK_OK(gpu_buffer);
+  return *gpu_buffer;
 }
 
 - (mediapipe::GpuBuffer)mediapipeGpuBufferWithWidth:(int)width
                                              height:(int)height
                                              format:(mediapipe::GpuBufferFormat)
                                                         format {
-  return _gpuResources->gpu_buffer_pool().GetBuffer(width, height, format);
+  auto gpu_buffer =
+      _gpuResources->gpu_buffer_pool().GetBuffer(width, height, format);
+  ABSL_CHECK_OK(gpu_buffer);
+  return *gpu_buffer;
 }
 
 - (id<MTLLibrary>)newLibraryWithResourceName:(NSString*)name

mediapipe/gpu/cv_pixel_buffer_pool_wrapper.cc

@@ -18,6 +18,7 @@
 
 #include "CoreFoundation/CFBase.h"
 #include "absl/log/absl_check.h"
+#include "absl/status/statusor.h"
 #include "mediapipe/framework/port/logging.h"
 #include "mediapipe/objc/CFHolder.h"
 #include "mediapipe/objc/util.h"
@@ -35,7 +36,8 @@ CvPixelBufferPoolWrapper::CvPixelBufferPoolWrapper(
   texture_caches_ = texture_caches;
 }
 
-CFHolder<CVPixelBufferRef> CvPixelBufferPoolWrapper::GetBuffer() {
+absl::StatusOr<CFHolder<CVPixelBufferRef>>
+CvPixelBufferPoolWrapper::GetBuffer() {
   CVPixelBufferRef buffer;
   int threshold = 1;
   NSMutableDictionary* auxAttributes =
@@ -71,7 +73,8 @@ std::string CvPixelBufferPoolWrapper::GetDebugString() const {
 
 void CvPixelBufferPoolWrapper::Flush() { CVPixelBufferPoolFlush(*pool_, 0); }
 
-CFHolder<CVPixelBufferRef> CvPixelBufferPoolWrapper::CreateBufferWithoutPool(
+absl::StatusOr<CFHolder<CVPixelBufferRef>>
+CvPixelBufferPoolWrapper::CreateBufferWithoutPool(
     const internal::GpuBufferSpec& spec) {
   OSType cv_format = CVPixelFormatForGpuBufferFormat(spec.format);
   ABSL_CHECK_NE(cv_format, -1) << "unsupported pixel format";

mediapipe/gpu/cv_pixel_buffer_pool_wrapper.h

@@ -23,6 +23,7 @@
 #define MEDIAPIPE_GPU_CV_PIXEL_BUFFER_POOL_WRAPPER_H_
 
 #include "CoreFoundation/CFBase.h"
+#include "absl/status/statusor.h"
 #include "mediapipe/gpu/cv_texture_cache_manager.h"
 #include "mediapipe/gpu/gpu_buffer_format.h"
 #include "mediapipe/gpu/multi_pool.h"
@@ -45,14 +46,14 @@ class CvPixelBufferPoolWrapper {
         texture_caches);
   }
 
-  CFHolder<CVPixelBufferRef> GetBuffer();
+  absl::StatusOr<CFHolder<CVPixelBufferRef>> GetBuffer();
 
   int GetBufferCount() const { return count_; }
   std::string GetDebugString() const;
 
   void Flush();
 
-  static CFHolder<CVPixelBufferRef> CreateBufferWithoutPool(
+  static absl::StatusOr<CFHolder<CVPixelBufferRef>> CreateBufferWithoutPool(
       const internal::GpuBufferSpec& spec);
 
  private:

mediapipe/gpu/gl_calculator_helper.cc

@@ -16,6 +16,7 @@
 
 #include "absl/log/absl_check.h"
 #include "absl/log/absl_log.h"
+#include "absl/status/statusor.h"
 #include "mediapipe/framework/formats/image.h"
 #include "mediapipe/framework/formats/image_frame.h"
 #include "mediapipe/framework/legacy_calculator_support.h"
@@ -192,7 +193,7 @@ GpuBuffer GlCalculatorHelper::GpuBufferCopyingImageFrame(
   auto maybe_buffer = CreateCVPixelBufferCopyingImageFrame(image_frame);
   // Converts absl::StatusOr to absl::Status since ABSL_CHECK_OK() currently
   // only deals with absl::Status in MediaPipe OSS.
-  ABSL_CHECK_OK(maybe_buffer.status());
+  ABSL_CHECK_OK(maybe_buffer);
   return GpuBuffer(std::move(maybe_buffer).value());
 #else
   return GpuBuffer(GlTextureBuffer::Create(image_frame));
@@ -213,9 +214,10 @@ GlTexture GlCalculatorHelper::CreateDestinationTexture(int width, int height,
     CreateFramebuffer();
   }
 
-  GpuBuffer gpu_buffer =
+  auto gpu_buffer =
       gpu_resources_->gpu_buffer_pool().GetBuffer(width, height, format);
-  return MapGpuBuffer(gpu_buffer, gpu_buffer.GetWriteView<GlTextureView>(0));
+  ABSL_CHECK_OK(gpu_buffer);
+  return MapGpuBuffer(*gpu_buffer, gpu_buffer->GetWriteView<GlTextureView>(0));
 }
 
 GlTexture GlCalculatorHelper::CreateDestinationTexture(