Metal: fix UBO data update undetected between draw calls.

This CL fixes an issue where an UBO was updated between two draw calls.
This would internally allocate a new metal buffer however ContextMtl
didn't know that hence it didn't rebind with new buffers.

Bug: b/338348430
Change-Id: I3d8ce22921811e49ae1b8016ae4a20d770f6f372
Reviewed-on: https://chromium-review.googlesource.com/c/angle/angle/+/5515858
Reviewed-by: Geoff Lang <geofflang@chromium.org>
Commit-Queue: Quyen Le <lehoangquyen@chromium.org>

3 files changed, 103 insertions, 19 deletions

diff --git a/src/libANGLE/renderer/metal/BufferMtl.h b/src/libANGLE/renderer/metal/BufferMtl.h
index 64f951815f..8162f10338 100644
--- a/src/libANGLE/renderer/metal/BufferMtl.h
+++ b/src/libANGLE/renderer/metal/BufferMtl.h
@@ -184,6 +184,11 @@ class BufferMtl : public BufferImpl, public BufferHolderMtl
         const mtl::BufferRef clientBuffer);
 
   private:
+    angle::Result allocateNewMetalBuffer(ContextMtl *contextMtl,
+                                         MTLStorageMode storageMode,
+                                         size_t size,
+                                         bool returnOldBufferImmediately);
+
     angle::Result setDataImpl(const gl::Context *context,
                               gl::BufferBinding target,
                               const void *data,
diff --git a/src/libANGLE/renderer/metal/BufferMtl.mm b/src/libANGLE/renderer/metal/BufferMtl.mm
index 08b4310997..60e1c179b0 100644
--- a/src/libANGLE/renderer/metal/BufferMtl.mm
+++ b/src/libANGLE/renderer/metal/BufferMtl.mm
@@ -504,6 +504,26 @@ const std::vector<IndexRange> BufferMtl::getRestartIndicesFromClientData(
     return restartIndices;
 }
 
+angle::Result BufferMtl::allocateNewMetalBuffer(ContextMtl *contextMtl,
+                                                MTLStorageMode storageMode,
+                                                size_t size,
+                                                bool returnOldBufferImmediately)
+{
+    mtl::BufferManager &bufferManager = contextMtl->getBufferManager();
+    if (returnOldBufferImmediately && mBuffer)
+    {
+        // Return the current buffer to the buffer manager
+        // It will not be re-used until it's no longer in use.
+        bufferManager.returnBuffer(contextMtl, mBuffer);
+        mBuffer = nullptr;
+    }
+    ANGLE_TRY(bufferManager.getBuffer(contextMtl, storageMode, size, mBuffer));
+
+    onStateChange(angle::SubjectMessage::InternalMemoryAllocationChanged);
+
+    return angle::Result::Continue;
+}
+
 angle::Result BufferMtl::setDataImpl(const gl::Context *context,
                                      gl::BufferBinding target,
                                      const void *data,
@@ -537,18 +557,9 @@ angle::Result BufferMtl::setDataImpl(const gl::Context *context,
     }
 
     // Re-create the buffer
-    mtl::BufferManager &bufferManager = contextMtl->getBufferManager();
-    if (mBuffer)
-    {
-        // Return the current buffer to the buffer manager
-        // It will not be re-used until it's no longer in use.
-        bufferManager.returnBuffer(contextMtl, mBuffer);
-        mBuffer = nullptr;
-    }
-
-    // Get a new buffer
     auto storageMode = mtl::Buffer::getStorageModeForUsage(contextMtl, usage);
-    ANGLE_TRY(bufferManager.getBuffer(contextMtl, storageMode, adjustedSize, mBuffer));
+    ANGLE_TRY(allocateNewMetalBuffer(contextMtl, storageMode, adjustedSize,
+                                     /*returnOldBufferImmediately=*/true));
 
 #ifndef NDEBUG
     ANGLE_MTL_OBJC_SCOPE
@@ -628,11 +639,11 @@ angle::Result BufferMtl::putDataInNewBufferAndStartUsingNewBuffer(ContextMtl *co
 {
     ASSERT(isOffsetAndSizeMetalBlitCompatible(offset, sizeToCopy));
 
-    mtl::BufferManager &bufferManager = contextMtl->getBufferManager();
-    mtl::BufferRef oldBuffer          = mBuffer;
-    auto storageMode                  = mtl::Buffer::getStorageModeForUsage(contextMtl, mUsage);
+    mtl::BufferRef oldBuffer = mBuffer;
+    auto storageMode         = mtl::Buffer::getStorageModeForUsage(contextMtl, mUsage);
 
-    ANGLE_TRY(bufferManager.getBuffer(contextMtl, storageMode, mGLSize, mBuffer));
+    ANGLE_TRY(allocateNewMetalBuffer(contextMtl, storageMode, mGLSize,
+                                     /*returnOldBufferImmediately=*/false));
     mBuffer->get().label = [NSString stringWithFormat:@"BufferMtl=%p(%lu)", this, ++mRevisionCount];
 
     uint8_t *ptr = mBuffer->mapWithOpt(contextMtl, false, true);
@@ -657,6 +668,7 @@ angle::Result BufferMtl::putDataInNewBufferAndStartUsingNewBuffer(ContextMtl *co
         }
     }
 
+    mtl::BufferManager &bufferManager = contextMtl->getBufferManager();
     bufferManager.returnBuffer(contextMtl, oldBuffer);
     return angle::Result::Continue;
 }
@@ -753,12 +765,11 @@ angle::Result BufferMtl::commitShadowCopy(ContextMtl *contextMtl)
 
 angle::Result BufferMtl::commitShadowCopy(ContextMtl *contextMtl, size_t size)
 {
-    mtl::BufferManager &bufferManager = contextMtl->getBufferManager();
-    auto storageMode                  = mtl::Buffer::getStorageModeForUsage(contextMtl, mUsage);
+    auto storageMode = mtl::Buffer::getStorageModeForUsage(contextMtl, mUsage);
 
-    bufferManager.returnBuffer(contextMtl, mBuffer);
     size_t bufferSize = (mGLSize == 0 ? mShadowCopy.size() : mGLSize);
-    ANGLE_TRY(bufferManager.getBuffer(contextMtl, storageMode, bufferSize, mBuffer));
+    ANGLE_TRY(allocateNewMetalBuffer(contextMtl, storageMode, bufferSize,
+                                     /*returnOldBufferImmediately=*/true));
 
     if (size)
     {
diff --git a/src/tests/gl_tests/UniformBufferTest.cpp b/src/tests/gl_tests/UniformBufferTest.cpp
index 1ea9d119bc..e99e047079 100644
--- a/src/tests/gl_tests/UniformBufferTest.cpp
+++ b/src/tests/gl_tests/UniformBufferTest.cpp
@@ -78,6 +78,74 @@ TEST_P(UniformBufferTest, Simple)
     EXPECT_PIXEL_NEAR(0, 0, 128, 191, 64, 255, 1);
 }
 
+// Test a scenario that draws then update UBO (using bufferData or bufferSubData or mapBuffer) then
+// draws with updated data.
+TEST_P(UniformBufferTest, DrawThenUpdateThenDraw)
+{
+    constexpr char kVS[] = R"(#version 300 es
+precision highp float;
+
+void main()
+{
+    vec2 position = vec2(float(gl_VertexID >> 1), float(gl_VertexID & 1));
+    position = 2.0 * position - 1.0;
+    gl_Position = vec4(position.x, position.y, 0.0, 1.0);
+})";
+
+    enum class BufferUpdateMethod
+    {
+        BUFFER_DATA,
+        BUFFER_SUB_DATA,
+        MAP_BUFFER,
+    };
+
+    ANGLE_GL_PROGRAM(program, kVS, mkFS);
+    GLint uniformBufferIndex = glGetUniformBlockIndex(program, "uni");
+    ASSERT_NE(uniformBufferIndex, -1);
+
+    for (BufferUpdateMethod method :
+         {BufferUpdateMethod::BUFFER_DATA, BufferUpdateMethod::BUFFER_SUB_DATA,
+          BufferUpdateMethod::MAP_BUFFER})
+    {
+        glClear(GL_COLOR_BUFFER_BIT);
+        float floatData1[4] = {0.25f, 0.75f, 0.125f, 1.0f};
+
+        GLBuffer uniformBuffer;
+        glBindBuffer(GL_UNIFORM_BUFFER, uniformBuffer);
+        glBufferData(GL_UNIFORM_BUFFER, sizeof(float) * 4, floatData1, GL_DYNAMIC_DRAW);
+
+        glBindBufferBase(GL_UNIFORM_BUFFER, 0, uniformBuffer);
+
+        glEnable(GL_BLEND);
+        glBlendFunc(GL_ONE, GL_ONE);
+
+        glUniformBlockBinding(program, uniformBufferIndex, 0);
+        glUseProgram(program);
+        glDrawArrays(GL_TRIANGLE_STRIP, 0, 4);
+
+        float floatData2[4] = {0.25f, 0.0f, 0.125f, 0.0f};
+        switch (method)
+        {
+            case BufferUpdateMethod::BUFFER_DATA:
+                glBufferData(GL_UNIFORM_BUFFER, sizeof(float) * 4, floatData2, GL_DYNAMIC_DRAW);
+                break;
+            case BufferUpdateMethod::BUFFER_SUB_DATA:
+                glBufferSubData(GL_UNIFORM_BUFFER, 0, sizeof(float) * 4, floatData2);
+                break;
+            case BufferUpdateMethod::MAP_BUFFER:
+                void *mappedBuffer =
+                    glMapBufferRange(GL_UNIFORM_BUFFER, 0, sizeof(float) * 4, GL_MAP_WRITE_BIT);
+                memcpy(mappedBuffer, floatData2, sizeof(floatData2));
+                glUnmapBuffer(GL_UNIFORM_BUFFER);
+                break;
+        }
+        glDrawArrays(GL_TRIANGLE_STRIP, 0, 4);
+
+        ASSERT_GL_NO_ERROR();
+        EXPECT_PIXEL_NEAR(0, 0, 128, 191, 64, 255, 1);
+    }
+}
+
 // Test that using a UBO with a non-zero offset and size actually works.
 // The first step of this test renders a color from a UBO with a zero offset.
 // The second step renders a color from a UBO with a non-zero offset.