Add a way to query if a device exists on the wire server

Chromium tracks the devices which live on the wire so it can
automatically call tick on devices that have pending work. This used
to be done by querying an (id, generation) pair and checking if it
resolves to a non-null device.

This CL adds a new way to query directly using the device, since a
refactor in Chrome will change creation such that the id and generation
is not known when a new device is requested.

Also fixes a bug for device callbacks where the required callback
userdata wasn't fully populated for devices created with
requestAdapter. Update a test to check for this as well.

Bug: chromium:1315260
Change-Id: I7468edc3e77bade191e1e9f3eaadebbf4441d88a
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/89520
Reviewed-by: Loko Kung <lokokung@google.com>
Kokoro: Kokoro <noreply+kokoro@google.com>
Commit-Queue: Austin Eng <enga@chromium.org>
diff --git a/include/dawn/wire/WireServer.h b/include/dawn/wire/WireServer.h
index 9905608..03d694e 100644
--- a/include/dawn/wire/WireServer.h
+++ b/include/dawn/wire/WireServer.h
@@ -65,6 +65,11 @@
     // previously injected devices, and observing if GetDevice(id, generation) returns non-null.
     WGPUDevice GetDevice(uint32_t id, uint32_t generation);
 
+    // Check if a device handle is known by the wire.
+    // In Chrome, we need to know the list of live devices so we can call device.Tick() on all of
+    // them periodically to ensure progress on asynchronous work is made.
+    bool IsDeviceKnown(WGPUDevice device) const;
+
   private:
     std::unique_ptr<server::Server> mImpl;
 };
diff --git a/src/dawn/tests/unittests/wire/WireAdapterTests.cpp b/src/dawn/tests/unittests/wire/WireAdapterTests.cpp
index aa56c39..e87414d 100644
--- a/src/dawn/tests/unittests/wire/WireAdapterTests.cpp
+++ b/src/dawn/tests/unittests/wire/WireAdapterTests.cpp
@@ -177,6 +177,10 @@
 
     // Expect the server to receive the message. Then, mock a fake reply.
     WGPUDevice apiDevice = api.GetNewDevice();
+    // The backend device should not be known by the wire server.
+    EXPECT_FALSE(GetWireServer()->IsDeviceKnown(apiDevice));
+
+    wgpu::Device device;
     EXPECT_CALL(api, OnAdapterRequestDevice(apiAdapter, NotNull(), NotNull(), NotNull()))
         .WillOnce(InvokeWithoutArgs([&]() {
             // Set on device creation to forward callbacks to the client.
@@ -203,15 +207,20 @@
                     return fakeFeatures.size();
                 })));
 
+            // The backend device should still not be known by the wire server since the
+            // callback has not been called yet.
+            EXPECT_FALSE(GetWireServer()->IsDeviceKnown(apiDevice));
             api.CallAdapterRequestDeviceCallback(apiAdapter, WGPURequestDeviceStatus_Success,
                                                  apiDevice, nullptr);
+            // After the callback is called, the backend device is now known by the server.
+            EXPECT_TRUE(GetWireServer()->IsDeviceKnown(apiDevice));
         }));
     FlushClient();
 
     // Expect the callback in the client and all the device information to match.
     EXPECT_CALL(cb, Call(WGPURequestDeviceStatus_Success, NotNull(), nullptr, this))
         .WillOnce(WithArg<1>(Invoke([&](WGPUDevice cDevice) {
-            wgpu::Device device = wgpu::Device::Acquire(cDevice);
+            device = wgpu::Device::Acquire(cDevice);
 
             wgpu::SupportedLimits limits;
             EXPECT_TRUE(device.GetLimits(&limits));
@@ -230,10 +239,28 @@
         })));
     FlushServer();
 
+    // Test that callbacks can propagate from server to client.
+    MockCallback<WGPUErrorCallback> errorCb;
+    device.SetUncapturedErrorCallback(errorCb.Callback(), errorCb.MakeUserdata(this));
+    api.CallDeviceSetUncapturedErrorCallbackCallback(apiDevice, WGPUErrorType_Validation,
+                                                     "Some error message");
+
+    EXPECT_CALL(errorCb, Call(WGPUErrorType_Validation, StrEq("Some error message"), this))
+        .Times(1);
+    FlushServer();
+
+    device = nullptr;
     // Cleared when the device is destroyed.
     EXPECT_CALL(api, OnDeviceSetUncapturedErrorCallback(apiDevice, nullptr, nullptr)).Times(1);
     EXPECT_CALL(api, OnDeviceSetLoggingCallback(apiDevice, nullptr, nullptr)).Times(1);
     EXPECT_CALL(api, OnDeviceSetDeviceLostCallback(apiDevice, nullptr, nullptr)).Times(1);
+    EXPECT_CALL(api, DeviceRelease(apiDevice));
+
+    // Server has not recevied the release yet, so the device should be known.
+    EXPECT_TRUE(GetWireServer()->IsDeviceKnown(apiDevice));
+    FlushClient();
+    // After receiving the release call, the device is no longer known by the server.
+    EXPECT_FALSE(GetWireServer()->IsDeviceKnown(apiDevice));
 }
 
 // Test that features requested that the implementation supports, but not the
diff --git a/src/dawn/wire/WireServer.cpp b/src/dawn/wire/WireServer.cpp
index c806fe0..b864112 100644
--- a/src/dawn/wire/WireServer.cpp
+++ b/src/dawn/wire/WireServer.cpp
@@ -58,6 +58,10 @@
     return mImpl->GetDevice(id, generation);
 }
 
+bool WireServer::IsDeviceKnown(WGPUDevice device) const {
+    return mImpl->IsDeviceKnown(device);
+}
+
 namespace server {
 MemoryTransferService::MemoryTransferService() = default;
 
diff --git a/src/dawn/wire/server/ObjectStorage.h b/src/dawn/wire/server/ObjectStorage.h
index 06fbc04..1545dc3 100644
--- a/src/dawn/wire/server/ObjectStorage.h
+++ b/src/dawn/wire/server/ObjectStorage.h
@@ -94,11 +94,11 @@
 
 // Keeps track of the mapping between client IDs and backend objects.
 template <typename T>
-class KnownObjects {
+class KnownObjectsBase {
   public:
     using Data = ObjectData<T>;
 
-    KnownObjects() {
+    KnownObjectsBase() {
         // Reserve ID 0 so that it can be used to represent nullptr for optional object values
         // in the wire format. However don't tag it as allocated so that it is an error to ask
         // KnownObjects for ID 0.
@@ -110,30 +110,35 @@
 
     // Get a backend objects for a given client ID.
     // Returns nullptr if the ID hasn't previously been allocated.
-    const Data* Get(uint32_t id, AllocationState expected = AllocationState::Allocated) const {
+    const Data* Get(uint32_t id) const {
         if (id >= mKnown.size()) {
             return nullptr;
         }
 
         const Data* data = &mKnown[id];
-
-        if (data->state != expected) {
+        if (data->state != AllocationState::Allocated) {
             return nullptr;
         }
-
         return data;
     }
-    Data* Get(uint32_t id, AllocationState expected = AllocationState::Allocated) {
+    Data* Get(uint32_t id) {
         if (id >= mKnown.size()) {
             return nullptr;
         }
 
         Data* data = &mKnown[id];
-
-        if (data->state != expected) {
+        if (data->state != AllocationState::Allocated) {
             return nullptr;
         }
+        return data;
+    }
 
+    Data* FillReservation(uint32_t id, T handle) {
+        ASSERT(id < mKnown.size());
+        Data* data = &mKnown[id];
+        ASSERT(data->state == AllocationState::Reserved);
+        data->handle = handle;
+        data->state = AllocationState::Allocated;
         return data;
     }
 
@@ -181,9 +186,9 @@
         return objects;
     }
 
-    std::vector<T> GetAllHandles() {
+    std::vector<T> GetAllHandles() const {
         std::vector<T> objects;
-        for (Data& data : mKnown) {
+        for (const Data& data : mKnown) {
             if (data.state == AllocationState::Allocated && data.handle != nullptr) {
                 objects.push_back(data.handle);
             }
@@ -192,10 +197,50 @@
         return objects;
     }
 
-  private:
+  protected:
     std::vector<Data> mKnown;
 };
 
+template <typename T>
+class KnownObjects : public KnownObjectsBase<T> {
+  public:
+    KnownObjects() = default;
+};
+
+template <>
+class KnownObjects<WGPUDevice> : public KnownObjectsBase<WGPUDevice> {
+  public:
+    KnownObjects() = default;
+
+    Data* Allocate(uint32_t id, AllocationState state = AllocationState::Allocated) {
+        Data* data = KnownObjectsBase<WGPUDevice>::Allocate(id, state);
+        AddToKnownSet(data);
+        return data;
+    }
+
+    Data* FillReservation(uint32_t id, WGPUDevice handle) {
+        Data* data = KnownObjectsBase<WGPUDevice>::FillReservation(id, handle);
+        AddToKnownSet(data);
+        return data;
+    }
+
+    void Free(uint32_t id) {
+        mKnownSet.erase(mKnown[id].handle);
+        KnownObjectsBase<WGPUDevice>::Free(id);
+    }
+
+    bool IsKnown(WGPUDevice device) const { return mKnownSet.count(device) != 0; }
+
+  private:
+    void AddToKnownSet(Data* data) {
+        if (data != nullptr && data->state == AllocationState::Allocated &&
+            data->handle != nullptr) {
+            mKnownSet.insert(data->handle);
+        }
+    }
+    std::unordered_set<WGPUDevice> mKnownSet;
+};
+
 // ObjectIds are lost in deserialization. Store the ids of deserialized
 // objects here so they can be used in command handlers. This is useful
 // for creating ReturnWireCmds which contain client ids
diff --git a/src/dawn/wire/server/Server.cpp b/src/dawn/wire/server/Server.cpp
index 6dc8358..cf6697d 100644
--- a/src/dawn/wire/server/Server.cpp
+++ b/src/dawn/wire/server/Server.cpp
@@ -155,6 +155,10 @@
     return data->handle;
 }
 
+bool Server::IsDeviceKnown(WGPUDevice device) const {
+    return DeviceObjects().IsKnown(device);
+}
+
 void Server::SetForwardingDeviceCallbacks(ObjectData<WGPUDevice>* deviceObject) {
     // Note: these callbacks are manually inlined here since they do not acquire and
     // free their userdata. Also unlike other callbacks, these are cleared and unset when
diff --git a/src/dawn/wire/server/Server.h b/src/dawn/wire/server/Server.h
index 9bf60c8..6aa0bea 100644
--- a/src/dawn/wire/server/Server.h
+++ b/src/dawn/wire/server/Server.h
@@ -170,6 +170,7 @@
     bool InjectInstance(WGPUInstance instance, uint32_t id, uint32_t generation);
 
     WGPUDevice GetDevice(uint32_t id, uint32_t generation);
+    bool IsDeviceKnown(WGPUDevice device) const;
 
     template <typename T,
               typename Enable = std::enable_if<std::is_base_of<CallbackUserdata, T>::value>>
diff --git a/src/dawn/wire/server/ServerAdapter.cpp b/src/dawn/wire/server/ServerAdapter.cpp
index 67e5d46..9735d26 100644
--- a/src/dawn/wire/server/ServerAdapter.cpp
+++ b/src/dawn/wire/server/ServerAdapter.cpp
@@ -50,11 +50,6 @@
                                      WGPURequestDeviceStatus status,
                                      WGPUDevice device,
                                      const char* message) {
-    auto* deviceObject = DeviceObjects().Get(data->deviceObjectId, AllocationState::Reserved);
-    // Should be impossible to fail. ObjectIds can't be freed by a destroy command until
-    // they move from Reserved to Allocated, or if they are destroyed here.
-    ASSERT(deviceObject != nullptr);
-
     ReturnAdapterRequestDeviceCallbackCmd cmd = {};
     cmd.adapter = data->adapter;
     cmd.requestSerial = data->requestSerial;
@@ -101,8 +96,10 @@
     cmd.limits = &limits;
 
     // Assign the handle and allocated status if the device is created successfully.
-    deviceObject->state = AllocationState::Allocated;
-    deviceObject->handle = device;
+    auto* deviceObject = DeviceObjects().FillReservation(data->deviceObjectId, device);
+    ASSERT(deviceObject != nullptr);
+    deviceObject->info->server = this;
+    deviceObject->info->self = ObjectHandle{data->deviceObjectId, deviceObject->generation};
     SetForwardingDeviceCallbacks(deviceObject);
 
     SerializeCommand(cmd);
diff --git a/src/dawn/wire/server/ServerDevice.cpp b/src/dawn/wire/server/ServerDevice.cpp
index 0a1fff5..91c70ea 100644
--- a/src/dawn/wire/server/ServerDevice.cpp
+++ b/src/dawn/wire/server/ServerDevice.cpp
@@ -19,21 +19,13 @@
 namespace {
 
 template <ObjectType objectType, typename Pipeline>
-void HandleCreateRenderPipelineAsyncCallbackResult(KnownObjects<Pipeline>* knownObjects,
-                                                   WGPUCreatePipelineAsyncStatus status,
-                                                   Pipeline pipeline,
-                                                   CreatePipelineAsyncUserData* data) {
-    // May be null if the device was destroyed. Device destruction destroys child
-    // objects on the wire.
-    auto* pipelineObject = knownObjects->Get(data->pipelineObjectID, AllocationState::Reserved);
-    // Should be impossible to fail. ObjectIds can't be freed by a destroy command until
-    // they move from Reserved to Allocated, or if they are destroyed here.
-    ASSERT(pipelineObject != nullptr);
-
+void HandleCreatePipelineAsyncCallback(KnownObjects<Pipeline>* knownObjects,
+                                       WGPUCreatePipelineAsyncStatus status,
+                                       Pipeline pipeline,
+                                       CreatePipelineAsyncUserData* data) {
     if (status == WGPUCreatePipelineAsyncStatus_Success) {
-        // Assign the handle and allocated status if the pipeline is created successfully.
-        pipelineObject->state = AllocationState::Allocated;
-        pipelineObject->handle = pipeline;
+        auto* pipelineObject = knownObjects->FillReservation(data->pipelineObjectID, pipeline);
+        ASSERT(pipelineObject != nullptr);
 
         // This should be impossible to fail. It would require a command to be sent that
         // creates a duplicate ObjectId, which would fail validation.
@@ -136,8 +128,8 @@
                                                   WGPUCreatePipelineAsyncStatus status,
                                                   WGPUComputePipeline pipeline,
                                                   const char* message) {
-    HandleCreateRenderPipelineAsyncCallbackResult<ObjectType::ComputePipeline>(
-        &ComputePipelineObjects(), status, pipeline, data);
+    HandleCreatePipelineAsyncCallback<ObjectType::ComputePipeline>(&ComputePipelineObjects(),
+                                                                   status, pipeline, data);
 
     ReturnDeviceCreateComputePipelineAsyncCallbackCmd cmd;
     cmd.device = data->device;
@@ -181,8 +173,8 @@
                                                  WGPUCreatePipelineAsyncStatus status,
                                                  WGPURenderPipeline pipeline,
                                                  const char* message) {
-    HandleCreateRenderPipelineAsyncCallbackResult<ObjectType::RenderPipeline>(
-        &RenderPipelineObjects(), status, pipeline, data);
+    HandleCreatePipelineAsyncCallback<ObjectType::RenderPipeline>(&RenderPipelineObjects(), status,
+                                                                  pipeline, data);
 
     ReturnDeviceCreateRenderPipelineAsyncCallbackCmd cmd;
     cmd.device = data->device;
diff --git a/src/dawn/wire/server/ServerInstance.cpp b/src/dawn/wire/server/ServerInstance.cpp
index d3998a3..dcdf62c 100644
--- a/src/dawn/wire/server/ServerInstance.cpp
+++ b/src/dawn/wire/server/ServerInstance.cpp
@@ -51,11 +51,6 @@
                                       WGPURequestAdapterStatus status,
                                       WGPUAdapter adapter,
                                       const char* message) {
-    auto* adapterObject = AdapterObjects().Get(data->adapterObjectId, AllocationState::Reserved);
-    // Should be impossible to fail. ObjectIds can't be freed by a destroy command until
-    // they move from Reserved to Allocated, or if they are destroyed here.
-    ASSERT(adapterObject != nullptr);
-
     ReturnInstanceRequestAdapterCallbackCmd cmd = {};
     cmd.instance = data->instance;
     cmd.requestSerial = data->requestSerial;
@@ -75,8 +70,8 @@
     std::vector<WGPUFeatureName> features;
 
     // Assign the handle and allocated status if the adapter is created successfully.
-    adapterObject->state = AllocationState::Allocated;
-    adapterObject->handle = adapter;
+    auto* adapterObject = AdapterObjects().FillReservation(data->adapterObjectId, adapter);
+    ASSERT(adapterObject != nullptr);
 
     size_t featuresCount = mProcs.adapterEnumerateFeatures(adapter, nullptr);
     features.resize(featuresCount);