webgpu: Factor out swapchain to separate file (#33367)

* Move some stuff to swapchain.rs

Signed-off-by: sagudev <16504129+sagudev@users.noreply.github.com>

* Use typed WebGPUContextId instead of u64

Signed-off-by: sagudev <16504129+sagudev@users.noreply.github.com>

* Extract create_swapchain function and move more stuff in it

Signed-off-by: sagudev <16504129+sagudev@users.noreply.github.com>

* extract destroy_swapchain

Signed-off-by: sagudev <16504129+sagudev@users.noreply.github.com>

* extract swapchain_present

Signed-off-by: sagudev <16504129+sagudev@users.noreply.github.com>

* extract update_wr_image callback

Signed-off-by: sagudev <16504129+sagudev@users.noreply.github.com>

* fixup

Signed-off-by: sagudev <16504129+sagudev@users.noreply.github.com>

---------

Signed-off-by: sagudev <16504129+sagudev@users.noreply.github.com>

components/constellation/constellation.rs

@@ -153,6 +153,7 @@ use servo_rand::{random, Rng, ServoRng, SliceRandom};
 use servo_url::{Host, ImmutableOrigin, ServoUrl};
 use style_traits::CSSPixel;
 use tracing::{span, Level};
+use webgpu::swapchain::WGPUImageMap;
 use webgpu::{self, WebGPU, WebGPURequest, WebGPUResponse};
 use webrender::{RenderApi, RenderApiSender};
 use webrender_api::DocumentId;
@@ -219,7 +220,7 @@ struct WebrenderWGPU {
     webrender_external_images: Arc<Mutex<WebrenderExternalImageRegistry>>,
 
     /// WebGPU data that supplied to Webrender for rendering
-    wgpu_image_map: Arc<Mutex<HashMap<u64, webgpu::PresentationData>>>,
+    wgpu_image_map: WGPUImageMap,
 }
 
 /// Servo supports multiple top-level browsing contexts or “webviews”, so `Constellation` needs to
@@ -552,7 +553,7 @@ pub struct InitialConstellationState {
     /// User agent string to report in network requests.
     pub user_agent: Cow<'static, str>,
 
-    pub wgpu_image_map: Arc<Mutex<HashMap<u64, webgpu::PresentationData>>>,
+    pub wgpu_image_map: WGPUImageMap,
 }
 
 /// Data needed for webdriver

components/script/dom/bindings/trace.rs

@@ -61,7 +61,6 @@ use crate::dom::bindings::refcounted::{Trusted, TrustedPromise};
 use crate::dom::bindings::reflector::{DomObject, Reflector};
 use crate::dom::bindings::root::{Dom, DomRoot};
 use crate::dom::bindings::str::{DOMString, USVString};
-use crate::dom::gpucanvascontext::WebGPUContextId;
 use crate::dom::htmlimageelement::SourceSet;
 use crate::dom::htmlmediaelement::HTMLMediaElementFetchContext;
 use crate::dom::windowproxy::WindowProxyHandler;
@@ -364,7 +363,6 @@ unsafe_no_jsmanaged_fields!(ContextForRequestInterrupt);
 unsafe_no_jsmanaged_fields!(WindowProxyHandler);
 unsafe_no_jsmanaged_fields!(DOMString);
 unsafe_no_jsmanaged_fields!(USVString);
-unsafe_no_jsmanaged_fields!(WebGPUContextId);
 unsafe_no_jsmanaged_fields!(SourceSet);
 unsafe_no_jsmanaged_fields!(HTMLMediaElementFetchContext);
 unsafe_no_jsmanaged_fields!(StreamConsumer);

components/script/dom/document.rs

@@ -65,6 +65,7 @@ use style::stylesheet_set::DocumentStylesheetSet;
 use style::stylesheets::{Origin, OriginSet, Stylesheet};
 use url::Host;
 use uuid::Uuid;
+use webgpu::swapchain::WebGPUContextId;
 use webrender_api::units::DeviceIntRect;
 
 use crate::animation_timeline::AnimationTimeline;
@@ -123,7 +124,7 @@ use crate::dom::eventtarget::EventTarget;
 use crate::dom::focusevent::FocusEvent;
 use crate::dom::fontfaceset::FontFaceSet;
 use crate::dom::globalscope::GlobalScope;
-use crate::dom::gpucanvascontext::{GPUCanvasContext, WebGPUContextId};
+use crate::dom::gpucanvascontext::GPUCanvasContext;
 use crate::dom::hashchangeevent::HashChangeEvent;
 use crate::dom::htmlanchorelement::HTMLAnchorElement;
 use crate::dom::htmlareaelement::HTMLAreaElement;
@@ -442,7 +443,7 @@ pub struct Document {
     dirty_webgl_contexts:
         DomRefCell<HashMapTracedValues<WebGLContextId, Dom<WebGLRenderingContext>>>,
     /// List of all WebGPU context IDs that need flushing.
-    dirty_webgpu_contexts: DomRefCell<HashMap<WebGPUContextId, Dom<GPUCanvasContext>>>,
+    dirty_webgpu_contexts: DomRefCell<HashMapTracedValues<WebGPUContextId, Dom<GPUCanvasContext>>>,
     /// <https://html.spec.whatwg.org/multipage/#concept-document-csp-list>
     #[ignore_malloc_size_of = "Defined in rust-content-security-policy"]
     #[no_trace]
@@ -3278,7 +3279,7 @@ impl Document {
             shadow_roots_styles_changed: Cell::new(false),
             media_controls: DomRefCell::new(HashMap::new()),
             dirty_webgl_contexts: DomRefCell::new(HashMapTracedValues::new()),
-            dirty_webgpu_contexts: DomRefCell::new(HashMap::new()),
+            dirty_webgpu_contexts: DomRefCell::new(HashMapTracedValues::new()),
             csp_list: DomRefCell::new(None),
             selection: MutNullableDom::new(None),
             animation_timeline: if pref!(layout.animations.test.enabled) {

components/script/dom/gpucanvascontext.rs

@@ -9,12 +9,10 @@ use dom_struct::dom_struct;
 use euclid::default::Size2D;
 use ipc_channel::ipc;
 use script_layout_interface::HTMLCanvasDataSource;
+use webgpu::swapchain::WebGPUContextId;
 use webgpu::wgc::id;
-use webgpu::{wgt, WebGPU, WebGPURequest, WebGPUTexture, PRESENTATION_BUFFER_COUNT};
+use webgpu::{WebGPU, WebGPURequest, WebGPUTexture, PRESENTATION_BUFFER_COUNT};
-use webrender_api::{
+use webrender_api::{units, ImageFormat, ImageKey};
-    units, ExternalImageData, ExternalImageId, ExternalImageType, ImageData, ImageDescriptor,
-    ImageDescriptorFlags, ImageFormat, ImageKey,
-};
 
 use super::bindings::codegen::Bindings::WebGPUBinding::GPUTextureUsageConstants;
 use super::bindings::codegen::UnionTypes::HTMLCanvasElementOrOffscreenCanvas;
@@ -35,9 +33,6 @@ use crate::dom::globalscope::GlobalScope;
 use crate::dom::htmlcanvaselement::{HTMLCanvasElement, LayoutCanvasRenderingContextHelpers};
 use crate::dom::node::{document_from_node, Node, NodeDamage};
 
-#[derive(Clone, Copy, Debug, Eq, Hash, MallocSizeOf, Ord, PartialEq, PartialOrd)]
-pub struct WebGPUContextId(pub u64);
-
 // TODO: make all this derivables available via new Bindings.conf option
 impl Clone for GPUCanvasConfiguration {
     fn clone(&self) -> Self {
@@ -104,6 +99,7 @@ pub struct GPUCanvasContext {
     #[ignore_malloc_size_of = "Defined in webrender"]
     #[no_trace]
     webrender_image: Cell<Option<webrender_api::ImageKey>>,
+    #[no_trace]
     context_id: WebGPUContextId,
     /// <https://gpuweb.github.io/gpuweb/#dom-gpucanvascontext-currenttexture-slot>
     texture: MutNullableDom<GPUTexture>,
@@ -154,7 +150,7 @@ impl GPUCanvasContext {
             .wgpu_id_hub()
             .create_command_encoder_id(texture_id.backend());
         if let Err(e) = self.channel.0.send(WebGPURequest::SwapChainPresent {
-            external_id: self.context_id.0,
+            context_id: self.context_id,
             texture_id,
             encoder_id,
         }) {
@@ -250,22 +246,6 @@ impl GPUCanvasContextMethods for GPUCanvasContext {
         };
 
         // Step 8
-        let image_desc = ImageDescriptor {
-            format,
-            size: units::DeviceIntSize::new(size.width as i32, size.height as i32),
-            stride: Some(
-                (((size.width as u32 * 4) | (wgt::COPY_BYTES_PER_ROW_ALIGNMENT - 1)) + 1) as i32,
-            ),
-            offset: 0,
-            flags: ImageDescriptorFlags::from_bits(1).unwrap(),
-        };
-
-        let image_data = ImageData::External(ExternalImageData {
-            id: ExternalImageId(self.context_id.0),
-            channel_index: 0,
-            image_type: ExternalImageType::Buffer,
-        });
-
         let (sender, receiver) = ipc::channel().unwrap();
 
         let mut buffer_ids = ArrayVec::<id::BufferId, PRESENTATION_BUFFER_COUNT>::new();
@@ -283,10 +263,10 @@ impl GPUCanvasContextMethods for GPUCanvasContext {
                 device_id: descriptor.device.id().0,
                 queue_id: descriptor.device.GetQueue().id().0,
                 buffer_ids,
-                external_id: self.context_id.0,
+                context_id: self.context_id,
                 sender,
-                image_desc,
+                format,
-                image_data,
+                size: units::DeviceIntSize::new(size.width as i32, size.height as i32),
             })
             .expect("Failed to create WebGPU SwapChain");
 
@@ -301,7 +281,7 @@ impl GPUCanvasContextMethods for GPUCanvasContext {
     fn Unconfigure(&self) {
         if let Some(image_key) = self.webrender_image.take() {
             if let Err(e) = self.channel.0.send(WebGPURequest::DestroySwapChain {
-                external_id: self.context_id.0,
+                context_id: self.context_id,
                 image_key,
             }) {
                 warn!(

components/servo/lib.rs

@@ -19,7 +19,6 @@
 
 use std::borrow::{BorrowMut, Cow};
 use std::cmp::max;
-use std::collections::HashMap;
 use std::path::PathBuf;
 use std::rc::Rc;
 use std::sync::{Arc, Mutex};
@@ -90,6 +89,7 @@ use surfman::platform::generic::multi::context::NativeContext as LinuxNativeCont
 use surfman::{GLApi, GLVersion};
 #[cfg(all(target_os = "linux", not(target_env = "ohos")))]
 use surfman::{NativeConnection, NativeContext};
+use webgpu::swapchain::WGPUImageMap;
 use webrender::{RenderApiSender, ShaderPrecacheFlags, UploadMethod, ONE_TIME_USAGE_HINT};
 use webrender_api::{
     ColorF, DocumentId, FontInstanceFlags, FontInstanceKey, FontKey, FramePublishId, ImageKey,
@@ -1017,7 +1017,7 @@ fn create_constellation(
     glplayer_threads: Option<GLPlayerThreads>,
     initial_window_size: WindowSizeData,
     external_images: Arc<Mutex<WebrenderExternalImageRegistry>>,
-    wgpu_image_map: Arc<Mutex<HashMap<u64, webgpu::PresentationData>>>,
+    wgpu_image_map: WGPUImageMap,
     protocols: ProtocolRegistry,
 ) -> Sender<ConstellationMsg> {
     // Global configuration options, parsed from the command line.

components/webgpu/ipc_messages/recv.rs

@@ -10,7 +10,8 @@ use base::id::PipelineId;
 use ipc_channel::ipc::{IpcSender, IpcSharedMemory};
 use serde::{Deserialize, Serialize};
 use smallvec::SmallVec;
-use webrender_api::{ExternalImageId, ImageData, ImageDescriptor, ImageKey};
+use webrender_api::units::DeviceIntSize;
+use webrender_api::{ImageFormat, ImageKey};
 use wgc::binding_model::{
     BindGroupDescriptor, BindGroupLayoutDescriptor, PipelineLayoutDescriptor,
 };
@@ -30,6 +31,7 @@ pub use {wgpu_core as wgc, wgpu_types as wgt};
 
 use crate::identity::*;
 use crate::render_commands::RenderCommand;
+use crate::swapchain::WebGPUContextId;
 use crate::{Error, ErrorFilter, WebGPUResponse, PRESENTATION_BUFFER_COUNT};
 
 #[derive(Debug, Deserialize, Serialize)]
@@ -101,7 +103,7 @@ pub enum WebGPURequest {
         /// present only on ASYNC versions
         async_sender: Option<IpcSender<WebGPUResponse>>,
     },
-    CreateContext(IpcSender<ExternalImageId>),
+    CreateContext(IpcSender<WebGPUContextId>),
     CreatePipelineLayout {
         device_id: id::DeviceId,
         pipeline_layout_id: id::PipelineLayoutId,
@@ -131,10 +133,10 @@ pub enum WebGPURequest {
         device_id: id::DeviceId,
         queue_id: id::QueueId,
         buffer_ids: ArrayVec<id::BufferId, PRESENTATION_BUFFER_COUNT>,
-        external_id: u64,
+        context_id: WebGPUContextId,
         sender: IpcSender<ImageKey>,
-        image_desc: ImageDescriptor,
+        format: ImageFormat,
-        image_data: ImageData,
+        size: DeviceIntSize,
     },
     CreateTexture {
         device_id: id::DeviceId,
@@ -154,7 +156,7 @@ pub enum WebGPURequest {
         texture_id: id::TextureId,
     },
     DestroySwapChain {
-        external_id: u64,
+        context_id: WebGPUContextId,
         image_key: ImageKey,
     },
     DropTexture(id::TextureId),
@@ -256,7 +258,7 @@ pub enum WebGPURequest {
         command_buffers: Vec<id::CommandBufferId>,
     },
     SwapChainPresent {
-        external_id: u64,
+        context_id: WebGPUContextId,
         texture_id: id::TextureId,
         encoder_id: id::CommandEncoderId,
     },

components/webgpu/lib.rs

@@ -3,6 +3,8 @@
  * file, You can obtain one at https://mozilla.org/MPL/2.0/. */
 
 use log::warn;
+use swapchain::WGPUImageMap;
+pub use swapchain::{PresentationData, WGPUExternalImages};
 use webrender::RenderApiSender;
 use wgpu_thread::WGPU;
 pub use {wgpu_core as wgc, wgpu_types as wgt};
@@ -12,30 +14,25 @@ mod poll_thread;
 mod wgpu_thread;
 
 use std::borrow::Cow;
-use std::collections::HashMap;
 use std::sync::{Arc, Mutex};
 
-use arrayvec::ArrayVec;
-use euclid::default::Size2D;
 pub use gpu_error::{Error, ErrorFilter, PopError};
 use ipc_channel::ipc::{self, IpcReceiver, IpcSender};
 pub use render_commands::RenderCommand;
 use serde::{Deserialize, Serialize};
 use servo_config::pref;
-use webrender_api::{DocumentId, ImageData, ImageDescriptor, ImageKey};
+use webrender_api::DocumentId;
-use webrender_traits::{
+use webrender_traits::WebrenderExternalImageRegistry;
-    WebrenderExternalImageApi, WebrenderExternalImageRegistry, WebrenderImageSource,
-};
-use wgc::id;
 
 mod gpu_error;
 mod ipc_messages;
 mod render_commands;
+pub mod swapchain;
 pub use identity::*;
 pub use ipc_messages::recv::*;
 pub use ipc_messages::to_dom::*;
 pub use ipc_messages::to_script::*;
-pub use wgpu_thread::PRESENTATION_BUFFER_COUNT;
+pub use swapchain::PRESENTATION_BUFFER_COUNT;
 
 #[derive(Clone, Debug, Deserialize, Serialize)]
 pub struct WebGPU(pub IpcSender<WebGPURequest>);
@@ -45,7 +42,7 @@ impl WebGPU {
         webrender_api_sender: RenderApiSender,
         webrender_document: DocumentId,
         external_images: Arc<Mutex<WebrenderExternalImageRegistry>>,
-        wgpu_image_map: Arc<Mutex<HashMap<u64, PresentationData>>>,
+        wgpu_image_map: WGPUImageMap,
     ) -> Option<(Self, IpcReceiver<WebGPUMsg>)> {
         if !pref!(dom.webgpu.enabled) {
             return None;
@@ -100,46 +97,3 @@ impl WebGPU {
             .map_err(|_| "Failed to send Exit message")
     }
 }
-
-#[derive(Default)]
-pub struct WGPUExternalImages {
-    pub images: Arc<Mutex<HashMap<u64, PresentationData>>>,
-    pub locked_ids: HashMap<u64, Vec<u8>>,
-}
-
-impl WebrenderExternalImageApi for WGPUExternalImages {
-    fn lock(&mut self, id: u64) -> (WebrenderImageSource, Size2D<i32>) {
-        let size;
-        let data;
-        if let Some(present_data) = self.images.lock().unwrap().get(&id) {
-            size = present_data.size;
-            data = present_data.data.clone();
-        } else {
-            size = Size2D::new(0, 0);
-            data = Vec::new();
-        }
-        let _ = self.locked_ids.insert(id, data);
-        (
-            WebrenderImageSource::Raw(self.locked_ids.get(&id).unwrap().as_slice()),
-            size,
-        )
-    }
-
-    fn unlock(&mut self, id: u64) {
-        let _ = self.locked_ids.remove(&id);
-    }
-}
-
-pub struct PresentationData {
-    device_id: id::DeviceId,
-    queue_id: id::QueueId,
-    pub data: Vec<u8>,
-    pub size: Size2D<i32>,
-    unassigned_buffer_ids: ArrayVec<id::BufferId, PRESENTATION_BUFFER_COUNT>,
-    available_buffer_ids: ArrayVec<id::BufferId, PRESENTATION_BUFFER_COUNT>,
-    queued_buffer_ids: ArrayVec<id::BufferId, PRESENTATION_BUFFER_COUNT>,
-    buffer_stride: u32,
-    image_key: ImageKey,
-    image_desc: ImageDescriptor,
-    image_data: ImageData,
-}

components/webgpu/swapchain.rs

@@ -0,0 +1,337 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at https://mozilla.org/MPL/2.0/. */
+
+use std::collections::HashMap;
+use std::ops::ControlFlow;
+use std::slice;
+use std::sync::{Arc, Mutex, MutexGuard};
+
+use arrayvec::ArrayVec;
+use euclid::default::Size2D;
+use log::{error, warn};
+use malloc_size_of::MallocSizeOf;
+use serde::{Deserialize, Serialize};
+use webrender::{RenderApi, Transaction};
+use webrender_api::units::DeviceIntSize;
+use webrender_api::{
+    DirtyRect, ExternalImageData, ExternalImageId, ExternalImageType, ImageData, ImageDescriptor,
+    ImageDescriptorFlags, ImageFormat, ImageKey,
+};
+use webrender_traits::{WebrenderExternalImageApi, WebrenderImageSource};
+use wgpu_core::device::HostMap;
+use wgpu_core::global::Global;
+use wgpu_core::id;
+use wgpu_core::resource::{BufferAccessError, BufferMapCallback, BufferMapOperation};
+
+use crate::{wgt, WebGPUMsg};
+
+pub const PRESENTATION_BUFFER_COUNT: usize = 10;
+
+#[derive(Clone, Copy, Debug, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)]
+pub struct WebGPUContextId(pub u64);
+
+impl MallocSizeOf for WebGPUContextId {
+    fn size_of(&self, _ops: &mut malloc_size_of::MallocSizeOfOps) -> usize {
+        0
+    }
+}
+
+pub type WGPUImageMap = Arc<Mutex<HashMap<WebGPUContextId, PresentationData>>>;
+
+#[derive(Default)]
+pub struct WGPUExternalImages {
+    pub images: WGPUImageMap,
+    pub locked_ids: HashMap<WebGPUContextId, Vec<u8>>,
+}
+
+impl WebrenderExternalImageApi for WGPUExternalImages {
+    fn lock(&mut self, id: u64) -> (WebrenderImageSource, Size2D<i32>) {
+        let id = WebGPUContextId(id);
+        let size;
+        let data;
+        if let Some(present_data) = self.images.lock().unwrap().get(&id) {
+            size = present_data.image_desc.size.cast_unit();
+            data = present_data.data.clone();
+        } else {
+            size = Size2D::new(0, 0);
+            data = Vec::new();
+        }
+        let _ = self.locked_ids.insert(id, data);
+        (
+            WebrenderImageSource::Raw(self.locked_ids.get(&id).unwrap().as_slice()),
+            size,
+        )
+    }
+
+    fn unlock(&mut self, id: u64) {
+        let id = WebGPUContextId(id);
+        let _ = self.locked_ids.remove(&id);
+    }
+}
+
+pub struct PresentationData {
+    pub device_id: id::DeviceId,
+    pub queue_id: id::QueueId,
+    pub data: Vec<u8>,
+    pub unassigned_buffer_ids: ArrayVec<id::BufferId, PRESENTATION_BUFFER_COUNT>,
+    pub available_buffer_ids: ArrayVec<id::BufferId, PRESENTATION_BUFFER_COUNT>,
+    pub queued_buffer_ids: ArrayVec<id::BufferId, PRESENTATION_BUFFER_COUNT>,
+    pub image_key: ImageKey,
+    pub image_desc: ImageDescriptor,
+    pub image_data: ImageData,
+}
+
+impl PresentationData {
+    pub fn new(
+        device_id: id::DeviceId,
+        queue_id: id::QueueId,
+        buffer_ids: ArrayVec<id::BufferId, PRESENTATION_BUFFER_COUNT>,
+        image_key: ImageKey,
+        image_desc: ImageDescriptor,
+        image_data: ImageData,
+    ) -> Self {
+        let height = image_desc.size.height;
+        Self {
+            device_id,
+            queue_id,
+            // TODO: transparent black image
+            data: vec![
+                255;
+                (image_desc
+                    .stride
+                    .expect("Stride should be set when creating swapchain") *
+                    height) as usize
+            ],
+            unassigned_buffer_ids: buffer_ids,
+            available_buffer_ids: ArrayVec::<id::BufferId, PRESENTATION_BUFFER_COUNT>::new(),
+            queued_buffer_ids: ArrayVec::<id::BufferId, PRESENTATION_BUFFER_COUNT>::new(),
+            image_key,
+            image_desc,
+            image_data,
+        }
+    }
+}
+
+impl crate::WGPU {
+    pub(crate) fn create_swapchain(
+        &self,
+        device_id: id::DeviceId,
+        queue_id: id::QueueId,
+        buffer_ids: ArrayVec<id::BufferId, PRESENTATION_BUFFER_COUNT>,
+        context_id: WebGPUContextId,
+        format: ImageFormat,
+        size: DeviceIntSize,
+        image_key: ImageKey,
+        mut wr: MutexGuard<RenderApi>,
+    ) {
+        let image_desc = ImageDescriptor {
+            format,
+            size,
+            stride: Some(
+                (((size.width as u32 * 4) | (wgt::COPY_BYTES_PER_ROW_ALIGNMENT - 1)) + 1) as i32,
+            ),
+            offset: 0,
+            flags: ImageDescriptorFlags::IS_OPAQUE,
+        };
+
+        let image_data = ImageData::External(ExternalImageData {
+            id: ExternalImageId(context_id.0),
+            channel_index: 0,
+            image_type: ExternalImageType::Buffer,
+        });
+        let _ = self.wgpu_image_map.lock().unwrap().insert(
+            context_id,
+            PresentationData::new(
+                device_id,
+                queue_id,
+                buffer_ids,
+                image_key,
+                image_desc,
+                image_data.clone(),
+            ),
+        );
+
+        let mut txn = Transaction::new();
+        txn.add_image(image_key, image_desc, image_data, None);
+        wr.send_transaction(self.webrender_document, txn);
+    }
+
+    pub(crate) fn swapchain_present(
+        &mut self,
+        context_id: WebGPUContextId,
+        encoder_id: id::Id<id::markers::CommandEncoder>,
+        texture_id: id::Id<id::markers::Texture>,
+    ) -> ControlFlow<()> {
+        let global = &self.global;
+        let device_id;
+        let queue_id;
+        let size;
+        let buffer_id;
+        let buffer_stride;
+        {
+            if let Some(present_data) = self.wgpu_image_map.lock().unwrap().get_mut(&context_id) {
+                size = present_data.image_desc.size;
+                device_id = present_data.device_id;
+                queue_id = present_data.queue_id;
+                buffer_stride = present_data
+                    .image_desc
+                    .stride
+                    .expect("Stride should be set when creating swapchain");
+                buffer_id = if let Some(b_id) = present_data.available_buffer_ids.pop() {
+                    b_id
+                } else if let Some(b_id) = present_data.unassigned_buffer_ids.pop() {
+                    let buffer_size = (buffer_stride * size.height) as wgt::BufferAddress;
+                    let buffer_desc = wgt::BufferDescriptor {
+                        label: None,
+                        size: buffer_size,
+                        usage: wgt::BufferUsages::MAP_READ | wgt::BufferUsages::COPY_DST,
+                        mapped_at_creation: false,
+                    };
+                    let _ = global.device_create_buffer(device_id, &buffer_desc, Some(b_id));
+                    b_id
+                } else {
+                    error!("No staging buffer available for {:?}", context_id);
+                    return ControlFlow::Break(());
+                };
+                present_data.queued_buffer_ids.push(buffer_id);
+            } else {
+                error!("Data not found for {:?}", context_id);
+                return ControlFlow::Break(());
+            }
+        }
+        let buffer_size = (size.height * buffer_stride) as wgt::BufferAddress;
+        let comm_desc = wgt::CommandEncoderDescriptor { label: None };
+        let _ = global.device_create_command_encoder(device_id, &comm_desc, Some(encoder_id));
+        let buffer_cv = wgt::ImageCopyBuffer {
+            buffer: buffer_id,
+            layout: wgt::ImageDataLayout {
+                offset: 0,
+                bytes_per_row: Some(buffer_stride as u32),
+                rows_per_image: None,
+            },
+        };
+        let texture_cv = wgt::ImageCopyTexture {
+            texture: texture_id,
+            mip_level: 0,
+            origin: wgt::Origin3d::ZERO,
+            aspect: wgt::TextureAspect::All,
+        };
+        let copy_size = wgt::Extent3d {
+            width: size.width as u32,
+            height: size.height as u32,
+            depth_or_array_layers: 1,
+        };
+        let _ = global.command_encoder_copy_texture_to_buffer(
+            encoder_id,
+            &texture_cv,
+            &buffer_cv,
+            &copy_size,
+        );
+        let _ = global.command_encoder_finish(encoder_id, &wgt::CommandBufferDescriptor::default());
+        let _ = global.queue_submit(queue_id, &[encoder_id.into_command_buffer_id()]);
+        let callback = {
+            let global = Arc::clone(&self.global);
+            let wgpu_image_map = Arc::clone(&self.wgpu_image_map);
+            let webrender_api = Arc::clone(&self.webrender_api);
+            let webrender_document = self.webrender_document;
+            let token = self.poller.token();
+            BufferMapCallback::from_rust(Box::from(move |result| {
+                drop(token);
+                update_wr_image(
+                    result,
+                    global,
+                    buffer_id,
+                    buffer_size,
+                    wgpu_image_map,
+                    context_id,
+                    webrender_api,
+                    webrender_document,
+                );
+            }))
+        };
+        let map_op = BufferMapOperation {
+            host: HostMap::Read,
+            callback: Some(callback),
+        };
+        let _ = global.buffer_map_async(buffer_id, 0, Some(buffer_size), map_op);
+        self.poller.wake();
+
+        ControlFlow::Continue(())
+    }
+
+    pub(crate) fn destroy_swapchain(
+        &mut self,
+        context_id: WebGPUContextId,
+        image_key: webrender_api::ImageKey,
+    ) {
+        let data = self
+            .wgpu_image_map
+            .lock()
+            .unwrap()
+            .remove(&context_id)
+            .unwrap();
+        let global = &self.global;
+        for b_id in data.available_buffer_ids.iter() {
+            global.buffer_drop(*b_id);
+        }
+        for b_id in data.queued_buffer_ids.iter() {
+            global.buffer_drop(*b_id);
+        }
+        for b_id in data.unassigned_buffer_ids.iter() {
+            if let Err(e) = self.script_sender.send(WebGPUMsg::FreeBuffer(*b_id)) {
+                warn!("Unable to send FreeBuffer({:?}) ({:?})", *b_id, e);
+            };
+        }
+        let mut txn = Transaction::new();
+        txn.delete_image(image_key);
+        self.webrender_api
+            .lock()
+            .unwrap()
+            .send_transaction(self.webrender_document, txn);
+    }
+}
+
+fn update_wr_image(
+    result: Result<(), BufferAccessError>,
+    global: Arc<Global>,
+    buffer_id: id::BufferId,
+    buffer_size: u64,
+    wgpu_image_map: WGPUImageMap,
+    context_id: WebGPUContextId,
+    webrender_api: Arc<Mutex<RenderApi>>,
+    webrender_document: webrender_api::DocumentId,
+) {
+    match result {
+        Ok(()) => {
+            let (slice_pointer, range_size) = global
+                .buffer_get_mapped_range(buffer_id, 0, Some(buffer_size as u64))
+                .unwrap();
+            let data =
+                unsafe { slice::from_raw_parts(slice_pointer.as_ptr(), range_size as usize) }
+                    .to_vec();
+            if let Some(present_data) = wgpu_image_map.lock().unwrap().get_mut(&context_id) {
+                present_data.data = data;
+                let mut txn = Transaction::new();
+                txn.update_image(
+                    present_data.image_key,
+                    present_data.image_desc,
+                    present_data.image_data.clone(),
+                    &DirtyRect::All,
+                );
+                webrender_api
+                    .lock()
+                    .unwrap()
+                    .send_transaction(webrender_document, txn);
+                present_data
+                    .queued_buffer_ids
+                    .retain(|b_id| *b_id != buffer_id);
+                present_data.available_buffer_ids.push(buffer_id);
+            } else {
+                error!("Data not found for {:?}", context_id);
+            }
+            let _ = global.buffer_unmap(buffer_id);
+        },
+        _ => error!("Could not map buffer({:?})", buffer_id),
+    }
+}

components/webgpu/wgpu_thread.rs

@@ -6,23 +6,20 @@
 
 use std::borrow::Cow;
 use std::collections::HashMap;
+use std::ops::ControlFlow;
 use std::slice;
 use std::sync::{Arc, Mutex};
 
-use arrayvec::ArrayVec;
 use base::id::PipelineId;
-use euclid::default::Size2D;
 use ipc_channel::ipc::{IpcReceiver, IpcSender, IpcSharedMemory};
-use log::{error, info, warn};
+use log::{info, warn};
 use servo_config::pref;
-use webrender::{RenderApi, RenderApiSender, Transaction};
+use webrender::{RenderApi, RenderApiSender};
-use webrender_api::{DirtyRect, DocumentId};
+use webrender_api::DocumentId;
 use webrender_traits::{WebrenderExternalImageRegistry, WebrenderImageHandlerType};
-use wgc::command::{
+use wgc::command::{ComputePass, ComputePassDescriptor, RenderPass};
-    ComputePass, ComputePassDescriptor, ImageCopyBuffer, ImageCopyTexture, RenderPass,
-};
 use wgc::device::queue::SubmittedWorkDoneClosure;
-use wgc::device::{DeviceDescriptor, DeviceLostClosure, HostMap, ImplicitPipelineIds};
+use wgc::device::{DeviceDescriptor, DeviceLostClosure, ImplicitPipelineIds};
 use wgc::id;
 use wgc::id::DeviceId;
 use wgc::instance::parse_backends_from_comma_list;
@@ -39,13 +36,12 @@ pub use {wgpu_core as wgc, wgpu_types as wgt};
 use crate::gpu_error::ErrorScope;
 use crate::poll_thread::Poller;
 use crate::render_commands::apply_render_command;
+use crate::swapchain::{WGPUImageMap, WebGPUContextId};
 use crate::{
-    Adapter, ComputePassId, Error, Mapping, Pipeline, PopError, PresentationData, RenderPassId,
+    Adapter, ComputePassId, Error, Mapping, Pipeline, PopError, RenderPassId, WebGPU,
-    WebGPU, WebGPUAdapter, WebGPUDevice, WebGPUMsg, WebGPUQueue, WebGPURequest, WebGPUResponse,
+    WebGPUAdapter, WebGPUDevice, WebGPUMsg, WebGPUQueue, WebGPURequest, WebGPUResponse,
 };
 
-pub const PRESENTATION_BUFFER_COUNT: usize = 10;
-
 #[derive(Eq, Hash, PartialEq)]
 pub(crate) struct DeviceScope {
     pub device_id: DeviceId,
@@ -103,8 +99,8 @@ impl<P> Pass<P> {
 pub(crate) struct WGPU {
     receiver: IpcReceiver<WebGPURequest>,
     sender: IpcSender<WebGPURequest>,
-    script_sender: IpcSender<WebGPUMsg>,
+    pub(crate) script_sender: IpcSender<WebGPUMsg>,
-    global: Arc<wgc::global::Global>,
+    pub(crate) global: Arc<wgc::global::Global>,
     adapters: Vec<WebGPUAdapter>,
     devices: Arc<Mutex<HashMap<DeviceId, DeviceScope>>>,
     // Track invalid adapters https://gpuweb.github.io/gpuweb/#invalid
@@ -114,12 +110,12 @@ pub(crate) struct WGPU {
     /// because wgpu does not invalidate command encoder object
     /// (this is also reused for invalidation of command buffers)
     error_command_encoders: HashMap<id::CommandEncoderId, String>,
-    webrender_api: Arc<Mutex<RenderApi>>,
+    pub(crate) webrender_api: Arc<Mutex<RenderApi>>,
-    webrender_document: DocumentId,
+    pub(crate) webrender_document: DocumentId,
-    external_images: Arc<Mutex<WebrenderExternalImageRegistry>>,
+    pub(crate) external_images: Arc<Mutex<WebrenderExternalImageRegistry>>,
-    wgpu_image_map: Arc<Mutex<HashMap<u64, PresentationData>>>,
+    pub(crate) wgpu_image_map: WGPUImageMap,
     /// Provides access to poller thread
-    poller: Poller,
+    pub(crate) poller: Poller,
     /// Store compute passes
     compute_passes: HashMap<ComputePassId, Pass<ComputePass>>,
     /// Store render passes
@@ -134,7 +130,7 @@ impl WGPU {
         webrender_api_sender: RenderApiSender,
         webrender_document: DocumentId,
         external_images: Arc<Mutex<WebrenderExternalImageRegistry>>,
-        wgpu_image_map: Arc<Mutex<HashMap<u64, PresentationData>>>,
+        wgpu_image_map: WGPUImageMap,
     ) -> Self {
         let backend_pref = pref!(dom.webgpu.wgpu_backend);
         let backends = if backend_pref.is_empty() {
@@ -419,7 +415,7 @@ impl WGPU {
                             .lock()
                             .expect("Lock poisoned?")
                             .next_id(WebrenderImageHandlerType::WebGPU);
-                        if let Err(e) = sender.send(id) {
+                        if let Err(e) = sender.send(WebGPUContextId(id.0)) {
                             warn!("Failed to send ExternalImageId to new context ({})", e);
                         };
                     },
@@ -525,46 +521,20 @@ impl WGPU {
                         device_id,
                         queue_id,
                         buffer_ids,
-                        external_id,
+                        context_id,
                         sender,
-                        image_desc,
+                        size,
-                        image_data,
+                        format,
                     } => {
-                        let height = image_desc.size.height;
+                        let wr = self.webrender_api.lock().unwrap();
-                        let width = image_desc.size.width;
-                        let buffer_stride =
-                            ((width * 4) as u32 | (wgt::COPY_BYTES_PER_ROW_ALIGNMENT - 1)) + 1;
-                        let mut wr = self.webrender_api.lock().unwrap();
                         let image_key = wr.generate_image_key();
                         if let Err(e) = sender.send(image_key) {
                             warn!("Failed to send ImageKey ({})", e);
                         }
-                        let _ = self.wgpu_image_map.lock().unwrap().insert(
+                        self.create_swapchain(
-                            external_id,
+                            device_id, queue_id, buffer_ids, context_id, format, size, image_key,
-                            PresentationData {
+                            wr,
-                                device_id,
+                        )
-                                queue_id,
-                                data: vec![255; (buffer_stride * height as u32) as usize],
-                                size: Size2D::new(width, height),
-                                unassigned_buffer_ids: buffer_ids,
-                                available_buffer_ids: ArrayVec::<
-                                    id::BufferId,
-                                    PRESENTATION_BUFFER_COUNT,
-                                >::new(),
-                                queued_buffer_ids: ArrayVec::<
-                                    id::BufferId,
-                                    PRESENTATION_BUFFER_COUNT,
-                                >::new(),
-                                buffer_stride,
-                                image_key,
-                                image_desc,
-                                image_data: image_data.clone(),
-                            },
-                        );
-
-                        let mut txn = Transaction::new();
-                        txn.add_image(image_key, image_desc, image_data, None);
-                        wr.send_transaction(self.webrender_document, txn);
                     },
                     WebGPURequest::CreateTexture {
                         device_id,
@@ -604,33 +574,10 @@ impl WGPU {
                         self.poller.wake();
                     },
                     WebGPURequest::DestroySwapChain {
-                        external_id,
+                        context_id,
                         image_key,
                     } => {
-                        let data = self
+                        self.destroy_swapchain(context_id, image_key);
-                            .wgpu_image_map
-                            .lock()
-                            .unwrap()
-                            .remove(&external_id)
-                            .unwrap();
-                        let global = &self.global;
-                        for b_id in data.available_buffer_ids.iter() {
-                            global.buffer_drop(*b_id);
-                        }
-                        for b_id in data.queued_buffer_ids.iter() {
-                            global.buffer_drop(*b_id);
-                        }
-                        for b_id in data.unassigned_buffer_ids.iter() {
-                            if let Err(e) = self.script_sender.send(WebGPUMsg::FreeBuffer(*b_id)) {
-                                warn!("Unable to send FreeBuffer({:?}) ({:?})", *b_id, e);
-                            };
-                        }
-                        let mut txn = Transaction::new();
-                        txn.delete_image(image_key);
-                        self.webrender_api
-                            .lock()
-                            .unwrap()
-                            .send_transaction(self.webrender_document, txn);
                     },
                     WebGPURequest::DestroyTexture {
                         device_id,
@@ -1040,160 +987,15 @@ impl WGPU {
                         self.maybe_dispatch_error(device_id, result.err());
                     },
                     WebGPURequest::SwapChainPresent {
-                        external_id,
+                        context_id,
                         texture_id,
                         encoder_id,
                     } => {
-                        let global = &self.global;
+                        if let ControlFlow::Break(_) =
-                        let device_id;
+                            self.swapchain_present(context_id, encoder_id, texture_id)
-                        let queue_id;
-                        let size;
-                        let buffer_id;
-                        let buffer_stride;
                         {
-                            if let Some(present_data) =
+                            continue;
-                                self.wgpu_image_map.lock().unwrap().get_mut(&external_id)
-                            {
-                                size = present_data.size;
-                                device_id = present_data.device_id;
-                                queue_id = present_data.queue_id;
-                                buffer_stride = present_data.buffer_stride;
-                                buffer_id = if let Some(b_id) =
-                                    present_data.available_buffer_ids.pop()
-                                {
-                                    b_id
-                                } else if let Some(b_id) = present_data.unassigned_buffer_ids.pop()
-                                {
-                                    let buffer_size =
-                                        (buffer_stride * size.height as u32) as wgt::BufferAddress;
-                                    let buffer_desc = wgt::BufferDescriptor {
-                                        label: None,
-                                        size: buffer_size,
-                                        usage: wgt::BufferUsages::MAP_READ |
-                                            wgt::BufferUsages::COPY_DST,
-                                        mapped_at_creation: false,
-                                    };
-                                    let _ = global.device_create_buffer(
-                                        device_id,
-                                        &buffer_desc,
-                                        Some(b_id),
-                                    );
-                                    b_id
-                                } else {
-                                    warn!(
-                                        "No staging buffer available for ExternalImageId({:?})",
-                                        external_id
-                                    );
-                                    continue;
-                                };
-                                present_data.queued_buffer_ids.push(buffer_id);
-                            } else {
-                                warn!("Data not found for ExternalImageId({:?})", external_id);
-                                continue;
-                            }
                         }
-
-                        let buffer_size =
-                            (size.height as u32 * buffer_stride) as wgt::BufferAddress;
-                        let comm_desc = wgt::CommandEncoderDescriptor { label: None };
-                        let _ = global.device_create_command_encoder(
-                            device_id,
-                            &comm_desc,
-                            Some(encoder_id),
-                        );
-
-                        let buffer_cv = ImageCopyBuffer {
-                            buffer: buffer_id,
-                            layout: wgt::ImageDataLayout {
-                                offset: 0,
-                                bytes_per_row: Some(buffer_stride),
-                                rows_per_image: None,
-                            },
-                        };
-                        let texture_cv = ImageCopyTexture {
-                            texture: texture_id,
-                            mip_level: 0,
-                            origin: wgt::Origin3d::ZERO,
-                            aspect: wgt::TextureAspect::All,
-                        };
-                        let copy_size = wgt::Extent3d {
-                            width: size.width as u32,
-                            height: size.height as u32,
-                            depth_or_array_layers: 1,
-                        };
-                        let _ = global.command_encoder_copy_texture_to_buffer(
-                            encoder_id,
-                            &texture_cv,
-                            &buffer_cv,
-                            &copy_size,
-                        );
-                        let _ = global.command_encoder_finish(
-                            encoder_id,
-                            &wgt::CommandBufferDescriptor::default(),
-                        );
-                        let _ =
-                            global.queue_submit(queue_id, &[encoder_id.into_command_buffer_id()]);
-
-                        let glob = Arc::clone(&self.global);
-                        let wgpu_image_map = Arc::clone(&self.wgpu_image_map);
-                        let webrender_api = Arc::clone(&self.webrender_api);
-                        let webrender_document = self.webrender_document;
-                        let token = self.poller.token();
-                        let callback = BufferMapCallback::from_rust(Box::from(move |result| {
-                            drop(token);
-                            match result {
-                                Ok(()) => {
-                                    let global = &glob;
-                                    let (slice_pointer, range_size) = global
-                                        .buffer_get_mapped_range(
-                                            buffer_id,
-                                            0,
-                                            Some(buffer_size as u64),
-                                        )
-                                        .unwrap();
-                                    let data = unsafe {
-                                        slice::from_raw_parts(
-                                            slice_pointer.as_ptr(),
-                                            range_size as usize,
-                                        )
-                                    }
-                                    .to_vec();
-                                    if let Some(present_data) =
-                                        wgpu_image_map.lock().unwrap().get_mut(&external_id)
-                                    {
-                                        present_data.data = data;
-                                        let mut txn = Transaction::new();
-                                        txn.update_image(
-                                            present_data.image_key,
-                                            present_data.image_desc,
-                                            present_data.image_data.clone(),
-                                            &DirtyRect::All,
-                                        );
-                                        webrender_api
-                                            .lock()
-                                            .unwrap()
-                                            .send_transaction(webrender_document, txn);
-                                        present_data
-                                            .queued_buffer_ids
-                                            .retain(|b_id| *b_id != buffer_id);
-                                        present_data.available_buffer_ids.push(buffer_id);
-                                    } else {
-                                        warn!(
-                                            "Data not found for ExternalImageId({:?})",
-                                            external_id
-                                        );
-                                    }
-                                    let _ = global.buffer_unmap(buffer_id);
-                                },
-                                _ => error!("Could not map buffer({:?})", buffer_id),
-                            }
-                        }));
-                        let map_op = BufferMapOperation {
-                            host: HostMap::Read,
-                            callback: Some(callback),
-                        };
-                        let _ = global.buffer_map_async(buffer_id, 0, Some(buffer_size), map_op);
-                        self.poller.wake();
                     },
                     WebGPURequest::UnmapBuffer {
                         buffer_id,