WebGPU: Handle dawn_wire allocation errors

Bug: chromium:951558
Change-Id: I68fcba0541f48bcb492b23ad34ec8c8b21dd1fd3
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2463865
Commit-Queue: Austin Eng <enga@chromium.org>
Reviewed-by: Corentin Wallez <cwallez@chromium.org>
Cr-Commit-Position: refs/heads/master@{#817100}

gpu/command_buffer/client/webgpu_implementation.cc

@@ -80,50 +80,55 @@ void WebGPUCommandSerializer::RequestDeviceCreation(
   helper_->Flush();
 }
 
+size_t WebGPUCommandSerializer::GetMaximumAllocationSize() const {
+  return c2s_transfer_buffer_->GetMaxSize();
+}
+
 void* WebGPUCommandSerializer::GetCmdSpace(size_t size) {
+  // Note: Dawn will never call this function with |size| >
+  // GetMaximumAllocationSize().
+  DCHECK_LE(size, GetMaximumAllocationSize());
+
   // The buffer size must be initialized before any commands are serialized.
-  if (c2s_buffer_default_size_ == 0u) {
-    NOTREACHED();
-    return nullptr;
-  }
+  DCHECK_NE(c2s_buffer_default_size_, 0u);
 
-  base::CheckedNumeric<uint32_t> checked_next_offset(c2s_put_offset_);
-  checked_next_offset += size;
+  DCHECK_LE(c2s_put_offset_, c2s_buffer_.size());
+  const bool overflows_remaining_space =
+      size > static_cast<size_t>(c2s_buffer_.size() - c2s_put_offset_);
 
-  uint32_t next_offset;
-  bool next_offset_valid = checked_next_offset.AssignIfValid(&next_offset);
+  if (LIKELY(c2s_buffer_.valid() && !overflows_remaining_space)) {
+    // If the buffer is valid and has sufficient space, return the
+    // pointer and increment the offset.
+    uint8_t* ptr = static_cast<uint8_t*>(c2s_buffer_.address());
+    ptr += c2s_put_offset_;
 
-  // If the buffer does not have enough space, or if the buffer is not
-  // initialized, flush and reset the command stream.
-  if (!next_offset_valid || next_offset > c2s_buffer_.size() ||
-      !c2s_buffer_.valid()) {
-    Flush();
+    c2s_put_offset_ += static_cast<uint32_t>(size);
+    return ptr;
+  }
 
-    uint32_t max_allocation = c2s_transfer_buffer_->GetMaxSize();
-    // TODO(crbug.com/951558): Handle command chunking or ensure commands aren't
-    // this large.
-    CHECK_LE(size, max_allocation);
+  if (!c2s_transfer_buffer_) {
+    // The serializer hit a fatal error and was disconnected.
+    return nullptr;
+  }
 
-    uint32_t allocation_size =
-        std::max(c2s_buffer_default_size_, static_cast<uint32_t>(size));
-    TRACE_EVENT1(TRACE_DISABLED_BY_DEFAULT("gpu.dawn"),
-                 "WebGPUCommandSerializer::GetCmdSpace", "bytes",
-                 allocation_size);
-    c2s_buffer_.Reset(allocation_size);
-    c2s_put_offset_ = 0;
-    next_offset = size;
+  // Otherwise, flush and reset the command stream.
+  Flush();
 
-    // TODO(crbug.com/951558): Handle OOM.
-    CHECK(c2s_buffer_.valid());
-    CHECK_LE(size, c2s_buffer_.size());
-  }
+  uint32_t allocation_size =
+      std::max(c2s_buffer_default_size_, static_cast<uint32_t>(size));
+  TRACE_EVENT1(TRACE_DISABLED_BY_DEFAULT("gpu.dawn"),
+               "WebGPUCommandSerializer::GetCmdSpace", "bytes",
+               allocation_size);
+  c2s_buffer_.Reset(allocation_size);
 
-  DCHECK(c2s_buffer_.valid());
-  uint8_t* ptr = static_cast<uint8_t*>(c2s_buffer_.address());
-  ptr += c2s_put_offset_;
+  if (!c2s_buffer_.valid() || c2s_buffer_.size() < size) {
+    DLOG(ERROR) << "Dawn wire transfer buffer allocation failed";
+    HandleGpuControlLostContext();
+    return nullptr;
+  }
 
-  c2s_put_offset_ = next_offset;
-  return ptr;
+  c2s_put_offset_ = size;
+  return c2s_buffer_.address();
 }
 
 bool WebGPUCommandSerializer::Flush() {
@@ -160,11 +165,9 @@ void WebGPUCommandSerializer::HandleGpuControlLostContext() {
   c2s_buffer_.Discard();
   c2s_transfer_buffer_ = nullptr;
 
-  // Disconnect the wire client. WebGPU commands will be serialized into dummy
-  // space owned by the wire client, and the device will receive a Lost event.
-  // No commands will be sent after this point.
+  // Disconnect the wire client. WebGPU commands will become a noop, and the
+  // device will receive a Lost event.
   // NOTE: This assumes single-threaded operation.
-  // TODO(enga): Implement context reset/recovery.
   wire_client_->Disconnect();
 }
 

gpu/command_buffer/client/webgpu_implementation.h

@@ -44,6 +44,7 @@ class WebGPUCommandSerializer final : public dawn_wire::CommandSerializer {
       const WGPUDeviceProperties& requested_device_properties);
 
   // dawn_wire::CommandSerializer implementation
+  size_t GetMaximumAllocationSize() const final;
   void* GetCmdSpace(size_t size) final;
   bool Flush() final;
 

gpu/command_buffer/service/webgpu_decoder_impl.cc

@@ -39,11 +39,17 @@ constexpr size_t kMaxWireBufferSize =
     std::min(IPC::Channel::kMaximumMessageSize,
              static_cast<size_t>(1024 * 1024));
 
+constexpr size_t kDawnReturnCmdsOffset =
+    offsetof(cmds::DawnReturnCommandsInfo, deserialized_buffer);
+
+static_assert(kDawnReturnCmdsOffset < kMaxWireBufferSize, "");
+
 class WireServerCommandSerializer : public dawn_wire::CommandSerializer {
  public:
   WireServerCommandSerializer(DecoderClient* client,
                               DawnDeviceClientID device_client_id);
   ~WireServerCommandSerializer() override = default;
+  size_t GetMaximumAllocationSize() const final;
   void* GetCmdSpace(size_t size) final;
   bool Flush() final;
 
@@ -68,37 +74,29 @@ WireServerCommandSerializer::WireServerCommandSerializer(
   header->device_client_id = device_client_id;
 }
 
-void* WireServerCommandSerializer::GetCmdSpace(size_t size) {
-  // TODO(enga): Handle chunking commands if size +
-  // offsetof(cmds::DawnReturnCommandsInfo, deserialized_buffer)>
-  // kMaxWireBufferSize.
-  size_t total_wire_buffer_size =
-      (base::CheckedNumeric<size_t>(size) +
-       base::CheckedNumeric<size_t>(
-           offsetof(cmds::DawnReturnCommandsInfo, deserialized_buffer)))
-          .ValueOrDie();
-  if (total_wire_buffer_size > kMaxWireBufferSize) {
-    NOTREACHED();
-    return nullptr;
-  }
+size_t WireServerCommandSerializer::GetMaximumAllocationSize() const {
+  return kMaxWireBufferSize - kDawnReturnCmdsOffset;
+}
 
-  // |next_offset| should never be more than kMaxWireBufferSize +
-  // kMaxWireBufferSize.
+void* WireServerCommandSerializer::GetCmdSpace(size_t size) {
+  // Note: Dawn will never call this function with |size| >
+  // GetMaximumAllocationSize().
   DCHECK_LE(put_offset_, kMaxWireBufferSize);
-  DCHECK_LE(size, kMaxWireBufferSize);
+  DCHECK_LE(size, GetMaximumAllocationSize());
+
+  // Statically check that kMaxWireBufferSize + kMaxWireBufferSize is
+  // a valid uint32_t. We can add put_offset_ and size without overflow.
   static_assert(base::CheckAdd(kMaxWireBufferSize, kMaxWireBufferSize)
                     .IsValid<uint32_t>(),
                 "");
-  uint32_t next_offset = put_offset_ + size;
-
+  uint32_t next_offset = put_offset_ + static_cast<uint32_t>(size);
   if (next_offset > buffer_.size()) {
     Flush();
     // TODO(enga): Keep track of how much command space the application is using
     // and adjust the buffer size accordingly.
 
-    DCHECK_EQ(put_offset_,
-              offsetof(cmds::DawnReturnCommandsInfo, deserialized_buffer));
-    next_offset = put_offset_ + size;
+    DCHECK_EQ(put_offset_, kDawnReturnCmdsOffset);
+    next_offset = put_offset_ + static_cast<uint32_t>(size);
   }
 
   uint8_t* ptr = &buffer_[put_offset_];
@@ -107,8 +105,7 @@ void* WireServerCommandSerializer::GetCmdSpace(size_t size) {
 }
 
 bool WireServerCommandSerializer::Flush() {
-  if (put_offset_ >
-      offsetof(cmds::DawnReturnCommandsInfo, deserialized_buffer)) {
+  if (put_offset_ > kDawnReturnCmdsOffset) {
     TRACE_EVENT1(TRACE_DISABLED_BY_DEFAULT("gpu.dawn"),
                  "WireServerCommandSerializer::Flush", "bytes", put_offset_);
 
@@ -117,7 +114,7 @@ bool WireServerCommandSerializer::Flush() {
                             "DawnReturnCommands", return_trace_id++);
 
     client_->HandleReturnData(base::make_span(buffer_.data(), put_offset_));
-    put_offset_ = offsetof(cmds::DawnReturnCommandsInfo, deserialized_buffer);
+    put_offset_ = kDawnReturnCmdsOffset;
   }
   return true;
 }