dispatcher.rs

  1use std::{
  2    sync::atomic::{AtomicBool, Ordering},
  3    thread::{ThreadId, current},
  4    time::{Duration, Instant},
  5};
  6
  7use flume::Sender;
  8use util::ResultExt;
  9use windows::{
 10    System::Threading::{ThreadPool, ThreadPoolTimer, TimerElapsedHandler, WorkItemHandler},
 11    Win32::{
 12        Foundation::{LPARAM, WPARAM},
 13        UI::WindowsAndMessaging::PostMessageW,
 14    },
 15};
 16
 17use crate::{
 18    GLOBAL_THREAD_TIMINGS, HWND, PlatformDispatcher, RunnableVariant, SafeHwnd, THREAD_TIMINGS,
 19    TaskLabel, TaskTiming, ThreadTaskTimings, WM_GPUI_TASK_DISPATCHED_ON_MAIN_THREAD,
 20};
 21
 22pub(crate) struct WindowsDispatcher {
 23    pub(crate) wake_posted: AtomicBool,
 24    main_sender: Sender<RunnableVariant>,
 25    main_thread_id: ThreadId,
 26    pub(crate) platform_window_handle: SafeHwnd,
 27    validation_number: usize,
 28}
 29
 30impl WindowsDispatcher {
 31    pub(crate) fn new(
 32        main_sender: Sender<RunnableVariant>,
 33        platform_window_handle: HWND,
 34        validation_number: usize,
 35    ) -> Self {
 36        let main_thread_id = current().id();
 37        let platform_window_handle = platform_window_handle.into();
 38
 39        WindowsDispatcher {
 40            main_sender,
 41            main_thread_id,
 42            platform_window_handle,
 43            validation_number,
 44            wake_posted: AtomicBool::new(false),
 45        }
 46    }
 47
 48    fn dispatch_on_threadpool(&self, runnable: RunnableVariant) {
 49        let handler = {
 50            let mut task_wrapper = Some(runnable);
 51            WorkItemHandler::new(move |_| {
 52                Self::execute_runnable(task_wrapper.take().unwrap());
 53                Ok(())
 54            })
 55        };
 56        ThreadPool::RunAsync(&handler).log_err();
 57    }
 58
 59    fn dispatch_on_threadpool_after(&self, runnable: RunnableVariant, duration: Duration) {
 60        let handler = {
 61            let mut task_wrapper = Some(runnable);
 62            TimerElapsedHandler::new(move |_| {
 63                Self::execute_runnable(task_wrapper.take().unwrap());
 64                Ok(())
 65            })
 66        };
 67        ThreadPoolTimer::CreateTimer(&handler, duration.into()).log_err();
 68    }
 69
 70    #[inline(always)]
 71    pub(crate) fn execute_runnable(runnable: RunnableVariant) {
 72        let start = Instant::now();
 73
 74        let mut timing = match runnable {
 75            RunnableVariant::Meta(runnable) => {
 76                let location = runnable.metadata().location;
 77                let timing = TaskTiming {
 78                    location,
 79                    start,
 80                    end: None,
 81                };
 82                Self::add_task_timing(timing);
 83
 84                runnable.run();
 85
 86                timing
 87            }
 88            RunnableVariant::Compat(runnable) => {
 89                let timing = TaskTiming {
 90                    location: core::panic::Location::caller(),
 91                    start,
 92                    end: None,
 93                };
 94                Self::add_task_timing(timing);
 95
 96                runnable.run();
 97
 98                timing
 99            }
100        };
101
102        let end = Instant::now();
103        timing.end = Some(end);
104
105        Self::add_task_timing(timing);
106    }
107
108    pub(crate) fn add_task_timing(timing: TaskTiming) {
109        THREAD_TIMINGS.with(|timings| {
110            let mut timings = timings.lock();
111            let timings = &mut timings.timings;
112
113            if let Some(last_timing) = timings.iter_mut().rev().next() {
114                if last_timing.location == timing.location {
115                    last_timing.end = timing.end;
116                    return;
117                }
118            }
119
120            timings.push_back(timing);
121        });
122    }
123}
124
125impl PlatformDispatcher for WindowsDispatcher {
126    fn get_all_timings(&self) -> Vec<ThreadTaskTimings> {
127        let global_thread_timings = GLOBAL_THREAD_TIMINGS.lock();
128        ThreadTaskTimings::convert(&global_thread_timings)
129    }
130
131    fn get_current_thread_timings(&self) -> Vec<crate::TaskTiming> {
132        THREAD_TIMINGS.with(|timings| {
133            let timings = timings.lock();
134            let timings = &timings.timings;
135
136            let mut vec = Vec::with_capacity(timings.len());
137
138            let (s1, s2) = timings.as_slices();
139            vec.extend_from_slice(s1);
140            vec.extend_from_slice(s2);
141            vec
142        })
143    }
144
145    fn is_main_thread(&self) -> bool {
146        current().id() == self.main_thread_id
147    }
148
149    fn dispatch(&self, runnable: RunnableVariant, label: Option<TaskLabel>) {
150        self.dispatch_on_threadpool(runnable);
151        if let Some(label) = label {
152            log::debug!("TaskLabel: {label:?}");
153        }
154    }
155
156    fn dispatch_on_main_thread(&self, runnable: RunnableVariant) {
157        match self.main_sender.send(runnable) {
158            Ok(_) => {
159                if !self.wake_posted.swap(true, Ordering::AcqRel) {
160                    unsafe {
161                        PostMessageW(
162                            Some(self.platform_window_handle.as_raw()),
163                            WM_GPUI_TASK_DISPATCHED_ON_MAIN_THREAD,
164                            WPARAM(self.validation_number),
165                            LPARAM(0),
166                        )
167                        .log_err();
168                    }
169                }
170            }
171            Err(runnable) => {
172                // NOTE: Runnable may wrap a Future that is !Send.
173                //
174                // This is usually safe because we only poll it on the main thread.
175                // However if the send fails, we know that:
176                // 1. main_receiver has been dropped (which implies the app is shutting down)
177                // 2. we are on a background thread.
178                // It is not safe to drop something !Send on the wrong thread, and
179                // the app will exit soon anyway, so we must forget the runnable.
180                std::mem::forget(runnable);
181            }
182        }
183    }
184
185    fn dispatch_after(&self, duration: Duration, runnable: RunnableVariant) {
186        self.dispatch_on_threadpool_after(runnable, duration);
187    }
188}