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>

include/dawn/wire/WireServer.h

@@ -65,6 +65,11 @@ class DAWN_WIRE_EXPORT WireServer : public CommandHandler {
     // 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;
 };

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

@@ -177,6 +177,10 @@ TEST_F(WireAdapterTests, RequestDeviceSuccess) {
 
     // 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 @@ TEST_F(WireAdapterTests, RequestDeviceSuccess) {
                     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 @@ TEST_F(WireAdapterTests, RequestDeviceSuccess) {
         })));
     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

src/dawn/wire/WireServer.cpp

@@ -58,6 +58,10 @@ WGPUDevice WireServer::GetDevice(uint32_t id, uint32_t generation) {
     return mImpl->GetDevice(id, generation);
 }
 
+bool WireServer::IsDeviceKnown(WGPUDevice device) const {
+    return mImpl->IsDeviceKnown(device);
+}
+
 namespace server {
 MemoryTransferService::MemoryTransferService() = default;
 

src/dawn/wire/server/ObjectStorage.h

@@ -94,11 +94,11 @@ struct ObjectData<WGPUDevice> : public ObjectDataBase<WGPUDevice> {
 
 // 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 @@ class KnownObjects {
 
     // 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 != AllocationState::Allocated) {
-        if (data->state != expected) {
             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 != AllocationState::Allocated) {
-        if (data->state != expected) {
             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 @@ class KnownObjects {
         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 @@ class KnownObjects {
         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

src/dawn/wire/server/Server.cpp

@@ -155,6 +155,10 @@ WGPUDevice Server::GetDevice(uint32_t id, uint32_t generation) {
     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

src/dawn/wire/server/Server.h

@@ -170,6 +170,7 @@ class Server : public ServerBase {
     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>>

src/dawn/wire/server/ServerAdapter.cpp

@@ -50,11 +50,6 @@ void Server::OnRequestDeviceCallback(RequestDeviceUserdata* data,
                                      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 @@ void Server::OnRequestDeviceCallback(RequestDeviceUserdata* data,
     cmd.limits = &limits;
 
     // Assign the handle and allocated status if the device is created successfully.
-    deviceObject->state = AllocationState::Allocated;
+    auto* deviceObject = DeviceObjects().FillReservation(data->deviceObjectId, device);
-    deviceObject->handle = device;
+    ASSERT(deviceObject != nullptr);
+    deviceObject->info->server = this;
+    deviceObject->info->self = ObjectHandle{data->deviceObjectId, deviceObject->generation};
     SetForwardingDeviceCallbacks(deviceObject);
 
     SerializeCommand(cmd);

src/dawn/wire/server/ServerDevice.cpp

@@ -19,21 +19,13 @@ namespace dawn::wire::server {
 namespace {
 
 template <ObjectType objectType, typename Pipeline>
-void HandleCreateRenderPipelineAsyncCallbackResult(KnownObjects<Pipeline>* knownObjects,
+void HandleCreatePipelineAsyncCallback(KnownObjects<Pipeline>* knownObjects,
-                                                   WGPUCreatePipelineAsyncStatus status,
+                                       WGPUCreatePipelineAsyncStatus status,
-                                                   Pipeline pipeline,
+                                       Pipeline pipeline,
-                                                   CreatePipelineAsyncUserData* data) {
+                                       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);
-
     if (status == WGPUCreatePipelineAsyncStatus_Success) {
-        // Assign the handle and allocated status if the pipeline is created successfully.
+        auto* pipelineObject = knownObjects->FillReservation(data->pipelineObjectID, pipeline);
-        pipelineObject->state = AllocationState::Allocated;
+        ASSERT(pipelineObject != nullptr);
-        pipelineObject->handle = pipeline;
 
         // 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 @@ void Server::OnCreateComputePipelineAsyncCallback(CreatePipelineAsyncUserData* d
                                                   WGPUCreatePipelineAsyncStatus status,
                                                   WGPUComputePipeline pipeline,
                                                   const char* message) {
-    HandleCreateRenderPipelineAsyncCallbackResult<ObjectType::ComputePipeline>(
+    HandleCreatePipelineAsyncCallback<ObjectType::ComputePipeline>(&ComputePipelineObjects(),
-        &ComputePipelineObjects(), status, pipeline, data);
+                                                                   status, pipeline, data);
 
     ReturnDeviceCreateComputePipelineAsyncCallbackCmd cmd;
     cmd.device = data->device;
@@ -181,8 +173,8 @@ void Server::OnCreateRenderPipelineAsyncCallback(CreatePipelineAsyncUserData* da
                                                  WGPUCreatePipelineAsyncStatus status,
                                                  WGPURenderPipeline pipeline,
                                                  const char* message) {
-    HandleCreateRenderPipelineAsyncCallbackResult<ObjectType::RenderPipeline>(
+    HandleCreatePipelineAsyncCallback<ObjectType::RenderPipeline>(&RenderPipelineObjects(), status,
-        &RenderPipelineObjects(), status, pipeline, data);
+                                                                  pipeline, data);
 
     ReturnDeviceCreateRenderPipelineAsyncCallbackCmd cmd;
     cmd.device = data->device;

src/dawn/wire/server/ServerInstance.cpp

@@ -51,11 +51,6 @@ void Server::OnRequestAdapterCallback(RequestAdapterUserdata* data,
                                       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 @@ void Server::OnRequestAdapterCallback(RequestAdapterUserdata* data,
     std::vector<WGPUFeatureName> features;
 
     // Assign the handle and allocated status if the adapter is created successfully.
-    adapterObject->state = AllocationState::Allocated;
+    auto* adapterObject = AdapterObjects().FillReservation(data->adapterObjectId, adapter);
-    adapterObject->handle = adapter;
+    ASSERT(adapterObject != nullptr);
 
     size_t featuresCount = mProcs.adapterEnumerateFeatures(adapter, nullptr);
     features.resize(featuresCount);