fix(desktop): reuse window capture items on Windows 10 (#3449)

fluxer_desktop/native/win-game-capture/Cargo.toml

@@ -19,7 +19,9 @@ napi-derive = "3.6.7"
 windows = {version = "0.62.2", features = [
     "Win32_Foundation",
     "Win32_System_LibraryLoader",
+    "Win32_System_Memory",
     "Win32_System_Threading",
+    "Win32_System_WinRT_Graphics_Capture",
 ]}
 
 [build-dependencies]

fluxer_desktop/native/win-game-capture/index.d.ts

@@ -2,4 +2,6 @@
 
 export declare function elevateGpuSchedulingPriority(processId?: number, priorityClass?: 'high' | 'realtime'): boolean;
 export declare function restoreGpuSchedulingPriority(processId?: number): boolean;
+export declare function reuseWindowCaptureItems(): boolean;
+export declare function windowCaptureItemCalls(): {forwarded: number; reused: number} | null;
 export declare const loadError: Error | null;

fluxer_desktop/native/win-game-capture/index.js

@@ -94,9 +94,27 @@ function restoreGpuSchedulingPriority(processId) {
 	}
 }
 
+function reuseWindowCaptureItems() {
+	if (!binding) return false;
+	try {
+		binding.reuseWindowCaptureItems();
+		return true;
+	} catch (error) {
+		console.warn('[win-game-capture] reuseWindowCaptureItems failed:', error?.message || error);
+		return false;
+	}
+}
+
+function windowCaptureItemCalls() {
+	if (!binding) return null;
+	return binding.windowCaptureItemCalls();
+}
+
 module.exports = {
 	elevateGpuSchedulingPriority,
 	restoreGpuSchedulingPriority,
+	reuseWindowCaptureItems,
+	windowCaptureItemCalls,
 	get loadError() {
 		return loadError;
 	},

fluxer_desktop/native/win-game-capture/package.json

@@ -1,7 +1,7 @@
 {
 	"name": "@fluxer/win-game-capture",
 	"version": "0.0.0",
-	"description": "Native Windows GPU scheduling priority helpers",
+	"description": "Native Windows GPU scheduling priority and window capture item helpers",
 	"private": true,
 	"license": "AGPL-3.0-or-later",
 	"os": [

fluxer_desktop/native/win-game-capture/src/capture_item_reuse.rs

@@ -0,0 +1,204 @@
+// SPDX-License-Identifier: AGPL-3.0-or-later
+
+use crate::live_items::{LiveItems, WeakItem};
+use std::{
+    cell::RefCell,
+    ffi::c_void,
+    mem::{offset_of, size_of},
+    sync::{
+        OnceLock,
+        atomic::{AtomicPtr, AtomicU64, Ordering},
+    },
+};
+use windows::{
+    Win32::{
+        Foundation::{HMODULE, HWND},
+        System::{
+            LibraryLoader::{
+                GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS, GET_MODULE_HANDLE_EX_FLAG_PIN,
+                GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT, GetModuleHandleExW,
+            },
+            Memory::{PAGE_PROTECTION_FLAGS, PAGE_READWRITE, VirtualProtect},
+            WinRT::{
+                Graphics::Capture::{
+                    IGraphicsCaptureItemInterop, IGraphicsCaptureItemInterop_Vtbl,
+                },
+                RoGetActivationFactory,
+            },
+        },
+    },
+    core::{GUID, HRESULT, IUnknown, Interface, PCWSTR, Weak, h, w},
+};
+
+type CreateForWindow =
+    unsafe extern "system" fn(*mut c_void, HWND, *const GUID, *mut *mut c_void) -> HRESULT;
+
+const _: () = assert!(
+    offset_of!(IGraphicsCaptureItemInterop_Vtbl, CreateForWindow) == 3 * size_of::<usize>()
+);
+const _: () = assert!(size_of::<CreateForWindow>() == size_of::<*mut c_void>());
+
+static INSTALLED: OnceLock<Result<(), String>> = OnceLock::new();
+static CREATE_FOR_WINDOW_ORIGINAL: OnceLock<CreateForWindow> = OnceLock::new();
+static CALLS_FORWARDED: AtomicU64 = AtomicU64::new(0);
+static CALLS_REUSED: AtomicU64 = AtomicU64::new(0);
+
+pub struct CreateForWindowCalls {
+    pub forwarded: u64,
+    pub reused: u64,
+}
+
+thread_local! {
+    static LIVE_ITEMS: RefCell<LiveItems<Weak<IUnknown>>> =
+        const { RefCell::new(LiveItems::new()) };
+}
+
+impl WeakItem for Weak<IUnknown> {
+    type Strong = IUnknown;
+
+    fn upgrade(&self) -> Option<IUnknown> {
+        Weak::upgrade(self)
+    }
+}
+
+fn live_item(window: isize) -> Option<IUnknown> {
+    LIVE_ITEMS
+        .try_with(|items| items.try_borrow_mut().ok()?.lookup(window))
+        .ok()
+        .flatten()
+}
+
+fn remember_item(window: isize, item: &IUnknown) {
+    let Ok(weak) = item.downgrade() else {
+        return;
+    };
+    let _ = LIVE_ITEMS.try_with(|items| {
+        if let Ok(mut items) = items.try_borrow_mut() {
+            items.remember(window, weak);
+        }
+    });
+}
+
+unsafe fn reuse_live_item(
+    window: isize,
+    iid: *const GUID,
+    result: *mut *mut c_void,
+) -> Option<HRESULT> {
+    if iid.is_null() {
+        return None;
+    }
+    if result.is_null() {
+        return None;
+    }
+    let item = live_item(window)?;
+    let reused = unsafe { item.query(iid, result) };
+    reused.is_ok().then_some(reused)
+}
+
+unsafe extern "system" fn create_for_window_reusing_live_item(
+    this: *mut c_void,
+    window: HWND,
+    iid: *const GUID,
+    result: *mut *mut c_void,
+) -> HRESULT {
+    let Some(original) = CREATE_FOR_WINDOW_ORIGINAL.get().copied() else {
+        unreachable!("the slot is replaced only after the original is recorded");
+    };
+    let window_key = window.0 as isize;
+    if let Some(reused) = unsafe { reuse_live_item(window_key, iid, result) } {
+        CALLS_REUSED.fetch_add(1, Ordering::Relaxed);
+        return reused;
+    }
+    CALLS_FORWARDED.fetch_add(1, Ordering::Relaxed);
+    let created = unsafe { original(this, window, iid, result) };
+    if created.is_err() {
+        return created;
+    }
+    if result.is_null() {
+        return created;
+    }
+    if let Some(item) = unsafe { IUnknown::from_raw_borrowed(&*result) } {
+        remember_item(window_key, item);
+    }
+    created
+}
+
+fn module_at(address: *const c_void, flags: u32) -> Result<HMODULE, String> {
+    let mut module = HMODULE::default();
+    unsafe {
+        GetModuleHandleExW(
+            GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS | flags,
+            PCWSTR(address.cast()),
+            &mut module,
+        )
+    }
+    .map_err(|error| format!("GetModuleHandleExW(address): {error}"))?;
+    Ok(module)
+}
+
+fn pin_modules(slot: *const c_void) -> Result<(), String> {
+    let mut graphics_capture = HMODULE::default();
+    unsafe {
+        GetModuleHandleExW(
+            GET_MODULE_HANDLE_EX_FLAG_PIN,
+            w!("GraphicsCapture.dll"),
+            &mut graphics_capture,
+        )
+    }
+    .map_err(|error| format!("GetModuleHandleExW(GraphicsCapture.dll): {error}"))?;
+    if module_at(slot, GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT)? != graphics_capture {
+        return Err("capture item factory vtable is outside GraphicsCapture.dll".into());
+    }
+    let replacement: CreateForWindow = create_for_window_reusing_live_item;
+    module_at(replacement as *const c_void, GET_MODULE_HANDLE_EX_FLAG_PIN)?;
+    Ok(())
+}
+
+fn replace_slot(slot: *mut CreateForWindow, original: CreateForWindow) -> Result<(), String> {
+    let replacement: CreateForWindow = create_for_window_reusing_live_item;
+    let size = size_of::<CreateForWindow>();
+    let mut protection = PAGE_PROTECTION_FLAGS(0);
+    unsafe { VirtualProtect(slot.cast(), size, PAGE_READWRITE, &mut protection) }
+        .map_err(|error| format!("VirtualProtect(PAGE_READWRITE): {error}"))?;
+    let slot_value = unsafe { AtomicPtr::<c_void>::from_ptr(slot.cast()) };
+    slot_value.store(replacement as *mut c_void, Ordering::Release);
+    let mut replaced = PAGE_PROTECTION_FLAGS(0);
+    let restored = unsafe { VirtualProtect(slot.cast(), size, protection, &mut replaced) };
+    if let Err(error) = restored {
+        slot_value.store(original as *mut c_void, Ordering::Release);
+        return Err(format!("VirtualProtect(restore): {error}"));
+    }
+    Ok(())
+}
+
+fn install_once() -> Result<(), String> {
+    let interop: IGraphicsCaptureItemInterop =
+        unsafe { RoGetActivationFactory(h!("Windows.Graphics.Capture.GraphicsCaptureItem")) }
+            .map_err(|error| format!("GraphicsCaptureItem activation factory: {error}"))?;
+    let vtable = unsafe {
+        *interop
+            .as_raw()
+            .cast::<*mut IGraphicsCaptureItemInterop_Vtbl>()
+    };
+    assert!(!vtable.is_null());
+    let slot = unsafe { &raw mut (*vtable).CreateForWindow };
+    pin_modules(slot.cast())?;
+    let original = unsafe { slot.read() };
+    let replacement: CreateForWindow = create_for_window_reusing_live_item;
+    assert!(!std::ptr::fn_addr_eq(original, replacement));
+    CREATE_FOR_WINDOW_ORIGINAL
+        .set(original)
+        .map_err(|_| "CreateForWindow original already recorded".to_string())?;
+    replace_slot(slot, original)
+}
+
+pub fn install() -> Result<(), String> {
+    INSTALLED.get_or_init(install_once).clone()
+}
+
+pub fn create_for_window_calls() -> CreateForWindowCalls {
+    CreateForWindowCalls {
+        forwarded: CALLS_FORWARDED.load(Ordering::Relaxed),
+        reused: CALLS_REUSED.load(Ordering::Relaxed),
+    }
+}

fluxer_desktop/native/win-game-capture/src/lib.rs

@@ -2,7 +2,11 @@
 
 #![deny(clippy::all)]
 
+#[cfg(target_os = "windows")]
+mod capture_item_reuse;
 mod gpu_priority;
+#[cfg(any(target_os = "windows", test))]
+mod live_items;
 
 use napi::bindgen_prelude::*;
 use napi_derive::napi;
@@ -19,3 +23,33 @@ pub fn elevate_gpu_scheduling_priority(
 pub fn restore_gpu_scheduling_priority(process_id: Option<u32>) -> Result<()> {
     gpu_priority::restore(process_id).map_err(napi::Error::from_reason)
 }
+
+#[napi(object)]
+pub struct WindowCaptureItemCalls {
+    pub forwarded: i64,
+    pub reused: i64,
+}
+
+#[napi(js_name = "reuseWindowCaptureItems")]
+pub fn reuse_window_capture_items() -> Result<()> {
+    #[cfg(target_os = "windows")]
+    capture_item_reuse::install().map_err(napi::Error::from_reason)?;
+    Ok(())
+}
+
+#[napi(js_name = "windowCaptureItemCalls")]
+pub fn window_capture_item_calls() -> WindowCaptureItemCalls {
+    #[cfg(target_os = "windows")]
+    {
+        let calls = capture_item_reuse::create_for_window_calls();
+        WindowCaptureItemCalls {
+            forwarded: i64::try_from(calls.forwarded).unwrap_or(i64::MAX),
+            reused: i64::try_from(calls.reused).unwrap_or(i64::MAX),
+        }
+    }
+    #[cfg(not(target_os = "windows"))]
+    WindowCaptureItemCalls {
+        forwarded: 0,
+        reused: 0,
+    }
+}

fluxer_desktop/native/win-game-capture/src/live_items.rs

@@ -0,0 +1,305 @@
+// SPDX-License-Identifier: AGPL-3.0-or-later
+
+pub const LIVE_ITEMS_MAX: usize = 8;
+
+pub trait WeakItem {
+    type Strong;
+
+    fn upgrade(&self) -> Option<Self::Strong>;
+}
+
+struct LiveItem<W> {
+    window: isize,
+    weak: W,
+    used_at: u64,
+}
+
+pub struct LiveItems<W> {
+    slots: [Option<LiveItem<W>>; LIVE_ITEMS_MAX],
+    uses: u64,
+}
+
+impl<W: WeakItem> LiveItems<W> {
+    pub const fn new() -> Self {
+        Self {
+            slots: [const { None }; LIVE_ITEMS_MAX],
+            uses: 0,
+        }
+    }
+
+    pub fn lookup(&mut self, window: isize) -> Option<W::Strong> {
+        let used_at = self.next_use();
+        for slot in &mut self.slots {
+            let Some(item) = slot else {
+                continue;
+            };
+            if item.window != window {
+                continue;
+            }
+            match item.weak.upgrade() {
+                Some(strong) => {
+                    item.used_at = used_at;
+                    return Some(strong);
+                }
+                None => *slot = None,
+            }
+        }
+        None
+    }
+
+    pub fn remember(&mut self, window: isize, weak: W) {
+        let used_at = self.next_use();
+        let index = self.vacant_index();
+        assert!(index < LIVE_ITEMS_MAX);
+        self.slots[index] = Some(LiveItem {
+            window,
+            weak,
+            used_at,
+        });
+    }
+
+    fn next_use(&mut self) -> u64 {
+        assert!(self.uses < u64::MAX);
+        self.uses += 1;
+        self.uses
+    }
+
+    fn vacant_index(&self) -> usize {
+        let mut least_recently_used = 0;
+        let mut least_recent_use = u64::MAX;
+        for (index, slot) in self.slots.iter().enumerate() {
+            let Some(item) = slot else {
+                return index;
+            };
+            if item.weak.upgrade().is_none() {
+                return index;
+            }
+            if item.used_at < least_recent_use {
+                least_recent_use = item.used_at;
+                least_recently_used = index;
+            }
+        }
+        least_recently_used
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::rc::{Rc, Weak};
+
+    const FRAMES_IN_HALF_AN_HOUR_AT_30_FPS: u32 = 30 * 60 * 30;
+
+    impl WeakItem for Weak<isize> {
+        type Strong = Rc<isize>;
+
+        fn upgrade(&self) -> Option<Rc<isize>> {
+            Weak::upgrade(self)
+        }
+    }
+
+    struct CaptureService {
+        items_created: u32,
+        live: LiveItems<Weak<isize>>,
+    }
+
+    impl CaptureService {
+        fn new() -> Self {
+            Self {
+                items_created: 0,
+                live: LiveItems::new(),
+            }
+        }
+
+        fn create_for_window(&mut self, window: isize) -> Rc<isize> {
+            if let Some(item) = self.live.lookup(window) {
+                return item;
+            }
+            self.items_created += 1;
+            let item = Rc::new(window);
+            self.live.remember(window, Rc::downgrade(&item));
+            item
+        }
+
+        fn occupied_count(&self) -> usize {
+            self.live.slots.iter().flatten().count()
+        }
+    }
+
+    #[test]
+    fn per_frame_probes_reuse_the_item_the_share_holds() {
+        let mut service = CaptureService::new();
+        drop(service.create_for_window(0x1234));
+        let share_item = service.create_for_window(0x1234);
+        assert_eq!(service.items_created, 2);
+        for _ in 0..FRAMES_IN_HALF_AN_HOUR_AT_30_FPS {
+            let probe = service.create_for_window(0x1234);
+            assert!(Rc::ptr_eq(&probe, &share_item));
+        }
+        assert_eq!(service.items_created, 2);
+        assert_eq!(Rc::strong_count(&share_item), 1);
+    }
+
+    #[test]
+    fn probes_without_a_held_item_each_create_one() {
+        let mut service = CaptureService::new();
+        for _ in 0..100 {
+            drop(service.create_for_window(7));
+        }
+        assert_eq!(service.items_created, 100);
+        assert_eq!(service.occupied_count(), 1);
+    }
+
+    #[test]
+    fn a_new_share_of_the_same_window_creates_a_new_item() {
+        let mut service = CaptureService::new();
+        drop(service.create_for_window(9));
+        let second_share_item = service.create_for_window(9);
+        assert_eq!(service.items_created, 2);
+        drop(second_share_item);
+        assert!(service.live.lookup(9).is_none());
+        assert_eq!(service.occupied_count(), 0);
+    }
+
+    #[test]
+    fn windows_never_share_an_item() {
+        let mut service = CaptureService::new();
+        let first = service.create_for_window(1);
+        let second = service.create_for_window(2);
+        assert!(!Rc::ptr_eq(&first, &second));
+        assert!(Rc::ptr_eq(&service.create_for_window(1), &first));
+        assert!(Rc::ptr_eq(&service.create_for_window(2), &second));
+        assert_eq!(service.items_created, 2);
+    }
+
+    #[test]
+    fn a_dead_slot_is_reused_before_a_live_one_is_evicted() {
+        let mut service = CaptureService::new();
+        let held: Vec<_> = (0..LIVE_ITEMS_MAX as isize)
+            .map(|window| service.create_for_window(window))
+            .collect();
+        drop(service.create_for_window(100));
+        let newest = service.create_for_window(101);
+        assert_eq!(service.occupied_count(), LIVE_ITEMS_MAX);
+        assert!(Rc::ptr_eq(&service.create_for_window(101), &newest));
+        for (window, item) in held.iter().enumerate().skip(1) {
+            assert!(Rc::ptr_eq(
+                &service.create_for_window(window as isize),
+                item
+            ));
+        }
+    }
+
+    #[test]
+    fn a_window_probed_every_frame_outlives_more_held_items_than_slots() {
+        let mut service = CaptureService::new();
+        drop(service.create_for_window(1));
+        let editor_item = service.create_for_window(1);
+        let mut slide_show_items = Vec::new();
+        for slide_show in 0..(LIVE_ITEMS_MAX as isize * 3) {
+            let slide_show_window = 100 + slide_show;
+            drop(service.create_for_window(slide_show_window));
+            let slide_show_item = service.create_for_window(slide_show_window);
+            for _ in 0..30 {
+                let probe = service.create_for_window(slide_show_window);
+                assert!(Rc::ptr_eq(&probe, &slide_show_item));
+            }
+            slide_show_items.push(slide_show_item);
+            let items_created = service.items_created;
+            for _ in 0..30 {
+                let probe = service.create_for_window(1);
+                assert!(Rc::ptr_eq(&probe, &editor_item));
+            }
+            assert_eq!(service.items_created, items_created);
+        }
+        assert_eq!(service.items_created, 2 + 2 * LIVE_ITEMS_MAX as u32 * 3);
+        assert_eq!(service.occupied_count(), LIVE_ITEMS_MAX);
+    }
+
+    #[test]
+    fn the_least_recently_used_held_item_is_the_one_evicted() {
+        for unused in 0..LIVE_ITEMS_MAX as isize {
+            let mut service = CaptureService::new();
+            let held: Vec<_> = (0..LIVE_ITEMS_MAX as isize)
+                .map(|window| service.create_for_window(window))
+                .collect();
+            for window in (0..LIVE_ITEMS_MAX as isize).filter(|window| *window != unused) {
+                drop(service.create_for_window(window));
+            }
+            let newest = service.create_for_window(100);
+            assert!(service.live.lookup(unused).is_none());
+            assert!(Rc::ptr_eq(&service.create_for_window(100), &newest));
+            for window in (0..LIVE_ITEMS_MAX as isize).filter(|window| *window != unused) {
+                assert!(Rc::ptr_eq(
+                    &service.create_for_window(window),
+                    &held[window as usize]
+                ));
+            }
+        }
+    }
+
+    #[test]
+    fn a_newly_held_item_is_evicted_after_the_older_ones() {
+        let mut service = CaptureService::new();
+        let mut held: Vec<_> = (0..LIVE_ITEMS_MAX as isize - 1)
+            .map(|window| service.create_for_window(window))
+            .collect();
+        for window in 0..LIVE_ITEMS_MAX as isize - 1 {
+            drop(service.create_for_window(window));
+        }
+        let newly_held_window = LIVE_ITEMS_MAX as isize - 1;
+        held.push(service.create_for_window(newly_held_window));
+        let over_the_bound = service.create_for_window(100);
+        assert!(service.live.lookup(0).is_none());
+        assert!(Rc::ptr_eq(
+            &service.create_for_window(newly_held_window),
+            &held[newly_held_window as usize]
+        ));
+        assert!(Rc::ptr_eq(&service.create_for_window(100), &over_the_bound));
+    }
+
+    #[test]
+    fn random_share_sequences_match_a_reference_model() {
+        for seed in 1..=32u64 {
+            let mut state = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15);
+            let mut service = CaptureService::new();
+            let mut held: [Option<Rc<isize>>; LIVE_ITEMS_MAX] = [const { None }; LIVE_ITEMS_MAX];
+            let mut items_expected = 0;
+            for _ in 0..5_000 {
+                state ^= state << 13;
+                state ^= state >> 7;
+                state ^= state << 17;
+                let window = (state % LIVE_ITEMS_MAX as u64) as usize;
+                let action = (state >> 32) % 8;
+                if action == 0 {
+                    held[window] = None;
+                } else if let Some(share_item) = &held[window] {
+                    let probe = service.create_for_window(window as isize);
+                    assert!(Rc::ptr_eq(&probe, share_item), "seed {seed}");
+                } else {
+                    items_expected += 1;
+                    let item = service.create_for_window(window as isize);
+                    assert_eq!(*item, window as isize, "seed {seed}");
+                    if action == 1 {
+                        held[window] = Some(item);
+                    }
+                }
+                assert_eq!(service.items_created, items_expected, "seed {seed}");
+                for share_item in held.iter().flatten() {
+                    assert_eq!(Rc::strong_count(share_item), 1, "seed {seed}");
+                }
+            }
+        }
+    }
+
+    #[test]
+    fn the_table_never_grows_past_its_bound() {
+        let mut service = CaptureService::new();
+        let held: Vec<_> = (0..(LIVE_ITEMS_MAX as isize * 4))
+            .map(|window| service.create_for_window(window))
+            .collect();
+        assert_eq!(service.occupied_count(), LIVE_ITEMS_MAX);
+        drop(held);
+        assert!(service.live.lookup(0).is_none());
+    }
+}