dawn_wire: Add support for injecting/reserving swapchains

This will help experiment using dawn_wire for remoting WebGPU to render
on the screen.

Bug: None
Change-Id: I9a60ff8c3889ec917f6fd56e4cbb1ffef639748d
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/47621
Auto-Submit: Corentin Wallez <cwallez@chromium.org>
Reviewed-by: Austin Eng <enga@chromium.org>
Reviewed-by: Brandon Jones <bajones@chromium.org>
Commit-Queue: Austin Eng <enga@chromium.org>
diff --git a/src/dawn_wire/WireClient.cpp b/src/dawn_wire/WireClient.cpp
index 481ced4..aa83013 100644
--- a/src/dawn_wire/WireClient.cpp
+++ b/src/dawn_wire/WireClient.cpp
@@ -33,6 +33,10 @@
         return mImpl->ReserveTexture(device);
     }
 
+    ReservedSwapChain WireClient::ReserveSwapChain(WGPUDevice device) {
+        return mImpl->ReserveSwapChain(device);
+    }
+
     ReservedDevice WireClient::ReserveDevice() {
         return mImpl->ReserveDevice();
     }
@@ -41,6 +45,10 @@
         mImpl->ReclaimTextureReservation(reservation);
     }
 
+    void WireClient::ReclaimSwapChainReservation(const ReservedSwapChain& reservation) {
+        mImpl->ReclaimSwapChainReservation(reservation);
+    }
+
     void WireClient::ReclaimDeviceReservation(const ReservedDevice& reservation) {
         mImpl->ReclaimDeviceReservation(reservation);
     }
diff --git a/src/dawn_wire/WireServer.cpp b/src/dawn_wire/WireServer.cpp
index a3599f4..bb3d7ba 100644
--- a/src/dawn_wire/WireServer.cpp
+++ b/src/dawn_wire/WireServer.cpp
@@ -39,6 +39,14 @@
         return mImpl->InjectTexture(texture, id, generation, deviceId, deviceGeneration);
     }
 
+    bool WireServer::InjectSwapChain(WGPUSwapChain swapchain,
+                                     uint32_t id,
+                                     uint32_t generation,
+                                     uint32_t deviceId,
+                                     uint32_t deviceGeneration) {
+        return mImpl->InjectSwapChain(swapchain, id, generation, deviceId, deviceGeneration);
+    }
+
     bool WireServer::InjectDevice(WGPUDevice device, uint32_t id, uint32_t generation) {
         return mImpl->InjectDevice(device, id, generation);
     }
diff --git a/src/dawn_wire/client/Client.cpp b/src/dawn_wire/client/Client.cpp
index 3b29f9e..a00bb5e 100644
--- a/src/dawn_wire/client/Client.cpp
+++ b/src/dawn_wire/client/Client.cpp
@@ -96,6 +96,18 @@
         return result;
     }
 
+    ReservedSwapChain Client::ReserveSwapChain(WGPUDevice device) {
+        auto* allocation = SwapChainAllocator().New(this);
+
+        ReservedSwapChain result;
+        result.swapchain = ToAPI(allocation->object.get());
+        result.id = allocation->object->id;
+        result.generation = allocation->generation;
+        result.deviceId = FromAPI(device)->id;
+        result.deviceGeneration = DeviceAllocator().GetGeneration(FromAPI(device)->id);
+        return result;
+    }
+
     ReservedDevice Client::ReserveDevice() {
         auto* allocation = DeviceAllocator().New(this);
 
@@ -110,6 +122,10 @@
         TextureAllocator().Free(FromAPI(reservation.texture));
     }
 
+    void Client::ReclaimSwapChainReservation(const ReservedSwapChain& reservation) {
+        SwapChainAllocator().Free(FromAPI(reservation.swapchain));
+    }
+
     void Client::ReclaimDeviceReservation(const ReservedDevice& reservation) {
         DeviceAllocator().Free(FromAPI(reservation.device));
     }
diff --git a/src/dawn_wire/client/Client.h b/src/dawn_wire/client/Client.h
index db9ed43..3616e37 100644
--- a/src/dawn_wire/client/Client.h
+++ b/src/dawn_wire/client/Client.h
@@ -44,9 +44,11 @@
         }
 
         ReservedTexture ReserveTexture(WGPUDevice device);
+        ReservedSwapChain ReserveSwapChain(WGPUDevice device);
         ReservedDevice ReserveDevice();
 
         void ReclaimTextureReservation(const ReservedTexture& reservation);
+        void ReclaimSwapChainReservation(const ReservedSwapChain& reservation);
         void ReclaimDeviceReservation(const ReservedDevice& reservation);
 
         template <typename Cmd>
diff --git a/src/dawn_wire/server/Server.cpp b/src/dawn_wire/server/Server.cpp
index ea0d8fc..24fbeda 100644
--- a/src/dawn_wire/server/Server.cpp
+++ b/src/dawn_wire/server/Server.cpp
@@ -72,6 +72,38 @@
         return true;
     }
 
+    bool Server::InjectSwapChain(WGPUSwapChain swapchain,
+                                 uint32_t id,
+                                 uint32_t generation,
+                                 uint32_t deviceId,
+                                 uint32_t deviceGeneration) {
+        ASSERT(swapchain != nullptr);
+        ObjectData<WGPUDevice>* device = DeviceObjects().Get(deviceId);
+        if (device == nullptr || device->generation != deviceGeneration) {
+            return false;
+        }
+
+        ObjectData<WGPUSwapChain>* data = SwapChainObjects().Allocate(id);
+        if (data == nullptr) {
+            return false;
+        }
+
+        data->handle = swapchain;
+        data->generation = generation;
+        data->state = AllocationState::Allocated;
+        data->deviceInfo = device->info.get();
+
+        if (!TrackDeviceChild(data->deviceInfo, ObjectType::SwapChain, id)) {
+            return false;
+        }
+
+        // The texture is externally owned so it shouldn't be destroyed when we receive a destroy
+        // message from the client. Add a reference to counterbalance the eventual release.
+        mProcs.swapChainReference(swapchain);
+
+        return true;
+    }
+
     bool Server::InjectDevice(WGPUDevice device, uint32_t id, uint32_t generation) {
         ASSERT(device != nullptr);
         ObjectData<WGPUDevice>* data = DeviceObjects().Allocate(id);
diff --git a/src/dawn_wire/server/Server.h b/src/dawn_wire/server/Server.h
index 5baea19..1979d87 100644
--- a/src/dawn_wire/server/Server.h
+++ b/src/dawn_wire/server/Server.h
@@ -180,6 +180,12 @@
                            uint32_t deviceId,
                            uint32_t deviceGeneration);
 
+        bool InjectSwapChain(WGPUSwapChain swapchain,
+                             uint32_t id,
+                             uint32_t generation,
+                             uint32_t deviceId,
+                             uint32_t deviceGeneration);
+
         bool InjectDevice(WGPUDevice device, uint32_t id, uint32_t generation);
 
         WGPUDevice GetDevice(uint32_t id, uint32_t generation);
diff --git a/src/include/dawn_wire/WireClient.h b/src/include/dawn_wire/WireClient.h
index 097fd1a..87f2bab 100644
--- a/src/include/dawn_wire/WireClient.h
+++ b/src/include/dawn_wire/WireClient.h
@@ -38,6 +38,14 @@
         uint32_t deviceGeneration;
     };
 
+    struct ReservedSwapChain {
+        WGPUSwapChain swapchain;
+        uint32_t id;
+        uint32_t generation;
+        uint32_t deviceId;
+        uint32_t deviceGeneration;
+    };
+
     struct ReservedDevice {
         WGPUDevice device;
         uint32_t id;
@@ -58,9 +66,11 @@
                                             size_t size) override final;
 
         ReservedTexture ReserveTexture(WGPUDevice device);
+        ReservedSwapChain ReserveSwapChain(WGPUDevice device);
         ReservedDevice ReserveDevice();
 
         void ReclaimTextureReservation(const ReservedTexture& reservation);
+        void ReclaimSwapChainReservation(const ReservedSwapChain& reservation);
         void ReclaimDeviceReservation(const ReservedDevice& reservation);
 
         // Disconnects the client.
diff --git a/src/include/dawn_wire/WireServer.h b/src/include/dawn_wire/WireServer.h
index 59df6a4..14d2354 100644
--- a/src/include/dawn_wire/WireServer.h
+++ b/src/include/dawn_wire/WireServer.h
@@ -42,12 +42,16 @@
         const volatile char* HandleCommands(const volatile char* commands,
                                             size_t size) override final;
 
-        // TODO(enga): Remove defaults after updating Chrome.
         bool InjectTexture(WGPUTexture texture,
                            uint32_t id,
                            uint32_t generation,
-                           uint32_t deviceId = 1,
-                           uint32_t deviceGeneration = 0);
+                           uint32_t deviceId,
+                           uint32_t deviceGeneration);
+        bool InjectSwapChain(WGPUSwapChain swapchain,
+                             uint32_t id,
+                             uint32_t generation,
+                             uint32_t deviceId,
+                             uint32_t deviceGeneration);
 
         bool InjectDevice(WGPUDevice device, uint32_t id, uint32_t generation);
 
diff --git a/src/tests/BUILD.gn b/src/tests/BUILD.gn
index 18d2fb5..b025a63 100644
--- a/src/tests/BUILD.gn
+++ b/src/tests/BUILD.gn
@@ -230,6 +230,7 @@
     "unittests/wire/WireExtensionTests.cpp",
     "unittests/wire/WireFenceTests.cpp",
     "unittests/wire/WireInjectDeviceTests.cpp",
+    "unittests/wire/WireInjectSwapChainTests.cpp",
     "unittests/wire/WireInjectTextureTests.cpp",
     "unittests/wire/WireMemoryTransferServiceTests.cpp",
     "unittests/wire/WireOptionalTests.cpp",
diff --git a/src/tests/unittests/wire/WireInjectSwapChainTests.cpp b/src/tests/unittests/wire/WireInjectSwapChainTests.cpp
new file mode 100644
index 0000000..6c3e104
--- /dev/null
+++ b/src/tests/unittests/wire/WireInjectSwapChainTests.cpp
@@ -0,0 +1,116 @@
+// Copyright 2021 The Dawn Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "tests/unittests/wire/WireTest.h"
+
+#include "dawn_wire/WireClient.h"
+#include "dawn_wire/WireServer.h"
+
+using namespace testing;
+using namespace dawn_wire;
+
+class WireInjectSwapChainTests : public WireTest {
+  public:
+    WireInjectSwapChainTests() {
+    }
+    ~WireInjectSwapChainTests() override = default;
+};
+
+// Test that reserving and injecting a swapchain makes calls on the client object forward to the
+// server object correctly.
+TEST_F(WireInjectSwapChainTests, CallAfterReserveInject) {
+    ReservedSwapChain reservation = GetWireClient()->ReserveSwapChain(device);
+
+    WGPUSwapChain apiSwapchain = api.GetNewSwapChain();
+    EXPECT_CALL(api, SwapChainReference(apiSwapchain));
+    ASSERT_TRUE(GetWireServer()->InjectSwapChain(apiSwapchain, reservation.id,
+                                                 reservation.generation, reservation.deviceId,
+                                                 reservation.deviceGeneration));
+
+    wgpuSwapChainPresent(reservation.swapchain);
+    EXPECT_CALL(api, SwapChainPresent(apiSwapchain));
+    FlushClient();
+}
+
+// Test that reserve correctly returns different IDs each time.
+TEST_F(WireInjectSwapChainTests, ReserveDifferentIDs) {
+    ReservedSwapChain reservation1 = GetWireClient()->ReserveSwapChain(device);
+    ReservedSwapChain reservation2 = GetWireClient()->ReserveSwapChain(device);
+
+    ASSERT_NE(reservation1.id, reservation2.id);
+    ASSERT_NE(reservation1.swapchain, reservation2.swapchain);
+}
+
+// Test that injecting the same id without a destroy first fails.
+TEST_F(WireInjectSwapChainTests, InjectExistingID) {
+    ReservedSwapChain reservation = GetWireClient()->ReserveSwapChain(device);
+
+    WGPUSwapChain apiSwapchain = api.GetNewSwapChain();
+    EXPECT_CALL(api, SwapChainReference(apiSwapchain));
+    ASSERT_TRUE(GetWireServer()->InjectSwapChain(apiSwapchain, reservation.id,
+                                                 reservation.generation, reservation.deviceId,
+                                                 reservation.deviceGeneration));
+
+    // ID already in use, call fails.
+    ASSERT_FALSE(GetWireServer()->InjectSwapChain(apiSwapchain, reservation.id,
+                                                  reservation.generation, reservation.deviceId,
+                                                  reservation.deviceGeneration));
+}
+
+// Test that the server only borrows the swapchain and does a single reference-release
+TEST_F(WireInjectSwapChainTests, InjectedSwapChainLifetime) {
+    ReservedSwapChain reservation = GetWireClient()->ReserveSwapChain(device);
+
+    // Injecting the swapchain adds a reference
+    WGPUSwapChain apiSwapchain = api.GetNewSwapChain();
+    EXPECT_CALL(api, SwapChainReference(apiSwapchain));
+    ASSERT_TRUE(GetWireServer()->InjectSwapChain(apiSwapchain, reservation.id,
+                                                 reservation.generation, reservation.deviceId,
+                                                 reservation.deviceGeneration));
+
+    // Releasing the swapchain removes a single reference.
+    wgpuSwapChainRelease(reservation.swapchain);
+    EXPECT_CALL(api, SwapChainRelease(apiSwapchain));
+    FlushClient();
+
+    // Deleting the server doesn't release a second reference.
+    DeleteServer();
+    Mock::VerifyAndClearExpectations(&api);
+}
+
+// Test that a swapchain reservation can be reclaimed. This is necessary to
+// avoid leaking ObjectIDs for reservations that are never injected.
+TEST_F(WireInjectSwapChainTests, ReclaimSwapChainReservation) {
+    // Test that doing a reservation and full release is an error.
+    {
+        ReservedSwapChain reservation = GetWireClient()->ReserveSwapChain(device);
+        wgpuSwapChainRelease(reservation.swapchain);
+        FlushClient(false);
+    }
+
+    // Test that doing a reservation and then reclaiming it recycles the ID.
+    {
+        ReservedSwapChain reservation1 = GetWireClient()->ReserveSwapChain(device);
+        GetWireClient()->ReclaimSwapChainReservation(reservation1);
+
+        ReservedSwapChain reservation2 = GetWireClient()->ReserveSwapChain(device);
+
+        // The ID is the same, but the generation is still different.
+        ASSERT_EQ(reservation1.id, reservation2.id);
+        ASSERT_NE(reservation1.generation, reservation2.generation);
+
+        // No errors should occur.
+        FlushClient();
+    }
+}