tenant.openagents/omega

1use std::{
2    ffi::c_void,
3    ptr::NonNull,
4    sync::atomic::{AtomicBool, Ordering},
5    thread::{ThreadId, current},
6    time::Duration,
7};
8
9use anyhow::Context;
10use gpui_util::ResultExt;
11use windows::Win32::{
12    Foundation::{FILETIME, LPARAM, WPARAM},
13    Media::{timeBeginPeriod, timeEndPeriod},
14    System::Threading::{
15        CloseThreadpoolTimer, CreateThreadpoolTimer, GetCurrentThread, PTP_CALLBACK_INSTANCE,
16        PTP_TIMER, SetThreadPriority, SetThreadpoolTimer, THREAD_PRIORITY_TIME_CRITICAL,
17        TP_CALLBACK_ENVIRON_V3, TP_CALLBACK_PRIORITY, TP_CALLBACK_PRIORITY_HIGH,
18        TP_CALLBACK_PRIORITY_LOW, TP_CALLBACK_PRIORITY_NORMAL, TrySubmitThreadpoolCallback,
19    },
20    UI::WindowsAndMessaging::PostMessageW,
21};
22
23use crate::{HWND, SafeHwnd, WM_GPUI_TASK_DISPATCHED_ON_MAIN_THREAD};
24use gpui::{
25    PlatformDispatcher, Priority, PriorityQueueSender, RunnableVariant, TimerResolutionGuard,
26};
27
28pub(crate) struct WindowsDispatcher {
29    pub(crate) wake_posted: AtomicBool,
30    main_sender: PriorityQueueSender<RunnableVariant>,
31    main_thread_id: ThreadId,
32    pub(crate) platform_window_handle: SafeHwnd,
33    validation_number: usize,
34}
35
36impl WindowsDispatcher {
37    pub(crate) fn new(
38        main_sender: PriorityQueueSender<RunnableVariant>,
39        platform_window_handle: HWND,
40        validation_number: usize,
41    ) -> Self {
42        let main_thread_id = current().id();
43        let platform_window_handle = platform_window_handle.into();
44
45        WindowsDispatcher {
46            main_sender,
47            main_thread_id,
48            platform_window_handle,
49            validation_number,
50            wake_posted: AtomicBool::new(false),
51        }
52    }
53
54    fn dispatch_on_threadpool(&self, priority: TP_CALLBACK_PRIORITY, runnable: RunnableVariant) {
55        let environ = TP_CALLBACK_ENVIRON_V3 {
56            Version: 3,
57            CallbackPriority: priority,
58            Size: size_of::<TP_CALLBACK_ENVIRON_V3>() as u32,
59            ..Default::default()
60        };
61
62        // If the thread pool never runs our callback, the matching `from_raw` is never called, which leaks the runnable.
63        // Dropping the scheduled runnable would cancel its task and make the next poll of any awaiter panic. Since we expect
64        // the scenario to usually happen during shutdown, this leak is acceptable.
65        let context = runnable.into_raw().as_ptr() as *mut c_void;
66
67        unsafe {
68            TrySubmitThreadpoolCallback(Some(run_work_callback), Some(context), Some(&environ))
69                .log_err();
70        }
71    }
72
73    fn dispatch_on_threadpool_after(&self, runnable: RunnableVariant, duration: Duration) {
74        let context = runnable.into_raw().as_ptr() as *mut c_void;
75
76        unsafe {
77            if let Ok(timer) = CreateThreadpoolTimer(Some(run_timer_callback), Some(context), None)
78            {
79                // Negative FILETIME expresses a relative delay in 100ns ticks
80                let ticks = (duration.as_nanos() / 100).min(i64::MAX as u128) as i64;
81                let due = (-ticks) as u64;
82                let due_time = FILETIME {
83                    dwLowDateTime: due as u32,
84                    dwHighDateTime: (due >> 32) as u32,
85                };
86                SetThreadpoolTimer(timer, Some(&due_time), 0, None);
87            }
88        }
89    }
90
91    #[inline(always)]
92    pub(crate) fn execute_runnable(runnable: RunnableVariant) {
93        let location = runnable.metadata().location;
94        let spawned = runnable.metadata().spawned;
95        gpui::profiler::update_running_task(spawned, location);
96        runnable.run();
97        gpui::profiler::save_task_timing();
98    }
99}
100
101impl PlatformDispatcher for WindowsDispatcher {
102    fn is_main_thread(&self) -> bool {
103        current().id() == self.main_thread_id
104    }
105
106    fn dispatch(&self, runnable: RunnableVariant, priority: Priority) {
107        let priority = match priority {
108            Priority::RealtimeAudio => {
109                panic!("RealtimeAudio priority should use spawn_realtime, not dispatch")
110            }
111            Priority::High => TP_CALLBACK_PRIORITY_HIGH,
112            Priority::Medium => TP_CALLBACK_PRIORITY_NORMAL,
113            Priority::Low => TP_CALLBACK_PRIORITY_LOW,
114        };
115        self.dispatch_on_threadpool(priority, runnable);
116    }
117
118    fn dispatch_on_main_thread(&self, runnable: RunnableVariant, priority: Priority) {
119        match self.main_sender.send(priority, runnable) {
120            Ok(_) => {
121                if !self.wake_posted.swap(true, Ordering::AcqRel) {
122                    unsafe {
123                        PostMessageW(
124                            Some(self.platform_window_handle.as_raw()),
125                            WM_GPUI_TASK_DISPATCHED_ON_MAIN_THREAD,
126                            WPARAM(self.validation_number),
127                            LPARAM(0),
128                        )
129                        .log_err();
130                    }
131                }
132            }
133            Err(runnable) => {
134                // NOTE: Runnable may wrap a Future that is !Send.
135                //
136                // This is usually safe because we only poll it on the main thread.
137                // However if the send fails, we know that:
138                // 1. main_receiver has been dropped (which implies the app is shutting down)
139                // 2. we are on a background thread.
140                // It is not safe to drop something !Send on the wrong thread, and
141                // the app will exit soon anyway, so we must forget the runnable.
142                std::mem::forget(runnable);
143            }
144        }
145    }
146
147    fn dispatch_after(&self, duration: Duration, runnable: RunnableVariant) {
148        self.dispatch_on_threadpool_after(runnable, duration);
149    }
150
151    fn spawn_realtime(&self, f: Box<dyn FnOnce() + Send>) {
152        std::thread::spawn(move || {
153            // SAFETY: always safe to call
154            let thread_handle = unsafe { GetCurrentThread() };
155
156            // SAFETY: thread_handle is a valid handle to the current thread
157            unsafe { SetThreadPriority(thread_handle, THREAD_PRIORITY_TIME_CRITICAL) }
158                .context("thread priority")
159                .log_err();
160
161            f();
162        });
163    }
164
165    fn increase_timer_resolution(&self) -> TimerResolutionGuard {
166        unsafe {
167            timeBeginPeriod(1);
168        }
169        gpui_util::defer(Box::new(|| unsafe {
170            timeEndPeriod(1);
171        }))
172    }
173}
174
175unsafe extern "system" fn run_work_callback(
176    _instance: PTP_CALLBACK_INSTANCE,
177    context: *mut c_void,
178) {
179    let runnable = unsafe { RunnableVariant::from_raw(NonNull::new_unchecked(context as *mut ())) };
180    WindowsDispatcher::execute_runnable(runnable);
181}
182
183unsafe extern "system" fn run_timer_callback(
184    _instance: PTP_CALLBACK_INSTANCE,
185    context: *mut c_void,
186    timer: PTP_TIMER,
187) {
188    let runnable = unsafe { RunnableVariant::from_raw(NonNull::new_unchecked(context as *mut ())) };
189    WindowsDispatcher::execute_runnable(runnable);
190    unsafe { CloseThreadpoolTimer(timer) };
191}
192