dawn::wire::client Use Refcounted for ObjectBase.

In future CLs this will let objects references each other, without
breaking the semantic that Instance/Device are destroyed/lost
immediately when their external refcount reaches 0.

This changes the lifetime semantics in the Dawn wire client to have each
object manage its own allocation lifetime through refcounting.
Previously the dawn::wire::client::Client managed their lifetime inside
the ObjectStores by getting a signal from the object when it reached a
refcount of 0.

This CL also changes Instance and Device to use
RefCountedWithExternalCount so that they can get a signal to unregister
themselves from the wire when the external refcount gets to 0.

Future cleanups will come in follow-up CLs:
 - Removing the logic to destroy Device first on Client disconnects.
   as it is no longer needed to preserve semantics.
 - Replacing the Device::mIsAlive weak_ptr with refs to the device.
 - Removing the LinkedList used to keep the list of all ObjectBase.
 - <your suggestion here?>

Bug: 344963953
Change-Id: I277ceef4ef30291d16c75df543eeae1d6552aff1
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/197235
Reviewed-by: Austin Eng <enga@chromium.org>
Reviewed-by: Loko Kung <lokokung@google.com>
Commit-Queue: Corentin Wallez <cwallez@chromium.org>
diff --git a/generator/templates/dawn/wire/client/ApiProcs.cpp b/generator/templates/dawn/wire/client/ApiProcs.cpp
index 5a28ecd..23ab85b 100644
--- a/generator/templates/dawn/wire/client/ApiProcs.cpp
+++ b/generator/templates/dawn/wire/client/ApiProcs.cpp
@@ -112,11 +112,11 @@
     //* When an object's refcount reaches 0, notify the server side of it and delete it.
     DAWN_WIRE_EXPORT void {{as_cMethodNamespaced(type.name, Name("release"), Name('dawn wire client'))}}({{cType}} cObj) {
         {{Type}}* obj = reinterpret_cast<{{Type}}*>(cObj);
-        obj->Release();
+        obj->APIRelease();
     }
 
     DAWN_WIRE_EXPORT void {{as_cMethodNamespaced(type.name, Name("add ref"), Name('dawn wire client'))}}({{cType}} cObj) {
-        reinterpret_cast<{{Type}}*>(cObj)->AddRef();
+        reinterpret_cast<{{Type}}*>(cObj)->APIAddRef();
     }
 
 {% endfor %}
diff --git a/src/dawn/wire/client/Buffer.cpp b/src/dawn/wire/client/Buffer.cpp
index c73a236..066e120 100644
--- a/src/dawn/wire/client/Buffer.cpp
+++ b/src/dawn/wire/client/Buffer.cpp
@@ -439,8 +439,9 @@
       mDestructWriteHandleOnUnmap(descriptor->mappedAtCreation &&
                                   ((descriptor->usage & WGPUBufferUsage_MapWrite) == 0)) {}
 
-Buffer::~Buffer() {
+void Buffer::DeleteThis() {
     FreeMappedData();
+    ObjectWithEventsBase::DeleteThis();
 }
 
 ObjectType Buffer::GetObjectType() const {
@@ -495,8 +496,6 @@
                              size_t offset,
                              size_t size,
                              const WGPUBufferMapCallbackInfo& callbackInfo) {
-    DAWN_ASSERT(GetRefcount() != 0);
-
     Client* client = GetClient();
     auto [futureIDInternal, tracked] =
         GetEventManager().TrackEvent(std::make_unique<MapAsyncEvent>(callbackInfo, this));
@@ -543,8 +542,6 @@
                              size_t offset,
                              size_t size,
                              const WGPUBufferMapCallbackInfo2& callbackInfo) {
-    DAWN_ASSERT(GetRefcount() != 0);
-
     Client* client = GetClient();
     auto [futureIDInternal, tracked] =
         GetEventManager().TrackEvent(std::make_unique<MapAsyncEvent2>(callbackInfo, this));
diff --git a/src/dawn/wire/client/Buffer.h b/src/dawn/wire/client/Buffer.h
index 2462f7d..fa166b4 100644
--- a/src/dawn/wire/client/Buffer.h
+++ b/src/dawn/wire/client/Buffer.h
@@ -50,7 +50,7 @@
     Buffer(const ObjectBaseParams& params,
            const ObjectHandle& eventManagerHandle,
            const WGPUBufferDescriptor* descriptor);
-    ~Buffer() override;
+    void DeleteThis() override;
 
     ObjectType GetObjectType() const override;
 
diff --git a/src/dawn/wire/client/Client.cpp b/src/dawn/wire/client/Client.cpp
index b9b8e29..6ded2cd 100644
--- a/src/dawn/wire/client/Client.cpp
+++ b/src/dawn/wire/client/Client.cpp
@@ -74,12 +74,7 @@
     // which would be invalid if the queue was already freed.
     while (!mObjects[ObjectType::Device].empty()) {
         ObjectBase* object = mObjects[ObjectType::Device].head()->value();
-
-        UnregisterObjectCmd cmd;
-        cmd.objectType = ObjectType::Device;
-        cmd.objectId = object->GetWireId();
-        SerializeCommand(cmd);
-        mObjectStores[ObjectType::Device].Free(object);
+        object->Unregister();
     }
 
     for (auto& objectList : mObjects) {
@@ -89,12 +84,7 @@
         }
         while (!objectList.empty()) {
             ObjectBase* object = objectList.head()->value();
-
-            UnregisterObjectCmd cmd;
-            cmd.objectType = objectType;
-            cmd.objectId = object->GetWireId();
-            SerializeCommand(cmd);
-            mObjectStores[objectType].Free(object);
+            object->Unregister();
         }
     }
 }
@@ -147,23 +137,23 @@
 }
 
 void Client::ReclaimBufferReservation(const ReservedBuffer& reservation) {
-    Free(FromAPI(reservation.buffer));
+    ReclaimReservation(FromAPI(reservation.buffer));
 }
 
 void Client::ReclaimTextureReservation(const ReservedTexture& reservation) {
-    Free(FromAPI(reservation.texture));
+    ReclaimReservation(FromAPI(reservation.texture));
 }
 
 void Client::ReclaimSwapChainReservation(const ReservedSwapChain& reservation) {
-    Free(FromAPI(reservation.swapchain));
+    ReclaimReservation(FromAPI(reservation.swapchain));
 }
 
 void Client::ReclaimDeviceReservation(const ReservedDevice& reservation) {
-    Free(FromAPI(reservation.device));
+    ReclaimReservation(FromAPI(reservation.device));
 }
 
 void Client::ReclaimInstanceReservation(const ReservedInstance& reservation) {
-    Free(FromAPI(reservation.instance));
+    ReclaimReservation(FromAPI(reservation.instance));
 }
 
 EventManager& Client::GetEventManager(const ObjectHandle& instance) {
@@ -201,8 +191,18 @@
     return mDisconnected;
 }
 
-void Client::Free(ObjectBase* obj, ObjectType type) {
-    mObjectStores[type].Free(obj);
+void Client::Unregister(ObjectBase* obj, ObjectType type) {
+    UnregisterObjectCmd cmd;
+    cmd.objectType = type;
+    cmd.objectId = obj->GetWireId();
+    SerializeCommand(cmd);
+
+    ReclaimReservation(obj, type);
+}
+
+void Client::ReclaimReservation(ObjectBase* obj, ObjectType type) {
+    mObjectStores[type].Remove(obj);
+    obj->RemoveFromList();
 }
 
 }  // namespace dawn::wire::client
diff --git a/src/dawn/wire/client/Client.h b/src/dawn/wire/client/Client.h
index e382f40..e3c0bf5 100644
--- a/src/dawn/wire/client/Client.h
+++ b/src/dawn/wire/client/Client.h
@@ -64,17 +64,15 @@
         constexpr ObjectType type = ObjectTypeToTypeEnum<T>;
 
         ObjectBaseParams params = {this, mObjectStores[type].ReserveHandle()};
-        T* object = new T(params, std::forward<Args>(args)...);
+        Ref<T> object = AcquireRef(new T(params, std::forward<Args>(args)...));
 
-        mObjects[type].Append(object);
-        mObjectStores[type].Insert(std::unique_ptr<T>(object));
+        mObjects[type].Append(object.Get());
+        mObjectStores[type].Insert(object.Get());
 
-        Ref<T> ref;
-        ref.Acquire(object);
-        return ref;
+        return object;
     }
 
-    void Free(ObjectBase* obj, ObjectType type);
+    void Unregister(ObjectBase* obj, ObjectType type);
 
     template <typename T>
     T* Get(ObjectId id) {
@@ -115,10 +113,11 @@
 
   private:
     void UnregisterAllObjects();
+    void ReclaimReservation(ObjectBase* obj, ObjectType type);
 
     template <typename T>
-    void Free(T* obj) {
-        Free(obj, ObjectTypeToTypeEnum<T>);
+    void ReclaimReservation(T* obj) {
+        ReclaimReservation(obj, ObjectTypeToTypeEnum<T>);
     }
 
 #include "dawn/wire/client/ClientPrototypes_autogen.inc"
diff --git a/src/dawn/wire/client/Device.cpp b/src/dawn/wire/client/Device.cpp
index 527f358..0996167 100644
--- a/src/dawn/wire/client/Device.cpp
+++ b/src/dawn/wire/client/Device.cpp
@@ -253,7 +253,8 @@
 Device::Device(const ObjectBaseParams& params,
                const ObjectHandle& eventManagerHandle,
                const WGPUDeviceDescriptor* descriptor)
-    : ObjectWithEventsBase(params, eventManagerHandle), mIsAlive(std::make_shared<bool>(true)) {
+    : RefCountedWithExternalCount<ObjectWithEventsBase>(params, eventManagerHandle),
+      mIsAlive(std::make_shared<bool>(true)) {
 #if defined(DAWN_ENABLE_ASSERTS)
     static constexpr WGPUDeviceLostCallbackInfo2 kDefaultDeviceLostCallbackInfo = {
         nullptr, WGPUCallbackMode_AllowSpontaneous,
@@ -321,15 +322,9 @@
     return ObjectType::Device;
 }
 
-uint32_t Device::Release() {
-    // The device always has a reference in it's DeviceLossEvent which is created at construction,
-    // so when we drop to 1, we want to set the event so that the device can be loss according to
-    // the callback mode.
-    uint32_t refCount = ObjectBase::Release();
-    if (refCount == 1) {
-        HandleDeviceLost(WGPUDeviceLostReason_Destroyed, "Device was destroyed.");
-    }
-    return refCount;
+void Device::WillDropLastExternalRef() {
+    HandleDeviceLost(WGPUDeviceLostReason_Destroyed, "Device was destroyed.");
+    Unregister();
 }
 
 WGPUStatus Device::GetLimits(WGPUSupportedLimits* limits) const {
diff --git a/src/dawn/wire/client/Device.h b/src/dawn/wire/client/Device.h
index 6212e1d..7a395ee 100644
--- a/src/dawn/wire/client/Device.h
+++ b/src/dawn/wire/client/Device.h
@@ -31,6 +31,7 @@
 #include <memory>
 
 #include "dawn/common/LinkedList.h"
+#include "dawn/common/RefCountedWithExternalCount.h"
 #include "dawn/webgpu.h"
 #include "dawn/wire/WireCmd_autogen.h"
 #include "dawn/wire/client/ApiObjects_autogen.h"
@@ -43,7 +44,7 @@
 class Client;
 class Queue;
 
-class Device final : public ObjectWithEventsBase {
+class Device final : public RefCountedWithExternalCount<ObjectWithEventsBase> {
   public:
     explicit Device(const ObjectBaseParams& params,
                     const ObjectHandle& eventManagerHandle,
@@ -51,9 +52,6 @@
 
     ObjectType GetObjectType() const override;
 
-    // Override the default Release implementation to handle the device lost event.
-    uint32_t Release();
-
     void SetUncapturedErrorCallback(WGPUErrorCallback errorCallback, void* errorUserdata);
     void SetLoggingCallback(WGPULoggingCallback errorCallback, void* errorUserdata);
     void SetDeviceLostCallback(WGPUDeviceLostCallback errorCallback, void* errorUserdata);
@@ -99,6 +97,7 @@
     class DeviceLostEvent;
 
   private:
+    void WillDropLastExternalRef() override;
     template <typename Event,
               typename Cmd,
               typename CallbackInfo = typename Event::CallbackInfo,
diff --git a/src/dawn/wire/client/Instance.cpp b/src/dawn/wire/client/Instance.cpp
index f322d0d..5c760cd 100644
--- a/src/dawn/wire/client/Instance.cpp
+++ b/src/dawn/wire/client/Instance.cpp
@@ -146,10 +146,12 @@
 
 // Instance
 
-Instance::Instance(const ObjectBaseParams& params) : ObjectWithEventsBase(params, params.handle) {}
+Instance::Instance(const ObjectBaseParams& params)
+    : RefCountedWithExternalCount<ObjectWithEventsBase>(params, params.handle) {}
 
-Instance::~Instance() {
+void Instance::WillDropLastExternalRef() {
     GetEventManager().TransitionTo(EventManager::State::InstanceDropped);
+    Unregister();
 }
 
 ObjectType Instance::GetObjectType() const {
diff --git a/src/dawn/wire/client/Instance.h b/src/dawn/wire/client/Instance.h
index c309047..672dac0 100644
--- a/src/dawn/wire/client/Instance.h
+++ b/src/dawn/wire/client/Instance.h
@@ -29,16 +29,16 @@
 #define SRC_DAWN_WIRE_CLIENT_INSTANCE_H_
 
 #include "absl/container/flat_hash_set.h"
+#include "dawn/common/RefCountedWithExternalCount.h"
 #include "dawn/wire/WireClient.h"
 #include "dawn/wire/WireCmd_autogen.h"
 #include "dawn/wire/client/ObjectBase.h"
 
 namespace dawn::wire::client {
 
-class Instance final : public ObjectWithEventsBase {
+class Instance final : public RefCountedWithExternalCount<ObjectWithEventsBase> {
   public:
     explicit Instance(const ObjectBaseParams& params);
-    ~Instance() override;
 
     ObjectType GetObjectType() const override;
 
@@ -61,6 +61,7 @@
     size_t EnumerateWGSLLanguageFeatures(WGPUWGSLFeatureName* features) const;
 
   private:
+    void WillDropLastExternalRef() override;
     void GatherWGSLFeatures(const WGPUDawnWireWGSLControl* wgslControl,
                             const WGPUDawnWGSLBlocklist* wgslBlocklist);
 
diff --git a/src/dawn/wire/client/ObjectBase.cpp b/src/dawn/wire/client/ObjectBase.cpp
index 6480b86..c9ef383 100644
--- a/src/dawn/wire/client/ObjectBase.cpp
+++ b/src/dawn/wire/client/ObjectBase.cpp
@@ -33,47 +33,46 @@
 namespace dawn::wire::client {
 
 ObjectBase::ObjectBase(const ObjectBaseParams& params)
-    : mClient(params.client), mHandle(params.handle), mRefcount(1) {}
+    : mClient(params.client), mHandle(params.handle) {
+    DAWN_ASSERT(IsRegistered());
+}
 
-ObjectBase::~ObjectBase() {
-    RemoveFromList();
+bool ObjectBase::IsRegistered() const {
+    return mClient != nullptr;
 }
 
 const ObjectHandle& ObjectBase::GetWireHandle() const {
+    DAWN_ASSERT(IsRegistered());
     return mHandle;
 }
 
 ObjectId ObjectBase::GetWireId() const {
+    DAWN_ASSERT(IsRegistered());
     return mHandle.id;
 }
 
 ObjectGeneration ObjectBase::GetWireGeneration() const {
+    DAWN_ASSERT(IsRegistered());
     return mHandle.generation;
 }
 
 Client* ObjectBase::GetClient() const {
+    DAWN_ASSERT(IsRegistered());
     return mClient;
 }
 
-void ObjectBase::AddRef() {
-    mRefcount++;
+void ObjectBase::DeleteThis() {
+    Unregister();
+    RefCounted::DeleteThis();
 }
 
-uint32_t ObjectBase::Release() {
-    DAWN_ASSERT(mRefcount != 0);
-
-    uint32_t refCount = --mRefcount;
-    if (refCount == 0) {
-        UnregisterObjectCmd cmd;
-        cmd.objectType = GetObjectType();
-        cmd.objectId = GetWireId();
-
-        Client* client = GetClient();
-        client->SerializeCommand(cmd);
-        client->Free(this, GetObjectType());
+void ObjectBase::Unregister() {
+    if (!IsRegistered()) {
+        return;
     }
 
-    return refCount;
+    mClient->Unregister(this, GetObjectType());
+    mClient = nullptr;
 }
 
 ObjectWithEventsBase::ObjectWithEventsBase(const ObjectBaseParams& params,
diff --git a/src/dawn/wire/client/ObjectBase.h b/src/dawn/wire/client/ObjectBase.h
index 99121a6..e1bf928 100644
--- a/src/dawn/wire/client/ObjectBase.h
+++ b/src/dawn/wire/client/ObjectBase.h
@@ -32,7 +32,8 @@
 #include "partition_alloc/pointers/raw_ptr.h"
 
 #include "dawn/common/LinkedList.h"
-#include "dawn/common/RefBase.h"
+#include "dawn/common/Ref.h"
+#include "dawn/common/RefCounted.h"
 #include "dawn/wire/ObjectHandle.h"
 #include "dawn/wire/ObjectType_autogen.h"
 #include "dawn/wire/client/EventManager.h"
@@ -51,29 +52,29 @@
 //  - The external reference count, starting at 1.
 //  - An ID that is used to refer to this object when talking with the server side
 //  - A next/prev pointer. They are part of a linked list of objects of the same type.
-class ObjectBase : public LinkNode<ObjectBase> {
+class ObjectBase : public RefCounted, public LinkNode<ObjectBase> {
   public:
     explicit ObjectBase(const ObjectBaseParams& params);
-    virtual ~ObjectBase();
 
     virtual void CancelCallbacksForDisconnect() {}
     virtual ObjectType GetObjectType() const = 0;
 
+    // Objects are assumed to be registered with the wire on creation but can be unregistered
+    // with Unregister(). After that is done, other ObjectBase getters are invalid to use.
+    bool IsRegistered() const;
+    void Unregister();
+
     const ObjectHandle& GetWireHandle() const;
     ObjectId GetWireId() const;
     ObjectGeneration GetWireGeneration() const;
     Client* GetClient() const;
 
-    void AddRef();
-    uint32_t Release();
-
   protected:
-    uint32_t GetRefcount() const { return mRefcount; }
+    void DeleteThis() override;
 
   private:
-    const raw_ptr<Client> mClient;
+    raw_ptr<Client> mClient;
     const ObjectHandle mHandle;
-    uint32_t mRefcount;
 };
 
 // Compositable functionality for objects on the client side that need to have access to the event
@@ -93,27 +94,14 @@
     ObjectHandle mEventManagerHandle;
 };
 
-// Ref<T> for a T that's an ObjectBase*
-namespace detail {
-
-template <typename T>
-struct ObjectBaseTraits {
-    static constexpr T* kNullValue = nullptr;
-    static void AddRef(T* value) { value->AddRef(); }
-    static void Release(T* value) { value->Release(); }
-};
-
-}  // namespace detail
-
-template <typename T>
-class Ref : public RefBase<T*, detail::ObjectBaseTraits<T>> {
-  public:
-    using RefBase<T*, detail::ObjectBaseTraits<T>>::RefBase;
-};
-
-template <typename T>
-auto ReturnToAPI(Ref<T>&& r) {
-    return ToAPI(r.Detach());
+template <class T>
+auto ReturnToAPI(Ref<T>&& object) {
+    if constexpr (T::HasExternalRefCount) {
+        // For an object which has external ref count, just need to increase the external ref count,
+        // and keep the total ref count unchanged.
+        object->IncrementExternalRefCount();
+    }
+    return ToAPI(object.Detach());
 }
 
 }  // namespace dawn::wire::client
diff --git a/src/dawn/wire/client/ObjectStore.cpp b/src/dawn/wire/client/ObjectStore.cpp
index 48f8a48..08ba876 100644
--- a/src/dawn/wire/client/ObjectStore.cpp
+++ b/src/dawn/wire/client/ObjectStore.cpp
@@ -47,22 +47,22 @@
     return handle;
 }
 
-void ObjectStore::Insert(std::unique_ptr<ObjectBase> obj) {
+void ObjectStore::Insert(ObjectBase* obj) {
     ObjectId id = obj->GetWireId();
 
     if (id >= mObjects.size()) {
         DAWN_ASSERT(id == mObjects.size());
-        mObjects.emplace_back(std::move(obj));
+        mObjects.emplace_back(obj);
     } else {
         // The generation should never overflow. We don't recycle ObjectIds that would
         // overflow their next generation.
         DAWN_ASSERT(obj->GetWireGeneration() != 0);
         DAWN_ASSERT(mObjects[id] == nullptr);
-        mObjects[id] = std::move(obj);
+        mObjects[id] = obj;
     }
 }
 
-void ObjectStore::Free(ObjectBase* obj) {
+void ObjectStore::Remove(ObjectBase* obj) {
     DAWN_ASSERT(obj->IsInList());
     // The wire reuses ID for objects to keep them in a packed array starting from 0.
     // To avoid issues with asynchronous server->client communication referring to an ID that's
@@ -79,7 +79,7 @@
     if (id >= mObjects.size()) {
         return nullptr;
     }
-    return mObjects[id].get();
+    return mObjects[id];
 }
 
 }  // namespace dawn::wire::client
diff --git a/src/dawn/wire/client/ObjectStore.h b/src/dawn/wire/client/ObjectStore.h
index 19e7a00..9bf3bd2 100644
--- a/src/dawn/wire/client/ObjectStore.h
+++ b/src/dawn/wire/client/ObjectStore.h
@@ -49,14 +49,14 @@
     ObjectStore();
 
     ObjectHandle ReserveHandle();
-    void Insert(std::unique_ptr<ObjectBase> obj);
-    void Free(ObjectBase* obj);
+    void Insert(ObjectBase* obj);
+    void Remove(ObjectBase* obj);
     ObjectBase* Get(ObjectId id) const;
 
   private:
     uint32_t mCurrentId;
     std::vector<ObjectHandle> mFreeHandles;
-    std::vector<std::unique_ptr<ObjectBase>> mObjects;
+    std::vector<raw_ptr<ObjectBase>> mObjects;
 };
 
 }  // namespace dawn::wire::client