webgpu: Fix HTML event loop integration (#34631)

* webgpu: Fix HTML event loop integration

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

* Bring back self.drawing_buffer.borrow().cleared

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

* Rc webgpu_contexts

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

---------

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

components/script/dom/bindings/codegen/Bindings.conf

@@ -183,6 +183,10 @@ DOMInterfaces = {
     'canGc': ['MapAsync'],
 },
 
+'GPUCanvasContext': {
+    'weakReferenceable': True,
+},
+
 'GPUDevice': {
     'inRealms': [
         'CreateComputePipelineAsync',

components/script/dom/document.rs

@@ -3,7 +3,7 @@
  * file, You can obtain one at https://mozilla.org/MPL/2.0/. */
 
 use std::borrow::Cow;
-use std::cell::Cell;
+use std::cell::{Cell, RefCell};
 use std::cmp::Ordering;
 use std::collections::hash_map::Entry::{Occupied, Vacant};
 use std::collections::{HashMap, HashSet, VecDeque};
@@ -106,6 +106,7 @@ use crate::dom::bindings::reflector::{reflect_dom_object_with_proto, DomObject};
 use crate::dom::bindings::root::{Dom, DomRoot, DomSlice, LayoutDom, MutNullableDom};
 use crate::dom::bindings::str::{DOMString, USVString};
 use crate::dom::bindings::trace::{HashMapTracedValues, NoTrace};
+use crate::dom::bindings::weakref::WeakRef;
 use crate::dom::bindings::xmlname::XMLName::Invalid;
 use crate::dom::bindings::xmlname::{
     namespace_from_domstring, validate_and_extract, xml_name_type,
@@ -255,6 +256,9 @@ pub enum DeclarativeRefresh {
     CreatedAfterLoad,
 }
 
+pub(crate) type WebGPUContextsMap =
+    Rc<RefCell<HashMapTracedValues<WebGPUContextId, WeakRef<GPUCanvasContext>>>>;
+
 /// <https://dom.spec.whatwg.org/#document>
 #[dom_struct]
 pub struct Document {
@@ -450,9 +454,10 @@ pub struct Document {
     /// List of all WebGL context IDs that need flushing.
     dirty_webgl_contexts:
         DomRefCell<HashMapTracedValues<WebGLContextId, Dom<WebGLRenderingContext>>>,
-    /// List of all WebGPU context IDs that need flushing.
+    /// List of all WebGPU contexts.
     #[cfg(feature = "webgpu")]
-    dirty_webgpu_contexts: DomRefCell<HashMapTracedValues<WebGPUContextId, Dom<GPUCanvasContext>>>,
+    #[ignore_malloc_size_of = "Rc are hard"]
+    webgpu_contexts: WebGPUContextsMap,
     /// <https://w3c.github.io/slection-api/#dfn-selection>
     selection: MutNullableDom<Selection>,
     /// A timeline for animations which is used for synchronizing animations.
@@ -2999,20 +3004,19 @@ impl Document {
     }
 
     #[cfg(feature = "webgpu")]
-    pub fn add_dirty_webgpu_canvas(&self, context: &GPUCanvasContext) {
+    pub fn webgpu_contexts(&self) -> WebGPUContextsMap {
-        self.dirty_webgpu_contexts
+        self.webgpu_contexts.clone()
-            .borrow_mut()
-            .entry(context.context_id())
-            .or_insert_with(|| Dom::from_ref(context));
     }
 
     #[allow(crown::unrooted_must_root)]
     #[cfg(feature = "webgpu")]
-    pub fn flush_dirty_webgpu_canvases(&self) {
+    pub fn update_rendering_of_webgpu_canvases(&self) {
-        self.dirty_webgpu_contexts
+        self.webgpu_contexts
             .borrow_mut()
-            .drain()
+            .iter()
-            .for_each(|(_, context)| context.update_rendering_of_webgpu_canvas());
+            .filter_map(|(_, context)| context.root())
+            .filter(|context| context.onscreen())
+            .for_each(|context| context.update_rendering_of_webgpu_canvas());
     }
 
     pub fn id_map(&self) -> Ref<HashMapTracedValues<Atom, Vec<Dom<Element>>>> {
@@ -3386,7 +3390,7 @@ impl Document {
             media_controls: DomRefCell::new(HashMap::new()),
             dirty_webgl_contexts: DomRefCell::new(HashMapTracedValues::new()),
             #[cfg(feature = "webgpu")]
-            dirty_webgpu_contexts: DomRefCell::new(HashMapTracedValues::new()),
+            webgpu_contexts: Rc::new(RefCell::new(HashMapTracedValues::new())),
             selection: MutNullableDom::new(None),
             animation_timeline: if pref!(layout.animations.test.enabled) {
                 DomRefCell::new(AnimationTimeline::new_for_testing())

components/script/dom/webgpu/gpucanvascontext.rs

@@ -34,6 +34,8 @@ use crate::dom::bindings::inheritance::Castable;
 use crate::dom::bindings::reflector::{reflect_dom_object, DomObject, Reflector};
 use crate::dom::bindings::root::{DomRoot, LayoutDom, MutNullableDom};
 use crate::dom::bindings::str::USVString;
+use crate::dom::bindings::weakref::WeakRef;
+use crate::dom::document::WebGPUContextsMap;
 use crate::dom::globalscope::GlobalScope;
 use crate::dom::htmlcanvaselement::{HTMLCanvasElement, LayoutCanvasRenderingContextHelpers};
 use crate::dom::node::{document_from_node, Node, NodeDamage};
@@ -94,6 +96,9 @@ pub struct GPUCanvasContext {
     drawing_buffer: RefCell<DrawingBuffer>,
     /// <https://gpuweb.github.io/gpuweb/#dom-gpucanvascontext-currenttexture-slot>
     current_texture: MutNullableDom<GPUTexture>,
+    /// This is used for clearing
+    #[ignore_malloc_size_of = "Rc are hard"]
+    webgpu_contexts: WebGPUContextsMap,
 }
 
 impl GPUCanvasContext {
@@ -101,6 +106,7 @@ impl GPUCanvasContext {
         global: &GlobalScope,
         canvas: HTMLCanvasElementOrOffscreenCanvas,
         channel: WebGPU,
+        webgpu_contexts: WebGPUContextsMap,
     ) -> Self {
         let (sender, receiver) = ipc::channel().unwrap();
         let size = canvas.size().cast().cast_unit();
@@ -130,19 +136,27 @@ impl GPUCanvasContext {
             configuration: RefCell::new(None),
             texture_descriptor: RefCell::new(None),
             current_texture: MutNullableDom::default(),
+            webgpu_contexts,
         }
     }
 
     pub fn new(global: &GlobalScope, canvas: &HTMLCanvasElement, channel: WebGPU) -> DomRoot<Self> {
-        reflect_dom_object(
+        let document = document_from_node(canvas);
+        let this = reflect_dom_object(
             Box::new(GPUCanvasContext::new_inherited(
                 global,
                 HTMLCanvasElementOrOffscreenCanvas::HTMLCanvasElement(DomRoot::from_ref(canvas)),
                 channel,
+                document.webgpu_contexts(),
             )),
             global,
             CanGc::note(),
-        )
+        );
+        this.webgpu_contexts
+            .borrow_mut()
+            .entry(this.context_id())
+            .or_insert_with(|| WeakRef::new(&this));
+        this
     }
 }
 
@@ -257,8 +271,16 @@ impl GPUCanvasContext {
     pub(crate) fn mark_as_dirty(&self) {
         if let HTMLCanvasElementOrOffscreenCanvas::HTMLCanvasElement(canvas) = &self.canvas {
             canvas.upcast::<Node>().dirty(NodeDamage::OtherNodeDamage);
-            let document = document_from_node(&**canvas);
+        }
-            document.add_dirty_webgpu_canvas(self);
+    }
+
+    pub(crate) fn onscreen(&self) -> bool {
+        match self.canvas {
+            HTMLCanvasElementOrOffscreenCanvas::HTMLCanvasElement(ref canvas) => {
+                canvas.upcast::<Node>().is_connected()
+            },
+            // FIXME(#34628): Handle this properly
+            HTMLCanvasElementOrOffscreenCanvas::OffscreenCanvas(_) => false,
         }
     }
 
@@ -351,26 +373,27 @@ impl GPUCanvasContextMethods<crate::DomTypeHolder> for GPUCanvasContext {
 
     /// <https://gpuweb.github.io/gpuweb/#dom-gpucanvascontext-getcurrenttexture>
     fn GetCurrentTexture(&self) -> Fallible<DomRoot<GPUTexture>> {
-        // Step 1
+        // Step 1: If this.[[configuration]] is null, throw an InvalidStateError and return.
         let configuration = self.configuration.borrow();
         let Some(configuration) = configuration.as_ref() else {
             return Err(Error::InvalidState);
         };
-        // Step 2
+        // Step 2: Assert this.[[textureDescriptor]] is not null.
         let texture_descriptor = self.texture_descriptor.borrow();
         let texture_descriptor = texture_descriptor.as_ref().unwrap();
         // Step 6
         let current_texture = if let Some(current_texture) = self.current_texture.get() {
             current_texture
         } else {
-            // Step 3&4
+            // Step 4.1
             self.replace_drawing_buffer();
+            // Step 4.2
             let current_texture = configuration.device.CreateTexture(texture_descriptor)?;
             self.current_texture.set(Some(&current_texture));
+            // We only need to mark new texture
+            self.mark_as_dirty();
             current_texture
         };
-        // Step 5
-        self.mark_as_dirty();
         // Step 6
         Ok(current_texture)
     }
@@ -378,6 +401,7 @@ impl GPUCanvasContextMethods<crate::DomTypeHolder> for GPUCanvasContext {
 
 impl Drop for GPUCanvasContext {
     fn drop(&mut self) {
+        self.webgpu_contexts.borrow_mut().remove(&self.context_id());
         if let Err(e) = self.channel.0.send(WebGPURequest::DestroyContext {
             context_id: self.context_id,
         }) {

components/script/dom/window.rs

@@ -1895,8 +1895,6 @@ impl Window {
         // up-to-date contents.
         let for_display = reflow_goal.needs_display();
         if for_display {
-            #[cfg(feature = "webgpu")]
-            document.flush_dirty_webgpu_canvases();
             document.flush_dirty_webgl_canvases();
         }
 

components/script/script_thread.rs

@@ -1621,6 +1621,9 @@ impl ScriptThread {
 
             // TODO(#31871): Update the rendering: consolidate all reflow calls into one here?
 
+            #[cfg(feature = "webgpu")]
+            document.update_rendering_of_webgpu_canvases();
+
             // > Step 22: For each doc of docs, update the rendering or user interface of
             // > doc and its node navigable to reflect the current state.
             let window = document.window();