Allow reentrant buffer mapping operations in the mapping callback.

WebGPU clears [pending_map] before resolving or rejecting promises,
which means that the callback for .then or .catch is allowed to unmap,
remap, or do anything with the buffer assuming previous mapping
operations are entirely done the buffer.

Mimic this in Dawn by finishing all operations related to mapping
before calling callbacks they might trigger.

bug: dawn:1619
Change-Id: I9e5b82789a68f28a496a54c31bf9fe0ffde23ccf
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/116220
Kokoro: Kokoro <noreply+kokoro@google.com>
Reviewed-by: Corentin Wallez <cwallez@chromium.org>
Commit-Queue: Takahiro <hogehoge@gachapin.jp>
Reviewed-by: Austin Eng <enga@chromium.org>

src/dawn/native/Buffer.cpp

@@ -134,6 +134,46 @@ MaybeError ValidateBufferDescriptor(DeviceBase* device, const BufferDescriptor*
     return {};
 }
 
+// BufferBase::PendingMappingCallback
+
+BufferBase::PendingMappingCallback::PendingMappingCallback()
+    : callback(nullptr), userdata(nullptr) {}
+
+// Ensure to call the callback.
+BufferBase::PendingMappingCallback::~PendingMappingCallback() {
+    ASSERT(callback == nullptr);
+    ASSERT(userdata == nullptr);
+}
+
+BufferBase::PendingMappingCallback::PendingMappingCallback(
+    BufferBase::PendingMappingCallback&& other) {
+    this->callback = std::move(other.callback);
+    this->userdata = std::move(other.userdata);
+    this->status = other.status;
+    other.callback = nullptr;
+    other.userdata = nullptr;
+}
+
+BufferBase::PendingMappingCallback& BufferBase::PendingMappingCallback::operator=(
+    PendingMappingCallback&& other) {
+    if (&other != this) {
+        this->callback = std::move(other.callback);
+        this->userdata = std::move(other.userdata);
+        this->status = other.status;
+        other.callback = nullptr;
+        other.userdata = nullptr;
+    }
+    return *this;
+}
+
+void BufferBase::PendingMappingCallback::Call() {
+    if (callback != nullptr) {
+        callback(status, userdata);
+        callback = nullptr;
+        userdata = nullptr;
+    }
+}
+
 // Buffer
 
 BufferBase::BufferBase(DeviceBase* device, const BufferDescriptor* descriptor)
@@ -192,16 +232,20 @@ BufferBase::~BufferBase() {
 }
 
 void BufferBase::DestroyImpl() {
+    PendingMappingCallback toCall;
+
     if (mState == BufferState::Mapped || mState == BufferState::PendingMap) {
-        UnmapInternal(WGPUBufferMapAsyncStatus_DestroyedBeforeCallback);
+        toCall = UnmapInternal(WGPUBufferMapAsyncStatus_DestroyedBeforeCallback);
     } else if (mState == BufferState::MappedAtCreation) {
         if (mStagingBuffer != nullptr) {
             mStagingBuffer.reset();
         } else if (mSize != 0) {
-            UnmapInternal(WGPUBufferMapAsyncStatus_DestroyedBeforeCallback);
+            toCall = UnmapInternal(WGPUBufferMapAsyncStatus_DestroyedBeforeCallback);
         }
     }
+
     mState = BufferState::Destroyed;
+    toCall.Call();
 }
 
 // static
@@ -330,22 +374,31 @@ MaybeError BufferBase::ValidateCanUseOnQueueNow() const {
     UNREACHABLE();
 }
 
-void BufferBase::CallMapCallback(MapRequestID mapID, WGPUBufferMapAsyncStatus status) {
+// Store the callback to be called in an intermediate struct that bubbles up the call stack
+// and is called by the top most function at the very end. It helps to make sure that
+// all code paths ensure that nothing happens after the callback.
+BufferBase::PendingMappingCallback BufferBase::WillCallMappingCallback(
+    MapRequestID mapID,
+    WGPUBufferMapAsyncStatus status) {
     ASSERT(!IsError());
-    if (mMapCallback != nullptr && mapID == mLastMapID) {
+    PendingMappingCallback toCall;
-        // Tag the callback as fired before firing it, otherwise it could fire a second time if
-        // for example buffer.Unmap() is called inside the application-provided callback.
-        WGPUBufferMapCallback callback = mMapCallback;
-        mMapCallback = nullptr;
 
+    if (mMapCallback != nullptr && mapID == mLastMapID) {
+        toCall.callback = std::move(mMapCallback);
+        toCall.userdata = std::move(mMapUserdata);
         if (GetDevice()->IsLost()) {
-            callback(WGPUBufferMapAsyncStatus_DeviceLost, mMapUserdata);
+            toCall.status = WGPUBufferMapAsyncStatus_DeviceLost;
         } else {
-            callback(status, mMapUserdata);
+            toCall.status = status;
         }
 
-        mMapUserdata = 0;
+        // Tag the callback as fired before firing it, otherwise it could fire a second time if
+        // for example buffer.Unmap() is called inside the application-provided callback.
+        mMapCallback = nullptr;
+        mMapUserdata = nullptr;
     }
+
+    return toCall;
 }
 
 void BufferBase::APIMapAsync(wgpu::MapMode mode,
@@ -389,7 +442,7 @@ void BufferBase::APIMapAsync(wgpu::MapMode mode,
     mState = BufferState::PendingMap;
 
     if (GetDevice()->ConsumedError(MapAsyncImpl(mode, offset, size))) {
-        CallMapCallback(mLastMapID, WGPUBufferMapAsyncStatus_DeviceLost);
+        WillCallMappingCallback(mLastMapID, WGPUBufferMapAsyncStatus_DeviceLost).Call();
         return;
     }
     std::unique_ptr<MapRequestTask> request =
@@ -458,12 +511,15 @@ void BufferBase::Unmap() {
     if (mState == BufferState::Destroyed) {
         return;
     }
-    UnmapInternal(WGPUBufferMapAsyncStatus_UnmappedBeforeCallback);
+    UnmapInternal(WGPUBufferMapAsyncStatus_UnmappedBeforeCallback).Call();
 }
 
-void BufferBase::UnmapInternal(WGPUBufferMapAsyncStatus callbackStatus) {
+BufferBase::PendingMappingCallback BufferBase::UnmapInternal(
+    WGPUBufferMapAsyncStatus callbackStatus) {
+    PendingMappingCallback toCall;
+
     if (mState == BufferState::PendingMap) {
-        CallMapCallback(mLastMapID, callbackStatus);
+        toCall = WillCallMappingCallback(mLastMapID, callbackStatus);
         UnmapImpl();
     } else if (mState == BufferState::Mapped) {
         UnmapImpl();
@@ -476,6 +532,7 @@ void BufferBase::UnmapInternal(WGPUBufferMapAsyncStatus callbackStatus) {
     }
 
     mState = BufferState::Unmapped;
+    return toCall;
 }
 
 MaybeError BufferBase::ValidateMapAsync(wgpu::MapMode mode,
@@ -577,11 +634,12 @@ MaybeError BufferBase::ValidateUnmap() const {
 }
 
 void BufferBase::OnMapRequestCompleted(MapRequestID mapID, WGPUBufferMapAsyncStatus status) {
+    PendingMappingCallback toCall = WillCallMappingCallback(mapID, status);
     if (mapID == mLastMapID && status == WGPUBufferMapAsyncStatus_Success &&
         mState == BufferState::PendingMap) {
         mState = BufferState::Mapped;
     }
-    CallMapCallback(mapID, status);
+    toCall.Call();
 }
 
 bool BufferBase::NeedsInitialization() const {

src/dawn/native/Buffer.h

@@ -17,6 +17,8 @@
 
 #include <memory>
 
+#include "dawn/common/NonCopyable.h"
+
 #include "dawn/native/Error.h"
 #include "dawn/native/Forward.h"
 #include "dawn/native/IntegerTypes.h"
@@ -104,6 +106,25 @@ class BufferBase : public ApiObjectBase {
     uint64_t mAllocatedSize = 0;
 
   private:
+    // A helper structure to enforce that the mapAsync callback is called only at the very end of
+    // methods that might trigger callbacks. Non-copyable but movable for the assertion in the
+    // destructor to ensure not to forget to call the callback
+    struct [[nodiscard]] PendingMappingCallback : public NonCopyable {
+        WGPUBufferMapCallback callback;
+        void* userdata;
+        WGPUBufferMapAsyncStatus status;
+
+        PendingMappingCallback();
+        ~PendingMappingCallback();
+
+        PendingMappingCallback(PendingMappingCallback&& other);
+        PendingMappingCallback& operator=(PendingMappingCallback&& other);
+
+        void Call();
+    };
+    PendingMappingCallback WillCallMappingCallback(MapRequestID mapID,
+                                                   WGPUBufferMapAsyncStatus status);
+
     virtual MaybeError MapAtCreationImpl() = 0;
     virtual MaybeError MapAsyncImpl(wgpu::MapMode mode, size_t offset, size_t size) = 0;
     virtual void UnmapImpl() = 0;
@@ -111,7 +132,6 @@ class BufferBase : public ApiObjectBase {
 
     virtual bool IsCPUWritableAtCreation() const = 0;
     MaybeError CopyFromStagingBuffer();
-    void CallMapCallback(MapRequestID mapID, WGPUBufferMapAsyncStatus status);
 
     MaybeError ValidateMapAsync(wgpu::MapMode mode,
                                 size_t offset,
@@ -119,7 +139,7 @@ class BufferBase : public ApiObjectBase {
                                 WGPUBufferMapAsyncStatus* status) const;
     MaybeError ValidateUnmap() const;
     bool CanGetMappedRange(bool writable, size_t offset, size_t size) const;
-    void UnmapInternal(WGPUBufferMapAsyncStatus callbackStatus);
+    PendingMappingCallback UnmapInternal(WGPUBufferMapAsyncStatus callbackStatus);
 
     uint64_t mSize = 0;
     wgpu::BufferUsage mUsage = wgpu::BufferUsage::None;
@@ -129,7 +149,7 @@ class BufferBase : public ApiObjectBase {
     std::unique_ptr<StagingBufferBase> mStagingBuffer;
 
     WGPUBufferMapCallback mMapCallback = nullptr;
-    void* mMapUserdata = 0;
+    void* mMapUserdata = nullptr;
     MapRequestID mLastMapID = MapRequestID(0);
     wgpu::MapMode mMapMode = wgpu::MapMode::None;
     size_t mMapOffset = 0;

src/dawn/tests/unittests/validation/BufferValidationTests.cpp

@@ -523,6 +523,89 @@ TEST_F(BufferValidationTest, MapAsync_DestroyCalledInCallback) {
     }
 }
 
+// Test MapAsync call in MapAsync success callback
+// This test is disabled now because there seems to be a reeantrancy bug in the
+// FlushWire call. See https://dawn-review.googlesource.com/c/dawn/+/116220 for the details.
+TEST_F(BufferValidationTest, DISABLED_MapAsync_MapAsyncInMapAsyncSuccessCallback) {
+    // Test MapAsync call in MapAsync validation success callback
+    {
+        wgpu::Buffer buf = CreateMapReadBuffer(4);
+
+        EXPECT_CALL(*mockBufferMapAsyncCallback, Call(WGPUBufferMapAsyncStatus_Success, _))
+            .WillOnce(InvokeWithoutArgs([&]() {
+                EXPECT_CALL(*mockBufferMapAsyncCallback, Call(WGPUBufferMapAsyncStatus_Error, _));
+                // Should cause validation error because of already mapped buffer
+                ASSERT_DEVICE_ERROR(
+                    buf.MapAsync(wgpu::MapMode::Read, 0, 4, ToMockBufferMapAsyncCallback, nullptr));
+            }));
+
+        buf.MapAsync(wgpu::MapMode::Read, 0, 4, ToMockBufferMapAsyncCallback, nullptr);
+        WaitForAllOperations(device);
+        // we need another wire flush to make the MapAsync in the callback to the server
+        WaitForAllOperations(device);
+    }
+}
+
+// Test MapAsync call in MapAsync rejection callback
+TEST_F(BufferValidationTest, MapAsync_MapAsyncInMapAsyncRejectionCallback) {
+    // Test MapAsync call in MapAsync validation error callback
+    {
+        wgpu::Buffer buf = CreateMapReadBuffer(4);
+
+        EXPECT_CALL(*mockBufferMapAsyncCallback, Call(WGPUBufferMapAsyncStatus_Error, _))
+            .WillOnce(InvokeWithoutArgs([&]() {
+                // Retry with valid parameter and it should succeed
+                EXPECT_CALL(*mockBufferMapAsyncCallback, Call(WGPUBufferMapAsyncStatus_Success, _));
+                buf.MapAsync(wgpu::MapMode::Read, 0, 4, ToMockBufferMapAsyncCallback, nullptr);
+            }));
+
+        // Write map mode on read buffer is invalid and it should reject with validation error
+        ASSERT_DEVICE_ERROR(
+            buf.MapAsync(wgpu::MapMode::Write, 0, 4, ToMockBufferMapAsyncCallback, nullptr));
+
+        WaitForAllOperations(device);
+        // we need another wire flush to make the MapAsync in the callback to the server
+        WaitForAllOperations(device);
+    }
+
+    // Test MapAsync call in MapAsync Unmapped before callback callback
+    {
+        wgpu::Buffer buf = CreateMapReadBuffer(4);
+
+        EXPECT_CALL(*mockBufferMapAsyncCallback,
+                    Call(WGPUBufferMapAsyncStatus_UnmappedBeforeCallback, _))
+            .WillOnce(InvokeWithoutArgs([&]() {
+                // MapAsync call on unmapped buffer should be valid
+                EXPECT_CALL(*mockBufferMapAsyncCallback, Call(WGPUBufferMapAsyncStatus_Success, _));
+                buf.MapAsync(wgpu::MapMode::Read, 0, 4, ToMockBufferMapAsyncCallback, nullptr);
+            }));
+
+        buf.MapAsync(wgpu::MapMode::Read, 0, 4, ToMockBufferMapAsyncCallback, nullptr);
+        buf.Unmap();
+        WaitForAllOperations(device);
+        WaitForAllOperations(device);
+    }
+
+    // Test MapAsync call in MapAsync Destroyed before callback callback
+    {
+        wgpu::Buffer buf = CreateMapReadBuffer(4);
+
+        EXPECT_CALL(*mockBufferMapAsyncCallback,
+                    Call(WGPUBufferMapAsyncStatus_DestroyedBeforeCallback, _))
+            .WillOnce(InvokeWithoutArgs([&]() {
+                // MapAsync call on destroyed buffer should be invalid
+                EXPECT_CALL(*mockBufferMapAsyncCallback, Call(WGPUBufferMapAsyncStatus_Error, _));
+                ASSERT_DEVICE_ERROR(
+                    buf.MapAsync(wgpu::MapMode::Read, 0, 4, ToMockBufferMapAsyncCallback, nullptr));
+            }));
+
+        buf.MapAsync(wgpu::MapMode::Read, 0, 4, ToMockBufferMapAsyncCallback, nullptr);
+        buf.Destroy();
+        WaitForAllOperations(device);
+        WaitForAllOperations(device);
+    }
+}
+
 // Test the success case for mappedAtCreation
 TEST_F(BufferValidationTest, MappedAtCreationSuccess) {
     BufferMappedAtCreation(4, wgpu::BufferUsage::MapWrite);

src/dawn/tests/unittests/wire/WireBufferMappingTests.cpp

@@ -859,9 +859,8 @@ TEST_F(WireBufferMappingTests, MapInsideCallbackBeforeDisconnect) {
 
     FlushClient();
 
-    EXPECT_CALL(*mockBufferMapCallback, Call(WGPUBufferMapAsyncStatus_Error, this))
+    EXPECT_CALL(*mockBufferMapCallback, Call(WGPUBufferMapAsyncStatus_DeviceLost, this))
-        .Times(testData.numRequests);
+        .Times(testData.numRequests + 1);
-    EXPECT_CALL(*mockBufferMapCallback, Call(WGPUBufferMapAsyncStatus_DeviceLost, this)).Times(1);
     GetWireClient()->Disconnect();
 }
 
@@ -879,8 +878,9 @@ TEST_F(WireBufferMappingWriteTests, MapInsideCallbackBeforeDestruction) {
 
     FlushClient();
 
+    // Maybe this should be assert errors, see dawn:1624
     EXPECT_CALL(*mockBufferMapCallback, Call(WGPUBufferMapAsyncStatus_Error, this))
-        .Times(testData.numRequests);
+        .Times(testData.numRequests - 1);
     EXPECT_CALL(*mockBufferMapCallback,
                 Call(WGPUBufferMapAsyncStatus_DestroyedBeforeCallback, this))
         .Times(1);

src/dawn/wire/client/Buffer.cpp

@@ -159,8 +159,8 @@ Buffer::Buffer(const ObjectBaseParams& params,
       mDeviceIsAlive(device->GetAliveWeakPtr()) {}
 
 Buffer::~Buffer() {
-    InvokeAndClearCallback(WGPUBufferMapAsyncStatus_DestroyedBeforeCallback);
     FreeMappedData();
+    InvokeAndClearCallback(WGPUBufferMapAsyncStatus_DestroyedBeforeCallback);
 }
 
 void Buffer::CancelCallbacksForDisconnect() {
@@ -172,10 +172,10 @@ void Buffer::InvokeAndClearCallback(WGPUBufferMapAsyncStatus status) {
     void* userdata = mRequest.userdata;
     mRequest.callback = nullptr;
     mRequest.userdata = nullptr;
+    mPendingMap = false;
     if (callback != nullptr) {
         callback(status, userdata);
     }
-    mPendingMap = false;
 }
 
 void Buffer::MapAsync(WGPUMapModeFlags mode,
@@ -358,11 +358,11 @@ void Buffer::Unmap() {
     mMapOffset = 0;
     mMapSize = 0;
 
-    InvokeAndClearCallback(WGPUBufferMapAsyncStatus_UnmappedBeforeCallback);
-
     BufferUnmapCmd cmd;
     cmd.self = ToAPI(this);
     client->SerializeCommand(cmd);
+
+    InvokeAndClearCallback(WGPUBufferMapAsyncStatus_UnmappedBeforeCallback);
 }
 
 void Buffer::Destroy() {
@@ -372,11 +372,11 @@ void Buffer::Destroy() {
     FreeMappedData();
     mMapState = MapState::Unmapped;
 
-    InvokeAndClearCallback(WGPUBufferMapAsyncStatus_DestroyedBeforeCallback);
-
     BufferDestroyCmd cmd;
     cmd.self = ToAPI(this);
     client->SerializeCommand(cmd);
+
+    InvokeAndClearCallback(WGPUBufferMapAsyncStatus_DestroyedBeforeCallback);
 }
 
 WGPUBufferUsage Buffer::GetUsage() const {