Use shared_from_this in GlTextureBuffer::GetReadView, GetWriteView

This ensures that the callbacks in GlTextureView won't call an expired object, even if user code holds a GlTextureView after releasing the buffer.

Note that GlTextureBuffer is not always held by a shared_ptr, but it always is when GpuBuffer calls GetRead/WriteView on it. An alternative solution would have been to have GpuBuffer pass its shared_ptr<GlTextureBuffer> to the view method, which could have been implemented with some compile-time logic to detect whether the method expects such an argument. However, that doesn't seem necessary.

PiperOrigin-RevId: 489611843

mediapipe/gpu/gl_texture_buffer.cc

@@ -260,13 +260,18 @@ GlTextureView GlTextureBuffer::GetReadView(internal::types<GlTextureView>,
   auto gl_context = GlContext::GetCurrent();
   CHECK(gl_context);
   CHECK_EQ(plane, 0);
+  // Note that this method is only supposed to be called by GpuBuffer, which
+  // ensures this condition is satisfied.
+  DCHECK(!weak_from_this().expired())
+      << "GlTextureBuffer must be held in shared_ptr to get a GlTextureView";
   // Insert wait call to sync with the producer.
   WaitOnGpu();
-  GlTextureView::DetachFn detach = [this](GlTextureView& texture) {
+  GlTextureView::DetachFn detach =
-    // Inform the GlTextureBuffer that we have finished accessing its
+      [texbuf = shared_from_this()](GlTextureView& texture) {
-    // contents, and create a consumer sync point.
+        // Inform the GlTextureBuffer that we have finished accessing its
-    DidRead(texture.gl_context()->CreateSyncToken());
+        // contents, and create a consumer sync point.
-  };
+        texbuf->DidRead(texture.gl_context()->CreateSyncToken());
+      };
   return GlTextureView(gl_context.get(), target(), name(), width(), height(),
                        plane, std::move(detach), nullptr);
 }
@@ -276,12 +281,18 @@ GlTextureView GlTextureBuffer::GetWriteView(internal::types<GlTextureView>,
   auto gl_context = GlContext::GetCurrent();
   CHECK(gl_context);
   CHECK_EQ(plane, 0);
+  // Note that this method is only supposed to be called by GpuBuffer, which
+  // ensures this condition is satisfied.
+  DCHECK(!weak_from_this().expired())
+      << "GlTextureBuffer must be held in shared_ptr to get a GlTextureView";
   // Insert wait call to sync with the producer.
   WaitOnGpu();
   Reuse();  // TODO: the producer wait should probably be part of Reuse in the
             // case when there are no consumers.
   GlTextureView::DoneWritingFn done_writing =
-      [this](const GlTextureView& texture) { ViewDoneWriting(texture); };
+      [texbuf = shared_from_this()](const GlTextureView& texture) {
+        texbuf->ViewDoneWriting(texture);
+      };
   return GlTextureView(gl_context.get(), target(), name(), width(), height(),
                        plane, nullptr, std::move(done_writing));
 }

mediapipe/gpu/gl_texture_buffer.h

@@ -35,7 +35,8 @@ class GlCalculatorHelperImpl;
 // Implements a GPU memory buffer as an OpenGL texture. For internal use.
 class GlTextureBuffer
     : public internal::GpuBufferStorageImpl<
-          GlTextureBuffer, internal::ViewProvider<GlTextureView>> {
+          GlTextureBuffer, internal::ViewProvider<GlTextureView>>,
+      public std::enable_shared_from_this<GlTextureBuffer> {
  public:
   // This is called when the texture buffer is deleted. It is passed a sync
   // token created at that time on the GlContext. If the GlTextureBuffer has

mediapipe/gpu/gpu_buffer_test.cc

@@ -14,6 +14,8 @@
 
 #include "mediapipe/gpu/gpu_buffer.h"
 
+#include <utility>
+
 #include "mediapipe/framework/formats/image_format.pb.h"
 #include "mediapipe/framework/port/gmock.h"
 #include "mediapipe/framework/port/gtest.h"
@@ -206,5 +208,25 @@ TEST_F(GpuBufferTest, Overwrite) {
   }
 }
 
+TEST_F(GpuBufferTest, GlTextureViewRetainsWhatItNeeds) {
+  GpuBuffer buffer(300, 200, GpuBufferFormat::kBGRA32);
+  {
+    std::shared_ptr<ImageFrame> view = buffer.GetWriteView<ImageFrame>();
+    EXPECT_EQ(view->Width(), 300);
+    EXPECT_EQ(view->Height(), 200);
+    FillImageFrameRGBA(*view, 255, 0, 0, 255);
+  }
+
+  RunInGlContext([buffer = std::move(buffer)]() mutable {
+    // This is not a recommended pattern, but let's make sure that we don't
+    // crash if the buffer is released before the view. The view can hold
+    // callbacks into its underlying storage.
+    auto view = buffer.GetReadView<GlTextureView>(0);
+    buffer = nullptr;
+  });
+  // We're really checking that we haven't crashed.
+  EXPECT_TRUE(true);
+}
+
 }  // anonymous namespace
 }  // namespace mediapipe