Bug 1940374 - Readback snapshot when last submitted TextureId is valid r=gfx-reviewers,lsalzman

Store last submitted TextureId in PresentationData for getting front buffer snapshot.

WebGPUParent::RecvQueueSubmit() stores the latest TextureId in the swap chain in PresentationData::LastSubmittedTextureId. To do so, it retrieves the OwnerId from ExternalTexture. Next, get PresentationData from OwnerId.

If readback is required, it is performed synchronously. Its implementation is borrowed from WebGPUParent::RecvSwapChainPresent().

Differential Revision: https://phabricator.services.mozilla.com/D233473

dom/webgpu/CanvasContext.cpp

@@ -331,7 +331,7 @@ NS_IMETHODIMP CanvasContext::GetInputStream(const char* aMimeType,
                                   aMimeType, aEncoderOptions, aStream);
 }
 
-already_AddRefed<mozilla::gfx::SourceSurface> CanvasContext::GetSurfaceSnapshot(
+already_AddRefed<gfx::SourceSurface> CanvasContext::GetSurfaceSnapshot(
     gfxAlphaType* aOutAlphaType) {
   if (aOutAlphaType) {
     *aOutAlphaType = gfxAlphaType::Premult;
@@ -347,9 +347,16 @@ already_AddRefed<mozilla::gfx::SourceSurface> CanvasContext::GetSurfaceSnapshot(
   }
 
   MOZ_ASSERT(mRemoteTextureOwnerId.isSome());
-  return cm->GetSnapshot(cm->Id(), mBridge->Id(), mRemoteTextureOwnerId,
-                         mGfxFormat, /* aPremultiply */ false,
+
+  // The parent side needs to create a command encoder which will be submitted
+  // and dropped right away so we create and release an encoder ID here.
+  RawId encoderId = ffi::wgpu_client_make_encoder_id(mBridge->GetClient());
+  RefPtr<gfx::SourceSurface> snapshot =
+      cm->GetSnapshot(cm->Id(), mBridge->Id(), mRemoteTextureOwnerId,
+                      Some(encoderId), mGfxFormat, /* aPremultiply */ false,
                       /* aYFlip */ false);
+  ffi::wgpu_client_free_command_encoder_id(mBridge->GetClient(), encoderId);
+  return snapshot.forget();
 }
 
 Maybe<layers::SurfaceDescriptor> CanvasContext::GetFrontBuffer(

dom/webgpu/CanvasContext.h

@@ -59,7 +59,7 @@ class CanvasContext final : public nsICanvasRenderingContextInternal,
   NS_IMETHOD GetInputStream(const char* aMimeType,
                             const nsAString& aEncoderOptions,
                             nsIInputStream** aStream) override;
-  already_AddRefed<mozilla::gfx::SourceSurface> GetSurfaceSnapshot(
+  already_AddRefed<gfx::SourceSurface> GetSurfaceSnapshot(
       gfxAlphaType* aOutAlphaType) override;
 
   void SetOpaqueValueFromOpaqueAttr(bool aOpaqueAttrValue) override {}

dom/webgpu/ExternalTexture.cpp

@@ -55,4 +55,21 @@ void ExternalTexture::SetSubmissionIndex(uint64_t aSubmissionIndex) {
   mSubmissionIndex = aSubmissionIndex;
 }
 
+UniquePtr<ExternalTextureReadBackPresent>
+ExternalTextureReadBackPresent::Create(
+    const uint32_t aWidth, const uint32_t aHeight,
+    const struct ffi::WGPUTextureFormat aFormat,
+    const ffi::WGPUTextureUsages aUsage) {
+  return MakeUnique<ExternalTextureReadBackPresent>(aWidth, aHeight, aFormat,
+                                                    aUsage);
+}
+
+ExternalTextureReadBackPresent::ExternalTextureReadBackPresent(
+    const uint32_t aWidth, const uint32_t aHeight,
+    const struct ffi::WGPUTextureFormat aFormat,
+    const ffi::WGPUTextureUsages aUsage)
+    : ExternalTexture(aWidth, aHeight, aFormat, aUsage) {}
+
+ExternalTextureReadBackPresent::~ExternalTextureReadBackPresent() {}
+
 }  // namespace mozilla::webgpu

dom/webgpu/ExternalTexture.h

@@ -57,6 +57,14 @@ class ExternalTexture {
   void SetSubmissionIndex(uint64_t aSubmissionIndex);
   uint64_t GetSubmissionIndex() const { return mSubmissionIndex; }
 
+  void SetOwnerId(const layers::RemoteTextureOwnerId aOwnerId) {
+    mOwnerId = aOwnerId;
+  }
+  layers::RemoteTextureOwnerId GetOwnerId() const {
+    MOZ_ASSERT(mOwnerId.IsValid());
+    return mOwnerId;
+  }
+
   const uint32_t mWidth;
   const uint32_t mHeight;
   const struct ffi::WGPUTextureFormat mFormat;
@@ -64,6 +72,26 @@ class ExternalTexture {
 
  protected:
   uint64_t mSubmissionIndex = 0;
+  layers::RemoteTextureOwnerId mOwnerId;
+};
+
+// Dummy class
+class ExternalTextureReadBackPresent final : public ExternalTexture {
+ public:
+  static UniquePtr<ExternalTextureReadBackPresent> Create(
+      const uint32_t aWidth, const uint32_t aHeight,
+      const struct ffi::WGPUTextureFormat aFormat,
+      const ffi::WGPUTextureUsages aUsage);
+
+  ExternalTextureReadBackPresent(const uint32_t aWidth, const uint32_t aHeight,
+                                 const struct ffi::WGPUTextureFormat aFormat,
+                                 const ffi::WGPUTextureUsages aUsage);
+  virtual ~ExternalTextureReadBackPresent();
+
+  Maybe<layers::SurfaceDescriptor> ToSurfaceDescriptor(
+      Maybe<gfx::FenceInfo>& aFenceInfo) override {
+    return Nothing();
+  }
 };
 
 }  // namespace webgpu

dom/webgpu/ipc/WebGPUParent.cpp

@@ -64,6 +64,16 @@ extern bool wgpu_server_ensure_external_texture_for_swap_chain(
       aSwapChainId, aDeviceId, aTextureId, aWidth, aHeight, aFormat, aUsage);
 }
 
+extern void wgpu_server_ensure_external_texture_for_readback(
+    void* aParam, WGPUSwapChainId aSwapChainId, WGPUDeviceId aDeviceId,
+    WGPUTextureId aTextureId, uint32_t aWidth, uint32_t aHeight,
+    struct WGPUTextureFormat aFormat, WGPUTextureUsages aUsage) {
+  auto* parent = static_cast<WebGPUParent*>(aParam);
+
+  parent->EnsureExternalTextureForReadBackPresent(
+      aSwapChainId, aDeviceId, aTextureId, aWidth, aHeight, aFormat, aUsage);
+}
+
 extern void* wgpu_server_get_external_texture_handle(void* aParam,
                                                      WGPUTextureId aId) {
   auto* parent = static_cast<WebGPUParent*>(aParam);
@@ -260,6 +270,7 @@ class PresentationData {
   bool mUseExternalTextureInSwapChain;
   const RawId mDeviceId;
   const RawId mQueueId;
+  Maybe<RawId> mLastSubmittedTextureId;
   const layers::RGBDescriptor mDesc;
 
   uint64_t mSubmissionIndex = 0;
@@ -275,6 +286,8 @@ class PresentationData {
   std::vector<RawId> mAvailableBufferIds;
   std::vector<RawId> mQueuedBufferIds;
 
+  bool mReadbackSnapshotCallbackCalled = false;
+
   PresentationData(WebGPUParent* aParent, bool aUseExternalTextureInSwapChain,
                    RawId aDeviceId, RawId aQueueId,
                    const layers::RGBDescriptor& aDesc, uint32_t aSourcePitch,
@@ -848,6 +861,13 @@ ipc::IPCResult WebGPUParent::RecvQueueSubmit(
         auto& externalTexture = it->second;
 
         externalTexture->SetSubmissionIndex(index);
+        // Update mLastSubmittedTextureId
+        auto ownerId = externalTexture->GetOwnerId();
+        const auto& lookup = mPresentationDataMap.find(ownerId);
+        if (lookup != mPresentationDataMap.end()) {
+          RefPtr<PresentationData> data = lookup->second.get();
+          data->mLastSubmittedTextureId = Some(textureId);
+        }
       }
     }
   }
@@ -1138,16 +1158,94 @@ static void ReadbackPresentCallback(uint8_t* userdata,
   }
 }
 
+struct ReadbackSnapshotRequest {
+  ReadbackSnapshotRequest(const ffi::WGPUGlobal* aContext,
+                          RefPtr<PresentationData>& aData,
+                          ffi::WGPUBufferId aBufferId,
+                          const ipc::Shmem& aDestShmem)
+      : mContext(aContext),
+        mData(aData),
+        mBufferId(aBufferId),
+        mDestShmem(aDestShmem) {}
+
+  const ffi::WGPUGlobal* mContext;
+  RefPtr<PresentationData> mData;
+  const ffi::WGPUBufferId mBufferId;
+  const ipc::Shmem& mDestShmem;
+};
+
+static void ReadbackSnapshotCallback(uint8_t* userdata,
+                                     ffi::WGPUBufferMapAsyncStatus status) {
+  UniquePtr<ReadbackSnapshotRequest> req(
+      reinterpret_cast<ReadbackSnapshotRequest*>(userdata));
+
+  RefPtr<PresentationData> data = req->mData;
+  data->mReadbackSnapshotCallbackCalled = true;
+
+  // Ensure we'll make the bufferId available for reuse
+  data->mAvailableBufferIds.push_back(req->mBufferId);
+
+  MOZ_LOG(sLogger, LogLevel::Info,
+          ("ReadbackSnapshotCallback for buffer %" PRIu64 " status=%d\n",
+           req->mBufferId, status));
+  if (status != ffi::WGPUBufferMapAsyncStatus_Success) {
+    return;
+  }
+  // copy the data
+  const auto bufferSize = data->mDesc.size().height * data->mSourcePitch;
+  ErrorBuffer getRangeError;
+  const auto mapped = ffi::wgpu_server_buffer_get_mapped_range(
+      req->mContext, req->mBufferId, 0, bufferSize, getRangeError.ToFFI());
+  getRangeError.CoerceValidationToInternal();
+  if (req->mData->mParent) {
+    req->mData->mParent->ForwardError(data->mDeviceId, getRangeError);
+  }
+  if (auto innerError = getRangeError.GetError()) {
+    MOZ_LOG(sLogger, LogLevel::Info,
+            ("WebGPU present: buffer get_mapped_range for internal "
+             "presentation readback failed: %s\n",
+             innerError->message.get()));
+    return;
+  }
+
+  MOZ_RELEASE_ASSERT(mapped.length >= bufferSize);
+
+  uint8_t* src = mapped.ptr;
+  uint8_t* dst = req->mDestShmem.get<uint8_t>();
+  const uint32_t stride = layers::ImageDataSerializer::ComputeRGBStride(
+      gfx::SurfaceFormat::B8G8R8A8, data->mDesc.size().width);
+
+  for (auto row = 0; row < data->mDesc.size().height; ++row) {
+    memcpy(dst, src, stride);
+    src += data->mSourcePitch;
+    dst += stride;
+  }
+
+  ErrorBuffer unmapError;
+  wgpu_server_buffer_unmap(req->mContext, req->mBufferId, unmapError.ToFFI());
+  unmapError.CoerceValidationToInternal();
+  if (req->mData->mParent) {
+    req->mData->mParent->ForwardError(data->mDeviceId, unmapError);
+  }
+  if (auto innerError = unmapError.GetError()) {
+    MOZ_LOG(sLogger, LogLevel::Info,
+            ("WebGPU snapshot: buffer unmap for internal presentation "
+             "readback failed: %s\n",
+             innerError->message.get()));
+  }
+}
+
 ipc::IPCResult WebGPUParent::GetFrontBufferSnapshot(
     IProtocol* aProtocol, const layers::RemoteTextureOwnerId& aOwnerId,
-    Maybe<Shmem>& aShmem, gfx::IntSize& aSize) {
+    const RawId& aCommandEncoderId, Maybe<Shmem>& aShmem, gfx::IntSize& aSize) {
   const auto& lookup = mPresentationDataMap.find(aOwnerId);
-  if (lookup == mPresentationDataMap.end() || !mRemoteTextureOwner ||
-      !mRemoteTextureOwner->IsRegistered(aOwnerId)) {
+  if (lookup == mPresentationDataMap.end()) {
+    MOZ_ASSERT_UNREACHABLE("unexpected to be called");
     return IPC_OK();
   }
 
   RefPtr<PresentationData> data = lookup->second.get();
+  data->mReadbackSnapshotCallbackCalled = false;
   aSize = data->mDesc.size();
   uint32_t stride = layers::ImageDataSerializer::ComputeRGBStride(
       data->mDesc.format(), aSize.width);
@@ -1157,13 +1255,151 @@ ipc::IPCResult WebGPUParent::GetFrontBufferSnapshot(
     return IPC_OK();
   }
 
+  if (data->mLastSubmittedTextureId.isNothing()) {
+    return IPC_OK();
+  }
+
+  auto it = mExternalTextures.find(data->mLastSubmittedTextureId.ref());
+
+  // ExternalTexture is still in use as no readback
+  if (data->mUseExternalTextureInSwapChain && it != mExternalTextures.end()) {
+    auto& externalTexture = it->second;
+    externalTexture->GetSnapshot(shmem, aSize);
+    return IPC_OK();
+  }
+
+  // External texture is already invalid and posted to RemoteTextureMap
+  if (it != mExternalTextures.end()) {
+    if (!mRemoteTextureOwner || !mRemoteTextureOwner->IsRegistered(aOwnerId)) {
+      MOZ_ASSERT_UNREACHABLE("unexpected to be called");
+      return IPC_OK();
+    }
     if (!data->mUseExternalTextureInSwapChain) {
       ffi::wgpu_server_device_poll(mContext.get(), data->mDeviceId, true);
     }
-
     mRemoteTextureOwner->GetLatestBufferSnapshot(aOwnerId, shmem, aSize);
     aShmem.emplace(std::move(shmem));
+    return IPC_OK();
+  }
 
+  // Readback synchronously
+
+  RawId bufferId = 0;
+  const auto& size = data->mDesc.size();
+  const auto bufferSize = data->mDesc.size().height * data->mSourcePitch;
+
+  // step 1: find an available staging buffer, or create one
+  {
+    if (!data->mAvailableBufferIds.empty()) {
+      bufferId = data->mAvailableBufferIds.back();
+      data->mAvailableBufferIds.pop_back();
+    } else if (!data->mUnassignedBufferIds.empty()) {
+      bufferId = data->mUnassignedBufferIds.back();
+      data->mUnassignedBufferIds.pop_back();
+
+      ffi::WGPUBufferUsages usage =
+          WGPUBufferUsages_COPY_DST | WGPUBufferUsages_MAP_READ;
+
+      ErrorBuffer error;
+      ffi::wgpu_server_device_create_buffer(mContext.get(), data->mDeviceId,
+                                            bufferId, nullptr, bufferSize,
+                                            usage, false, false, error.ToFFI());
+      if (ForwardError(data->mDeviceId, error)) {
+        return IPC_OK();
+      }
+    } else {
+      bufferId = 0;
+    }
+  }
+
+  MOZ_LOG(sLogger, LogLevel::Info,
+          ("GetFrontBufferSnapshot with buffer %" PRIu64 "\n", bufferId));
+  if (!bufferId) {
+    // TODO: add a warning - no buffer are available!
+    return IPC_OK();
+  }
+
+  // step 3: submit a copy command for the frame
+  ffi::WGPUCommandEncoderDescriptor encoderDesc = {};
+  {
+    ErrorBuffer error;
+    ffi::wgpu_server_device_create_encoder(mContext.get(), data->mDeviceId,
+                                           &encoderDesc, aCommandEncoderId,
+                                           error.ToFFI());
+    if (ForwardError(data->mDeviceId, error)) {
+      return IPC_OK();
+    }
+  }
+
+  if (data->mLastSubmittedTextureId.isNothing()) {
+    return IPC_OK();
+  }
+
+  const ffi::WGPUTexelCopyTextureInfo texView = {
+      data->mLastSubmittedTextureId.ref(),
+  };
+  const ffi::WGPUTexelCopyBufferLayout bufLayout = {
+      0,
+      &data->mSourcePitch,
+      nullptr,
+  };
+  const ffi::WGPUExtent3d extent = {
+      static_cast<uint32_t>(size.width),
+      static_cast<uint32_t>(size.height),
+      1,
+  };
+
+  {
+    ErrorBuffer error;
+    ffi::wgpu_server_encoder_copy_texture_to_buffer(
+        mContext.get(), aCommandEncoderId, &texView, bufferId, &bufLayout,
+        &extent, error.ToFFI());
+    if (ForwardError(data->mDeviceId, error)) {
+      return IPC_OK();
+    }
+  }
+  ffi::WGPUCommandBufferDescriptor commandDesc = {};
+  {
+    ErrorBuffer error;
+    ffi::wgpu_server_encoder_finish(mContext.get(), aCommandEncoderId,
+                                    &commandDesc, error.ToFFI());
+    if (ForwardError(data->mDeviceId, error)) {
+      ffi::wgpu_server_encoder_drop(mContext.get(), aCommandEncoderId);
+      return IPC_OK();
+    }
+  }
+
+  {
+    ErrorBuffer error;
+    ffi::wgpu_server_queue_submit(mContext.get(), data->mQueueId,
+                                  &aCommandEncoderId, 1, error.ToFFI());
+    ffi::wgpu_server_encoder_drop(mContext.get(), aCommandEncoderId);
+    if (ForwardError(data->mDeviceId, error)) {
+      return IPC_OK();
+    }
+  }
+
+  auto snapshotRequest = MakeUnique<ReadbackSnapshotRequest>(
+      mContext.get(), data, bufferId, shmem);
+
+  ffi::WGPUBufferMapClosure closure = {
+      &ReadbackSnapshotCallback,
+      reinterpret_cast<uint8_t*>(snapshotRequest.release())};
+
+  ErrorBuffer error;
+  ffi::wgpu_server_buffer_map(mContext.get(), bufferId, 0, bufferSize,
+                              ffi::WGPUHostMap_Read, closure, error.ToFFI());
+  if (ForwardError(data->mDeviceId, error)) {
+    return IPC_OK();
+  }
+
+  // Callback should be called during the poll.
+  ffi::wgpu_server_poll_all_devices(mContext.get(), true);
+
+  // Check if ReadbackSnapshotCallback is called.
+  MOZ_RELEASE_ASSERT(data->mReadbackSnapshotCallbackCalled == true);
+
+  aShmem.emplace(std::move(shmem));
   return IPC_OK();
 }
 
@@ -1231,6 +1467,8 @@ ipc::IPCResult WebGPUParent::RecvSwapChainPresent(
     std::shared_ptr<ExternalTexture> externalTexture = it->second;
     mExternalTextures.erase(it);
 
+    MOZ_ASSERT(externalTexture->GetOwnerId() == aOwnerId);
+
     PostExternalTexture(std::move(externalTexture), aRemoteTextureId, aOwnerId);
     return IPC_OK();
   }
@@ -1617,6 +1855,7 @@ bool WebGPUParent::EnsureExternalTextureForSwapChain(
     // Check if the texture is recyclable.
     if (texture->mWidth == aWidth && texture->mHeight == aHeight &&
         texture->mFormat.tag == aFormat.tag && texture->mUsage == aUsage) {
+      texture->SetOwnerId(ownerId);
       data->mRecycledExternalTextures.pop_front();
       mExternalTextures.emplace(aTextureId, texture);
       return true;
@@ -1624,14 +1863,44 @@ bool WebGPUParent::EnsureExternalTextureForSwapChain(
     data->mRecycledExternalTextures.clear();
   }
 
-  auto externalTexture = CreateExternalTexture(aDeviceId, aTextureId, aWidth,
-                                               aHeight, aFormat, aUsage);
+  auto externalTexture = CreateExternalTexture(
+      ownerId, aDeviceId, aTextureId, aWidth, aHeight, aFormat, aUsage);
   return static_cast<bool>(externalTexture);
 }
 
+void WebGPUParent::EnsureExternalTextureForReadBackPresent(
+    ffi::WGPUSwapChainId aSwapChainId, ffi::WGPUDeviceId aDeviceId,
+    ffi::WGPUTextureId aTextureId, uint32_t aWidth, uint32_t aHeight,
+    struct ffi::WGPUTextureFormat aFormat, ffi::WGPUTextureUsages aUsage) {
+  auto ownerId = layers::RemoteTextureOwnerId{aSwapChainId._0};
+  const auto& lookup = mPresentationDataMap.find(ownerId);
+  if (lookup == mPresentationDataMap.end()) {
+    MOZ_ASSERT_UNREACHABLE("unexpected to be called");
+    return;
+  }
+
+  RefPtr<PresentationData> data = lookup->second.get();
+  if (data->mUseExternalTextureInSwapChain) {
+    MOZ_ASSERT_UNREACHABLE("unexpected to be called");
+    return;
+  }
+
+  UniquePtr<ExternalTexture> texture =
+      ExternalTextureReadBackPresent::Create(aWidth, aHeight, aFormat, aUsage);
+  if (!texture) {
+    MOZ_ASSERT_UNREACHABLE("unexpected to be called");
+    return;
+  }
+
+  texture->SetOwnerId(ownerId);
+  std::shared_ptr<ExternalTexture> shared(texture.release());
+  mExternalTextures[aTextureId] = shared;
+}
+
 std::shared_ptr<ExternalTexture> WebGPUParent::CreateExternalTexture(
-    ffi::WGPUDeviceId aDeviceId, ffi::WGPUTextureId aTextureId, uint32_t aWidth,
-    uint32_t aHeight, const struct ffi::WGPUTextureFormat aFormat,
+    const layers::RemoteTextureOwnerId& aOwnerId, ffi::WGPUDeviceId aDeviceId,
+    ffi::WGPUTextureId aTextureId, uint32_t aWidth, uint32_t aHeight,
+    const struct ffi::WGPUTextureFormat aFormat,
     ffi::WGPUTextureUsages aUsage) {
   MOZ_RELEASE_ASSERT(mExternalTextures.find(aTextureId) ==
                      mExternalTextures.end());
@@ -1642,6 +1911,7 @@ std::shared_ptr<ExternalTexture> WebGPUParent::CreateExternalTexture(
     return nullptr;
   }
 
+  texture->SetOwnerId(aOwnerId);
   std::shared_ptr<ExternalTexture> shared(texture.release());
   mExternalTextures.emplace(aTextureId, shared);
 

dom/webgpu/ipc/WebGPUParent.h

@@ -135,7 +135,8 @@ class WebGPUParent final : public PWebGPUParent, public SupportsWeakPtr {
 
   ipc::IPCResult GetFrontBufferSnapshot(
       IProtocol* aProtocol, const layers::RemoteTextureOwnerId& aOwnerId,
-      Maybe<Shmem>& aShmem, gfx::IntSize& aSize);
+      const RawId& aCommandEncoderId, Maybe<Shmem>& aShmem,
+      gfx::IntSize& aSize);
 
   void ActorDestroy(ActorDestroyReason aWhy) override;
 
@@ -161,9 +162,14 @@ class WebGPUParent final : public PWebGPUParent, public SupportsWeakPtr {
                                          struct ffi::WGPUTextureFormat aFormat,
                                          ffi::WGPUTextureUsages aUsage);
 
+  void EnsureExternalTextureForReadBackPresent(
+      ffi::WGPUSwapChainId aSwapChainId, ffi::WGPUDeviceId aDeviceId,
+      ffi::WGPUTextureId aTextureId, uint32_t aWidth, uint32_t aHeight,
+      struct ffi::WGPUTextureFormat aFormat, ffi::WGPUTextureUsages aUsage);
+
   std::shared_ptr<ExternalTexture> CreateExternalTexture(
-      ffi::WGPUDeviceId aDeviceId, ffi::WGPUTextureId aTextureId,
-      uint32_t aWidth, uint32_t aHeight,
+      const layers::RemoteTextureOwnerId& aOwnerId, ffi::WGPUDeviceId aDeviceId,
+      ffi::WGPUTextureId aTextureId, uint32_t aWidth, uint32_t aHeight,
       const struct ffi::WGPUTextureFormat aFormat,
       ffi::WGPUTextureUsages aUsage);
 

gfx/ipc/CanvasManagerChild.cpp

@@ -230,14 +230,16 @@ layers::ActiveResourceTracker* CanvasManagerChild::GetActiveResourceTracker() {
 
 already_AddRefed<DataSourceSurface> CanvasManagerChild::GetSnapshot(
     uint32_t aManagerId, int32_t aProtocolId,
-    const Maybe<RemoteTextureOwnerId>& aOwnerId, SurfaceFormat aFormat,
+    const Maybe<RemoteTextureOwnerId>& aOwnerId,
+    const Maybe<RawId>& aCommandEncoderId, SurfaceFormat aFormat,
     bool aPremultiply, bool aYFlip) {
   if (!CanSend()) {
     return nullptr;
   }
 
   webgl::FrontBufferSnapshotIpc res;
-  if (!SendGetSnapshot(aManagerId, aProtocolId, aOwnerId, &res)) {
+  if (!SendGetSnapshot(aManagerId, aProtocolId, aOwnerId, aCommandEncoderId,
+                       &res)) {
     return nullptr;
   }
 

gfx/ipc/CanvasManagerChild.h

@@ -39,7 +39,8 @@ class CanvasManagerChild final : public PCanvasManagerChild {
   uint32_t Id() const { return mId; }
   already_AddRefed<DataSourceSurface> GetSnapshot(
       uint32_t aManagerId, int32_t aProtocolId,
-      const Maybe<RemoteTextureOwnerId>& aOwnerId, SurfaceFormat aFormat,
+      const Maybe<RemoteTextureOwnerId>& aOwnerId,
+      const Maybe<RawId>& aCommandEncoderId, SurfaceFormat aFormat,
       bool aPremultiply, bool aYFlip);
   void ActorDestroy(ActorDestroyReason aReason) override;
 

gfx/ipc/CanvasManagerParent.cpp

@@ -180,6 +180,7 @@ CanvasManagerParent::AllocPCanvasParent() {
 mozilla::ipc::IPCResult CanvasManagerParent::RecvGetSnapshot(
     const uint32_t& aManagerId, const int32_t& aProtocolId,
     const Maybe<RemoteTextureOwnerId>& aOwnerId,
+    const Maybe<RawId>& aCommandEncoderId,
     webgl::FrontBufferSnapshotIpc* aResult) {
   if (!aManagerId) {
     return IPC_FAIL(this, "invalid id");
@@ -217,8 +218,11 @@ mozilla::ipc::IPCResult CanvasManagerParent::RecvGetSnapshot(
       if (aOwnerId.isNothing()) {
         return IPC_FAIL(this, "invalid OwnerId");
       }
-      mozilla::ipc::IPCResult rv =
-          webgpu->GetFrontBufferSnapshot(this, *aOwnerId, buffer.shmem, size);
+      if (aCommandEncoderId.isNothing()) {
+        return IPC_FAIL(this, "invalid CommandEncoderId");
+      }
+      mozilla::ipc::IPCResult rv = webgpu->GetFrontBufferSnapshot(
+          this, *aOwnerId, *aCommandEncoderId, buffer.shmem, size);
       if (!rv) {
         return rv;
       }

gfx/ipc/CanvasManagerParent.h

@@ -49,6 +49,7 @@ class CanvasManagerParent final : public PCanvasManagerParent {
   mozilla::ipc::IPCResult RecvGetSnapshot(
       const uint32_t& aManagerId, const int32_t& aProtocolId,
       const Maybe<RemoteTextureOwnerId>& aOwnerId,
+      const Maybe<RawId>& aCommandEncoderId,
       webgl::FrontBufferSnapshotIpc* aResult);
 
  private:

gfx/ipc/PCanvasManager.ipdl

@@ -12,6 +12,7 @@ include protocol PWebGPU;
 
 using mozilla::layers::RemoteTextureOwnerId from "mozilla/layers/LayersTypes.h";
 using mozilla::webgl::FrontBufferSnapshotIpc from "mozilla/dom/WebGLIpdl.h";
+using mozilla::webgpu::RawId from "mozilla/webgpu/WebGPUTypes.h";
 
 namespace mozilla {
 namespace gfx {
@@ -45,7 +46,7 @@ parent:
   // intended to be used by the main thread in the content process to block
   // reading without having to block on the worker thread that owns the context
   // instance.
-  sync GetSnapshot(uint32_t aManagerId, int32_t aProtocolId, RemoteTextureOwnerId? ownerId) returns (FrontBufferSnapshotIpc ret);
+  sync GetSnapshot(uint32_t aManagerId, int32_t aProtocolId, RemoteTextureOwnerId? ownerId, RawId? commandEncoderId) returns (FrontBufferSnapshotIpc ret);
 };
 
 } // gfx

gfx/wgpu_bindings/src/server.rs

@@ -1414,6 +1414,17 @@ extern "C" {
         usage: wgt::TextureUsages,
     ) -> bool;
     #[allow(dead_code)]
+    fn wgpu_server_ensure_external_texture_for_readback(
+        param: *mut c_void,
+        swap_chain_id: SwapChainId,
+        device_id: id::DeviceId,
+        texture_id: id::TextureId,
+        width: u32,
+        height: u32,
+        format: wgt::TextureFormat,
+        usage: wgt::TextureUsages,
+    );
+    #[allow(dead_code)]
     fn wgpu_server_get_external_texture_handle(
         param: *mut c_void,
         id: id::TextureId,
@@ -2035,6 +2046,21 @@ impl Global {
                     };
                 }
 
+                if let Some(swap_chain_id) = swap_chain_id {
+                    unsafe {
+                        wgpu_server_ensure_external_texture_for_readback(
+                            self.owner,
+                            swap_chain_id,
+                            self_id,
+                            id,
+                            desc.size.width,
+                            desc.size.height,
+                            desc.format,
+                            desc.usage,
+                        )
+                    };
+                }
+
                 let (_, error) = self.device_create_texture(self_id, &desc, Some(id));
                 if let Some(err) = error {
                     error_buf.init(err);