dawn_wire: Support chunked commands

This CL adds support for chunking large commands by first serializing
large commands first into a separate buffer, and then sending the
buffer data chunk by chunk.

This code path is used for large writeBuffer and writeTexture, as well
as the inline memory transfer service for buffer mapping. The transfer
for writeBuffer and writeTexture will be optimized further in Chrome,
and the inline memory transfer service is currently used only in tests.

Bug: chromium:1123861, chromium:951558
Change-Id: I02491a44e653e2383174958d9c3d4a4db6fd7bde
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/28882
Commit-Queue: Austin Eng <enga@chromium.org>
Reviewed-by: Stephen White <senorblanco@chromium.org>

generator/templates/dawn_wire/WireCmd.cpp

@@ -99,11 +99,12 @@ namespace {
 {% macro write_record_serialization_helpers(record, name, members, is_cmd=False, is_return_command=False) %}
     {% set Return = "Return" if is_return_command else "" %}
     {% set Cmd = "Cmd" if is_cmd else "" %}
+    {% set Inherits = " : CmdHeader" if is_cmd else "" %}
 
     //* Structure for the wire format of each of the records. Members that are values
     //* are embedded directly in the structure. Other members are assumed to be in the
     //* memory directly following the structure in the buffer.
-    struct {{Return}}{{name}}Transfer {
+    struct {{Return}}{{name}}Transfer{{Inherits}} {
         static_assert({{[is_cmd, record.extensible, record.chained].count(True)}} <= 1,
                       "Record must be at most one of is_cmd, extensible, and chained.");
         {% if is_cmd %}
@@ -130,6 +131,11 @@ namespace {
         {% endfor %}
     };
 
+    {% if is_cmd %}
+        static_assert(offsetof({{Return}}{{name}}Transfer, commandSize) == 0, "");
+        static_assert(offsetof({{Return}}{{name}}Transfer, commandId) == sizeof(CmdHeader), "");
+    {% endif %}
+
     {% if record.chained %}
         static_assert(offsetof({{Return}}{{name}}Transfer, chain) == 0, "");
     {% endif %}
@@ -379,12 +385,13 @@ namespace {
         return size;
     }
 
-    void {{Cmd}}::Serialize(char* buffer
+    void {{Cmd}}::Serialize(size_t commandSize, char* buffer
         {%- if not is_return -%}
             , const ObjectIdProvider& objectIdProvider
         {%- endif -%}
     ) const {
         auto transfer = reinterpret_cast<{{Name}}Transfer*>(buffer);
+        transfer->commandSize = commandSize;
         buffer += sizeof({{Name}}Transfer);
 
         {{Name}}Serialize(*this, transfer, &buffer

generator/templates/dawn_wire/WireCmd.h

@@ -92,6 +92,10 @@ namespace dawn_wire {
         {% endfor %}
     };
 
+    struct CmdHeader {
+        uint64_t commandSize;
+    };
+
 {% macro write_command_struct(command, is_return_command) %}
     {% set Return = "Return" if is_return_command else "" %}
     {% set Cmd = command.name.CamelCase() + "Cmd" %}
@@ -101,7 +105,7 @@ namespace dawn_wire {
 
         //* Serialize the structure and everything it points to into serializeBuffer which must be
         //* big enough to contain all the data (as queried from GetRequiredSize).
-        void Serialize(char* serializeBuffer
+        void Serialize(size_t commandSize, char* serializeBuffer
             {%- if not is_return_command -%}
                 , const ObjectIdProvider& objectIdProvider
             {%- endif -%}

generator/templates/dawn_wire/client/ClientBase.h

@@ -15,13 +15,14 @@
 #ifndef DAWNWIRE_CLIENT_CLIENTBASE_AUTOGEN_H_
 #define DAWNWIRE_CLIENT_CLIENTBASE_AUTOGEN_H_
 
+#include "dawn_wire/ChunkedCommandHandler.h"
 #include "dawn_wire/WireCmd_autogen.h"
 #include "dawn_wire/client/ApiObjects.h"
 #include "dawn_wire/client/ObjectAllocator.h"
 
 namespace dawn_wire { namespace client {
 
-    class ClientBase : public ObjectIdProvider {
+    class ClientBase : public ChunkedCommandHandler, public ObjectIdProvider {
       public:
         ClientBase() {
         }

generator/templates/dawn_wire/client/ClientHandlers.cpp

@@ -53,10 +53,20 @@ namespace dawn_wire { namespace client {
         }
     {% endfor %}
 
-    const volatile char* Client::HandleCommands(const volatile char* commands, size_t size) {
-        while (size >= sizeof(ReturnWireCmd)) {
-            ReturnWireCmd cmdId = *reinterpret_cast<const volatile ReturnWireCmd*>(commands);
+    const volatile char* Client::HandleCommandsImpl(const volatile char* commands, size_t size) {
+        while (size >= 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(commands, size)) {
+                case ChunkedCommandsResult::Consumed:
+                    return commands + size;
+                case ChunkedCommandsResult::Error:
+                    return nullptr;
+                case ChunkedCommandsResult::Passthrough:
+                    break;
+            }
 
+            ReturnWireCmd cmdId = *reinterpret_cast<const volatile ReturnWireCmd*>(commands + sizeof(CmdHeader));
             bool success = false;
             switch (cmdId) {
                 {% for command in cmd_records["return command"] %}

generator/templates/dawn_wire/server/ServerBase.h

@@ -16,6 +16,7 @@
 #define DAWNWIRE_SERVER_SERVERBASE_H_
 
 #include "dawn/dawn_proc_table.h"
+#include "dawn_wire/ChunkedCommandHandler.h"
 #include "dawn_wire/Wire.h"
 #include "dawn_wire/WireCmd_autogen.h"
 #include "dawn_wire/WireDeserializeAllocator.h"
@@ -23,7 +24,7 @@
 
 namespace dawn_wire { namespace server {
 
-    class ServerBase : public ObjectIdResolver {
+    class ServerBase : public ChunkedCommandHandler, public ObjectIdResolver {
       public:
         ServerBase() = default;
         virtual ~ServerBase() = default;

generator/templates/dawn_wire/server/ServerHandlers.cpp

@@ -89,12 +89,22 @@ namespace dawn_wire { namespace server {
         }
     {% endfor %}
 
-    const volatile char* Server::HandleCommands(const volatile char* commands, size_t size) {
+    const volatile char* Server::HandleCommandsImpl(const volatile char* commands, size_t size) {
         mProcs.deviceTick(DeviceObjects().Get(1)->handle);
 
-        while (size >= sizeof(WireCmd)) {
-            WireCmd cmdId = *reinterpret_cast<const volatile WireCmd*>(commands);
+        while (size >= 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(commands, size)) {
+                case ChunkedCommandsResult::Consumed:
+                    return commands + size;
+                case ChunkedCommandsResult::Error:
+                    return nullptr;
+                case ChunkedCommandsResult::Passthrough:
+                    break;
+            }
 
+            WireCmd cmdId = *reinterpret_cast<const volatile WireCmd*>(commands + sizeof(CmdHeader));
             bool success = false;
             switch (cmdId) {
                 {% for command in cmd_records["command"] %}

src/dawn_wire/BUILD.gn

@@ -58,6 +58,10 @@ dawn_component("dawn_wire") {
   configs = [ "${dawn_root}/src/common:dawn_internal" ]
   sources = get_target_outputs(":dawn_wire_gen")
   sources += [
+    "ChunkedCommandHandler.cpp",
+    "ChunkedCommandHandler.h",
+    "ChunkedCommandSerializer.cpp",
+    "ChunkedCommandSerializer.h",
     "WireClient.cpp",
     "WireDeserializeAllocator.cpp",
     "WireDeserializeAllocator.h",

src/dawn_wire/CMakeLists.txt

@@ -31,6 +31,10 @@ target_sources(dawn_wire PRIVATE
     "${DAWN_INCLUDE_DIR}/dawn_wire/WireServer.h"
     "${DAWN_INCLUDE_DIR}/dawn_wire/dawn_wire_export.h"
     ${DAWN_WIRE_GEN_SOURCES}
+    "ChunkedCommandHandler.cpp"
+    "ChunkedCommandHandler.h"
+    "ChunkedCommandSerializer.cpp"
+    "ChunkedCommandSerializer.h"
     "WireClient.cpp"
     "WireDeserializeAllocator.cpp"
     "WireDeserializeAllocator.h"

src/dawn_wire/ChunkedCommandHandler.cpp

@@ -0,0 +1,64 @@
+// Copyright 2020 The Dawn Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "dawn_wire/ChunkedCommandHandler.h"
+
+namespace dawn_wire {
+
+    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);
+            mChunkedCommandData.insert(mChunkedCommandData.end(), commands, commands + chunkSize);
+
+            commands += chunkSize;
+            mChunkedCommandRemainingSize -= chunkSize;
+            size -= chunkSize;
+
+            if (mChunkedCommandRemainingSize == 0) {
+                // Once the chunked command is complete, pass the data to the command handler
+                // implemenation.
+                const char* chunkedCommands = mChunkedCommandData.data();
+                size_t chunkedSize = mChunkedCommandData.size();
+                if (HandleCommandsImpl(chunkedCommands, chunkedSize) == nullptr) {
+                    // |HandleCommandsImpl| returns nullptr on error. Forward any errors
+                    // out.
+                    return nullptr;
+                }
+                mChunkedCommandData.clear();
+            }
+        }
+
+        return HandleCommandsImpl(commands, size);
+    }
+
+    void ChunkedCommandHandler::BeginChunkedCommandData(const volatile char* commands,
+                                                        size_t commandSize,
+                                                        size_t initialSize) {
+        ASSERT(mChunkedCommandData.empty());
+
+        // Reserve space for all the command data we're expecting, and append the initial data
+        // to the end of the vector.
+        mChunkedCommandRemainingSize = commandSize - initialSize;
+        mChunkedCommandData.reserve(commandSize);
+        mChunkedCommandData.insert(mChunkedCommandData.end(), commands, commands + initialSize);
+    }
+
+}  // namespace dawn_wire

src/dawn_wire/ChunkedCommandHandler.h

@@ -0,0 +1,71 @@
+// Copyright 2020 The Dawn Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef DAWNWIRE_CHUNKEDCOMMANDHANDLER_H_
+#define DAWNWIRE_CHUNKEDCOMMANDHANDLER_H_
+
+#include "common/Assert.h"
+#include "dawn_wire/Wire.h"
+#include "dawn_wire/WireCmd_autogen.h"
+
+#include <cstdint>
+#include <vector>
+
+namespace dawn_wire {
+
+    class ChunkedCommandHandler : public CommandHandler {
+      public:
+        const volatile char* HandleCommands(const volatile char* commands, size_t size) override;
+        ~ChunkedCommandHandler() override;
+
+      protected:
+        enum class ChunkedCommandsResult {
+            Passthrough,
+            Consumed,
+            Error,
+        };
+
+        // 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) {
+                BeginChunkedCommandData(commands, commandSize, size);
+                return ChunkedCommandsResult::Consumed;
+            }
+            return ChunkedCommandsResult::Passthrough;
+        }
+
+      private:
+        virtual const volatile char* HandleCommandsImpl(const volatile char* commands,
+                                                        size_t size) = 0;
+
+        void BeginChunkedCommandData(const volatile char* commands,
+                                     size_t commandSize,
+                                     size_t initialSize);
+
+        size_t mChunkedCommandRemainingSize = 0;
+        std::vector<char> mChunkedCommandData;
+    };
+
+}  // namespace dawn_wire
+
+#endif  // DAWNWIRE_CHUNKEDCOMMANDHANDLER_H_

src/dawn_wire/ChunkedCommandSerializer.cpp

@@ -0,0 +1,38 @@
+// Copyright 2020 The Dawn Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "dawn_wire/ChunkedCommandSerializer.h"
+
+namespace dawn_wire {
+
+    ChunkedCommandSerializer::ChunkedCommandSerializer(CommandSerializer* serializer)
+        : mSerializer(serializer), mMaxAllocationSize(serializer->GetMaximumAllocationSize()) {
+    }
+
+    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);
+
+            allocatedBuffer += chunkSize;
+            remainingSize -= chunkSize;
+        }
+    }
+
+}  // namespace dawn_wire

src/dawn_wire/ChunkedCommandSerializer.h

@@ -0,0 +1,101 @@
+// Copyright 2020 The Dawn Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef DAWNWIRE_CHUNKEDCOMMANDSERIALIZER_H_
+#define DAWNWIRE_CHUNKEDCOMMANDSERIALIZER_H_
+
+#include "common/Compiler.h"
+#include "dawn_wire/Wire.h"
+#include "dawn_wire/WireCmd_autogen.h"
+
+#include <algorithm>
+#include <cstring>
+
+namespace dawn_wire {
+
+    class ChunkedCommandSerializer {
+      public:
+        ChunkedCommandSerializer(CommandSerializer* serializer);
+
+        template <typename Cmd>
+        void SerializeCommand(const Cmd& cmd) {
+            SerializeCommand(cmd, 0, [](char*) {});
+        }
+
+        template <typename Cmd, typename ExtraSizeSerializeFn>
+        void SerializeCommand(const Cmd& cmd,
+                              size_t extraSize,
+                              ExtraSizeSerializeFn&& SerializeExtraSize) {
+            SerializeCommandImpl(
+                cmd,
+                [](const Cmd& cmd, size_t requiredSize, char* allocatedBuffer) {
+                    cmd.Serialize(requiredSize, allocatedBuffer);
+                },
+                extraSize, std::forward<ExtraSizeSerializeFn>(SerializeExtraSize));
+        }
+
+        template <typename Cmd>
+        void SerializeCommand(const Cmd& cmd, const ObjectIdProvider& objectIdProvider) {
+            SerializeCommand(cmd, objectIdProvider, 0, [](char*) {});
+        }
+
+        template <typename Cmd, typename ExtraSizeSerializeFn>
+        void SerializeCommand(const Cmd& cmd,
+                              const ObjectIdProvider& objectIdProvider,
+                              size_t extraSize,
+                              ExtraSizeSerializeFn&& SerializeExtraSize) {
+            SerializeCommandImpl(
+                cmd,
+                [&objectIdProvider](const Cmd& cmd, size_t requiredSize, char* allocatedBuffer) {
+                    cmd.Serialize(requiredSize, allocatedBuffer, objectIdProvider);
+                },
+                extraSize, std::forward<ExtraSizeSerializeFn>(SerializeExtraSize));
+        }
+
+      private:
+        template <typename Cmd, typename SerializeCmdFn, typename ExtraSizeSerializeFn>
+        void SerializeCommandImpl(const Cmd& cmd,
+                                  SerializeCmdFn&& SerializeCmd,
+                                  size_t extraSize,
+                                  ExtraSizeSerializeFn&& SerializeExtraSize) {
+            size_t commandSize = cmd.GetRequiredSize();
+            size_t requiredSize = commandSize + extraSize;
+
+            if (requiredSize <= mMaxAllocationSize) {
+                char* allocatedBuffer = static_cast<char*>(mSerializer->GetCmdSpace(requiredSize));
+                if (allocatedBuffer != nullptr) {
+                    SerializeCmd(cmd, requiredSize, allocatedBuffer);
+                    SerializeExtraSize(allocatedBuffer + commandSize);
+                }
+                return;
+            }
+
+            auto cmdSpace = std::unique_ptr<char[]>(new (std::nothrow) char[requiredSize]);
+            if (!cmdSpace) {
+                return;
+            }
+            SerializeCmd(cmd, requiredSize, cmdSpace.get());
+            SerializeExtraSize(cmdSpace.get() + commandSize);
+            SerializeChunkedCommand(cmdSpace.get(), requiredSize);
+        }
+
+        void SerializeChunkedCommand(const char* allocatedBuffer, size_t remainingSize);
+
+        CommandSerializer* mSerializer;
+        size_t mMaxAllocationSize;
+    };
+
+}  // namespace dawn_wire
+
+#endif  // DAWNWIRE_CHUNKEDCOMMANDSERIALIZER_H_

src/dawn_wire/client/Buffer.cpp

@@ -71,20 +71,19 @@ namespace dawn_wire { namespace client {
         cmd.handleCreateInfoLength = writeHandleCreateInfoLength;
         cmd.handleCreateInfo = nullptr;
 
-        char* writeHandleSpace = wireClient->SerializeCommand(cmd, writeHandleCreateInfoLength);
-
-        if (descriptor->mappedAtCreation) {
-            // Serialize the WriteHandle into the space after the command.
-            writeHandle->SerializeCreate(writeHandleSpace);
-
-            // Set the buffer state for the mapping at creation. The buffer now owns the write
-            // handle..
-            buffer->mWriteHandle = std::move(writeHandle);
-            buffer->mMappedData = writeData;
-            buffer->mMapOffset = 0;
-            buffer->mMapSize = buffer->mSize;
-        }
+        wireClient->SerializeCommand(cmd, writeHandleCreateInfoLength, [&](char* cmdSpace) {
+            if (descriptor->mappedAtCreation) {
+                // Serialize the WriteHandle into the space after the command.
+                writeHandle->SerializeCreate(cmdSpace);
 
+                // Set the buffer state for the mapping at creation. The buffer now owns the write
+                // handle..
+                buffer->mWriteHandle = std::move(writeHandle);
+                buffer->mMappedData = writeData;
+                buffer->mMapOffset = 0;
+                buffer->mMapSize = buffer->mSize;
+            }
+        });
         return ToAPI(buffer);
     }
 
@@ -176,15 +175,15 @@ namespace dawn_wire { namespace client {
         // Step 3a. Fill the handle create info in the command.
         if (isReadMode) {
             cmd.handleCreateInfoLength = request.readHandle->SerializeCreateSize();
-            char* handleCreateInfoSpace =
-                device->GetClient()->SerializeCommand(cmd, cmd.handleCreateInfoLength);
-            request.readHandle->SerializeCreate(handleCreateInfoSpace);
+            device->GetClient()->SerializeCommand(
+                cmd, cmd.handleCreateInfoLength,
+                [&](char* cmdSpace) { request.readHandle->SerializeCreate(cmdSpace); });
         } else {
             ASSERT(isWriteMode);
             cmd.handleCreateInfoLength = request.writeHandle->SerializeCreateSize();
-            char* handleCreateInfoSpace =
-                device->GetClient()->SerializeCommand(cmd, cmd.handleCreateInfoLength);
-            request.writeHandle->SerializeCreate(handleCreateInfoSpace);
+            device->GetClient()->SerializeCommand(
+                cmd, cmd.handleCreateInfoLength,
+                [&](char* cmdSpace) { request.writeHandle->SerializeCreate(cmdSpace); });
         }
 
         // Step 4. Register this request so that we can retrieve it from its serial when the server
@@ -311,12 +310,11 @@ namespace dawn_wire { namespace client {
             cmd.writeFlushInfoLength = writeFlushInfoLength;
             cmd.writeFlushInfo = nullptr;
 
-            char* writeHandleSpace =
-                device->GetClient()->SerializeCommand(cmd, writeFlushInfoLength);
-
-            // Serialize flush metadata into the space after the command.
-            // This closes the handle for writing.
-            mWriteHandle->SerializeFlush(writeHandleSpace);
+            device->GetClient()->SerializeCommand(cmd, writeFlushInfoLength, [&](char* cmdSpace) {
+                // Serialize flush metadata into the space after the command.
+                // This closes the handle for writing.
+                mWriteHandle->SerializeFlush(cmdSpace);
+            });
             mWriteHandle = nullptr;
 
         } else if (mReadHandle) {

src/dawn_wire/client/Client.cpp

@@ -52,22 +52,12 @@ namespace dawn_wire { namespace client {
         return result;
     }
 
-    char* Client::GetCmdSpace(size_t size) {
-        if (DAWN_UNLIKELY(mIsDisconnected)) {
-            if (size > mDummyCmdSpace.size()) {
-                mDummyCmdSpace.resize(size);
-            }
-            return mDummyCmdSpace.data();
-        }
-        return static_cast<char*>(mSerializer->GetCmdSpace(size));
-    }
-
     void Client::Disconnect() {
-        if (mIsDisconnected) {
+        if (mDisconnected) {
             return;
         }
 
-        mIsDisconnected = true;
+        mDisconnected = true;
         if (mDevice != nullptr) {
             mDevice->HandleDeviceLost("GPU connection lost");
         }

src/dawn_wire/client/Client.h

@@ -18,6 +18,7 @@
 #include <dawn/webgpu.h>
 #include <dawn_wire/Wire.h>
 
+#include "dawn_wire/ChunkedCommandSerializer.h"
 #include "dawn_wire/WireClient.h"
 #include "dawn_wire/WireCmd_autogen.h"
 #include "dawn_wire/WireDeserializeAllocator.h"
@@ -31,7 +32,11 @@ namespace dawn_wire { namespace client {
     class Client : public ClientBase {
       public:
         Client(CommandSerializer* serializer, MemoryTransferService* memoryTransferService);
-        ~Client();
+        ~Client() override;
+
+        // ChunkedCommandHandler implementation
+        const volatile char* HandleCommandsImpl(const volatile char* commands,
+                                                size_t size) override;
 
         WGPUDevice GetDevice();
 
@@ -39,16 +44,26 @@ namespace dawn_wire { namespace client {
             return mMemoryTransferService;
         }
 
-        const volatile char* HandleCommands(const volatile char* commands, size_t size);
         ReservedTexture ReserveTexture(WGPUDevice device);
 
         template <typename Cmd>
-        char* SerializeCommand(const Cmd& cmd, size_t extraSize = 0) {
-            size_t requiredSize = cmd.GetRequiredSize();
-            // TODO(cwallez@chromium.org): Check for overflows and allocation success?
-            char* allocatedBuffer = GetCmdSpace(requiredSize + extraSize);
-            cmd.Serialize(allocatedBuffer, *this);
-            return allocatedBuffer + requiredSize;
+        void SerializeCommand(const Cmd& cmd) {
+            // TODO(enga): Swap out the serializer with a no-op one on disconnect.
+            if (mDisconnected) {
+                return;
+            }
+            mSerializer.SerializeCommand(cmd, *this);
+        }
+
+        template <typename Cmd, typename ExtraSizeSerializeFn>
+        void SerializeCommand(const Cmd& cmd,
+                              size_t extraSize,
+                              ExtraSizeSerializeFn&& SerializeExtraSize) {
+            // TODO(enga): Swap out the serializer with a no-op one on disconnect.
+            if (mDisconnected) {
+                return;
+            }
+            mSerializer.SerializeCommand(cmd, *this, extraSize, SerializeExtraSize);
         }
 
         void Disconnect();
@@ -56,16 +71,12 @@ namespace dawn_wire { namespace client {
       private:
 #include "dawn_wire/client/ClientPrototypes_autogen.inc"
 
-        char* GetCmdSpace(size_t size);
-
         Device* mDevice = nullptr;
-        CommandSerializer* mSerializer = nullptr;
+        ChunkedCommandSerializer mSerializer;
         WireDeserializeAllocator mAllocator;
         MemoryTransferService* mMemoryTransferService = nullptr;
         std::unique_ptr<MemoryTransferService> mOwnedMemoryTransferService = nullptr;
-
-        std::vector<char> mDummyCmdSpace;
-        bool mIsDisconnected = false;
+        bool mDisconnected = false;
     };
 
     std::unique_ptr<MemoryTransferService> CreateInlineMemoryTransferService();

src/dawn_wire/server/Server.cpp

@@ -39,10 +39,6 @@ namespace dawn_wire { namespace server {
         DestroyAllObjects(mProcs);
     }
 
-    char* Server::GetCmdSpace(size_t size) {
-        return static_cast<char*>(mSerializer->GetCmdSpace(size));
-    }
-
     bool Server::InjectTexture(WGPUTexture texture, uint32_t id, uint32_t generation) {
         ObjectData<WGPUTexture>* data = TextureObjects().Allocate(id);
         if (data == nullptr) {

src/dawn_wire/server/Server.h

@@ -15,6 +15,7 @@
 #ifndef DAWNWIRE_SERVER_SERVER_H_
 #define DAWNWIRE_SERVER_SERVER_H_
 
+#include "dawn_wire/ChunkedCommandSerializer.h"
 #include "dawn_wire/server/ServerBase_autogen.h"
 
 namespace dawn_wire { namespace server {
@@ -60,22 +61,26 @@ namespace dawn_wire { namespace server {
                const DawnProcTable& procs,
                CommandSerializer* serializer,
                MemoryTransferService* memoryTransferService);
-        ~Server();
+        ~Server() override;
 
-        const volatile char* HandleCommands(const volatile char* commands, size_t size);
+        // ChunkedCommandHandler implementation
+        const volatile char* HandleCommandsImpl(const volatile char* commands,
+                                                size_t size) override;
 
         bool InjectTexture(WGPUTexture texture, uint32_t id, uint32_t generation);
 
       private:
         template <typename Cmd>
-        char* SerializeCommand(const Cmd& cmd, size_t extraSize = 0) {
-            size_t requiredSize = cmd.GetRequiredSize();
-            // TODO(cwallez@chromium.org): Check for overflows and allocation success?
-            char* allocatedBuffer = GetCmdSpace(requiredSize + extraSize);
-            cmd.Serialize(allocatedBuffer);
-            return allocatedBuffer + requiredSize;
+        void SerializeCommand(const Cmd& cmd) {
+            mSerializer.SerializeCommand(cmd);
+        }
+
+        template <typename Cmd, typename ExtraSizeSerializeFn>
+        void SerializeCommand(const Cmd& cmd,
+                              size_t extraSize,
+                              ExtraSizeSerializeFn&& SerializeExtraSize) {
+            mSerializer.SerializeCommand(cmd, extraSize, SerializeExtraSize);
         }
-        char* GetCmdSpace(size_t size);
 
         // Forwarding callbacks
         static void ForwardUncapturedError(WGPUErrorType type, const char* message, void* userdata);
@@ -99,8 +104,8 @@ namespace dawn_wire { namespace server {
 
 #include "dawn_wire/server/ServerPrototypes_autogen.inc"
 
-        CommandSerializer* mSerializer = nullptr;
         WireDeserializeAllocator mAllocator;
+        ChunkedCommandSerializer mSerializer;
         DawnProcTable mProcs;
         std::unique_ptr<MemoryTransferService> mOwnedMemoryTransferService = nullptr;
         MemoryTransferService* mMemoryTransferService = nullptr;

src/dawn_wire/server/ServerBuffer.cpp

@@ -238,26 +238,26 @@ namespace dawn_wire { namespace server {
                 data->readHandle->SerializeInitialDataSize(readData, data->size);
         }
 
-        char* readHandleSpace = SerializeCommand(cmd, cmd.readInitialDataInfoLength);
-
-        if (isSuccess) {
-            if (isRead) {
-                // Serialize the initialization message into the space after the command.
-                data->readHandle->SerializeInitialData(readData, data->size, readHandleSpace);
-                // The in-flight map request returned successfully.
-                // Move the ReadHandle so it is owned by the buffer.
-                bufferData->readHandle = std::move(data->readHandle);
-            } else {
-                // The in-flight map request returned successfully.
-                // Move the WriteHandle so it is owned by the buffer.
-                bufferData->writeHandle = std::move(data->writeHandle);
-                bufferData->mapWriteState = BufferMapWriteState::Mapped;
-                // Set the target of the WriteHandle to the mapped buffer data.
-                bufferData->writeHandle->SetTarget(
-                    mProcs.bufferGetMappedRange(data->bufferObj, data->offset, data->size),
-                    data->size);
+        SerializeCommand(cmd, cmd.readInitialDataInfoLength, [&](char* cmdSpace) {
+            if (isSuccess) {
+                if (isRead) {
+                    // Serialize the initialization message into the space after the command.
+                    data->readHandle->SerializeInitialData(readData, data->size, cmdSpace);
+                    // The in-flight map request returned successfully.
+                    // Move the ReadHandle so it is owned by the buffer.
+                    bufferData->readHandle = std::move(data->readHandle);
+                } else {
+                    // The in-flight map request returned successfully.
+                    // Move the WriteHandle so it is owned by the buffer.
+                    bufferData->writeHandle = std::move(data->writeHandle);
+                    bufferData->mapWriteState = BufferMapWriteState::Mapped;
+                    // Set the target of the WriteHandle to the mapped buffer data.
+                    bufferData->writeHandle->SetTarget(
+                        mProcs.bufferGetMappedRange(data->bufferObj, data->offset, data->size),
+                        data->size);
+                }
             }
-        }
+        });
     }
 
 }}  // namespace dawn_wire::server

src/include/dawn_wire/Wire.h

@@ -16,6 +16,7 @@
 #define DAWNWIRE_WIRE_H_
 
 #include <cstdint>
+#include <limits>
 
 #include "dawn/webgpu.h"
 #include "dawn_wire/dawn_wire_export.h"
@@ -25,8 +26,18 @@ namespace dawn_wire {
     class DAWN_WIRE_EXPORT CommandSerializer {
       public:
         virtual ~CommandSerializer() = default;
+
+        // Get space for serializing commands.
+        // GetCmdSpace will never be called with a value larger than
+        // what GetMaximumAllocationSize returns. Return nullptr to indicate
+        // a fatal error.
         virtual void* GetCmdSpace(size_t size) = 0;
         virtual bool Flush() = 0;
+
+        // TODO(enga): Make pure virtual after updating Chromium.
+        virtual size_t GetMaximumAllocationSize() const {
+            return std::numeric_limits<size_t>::max();
+        }
     };
 
     class DAWN_WIRE_EXPORT CommandHandler {

src/tests/end2end/QueueTests.cpp

@@ -331,6 +331,16 @@ TEST_P(QueueWriteTextureTests, VaryingTextureSize) {
     }
 }
 
+// Test uploading a large amount of data with writeTexture.
+TEST_P(QueueWriteTextureTests, LargeWriteTexture) {
+    TextureSpec textureSpec;
+    textureSpec.textureSize = {2048, 2048, 2};
+    textureSpec.copyOrigin = {0, 0, 0};
+    textureSpec.level = 0;
+
+    DoTest(textureSpec, MinimumDataSpec(textureSpec.textureSize), textureSpec.textureSize);
+}
+
 // Test writing a pixel with an offset.
 TEST_P(QueueWriteTextureTests, VaryingTextureOffset) {
     constexpr uint32_t kWidth = 259;

src/utils/TerribleCommandBuffer.cpp

@@ -29,48 +29,18 @@ namespace utils {
         mHandler = handler;
     }
 
+    size_t TerribleCommandBuffer::GetMaximumAllocationSize() const {
+        return sizeof(mBuffer);
+    }
+
     void* TerribleCommandBuffer::GetCmdSpace(size_t size) {
         // TODO(kainino@chromium.org): Should we early-out if size is 0?
         //   (Here and/or in the caller?) It might be good to make the wire receiver get a nullptr
         //   instead of pointer to zero-sized allocation in mBuffer.
-
-        // Cannot have commands in mBuffer and mLargeBuffer at same time.
-        ASSERT(mOffset == 0 || mLargeBufferCmdSize == 0);
-
         if (size > sizeof(mBuffer)) {
-            // Flush current cmds in mBuffer to keep order.
-            if (mOffset > 0) {
-                if (!Flush()) {
-                    return nullptr;
-                }
-                return GetCmdSpace(size);
-            }
-
-            // Resize large buffer to the size that can
-            // contain incoming command if needed.
-            if (mLargeBuffer.size() < size) {
-                mLargeBuffer.resize(size);
-            }
-
-            // Record whole cmd space.
-            mLargeBufferCmdSize = size;
-
-            return mLargeBuffer.data();
+            return nullptr;
         }
-
-        // Trigger flush if large buffer contain cmds.
-        if (mLargeBufferCmdSize > 0) {
-            if (!Flush()) {
-                return nullptr;
-            }
-            return GetCmdSpace(size);
-        }
-
-        // Need to flush large buffer first.
-        ASSERT(mLargeBufferCmdSize == 0);
-
         char* result = &mBuffer[mOffset];
-
         if (sizeof(mBuffer) - size < mOffset) {
             if (!Flush()) {
                 return nullptr;
@@ -79,26 +49,12 @@ namespace utils {
         }
 
         mOffset += size;
-
         return result;
     }
 
     bool TerribleCommandBuffer::Flush() {
-        // Cannot have commands in mBuffer and mLargeBuffer at same time.
-        ASSERT(mOffset == 0 || mLargeBufferCmdSize == 0);
-
-        bool success = false;
-        // Big buffer not empty, flush it!
-        if (mLargeBufferCmdSize > 0) {
-            success = mHandler->HandleCommands(mLargeBuffer.data(), mLargeBufferCmdSize) != nullptr;
-            // Clear big command buffers.
-            mLargeBufferCmdSize = 0;
-            return success;
-        }
-
-        success = mHandler->HandleCommands(mBuffer, mOffset) != nullptr;
+        bool success = mHandler->HandleCommands(mBuffer, mOffset) != nullptr;
         mOffset = 0;
-
         return success;
     }
 

src/utils/TerribleCommandBuffer.h

@@ -15,8 +15,6 @@
 #ifndef UTILS_TERRIBLE_COMMAND_BUFFER_H_
 #define UTILS_TERRIBLE_COMMAND_BUFFER_H_
 
-#include <vector>
-
 #include "dawn_wire/Wire.h"
 
 namespace utils {
@@ -28,17 +26,15 @@ namespace utils {
 
         void SetHandler(dawn_wire::CommandHandler* handler);
 
+        size_t GetMaximumAllocationSize() const override;
+
         void* GetCmdSpace(size_t size) override;
         bool Flush() override;
 
       private:
         dawn_wire::CommandHandler* mHandler = nullptr;
         size_t mOffset = 0;
-        // Cannot have commands in mBuffer and mLargeBuffer
-        // at the same time to ensure commands order.
         char mBuffer[1000000];
-        std::vector<char> mLargeBuffer;
-        size_t mLargeBufferCmdSize = 0;
     };
 
 }  // namespace utils