diff --git a/examples/interrmitent-failure/src/main.rs b/examples/interrmitent-failure/src/main.rs index 31e9eb5..eeb3dba 100644 --- a/examples/interrmitent-failure/src/main.rs +++ b/examples/interrmitent-failure/src/main.rs @@ -1,6 +1,6 @@ use std::time::Duration; -use poolparty::{buffer::RingBuffer, Supervisor, Task, Workable}; +use poolparty::{buffer::RingBuffer, worker::Worker, Supervisor, Task, Workable}; use tokio::sync::mpsc::Sender; use tracing_subscriber::{fmt, prelude::*, EnvFilter}; diff --git a/poolparty/src/lib.rs b/poolparty/src/lib.rs index 7dc639f..0c0070a 100644 --- a/poolparty/src/lib.rs +++ b/poolparty/src/lib.rs @@ -4,7 +4,7 @@ pub mod supervisor; pub mod worker; pub use crate::{ - message::Response, + message::WorkerMessage, supervisor::Supervisor, worker::{Task, Workable}, }; diff --git a/poolparty/src/message.rs b/poolparty/src/message.rs index 02a286f..af4730a 100644 --- a/poolparty/src/message.rs +++ b/poolparty/src/message.rs @@ -1,15 +1,22 @@ use crate::worker::Workable; +// FIXME(jdb): Should really revise the naming. This feels +// half-arsed. #[derive(Debug, Clone)] -pub enum Request { +pub enum SupervisorMessage { State, Task(W::Task), Cancel, + + /// Acknowledge + SubscribeAck, Shutdown, } #[derive(Debug)] -pub enum Response { - ShutdownAck, // Acknowledge shutdown +pub enum WorkerMessage { Complete(Result), + /// Worker subscribing to a running Supervisor. + Subscribe, + ShutdownAck, // Acknowledge shutdown } diff --git a/poolparty/src/supervisor.rs b/poolparty/src/supervisor.rs index e7f7ed1..58e6427 100644 --- a/poolparty/src/supervisor.rs +++ b/poolparty/src/supervisor.rs @@ -1,6 +1,6 @@ use crate::{ buffer::RingBuffer, - message::{Request, Response}, + message::{SupervisorMessage, WorkerMessage}, worker::{Workable, Worker}, Pid, }; @@ -19,12 +19,12 @@ use tokio::{ pub struct Supervisor<'buf, W: Workable> { /// Internal worker pool, containing the queue of workers that are ready /// to receive a task (i.e., checkout). - pool: BTreeMap>, JoinHandle<()>)>, + pool: BTreeMap>, JoinHandle<()>)>, /// An internal pool containing the list of checked out workers. We need /// to do this in order to keep channels alive and keep communication with /// workers even as they are running. - checked: BTreeMap>, JoinHandle<()>)>, + checked: BTreeMap>, JoinHandle<()>)>, /// Pending queue of tasks to be executed tasks: VecDeque, @@ -36,11 +36,11 @@ pub struct Supervisor<'buf, W: Workable> { pub results: &'buf RingBuffer>, /// Sender end, mostly kept here for the time being to simplify spawning. - sender: Sender<(Pid, Response)>, + pub sender: Sender<(Pid, WorkerMessage)>, /// Receiver end of the channel between all workers and the supervisor. This /// allows workers to emit messages back to the supervisor efficiently. - receiver: Receiver<(Pid, Response)>, + receiver: Receiver<(Pid, WorkerMessage)>, } impl<'buf, W: Workable + 'static> Supervisor<'buf, W> { @@ -61,29 +61,9 @@ impl<'buf, W: Workable + 'static> Supervisor<'buf, W> { receiver: supervisor_rx, }; - supervisor.spawn(size); - supervisor } - fn spawn(&mut self, n: usize) { - let id = uuid::Uuid::new_v4(); - - for _ in 0..n { - let supervisor_tx = self.sender.clone(); - let (worker_tx, worker_rx) = mpsc::channel(1024); - - let handle = tokio::spawn(async move { - Worker::new(id, supervisor_tx.clone(), worker_rx) - .run() - .await; - }); - - tracing::info!("spawning worker with id {id}"); - self.pool.insert(id, (worker_tx, handle)); - } - } - #[tracing::instrument(skip_all)] pub async fn run(&mut self) { // Start running the supervisor manage worker lifecycle @@ -126,7 +106,7 @@ impl<'buf, W: Workable + 'static> Supervisor<'buf, W> { // possible from the task queue to our worker pool. Currently // we are only allocating one task at a time. if let (Some(worker), Some(task)) = (self.pool.pop_first(), self.tasks.pop_front()) { - let msg = Request::Task(task); + let msg = SupervisorMessage::Task(task); if worker.1.0.send(msg).await.is_ok() { // Move the worker to the checked out pool so that the channel @@ -138,7 +118,7 @@ impl<'buf, W: Workable + 'static> Supervisor<'buf, W> { // Received a message from one of our workers msg = self.receiver.recv() => { match msg { - Some((pid, Response::Complete(result))) => { + Some((pid, WorkerMessage::Complete(result))) => { tracing::debug!("received msg from worker: {result:?}"); match &result { @@ -151,12 +131,8 @@ impl<'buf, W: Workable + 'static> Supervisor<'buf, W> { Err((err, task)) => { tracing::info!("received error from worker {pid}, err: {err:?}, retrying task: {task:?}"); - // Remove the worker from the checked pool and kill/drop it. - // - // We should then spawn a new worker in the pool to be able to - // take it's place. + // Remove the worker from the checked pool and drop it. if let Some(worker) = self.checked.remove_entry(&pid) { - self.spawn(1); drop(worker); } @@ -176,6 +152,9 @@ impl<'buf, W: Workable + 'static> Supervisor<'buf, W> { } }, + Some((pid, WorkerMessage::Subscribe)) => { + tracing::info!("worker {pid} attempting to subscribe"); + } Some(res) => { tracing::debug!("received res from worker: {res:?}"); } @@ -200,14 +179,15 @@ impl<'buf, W: Workable + 'static> Supervisor<'buf, W> { // workers to ack or timeout. tracing::info!("shutting down supervisor"); - let mut workers: BTreeMap>, JoinHandle<()>)> = BTreeMap::new(); + let mut workers: BTreeMap>, JoinHandle<()>)> = + BTreeMap::new(); workers.append(&mut self.pool); workers.append(&mut self.checked); for worker in workers.iter() { let sender = &worker.1 .0; // FIXME(jdb): terrible, but lazy sender - .send(Request::Shutdown) + .send(SupervisorMessage::Shutdown) .await .expect("unable to send shutdown to worker {worker}"); } @@ -223,7 +203,7 @@ impl<'buf, W: Workable + 'static> Supervisor<'buf, W> { tokio::select! { msg = self.receiver.recv() => { - if let Some((worker_id, Response::ShutdownAck)) = msg { + if let Some((worker_id, WorkerMessage::ShutdownAck)) = msg { tracing::debug!("got shutdown ack from {worker_id}"); workers.remove(&worker_id); } diff --git a/poolparty/src/worker.rs b/poolparty/src/worker.rs index 0883230..69eb568 100644 --- a/poolparty/src/worker.rs +++ b/poolparty/src/worker.rs @@ -1,7 +1,7 @@ use std::{fmt::Debug, future::Future}; use crate::{ - message::{Request, Response}, + message::{SupervisorMessage, WorkerMessage}, Pid, }; use tokio::sync::mpsc::{Receiver, Sender}; @@ -26,9 +26,9 @@ pub trait Workable: Debug + Send + Sync + Sized { #[derive(Debug)] pub struct Worker { - id: Pid, - tx: Sender<(Pid, Response)>, - rx: Receiver>, + pid: Pid, + tx: Sender<(Pid, WorkerMessage)>, + rx: Receiver>, state: State, } @@ -43,16 +43,32 @@ impl Drop for Worker { } impl Worker { - pub fn new(id: Pid, tx: Sender<(Pid, Response)>, rx: Receiver>) -> Self { + pub async fn new( + tx: Sender<(Pid, WorkerMessage)>, + rx: Receiver>, + ) -> Self { + let pid = uuid::Uuid::new_v4(); + + // NOTE(jdb): For now we are bootstrapping the subscription to the + // supervisor. In the future we would realistically want to: + // + // 1. Allow the worker to subscribe at a later stage when the client + // decides the worker is ready to `run`; + // 2. We would want to chang the subscription to an rpc call. For the + // time being all communication is made through internal channels. + tx.send((pid, WorkerMessage::Subscribe)) + .await + .expect("fatal error subscribing to supervisor, is the supervisor dropped?"); + Self { - id, + pid, tx, rx, state: State::Idle, } } - #[tracing::instrument(skip(self), fields(worker_id = self.id.to_string()))] + #[tracing::instrument(skip(self), fields(worker_id = self.pid.to_string()))] pub async fn run(mut self) { loop { tokio::select! { @@ -82,12 +98,12 @@ impl Worker { let message = match result { - Ok(result) => Response::Complete(Ok(result)), - Err(e) => Response::Complete(Err((e, task.clone()))) + Ok(result) => WorkerMessage::Complete(Ok(result)), + Err(e) => WorkerMessage::Complete(Err((e, task.clone()))) }; self.state = State::Idle; - let _ = self.tx.send((self.id, message)).await; + let _ = self.tx.send((self.pid, message)).await; } State::Idle => { @@ -99,7 +115,7 @@ impl Worker { } State::Stop => { tracing::debug!("received shutdown signal from supervisor"); - let _ = self.tx.send((self.id, Response::ShutdownAck)).await; + let _ = self.tx.send((self.pid, WorkerMessage::ShutdownAck)).await; return; } } @@ -155,12 +171,12 @@ impl State { // the worker, and spawn a new one with the same task. This is done to // ensure that we are able to capture error state before terminating a // worker. - fn next(&self, event: Request) -> Result, String> { + fn next(&self, event: SupervisorMessage) -> Result, String> { match (self, &event) { - (State::Idle, Request::Task(t)) => Ok(State::Running { task: t.clone() }), - (State::Idle, Request::Cancel) => Ok(State::Idle), - (State::Running { task: _ }, Request::Cancel) => Ok(State::Idle), - (_, Request::Shutdown) => Ok(State::Stop), + (State::Idle, SupervisorMessage::Task(t)) => Ok(State::Running { task: t.clone() }), + (State::Idle, SupervisorMessage::Cancel) => Ok(State::Idle), + (State::Running { task: _ }, SupervisorMessage::Cancel) => Ok(State::Idle), + (_, SupervisorMessage::Shutdown) => Ok(State::Stop), _ => Err(format!( "invalid transition, event: {event:?}, state: {self:?}" )),