Worker listening loop, worker commands

Implement a message listening loop where the worker checks for an
executes messages. Also implement commands that can be queued for the
workers, like a command to stop it processing.

Update the simple example accordingly.

Author: magnunm

README.md

@@ -12,3 +12,25 @@ Spin up a Redis instance and pass the connection URL as the first argument.
 ```sh
 cargo r --example simple_redis_queue 'redis://localhost:6379/1'
 ```
+
+This example uses two threads: a "worker thread" and a "controller thread". The
+controller queues tasks and commands that the worker picks up and executes. In
+the example code the controller queues a single tasks to print "Hello, World!"
+and return 42 and then waits for the worker to finish executing it. It then
+commands the worker to stop before terminating. The worker thread listens for
+messages and commands and executes them until it gets a stop command. It then
+terminates. A typical run looks something like this:
+
+```
+Controller thread: Queueing task
+Controller thread: Polling for task run result
+Worker thread: Registered worker with ID 01H3FV14GHSWEB48PPDB4Q506S
+Worker thread: Listening for messages...
+Hello, World!
+Controller thread: Polling for task run result
+Controller thread: Got task run result, 42
+Controller thread: Looking up all registered workers
+Controller thread: Sending command to stop worker 01H3FV14GHSWEB48PPDB4Q506S
+Controller thread: Done
+Worker thread: Done
+```

examples/simple_redis_queue.rs

@@ -1,14 +1,17 @@
 use anyhow;
-use parsnip::{self, brokers::redis::RedisBroker, task::Signature, task::Task, worker::Worker};
-use std::env;
+use parsnip::{
+    self, brokers::redis::RedisBroker, messages::Command, task::Signature, task::Task,
+    worker::Worker,
+};
+use std::{env, thread, time};
 
 struct HelloWorldTask {
     called_with_signature: Signature<Self>,
 }
 
 impl Task for HelloWorldTask {
     type ArgumentType = ();
-    type ReturnType = ();
+    type ReturnType = usize;
 
     const ID: &'static str = "HelloWorldTask";
 
@@ -20,27 +23,77 @@ impl Task for HelloWorldTask {
 
     fn run(_: &Self::ArgumentType) -> Self::ReturnType {
         println!("Hello, World!");
+        42
     }
 
     fn signature(&self) -> &Signature<Self> {
         &self.called_with_signature
     }
 }
 
-fn main() -> Result<(), anyhow::Error> {
-    let connect_url: String = env::args()
-        .nth(1)
-        .expect("No Redis connect URL passed the first an argument");
+fn controller_main(connect_url: String) -> () {
+    let mut broker = RedisBroker::new(&connect_url).expect("Can not connect to Redis");
+    let mut app = parsnip::App::new(&mut broker);
+    app.register_task::<HelloWorldTask>();
+
+    println!("Controller thread: Queueing task");
+    let signature_id = app.queue_task::<HelloWorldTask>(()).unwrap();
+
+    println!("Controller thread: Polling for task run result");
+    let mut result = app.get_task_result(&signature_id).unwrap();
+    while result.is_none() {
+        thread::sleep(time::Duration::from_millis(500));
+        println!("Controller thread: Polling for task run result");
+        result = app.get_task_result(&signature_id).unwrap();
+    }
 
-    let mut broker = RedisBroker::new(&connect_url)?;
+    println!(
+        "Controller thread: Got task run result, {}",
+        result.unwrap().result
+    );
 
+    println!("Controller thread: Looking up all registered workers");
+    match app.list_workers().unwrap() {
+        None => panic!(
+            "The task was picked up and run by a worker, but none are registered. This is a bug."
+        ),
+        Some(v) => v.into_iter().for_each(|worker| {
+            println!(
+                "Controller thread: Sending command to stop worker {}",
+                &worker.id
+            );
+            app.queue_command(&Command::StopWorker, &worker.id).unwrap();
+        }),
+    };
+    println!("Controller thread: Done");
+}
+
+fn worker_main(connect_url: String) -> () {
+    let mut broker = RedisBroker::new(&connect_url).expect("Can not connect to Redis");
     let mut app = parsnip::App::new(&mut broker);
     app.register_task::<HelloWorldTask>();
 
-    let worker = Worker::new(&app);
+    let worker = Worker::new(&app).expect("Worker initalization failed");
+    println!("Worker thread: Registered worker with ID {}", worker.id);
+    println!("Worker thread: Listening for messages...");
+    worker
+        .listen_for_messages()
+        .expect("Malfunction while handling messages.");
+    println!("Worker thread: Done")
+}
+
+fn main() -> Result<(), anyhow::Error> {
+    let connect_url: String = env::args()
+        .nth(1)
+        .expect("No Redis connect URL passed the first an argument");
+
+    let connect_url_controller = connect_url.clone();
+    let handle_controller_thread = thread::spawn(|| controller_main(connect_url_controller));
 
-    app.queue_task::<HelloWorldTask>(())?;
-    worker.take_first_task_in_queue()?;
+    let connect_url_worker = connect_url.clone();
+    let handler_worker_thread = thread::spawn(|| worker_main(connect_url_worker));
 
+    handle_controller_thread.join().unwrap();
+    handler_worker_thread.join().unwrap();
     Ok(())
 }

src/broker.rs

@@ -1,10 +1,38 @@
-use super::messages::{Message, ResultMessage};
+use super::messages::{Command, Message, ResultMessage};
 use anyhow::Result;
+use serde::{Deserialize, Serialize};
+
+#[derive(Serialize, Deserialize, Clone)]
+pub enum WorkerState {
+    Pending,
+    Running,
+    Stopped,
+}
+
+#[derive(Serialize, Deserialize, Clone)]
+pub struct WorkerInfo {
+    pub state: WorkerState,
+    pub id: String,
+}
 
 pub trait Broker {
     fn push_message(&self, message: &Message) -> Result<()>;
 
     fn pop_message(&self) -> Result<Option<Message>>;
 
+    fn push_command(&self, command: &Command, worker_id: &str) -> Result<()>;
+
+    fn pop_command(&self, worker_id: &str) -> Result<Option<Command>>;
+
     fn store_result(&self, result_message: ResultMessage) -> Result<()>;
+
+    fn get_result(&self, signature_id: &str) -> Result<Option<ResultMessage>>;
+
+    fn update_worker_info(&self, info: WorkerInfo) -> Result<()>;
+
+    fn remove_worker_info(&self, worker_id: &str) -> Result<()>;
+
+    fn get_worker_info(&self, worker_id: &str) -> Result<Option<WorkerInfo>>;
+
+    fn all_workers(&self) -> Result<Option<Vec<WorkerInfo>>>;
 }

src/brokers/redis.rs

@@ -1,4 +1,4 @@
-use crate::broker::Broker;
+use crate::broker::{Broker, WorkerInfo};
 
 use anyhow::Result;
 use redis::{self, Commands};
@@ -7,7 +7,9 @@ use serde_json;
 pub struct RedisBroker {
     redis_client: redis::Client,
     queue: String,
+    command_queue_prefix: String,
     result_hash_map: String,
+    worker_register: String,
 }
 
 impl RedisBroker {
@@ -17,7 +19,9 @@ impl RedisBroker {
         Ok(Self {
             redis_client,
             queue: "parsnip_queue".to_string(),
+            command_queue_prefix: "parsnip_command_queue".to_string(),
             result_hash_map: "parsnip_task_results".to_string(),
+            worker_register: "worker_register".to_string(),
         })
     }
 }
@@ -47,4 +51,69 @@ impl Broker for RedisBroker {
         )?;
         Ok(())
     }
+
+    fn get_result(&self, signature_id: &str) -> Result<Option<crate::messages::ResultMessage>> {
+        let mut con = self.redis_client.get_connection()?;
+        let serialized_result: Option<String> = con.hget(&self.result_hash_map, signature_id)?;
+        serialized_result.map_or(Ok(None), |v| {
+            serde_json::from_str(&v).map_err(|e| anyhow::anyhow!("{}", e))
+        })
+    }
+
+    fn push_command(&self, command: &crate::messages::Command, worker_id: &str) -> Result<()> {
+        let mut con = self.redis_client.get_connection()?;
+        con.lpush(
+            &format!("{}_{}", self.command_queue_prefix, worker_id),
+            serde_json::to_string(&command)?,
+        )?;
+        Ok(())
+    }
+
+    fn pop_command(&self, worker_id: &str) -> Result<Option<crate::messages::Command>> {
+        let mut con = self.redis_client.get_connection()?;
+        let serialized_command: Option<String> = con.rpop(
+            &format!("{}_{}", self.command_queue_prefix, worker_id),
+            None,
+        )?;
+        match serialized_command {
+            Some(v) => Ok(Some(serde_json::from_str(&v)?)),
+            None => Ok(None),
+        }
+    }
+
+    fn update_worker_info(&self, info: WorkerInfo) -> Result<()> {
+        let mut con = self.redis_client.get_connection()?;
+        con.hset(
+            &self.worker_register,
+            &info.id,
+            serde_json::to_string(&info)?,
+        )?;
+        Ok(())
+    }
+
+    fn get_worker_info(&self, worker_id: &str) -> Result<Option<WorkerInfo>> {
+        let mut con = self.redis_client.get_connection()?;
+        let serialized_info: Option<String> = con.hget(&self.worker_register, worker_id)?;
+        serialized_info.map_or(Ok(None), |v| {
+            serde_json::from_str(&v).map_err(|e| anyhow::anyhow!("{}", e))
+        })
+    }
+
+    fn remove_worker_info(&self, worker_id: &str) -> Result<()> {
+        let mut con = self.redis_client.get_connection()?;
+        con.hdel(&self.worker_register, worker_id)
+            .map_err(|e| anyhow::anyhow!("{}", e))
+    }
+
+    fn all_workers(&self) -> Result<Option<Vec<WorkerInfo>>> {
+        let mut con = self.redis_client.get_connection()?;
+
+        let serialized_info: Option<Vec<String>> = con.hvals(&self.worker_register)?;
+        serialized_info.map_or(Ok(None), |info_vec| {
+            info_vec
+                .iter()
+                .map(|v| serde_json::from_str(v).map_err(|e| anyhow::anyhow!("{}", e)))
+                .collect()
+        })
+    }
 }

src/lib.rs

@@ -5,8 +5,8 @@ mod runner;
 pub mod task;
 pub mod worker;
 
-use broker::Broker;
-use messages::{Message, ResultMessage};
+use broker::{Broker, WorkerInfo};
+use messages::{Command, Message, ResultMessage};
 use runner::TaskRunnerBuilder;
 use task::{Signature, Task};
 
@@ -33,25 +33,46 @@ impl<'a, B: Broker + 'static> App<'a, B> {
             .insert(T::ID.into(), Box::new(runner::build_task_runner::<T, B>));
     }
 
-    pub fn queue_task<T: Task + 'static>(&self, arg: T::ArgumentType) -> Result<(), Error> {
+    /// Queue a task for pickup by a worker.
+    ///
+    /// Returns the signature_id for the task invocation, which can be used to
+    /// lookup the result of running the task.
+    pub fn queue_task<T: Task + 'static>(&self, arg: T::ArgumentType) -> Result<String, Error> {
         if !self.task_runner_builders.contains_key(T::ID.into()) {
             anyhow::bail!(
                 "Can not queue task with ID '{}' as it is not registered.",
                 T::ID
             );
         }
 
+        let signature_id = Ulid::new().to_string();
         let signature = Signature::<T> {
             arg,
-            id: Ulid::new().to_string(),
+            id: signature_id.clone(),
         };
         self.broker
             .push_message(&Message {
                 task_id: T::ID.into(),
                 signature: serde_json::to_string(&signature)?,
             })
             .context("Failed to put task invocation on the queue.")?;
-        Ok(())
+        Ok(signature_id)
+    }
+
+    pub fn get_task_result(&self, signatrue_id: &str) -> Result<Option<ResultMessage>, Error> {
+        self.broker.get_result(signatrue_id)
+    }
+
+    pub fn queue_command(&self, command: &Command, worker_id: &str) -> Result<(), Error> {
+        self.broker.push_command(command, worker_id)
+    }
+
+    pub fn get_worker_info(&self, worker_id: &str) -> Result<Option<WorkerInfo>, Error> {
+        self.broker.get_worker_info(worker_id)
+    }
+
+    pub fn list_workers(&self) -> Result<Option<Vec<WorkerInfo>>, Error> {
+        self.broker.all_workers()
     }
 
     fn handle_message(&self, message: &Message) -> Result<(), Error> {
@@ -69,7 +90,14 @@ impl<'a, B: Broker + 'static> App<'a, B> {
     }
 
     fn store_task_result(&self, result: ResultMessage) -> Result<(), Error> {
-        self.broker.store_result(result)?;
-        Ok(())
+        self.broker.store_result(result)
+    }
+
+    fn update_worker_info(&self, info: WorkerInfo) -> Result<(), Error> {
+        self.broker.update_worker_info(info)
+    }
+
+    fn remove_worker_info(&self, worker_id: &str) -> Result<(), Error> {
+        self.broker.remove_worker_info(worker_id)
     }
 }

src/messages.rs

@@ -11,3 +11,8 @@ pub struct ResultMessage {
     pub signature_id: String,
     pub result: String,
 }
+
+#[derive(Serialize, Deserialize, Clone)]
+pub enum Command {
+    StopWorker,
+}