From 79e5352b545898e6797e235996c43b0103d76dd9 Mon Sep 17 00:00:00 2001
From: Bo Liu <boliu@chromium.org>
Date: Wed, 12 Aug 2026 05:40:26 -0700
Subject: [PATCH] [webgpu] Check Skia flush success before calling SetCleared

When using Skia to clear a shared image before, WebGPU decoder doesn't
check if the Skia commands succeeded with Graphite. This CL fixes that
by making the Graphite flush and submit functions return a success value
and checking it in WebGPU decoder. Also add a unit test that simulates
Graphite insert recording failure and assert that the shared image is
not cleared incorrectly.

Fixed: 517527943
Change-Id: I4fcf826b5296aa60d1661255e57e1fbf6a6a6964
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/8027742
Reviewed-by: Kai Ninomiya <kainino@chromium.org>
Reviewed-by: Kartar Singh <kartarsingh@google.com>
Commit-Queue: Bo Liu <boliu@chromium.org>
Cr-Commit-Position: refs/heads/main@{#1677962}
---
diff --git a/gpu/command_buffer/service/graphite_shared_context.cc b/gpu/command_buffer/service/graphite_shared_context.cc
index 2c90ae4f..44a5f62 100644
--- a/gpu/command_buffer/service/graphite_shared_context.cc
+++ b/gpu/command_buffer/service/graphite_shared_context.cc
@@ -282,6 +282,12 @@
   return graphite_context_->makePrecompileContext();
 }
 
+void GraphiteSharedContext::set_simulated_insert_status(
+    skgpu::graphite::InsertStatus status) {
+  AutoLock auto_lock(this);
+  simulated_insert_status_ = status;
+}
+
 bool GraphiteSharedContext::insertRecording(
     const skgpu::graphite::InsertRecordingInfo& info) {
   AutoLock auto_lock(this);
@@ -313,12 +319,17 @@
   // graphite::Context.
   std::optional<skgpu::graphite::InsertRecordingInfo> info_copy;
   if (info.fFinishedProc && task_runner) {
-    info_copy = info;
+    info_copy = *info_ptr;
     std::tie(info_copy->fFinishedProc, info_copy->fFinishedContext) =
         CreateFinishedProcThreadSafe(info.fFinishedProc, info.fFinishedContext,
                                      std::move(task_runner));
     info_ptr = &info_copy.value();
   }
+  if (simulated_insert_status_ != skgpu::graphite::InsertStatus::kSuccess) {
+    info_copy = *info_ptr;
+    info_copy->fSimulatedStatus = simulated_insert_status_;
+    info_ptr = &info_copy.value();
+  }
 
   auto insert_status = graphite_context_->insertRecording(*info_ptr);
 
diff --git a/gpu/command_buffer/service/graphite_shared_context.h b/gpu/command_buffer/service/graphite_shared_context.h
index 9f4c82a..5230d7fe 100644
--- a/gpu/command_buffer/service/graphite_shared_context.h
+++ b/gpu/command_buffer/service/graphite_shared_context.h
@@ -226,6 +226,10 @@
 
   skgpu::GpuStatsFlags supportedGpuStats() const;
 
+  // Overrides `fSimulatedStatus` for any `insertRecording` call when passed a
+  // `status` other than `kSuccess`.
+  void set_simulated_insert_status(skgpu::graphite::InsertStatus status);
+
  private:
   class AutoLock;
 
@@ -254,6 +258,9 @@
   size_t num_pending_recordings_ = 0;
 
   raw_ptr<Delegate> delegate_ = nullptr;
+
+  skgpu::graphite::InsertStatus simulated_insert_status_ =
+      skgpu::graphite::InsertStatus::kSuccess;
 };
 
 }  // namespace gpu
diff --git a/gpu/command_buffer/service/graphite_utils.cc b/gpu/command_buffer/service/graphite_utils.cc
index 32f781c..9dd217e 100644
--- a/gpu/command_buffer/service/graphite_utils.cc
+++ b/gpu/command_buffer/service/graphite_utils.cc
@@ -39,7 +39,9 @@
                                 size_t dst_bytes_per_row,
                                 int src_x,
                                 int src_y) {
-  GraphiteFlush(context, recorder);
+  if (!GraphiteFlush(context, recorder)) {
+    return false;
+  }
 
   ReadPixelsContext read_context;
   const SkIRect src_rect =
@@ -75,6 +77,9 @@
 bool GraphiteFlushAndSubmit(GraphiteSharedContext* context,
                             skgpu::graphite::Recorder* recorder) {
   bool success = GraphiteFlush(context, recorder);
+  // We submit any pending GPU work despite insertRecording failing since the
+  // caller can expect any resources used to be eventually released when the
+  // submitted work is done on the GPU.
   context->submit();
   return success;
 }
diff --git a/gpu/command_buffer/service/shared_context_state.cc b/gpu/command_buffer/service/shared_context_state.cc
index bcbf59bd..ea954f6 100644
--- a/gpu/command_buffer/service/shared_context_state.cc
+++ b/gpu/command_buffer/service/shared_context_state.cc
@@ -5,6 +5,7 @@
 #include "gpu/command_buffer/service/shared_context_state.h"
 
 #include "base/compiler_specific.h"
+#include "base/debug/crash_logging.h"
 #include "base/debug/dump_without_crashing.h"
 #include "base/immediate_crash.h"
 #include "base/metrics/histogram_functions.h"
@@ -831,16 +832,21 @@
   return graphite_shared_context()->insertRecording(info);
 }
 
-void SharedContextState::FlushAndSubmit(bool sync_to_cpu) {
+bool SharedContextState::FlushAndSubmit(bool sync_to_cpu) {
   if (graphite_shared_context()) {
-    FlushGraphiteRecorder();
+    bool flush_succeeded = FlushGraphiteRecorder();
+    // We submit any pending GPU work despite insertRecording failing since the
+    // caller can expect any resources used to be eventually released when the
+    // submitted work is done on the GPU.
     graphite_shared_context()->submit(sync_to_cpu
                                           ? skgpu::graphite::SyncToCpu::kYes
                                           : skgpu::graphite::SyncToCpu::kNo);
+    return flush_succeeded;
   } else if (gr_context()) {
-    gr_context()->flushAndSubmit(sync_to_cpu ? GrSyncCpu::kYes
-                                             : GrSyncCpu::kNo);
+    gr_context()->flush();
+    return gr_context()->submit(sync_to_cpu ? GrSyncCpu::kYes : GrSyncCpu::kNo);
   }
+  return true;
 }
 
 bool SharedContextState::FlushWriteAccess(
@@ -877,7 +883,7 @@
   return success;
 }
 
-void SharedContextState::SubmitIfNecessary(
+bool SharedContextState::SubmitIfNecessary(
     std::vector<GrBackendSemaphore> signal_semaphores,
     bool need_graphite_submit) {
   if (graphite_shared_context() && need_graphite_submit) {
@@ -889,33 +895,35 @@
     // and DrDC is not enabled.
     CHECK(signal_semaphores.empty());
     graphite_shared_context()->submit(skgpu::graphite::SyncToCpu::kNo);
-    return;
+    return true;
   }
 
-  // Do nothing here if there is no context.
+  // Do nothing here if there is no Skia Ganesh GrContext.
   if (!gr_context()) {
-    return;
+    return true;
   }
 
   // Note that when DrDc is enabled, we need to call
   // AddVulkanCleanupTaskForSkiaFlush() on gpu main thread and do skia flush.
   // This will ensure that vulkan memory allocated on gpu main thread will be
   // cleaned up.
+  SCOPED_CRASH_KEY_BOOL("gpu", "DrDcEnabled", is_drdc_enabled_);
+  SCOPED_CRASH_KEY_NUMBER("gpu", "SignalSemaphores", signal_semaphores.size());
   if (!signal_semaphores.empty() || is_drdc_enabled_) {
-    // NOTE: The Graphite SharedImage representation does not set semaphores,
-    // and we are not enabling DrDC with Graphite.
-    CHECK(gr_context());
     GrFlushInfo flush_info = {
         .fNumSemaphores = signal_semaphores.size(),
         .fSignalSemaphores = signal_semaphores.data(),
     };
     gpu::AddVulkanCleanupTaskForSkiaFlush(vk_context_provider(), &flush_info);
 
-    auto result = gr_context()->flush(flush_info);
-    DCHECK_EQ(result, GrSemaphoresSubmitted::kYes);
+    if (gr_context()->flush(flush_info) != GrSemaphoresSubmitted::kYes) {
+      base::debug::DumpWithoutCrashing();
+      return false;
+    }
   }
 
   bool sync_cpu = gpu::ShouldVulkanSyncCpuForSkiaSubmit(vk_context_provider());
+  SCOPED_CRASH_KEY_BOOL("gpu", "SubmitSyncCpu", sync_cpu);
 
   // If DrDc is enabled, submit the gr_context() to ensure correct ordering
   // of vulkan commands between raster and display compositor.
@@ -927,10 +935,12 @@
   const bool need_submit =
       sync_cpu || !signal_semaphores.empty() || is_drdc_enabled_;
 
-  if (need_submit) {
-    CHECK(gr_context());
-    gr_context()->submit(sync_cpu ? GrSyncCpu::kYes : GrSyncCpu::kNo);
+  if (need_submit &&
+      !gr_context()->submit(sync_cpu ? GrSyncCpu::kYes : GrSyncCpu::kNo)) {
+    base::debug::DumpWithoutCrashing();
+    return false;
   }
+  return true;
 }
 
 bool SharedContextState::MakeCurrent(gl::GLSurface* surface, bool needs_gl) {
diff --git a/gpu/command_buffer/service/shared_context_state.h b/gpu/command_buffer/service/shared_context_state.h
index 84cb34f9..5f14710 100644
--- a/gpu/command_buffer/service/shared_context_state.h
+++ b/gpu/command_buffer/service/shared_context_state.h
@@ -146,10 +146,10 @@
 
   bool IsGLInitialized() const { return !!feature_info_; }
 
-  void FlushAndSubmit(bool sync_to_cpu);
+  bool FlushAndSubmit(bool sync_to_cpu);
   bool FlushGraphiteRecorder();
   bool FlushWriteAccess(SkiaImageRepresentation::ScopedWriteAccess* access);
-  void SubmitIfNecessary(std::vector<GrBackendSemaphore> signal_semaphores,
+  bool SubmitIfNecessary(std::vector<GrBackendSemaphore> signal_semaphores,
                          bool need_graphite_submit);
 
   // Returns true if context state is using GL, either for Skia to run on
diff --git a/gpu/command_buffer/service/shared_image/wrapped_graphite_texture_holder.cc b/gpu/command_buffer/service/shared_image/wrapped_graphite_texture_holder.cc
index 2e37279..9e8f001 100644
--- a/gpu/command_buffer/service/shared_image/wrapped_graphite_texture_holder.cc
+++ b/gpu/command_buffer/service/shared_image/wrapped_graphite_texture_holder.cc
@@ -21,7 +21,10 @@
 WrappedGraphiteTextureHolder::~WrappedGraphiteTextureHolder() {
   auto destroy_resource = [](scoped_refptr<SharedContextState> context_state,
                              skgpu::graphite::BackendTexture texture) {
-    if (texture.isValid()) {
+    // It's possible in some tests that teardown with unsubmitted commands for
+    // the WrappedGraphiteTextureHolder to be released after the GraphiteDawn
+    // context has been nulled out and is in the middle of being torn down.
+    if (context_state->graphite_shared_context() && texture.isValid()) {
       context_state->graphite_shared_context()->deleteBackendTexture(texture);
     }
   };
diff --git a/gpu/command_buffer/service/webgpu_decoder_impl.cc b/gpu/command_buffer/service/webgpu_decoder_impl.cc
index d407e31a..74ed3a37 100644
--- a/gpu/command_buffer/service/webgpu_decoder_impl.cc
+++ b/gpu/command_buffer/service/webgpu_decoder_impl.cc
@@ -2450,10 +2450,12 @@
   if (!begin_semaphores.empty()) {
     // gr_context() must not be null when begin_semaphores is not empty.
     DCHECK(shared_context_state_->gr_context());
-    bool wait_result = shared_context_state_->gr_context()->wait(
-        begin_semaphores.size(), begin_semaphores.data(),
-        /*deleteSemaphoresAfterWait=*/false);
-    DCHECK(wait_result);
+    if (!shared_context_state_->gr_context()->wait(
+            begin_semaphores.size(), begin_semaphores.data(),
+            /*deleteSemaphoresAfterWait=*/false)) {
+      DLOG(ERROR) << "ClearSharedImage: Skia GrContext wait failed";
+      return false;
+    }
   }
   auto* canvas = surface->getCanvas();
   SkColor4f clear_color;
@@ -2466,38 +2468,27 @@
 
   // It's ok to pass in empty GrFlushInfo here since SignalSemaphores()
   // will populate it with semaphores and call GrDirectContext::flush.
-  bool success = true;
   if (shared_context_state_->gr_context()) {
     skgpu::ganesh::Flush(surface);
   } else {
     DCHECK(shared_context_state_->graphite_shared_context());
     DCHECK(shared_context_state_->gpu_main_graphite_recorder());
-    success = GraphiteFlushAndSubmit(
-        shared_context_state_->graphite_shared_context(),
-        shared_context_state_->gpu_main_graphite_recorder());
-  }
-  if (success) {
-    representation->SetCleared();
+    if (!GraphiteFlushAndSubmit(
+            shared_context_state_->graphite_shared_context(),
+            shared_context_state_->gpu_main_graphite_recorder())) {
+      DLOG(ERROR) << "ClearSharedImage: GraphiteFlushAndSubmit failed";
+      return false;
+    }
   }
   // Transition the image back to the desired end state. This is used for
   // transitioning the image to the external queue for Vulkan/GL interop.
   scoped_write_access->ApplyBackendSurfaceEndState();
-
-  if (!end_semaphores.empty()) {
-    // gr_context() must not be null when end_semaphores is not empty.
-    DCHECK(shared_context_state_->gr_context());
-    GrFlushInfo flush_info = {
-        .fNumSemaphores = end_semaphores.size(),
-        .fSignalSemaphores = end_semaphores.data(),
-    };
-    // Note: this is a no-op if vk_context_provider is null.
-    AddVulkanCleanupTaskForSkiaFlush(
-        shared_context_state_->vk_context_provider(), &flush_info);
-    auto flush_result = shared_context_state_->gr_context()->flush(flush_info);
-    DCHECK(flush_result == GrSemaphoresSubmitted::kYes);
-    shared_context_state_->gr_context()->submit();
+  if (!shared_context_state_->SubmitIfNecessary(
+          std::move(end_semaphores), /*need_graphite_submit=*/false)) {
+    DLOG(ERROR) << "ClearSharedImage: SharedContextState submit failed";
+    return false;
   }
-
+  representation->SetCleared();
   return true;
 }
 
diff --git a/components/viz/test/test_gpu_service_holder.cc b/components/viz/test/test_gpu_service_holder.cc
index 3676b225..5774f25 100644
--- a/components/viz/test/test_gpu_service_holder.cc
+++ b/components/viz/test/test_gpu_service_holder.cc
@@ -266,6 +266,16 @@
   return gpu_service_->share_group();
 }
 
+gpu::GraphiteSharedContext* TestGpuServiceHolder::GetGraphiteSharedContext()
+    const {
+#if BUILDFLAG(SKIA_USE_DAWN)
+  if (gpu_service_->dawn_context_provider()) {
+    return gpu_service_->dawn_context_provider()->GetGraphiteSharedContext();
+  }
+#endif
+  return nullptr;
+}
+
 void TestGpuServiceHolder::ScheduleGpuMainTask(base::OnceClosure callback) {
   DCHECK(gpu_main_task_sequence_);
   gpu_main_task_sequence_->ScheduleTask(
diff --git a/components/viz/test/test_gpu_service_holder.h b/components/viz/test/test_gpu_service_holder.h
index c4f3e61..0928e83 100644
--- a/components/viz/test/test_gpu_service_holder.h
+++ b/components/viz/test/test_gpu_service_holder.h
@@ -31,6 +31,7 @@
 
 namespace gpu {
 class CommandBufferTaskExecutor;
+class GraphiteSharedContext;
 class SingleTaskSequence;
 #if BUILDFLAG(ENABLE_VULKAN)
 class VulkanImplementation;
@@ -131,6 +132,8 @@
   scoped_refptr<gpu::SharedContextState> GetSharedContextState() override;
   scoped_refptr<gl::GLShareGroup> GetShareGroup() override;
 
+  gpu::GraphiteSharedContext* GetGraphiteSharedContext() const;
+
  private:
   void InitializeOnGpuThread(const gpu::GpuPreferences& preferences,
                              base::WaitableEvent* completion);
diff --git a/gpu/command_buffer/tests/webgpu_mailbox_unittest.cc b/gpu/command_buffer/tests/webgpu_mailbox_unittest.cc
index adc1b4c8..b25bdcc 100644
--- a/gpu/command_buffer/tests/webgpu_mailbox_unittest.cc
+++ b/gpu/command_buffer/tests/webgpu_mailbox_unittest.cc
@@ -4,6 +4,7 @@
 
 #include "base/compiler_specific.h"
 #include "base/strings/stringprintf.h"
+#include "base/test/bind.h"
 #include "build/build_config.h"
 #include "components/viz/test/test_gpu_service_holder.h"
 #include "gpu/command_buffer/client/client_shared_image.h"
@@ -11,6 +12,7 @@
 #include "gpu/command_buffer/client/webgpu_implementation.h"
 #include "gpu/command_buffer/common/mailbox.h"
 #include "gpu/command_buffer/common/shared_image_usage.h"
+#include "gpu/command_buffer/service/graphite_shared_context.h"
 #include "gpu/command_buffer/service/webgpu_decoder.h"
 #include "gpu/command_buffer/tests/webgpu_test.h"
 #include "gpu/config/gpu_finch_features.h"
@@ -1471,6 +1473,51 @@
                      std::move(gl_surface1), std::move(gl_surface2)));
 }
 
+// Test that a SharedImage that is presented while still uninitialized stays
+// reported as uncleared if the Skia clear could not be inserted into the
+// Graphite context, so subsequent reads see lazy-cleared contents rather than
+// stale data.
+TEST_P(WebGPUMailboxTextureTest,
+       PresentClearLeavesSharedImageUnclearedOnInsertFailure) {
+  auto* graphite_shared_context =
+      GetGpuServiceHolder()->GetGraphiteSharedContext();
+  SKIP_TEST_IF(!graphite_shared_context);
+
+  // Create the shared image and fill it with a non-zero sentinel.
+  SharedImageInterface* sii = GetSharedImageInterface();
+  scoped_refptr<gpu::ClientSharedImage> shared_image =
+      sii->CreateSharedImage({GetParam().format,
+                              {1, 1},
+                              gfx::ColorSpace::CreateSRGB(),
+                              GetSharedImageUsage(AccessType::ReadWrite),
... (file diff truncated)

diff --git a/components/viz/test/test_gpu_service_holder.cc b/components/viz/test/test_gpu_service_holder.cc
index 3676b225..5774f25 100644
--- a/components/viz/test/test_gpu_service_holder.cc
+++ b/components/viz/test/test_gpu_service_holder.cc
@@ -266,6 +266,16 @@
   return gpu_service_->share_group();
 }
 
+gpu::GraphiteSharedContext* TestGpuServiceHolder::GetGraphiteSharedContext()
+    const {
+#if BUILDFLAG(SKIA_USE_DAWN)
+  if (gpu_service_->dawn_context_provider()) {
+    return gpu_service_->dawn_context_provider()->GetGraphiteSharedContext();
+  }
+#endif
+  return nullptr;
+}
+
 void TestGpuServiceHolder::ScheduleGpuMainTask(base::OnceClosure callback) {
   DCHECK(gpu_main_task_sequence_);
   gpu_main_task_sequence_->ScheduleTask(
diff --git a/components/viz/test/test_gpu_service_holder.h b/components/viz/test/test_gpu_service_holder.h
index c4f3e61..0928e83 100644
--- a/components/viz/test/test_gpu_service_holder.h
+++ b/components/viz/test/test_gpu_service_holder.h
@@ -31,6 +31,7 @@
 
 namespace gpu {
 class CommandBufferTaskExecutor;
+class GraphiteSharedContext;
 class SingleTaskSequence;
 #if BUILDFLAG(ENABLE_VULKAN)
 class VulkanImplementation;
@@ -131,6 +132,8 @@
   scoped_refptr<gpu::SharedContextState> GetSharedContextState() override;
   scoped_refptr<gl::GLShareGroup> GetShareGroup() override;
 
+  gpu::GraphiteSharedContext* GetGraphiteSharedContext() const;
+
  private:
   void InitializeOnGpuThread(const gpu::GpuPreferences& preferences,
                              base::WaitableEvent* completion);
diff --git a/gpu/command_buffer/tests/webgpu_mailbox_unittest.cc b/gpu/command_buffer/tests/webgpu_mailbox_unittest.cc
index adc1b4c8..b25bdcc 100644
--- a/gpu/command_buffer/tests/webgpu_mailbox_unittest.cc
+++ b/gpu/command_buffer/tests/webgpu_mailbox_unittest.cc
@@ -4,6 +4,7 @@
 
 #include "base/compiler_specific.h"
 #include "base/strings/stringprintf.h"
+#include "base/test/bind.h"
 #include "build/build_config.h"
 #include "components/viz/test/test_gpu_service_holder.h"
 #include "gpu/command_buffer/client/client_shared_image.h"
@@ -11,6 +12,7 @@
 #include "gpu/command_buffer/client/webgpu_implementation.h"
 #include "gpu/command_buffer/common/mailbox.h"
 #include "gpu/command_buffer/common/shared_image_usage.h"
+#include "gpu/command_buffer/service/graphite_shared_context.h"
 #include "gpu/command_buffer/service/webgpu_decoder.h"
 #include "gpu/command_buffer/tests/webgpu_test.h"
 #include "gpu/config/gpu_finch_features.h"
@@ -1471,6 +1473,51 @@
                      std::move(gl_surface1), std::move(gl_surface2)));
 }
 
+// Test that a SharedImage that is presented while still uninitialized stays
+// reported as uncleared if the Skia clear could not be inserted into the
+// Graphite context, so subsequent reads see lazy-cleared contents rather than
+// stale data.
+TEST_P(WebGPUMailboxTextureTest,
+       PresentClearLeavesSharedImageUnclearedOnInsertFailure) {
+  auto* graphite_shared_context =
+      GetGpuServiceHolder()->GetGraphiteSharedContext();
+  SKIP_TEST_IF(!graphite_shared_context);
+
+  // Create the shared image and fill it with a non-zero sentinel.
+  SharedImageInterface* sii = GetSharedImageInterface();
+  scoped_refptr<gpu::ClientSharedImage> shared_image =
+      sii->CreateSharedImage({GetParam().format,
+                              {1, 1},
+                              gfx::ColorSpace::CreateSRGB(),
+                              GetSharedImageUsage(AccessType::ReadWrite),
+                              "TestLabel"},
+                             kNullSurfaceHandle);
+  InitializeTextureColor(device_, shared_image, {1.0, 0, 0, 1.0});
+  WaitForCompletion(device_);
+
+  // Simulate a recoverable Graphite insertRecording() failure so any recorded
+  // work is dropped without losing the Skia context.
+  graphite_shared_context->set_simulated_insert_status(
+      skgpu::graphite::InsertStatus::kPromiseImageInstantiationFailed);
+
+  // Associate with DISCARD so the texture is treated as uninitialized, then
+  // present without writing. The decoder will attempt to clear via Skia and
+  // the recording insertion will be dropped.
+  wgpu::TextureDescriptor desc = {
+      .usage = wgpu::TextureUsage::RenderAttachment,
+  };
+  std::unique_ptr<WebGPUTextureScopedAccess> webgpu_scoped_access =
+      shared_image->BeginWebGPUTextureAccess(webgpu(), gpu::SyncToken(),
+                                             device_, desc, /*usage=*/0,
+                                             webgpu::WEBGPU_MAILBOX_DISCARD);
+  webgpu_scoped_access->SetNeedsPresent(true);
+
+  webgpu_impl()->SetLostContextCallback(
+      base::MakeExpectedRunClosure(FROM_HERE));
+  EXPECT_WEBGPU_DEVICE_LOST(device_, WebGPUTextureScopedAccess::EndAccess(
+                                         std::move(webgpu_scoped_access)));
+}
+
 INSTANTIATE_TEST_SUITE_P(
     ,
     WebGPUMailboxTextureTest,
diff --git a/gpu/command_buffer/tests/webgpu_test.cc b/gpu/command_buffer/tests/webgpu_test.cc
index 7f595f9..f4c6c94 100644
--- a/gpu/command_buffer/tests/webgpu_test.cc
+++ b/gpu/command_buffer/tests/webgpu_test.cc
@@ -42,9 +42,13 @@
 
 WebGPUTest::Options::Options() = default;
 
-std::map<std::pair<WGPUDevice, wgpu::ErrorType>, /* matched */ bool>
+// static
+std::map<std::pair<WGPUDevice, wgpu::ErrorType>, bool>
     WebGPUTest::s_expected_errors = {};
 
+// static
+std::map<WGPUDevice, bool> WebGPUTest::s_expected_devices_lost = {};
+
 WebGPUTest::WebGPUTest() = default;
 WebGPUTest::~WebGPUTest() = default;
 
@@ -243,8 +247,13 @@
 
   device_desc.SetDeviceLostCallback(
       wgpu::CallbackMode::AllowSpontaneous,
-      [](const wgpu::Device&, wgpu::DeviceLostReason reason,
+      [](const wgpu::Device& device, wgpu::DeviceLostReason reason,
          wgpu::StringView message) {
+        auto it = s_expected_devices_lost.find(device.Get());
+        if (it != s_expected_devices_lost.end() && !it->second) {
+          it->second = true;
+          return;
+        }
         if (reason == wgpu::DeviceLostReason::Destroyed) {
           return;
         }
diff --git a/gpu/command_buffer/tests/webgpu_test.h b/gpu/command_buffer/tests/webgpu_test.h
index fff847f..e68dc265 100644
--- a/gpu/command_buffer/tests/webgpu_test.h
+++ b/gpu/command_buffer/tests/webgpu_test.h
@@ -8,6 +8,7 @@
 #include <dawn/wire/client/webgpu_cpp.h>
 #include <dawn/wire/client/webgpu_cpp_print.h>
 
+#include <map>
 #include <memory>
 
 #include "build/build_config.h"
@@ -87,6 +88,8 @@
   static std::map<std::pair<WGPUDevice, wgpu::ErrorType>, /* matched */ bool>
       s_expected_errors;
 
+  static std::map<WGPUDevice, /* matched */ bool> s_expected_devices_lost;
+
   wgpu::Instance instance_ = nullptr;
   wgpu::Adapter adapter_ = nullptr;
 
@@ -110,6 +113,18 @@
     s_expected_errors.erase(it.first);                                         \
   } while (0)
 
+#define EXPECT_WEBGPU_DEVICE_LOST(device, statement)                        \
+  do {                                                                      \
+    PollUntilIdle();                                                        \
+    auto it = s_expected_devices_lost.insert({device.Get(), false});        \
+    EXPECT_TRUE(it.second) << "Only one expectation per-device supported."; \
+    statement;                                                              \
+    PollUntilIdle();                                                        \
+    EXPECT_TRUE(it.first->second)                                           \
+        << "Expected device lost in `" #statement "`";                      \
+    s_expected_devices_lost.erase(it.first);                                \
+  } while (0)
+
 }  // namespace gpu
 
 #endif  // GPU_COMMAND_BUFFER_TESTS_WEBGPU_TEST_H_