cranpose-render-wgpu 0.9.5

#[cfg(not(target_arch = "wasm32"))]
use std::sync::{
    Arc, Mutex, PoisonError,
    atomic::{AtomicBool, Ordering},
    mpsc::{self, Receiver, Sender},
};

/// How long the last handle waits for a lane to finish the pipeline it is
/// building. A process that exits while a thread is inside the driver's
/// pipeline compile can crash in the driver's own teardown (NVIDIA's does,
/// #859), so a dropped compiler lets that compile end; a driver slower than
/// this is left to it rather than holding the dropping thread.
#[cfg(not(target_arch = "wasm32"))]
const LANE_FINISH_WAIT: web_time::Duration = web_time::Duration::from_secs(5);

/// The compiler's threads, one per [`CompileLane`].
#[cfg(not(target_arch = "wasm32"))]
const LANES: usize = 2;

#[cfg(not(target_arch = "wasm32"))]
type Job = Box<dyn FnOnce() + Send + 'static>;

/// What a value crossing to the compiler threads must be: `Send` where
/// those threads exist, nothing on the web, where wgpu handles are single-threaded
/// and the compiler runs nothing.
#[cfg(not(target_arch = "wasm32"))]
pub(crate) trait CompilerSend: Send {}
#[cfg(not(target_arch = "wasm32"))]
impl<T: Send + ?Sized> CompilerSend for T {}
#[cfg(target_arch = "wasm32")]
pub(crate) trait CompilerSend {}
#[cfg(target_arch = "wasm32")]
impl<T: ?Sized> CompilerSend for T {}

/// What a value shared with the compiler threads must be: `Sync` where
/// those threads exist, nothing on the web.
#[cfg(not(target_arch = "wasm32"))]
pub(crate) trait CompilerSync: Sync {}
#[cfg(not(target_arch = "wasm32"))]
impl<T: Sync + ?Sized> CompilerSync for T {}
#[cfg(target_arch = "wasm32")]
pub(crate) trait CompilerSync {}
#[cfg(target_arch = "wasm32")]
impl<T: ?Sized> CompilerSync for T {}

/// Why a pipeline is queued, which decides the thread that builds it.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum CompileLane {
    /// A frame is drawing with a stand-in until this pipeline is ready.
    Demanded,
    /// No frame has asked for this pipeline yet; it is built ahead of its
    /// first use.
    WarmUp,
}

/// Runs pipeline creation on background threads, so a draw finds its
/// pipeline ready instead of paying the driver's compile inside the frame.
/// Each [`CompileLane`] has its own thread and runs its jobs in the order
/// they were queued, so a pipeline a frame is waiting for never queues
/// behind warm-ups, which take seconds each on a slow device's driver. The
/// threads end with the last handle, which skips the jobs they had not
/// started and waits, up to [`LANE_FINISH_WAIT`], for the ones they had.
#[derive(Clone, Default)]
pub(crate) struct PipelineCompiler {
    #[cfg(not(target_arch = "wasm32"))]
    workers: Option<Arc<Workers>>,
}

#[cfg(not(target_arch = "wasm32"))]
struct Workers {
    /// `None` once dropping has closed the lanes.
    demanded: Option<Sender<Job>>,
    warm_up: Option<Sender<Job>>,
    stopped: Arc<AtomicBool>,
    /// Each lane's thread says here that it has ended. Behind a mutex only
    /// so the handle is `Sync`; the last handle reads it without locking.
    finished: Mutex<Receiver<()>>,
}

#[cfg(not(target_arch = "wasm32"))]
impl Drop for Workers {
    fn drop(&mut self) {
        self.stopped.store(true, Ordering::Release);
        self.demanded = None;
        self.warm_up = None;
        let finished = self
            .finished
            .get_mut()
            .unwrap_or_else(PoisonError::into_inner);
        let deadline = web_time::Instant::now() + LANE_FINISH_WAIT;
        for _ in 0..LANES {
            let left = deadline.saturating_duration_since(web_time::Instant::now());
            if finished.recv_timeout(left).is_err() {
                log::warn!(
                    "[gpu-pipeline] a compile outlived its renderer by {LANE_FINISH_WAIT:?}"
                );
                break;
            }
        }
    }
}

#[cfg(not(target_arch = "wasm32"))]
fn spawn_lane(
    name: &str,
    stopped: &Arc<AtomicBool>,
    finished: &Sender<()>,
) -> std::io::Result<Sender<Job>> {
    let (jobs, queued) = mpsc::channel::<Job>();
    let stopped = Arc::clone(stopped);
    let finished = finished.clone();
    std::thread::Builder::new()
        .name(name.into())
        .spawn(move || {
            crate::render::mark_thread_off_frame();
            while let Ok(job) = queued.recv() {
                if stopped.load(Ordering::Acquire) {
                    break;
                }
                job();
            }
            let _ = finished.send(());
        })?;
    Ok(jobs)
}

/// Where a renderer compiles its pipelines.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum PipelineCompilation {
    /// On background threads, a general pipeline standing in until the
    /// specialized one is ready.
    Background,
    /// Where each is first needed, so no draw uses a stand-in.
    Inline,
}

impl PipelineCompiler {
    pub(crate) fn for_compilation(compilation: PipelineCompilation) -> Self {
        match compilation {
            PipelineCompilation::Background => Self::spawn(),
            PipelineCompilation::Inline => Self::inactive(),
        }
    }

    /// A compiler that runs nothing: every resource compiles where it is
    /// first needed.
    pub(crate) fn inactive() -> Self {
        Self::default()
    }

    pub(crate) fn spawn() -> Self {
        #[cfg(not(target_arch = "wasm32"))]
        {
            static BACKGROUND_PIPELINES: crate::debug_toggles::DebugToggle =
                crate::debug_toggles::DebugToggle::new("CRANPOSE_BACKGROUND_PIPELINES");
            if BACKGROUND_PIPELINES.equals("0") {
                return Self::inactive();
            }
            let stopped = Arc::new(AtomicBool::new(false));
            let (finished_tx, finished) = mpsc::channel();
            let lanes =
                spawn_lane("cranpose-pipelines", &stopped, &finished_tx).and_then(|demanded| {
                    spawn_lane("cranpose-warm-up", &stopped, &finished_tx)
                        .map(|warm_up| (demanded, warm_up))
                });
            match lanes {
                Ok((demanded, warm_up)) => Self {
                    workers: Some(Arc::new(Workers {
                        demanded: Some(demanded),
                        warm_up: Some(warm_up),
                        stopped,
                        finished: Mutex::new(finished),
                    })),
                },
                Err(error) => {
                    log::error!("[gpu-pipeline] background compiler failed to spawn: {error}");
                    Self::inactive()
                }
            }
        }
        #[cfg(target_arch = "wasm32")]
        {
            Self::inactive()
        }
    }

    pub(crate) fn is_active(&self) -> bool {
        #[cfg(not(target_arch = "wasm32"))]
        {
            self.workers.is_some()
        }
        #[cfg(target_arch = "wasm32")]
        {
            false
        }
    }

    /// Queues `job` on `lane`'s thread, behind the jobs already waiting
    /// there; an inactive compiler drops it, and the caller compiles at
    /// first use as before.
    pub(crate) fn enqueue(&self, lane: CompileLane, job: impl FnOnce() + CompilerSend + 'static) {
        #[cfg(not(target_arch = "wasm32"))]
        if let Some(workers) = self.workers.as_ref() {
            let jobs = match lane {
                CompileLane::Demanded => workers.demanded.as_ref(),
                CompileLane::WarmUp => workers.warm_up.as_ref(),
            };
            if jobs.is_none_or(|jobs| jobs.send(Box::new(job)).is_err()) {
                log::error!("[gpu-pipeline] background compiler stopped unexpectedly");
            }
        }
        #[cfg(target_arch = "wasm32")]
        drop((lane, job));
    }
}

#[cfg(all(test, not(target_arch = "wasm32")))]
#[path = "tests/pipeline_compiler_tests.rs"]
mod tests;