use std::collections::HashMap;
use std::net::SocketAddr;
use std::sync::{Arc, Mutex};
use futures::{Future, Poll, Sink, Stream};
use futures::future;
use futures::future::Executor;
use futures::sync::{mpsc, oneshot};
use error;
use resp;
use resp::FromResp;
use super::connect::{connect, ClientConnection};
fn err<S>(msg: S) -> Box<Future<Item = (), Error = error::Error> + Send>
where
S: Into<String>,
{
Box::new(future::err(error::internal(msg)))
}
pub fn pubsub_connect<E>(
addr: &SocketAddr,
executor: E,
) -> Box<Future<Item = PubsubConnection, Error = error::Error>>
where
E: Executor<Box<Future<Item = (), Error = ()> + Send>> + 'static,
{
let pubsub_con = connect(addr)
.map_err(|e| e.into())
.and_then(move |connection| {
let ClientConnection { sender, receiver } = connection;
let (out_tx, out_rx) = mpsc::unbounded();
let sender = Box::new(
sender
.sink_map_err(|e| error!("Sending socket error: {}", e))
.send_all(out_rx.map_err(|()| error!("Outgoing message channel error")))
.map(|_| debug!("Send all completed successfully")),
) as Box<Future<Item = (), Error = ()> + Send>;
let subs = Arc::new(Mutex::new(PubsubSubscriptions {
pending: HashMap::new(),
confirmed: HashMap::new(),
}));
let subs_reader = subs.clone();
let receiver_raw = receiver.for_each(move |msg| {
let (message_type, topic, msg) = match msg {
resp::RespValue::Array(mut messages) => match (
messages.pop(),
messages.pop(),
messages.pop(),
messages.pop(),
) {
(Some(msg), Some(topic), Some(message_type), None) => {
match (msg, String::from_resp(topic), message_type) {
(msg, Ok(topic), resp::RespValue::BulkString(bytes)) => {
(bytes, topic, msg)
}
_ => return err("RESP Message structure invalid"),
}
}
_ => {
return err(format!(
"Incorrect number of records in PUBSUB message: {}",
messages.len()
))
}
},
_ => return err("Incorrect message type"),
};
let mut subs = subs_reader.lock().expect("Lock is tainted");
if message_type == b"subscribe" {
if let Some((tx, notification_tx)) = subs.pending.remove(&topic) {
subs.confirmed.insert(topic, tx);
match notification_tx.send(()) {
Ok(()) => {
let ok =
future::ok(()).map_err(|_: ()| error::internal("unreachable"));
Box::new(ok)
}
Err(()) => err("Cannot send notification of subscription"),
}
} else {
err("Received subscribe notification that wasn't expected")
}
} else if message_type == b"unsubscribe" {
if let Some(_) = subs.confirmed.remove(&topic) {
let result = if subs.confirmed.is_empty() {
future::err(error::Error::EndOfStream)
} else {
future::ok(())
};
Box::new(result)
} else {
err("Receieved unsubscription notification from a topic we're not subscribed to")
}
} else if message_type == b"message" {
match subs.confirmed.get(&topic) {
Some(tx) => {
let future = tx.clone().send(msg).map(|_| ()).map_err(|e| e.into());
Box::new(future)
}
None => {
let ok = future::ok(()).map_err(|_: ()| error::internal("unreachable"));
Box::new(ok)
}
}
} else {
return err(format!("Unexpected message type: {:?}", message_type));
}
});
let receiver = Box::new(
receiver_raw
.then(|result| match result {
Ok(_) => future::ok(()),
Err(error::Error::EndOfStream) => future::ok(()),
Err(e) => future::err(e),
})
.map(|_| debug!("End of PUBSUB connection, successful"))
.map_err(|e| error!("PUBSUB Receiver error: {}", e)),
) as Box<Future<Item = (), Error = ()> + Send>;
match executor
.execute(sender)
.and_then(|_| executor.execute(receiver))
{
Ok(()) => future::ok(PubsubConnection {
out_tx: out_tx,
subscriptions: subs,
}),
Err(e) => future::err(error::internal(format!(
"Cannot execute background task: {:?}",
e
))),
}
});
Box::new(pubsub_con)
}
struct PubsubSubscriptions {
pending: HashMap<String, (mpsc::Sender<resp::RespValue>, oneshot::Sender<()>)>,
confirmed: HashMap<String, mpsc::Sender<resp::RespValue>>,
}
#[derive(Clone)]
pub struct PubsubConnection {
out_tx: mpsc::UnboundedSender<resp::RespValue>,
subscriptions: Arc<Mutex<PubsubSubscriptions>>,
}
impl PubsubConnection {
pub fn subscribe<T: Into<String>>(
&self,
topic: T,
) -> Box<Future<Item = PubsubStream, Error = error::Error> + Send> {
let topic = topic.into();
let mut subs = self.subscriptions.lock().expect("Lock is tainted");
let (tx, rx) = mpsc::channel(10);
if subs.confirmed.contains_key(&topic) || subs.pending.contains_key(&topic) {
return Box::new(future::err(error::internal("Already subscribed")));
}
let (notification_tx, notification_rx) = oneshot::channel();
let subscribe_msg = resp_array!["SUBSCRIBE", &topic];
subs.pending.insert(topic.clone(), (tx, notification_tx));
let new_out_tx = self.out_tx.clone();
match self.out_tx.unbounded_send(subscribe_msg) {
Ok(_) => {
let done = notification_rx
.map(|_| PubsubStream::new(topic, rx, new_out_tx))
.map_err(|e| e.into());
Box::new(done)
}
Err(e) => Box::new(future::err(e.into())),
}
}
}
pub struct PubsubStream {
topic: String,
underlying: mpsc::Receiver<resp::RespValue>,
out_tx: mpsc::UnboundedSender<resp::RespValue>,
}
impl PubsubStream {
fn new(
topic: String,
stream: mpsc::Receiver<resp::RespValue>,
out_tx: mpsc::UnboundedSender<resp::RespValue>,
) -> Self {
PubsubStream {
topic: topic,
underlying: stream,
out_tx: out_tx,
}
}
}
impl Stream for PubsubStream {
type Item = resp::RespValue;
type Error = ();
fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
self.underlying.poll()
}
}
impl Drop for PubsubStream {
fn drop(&mut self) {
let send_f = self.out_tx
.clone()
.send(resp_array!["UNSUBSCRIBE", &self.topic]);
match send_f.wait() {
Ok(_) => debug!("Send unsubscribe command for topic: {}", self.topic),
Err(e) => error!(
"Failed to send unsubscribe command for topic: {}, error: {}",
self.topic, e
),
}
}
}