1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144

/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/. */

//! Data and main loop of WGPU poll thread.
//!
//! This is roughly based on <https://github.com/LucentFlux/wgpu-async/blob/1322c7e3fcdfc1865a472c7bbbf0e2e06dcf4da8/src/wgpu_future.rs>

use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex, MutexGuard};
use std::thread::JoinHandle;

use log::warn;

use crate::wgc::global::Global;

/// Polls devices while there is something to poll.
///
/// This objects corresponds to a thread that parks itself when there is no work,
/// waiting on it, and then calls `poll_all_devices` repeatedly to block.
///
/// The thread dies when this object is dropped, and all work in submission is done.
///
/// ## Example
/// ```no_run
/// let token = self.poller.token(); // create a new token
/// let callback = SubmittedWorkDoneClosure::from_rust(Box::from(move || {
///    drop(token); // drop token as closure has been fired
///    // ...
/// }));
/// let result = gfx_select!(queue_id => global.queue_on_submitted_work_done(queue_id, callback));
/// self.poller.wake(); // wake poller thread to actually poll
/// ```
#[derive(Debug)]
pub(crate) struct Poller {
    /// The number of closures that still needs to be fired.
    /// When this is 0, the thread can park itself.
    work_count: Arc<AtomicUsize>,
    /// True if thread should die after all work in submission is done
    is_done: Arc<AtomicBool>,
    /// Handle to the WGPU poller thread (to be used for unparking the thread)
    handle: Option<JoinHandle<()>>,
    /// Lock for device maintain calls (in poll_all_devices and queue_submit)
    ///
    /// This is workaround for wgpu deadlocks: <https://github.com/gfx-rs/wgpu/issues/5572>
    lock: Arc<Mutex<()>>,
}

#[inline]
fn poll_all_devices(
    global: &Arc<Global>,
    more_work: &mut bool,
    force_wait: bool,
    lock: &Mutex<()>,
) {
    let _guard = lock.lock().unwrap();
    match global.poll_all_devices(force_wait) {
        Ok(all_queue_empty) => *more_work = !all_queue_empty,
        Err(e) => warn!("Poller thread got `{e}` on poll_all_devices."),
    }
    // drop guard
}

impl Poller {
    pub(crate) fn new(global: Arc<Global>) -> Self {
        let work_count = Arc::new(AtomicUsize::new(0));
        let is_done = Arc::new(AtomicBool::new(false));
        let work = work_count.clone();
        let done = is_done.clone();
        let lock = Arc::new(Mutex::new(()));
        Self {
            work_count,
            is_done,
            lock: Arc::clone(&lock),
            handle: Some(
                std::thread::Builder::new()
                    .name("WGPU poller".into())
                    .spawn(move || {
                        while !done.load(Ordering::Acquire) {
                            let mut more_work = false;
                            // Do non-blocking poll unconditionally
                            // so every `ẁake` (even spurious) will do at least one poll.
                            // this is mostly useful for stuff that is deferred
                            // to maintain calls in wgpu (device resource destruction)
                            poll_all_devices(&global, &mut more_work, false, &lock);
                            while more_work || work.load(Ordering::Acquire) != 0 {
                                poll_all_devices(&global, &mut more_work, true, &lock);
                            }
                            std::thread::park(); //TODO: should we use timeout here
                        }
                    })
                    .expect("Spawning thread should not fail"),
            ),
        }
    }

    /// Creates a token of work
    pub(crate) fn token(&self) -> WorkToken {
        let prev = self.work_count.fetch_add(1, Ordering::AcqRel);
        debug_assert!(
            prev < usize::MAX,
            "cannot have more than `usize::MAX` outstanding operations on the GPU"
        );
        WorkToken {
            work_count: Arc::clone(&self.work_count),
        }
    }

    /// Wakes the poller thread to start polling.
    pub(crate) fn wake(&self) {
        self.handle
            .as_ref()
            .expect("Poller thread does not exist!")
            .thread()
            .unpark();
    }

    /// Lock for device maintain calls (in poll_all_devices and queue_submit)
    pub(crate) fn lock(&self) -> MutexGuard<()> {
        self.lock.lock().unwrap()
    }
}

impl Drop for Poller {
    fn drop(&mut self) {
        self.is_done.store(true, Ordering::Release);

        let handle = self.handle.take().expect("Poller dropped twice");
        handle.thread().unpark();
        handle.join().expect("Poller thread panicked");
    }
}

/// RAII indicating that there is some work enqueued (closure to be fired),
/// while this token is held.
pub(crate) struct WorkToken {
    work_count: Arc<AtomicUsize>,
}

impl Drop for WorkToken {
    fn drop(&mut self) {
        self.work_count.fetch_sub(1, Ordering::AcqRel);
    }
}