dawn_wire: Fix a bug with multiple injected devices

Device child objects were storing an *unstable* pointer to device
specific tracking information. Fix this by moving the tracking
information to a stable heap allocation.

Bug: dawn:565
Change-Id: I00ad72563ac66e29314603e77698718953fcbf15
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/38280
Commit-Queue: Austin Eng <enga@chromium.org>
Reviewed-by: Corentin Wallez <cwallez@chromium.org>

src/dawn_wire/server/ObjectStorage.h

@@ -24,6 +24,10 @@
 
 namespace dawn_wire { namespace server {
 
+    struct DeviceInfo {
+        std::unordered_set<uint64_t> childObjectTypesAndIds;
+    };
+
     template <typename T>
     struct ObjectDataBase {
         // The backend-provided handle and generation to this object.
@@ -34,7 +38,8 @@ namespace dawn_wire { namespace server {
         // TODO(cwallez@chromium.org): make this an internal bit vector in KnownObjects.
         bool allocated;
 
-        ObjectDataBase<WGPUDevice>* device = nullptr;
+        // This points to an allocation that is owned by the device.
+        DeviceInfo* deviceInfo = nullptr;
     };
 
     // Stores what the backend knows about the type.
@@ -68,7 +73,9 @@ namespace dawn_wire { namespace server {
 
     template <>
     struct ObjectData<WGPUDevice> : public ObjectDataBase<WGPUDevice> {
-        std::unordered_set<uint64_t> childObjectTypesAndIds;
+        // Store |info| as a separate allocation so that its address does not move.
+        // The pointer to |info| is stored in device child objects.
+        std::unique_ptr<DeviceInfo> info = std::make_unique<DeviceInfo>();
     };
 
     // Keeps track of the mapping between client IDs and backend objects.

src/dawn_wire/server/Server.cpp

@@ -64,14 +64,14 @@ namespace dawn_wire { namespace server {
             return false;
         }
 
-        if (!TrackDeviceChild(device, ObjectType::Texture, id)) {
-            return false;
-        }
-
         data->handle = texture;
         data->generation = generation;
         data->allocated = true;
-        data->device = device;
+        data->deviceInfo = device->info.get();
+
+        if (!TrackDeviceChild(data->deviceInfo, ObjectType::Texture, 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.
@@ -131,9 +131,8 @@ namespace dawn_wire { namespace server {
         mProcs.deviceSetDeviceLostCallback(device, nullptr, nullptr);
     }
 
-    bool TrackDeviceChild(ObjectDataBase<WGPUDevice>* device, ObjectType type, ObjectId id) {
-        auto it = static_cast<ObjectData<WGPUDevice>*>(device)->childObjectTypesAndIds.insert(
-            PackObjectTypeAndId(type, id));
+    bool TrackDeviceChild(DeviceInfo* info, ObjectType type, ObjectId id) {
+        auto it = info->childObjectTypesAndIds.insert(PackObjectTypeAndId(type, id));
         if (!it.second) {
             // An object of this type and id already exists.
             return false;
@@ -141,8 +140,8 @@ namespace dawn_wire { namespace server {
         return true;
     }
 
-    bool UntrackDeviceChild(ObjectDataBase<WGPUDevice>* device, ObjectType type, ObjectId id) {
-        auto& children = static_cast<ObjectData<WGPUDevice>*>(device)->childObjectTypesAndIds;
+    bool UntrackDeviceChild(DeviceInfo* info, ObjectType type, ObjectId id) {
+        auto& children = info->childObjectTypesAndIds;
         auto it = children.find(PackObjectTypeAndId(type, id));
         if (it == children.end()) {
             // An object of this type and id was already deleted.

src/dawn_wire/server/Server.h

@@ -223,8 +223,8 @@ namespace dawn_wire { namespace server {
         std::shared_ptr<bool> mIsAlive;
     };
 
-    bool TrackDeviceChild(ObjectDataBase<WGPUDevice>* device, ObjectType type, ObjectId id);
-    bool UntrackDeviceChild(ObjectDataBase<WGPUDevice>* device, ObjectType type, ObjectId id);
+    bool TrackDeviceChild(DeviceInfo* device, ObjectType type, ObjectId id);
+    bool UntrackDeviceChild(DeviceInfo* device, ObjectType type, ObjectId id);
 
     std::unique_ptr<MemoryTransferService> CreateInlineMemoryTransferService();
 

src/dawn_wire/server/ServerBuffer.cpp

@@ -137,8 +137,8 @@ namespace dawn_wire { namespace server {
         }
         resultData->generation = bufferResult.generation;
         resultData->handle = mProcs.deviceCreateBuffer(device->handle, descriptor);
-        resultData->device = device;
-        if (!TrackDeviceChild(device, ObjectType::Buffer, bufferResult.id)) {
+        resultData->deviceInfo = device->info.get();
+        if (!TrackDeviceChild(resultData->deviceInfo, ObjectType::Buffer, bufferResult.id)) {
             return false;
         }
 

src/dawn_wire/server/ServerDevice.cpp

@@ -74,8 +74,9 @@ namespace dawn_wire { namespace server {
         }
 
         resultData->generation = pipelineObjectHandle.generation;
-        resultData->device = device;
-        if (!TrackDeviceChild(device, ObjectType::ComputePipeline, pipelineObjectHandle.id)) {
+        resultData->deviceInfo = device->info.get();
+        if (!TrackDeviceChild(resultData->deviceInfo, ObjectType::ComputePipeline,
+                              pipelineObjectHandle.id)) {
             return false;
         }
 
@@ -140,8 +141,9 @@ namespace dawn_wire { namespace server {
         }
 
         resultData->generation = pipelineObjectHandle.generation;
-        resultData->device = device;
-        if (!TrackDeviceChild(device, ObjectType::RenderPipeline, pipelineObjectHandle.id)) {
+        resultData->deviceInfo = device->info.get();
+        if (!TrackDeviceChild(resultData->deviceInfo, ObjectType::RenderPipeline,
+                              pipelineObjectHandle.id)) {
             return false;
         }
 

src/tests/unittests/wire/WireInjectDeviceTests.cpp

@@ -182,3 +182,49 @@ TEST_F(WireInjectDeviceTests, ReflectLiveDevices) {
     EXPECT_CALL(api, OnDeviceSetUncapturedErrorCallback(serverDevice2, nullptr, nullptr)).Times(1);
     EXPECT_CALL(api, OnDeviceSetDeviceLostCallback(serverDevice2, nullptr, nullptr)).Times(1);
 }
+
+// This is a regression test where a second device reservation invalidated pointers into the
+// KnownObjects std::vector of devices. The fix was to store pointers to heap allocated
+// objects instead.
+TEST_F(WireInjectDeviceTests, TrackChildObjectsWithTwoReservedDevices) {
+    // Reserve one device, inject it, and get the default queue.
+    ReservedDevice reservation1 = GetWireClient()->ReserveDevice();
+
+    WGPUDevice serverDevice1 = api.GetNewDevice();
+    EXPECT_CALL(api, DeviceReference(serverDevice1));
+    EXPECT_CALL(api, OnDeviceSetUncapturedErrorCallback(serverDevice1, _, _));
+    EXPECT_CALL(api, OnDeviceSetDeviceLostCallback(serverDevice1, _, _));
+    ASSERT_TRUE(
+        GetWireServer()->InjectDevice(serverDevice1, reservation1.id, reservation1.generation));
+
+    WGPUCommandEncoder commandEncoder =
+        wgpuDeviceCreateCommandEncoder(reservation1.device, nullptr);
+
+    WGPUCommandEncoder serverCommandEncoder = api.GetNewCommandEncoder();
+    EXPECT_CALL(api, DeviceCreateCommandEncoder(serverDevice1, _))
+        .WillOnce(Return(serverCommandEncoder));
+    FlushClient();
+
+    // Reserve a second device, and inject it.
+    ReservedDevice reservation2 = GetWireClient()->ReserveDevice();
+
+    WGPUDevice serverDevice2 = api.GetNewDevice();
+    EXPECT_CALL(api, DeviceReference(serverDevice2));
+    EXPECT_CALL(api, OnDeviceSetUncapturedErrorCallback(serverDevice2, _, _));
+    EXPECT_CALL(api, OnDeviceSetDeviceLostCallback(serverDevice2, _, _));
+    ASSERT_TRUE(
+        GetWireServer()->InjectDevice(serverDevice2, reservation2.id, reservation2.generation));
+
+    // Release the encoder. This should work without error because it stores a stable
+    // pointer to its device's list of child objects. On destruction, it removes itself from the
+    // list.
+    wgpuCommandEncoderRelease(commandEncoder);
+    EXPECT_CALL(api, CommandEncoderRelease(serverCommandEncoder));
+    FlushClient();
+
+    // Called on shutdown.
+    EXPECT_CALL(api, OnDeviceSetUncapturedErrorCallback(serverDevice1, nullptr, nullptr)).Times(1);
+    EXPECT_CALL(api, OnDeviceSetDeviceLostCallback(serverDevice1, nullptr, nullptr)).Times(1);
+    EXPECT_CALL(api, OnDeviceSetUncapturedErrorCallback(serverDevice2, nullptr, nullptr)).Times(1);
+    EXPECT_CALL(api, OnDeviceSetDeviceLostCallback(serverDevice2, nullptr, nullptr)).Times(1);
+}