use std::{
collections::VecDeque,
sync::{Arc, Mutex},
};
use crate::{actor_task::ActorTask, doc_actor_inner::DocActorInner};
pub(crate) struct ThreadpoolActor {
inner: Arc<Mutex<DocActorInner>>,
pool: Arc<rayon::ThreadPool>,
mailbox: Mutex<Mailbox>,
}
#[derive(Default)]
struct Mailbox {
tasks: VecDeque<ActorTask>,
running: bool,
}
impl ThreadpoolActor {
pub(crate) fn new(inner: Arc<Mutex<DocActorInner>>, pool: Arc<rayon::ThreadPool>) -> Arc<Self> {
Arc::new(Self {
inner,
pool,
mailbox: Mutex::new(Mailbox::default()),
})
}
pub(crate) fn dispatch(self: &Arc<Self>, task: ActorTask) {
let mut mailbox = self.mailbox.lock().unwrap();
mailbox.tasks.push_back(task);
if mailbox.running {
return;
}
mailbox.running = true;
drop(mailbox);
self.schedule();
}
fn schedule(self: &Arc<Self>) {
let actor = self.clone();
self.pool.spawn_fifo(move || actor.run());
}
fn run(self: Arc<Self>) {
const BATCH_SIZE: usize = 64;
for _ in 0..BATCH_SIZE {
let task = {
let mut mailbox = self.mailbox.lock().unwrap();
let Some(task) = mailbox.tasks.pop_front() else {
mailbox.running = false;
return;
};
task
};
self.inner.lock().unwrap().handle_task(task);
}
self.schedule();
}
}
#[cfg(all(test, feature = "tokio"))]
mod tests {
use std::{sync::mpsc, time::Duration};
use automerge::Automerge;
use futures::{
FutureExt, StreamExt,
channel::{mpsc as async_mpsc, oneshot},
};
use super::*;
use crate::{ConcurrencyConfig, Repo, doc_runner::DocRunner};
async fn repo(threads: usize) -> Repo {
Repo::build_tokio()
.with_concurrency(ConcurrencyConfig::Threadpool(
rayon::ThreadPoolBuilder::new()
.num_threads(threads)
.build()
.unwrap(),
))
.load()
.await
}
#[tokio::test]
async fn hub_and_handle_tasks_share_fifo_order_even_when_submitted_from_a_worker() {
tokio::time::timeout(Duration::from_secs(5), async {
let repo = repo(1).await;
let doc = repo.create(Automerge::new()).await.unwrap();
let (pool, actor_id) = {
let inner = repo.inner.lock().unwrap();
let DocRunner::Threadpool(pool) = &inner.doc_runner else {
unreachable!()
};
(pool.clone(), *inner.actors.keys().next().unwrap())
};
for _ in 0..3 {
let (tx, mut rx) = async_mpsc::unbounded();
let inner = repo.inner.clone();
let doc = doc.clone();
pool.spawn(move || {
for index in 0..200 {
let tx = tx.clone();
if index % 2 == 0 {
assert!(
doc.with_document_async(move |_| {
tx.unbounded_send(index).unwrap();
})
.now_or_never()
.is_none()
);
} else {
inner.lock().unwrap().dispatch_task(
actor_id,
ActorTask::WithDocument(Box::new(move |_| {
tx.unbounded_send(index).unwrap();
})),
);
}
}
});
let mut actual = Vec::new();
while let Some(index) = rx.next().await {
actual.push(index);
}
assert_eq!(actual, (0..200).collect::<Vec<_>>());
}
repo.stop().await;
})
.await
.unwrap();
}
#[tokio::test]
async fn queued_tasks_for_a_busy_actor_do_not_occupy_other_workers() {
tokio::time::timeout(Duration::from_secs(5), async {
let repo = repo(2).await;
let busy = repo.create(Automerge::new()).await.unwrap();
let other = repo.create(Automerge::new()).await.unwrap();
let (started, ready) = oneshot::channel();
let (release, wait) = mpsc::channel();
let mut blocked = Box::pin(busy.with_document_async(move |_| {
let _ = started.send(());
let _ = wait.recv_timeout(Duration::from_secs(5));
}));
assert!(futures::poll!(&mut blocked).is_pending());
ready.await.unwrap();
let mut queued = Vec::new();
for _ in 0..100 {
let mut task = Box::pin(busy.with_document_async(|_| ()));
assert!(futures::poll!(&mut task).is_pending());
queued.push(task);
}
other.with_document_async(|_| ()).await.unwrap();
release.send(()).unwrap();
blocked.await.unwrap();
for task in queued {
task.await.unwrap();
}
repo.stop().await;
})
.await
.unwrap();
}
}