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
This commit is contained in:
Camillo Lugaresi 2022-11-18 19:41:05 -08:00 committed by Copybara-Service
parent bbd5da7971
commit eb8ef1ace0
3 changed files with 41 additions and 7 deletions

View File

@ -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) {
// Inform the GlTextureBuffer that we have finished accessing its
// contents, and create a consumer sync point.
DidRead(texture.gl_context()->CreateSyncToken());
};
GlTextureView::DetachFn detach =
[texbuf = shared_from_this()](GlTextureView& texture) {
// Inform the GlTextureBuffer that we have finished accessing its
// 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));
}

View File

@ -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

View File

@ -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