Potentially refactor library structure

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};
 

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},
 };

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<W: Workable> {
+pub enum SupervisorMessage<W: Workable> {
     State,
     Task(W::Task),
     Cancel,
+
+    /// Acknowledge
+    SubscribeAck,
     Shutdown,
 }
 
 #[derive(Debug)]
-pub enum Response<W: Workable> {
-    ShutdownAck, // Acknowledge shutdown
+pub enum WorkerMessage<W: Workable> {
     Complete(Result<W::Output, (W::Error, W::Task)>),
+    /// Worker subscribing to a running Supervisor.
+    Subscribe,
+    ShutdownAck, // Acknowledge shutdown
 }

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<Pid, (Sender<Request<W>>, JoinHandle<()>)>,
+    pool: BTreeMap<Pid, (Sender<SupervisorMessage<W>>, 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<Pid, (Sender<Request<W>>, JoinHandle<()>)>,
+    checked: BTreeMap<Pid, (Sender<SupervisorMessage<W>>, JoinHandle<()>)>,
 
     /// Pending queue of tasks to be executed
     tasks: VecDeque<W::Task>,
@@ -36,11 +36,11 @@ pub struct Supervisor<'buf, W: Workable> {
     pub results: &'buf RingBuffer<Result<W::Output, (W::Error, W::Task)>>,
 
     /// Sender end, mostly kept here for the time being to simplify spawning.
-    sender: Sender<(Pid, Response<W>)>,
+    pub sender: Sender<(Pid, WorkerMessage<W>)>,
 
     /// 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<W>)>,
+    receiver: Receiver<(Pid, WorkerMessage<W>)>,
 }
 
 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<Pid, (Sender<Request<W>>, JoinHandle<()>)> = BTreeMap::new();
+        let mut workers: BTreeMap<Pid, (Sender<SupervisorMessage<W>>, 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);
                     }

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<W: Workable> {
-    id: Pid,
-    tx: Sender<(Pid, Response<W>)>,
-    rx: Receiver<Request<W>>,
+    pid: Pid,
+    tx: Sender<(Pid, WorkerMessage<W>)>,
+    rx: Receiver<SupervisorMessage<W>>,
     state: State<W>,
 }
 
@@ -43,16 +43,32 @@ impl<W: Workable> Drop for Worker<W> {
 }
 
 impl<W: Workable> Worker<W> {
-    pub fn new(id: Pid, tx: Sender<(Pid, Response<W>)>, rx: Receiver<Request<W>>) -> Self {
+    pub async fn new(
+        tx: Sender<(Pid, WorkerMessage<W>)>,
+        rx: Receiver<SupervisorMessage<W>>,
+    ) -> 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<W: Workable> Worker<W> {
 
 
                                 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<W: Workable> Worker<W> {
                             }
                             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<W: Workable> State<W> {
     //  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<W>) -> Result<State<W>, String> {
+    fn next(&self, event: SupervisorMessage<W>) -> Result<State<W>, 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:?}"
             )),