Auto merge of #26619 - kunalmohan:gpu-client-id, r=kvark

Implement client-side logic for WebGPU id recycling

<!-- Please describe your changes on the following line: -->
1. Set up `WebGPUMsg` receiver in script thread and add respective message handler function.
2. Add `kill_xxx_id()` methods to `Identities`.

r?@kvark

---
<!-- Thank you for contributing to Servo! Please replace each `[ ]` by `[X]` when the step is complete, and replace `___` with appropriate data: -->
- [X] `./mach build -d` does not report any errors
- [X] `./mach test-tidy` does not report any errors
- [ ] These changes fix #___ (GitHub issue number if applicable)

<!-- Either: -->
- [ ] There are tests for these changes OR
- [ ] These changes do not require tests because ___

<!-- Also, please make sure that "Allow edits from maintainers" checkbox is checked, so that we can help you if you get stuck somewhere along the way.-->

<!-- Pull requests that do not address these steps are welcome, but they will require additional verification as part of the review process. -->

components/constellation/constellation.rs

@@ -2209,17 +2209,14 @@ where
             Some(bc) => &bc.bc_group_id,
             None => return warn!("Browsing context not found"),
         };
-        let host = match self
+        let source_pipeline = match self.pipelines.get(&source_pipeline_id) {
-            .pipelines
+            Some(pipeline) => pipeline,
-            .get(&source_pipeline_id)
-            .map(|pipeline| &pipeline.url)
-        {
-            Some(ref url) => match reg_host(&url) {
-                Some(host) => host,
-                None => return warn!("Invalid host url"),
-            },
             None => return warn!("ScriptMsg from closed pipeline {:?}.", source_pipeline_id),
         };
+        let host = match reg_host(&source_pipeline.url) {
+            Some(host) => host,
+            None => return warn!("Invalid host url"),
+        };
         match self
             .browsing_context_group_set
             .get_mut(&browsing_context_group_id)
@@ -2238,7 +2235,16 @@ where
                 let send = match browsing_context_group.webgpus.entry(host) {
                     Entry::Vacant(v) => v
                         .insert(match WebGPU::new() {
-                            Some(webgpu) => webgpu,
+                            Some(webgpu) => {
+                                let msg = ConstellationControlMsg::SetWebGPUPort(webgpu.1);
+                                if let Err(e) = source_pipeline.event_loop.send(msg) {
+                                    warn!(
+                                        "Failed to send SetWebGPUPort to pipeline {} ({:?})",
+                                        source_pipeline_id, e
+                                    );
+                                }
+                                webgpu.0
+                            },
                             None => return warn!("Failed to create new WebGPU thread"),
                         })
                         .0

components/msg/constellation_msg.rs

@@ -655,6 +655,7 @@ pub enum ScriptHangAnnotation {
     WebVREvent,
     PerformanceTimelineTask,
     PortMessage,
+    WebGPUMsg,
 }
 
 #[derive(Clone, Copy, Debug, Deserialize, Serialize)]

components/profile/time.rs

@@ -151,6 +151,7 @@ impl Formattable for ProfilerCategory {
             ProfilerCategory::ScriptWebVREvent => "Script WebVR Event",
             ProfilerCategory::ScriptWorkletEvent => "Script Worklet Event",
             ProfilerCategory::ScriptPerformanceEvent => "Script Performance Event",
+            ProfilerCategory::ScriptWebGPUMsg => "Script WebGPU Message",
             ProfilerCategory::TimeToFirstPaint => "Time To First Paint",
             ProfilerCategory::TimeToFirstContentfulPaint => "Time To First Contentful Paint",
             ProfilerCategory::TimeToInteractive => "Time to Interactive",

components/profile_traits/time.rs

@@ -109,6 +109,7 @@ pub enum ProfilerCategory {
     ScriptPerformanceEvent = 0x7b,
     ScriptHistoryEvent = 0x7c,
     ScriptPortMessage = 0x7d,
+    ScriptWebGPUMsg = 0x7e,
     TimeToFirstPaint = 0x80,
     TimeToFirstContentfulPaint = 0x81,
     TimeToInteractive = 0x82,

components/script/dom/identityhub.rs

@@ -74,7 +74,7 @@ impl IdentityHub {
         self.shader_modules.alloc(self.backend)
     }
 
-    pub fn create_command_encoder_id(&mut self) -> CommandEncoderId {
+    fn create_command_encoder_id(&mut self) -> CommandEncoderId {
         self.command_encoders.alloc(self.backend)
     }
 }
@@ -141,6 +141,10 @@ impl Identities {
         self.select(backend).create_device_id()
     }
 
+    pub fn kill_device_id(&mut self, id: DeviceId) {
+        self.select(id.backend()).devices.free(id);
+    }
+
     pub fn create_adapter_ids(&mut self) -> SmallVec<[AdapterId; 4]> {
         let mut ids = SmallVec::new();
         for hub in self.hubs() {
@@ -149,31 +153,63 @@ impl Identities {
         ids
     }
 
+    pub fn kill_adapter_id(&mut self, id: AdapterId) {
+        self.select(id.backend()).adapters.free(id);
+    }
+
     pub fn create_buffer_id(&mut self, backend: Backend) -> BufferId {
         self.select(backend).create_buffer_id()
     }
 
+    pub fn kill_buffer_id(&mut self, id: BufferId) {
+        self.select(id.backend()).buffers.free(id);
+    }
+
     pub fn create_bind_group_id(&mut self, backend: Backend) -> BindGroupId {
         self.select(backend).create_bind_group_id()
     }
 
+    pub fn kill_bind_group_id(&mut self, id: BindGroupId) {
+        self.select(id.backend()).bind_groups.free(id);
+    }
+
     pub fn create_bind_group_layout_id(&mut self, backend: Backend) -> BindGroupLayoutId {
         self.select(backend).create_bind_group_layout_id()
     }
 
+    pub fn kill_bind_group_layout_id(&mut self, id: BindGroupLayoutId) {
+        self.select(id.backend()).bind_group_layouts.free(id);
+    }
+
     pub fn create_compute_pipeline_id(&mut self, backend: Backend) -> ComputePipelineId {
         self.select(backend).create_compute_pipeline_id()
     }
 
+    pub fn kill_compute_pipeline_id(&mut self, id: ComputePipelineId) {
+        self.select(id.backend()).compute_pipelines.free(id);
+    }
+
     pub fn create_pipeline_layout_id(&mut self, backend: Backend) -> PipelineLayoutId {
         self.select(backend).create_pipeline_layout_id()
     }
 
+    pub fn kill_pipeline_layout_id(&mut self, id: PipelineLayoutId) {
+        self.select(id.backend()).pipeline_layouts.free(id);
+    }
+
     pub fn create_shader_module_id(&mut self, backend: Backend) -> ShaderModuleId {
         self.select(backend).create_shader_module_id()
     }
 
+    pub fn kill_shader_module_id(&mut self, id: ShaderModuleId) {
+        self.select(id.backend()).shader_modules.free(id);
+    }
+
     pub fn create_command_encoder_id(&mut self, backend: Backend) -> CommandEncoderId {
         self.select(backend).create_command_encoder_id()
     }
+
+    pub fn kill_command_buffer_id(&mut self, id: CommandEncoderId) {
+        self.select(id.backend()).command_encoders.free(id);
+    }
 }

components/script/dom/workletglobalscope.rs

@@ -159,7 +159,7 @@ pub struct WorkletGlobalScopeInit {
     pub is_headless: bool,
     /// An optional string allowing the user agent to be set for testing
     pub user_agent: Cow<'static, str>,
-    /// Channel to WebGPU
+    /// Identity manager for WebGPU resources
     pub gpu_id_hub: Arc<Mutex<Identities>>,
 }
 

components/script/script_runtime.rs

@@ -157,6 +157,7 @@ pub enum ScriptThreadEventCategory {
     EnterFullscreen,
     ExitFullscreen,
     PerformanceTimelineTask,
+    WebGPUMsg,
 }
 
 /// An interface for receiving ScriptMsg values in an event loop. Used for synchronous DOM

components/script/script_thread.rs

@@ -168,6 +168,7 @@ use style::dom::OpaqueNode;
 use style::thread_state::{self, ThreadState};
 use time::{at_utc, get_time, precise_time_ns, Timespec};
 use url::Position;
+use webgpu::identity::WebGPUMsg;
 use webrender_api::units::LayoutPixel;
 use webrender_api::DocumentId;
 
@@ -267,6 +268,7 @@ enum MixedMessage {
     FromScript(MainThreadScriptMsg),
     FromDevtools(DevtoolScriptControlMsg),
     FromImageCache((PipelineId, PendingImageResponse)),
+    FromWebGPUServer(WebGPUMsg),
 }
 
 /// Messages used to control the script event loop.
@@ -692,8 +694,11 @@ pub struct ScriptThread {
     /// Code is running as a consequence of a user interaction
     is_user_interacting: Cell<bool>,
 
-    /// Channel to WebGPU
+    /// Identity manager for WebGPU resources
     gpu_id_hub: Arc<Mutex<Identities>>,
+
+    /// Receiver to receive commands from optional WebGPU server.
+    webgpu_port: RefCell<Option<Receiver<WebGPUMsg>>>,
 }
 
 /// In the event of thread panic, all data on the stack runs its destructor. However, there
@@ -1383,6 +1388,7 @@ impl ScriptThread {
             node_ids: Default::default(),
             is_user_interacting: Cell::new(false),
             gpu_id_hub: Arc::new(Mutex::new(Identities::new())),
+            webgpu_port: RefCell::new(None),
         }
     }
 
@@ -1405,7 +1411,9 @@ impl ScriptThread {
     /// Handle incoming control messages.
     fn handle_msgs(&self) -> bool {
         use self::MixedMessage::FromScript;
-        use self::MixedMessage::{FromConstellation, FromDevtools, FromImageCache};
+        use self::MixedMessage::{
+            FromConstellation, FromDevtools, FromImageCache, FromWebGPUServer,
+        };
 
         // Handle pending resize events.
         // Gather them first to avoid a double mut borrow on self.
@@ -1445,6 +1453,8 @@ impl ScriptThread {
             recv(self.devtools_chan.as_ref().map(|_| &self.devtools_port).unwrap_or(&crossbeam_channel::never())) -> msg
                 => FromDevtools(msg.unwrap()),
             recv(self.image_cache_port) -> msg => FromImageCache(msg.unwrap()),
+            recv(self.webgpu_port.borrow().as_ref().unwrap_or(&crossbeam_channel::never())) -> msg
+                => FromWebGPUServer(msg.unwrap()),
         };
         debug!("Got event.");
 
@@ -1540,7 +1550,13 @@ impl ScriptThread {
                 Err(_) => match self.task_queue.try_recv() {
                     Err(_) => match self.devtools_port.try_recv() {
                         Err(_) => match self.image_cache_port.try_recv() {
-                            Err(_) => break,
+                            Err(_) => match &*self.webgpu_port.borrow() {
+                                Some(p) => match p.try_recv() {
+                                    Err(_) => break,
+                                    Ok(ev) => event = FromWebGPUServer(ev),
+                                },
+                                None => break,
+                            },
                             Ok(ev) => event = FromImageCache(ev),
                         },
                         Ok(ev) => event = FromDevtools(ev),
@@ -1572,6 +1588,7 @@ impl ScriptThread {
                     FromScript(inner_msg) => self.handle_msg_from_script(inner_msg),
                     FromDevtools(inner_msg) => self.handle_msg_from_devtools(inner_msg),
                     FromImageCache(inner_msg) => self.handle_msg_from_image_cache(inner_msg),
+                    FromWebGPUServer(inner_msg) => self.handle_msg_from_webgpu_server(inner_msg),
                 }
 
                 None
@@ -1659,6 +1676,7 @@ impl ScriptThread {
                 },
                 _ => ScriptThreadEventCategory::ScriptEvent,
             },
+            MixedMessage::FromWebGPUServer(_) => ScriptThreadEventCategory::WebGPUMsg,
         }
     }
 
@@ -1698,6 +1716,7 @@ impl ScriptThread {
                 ScriptHangAnnotation::PerformanceTimelineTask
             },
             ScriptThreadEventCategory::PortMessage => ScriptHangAnnotation::PortMessage,
+            ScriptThreadEventCategory::WebGPUMsg => ScriptHangAnnotation::WebGPUMsg,
         };
         self.background_hang_monitor
             .as_ref()
@@ -1743,6 +1762,7 @@ impl ScriptThread {
                 PaintMetric(..) => None,
                 ExitFullScreen(id, ..) => Some(id),
                 MediaSessionAction(..) => None,
+                SetWebGPUPort(..) => None,
             },
             MixedMessage::FromDevtools(_) => None,
             MixedMessage::FromScript(ref inner_msg) => match *inner_msg {
@@ -1756,6 +1776,7 @@ impl ScriptThread {
                 MainThreadScriptMsg::WakeUp => None,
             },
             MixedMessage::FromImageCache((pipeline_id, _)) => Some(pipeline_id),
+            MixedMessage::FromWebGPUServer(..) => None,
         }
     }
 
@@ -1810,6 +1831,7 @@ impl ScriptThread {
                 ScriptThreadEventCategory::PerformanceTimelineTask => {
                     ProfilerCategory::ScriptPerformanceEvent
                 },
+                ScriptThreadEventCategory::WebGPUMsg => ProfilerCategory::ScriptWebGPUMsg,
             };
             profile(profiler_cat, None, self.time_profiler_chan.clone(), f)
         } else {
@@ -1959,6 +1981,14 @@ impl ScriptThread {
             ConstellationControlMsg::MediaSessionAction(pipeline_id, action) => {
                 self.handle_media_session_action(pipeline_id, action)
             },
+            ConstellationControlMsg::SetWebGPUPort(port) => {
+                if self.webgpu_port.borrow().is_some() {
+                    warn!("WebGPU port already exists for this content process");
+                } else {
+                    let p = ROUTER.route_ipc_receiver_to_new_crossbeam_receiver(port);
+                    *self.webgpu_port.borrow_mut() = Some(p);
+                }
+            },
             msg @ ConstellationControlMsg::AttachLayout(..) |
             msg @ ConstellationControlMsg::Viewport(..) |
             msg @ ConstellationControlMsg::SetScrollState(..) |
@@ -1970,6 +2000,25 @@ impl ScriptThread {
         }
     }
 
+    fn handle_msg_from_webgpu_server(&self, msg: WebGPUMsg) {
+        match msg {
+            WebGPUMsg::FreeAdapter(id) => self.gpu_id_hub.lock().kill_adapter_id(id),
+            WebGPUMsg::FreeDevice(id) => self.gpu_id_hub.lock().kill_device_id(id),
+            WebGPUMsg::FreeBuffer(id) => self.gpu_id_hub.lock().kill_buffer_id(id),
+            WebGPUMsg::FreePipelineLayout(id) => self.gpu_id_hub.lock().kill_pipeline_layout_id(id),
+            WebGPUMsg::FreeComputePipeline(id) => {
+                self.gpu_id_hub.lock().kill_compute_pipeline_id(id)
+            },
+            WebGPUMsg::FreeBindGroup(id) => self.gpu_id_hub.lock().kill_bind_group_id(id),
+            WebGPUMsg::FreeBindGroupLayout(id) => {
+                self.gpu_id_hub.lock().kill_bind_group_layout_id(id)
+            },
+            WebGPUMsg::FreeCommandBuffer(id) => self.gpu_id_hub.lock().kill_command_buffer_id(id),
+            WebGPUMsg::FreeShaderModule(id) => self.gpu_id_hub.lock().kill_shader_module_id(id),
+            _ => {},
+        }
+    }
+
     fn handle_msg_from_script(&self, msg: MainThreadScriptMsg) {
         match msg {
             MainThreadScriptMsg::Common(CommonScriptMsg::Task(_, task, _, _)) => task.run_box(),

components/script_traits/lib.rs

@@ -66,6 +66,7 @@ use std::sync::atomic::AtomicBool;
 use std::sync::Arc;
 use style_traits::CSSPixel;
 use style_traits::SpeculativePainter;
+use webgpu::identity::WebGPUMsg;
 use webrender_api::units::{
     DeviceIntSize, DevicePixel, LayoutPixel, LayoutPoint, LayoutSize, WorldPoint,
 };
@@ -401,6 +402,8 @@ pub enum ConstellationControlMsg {
     PaintMetric(PipelineId, ProgressiveWebMetricType, u64),
     /// Notifies the media session about a user requested media session action.
     MediaSessionAction(PipelineId, MediaSessionActionType),
+    /// Notifies script thread that WebGPU server has started
+    SetWebGPUPort(IpcReceiver<WebGPUMsg>),
 }
 
 impl fmt::Debug for ConstellationControlMsg {
@@ -438,6 +441,7 @@ impl fmt::Debug for ConstellationControlMsg {
             PaintMetric(..) => "PaintMetric",
             ExitFullScreen(..) => "ExitFullScreen",
             MediaSessionAction(..) => "MediaSessionAction",
+            SetWebGPUPort(..) => "SetWebGPUPort",
         };
         write!(formatter, "ConstellationControlMsg::{}", variant)
     }

components/webgpu/lib.rs

@@ -8,7 +8,7 @@ extern crate log;
 pub extern crate wgpu_core as wgpu;
 pub extern crate wgpu_types as wgt;
 
-mod identity;
+pub mod identity;
 
 use identity::{IdentityRecyclerFactory, WebGPUMsg};
 use ipc_channel::ipc::{self, IpcReceiver, IpcSender};
@@ -145,7 +145,7 @@ pub enum WebGPURequest {
 pub struct WebGPU(pub IpcSender<WebGPURequest>);
 
 impl WebGPU {
-    pub fn new() -> Option<Self> {
+    pub fn new() -> Option<(Self, IpcReceiver<WebGPUMsg>)> {
         if !pref!(dom.webgpu.enabled) {
             return None;
         }
@@ -161,7 +161,7 @@ impl WebGPU {
         };
         let sender_clone = sender.clone();
 
-        let (script_sender, _script_recv) = match ipc::channel() {
+        let (script_sender, script_recv) = match ipc::channel() {
             Ok(sender_and_receiver) => sender_and_receiver,
             Err(e) => {
                 warn!(
@@ -181,7 +181,7 @@ impl WebGPU {
             warn!("Failed to spwan WGPU thread ({})", e);
             return None;
         }
-        Some(WebGPU(sender))
+        Some((WebGPU(sender), script_recv))
     }
 
     pub fn exit(&self, sender: IpcSender<()>) -> Result<(), &'static str> {