Reland "[dawn][wire] Handle command chunking in the wire in more robustly."

The new assertion added is too strict for the fuzzer, so removes it.
Note that we already handled the case, I just thought that it would be
ok to assert as well, but that causes the fuzzer to crash which is
undesirable.

This is a reland of commit 11b115b26569c8aa46781b4dac589ce17f5d6c89

Original change's description:
> [dawn][wire] Handle command chunking in the wire in more robustly.
>
> - With the changes to make the wire more spontaneous, chunked
>   command serialization and handling need to be updated to work
>   in a multithreaded manner. Previously, the handler assumes that
>   once a chunked command is received, it will receive the entire
>   chunked command before it receives any other commands. This may
>   no longer be true, so we upgrade chunked commands into its own
>   command so with unique id's to ensure that we handle them
>   appropriately.
> - This change also unifies the command id enums in the wire to a
>   single enum since chunked command handling may occur on both
>   the server and the client (though in practice in Chromium,
>   only the server->client may be multithreaded for now.) This
>   allows us to handle commands in a unified way and also ensures
>   that the command ids are now unique both direction.
>
> Bug: 412761856
> Change-Id: Ia1eb34d57c2e2fe0457c70aae9348bc051fce573
> Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/266475
> Reviewed-by: Corentin Wallez <cwallez@chromium.org>
> Reviewed-by: Kai Ninomiya <kainino@chromium.org>
> Commit-Queue: Loko Kung <lokokung@google.com>

Bug: b:412761856 b:452575709
Change-Id: Ie9c6554908597362cc2576182e09bf49fc35a4b9
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/267295
Commit-Queue: Kai Ninomiya <kainino@chromium.org>
Reviewed-by: Kai Ninomiya <kainino@chromium.org>
diff --git a/generator/dawn_json_generator.py b/generator/dawn_json_generator.py
index 4bd6757..dd55164 100644
--- a/generator/dawn_json_generator.py
+++ b/generator/dawn_json_generator.py
@@ -701,6 +701,7 @@
 
     commands = []
     return_commands = []
+    special_commands = []
 
     wire_json['special items']['client_handwritten_commands'] += wire_json[
         'special items']['client_side_commands']
@@ -766,9 +767,14 @@
         return_commands.append(
             Command(name, linked_record_members(json_data, types)))
 
+    for (name, json_data) in wire_json['special commands'].items():
+        special_commands.append(
+            Command(name, linked_record_members(json_data, types)))
+
     wire_params['cmd_records'] = {
         'command': commands,
-        'return command': return_commands
+        'return command': return_commands,
+        'special command': special_commands
     }
 
     for commands in wire_params['cmd_records'].values():
diff --git a/generator/templates/dawn/wire/WireCmd.cpp b/generator/templates/dawn/wire/WireCmd.cpp
index f12db7f..9714952 100644
--- a/generator/templates/dawn/wire/WireCmd.cpp
+++ b/generator/templates/dawn/wire/WireCmd.cpp
@@ -136,7 +136,7 @@
                       "Record must be at most one of is_cmd, extensible, and chained.");
         {% if is_cmd %}
             //* Start the transfer structure with the command ID, so that casting to WireCmd gives the ID.
-            {{Return}}WireCmd commandId;
+            WireCmd commandId;
         {% elif record.extensible %}
             WGPUBool hasNextInChain;
         {% elif record.chained %}
@@ -242,7 +242,7 @@
     ) {
         //* Handle special transfer members of methods.
         {% if is_cmd %}
-            transfer->commandId = {{Return}}WireCmd::{{name}};
+            transfer->commandId = WireCmd::{{Return}}{{name}};
         {% endif %}
 
         {% if record.extensible %}
@@ -348,7 +348,7 @@
             , const ObjectIdResolver& resolver
         {%- endif -%}) {
         {% if is_cmd %}
-            DAWN_ASSERT(transfer->commandId == {{Return}}WireCmd::{{name}});
+            DAWN_ASSERT(transfer->commandId == WireCmd::{{Return}}{{name}});
         {% endif %}
         {% if record.derived_method %}
             record->selfId = transfer->self;
@@ -690,6 +690,12 @@
     {% do sTypes.append(sType) %}
 {% endfor %}
 
+//* Output [de]serialization helpers for special commands
+{% for command in cmd_records["special command"] %}
+    {%- set name = command.name.CamelCase() -%}
+    {{write_record_serialization_helpers(command, name, command.members, is_cmd=True)}}
+{% endfor %}
+
 //* Output [de]serialization helpers for commands
 {% for command in cmd_records["command"] %}
     {%- set name = command.name.CamelCase() -%}
@@ -735,6 +741,10 @@
 
 }  // anonymous namespace
 
+{% for command in cmd_records["special command"] -%}
+    {{write_command_serialization_methods(command, False)}}
+{% endfor %}
+
 {% for command in cmd_records["command"] -%}
     {{write_command_serialization_methods(command, False)}}
 {% endfor %}
diff --git a/generator/templates/dawn/wire/WireCmd.h b/generator/templates/dawn/wire/WireCmd.h
index ba1542a..bd42e2a 100644
--- a/generator/templates/dawn/wire/WireCmd.h
+++ b/generator/templates/dawn/wire/WireCmd.h
@@ -63,17 +63,20 @@
             {% endfor %}
     };
 
-    //* Enum used as a prefix to each command on the wire format.
     enum class WireCmd : uint32_t {
+        // Enums used in special wire commands.
+        {% for command in cmd_records["special command"] %}
+            {{command.name.CamelCase()}},
+        {% endfor %}
+
+        // Enums used on the wire format.
         {% for command in cmd_records["command"] %}
             {{command.name.CamelCase()}},
         {% endfor %}
-    };
 
-    //* Enum used as a prefix to each command on the return wire format.
-    enum class ReturnWireCmd : uint32_t {
+        // Enums used on the return wire format.
         {% for command in cmd_records["return command"] %}
-            {{command.name.CamelCase()}},
+            Return{{command.name.CamelCase()}},
         {% endfor %}
     };
 
@@ -117,6 +120,10 @@
     };
 {% endmacro %}
 
+    {% for command in cmd_records["special command"] %}
+        {{write_command_struct(command, False)}}
+    {% endfor %}
+
     {% for command in cmd_records["command"] %}
         {{write_command_struct(command, False)}}
     {% endfor %}
diff --git a/generator/templates/dawn/wire/client/ClientHandlers.cpp b/generator/templates/dawn/wire/client/ClientHandlers.cpp
index dc32d84..5c741fc 100644
--- a/generator/templates/dawn/wire/client/ClientHandlers.cpp
+++ b/generator/templates/dawn/wire/client/ClientHandlers.cpp
@@ -34,7 +34,7 @@
     {% for command in cmd_records["return command"] %}
         WireResult Client::Handle{{command.name.CamelCase()}}(DeserializeBuffer* deserializeBuffer) {
             Return{{command.name.CamelCase()}}Cmd cmd;
-            WIRE_TRY(cmd.Deserialize(deserializeBuffer, &mWireCommandAllocator));
+            WIRE_TRY(cmd.Deserialize(deserializeBuffer, &mAllocator));
 
             {% for member in command.members if member.handle_type %}
                 {% set Type = member.handle_type.name.CamelCase() %}
@@ -61,37 +61,34 @@
         }
     {% endfor %}
 
-    const volatile char* Client::HandleCommandsImpl(const volatile char* commands, size_t size) {
+    const volatile char* Client::HandleCommands(const volatile char* commands, size_t size) {
         DeserializeBuffer deserializeBuffer(commands, size);
 
-        while (deserializeBuffer.AvailableSize() >= sizeof(CmdHeader) + sizeof(ReturnWireCmd)) {
-            // Start by chunked command handling, if it is done, then it means the whole buffer
-            // was consumed by it, so we return a pointer to the end of the commands.
-            switch (HandleChunkedCommands(deserializeBuffer.Buffer(), deserializeBuffer.AvailableSize())) {
-                case ChunkedCommandsResult::Consumed:
-                    return commands + size;
-                case ChunkedCommandsResult::Error:
-                    return nullptr;
-                case ChunkedCommandsResult::Passthrough:
-                    break;
-            }
-
-            ReturnWireCmd cmdId = *static_cast<const volatile ReturnWireCmd*>(static_cast<const volatile void*>(
+        while (deserializeBuffer.AvailableSize() >= sizeof(CmdHeader) + sizeof(WireCmd)) {
+            WireCmd cmdId = *static_cast<const volatile WireCmd*>(static_cast<const volatile void*>(
                 deserializeBuffer.Buffer() + sizeof(CmdHeader)));
             WireResult result = WireResult::FatalError;
             switch (cmdId) {
-                {% for command in cmd_records["return command"] %}
+                {% for command in cmd_records["special command"] %}
                     {% set Suffix = command.name.CamelCase() %}
-                    case ReturnWireCmd::{{Suffix}}:
+                    case WireCmd::{{Suffix}}:
                         result = Handle{{Suffix}}(&deserializeBuffer);
                         break;
                 {% endfor %}
+                {% for command in cmd_records["return command"] %}
+                    {% set Suffix = command.name.CamelCase() %}
+                    case WireCmd::Return{{Suffix}}:
+                        result = Handle{{Suffix}}(&deserializeBuffer);
+                        break;
+                {% endfor %}
+                default:
+                    result = WireResult::FatalError;
             }
 
             if (result != WireResult::Success) {
                 return nullptr;
             }
-            mWireCommandAllocator.Reset();
+            mAllocator.Reset();
         }
 
         if (deserializeBuffer.AvailableSize() != 0) {
diff --git a/generator/templates/dawn/wire/client/ClientPrototypes.inc b/generator/templates/dawn/wire/client/ClientPrototypes.inc
index 12907e1..576a0a3 100644
--- a/generator/templates/dawn/wire/client/ClientPrototypes.inc
+++ b/generator/templates/dawn/wire/client/ClientPrototypes.inc
@@ -25,14 +25,11 @@
 //* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 //* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
-//* Return command handlers
+// Command handlers & doers
 {% for command in cmd_records["return command"] %}
-    WireResult Handle{{command.name.CamelCase()}}(DeserializeBuffer* deserializeBuffer);
-{% endfor %}
-
-//* Return command doers
-{% for command in cmd_records["return command"] %}
-    WireResult Do{{command.name.CamelCase()}}(
+    {% set Suffix = command.name.CamelCase() %}
+    WireResult Handle{{Suffix}}(DeserializeBuffer* deserializeBuffer);
+    WireResult Do{{Suffix}}(
         {%- for member in command.members -%}
             {%- if member.handle_type -%}
                 {{as_wireType(member.handle_type)}} {{as_varName(member.name)}}
@@ -42,4 +39,5 @@
             {%- if not loop.last -%}, {% endif %}
         {%- endfor -%}
     );
+
 {% endfor %}
diff --git a/generator/templates/dawn/wire/server/ServerHandlers.cpp b/generator/templates/dawn/wire/server/ServerHandlers.cpp
index c954823..57876a2 100644
--- a/generator/templates/dawn/wire/server/ServerHandlers.cpp
+++ b/generator/templates/dawn/wire/server/ServerHandlers.cpp
@@ -87,26 +87,15 @@
         }
     {% endfor %}
 
-    const volatile char* Server::HandleCommandsImpl(const volatile char* commands, size_t size) {
+    const volatile char* Server::HandleCommands(const volatile char* commands, size_t size) {
         DeserializeBuffer deserializeBuffer(commands, size);
 
         while (deserializeBuffer.AvailableSize() >= sizeof(CmdHeader) + sizeof(WireCmd)) {
-            // Start by chunked command handling, if it is done, then it means the whole buffer
-            // was consumed by it, so we return a pointer to the end of the commands.
-            switch (HandleChunkedCommands(deserializeBuffer.Buffer(), deserializeBuffer.AvailableSize())) {
-                case ChunkedCommandsResult::Consumed:
-                    return commands + size;
-                case ChunkedCommandsResult::Error:
-                    return nullptr;
-                case ChunkedCommandsResult::Passthrough:
-                    break;
-            }
-
             WireCmd cmdId = *static_cast<const volatile WireCmd*>(static_cast<const volatile void*>(
                 deserializeBuffer.Buffer() + sizeof(CmdHeader)));
             WireResult result;
             switch (cmdId) {
-                {% for command in cmd_records["command"] %}
+                {% for command in cmd_records["special command"] + cmd_records["command"] %}
                     case WireCmd::{{command.name.CamelCase()}}:
                         result = Handle{{command.name.CamelCase()}}(&deserializeBuffer);
                         break;
diff --git a/generator/templates/dawn/wire/server/ServerPrototypes.inc b/generator/templates/dawn/wire/server/ServerPrototypes.inc
index f7404df..7cdc7c3 100644
--- a/generator/templates/dawn/wire/server/ServerPrototypes.inc
+++ b/generator/templates/dawn/wire/server/ServerPrototypes.inc
@@ -29,7 +29,6 @@
 {% for command in cmd_records["command"] %}
     {% set Suffix = command.name.CamelCase() %}
     WireResult Handle{{Suffix}}(DeserializeBuffer* deserializeBuffer);
-
     WireResult Do{{Suffix}}(
         {%- for member in command.members -%}
             {%- if member.is_return_value -%}
@@ -46,6 +45,7 @@
             {%- if not loop.last -%}, {% endif %}
         {%- endfor -%}
     );
+
 {% endfor %}
 
 {% for CommandName in server_custom_pre_handler_commands %}
diff --git a/src/dawn/dawn_wire.json b/src/dawn/dawn_wire.json
index ad848fd..0fd2c01 100644
--- a/src/dawn/dawn_wire.json
+++ b/src/dawn/dawn_wire.json
@@ -221,6 +221,14 @@
             { "name": "features", "type": "feature name", "annotation": "const*", "length": "features count"}
         ]
     },
+    "special commands": {
+        "chunked command": [
+            {"name": "id", "type": "uint64_t"},
+            {"name": "size", "type": "uint64_t"},
+            {"name": "chunk data", "type": "char", "annotation": "const*", "length": "chunk size", "wire_is_data_only": true},
+            {"name": "chunk size", "type": "uint32_t"}
+        ]
+    },
     "special items": {
         "client_side_structures": [
             "FutureWaitInfo",
diff --git a/src/dawn/tests/unittests/wire/WireQueueTests.cpp b/src/dawn/tests/unittests/wire/WireQueueTests.cpp
index 1751551..aab12c6 100644
--- a/src/dawn/tests/unittests/wire/WireQueueTests.cpp
+++ b/src/dawn/tests/unittests/wire/WireQueueTests.cpp
@@ -25,6 +25,7 @@
 // OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
+#include <cstring>
 #include <memory>
 
 #include "dawn/common/StringViewUtils.h"
@@ -45,6 +46,39 @@
 using testing::Return;
 using testing::SizedString;
 
+class WireWriteBufferTests : public WireTest {};
+
+// Tests that commands are serialized properly when they are too large and need to be chunked.
+TEST_F(WireWriteBufferTests, WriteBufferChunkedCommands) {
+    WGPUBuffer apiBuffer = api.GetNewBuffer();
+
+    // In order to create a command larger than the the maximum that is allowed to be serialized at
+    // a time for a single command to force command chunking, use a value larger than the maximum
+    // allocation size.
+    static size_t kLargeAllocationSize = GetC2SMaxAllocationSize() + 16u;
+
+    wgpu::BufferDescriptor desc = {};
+    desc.usage = wgpu::BufferUsage::CopyDst | wgpu::BufferUsage::MapRead;
+    desc.size = kLargeAllocationSize;
+    wgpu::Buffer buffer = device.CreateBuffer(&desc);
+    EXPECT_CALL(api, DeviceCreateBuffer(apiDevice, _))
+        .WillOnce(Return(apiBuffer))
+        .RetiresOnSaturation();
+    FlushClient();
+
+    auto expected = std::make_unique<uint8_t[]>(kLargeAllocationSize);
+    std::memset(expected.get(), 0b10101010, kLargeAllocationSize);
+    queue.WriteBuffer(buffer, 0, expected.get(), kLargeAllocationSize);
+
+    EXPECT_CALL(
+        api, QueueWriteBuffer(apiQueue, apiBuffer, 0, MatchesLambda([&](void const* actual) {
+                                  return !std::memcmp(expected.get(), actual, kLargeAllocationSize);
+                              }),
+                              kLargeAllocationSize))
+        .Times(1);
+    FlushClient();
+}
+
 using WireQueueTestBase = WireFutureTest<wgpu::QueueWorkDoneCallback<void>*>;
 class WireQueueTests : public WireQueueTestBase {
   protected:
diff --git a/src/dawn/tests/unittests/wire/WireTest.cpp b/src/dawn/tests/unittests/wire/WireTest.cpp
index 3d60204..41ef354 100644
--- a/src/dawn/tests/unittests/wire/WireTest.cpp
+++ b/src/dawn/tests/unittests/wire/WireTest.cpp
@@ -260,6 +260,10 @@
     return mWireClient.get();
 }
 
+size_t WireTest::GetC2SMaxAllocationSize() {
+    return mC2sBuf->GetMaximumAllocationSize();
+}
+
 void WireTest::DeleteServer() {
     EXPECT_CALL(api, QueueRelease(apiQueue)).Times(1);
     EXPECT_CALL(api, DeviceRelease(apiDevice)).Times(1);
diff --git a/src/dawn/tests/unittests/wire/WireTest.h b/src/dawn/tests/unittests/wire/WireTest.h
index 41fdbac..ab62ea6 100644
--- a/src/dawn/tests/unittests/wire/WireTest.h
+++ b/src/dawn/tests/unittests/wire/WireTest.h
@@ -180,6 +180,8 @@
     dawn::wire::WireServer* GetWireServer();
     dawn::wire::WireClient* GetWireClient();
 
+    size_t GetC2SMaxAllocationSize();
+
     void DeleteServer();
     void DeleteClient();
 
diff --git a/src/dawn/wire/ChunkedCommandHandler.cpp b/src/dawn/wire/ChunkedCommandHandler.cpp
index aa3de33..671d94f 100644
--- a/src/dawn/wire/ChunkedCommandHandler.cpp
+++ b/src/dawn/wire/ChunkedCommandHandler.cpp
@@ -44,56 +44,47 @@
 
 ChunkedCommandHandler::~ChunkedCommandHandler() = default;
 
-const volatile char* ChunkedCommandHandler::HandleCommands(const volatile char* commands,
-                                                           size_t size) {
-    if (mChunkedCommandRemainingSize > 0) {
-        // If there is a chunked command in flight, append the command data.
-        // We append at most |mChunkedCommandRemainingSize| which is enough to finish the
-        // in-flight chunked command, and then pass the rest along to a second call to
-        // |HandleCommandsImpl|.
-        size_t chunkSize = std::min(size, mChunkedCommandRemainingSize);
+WireResult ChunkedCommandHandler::HandleChunkedCommand(DeserializeBuffer* deserializeBuffer) {
+    ChunkedCommandCmd cmd;
+    WIRE_TRY(cmd.Deserialize(deserializeBuffer, &mAllocator));
 
-        memcpy(mChunkedCommandData.get() + mChunkedCommandPutOffset,
-               const_cast<const char*>(commands), chunkSize);
-        mChunkedCommandPutOffset += chunkSize;
-        mChunkedCommandRemainingSize -= chunkSize;
+    ChunkedCommand* chunkedCommand = nullptr;
+    if (auto it = mChunkedCommands.find(cmd.id); it != mChunkedCommands.end()) {
+        chunkedCommand = &(it->second);
+    } else {
+        ChunkedCommand newChunkedCommand = {};
+        newChunkedCommand.remainingSize = cmd.size;
+        newChunkedCommand.data.reset(AllocNoThrow<char>(cmd.size));
+        if (newChunkedCommand.data.get() == nullptr) {
+            return WireResult::FatalError;
+        }
 
-        commands += chunkSize;
-        size -= chunkSize;
+        const auto& [newIt, inserted] =
+            mChunkedCommands.insert({cmd.id, std::move(newChunkedCommand)});
+        DAWN_ASSERT(inserted);
+        chunkedCommand = &(newIt->second);
+    }
+    DAWN_ASSERT(chunkedCommand);
 
-        if (mChunkedCommandRemainingSize == 0) {
-            // Once the chunked command is complete, pass the data to the command handler
-            // implemenation.
-            auto chunkedCommandData = std::move(mChunkedCommandData);
-            if (HandleCommandsImpl(chunkedCommandData.get(), mChunkedCommandPutOffset) == nullptr) {
-                // |HandleCommandsImpl| returns nullptr on error. Forward any errors
-                // out.
-                return nullptr;
-            }
+    if (cmd.chunkSize > chunkedCommand->remainingSize) {
+        // If the chunk size is greater than the remaining size, something is wrong and we can no
+        // longer handle it, so just return a FatalError.
+        return WireResult::FatalError;
+    }
+    memcpy(chunkedCommand->data.get() + chunkedCommand->putOffset,
+           const_cast<const char*>(cmd.chunkData), cmd.chunkSize);
+    chunkedCommand->putOffset += cmd.chunkSize;
+    chunkedCommand->remainingSize -= cmd.chunkSize;
+
+    if (chunkedCommand->remainingSize == 0) {
+        ChunkedCommand fullCommand = std::move(*chunkedCommand);
+        mChunkedCommands.erase(cmd.id);
+        if (HandleCommands(fullCommand.data.get(), fullCommand.putOffset) == nullptr) {
+            return WireResult::FatalError;
         }
     }
 
-    return HandleCommandsImpl(commands, size);
-}
-
-ChunkedCommandHandler::ChunkedCommandsResult ChunkedCommandHandler::BeginChunkedCommandData(
-    const volatile char* commands,
-    size_t commandSize,
-    size_t initialSize) {
-    DAWN_ASSERT(!mChunkedCommandData);
-
-    // Reserve space for all the command data we're expecting, and copy the initial data
-    // to the start of the memory.
-    mChunkedCommandData.reset(AllocNoThrow<char>(commandSize));
-    if (!mChunkedCommandData) {
-        return ChunkedCommandsResult::Error;
-    }
-
-    memcpy(mChunkedCommandData.get(), const_cast<const char*>(commands), initialSize);
-    mChunkedCommandPutOffset = initialSize;
-    mChunkedCommandRemainingSize = commandSize - initialSize;
-
-    return ChunkedCommandsResult::Consumed;
+    return WireResult::Success;
 }
 
 }  // namespace dawn::wire
diff --git a/src/dawn/wire/ChunkedCommandHandler.h b/src/dawn/wire/ChunkedCommandHandler.h
index d5ab08e..050c813 100644
--- a/src/dawn/wire/ChunkedCommandHandler.h
+++ b/src/dawn/wire/ChunkedCommandHandler.h
@@ -32,9 +32,11 @@
 #include <limits>
 #include <memory>
 
+#include "absl/container/flat_hash_map.h"
 #include "dawn/common/Assert.h"
 #include "dawn/wire/Wire.h"
 #include "dawn/wire/WireCmd_autogen.h"
+#include "dawn/wire/WireDeserializeAllocator.h"
 
 namespace dawn::wire {
 
@@ -43,41 +45,21 @@
     ChunkedCommandHandler();
     ~ChunkedCommandHandler() override;
 
-    const volatile char* HandleCommands(const volatile char* commands, size_t size) override;
-
   protected:
-    enum class ChunkedCommandsResult {
-        Passthrough,
-        Consumed,
-        Error,
-    };
+    WireResult HandleChunkedCommand(DeserializeBuffer* deserializeBuffer);
 
-    // Returns |true| if the commands were entirely consumed into the chunked command vector
-    // and should be handled later once we receive all the command data.
-    // Returns |false| if commands should be handled now immediately.
-    ChunkedCommandsResult HandleChunkedCommands(const volatile char* commands, size_t size) {
-        uint64_t commandSize64 = reinterpret_cast<const volatile CmdHeader*>(commands)->commandSize;
-
-        if (commandSize64 > std::numeric_limits<size_t>::max()) {
-            return ChunkedCommandsResult::Error;
-        }
-        size_t commandSize = static_cast<size_t>(commandSize64);
-        if (size < commandSize) {
-            return BeginChunkedCommandData(commands, commandSize, size);
-        }
-        return ChunkedCommandsResult::Passthrough;
-    }
+    WireDeserializeAllocator mAllocator;
 
   private:
-    virtual const volatile char* HandleCommandsImpl(const volatile char* commands, size_t size) = 0;
-
-    ChunkedCommandsResult BeginChunkedCommandData(const volatile char* commands,
-                                                  size_t commandSize,
-                                                  size_t initialSize);
-
-    size_t mChunkedCommandRemainingSize = 0;
-    size_t mChunkedCommandPutOffset = 0;
-    std::unique_ptr<char[]> mChunkedCommandData;
+    // This map keeps track of all in-flight chunked commands. Note that because |HandleCommands|
+    // must be called in a thread-safe manner, we do not need to explicitly synchronize access to
+    // this map.
+    struct ChunkedCommand {
+        size_t remainingSize = 0;
+        size_t putOffset = 0;
+        std::unique_ptr<char[]> data = nullptr;
+    };
+    absl::flat_hash_map<uint64_t, ChunkedCommand> mChunkedCommands;
 };
 
 }  // namespace dawn::wire
diff --git a/src/dawn/wire/ChunkedCommandSerializer.cpp b/src/dawn/wire/ChunkedCommandSerializer.cpp
index f306495..51de9d5 100644
--- a/src/dawn/wire/ChunkedCommandSerializer.cpp
+++ b/src/dawn/wire/ChunkedCommandSerializer.cpp
@@ -32,6 +32,8 @@
 
 #include "dawn/wire/ChunkedCommandSerializer.h"
 
+#include "dawn/common/Assert.h"
+
 namespace dawn::wire {
 
 ChunkedCommandSerializer::ChunkedCommandSerializer(CommandSerializer* serializer)
@@ -39,22 +41,35 @@
     DAWN_ASSERT(mMaxAllocationSize > 0);
 }
 
+void ChunkedCommandSerializer::SetCommandSerializerForDisconnect(CommandSerializer* serializer) {
+    mSerializer = serializer;
+    mMaxAllocationSize = serializer->GetMaximumAllocationSize();
+    DAWN_ASSERT(mMaxAllocationSize > 0);
+}
+
 void ChunkedCommandSerializer::Flush() {
     mSerializer->Flush();
 }
 
 void ChunkedCommandSerializer::SerializeChunkedCommand(const char* allocatedBuffer,
-                                                       size_t remainingSize) {
-    while (remainingSize > 0) {
-        size_t chunkSize = std::min(remainingSize, mMaxAllocationSize);
-        void* dst = mSerializer->GetCmdSpace(chunkSize);
-        if (dst == nullptr) {
-            return;
-        }
-        memcpy(dst, allocatedBuffer, chunkSize);
+                                                       size_t totalSize) {
+    // Constant regarding the size of the WireChunkedCommandCmd that can be computed once.
+    static size_t kWireChunkedCmdPrefixSize = ChunkedCommandCmd{0, 0, nullptr, 0}.GetRequiredSize();
 
-        allocatedBuffer += chunkSize;
-        remainingSize -= chunkSize;
+    ChunkedCommandCmd cmd;
+    cmd.id = mNextChunkedCommandId++;
+    cmd.size = totalSize;
+
+    size_t remainingSize = totalSize;
+    while (remainingSize > 0) {
+        cmd.chunkData = allocatedBuffer;
+        cmd.chunkSize = std::min(remainingSize, mMaxAllocationSize - kWireChunkedCmdPrefixSize);
+        DAWN_ASSERT(cmd.GetRequiredSize() <= mMaxAllocationSize);
+
+        SerializeCommand(cmd);
+
+        allocatedBuffer += cmd.chunkSize;
+        remainingSize -= cmd.chunkSize;
     }
 }
 
diff --git a/src/dawn/wire/ChunkedCommandSerializer.h b/src/dawn/wire/ChunkedCommandSerializer.h
index 6194315..bc15d50 100644
--- a/src/dawn/wire/ChunkedCommandSerializer.h
+++ b/src/dawn/wire/ChunkedCommandSerializer.h
@@ -29,6 +29,7 @@
 #define SRC_DAWN_WIRE_CHUNKEDCOMMANDSERIALIZER_H_
 
 #include <algorithm>
+#include <atomic>
 #include <cstring>
 #include <functional>
 #include <memory>
@@ -75,6 +76,10 @@
   public:
     explicit ChunkedCommandSerializer(CommandSerializer* serializer);
 
+    // This utility function is intended only for disconnect situations where we want the serializer
+    // to appear to keep working even though we are no longer serializing and flushing commands.
+    void SetCommandSerializerForDisconnect(CommandSerializer* serializer);
+
     template <typename Cmd>
     void SerializeCommand(const Cmd& cmd) {
         SerializeCommandImpl(
@@ -144,10 +149,11 @@
         SerializeChunkedCommand(cmdSpace.get(), requiredSize);
     }
 
-    void SerializeChunkedCommand(const char* allocatedBuffer, size_t remainingSize);
+    void SerializeChunkedCommand(const char* allocatedBuffer, size_t totalSize);
 
     raw_ptr<CommandSerializer> mSerializer;
     size_t mMaxAllocationSize;
+    std::atomic<uint64_t> mNextChunkedCommandId = 0;
 };
 
 }  // namespace dawn::wire
diff --git a/src/dawn/wire/client/Client.cpp b/src/dawn/wire/client/Client.cpp
index 6d7c7ad..f3e12b1 100644
--- a/src/dawn/wire/client/Client.cpp
+++ b/src/dawn/wire/client/Client.cpp
@@ -174,7 +174,7 @@
 
 void Client::Disconnect() {
     mDisconnected = true;
-    mSerializer = ChunkedCommandSerializer(NoopCommandSerializer::GetInstance());
+    mSerializer.SetCommandSerializerForDisconnect(NoopCommandSerializer::GetInstance());
 
     // Transition all event managers to ClientDropped state.
     for (auto& [_, eventManager] : mEventManagers) {
diff --git a/src/dawn/wire/client/Client.h b/src/dawn/wire/client/Client.h
index 3732454..1406fa4 100644
--- a/src/dawn/wire/client/Client.h
+++ b/src/dawn/wire/client/Client.h
@@ -82,7 +82,7 @@
     }
 
     // ChunkedCommandHandler implementation
-    const volatile char* HandleCommandsImpl(const volatile char* commands, size_t size) override;
+    const volatile char* HandleCommands(const volatile char* commands, size_t size) override;
 
     MemoryTransferService* GetMemoryTransferService() const { return mMemoryTransferService; }
 
@@ -143,7 +143,6 @@
 #include "dawn/wire/client/ClientPrototypes_autogen.inc"
 
     ChunkedCommandSerializer mSerializer;
-    WireDeserializeAllocator mWireCommandAllocator;
     PerObjectType<ObjectStore> mObjects;
     std::unique_ptr<MemoryTransferService> mOwnedMemoryTransferService = nullptr;
     raw_ptr<MemoryTransferService> mMemoryTransferService = nullptr;
diff --git a/src/dawn/wire/server/Server.h b/src/dawn/wire/server/Server.h
index a8896b4..6ded567 100644
--- a/src/dawn/wire/server/Server.h
+++ b/src/dawn/wire/server/Server.h
@@ -173,7 +173,7 @@
     ~Server() override;
 
     // ChunkedCommandHandler implementation
-    const volatile char* HandleCommandsImpl(const volatile char* commands, size_t size) override;
+    const volatile char* HandleCommands(const volatile char* commands, size_t size) override;
 
     WireResult InjectBuffer(WGPUBuffer buffer, const Handle& handle, const Handle& deviceHandle);
     WireResult InjectTexture(WGPUTexture texture, const Handle& handle, const Handle& deviceHandle);
@@ -281,7 +281,6 @@
 
 #include "dawn/wire/server/ServerPrototypes_autogen.inc"
 
-    WireDeserializeAllocator mAllocator;
     MutexProtected<ChunkedCommandSerializer> mSerializer;
     std::unique_ptr<MemoryTransferService> mOwnedMemoryTransferService = nullptr;
     raw_ptr<MemoryTransferService> mMemoryTransferService = nullptr;