Batches and trackable jobs

src/consumer/mod.rs

@@ -1,13 +1,14 @@
 use crate::error::Error;
-use crate::proto::{self, Client, ClientOptions, HeartbeatStatus, Reconnect};
+use crate::proto::{
+    self, parse_provided_or_from_env, Ack, Client, ClientOptions, Fail, HeartbeatStatus, Job,
+    Reconnect,
+};
 use fnv::FnvHashMap;
 use std::error::Error as StdError;
 use std::io::prelude::*;
 use std::net::TcpStream;
 use std::sync::{atomic, Arc, Mutex};
 
-use crate::proto::{Ack, Fail, Job};
-
 const STATUS_RUNNING: usize = 0;
 const STATUS_QUIET: usize = 1;
 const STATUS_TERMINATING: usize = 2;
@@ -213,10 +214,7 @@ impl<E> ConsumerBuilder<E> {
     ///
     /// If `url` is given, but does not specify a port, it defaults to 7419.
     pub fn connect(self, url: Option<&str>) -> Result<Consumer<TcpStream, E>, Error> {
-        let url = match url {
-            Some(url) => proto::url_parse(url),
-            None => proto::url_parse(&proto::get_env_url()),
-        }?;
+        let url = parse_provided_or_from_env(url)?;
         let stream = TcpStream::connect(proto::host_from_url(&url))?;
         Self::connect_with(self, stream, url.password().map(|p| p.to_string()))
     }
@@ -228,6 +226,7 @@ impl<E> ConsumerBuilder<E> {
         pwd: Option<String>,
     ) -> Result<Consumer<S, E>, Error> {
         self.opts.password = pwd;
+        self.opts.is_worker = true;
         Ok(Consumer::new(
             Client::new(stream, self.opts)?,
             self.workers,

src/lib.rs

@@ -76,5 +76,11 @@ pub use tls::TlsStream;
 pub use crate::consumer::{Consumer, ConsumerBuilder};
 pub use crate::error::Error;
 pub use crate::producer::Producer;
-pub use crate::proto::Reconnect;
-pub use crate::proto::{Job, JobBuilder};
+pub use crate::proto::{Client, Job, JobBuilder, Reconnect};
+
+#[cfg(feature = "ent")]
+#[cfg_attr(docsrs, doc(cfg(feature = "ent")))]
+pub use crate::proto::{
+    Batch, BatchBuilder, BatchHandle, BatchStatus, CallbackState, JobState, Progress,
+    ProgressUpdate, ProgressUpdateBuilder,
+};

src/producer/mod.rs

@@ -1,8 +1,11 @@
 use crate::error::Error;
-use crate::proto::{self, Client, Info, Job, Push, QueueAction, QueueControl};
+use crate::proto::{Client, Info, Job, Push, QueueAction, QueueControl};
 use std::io::prelude::*;
 use std::net::TcpStream;
 
+#[cfg(feature = "ent")]
+use crate::proto::{Batch, BatchHandle, CommitBatch, OpenBatch};
+
 /// `Producer` is used to enqueue new jobs that will in turn be processed by Faktory workers.
 ///
 /// # Connecting to Faktory
@@ -82,21 +85,16 @@ impl Producer<TcpStream> {
     ///
     /// If `url` is given, but does not specify a port, it defaults to 7419.
     pub fn connect(url: Option<&str>) -> Result<Self, Error> {
-        let url = match url {
-            Some(url) => proto::url_parse(url),
-            None => proto::url_parse(&proto::get_env_url()),
-        }?;
-        let stream = TcpStream::connect(proto::host_from_url(&url))?;
-        Self::connect_with(stream, url.password().map(|p| p.to_string()))
+        let c = Client::connect(url)?;
+        Ok(Producer { c })
     }
 }
 
 impl<S: Read + Write> Producer<S> {
     /// Connect to a Faktory server with a non-standard stream.
     pub fn connect_with(stream: S, pwd: Option<String>) -> Result<Producer<S>, Error> {
-        Ok(Producer {
-            c: Client::new_producer(stream, pwd)?,
-        })
+        let c = Client::connect_with(stream, pwd)?;
+        Ok(Producer { c })
     }
 
     /// Enqueue the given job on the Faktory server.
@@ -129,6 +127,30 @@ impl<S: Read + Write> Producer<S> {
             .issue(&QueueControl::new(QueueAction::Resume, queues))?
             .await_ok()
     }
+
+    /// Initiate a new batch of jobs.
+    #[cfg(feature = "ent")]
+    #[cfg_attr(docsrs, doc(cfg(feature = "ent")))]
+    pub fn start_batch(&mut self, batch: Batch) -> Result<BatchHandle<'_, S>, Error> {
+        let bid = self.c.issue(&batch)?.read_bid()?;
+        Ok(BatchHandle::new(bid, self))
+    }
+
+    /// Open an already existing batch of jobs.
+    ///
+    /// This will not error if a batch with the provided `bid` does not exist,
+    /// rather `Ok(None)` will be returned.
+    #[cfg(feature = "ent")]
+    #[cfg_attr(docsrs, doc(cfg(feature = "ent")))]
+    pub fn open_batch(&mut self, bid: String) -> Result<Option<BatchHandle<'_, S>>, Error> {
+        let bid = self.c.issue(&OpenBatch::from(bid))?.maybe_bid()?;
+        Ok(bid.map(|bid| BatchHandle::new(bid, self)))
+    }
+
+    #[cfg(feature = "ent")]
+    pub(crate) fn commit_batch(&mut self, bid: String) -> Result<(), Error> {
+        self.c.issue(&CommitBatch::from(bid))?.await_ok()
+    }
 }
 
 #[cfg(test)]

src/proto/batch/cmd.rs

@@ -0,0 +1,65 @@
+use crate::proto::single::FaktoryCommand;
+use crate::{Batch, Error};
+use std::io::Write;
+
+impl FaktoryCommand for Batch {
+    fn issue<W: Write>(&self, w: &mut W) -> Result<(), Error> {
+        w.write_all(b"BATCH NEW ")?;
+        serde_json::to_writer(&mut *w, self).map_err(Error::Serialization)?;
+        Ok(w.write_all(b"\r\n")?)
+    }
+}
+
+// ----------------------------------------------
+
+pub struct CommitBatch(String);
+
+impl From<String> for CommitBatch {
+    fn from(value: String) -> Self {
+        CommitBatch(value)
+    }
+}
+
+impl FaktoryCommand for CommitBatch {
+    fn issue<W: Write>(&self, w: &mut W) -> Result<(), Error> {
+        w.write_all(b"BATCH COMMIT ")?;
+        w.write_all(self.0.as_bytes())?;
+        Ok(w.write_all(b"\r\n")?)
+    }
+}
+
+// ----------------------------------------------
+
+pub struct GetBatchStatus(String);
+
+impl From<String> for GetBatchStatus {
+    fn from(value: String) -> Self {
+        GetBatchStatus(value)
+    }
+}
+
+impl FaktoryCommand for GetBatchStatus {
+    fn issue<W: Write>(&self, w: &mut W) -> Result<(), Error> {
+        w.write_all(b"BATCH STATUS ")?;
+        w.write_all(self.0.as_bytes())?;
+        Ok(w.write_all(b"\r\n")?)
+    }
+}
+
+// ----------------------------------------------
+
+pub struct OpenBatch(String);
+
+impl From<String> for OpenBatch {
+    fn from(value: String) -> Self {
+        OpenBatch(value)
+    }
+}
+
+impl FaktoryCommand for OpenBatch {
+    fn issue<W: Write>(&self, w: &mut W) -> Result<(), Error> {
+        w.write_all(b"BATCH OPEN ")?;
+        w.write_all(self.0.as_bytes())?;
+        Ok(w.write_all(b"\r\n")?)
+    }
+}