Source code
Revision control
Copy as Markdown
Other Tools
// 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
//! Spreads one JXL decode's work across a thread pool dedicated to JXL decoding.
//!
//! # Why a pool of our own
//!
//! The helpers run on a different pool from the decode itself, which is what
//! makes waiting here safe: the thread that blocks and the threads it waits on
//! belong to different pools, so a blocked decode can never be holding a thread
//! that a helper needs. An AVIF decode does the same thing, fanning out to
//! dav1d's own threads.
//!
//! The pool is a `SharedThreadPool`, so it costs nothing until the first
//! parallel decode in the process, it is shut down for us at
//! xpcom-shutdown-threads, and its threads are reaped once they go idle.
//!
//! # Why the closure's lifetime is erased
//!
//! The closure jxl-rs hands us borrows its decode state, so it is only valid for
//! the duration of the call. A helper is an `nsIRunnable`, which has to be
//! `'static`. Both cannot be true, so the lifetime is transmuted away and the
//! closure travels behind a raw pointer.
//!
//! What puts the guarantee back is that the call does not return until every
//! helper it queued has been accounted for. `Helper` decrements the outstanding
//! count from its `Drop`, which runs whether the helper ran or was thrown away
//! without running, and `WaitForHelpers` blocks until that count is zero.
//!
//! `WaitForHelpers` does its waiting from a destructor rather than from a
//! statement after the work, so it also covers the calling thread's own
//! invocation of the closure panicking, which would otherwise unwind past the
//! wait and leave helpers dereferencing a closure jxl-rs has taken back. Gecko
//! builds Rust with `panic=abort` so that unwind does not arise in Firefox, but
//! rusttests builds with unwinding.
//!
//! Helpers are dispatched with `DISPATCH_FALLIBLE`, which asks the pool to
//! release a helper it cannot queue instead of leaking it, so a failed dispatch
//! runs the `Drop` as well.
//!
//! # What this asks of the pool
//!
//! This depends on a queued helper eventually either running or being released,
//! since the wait is for its payload's destructor. `nsThreadPool` has that
//! property: its threads drain the queue to empty before exiting, even once it
//! is shutting down, and `Shutdown` joins them.
use std::sync::{Arc, Condvar, Mutex};
use jxl::api::{JxlParallelRunner, JxlParallelRunnerFun};
use jxl::error::{Error, Result};
use xpcom::interfaces::nsIThreadPool;
use xpcom::RefPtr;
extern "C" {
/// Gets the JXL decode pool, creating it on the first call.
///
/// `thread_limit` only has an effect on the call that creates the pool.
///
/// Null late in shutdown, once there is no longer a pool to be had. Earlier
/// than that it can hand back a pool that refuses work, which shows up as a
/// dispatch failure instead.
fn JxlGetDecodePool(thread_limit: u32) -> *const nsIThreadPool;
}
/// How many threads may work on one decode, counting the calling thread. One
/// participant means no helpers, which is why 1 turns parallel decoding off.
///
/// At 0 this is the machine's number of cpus clamped to between 2 and 8. The
/// lower bound so a single-core machine still gets a helper; the upper because
/// there are diminishing returns.
///
/// Any other value is taken literally.
fn participant_count() -> usize {
match static_prefs::pref!("image.jxl.decode_participants") {
0 => std::thread::available_parallelism()
.map_or(2, std::num::NonZero::get)
.clamp(2, 8),
n => n as usize,
}
}
/// The calling thread does work and so is included in the count so ask for 1
/// fewer threads from the pool to account for that.
fn pool_thread_limit() -> u32 {
participant_count().saturating_sub(1) as u32
}
fn acquire_pool() -> Option<RefPtr<nsIThreadPool>> {
// SAFETY: the callee returns one already-addrefed pointer, or null.
unsafe { RefPtr::from_raw_dont_addref(JxlGetDecodePool(pool_thread_limit())) }
}
/// Everything the participants share, behind one mutex, which is never held
/// while the closure runs.
struct State {
/// Next index to hand out, counting up to `num`.
next: usize,
/// Helpers queued but not yet finished. The caller waits for zero.
outstanding: usize,
/// First error reported by any participant. Once set, no further indices are
/// handed out, so the remaining work is abandoned.
err: Option<Error>,
}
struct Shared {
/// The closure, with its real lifetime erased. Valid until this call returns.
/// The caller waiting until `outstanding` reaches zero guarantees no one uses
/// the closure past its real lifetime.
fun: *const JxlParallelRunnerFun<'static>,
/// Number of indices jxl-rs asked for: the closure runs once for each index
/// below it.
num: usize,
state: Mutex<State>,
/// Notified when outstanding becomes 0.
idle: Condvar,
}
// SAFETY: only the raw pointer makes these unsafe. It is dereferenced solely as a
// shared reference, and `JxlParallelRunnerFun` carries a `Sync` bound, so several
// threads holding one at once is sound. What keeps the pointee alive that long is
// argued at the transmute in `run_in_parallel`.
unsafe impl Send for Shared {}
// SAFETY: same reason.
unsafe impl Sync for Shared {}
impl Shared {
/// The next index to run, or `None` once the work is exhausted or a
/// participant has failed.
fn claim(&self) -> Option<usize> {
let mut state = self.state.lock().unwrap();
if state.err.is_some() || state.next >= self.num {
return None;
}
let i = state.next;
state.next += 1;
Some(i)
}
/// Records an error if it is the first error, which also stops any further
/// indices being handed out.
fn fail(&self, e: Error) {
let mut state = self.state.lock().unwrap();
if state.err.is_none() {
state.err = Some(e);
}
}
/// The first error any participant reported. Only meaningful once they have
/// all finished.
fn take_err(&self) -> Option<Error> {
self.state.lock().unwrap().err.take()
}
/// Take and run indices until they are exhausted or a participant failed.
/// Runs on either kind of participant: a helper, or the calling thread.
fn drain(&self) {
while let Some(i) = self.claim() {
// SAFETY: the closure outlives every use of it here. The calling
// thread is inside the call; a helper is counted in `outstanding`
// from before it is queued until its `Drop`, and the call does not
// return while that count is above zero.
let fun = unsafe { &*self.fun };
if let Err(e) = fun(i) {
self.fail(e);
return;
}
}
}
fn increase_count(&self) {
self.state.lock().unwrap().outstanding += 1;
}
fn release_one(&self) {
let mut state = self.state.lock().unwrap();
// An underflow hangs `WaitForHelpers`, so crash instead.
state.outstanding = state
.outstanding
.checked_sub(1)
.expect("a JXL decode helper was accounted for twice");
if state.outstanding == 0 {
self.idle.notify_all();
}
}
}
/// One queued helper. Dropping it is what tells the caller this helper is done
/// with the closure, whether the helper ran or was discarded.
struct Helper(Arc<Shared>);
impl Drop for Helper {
fn drop(&mut self) {
self.0.release_one();
}
}
/// Blocks until every queued helper has been accounted for.
///
/// Waits from a destructor rather than from a call sited after the work, so it
/// happens however the caller leaves: a normal return, or a panic in the closure
/// unwinding past it.
struct WaitForHelpers<'a>(&'a Shared);
impl Drop for WaitForHelpers<'_> {
fn drop(&mut self) {
let mut state = self.0.state.lock().unwrap();
while state.outstanding > 0 {
state = self.0.idle.wait(state).unwrap();
}
}
}
/// Runs `fun(0..num)`, spread over the decode pool plus the calling thread.
/// Blocks until every index has completed, or returns the first error.
fn run_in_parallel(pool: &nsIThreadPool, num: usize, fun: &JxlParallelRunnerFun<'_>) -> Result<()> {
// The whole design rests on the caller not being one of this pool's threads:
// it blocks until the helpers finish, so a caller that was a pool thread could
// be holding the thread a helper needs.
debug_assert!(
!moz_task::is_on_current_thread(pool.coerce()),
"a JXL decode must not run on the JXL decode pool: it waits for helpers there"
);
// Starting with rustc 1.94 an `as` cast cannot extend the lifetime bound of a
// trait object pointer, so this has to be a transmute.
//
// SAFETY: the `'static` is not true; the closure only lives for this call.
// Sound because we do not return while a helper is still counted in
// `outstanding`.
#[allow(clippy::transmute_ptr_to_ptr)]
let fun_static: *const JxlParallelRunnerFun<'static> =
unsafe { std::mem::transmute(std::ptr::from_ref(fun)) };
let shared = Arc::new(Shared {
fun: fun_static,
num,
state: Mutex::new(State {
next: 0,
outstanding: 0,
err: None,
}),
idle: Condvar::new(),
});
// Declared before anything is queued and dropped after the work, so the wait
// covers every helper on every path out of here.
let wait = WaitForHelpers(&shared);
let helpers = participant_count().saturating_sub(1).min(num - 1);
for _ in 0..helpers {
// Increased before the helper exists, helper's decrement will be sound.
shared.increase_count();
let helper = Helper(Arc::clone(&shared));
let queued = moz_task::RunnableBuilder::new("JxlDecodeHelper", move || {
helper.0.drain();
})
.options(moz_task::DispatchOptions::default().fallible(true))
.dispatch(pool.coerce());
if queued.is_err() {
// Dispatching fallibly means the pool released the helper rather than
// leaking it, so its `Drop` has already run.
break;
}
}
// Also a participant, not just a waiter, so the work gets done even if no
// helper was queued.
shared.drain();
// Explicit, because the error has to be read after every helper has finished
// reporting one.
drop(wait);
match shared.take_err() {
Some(e) => Err(e),
None => Ok(()),
}
}
/// Spreads the work jxl-rs hands us across the JXL decode pool.
pub struct PoolRunner {
/// Left as `None` until there is some work big enough to need it.
pool: Option<RefPtr<nsIThreadPool>>,
}
impl PoolRunner {
/// `None` at one participant: with no helpers there is nothing to spread, so
/// let jxl-rs use its own sequential runner instead of ours.
pub fn new() -> Option<Self> {
if participant_count() == 1 {
return None;
}
Some(Self { pool: None })
}
}
impl JxlParallelRunner for PoolRunner {
fn run(&mut self, num: usize, fun: &JxlParallelRunnerFun<'_>) -> Result<()> {
if num == 0 {
return Ok(());
}
// Nothing to spread: stay on this thread and leave the pool alone.
if num == 1 {
return fun(0);
}
if self.pool.is_none() {
self.pool = acquire_pool();
}
match &self.pool {
Some(pool) => run_in_parallel(pool, num, fun),
None => (0..num).try_for_each(fun),
}
}
}
/// Builds the `Option<&mut dyn _>` argument the jxl-rs entry points take, borrowed
/// from an owned runner. One value cannot be reused across calls: `Option<&mut _>`
/// moves when passed, and the implicit reborrow that lets a bare `&mut` be passed
/// again does not reach inside the `Option`.
pub fn as_dyn(runner: &mut Option<PoolRunner>) -> Option<&mut dyn JxlParallelRunner> {
runner.as_mut().map(|r| r as &mut dyn JxlParallelRunner)
}