[WGPUFuture] Update buffer mapping in wire to keep reference to Buffer.

- Also updates buffer mapping on the wire to be more consistent with
  other async future functions.
- Note that buffer release no longer causes map async callbacks to
  trigger since the event holds a reference to the buffer now.

Bug: dawn:2061
Change-Id: Ifd0d0588f61ca531f0e75d1dd4082a3dc3e257a4
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/168800
Commit-Queue: Loko Kung <lokokung@google.com>
Reviewed-by: Kai Ninomiya <kainino@chromium.org>
Reviewed-by: Corentin Wallez <cwallez@chromium.org>
Kokoro: Kokoro <noreply+kokoro@google.com>
diff --git a/src/dawn/dawn_wire.json b/src/dawn/dawn_wire.json
index f37d6ac..5e92ad7 100644
--- a/src/dawn/dawn_wire.json
+++ b/src/dawn/dawn_wire.json
@@ -116,10 +116,9 @@
     },
     "return commands": {
         "buffer map async callback": [
-            { "name": "buffer", "type": "ObjectHandle", "handle_type": "buffer", "_comment": "TODO(dawn:2061) Remove this field once mapping is updated." },
             { "name": "event manager", "type": "ObjectHandle" },
             { "name": "future", "type": "future" },
-            { "name": "status", "type": "uint32_t" },
+            { "name": "status", "type": "buffer map async status" },
             { "name": "read data update info length", "type": "uint64_t" },
             { "name": "read data update info", "type": "uint8_t", "annotation": "const*", "length": "read data update info length", "skip_serialize": true }
         ],
diff --git a/src/dawn/tests/unittests/wire/WireBufferMappingTests.cpp b/src/dawn/tests/unittests/wire/WireBufferMappingTests.cpp
index c75ea3a..d72f166 100644
--- a/src/dawn/tests/unittests/wire/WireBufferMappingTests.cpp
+++ b/src/dawn/tests/unittests/wire/WireBufferMappingTests.cpp
@@ -303,22 +303,6 @@
         WGPUBufferMapAsyncStatus_DestroyedBeforeCallback);
 }
 
-// Check the map callback is called with "DestroyedBeforeCallback" when the map request would have
-// worked, but Release() was called on the last ref.
-TEST_P(WireBufferMappingTests, ReleaseCalledTooEarly) {
-    TestEarlyMapCancelled(
-        &wgpuBufferRelease, [&]() { EXPECT_CALL(api, BufferRelease(apiBuffer)); },
-        WGPUBufferMapAsyncStatus_DestroyedBeforeCallback);
-}
-
-// Check that if Release() was called early client-side on the last ref, we disregard server-side
-// validation errors.
-TEST_P(WireBufferMappingTests, ReleaseCalledTooEarlyServerSideError) {
-    TestEarlyMapErrorCancelled(
-        &wgpuBufferRelease, [&]() { EXPECT_CALL(api, BufferRelease(apiBuffer)); },
-        WGPUBufferMapAsyncStatus_DestroyedBeforeCallback);
-}
-
 // Test that the callback isn't fired twice when Unmap() is called inside the callback.
 TEST_P(WireBufferMappingTests, UnmapInsideMapCallback) {
     TestCancelInCallback(&wgpuBufferUnmap, [&]() { EXPECT_CALL(api, BufferUnmap(apiBuffer)); });
@@ -831,52 +815,6 @@
     }
 }
 
-#if defined(DAWN_ENABLE_ASSERTS)
-// Test that request inside user callbacks after object release is called.
-TEST_P(WireBufferMappingTests, MapInsideCallbackAfterRelease) {
-    WGPUMapMode mapMode = GetMapMode();
-    SetupBuffer(mapMode);
-    BufferMapAsync(buffer, mapMode, 0, kBufferSize, nullptr);
-
-    uint32_t bufferContent = 31337;
-    EXPECT_CALL(api, OnBufferMapAsync(apiBuffer, mapMode, 0, kBufferSize, _, _))
-        .WillOnce(InvokeWithoutArgs(
-            [&] { api.CallBufferMapAsyncCallback(apiBuffer, WGPUBufferMapAsyncStatus_Success); }));
-    ExpectMappedRangeCall(kBufferSize, &bufferContent);
-    EXPECT_CALL(api, BufferRelease(apiBuffer));
-
-    // By releasing the buffer the refcount reaches zero and pending map async
-    // should fail with destroyed before callback status.
-    if (IsSpontaneous()) {
-        // In spontaneous mode, the callback gets called as a part of the release.
-        ExpectWireCallbacksWhen([&](auto& mockCb) {
-            EXPECT_CALL(mockCb, Call(WGPUBufferMapAsyncStatus_DestroyedBeforeCallback, _))
-                .WillOnce([&]() {
-                    ASSERT_DEATH_IF_SUPPORTED(
-                        { BufferMapAsync(buffer, mapMode, 0, kBufferSize, nullptr); }, "");
-                });
-
-            wgpuBufferRelease(buffer);
-        });
-        FlushClient();
-        FlushCallbacks();
-    } else {
-        // Otherwise, the callback happens when we flush it.
-        wgpuBufferRelease(buffer);
-        FlushClient();
-        ExpectWireCallbacksWhen([&](auto& mockCb) {
-            EXPECT_CALL(mockCb, Call(WGPUBufferMapAsyncStatus_DestroyedBeforeCallback, _))
-                .WillOnce([&]() {
-                    ASSERT_DEATH_IF_SUPPORTED(
-                        { BufferMapAsync(buffer, mapMode, 0, kBufferSize, nullptr); }, "");
-                });
-
-            FlushCallbacks();
-        });
-    }
-}
-#endif  // defined(DAWN_ENABLE_ASSERTS)
-
 // Test that requests inside user callbacks before disconnect are called.
 TEST_P(WireBufferMappingTests, MapInsideCallbackBeforeDisconnect) {
     WGPUMapMode mapMode = GetMapMode();
@@ -906,8 +844,8 @@
     });
 }
 
-// Test that requests inside user callbacks before object release are called.
-TEST_P(WireBufferMappingTests, MapInsideCallbackBeforeRelease) {
+// Test that requests inside user callbacks before buffer destroy are called.
+TEST_P(WireBufferMappingTests, MapInsideCallbackBeforeDestroy) {
     WGPUMapMode mapMode = GetMapMode();
     SetupBuffer(mapMode);
     BufferMapAsync(buffer, mapMode, 0, kBufferSize, nullptr);
@@ -925,7 +863,7 @@
     if (IsSpontaneous()) {
         // In spontaneous mode, when the success callback fires, the first MapAsync request
         // generated by the callback is queued, then all subsequent requests' callbacks are
-        // immediately called with MappingAlreadyPending. Finally, when we call Release, the queued
+        // immediately called with MappingAlreadyPending. Finally, when we call Destroy, the queued
         // request's callback is then called with DestroyedBeforeCallback.
         ExpectWireCallbacksWhen([&](auto& mockCb) {
             EXPECT_CALL(mockCb, Call(WGPUBufferMapAsyncStatus_Success, _)).WillOnce([&]() {
@@ -941,13 +879,13 @@
         ExpectWireCallbacksWhen([&](auto& mockCb) {
             EXPECT_CALL(mockCb, Call(WGPUBufferMapAsyncStatus_DestroyedBeforeCallback, _)).Times(1);
 
-            wgpuBufferRelease(buffer);
+            wgpuBufferDestroy(buffer);
         });
         FlushCallbacks();
     } else {
         // In non-spontaneous modes, the first callback doesn't trigger any other immediate
         // callbacks, but internally, all but the first MapAsync call's callback is set to be ready
-        // with MappingAlreadyPending. When we call Release, the first pending request is then
+        // with MappingAlreadyPending. When we call Destroy, the first pending request is then
         // marked ready with DestroyedBeforeCallback. The callbacks all run when we flush them.
         ExpectWireCallbacksWhen([&](auto& mockCb) {
             EXPECT_CALL(mockCb, Call(WGPUBufferMapAsyncStatus_Success, _)).WillOnce([&]() {
@@ -958,7 +896,7 @@
 
             FlushCallbacks();
         });
-        wgpuBufferRelease(buffer);
+        wgpuBufferDestroy(buffer);
         ExpectWireCallbacksWhen([&](auto& mockCb) {
             EXPECT_CALL(mockCb, Call(WGPUBufferMapAsyncStatus_MappingAlreadyPending, _))
                 .Times(kNumRequests - 1);
diff --git a/src/dawn/tests/unittests/wire/WireMemoryTransferServiceTests.cpp b/src/dawn/tests/unittests/wire/WireMemoryTransferServiceTests.cpp
index e038b01..79ea2b8 100644
--- a/src/dawn/tests/unittests/wire/WireMemoryTransferServiceTests.cpp
+++ b/src/dawn/tests/unittests/wire/WireMemoryTransferServiceTests.cpp
@@ -526,7 +526,7 @@
 
     // Failed deserialization is a fatal failure and the client synchronously receives a
     // DEVICE_LOST callback.
-    EXPECT_CALL(*mockBufferMapCallback, Call(WGPUBufferMapAsyncStatus_DeviceLost, _)).Times(1);
+    EXPECT_CALL(*mockBufferMapCallback, Call(WGPUBufferMapAsyncStatus_Unknown, _)).Times(1);
 
     FlushServer(false);
 
diff --git a/src/dawn/wire/client/Buffer.cpp b/src/dawn/wire/client/Buffer.cpp
index 4dcf340..238f19f 100644
--- a/src/dawn/wire/client/Buffer.cpp
+++ b/src/dawn/wire/client/Buffer.cpp
@@ -40,7 +40,6 @@
 
 namespace dawn::wire::client {
 namespace {
-
 WGPUBuffer CreateErrorBufferOOMAtClient(Device* device, const WGPUBufferDescriptor* descriptor) {
     if (descriptor->mappedAtCreation) {
         return nullptr;
@@ -52,23 +51,112 @@
     errorBufferDescriptor.nextInChain = &errorInfo.chain;
     return GetProcs().deviceCreateErrorBuffer(ToAPI(device), &errorBufferDescriptor);
 }
+}  // anonymous namespace
 
-class MapAsyncEvent : public TrackedEvent {
+class Buffer::MapAsyncEvent : public TrackedEvent {
   public:
     static constexpr EventType kType = EventType::MapAsync;
 
-    explicit MapAsyncEvent(const WGPUBufferMapCallbackInfo& callbackInfo,
-                           const Ref<MapStateData>& mapStateData)
+    MapAsyncEvent(const WGPUBufferMapCallbackInfo& callbackInfo, Buffer* buffer)
         : TrackedEvent(callbackInfo.mode),
           mCallback(callbackInfo.callback),
           mUserdata(callbackInfo.userdata),
-          mMapStateData(mapStateData) {
-        DAWN_ASSERT(mMapStateData.Get() != nullptr);
+          mBuffer(buffer) {
+        DAWN_ASSERT(buffer != nullptr);
+        GetProcs().bufferReference(ToAPI(mBuffer));
     }
 
+    ~MapAsyncEvent() override { GetProcs().bufferRelease(ToAPI(mBuffer)); }
+
     EventType GetType() override { return kType; }
 
-    void ReadyHook(WGPUBufferMapAsyncStatus status) { mStatus = status; }
+    bool IsPendingRequest(FutureID futureID) {
+        return mBuffer->mPendingMapRequest && mBuffer->mPendingMapRequest->futureID == futureID;
+    }
+
+    WireResult ReadyHook(FutureID futureID,
+                         WGPUBufferMapAsyncStatus status,
+                         uint64_t readDataUpdateInfoLength = 0,
+                         const uint8_t* readDataUpdateInfo = nullptr) {
+        auto FailRequest = [this]() -> WireResult {
+            mStatus = WGPUBufferMapAsyncStatus_Unknown;
+            return WireResult::FatalError;
+        };
+
+        // Handling for different statuses.
+        switch (status) {
+            case WGPUBufferMapAsyncStatus_MappingAlreadyPending: {
+                DAWN_ASSERT(!IsPendingRequest(futureID));
+                mStatus = status;
+                break;
+            }
+
+            // For client-side rejection errors, we clear the pending request now since they always
+            // take precedence.
+            case WGPUBufferMapAsyncStatus_DestroyedBeforeCallback:
+            case WGPUBufferMapAsyncStatus_UnmappedBeforeCallback: {
+                mStatus = status;
+                mBuffer->mPendingMapRequest = std::nullopt;
+                break;
+            }
+
+            case WGPUBufferMapAsyncStatus_Success: {
+                if (!IsPendingRequest(futureID)) {
+                    // If a success occurs (which must come from the server), but it does not
+                    // correspond to the pending request, the pending request must have been
+                    // rejected early and hence the status must be set.
+                    DAWN_ASSERT(mStatus);
+                    break;
+                }
+                mStatus = status;
+                auto& pending = mBuffer->mPendingMapRequest.value();
+                if (!pending.type) {
+                    return FailRequest();
+                }
+                switch (*pending.type) {
+                    case MapRequestType::Read: {
+                        if (readDataUpdateInfoLength > std::numeric_limits<size_t>::max()) {
+                            // This is the size of data deserialized from the command stream, which
+                            // must be CPU-addressable.
+                            return FailRequest();
+                        }
+
+                        // Validate to prevent bad map request; buffer destroyed during map request
+                        if (mBuffer->mReadHandle == nullptr) {
+                            return FailRequest();
+                        }
+                        // Update user map data with server returned data
+                        if (!mBuffer->mReadHandle->DeserializeDataUpdate(
+                                readDataUpdateInfo, static_cast<size_t>(readDataUpdateInfoLength),
+                                pending.offset, pending.size)) {
+                            return FailRequest();
+                        }
+                        mBuffer->mMappedData = const_cast<void*>(mBuffer->mReadHandle->GetData());
+                        break;
+                    }
+                    case MapRequestType::Write: {
+                        if (mBuffer->mWriteHandle == nullptr) {
+                            return FailRequest();
+                        }
+                        mBuffer->mMappedData = mBuffer->mWriteHandle->GetData();
+                        break;
+                    }
+                }
+                mBuffer->mMappedOffset = pending.offset;
+                mBuffer->mMappedSize = pending.size;
+                break;
+            }
+
+            // All other statuses are server-side status.
+            default: {
+                if (!IsPendingRequest(futureID)) {
+                    break;
+                }
+                mStatus = status;
+            }
+        }
+        return WireResult::Success;
+    }
 
   private:
     void CompleteImpl(FutureID futureID, EventCompletionType completionType) override {
@@ -78,24 +166,38 @@
         if (mStatus) {
             status = *mStatus;
         }
-        if (mMapStateData->pendingRequest && futureID == mMapStateData->pendingRequest->futureID) {
-            if (status == WGPUBufferMapAsyncStatus_Success) {
-                switch (mMapStateData->pendingRequest->type) {
-                    case MapRequestType::Read:
-                        mMapStateData->mapState = MapState::MappedForRead;
-                        break;
-                    case MapRequestType::Write:
-                        mMapStateData->mapState = MapState::MappedForWrite;
-                        break;
-                    default:
-                        DAWN_UNREACHABLE();
-                }
+
+        auto Callback = [this, &status]() {
+            if (mCallback) {
+                mCallback(status, mUserdata);
             }
-            mMapStateData->pendingRequest = std::nullopt;
+        };
+
+        if (!IsPendingRequest(futureID)) {
+            DAWN_ASSERT(status != WGPUBufferMapAsyncStatus_Success);
+            return Callback();
         }
-        if (mCallback) {
-            mCallback(status, mUserdata);
+
+        if (status == WGPUBufferMapAsyncStatus_Success) {
+            if (mBuffer->mIsDeviceAlive.expired()) {
+                // If the device lost its last ref before this callback was resolved, we want to
+                // overwrite the status. This is necessary because otherwise dropping the last
+                // device reference could race w.r.t what this callback would see.
+                status = WGPUBufferMapAsyncStatus_DestroyedBeforeCallback;
+                return Callback();
+            }
+            DAWN_ASSERT(mBuffer->mPendingMapRequest->type);
+            switch (*mBuffer->mPendingMapRequest->type) {
+                case MapRequestType::Read:
+                    mBuffer->mMappedState = MapState::MappedForRead;
+                    break;
+                case MapRequestType::Write:
+                    mBuffer->mMappedState = MapState::MappedForWrite;
+                    break;
+            }
         }
+        mBuffer->mPendingMapRequest = std::nullopt;
+        return Callback();
     }
 
     WGPUBufferMapCallback mCallback;
@@ -104,12 +206,10 @@
 
     std::optional<WGPUBufferMapAsyncStatus> mStatus;
 
-    // Shared data with the Buffer that may be modified between event handling and user inputs.
-    Ref<MapStateData> mMapStateData;
+    // Strong reference to the buffer so that when we call the callback we can pass the buffer.
+    Buffer* const mBuffer;
 };
 
-}  // anonymous namespace
-
 // static
 WGPUBuffer Buffer::Create(Device* device, const WGPUBufferDescriptor* descriptor) {
     Client* wireClient = device->GetClient();
@@ -162,20 +262,15 @@
     // This must happen after any potential error buffer creation
     // as server expects allocating ids to be monotonically increasing
     Buffer* buffer = wireClient->Make<Buffer>(device->GetEventManagerHandle(), descriptor);
-    buffer->mDestructWriteHandleOnUnmap = false;
+    buffer->mIsDeviceAlive = device->GetAliveWeakPtr();
 
     if (descriptor->mappedAtCreation) {
         // If the buffer is mapped at creation, a write handle is created and will be
         // destructed on unmap if the buffer doesn't have MapWrite usage
         // The buffer is mapped right now.
-        buffer->mMapStateData->mapState = MapState::MappedAtCreation;
-
-        // This flag is for write handle created by mappedAtCreation
-        // instead of MapWrite usage. We don't have such a case for read handle
-        buffer->mDestructWriteHandleOnUnmap = (descriptor->usage & WGPUBufferUsage_MapWrite) == 0;
-
-        buffer->mMapOffset = 0;
-        buffer->mMapSize = buffer->mSize;
+        buffer->mMappedState = MapState::MappedAtCreation;
+        buffer->mMappedOffset = 0;
+        buffer->mMappedSize = buffer->mSize;
         DAWN_ASSERT(writeHandle != nullptr);
         buffer->mMappedData = writeHandle->GetData();
     }
@@ -211,32 +306,23 @@
                const ObjectHandle& eventManagerHandle,
                const WGPUBufferDescriptor* descriptor)
     : ObjectWithEventsBase(params, eventManagerHandle),
-      mMapStateData(AcquireRef(new MapStateData{})),
       mSize(descriptor->size),
-      mUsage(static_cast<WGPUBufferUsage>(descriptor->usage)) {}
+      mUsage(static_cast<WGPUBufferUsage>(descriptor->usage)),
+      // This flag is for the write handle created by mappedAtCreation
+      // instead of MapWrite usage. We don't have such a case for read handle.
+      mDestructWriteHandleOnUnmap(descriptor->mappedAtCreation &&
+                                  ((descriptor->usage & WGPUBufferUsage_MapWrite) == 0)) {}
 
 Buffer::~Buffer() {
     FreeMappedData();
-    SetFutureStatusAndClearPending(WGPUBufferMapAsyncStatus_DestroyedBeforeCallback);
 }
 
-bool Buffer::SetFutureStatus(WGPUBufferMapAsyncStatus status) {
-    DAWN_ASSERT(mMapStateData->pendingRequest);
-    return GetEventManager().SetFutureReady<MapAsyncEvent>(mMapStateData->pendingRequest->futureID,
-                                                           status) == WireResult::Success;
-}
-
-void Buffer::SetFutureStatusAndClearPending(WGPUBufferMapAsyncStatus status) {
-    if (!mMapStateData->pendingRequest) {
-        // Since this is unconditionally called on destruction, we might not have a pending map
-        // request all the time.
+void Buffer::SetFutureStatus(WGPUBufferMapAsyncStatus status) {
+    if (!mPendingMapRequest) {
         return;
     }
-
-    FutureID futureID = mMapStateData->pendingRequest->futureID;
-    mMapStateData->pendingRequest = std::nullopt;
-    DAWN_UNUSED(GetEventManager().SetFutureReady<MapAsyncEvent>(futureID, status));
-    return;
+    DAWN_CHECK(GetEventManager().SetFutureReady<MapAsyncEvent>(mPendingMapRequest->futureID,
+                                                               status) == WireResult::Success);
 }
 
 void Buffer::MapAsync(WGPUMapModeFlags mode,
@@ -259,12 +345,12 @@
 
     Client* client = GetClient();
     auto [futureIDInternal, tracked] =
-        GetEventManager().TrackEvent(std::make_unique<MapAsyncEvent>(callbackInfo, mMapStateData));
+        GetEventManager().TrackEvent(std::make_unique<MapAsyncEvent>(callbackInfo, this));
     if (!tracked) {
         return {futureIDInternal};
     }
 
-    if (mMapStateData->pendingRequest) {
+    if (mPendingMapRequest) {
         DAWN_UNUSED(GetEventManager().SetFutureReady<MapAsyncEvent>(
             futureIDInternal, WGPUBufferMapAsyncStatus_MappingAlreadyPending));
         return {futureIDInternal};
@@ -276,14 +362,14 @@
     }
 
     // Set up the request structure that will hold information while this mapping is in flight.
-    MapRequestType mapMode = MapRequestType::None;
+    std::optional<MapRequestType> mapMode;
     if (mode & WGPUMapMode_Read) {
         mapMode = MapRequestType::Read;
     } else if (mode & WGPUMapMode_Write) {
         mapMode = MapRequestType::Write;
     }
 
-    mMapStateData->pendingRequest = {futureIDInternal, offset, size, mapMode};
+    mPendingMapRequest = {futureIDInternal, offset, size, mapMode};
 
     // Serialize the command to send to the server.
     BufferMapAsyncCmd cmd;
@@ -298,63 +384,14 @@
     return {futureIDInternal};
 }
 
-bool Buffer::OnMapAsyncCallback(WGPUFuture future,
-                                uint32_t status,
-                                uint64_t readDataUpdateInfoLength,
-                                const uint8_t* readDataUpdateInfo) {
-    // Check that the response doesn't correspond to a request that has already been rejected by
-    // unmap or destroy.
-    if (!mMapStateData->pendingRequest) {
-        return true;
-    }
-    MapRequestData& pendingRequest = mMapStateData->pendingRequest.value();
-    if (pendingRequest.futureID != future.id) {
-        return true;
-    }
-
-    auto FailRequest = [this]() -> bool {
-        SetFutureStatus(WGPUBufferMapAsyncStatus_DeviceLost);
-        return false;
-    };
-
-    if (status == WGPUBufferMapAsyncStatus_Success) {
-        switch (pendingRequest.type) {
-            case MapRequestType::Read: {
-                if (readDataUpdateInfoLength > std::numeric_limits<size_t>::max()) {
-                    // This is the size of data deserialized from the command stream, which must
-                    // be CPU-addressable.
-                    return FailRequest();
-                }
-
-                // Validate to prevent bad map request; buffer destroyed during map request
-                if (mReadHandle == nullptr) {
-                    return FailRequest();
-                }
-                // Update user map data with server returned data
-                if (!mReadHandle->DeserializeDataUpdate(
-                        readDataUpdateInfo, static_cast<size_t>(readDataUpdateInfoLength),
-                        pendingRequest.offset, pendingRequest.size)) {
-                    return FailRequest();
-                }
-                mMappedData = const_cast<void*>(mReadHandle->GetData());
-                break;
-            }
-            case MapRequestType::Write: {
-                if (mWriteHandle == nullptr) {
-                    return FailRequest();
-                }
-                mMappedData = mWriteHandle->GetData();
-                break;
-            }
-            default:
-                DAWN_UNREACHABLE();
-        }
-
-        mMapOffset = pendingRequest.offset;
-        mMapSize = pendingRequest.size;
-    }
-
-    return SetFutureStatus(static_cast<WGPUBufferMapAsyncStatus>(status));
+bool Client::DoBufferMapAsyncCallback(ObjectHandle eventManager,
+                                      WGPUFuture future,
+                                      WGPUBufferMapAsyncStatus status,
+                                      uint64_t readDataUpdateInfoLength,
+                                      const uint8_t* readDataUpdateInfo) {
+    return GetEventManager(eventManager)
+               .SetFutureReady<Buffer::MapAsyncEvent>(future.id, status, readDataUpdateInfoLength,
+                                                      readDataUpdateInfo) == WireResult::Success;
 }
 
 void* Buffer::GetMappedRange(size_t offset, size_t size) {
@@ -383,23 +420,20 @@
     //   - Server -> Client: Result of MapRequest2
     Client* client = GetClient();
 
-    // mWriteHandle can still be nullptr if buffer has been destroyed before unmap
-    if ((mMapStateData->mapState == MapState::MappedForWrite ||
-         mMapStateData->mapState == MapState::MappedAtCreation) &&
-        mWriteHandle != nullptr) {
+    if (IsMappedForWriting()) {
         // Writes need to be flushed before Unmap is sent. Unmap calls all associated
         // in-flight callbacks which may read the updated data.
 
         // Get the serialization size of data update writes.
         size_t writeDataUpdateInfoLength =
-            mWriteHandle->SizeOfSerializeDataUpdate(mMapOffset, mMapSize);
+            mWriteHandle->SizeOfSerializeDataUpdate(mMappedOffset, mMappedSize);
 
         BufferUpdateMappedDataCmd cmd;
         cmd.bufferId = GetWireId();
         cmd.writeDataUpdateInfoLength = writeDataUpdateInfoLength;
         cmd.writeDataUpdateInfo = nullptr;
-        cmd.offset = mMapOffset;
-        cmd.size = mMapSize;
+        cmd.offset = mMappedOffset;
+        cmd.size = mMappedSize;
 
         client->SerializeCommand(
             cmd, CommandExtension{writeDataUpdateInfoLength, [&](char* writeHandleBuffer) {
@@ -412,11 +446,11 @@
         // If mDestructWriteHandleOnUnmap is true, that means the write handle is merely
         // for mappedAtCreation usage. It is destroyed on unmap after flush to server
         // instead of at buffer destruction.
-        if (mMapStateData->mapState == MapState::MappedAtCreation && mDestructWriteHandleOnUnmap) {
+        if (mDestructWriteHandleOnUnmap) {
             mWriteHandle = nullptr;
             if (mReadHandle) {
                 // If it's both mappedAtCreation and MapRead we need to reset
-                // mMappedData to readHandle's GetData(). This could be changed to
+                // mData to readHandle's GetData(). This could be changed to
                 // merging read/write handle in future
                 mMappedData = const_cast<void*>(mReadHandle->GetData());
             }
@@ -424,15 +458,15 @@
     }
 
     // Free map access tokens
-    mMapStateData->mapState = MapState::Unmapped;
-    mMapOffset = 0;
-    mMapSize = 0;
+    mMappedState = MapState::Unmapped;
+    mMappedOffset = 0;
+    mMappedSize = 0;
 
     BufferUnmapCmd cmd;
     cmd.self = ToAPI(this);
     client->SerializeCommand(cmd);
 
-    SetFutureStatusAndClearPending(WGPUBufferMapAsyncStatus_UnmappedBeforeCallback);
+    SetFutureStatus(WGPUBufferMapAsyncStatus_UnmappedBeforeCallback);
 }
 
 void Buffer::Destroy() {
@@ -440,13 +474,12 @@
 
     // Remove the current mapping and destroy Read/WriteHandles.
     FreeMappedData();
-    mMapStateData->mapState = MapState::Unmapped;
 
     BufferDestroyCmd cmd;
     cmd.self = ToAPI(this);
     client->SerializeCommand(cmd);
 
-    SetFutureStatusAndClearPending(WGPUBufferMapAsyncStatus_DestroyedBeforeCallback);
+    SetFutureStatus(WGPUBufferMapAsyncStatus_DestroyedBeforeCallback);
 }
 
 WGPUBufferUsage Buffer::GetUsage() const {
@@ -458,13 +491,13 @@
 }
 
 WGPUBufferMapState Buffer::GetMapState() const {
-    switch (mMapStateData->mapState) {
+    switch (mMappedState) {
         case MapState::MappedForRead:
         case MapState::MappedForWrite:
         case MapState::MappedAtCreation:
             return WGPUBufferMapState_Mapped;
         case MapState::Unmapped:
-            if (mMapStateData->pendingRequest) {
+            if (mPendingMapRequest) {
                 return WGPUBufferMapState_Pending;
             } else {
                 return WGPUBufferMapState_Unmapped;
@@ -474,27 +507,26 @@
 }
 
 bool Buffer::IsMappedForReading() const {
-    return mMapStateData->mapState == MapState::MappedForRead;
+    return mMappedState == MapState::MappedForRead;
 }
 
 bool Buffer::IsMappedForWriting() const {
-    return mMapStateData->mapState == MapState::MappedForWrite ||
-           mMapStateData->mapState == MapState::MappedAtCreation;
+    return mMappedState == MapState::MappedForWrite || mMappedState == MapState::MappedAtCreation;
 }
 
 bool Buffer::CheckGetMappedRangeOffsetSize(size_t offset, size_t size) const {
-    if (offset % 8 != 0 || offset < mMapOffset || offset > mSize) {
+    if (offset % 8 != 0 || offset < mMappedOffset || offset > mSize) {
         return false;
     }
 
     size_t rangeSize = size == WGPU_WHOLE_MAP_SIZE ? mSize - offset : size;
 
-    if (rangeSize % 4 != 0 || rangeSize > mMapSize) {
+    if (rangeSize % 4 != 0 || rangeSize > mMappedSize) {
         return false;
     }
 
-    size_t offsetInMappedRange = offset - mMapOffset;
-    return offsetInMappedRange <= mMapSize - rangeSize;
+    size_t offsetInMappedRange = offset - mMappedOffset;
+    return offsetInMappedRange <= mMappedSize - rangeSize;
 }
 
 void Buffer::FreeMappedData() {
@@ -503,15 +535,16 @@
     // use-after-free of the mapped data. This is particularly useful for WebGPU test about the
     // interaction of mapping and GC.
     if (mMappedData) {
-        memset(static_cast<uint8_t*>(mMappedData) + mMapOffset, 0xCA, mMapSize);
+        memset(static_cast<uint8_t*>(mMappedData) + mMappedOffset, 0xCA, mMappedSize);
     }
 #endif  // defined(DAWN_ENABLE_ASSERTS)
 
-    mMapOffset = 0;
-    mMapSize = 0;
+    mMappedOffset = 0;
+    mMappedSize = 0;
     mReadHandle = nullptr;
     mWriteHandle = nullptr;
     mMappedData = nullptr;
+    mMappedState = MapState::Unmapped;
 }
 
 }  // namespace dawn::wire::client
diff --git a/src/dawn/wire/client/Buffer.h b/src/dawn/wire/client/Buffer.h
index 1c52d57..7747790 100644
--- a/src/dawn/wire/client/Buffer.h
+++ b/src/dawn/wire/client/Buffer.h
@@ -43,28 +43,6 @@
 
 class Device;
 
-enum class MapRequestType { None, Read, Write };
-
-enum class MapState {
-    Unmapped,
-    MappedForRead,
-    MappedForWrite,
-    MappedAtCreation,
-};
-
-struct MapRequestData {
-    FutureID futureID = kNullFutureID;
-    size_t offset = 0;
-    size_t size = 0;
-    MapRequestType type = MapRequestType::None;
-};
-
-struct MapStateData : public RefCounted {
-    // Up to only one request can exist at a single time. Other requests are rejected.
-    std::optional<MapRequestData> pendingRequest = std::nullopt;
-    MapState mapState = MapState::Unmapped;
-};
-
 class Buffer final : public ObjectWithEventsBase {
   public:
     static WGPUBuffer Create(Device* device, const WGPUBufferDescriptor* descriptor);
@@ -74,10 +52,6 @@
            const WGPUBufferDescriptor* descriptor);
     ~Buffer() override;
 
-    bool OnMapAsyncCallback(WGPUFuture future,
-                            uint32_t status,
-                            uint64_t readDataUpdateInfoLength,
-                            const uint8_t* readDataUpdateInfo);
     void MapAsync(WGPUMapModeFlags mode,
                   size_t offset,
                   size_t size,
@@ -100,9 +74,11 @@
     WGPUBufferMapState GetMapState() const;
 
   private:
+    friend class Client;
+    class MapAsyncEvent;
+
     // Prepares the callbacks to be called and potentially calls them
-    bool SetFutureStatus(WGPUBufferMapAsyncStatus status);
-    void SetFutureStatusAndClearPending(WGPUBufferMapAsyncStatus status);
+    void SetFutureStatus(WGPUBufferMapAsyncStatus status);
 
     bool IsMappedForReading() const;
     bool IsMappedForWriting() const;
@@ -110,25 +86,39 @@
 
     void FreeMappedData();
 
-    // The map state is a shared resource with the TrackedEvent so that it is updated only when the
-    // callback is actually called. This is important for WaitAny and ProcessEvents cases where the
-    // server may have responded, but due to an early Unmap or Destroy before the corresponding
-    // WaitAny or ProcessEvents call, we need to update the callback result.
-    Ref<MapStateData> mMapStateData;
+    const uint64_t mSize = 0;
+    const WGPUBufferUsage mUsage;
+    const bool mDestructWriteHandleOnUnmap;
 
-    uint64_t mSize = 0;
-    WGPUBufferUsage mUsage;
+    std::weak_ptr<bool> mIsDeviceAlive;
+
+    // Mapping members are mutable depending on the current map state.
+    enum class MapRequestType { Read, Write };
+    struct MapRequest {
+        FutureID futureID = kNullFutureID;
+        size_t offset = 0;
+        size_t size = 0;
+        // Because validation for request type is validated via the backend, we use an optional type
+        // here. This is nullopt when an invalid request type is passed to the wire.
+        std::optional<MapRequestType> type;
+    };
+    enum class MapState {
+        Unmapped,
+        MappedForRead,
+        MappedForWrite,
+        MappedAtCreation,
+    };
+    std::optional<MapRequest> mPendingMapRequest = std::nullopt;
+    MapState mMappedState = MapState::Unmapped;
+    // TODO(https://crbug.com/dawn/2345): Investigate `DanglingUntriaged` in dawn/wire.
+    raw_ptr<void, DanglingUntriaged> mMappedData = nullptr;
+    size_t mMappedOffset = 0;
+    size_t mMappedSize = 0;
 
     // Only one mapped pointer can be active at a time
     // TODO(enga): Use a tagged pointer to save space.
     std::unique_ptr<MemoryTransferService::ReadHandle> mReadHandle = nullptr;
     std::unique_ptr<MemoryTransferService::WriteHandle> mWriteHandle = nullptr;
-    bool mDestructWriteHandleOnUnmap = false;
-
-    // TODO(https://crbug.com/dawn/2345): Investigate `DanglingUntriaged` in dawn/wire.
-    raw_ptr<void, DanglingUntriaged> mMappedData = nullptr;
-    size_t mMapOffset = 0;
-    size_t mMapSize = 0;
 };
 
 }  // namespace dawn::wire::client
diff --git a/src/dawn/wire/client/ClientDoers.cpp b/src/dawn/wire/client/ClientDoers.cpp
index 210cc53..2ed46d8 100644
--- a/src/dawn/wire/client/ClientDoers.cpp
+++ b/src/dawn/wire/client/ClientDoers.cpp
@@ -88,20 +88,6 @@
     return device->OnPopErrorScopeCallback(requestSerial, errorType, message);
 }
 
-// TODO(dawn:2061) May be able to move this to Buffer.cpp once we move all mapping logic.
-bool Client::DoBufferMapAsyncCallback(Buffer* buffer,
-                                      ObjectHandle eventManager,
-                                      WGPUFuture future,
-                                      uint32_t status,
-                                      uint64_t readDataUpdateInfoLength,
-                                      const uint8_t* readDataUpdateInfo) {
-    // The buffer might have been deleted or recreated so this isn't an error.
-    if (buffer == nullptr) {
-        return true;
-    }
-    return buffer->OnMapAsyncCallback(future, status, readDataUpdateInfoLength, readDataUpdateInfo);
-}
-
 bool Client::DoDeviceCreateComputePipelineAsyncCallback(Device* device,
                                                         uint64_t requestSerial,
                                                         WGPUCreatePipelineAsyncStatus status,
diff --git a/src/dawn/wire/client/EventManager.h b/src/dawn/wire/client/EventManager.h
index 071b825..8ea33ee 100644
--- a/src/dawn/wire/client/EventManager.h
+++ b/src/dawn/wire/client/EventManager.h
@@ -114,7 +114,7 @@
         }
 
         std::unique_ptr<TrackedEvent> spontaneousEvent;
-        WIRE_TRY(mTrackedEvents.Use([&](auto trackedEvents) {
+        WireResult result = mTrackedEvents.Use([&](auto trackedEvents) {
             auto it = trackedEvents->find(futureID);
             if (it == trackedEvents->end()) {
                 // If the future is not found, it must've already been completed.
@@ -128,8 +128,9 @@
                 DAWN_ASSERT(trackedEvent->GetType() == Event::kType);
                 return WireResult::FatalError;
             }
-            static_cast<Event*>(trackedEvent.get())
-                ->ReadyHook(std::forward<ReadyArgs>(readyArgs)...);
+
+            WireResult result = static_cast<Event*>(trackedEvent.get())
+                                    ->ReadyHook(futureID, std::forward<ReadyArgs>(readyArgs)...);
             trackedEvent->SetReady();
 
             // If the event can be spontaneously completed, prepare to do so now.
@@ -137,14 +138,14 @@
                 spontaneousEvent = std::move(trackedEvent);
                 trackedEvents->erase(futureID);
             }
-            return WireResult::Success;
-        }));
+            return result;
+        });
 
         // Handle spontaneous completions.
         if (spontaneousEvent) {
             spontaneousEvent->Complete(futureID, EventCompletionType::Ready);
         }
-        return WireResult::Success;
+        return result;
     }
 
     void ProcessPollEvents();
diff --git a/src/dawn/wire/client/Instance.cpp b/src/dawn/wire/client/Instance.cpp
index e67fc83..95718ad 100644
--- a/src/dawn/wire/client/Instance.cpp
+++ b/src/dawn/wire/client/Instance.cpp
@@ -54,12 +54,13 @@
 
     EventType GetType() override { return kType; }
 
-    void ReadyHook(WGPURequestAdapterStatus status,
-                   const char* message,
-                   const WGPUAdapterProperties* properties,
-                   const WGPUSupportedLimits* limits,
-                   uint32_t featuresCount,
-                   const WGPUFeatureName* features) {
+    WireResult ReadyHook(FutureID futureID,
+                         WGPURequestAdapterStatus status,
+                         const char* message,
+                         const WGPUAdapterProperties* properties,
+                         const WGPUSupportedLimits* limits,
+                         uint32_t featuresCount,
+                         const WGPUFeatureName* features) {
         DAWN_ASSERT(mAdapter != nullptr);
         mStatus = status;
         if (message != nullptr) {
@@ -70,6 +71,7 @@
             mAdapter->SetLimits(limits);
             mAdapter->SetFeatures(features, featuresCount);
         }
+        return WireResult::Success;
     }
 
   private:
diff --git a/src/dawn/wire/client/Queue.cpp b/src/dawn/wire/client/Queue.cpp
index 5337853..d897ba7 100644
--- a/src/dawn/wire/client/Queue.cpp
+++ b/src/dawn/wire/client/Queue.cpp
@@ -48,7 +48,10 @@
 
     EventType GetType() override { return kType; }
 
-    void ReadyHook(WGPUQueueWorkDoneStatus status) { mStatus = status; }
+    WireResult ReadyHook(FutureID futureID, WGPUQueueWorkDoneStatus status) {
+        mStatus = status;
+        return WireResult::Success;
+    }
 
   private:
     void CompleteImpl(FutureID futureID, EventCompletionType completionType) override {
diff --git a/src/dawn/wire/server/ServerBuffer.cpp b/src/dawn/wire/server/ServerBuffer.cpp
index ca1471c..22e2769 100644
--- a/src/dawn/wire/server/ServerBuffer.cpp
+++ b/src/dawn/wire/server/ServerBuffer.cpp
@@ -225,8 +225,6 @@
     bool isSuccess = status == WGPUBufferMapAsyncStatus_Success;
 
     ReturnBufferMapAsyncCallbackCmd cmd;
-    // TODO(dawn:2061) Should be able to remove buffer once mapping is updated.
-    cmd.buffer = data->buffer;
     cmd.eventManager = data->eventManager;
     cmd.future = data->future;
     cmd.status = status;