Add ThreadPool::broadcast

[cuviper]

A broadcast runs the closure on every thread in the pool, then collects
the results. It's scheduled somewhat like a very soft interrupt -- it
won't preempt a thread's local work, but will run before it goes to
steal from any other threads.

This can be used when you want to precisely split your work per-thread,
or to set or retrieve some thread-local data in the pool, e.g. #483.

[nikomatsakis]

I like this overall. In terms of what we are committing to publicly, I suppose it's this:

You have the ability to schedule a closure to run exactly once on each thread. The priority is sort of undefined (as is typical for Rayon).

That doesn't seem like too much of a commitment. The main question is whether that behavior is well-defined if the size of the threadpool were to ever become dynamic, but I am more and more dubious of such a thing ever happening, and certainly not without some form of opt-in.

I wonder if it'd be useful to give the closure the number of threads as well? I guess that is readily accessible from the pool, so that's why we don't, right?

[cuviper]

In terms of what we are committing to publicly, I suppose it's this:

You have the ability to schedule a closure to run exactly once on each thread. The priority is sort of undefined (as is typical for Rayon).

That doesn't seem like too much of a commitment.

One thing to strengthen this is that racing broadcasts will be consistently ordered. That is, given simultaneous broadcasts A and B, then if one thread sees A before B, they all will, and vice versa. So for example, if both are setting the same TLS, every thread will set them in the same order.

It doesn't guarantee that every A will complete before any B starts, but you could use a Barrier in A if needed. Consistent order means you won't have to worry about whether a barrier in B is also waiting, interleaved with A to deadlock.

We don't have to promise this, but I think it's powerful if we do.

The main question is whether that behavior is well-defined if the size of the threadpool were to ever become dynamic, but I am more and more dubious of such a thing ever happening, and certainly not without some form of opt-in.

I'm also skeptical of this happening. Internally, we could at least make sure that threads are never removed while they have a broadcast waiting, but I guess new threads would just be left out.

I wonder if it'd be useful to give the closure the number of threads as well? I guess that is readily accessible from the pool, so that's why we don't, right?

Yeah, it's easy to get, although that could be racy if there are dynamic threads. We could supply that value to indicate the number of times the broadcast will be called -- the number we're actually queuing -- independent of any new threads that may pop up.

[nikomatsakis]

@cuviper

Yeah, it's easy to get, although that could be racy if there are dynamic threads. We could supply that value to indicate the number of times the broadcast will be called -- the number we're actually queuing -- independent of any new threads that may pop up.

So actually, I think this is a case where it would make sense for us to supply a &BroadcastContext (or BroadcastContext) so that we can easily add more contextual information later. If we really want to future proof, it'd probably be &BroadcastContext<'_>, where the lifetime is currently just phantom.

[cuviper]

I've rebased and added a context type. Let me know what you think!

[Zoxc]

I'd like to use this in rustc to collect thread local data (complements my WorkerLocal type well).

[cuviper]

@Zoxc I've rebased this again, if you'd like to try it out with rustc.

[Zoxc]

There's 2 variations on this which may be useful too. spawn_broadcast which would spawn tasks on all threads without waiting for them. broadcast on a Scope which would spawn tasks on the Scope and could be considered a more general form of this PR.

I wonder if there's room for some generic code here. spawn_broadcast would wait until 'static ends before blocking (a.k.a. never), Scope::broadcast would wait until 'scope ends before blocking and the broadcast in this PR would wait on some internal lifetime before blocking. If you consider join as spawning 2 tasks, you can draw similar parallels between spawn, Scope::spawn and join.

[cuviper]

Those variations make sense to me. The Scope::broadcast idea should be a pretty easy extension, just adding to the scope counter so they're part of the blocking set. For the static/unblocked version, I guess we could wrap the closure in an Arc and basically just spawn it out.

They could all use the same injection queues, at least.

[cuviper]

I rebased again and added spawn_broadcast, including a scoped version. I need more docs/examples on that, but otherwise I think the code is in good shape here.

The last thing I'm thinking about is a specification of when a broadcast will run. Currently, it runs when the local deque is empty, before looking for jobs elsewhere. I can think of multiple options, from high to low "priority":

• ASAP -- if we've blocked for any reason, run broadcasts first even once the blocked latch is ready.
• After the seeing latch is still blocked, but before popping from the local deque.
• After the local deque, but before stealing from other threads. This is what we have now.
• After thread stealing, but before popping from the global injector.
• After the global injector, when there are no other sources of pending work.
• After the whole pool is otherwise blocked/idle.

I'm not certain that the current choice is the best. These could all be supported if we kept distinct queues, and perhaps an enum argument on the broadcast methods to indicate the user's choice, but maybe that's overkill.

[LoganDark]

These could all be supported if we kept distinct queues, and perhaps an enum argument on the broadcast methods to indicate the user's choice, but maybe that's overkill.

It sounds like something that could be tuned, but I agree that it sounds like overkill because people won't know what to pick to get it to Just Work.

You could support an enum, or not, but the most important thing is that it's as fast as possible both when it's the only thing executing on the thread pool and when the thread pool is juggling lots and lots of work. Basically, even if you do have an enum for user choice, please have (and recommend) an API where the user does not have to choose, so that they can benefit from a tried and tested compromise.

"as fast as possible" may mean many things. Broadcasts completing before other work, or a compromise between timely execution and not ruining the efficiency of other tasks running on the thread pool. I'd personally advocate for the latter as a sane default.

[cuviper]

Here's the current API summary:

pub fn broadcast<OP, R>(op: OP) -> Vec<R>
where
    OP: Fn(BroadcastContext<'_>) -> R + Sync,
    R: Send;

pub fn spawn_broadcast<OP>(op: OP)
where
    OP: Fn(BroadcastContext<'_>) + Send + Sync + 'static;

pub struct BroadcastContext<'a> { .. }

impl<'a> BroadcastContext<'a> {
    pub fn index(&self) -> usize;
    pub fn num_threads(&self) -> usize;
}

impl<'a> fmt::Debug for BroadcastContext<'a> { .. }

impl<'scope> Scope<'scope> {
    pub fn spawn_broadcast<BODY>(&self, body: BODY)
    where
        BODY: Fn(&Scope<'scope>, BroadcastContext<'_>) + Send + Sync + 'scope;
}

impl<'scope> ScopeFifo<'scope> {
    pub fn spawn_broadcast<BODY>(&self, body: BODY)
    where
        BODY: Fn(&ScopeFifo<'scope>, BroadcastContext<'_>) + Send + Sync + 'scope;
}

impl ThreadPool {
    pub fn broadcast<OP, R>(&self, op: OP) -> Vec<R>
    where
        OP: Fn(BroadcastContext<'_>) -> R + Sync,
        R: Send;

    pub fn spawn_broadcast<OP>(&self, op: OP)
    where
        OP: Fn(BroadcastContext<'_>) + Send + Sync + 'static;
}

I think that's pretty safe, and the current "priority" (before remote work-stealing) still feels like a reasonable default.

[cuviper]

bors r+

[bors]

Build succeeded:

• Check (1.56.0)
• Demo (nightly)
• Demo (stable)
• Format
• Test (macos-latest, beta)
• Test (macos-latest, nightly)
• Test (macos-latest, stable)
• Test (ubuntu-latest, beta)
• Test (ubuntu-latest, nightly)
• Test (ubuntu-latest, stable)
• Test (ubuntu-latest, stable-i686)
• Test (windows-latest, beta)
• Test (windows-latest, nightly)
• Test (windows-latest, stable)
• WebAssembly