├── .gitignore ├── .travis.yml ├── LICENSE ├── src ├── commands.rs ├── lib.rs ├── error.rs ├── response.rs ├── config.rs ├── producer.rs ├── consumer.rs ├── protocol.rs └── codec.rs ├── examples ├── pub.rs ├── mpub.rs └── sub.rs ├── Cargo.toml └── README.md /.gitignore: -------------------------------------------------------------------------------- 1 | target 2 | Cargo.lock 3 | -------------------------------------------------------------------------------- /.travis.yml: -------------------------------------------------------------------------------- 1 | language: rust 2 | rust: 3 | - nightly 4 | - stable 5 | sudo: false 6 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright (c) 2017 Flavio Oliveira 2 | 3 | Licensed under either of 4 | 5 | * Apache License, Version 2.0, (http://www.apache.org/licenses/LICENSE-2.0) 6 | * MIT license (http://opensource.org/licenses/MIT) 7 | 8 | at your option. 9 | -------------------------------------------------------------------------------- /src/commands.rs: -------------------------------------------------------------------------------- 1 | pub mod commands { 2 | pub const VERSION_2: &'static str = " V2"; 3 | 4 | pub const PUB: &'static str = "PUB"; 5 | pub const MPUB: &'static str = "MPUB"; 6 | pub const DPUB: &'static str = "DPUB"; 7 | 8 | pub const SUB: &'static str = "SUB"; 9 | 10 | pub const RDY: &'static str = "RDY"; 11 | pub const FIN: &'static str = "FIN"; 12 | 13 | pub const NOP: &'static str = "NOP"; 14 | pub const IDENTIFY: &'static str = "IDENTIFY"; 15 | } -------------------------------------------------------------------------------- /src/lib.rs: -------------------------------------------------------------------------------- 1 | extern crate serde; 2 | extern crate serde_json; 3 | extern crate futures; 4 | extern crate log; 5 | extern crate tokio_io; 6 | extern crate tokio_core; 7 | extern crate tokio_service; 8 | extern crate tokio_proto; 9 | extern crate bytes; 10 | extern crate byteorder; 11 | extern crate hostname; 12 | 13 | #[macro_use] 14 | extern crate serde_derive; 15 | 16 | mod codec; 17 | mod commands; 18 | mod protocol; 19 | pub mod response; 20 | pub mod error; 21 | pub mod config; 22 | pub mod consumer; 23 | pub mod producer; -------------------------------------------------------------------------------- /examples/pub.rs: -------------------------------------------------------------------------------- 1 | extern crate futures; 2 | extern crate tokio_core; 3 | extern crate nsqueue; 4 | 5 | use futures::Future; 6 | use tokio_core::reactor::Core; 7 | 8 | use nsqueue::config::*; 9 | use nsqueue::producer::*; 10 | 11 | fn main() { 12 | let mut core = Core::new().unwrap(); 13 | let handle = core.handle(); 14 | 15 | let addr = "127.0.0.1:4150".parse().unwrap(); 16 | 17 | let res = Producer::connect(&addr, &handle, Config::default()) 18 | .and_then(|conn| { 19 | conn.publish("some_topic".into(), "some_message".into()) 20 | .and_then(move |response| { 21 | println!("Response: {:?}", response); 22 | Ok(()) 23 | }) 24 | }); 25 | core.run(res).unwrap(); 26 | } -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "nsqueue" 3 | version = "0.1.5" 4 | authors = ["Flavio Oliveira "] 5 | 6 | description = "Rust client for the NSQ realtime message processing system" 7 | license = "MIT OR Apache-2.0" 8 | keywords = ["nsq", "queue", "tokio", "asynchronous"] 9 | categories = ["caching"] 10 | repository = "https://github.com/wisespace-io/nsqueue" 11 | readme = "README.md" 12 | 13 | [badges] 14 | travis-ci = { repository = "wisespace-io/nsqueue" } 15 | 16 | [lib] 17 | name = "nsqueue" 18 | path = "src/lib.rs" 19 | 20 | [dependencies] 21 | log = "^0.3" 22 | bytes = "^0.4" 23 | futures = "^0.1" 24 | tokio-io = "^0.1" 25 | tokio-core = "^0.1" 26 | tokio-proto = "^0.1" 27 | tokio-service = "^0.1" 28 | byteorder = "1.0.0" 29 | hostname = "^0.1" 30 | serde = "1.0" 31 | serde_json = "1.0" 32 | serde_derive = "1.0" 33 | -------------------------------------------------------------------------------- /examples/mpub.rs: -------------------------------------------------------------------------------- 1 | extern crate futures; 2 | extern crate tokio_core; 3 | extern crate nsqueue; 4 | 5 | use futures::Future; 6 | use tokio_core::reactor::Core; 7 | 8 | use nsqueue::config::*; 9 | use nsqueue::producer::*; 10 | 11 | fn main() { 12 | let mut core = Core::new().unwrap(); 13 | let handle = core.handle(); 14 | 15 | let addr = "127.0.0.1:4150".parse().unwrap(); 16 | 17 | let mut messages: Vec = Vec::new(); 18 | messages.push("First message".into()); 19 | messages.push("Second message".into()); 20 | 21 | let res = Producer::connect(&addr, &handle, Config::default()) 22 | .and_then(|conn| { 23 | conn.mpublish("some_topic".into(), messages) 24 | .and_then(move |response| { 25 | println!("Response: {:?}", response); 26 | Ok(()) 27 | }) 28 | }); 29 | core.run(res).unwrap(); 30 | } -------------------------------------------------------------------------------- /src/error.rs: -------------------------------------------------------------------------------- 1 | use std::io::Error as ioError; 2 | use std::fmt; 3 | use std::str; 4 | use std::error; 5 | 6 | #[derive(Debug)] 7 | pub enum NsqError { 8 | IOError(ioError), 9 | } 10 | 11 | impl fmt::Display for NsqError { 12 | fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { 13 | match *self { 14 | NsqError::IOError(ref err) => write!(f, "IO error: {}", err), 15 | } 16 | } 17 | } 18 | 19 | impl error::Error for NsqError { 20 | fn description(&self) -> &str { 21 | match *self { 22 | NsqError::IOError(ref err) => err.description(), 23 | } 24 | } 25 | 26 | fn cause(&self) -> Option<&error::Error> { 27 | match *self { 28 | NsqError::IOError(ref err) => Some(err), 29 | } 30 | } 31 | } 32 | 33 | impl From for NsqError { 34 | fn from(err: ioError) -> NsqError { 35 | NsqError::IOError(err) 36 | } 37 | } -------------------------------------------------------------------------------- /examples/sub.rs: -------------------------------------------------------------------------------- 1 | extern crate futures; 2 | extern crate tokio_core; 3 | extern crate nsqueue; 4 | 5 | use futures::{Stream, Future}; 6 | use tokio_core::reactor::Core; 7 | use nsqueue::config::*; 8 | use nsqueue::consumer::*; 9 | 10 | fn main() { 11 | let mut core = Core::new().unwrap(); 12 | let handle = core.handle(); 13 | let addr = "127.0.0.1:4150".parse().unwrap(); 14 | 15 | core.run( 16 | Consumer::connect(&addr, &handle, Config::default()) 17 | .and_then(|conn| { 18 | conn.subscribe("some_topic".into(), "some_channel".into()) 19 | .and_then(move |response| { 20 | let ret = response.for_each(move |message| { 21 | if message.message_id == "_heartbeat_" { 22 | conn.nop(); 23 | } else { 24 | println!("Response {:?} {:?}", message.message_id, message.message_body); 25 | conn.fin(message.message_id); // Inform NSQ (Message consumed) 26 | } 27 | Ok(()) 28 | }); 29 | ret 30 | }) 31 | }) 32 | ).unwrap(); 33 | } -------------------------------------------------------------------------------- /src/response.rs: -------------------------------------------------------------------------------- 1 | use std::io; 2 | use futures::{Stream, Poll, Async}; 3 | use tokio_proto::streaming::Body; 4 | 5 | #[derive(Debug)] 6 | pub struct ResponseStream { 7 | pub inner: Body, 8 | } 9 | 10 | impl Stream for ResponseStream { 11 | type Item = Message; 12 | type Error = io::Error; 13 | 14 | fn poll(&mut self) -> Poll, io::Error> { 15 | match self.inner.poll().unwrap() { 16 | Async::Ready(Some(request)) => { 17 | Ok(Async::Ready(Some(request))) 18 | } 19 | Async::Ready(None) => { 20 | // the stream finished. 21 | Ok(Async::Ready(None)) 22 | } 23 | Async::NotReady => { 24 | // no more messages to read 25 | Ok(Async::NotReady) 26 | } 27 | } 28 | } 29 | } 30 | 31 | /* 32 | impl Sink for ResponseStream 33 | where T: Sink, 34 | { 35 | type SinkItem = String; 36 | type SinkError = io::Error; 37 | 38 | fn start_send(&mut self, item: String) -> StartSend { 39 | self.upstream.start_send(item) 40 | } 41 | 42 | fn poll_complete(&mut self) -> Poll<(), io::Error> { 43 | self.upstream.poll_complete() 44 | } 45 | } 46 | */ 47 | 48 | pub struct Message { 49 | pub timestamp: i64, 50 | pub message_id: String, 51 | pub message_body: String 52 | } -------------------------------------------------------------------------------- /src/config.rs: -------------------------------------------------------------------------------- 1 | #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] 2 | pub struct Config { 3 | // Identifiers sent to nsqd representing this client 4 | client_id: Option, 5 | short_id: Option, 6 | long_id: Option, 7 | hostname: Option, 8 | user_agent: String, 9 | 10 | // Compression Settings 11 | deflate: bool, 12 | deflate_level: u16, 13 | snappy: bool, 14 | 15 | feature_negotiation: bool, 16 | 17 | // Duration of time between heartbeats. 18 | heartbeat_interval: i64, 19 | 20 | // Timeout used by nsqd before flushing buffered writes (set to 0 to disable). 21 | message_timeout: u32, 22 | 23 | // Size of the buffer (in bytes) used by nsqd for buffering writes to this connection 24 | output_buffer_size: u64, 25 | output_buffer_timeout: u32, 26 | 27 | // Integer percentage to sample the channel (requires nsqd 0.2.25+) 28 | sample_rate: u16, 29 | 30 | // tls_v1 - Bool enable TLS negotiation 31 | tls_v1: bool, 32 | } 33 | use hostname::get_hostname; 34 | 35 | #[allow(dead_code)] 36 | impl Config { 37 | pub fn default() -> Config { 38 | Config { 39 | client_id: get_hostname(), 40 | short_id: get_hostname(), 41 | long_id: get_hostname(), 42 | user_agent: String::from("github.com/wisespace-io/nsqueue"), 43 | hostname: get_hostname(), 44 | deflate: false, 45 | deflate_level: 6, 46 | snappy: false, 47 | feature_negotiation: true, 48 | heartbeat_interval: 30000, 49 | message_timeout: 0, 50 | output_buffer_size: 16384, 51 | output_buffer_timeout: 250, 52 | sample_rate: 0, 53 | tls_v1: false, 54 | } 55 | } 56 | 57 | pub fn client_id(mut self, client_id: String) -> Self { 58 | self.client_id = Some(client_id); 59 | self 60 | } 61 | 62 | pub fn hostname(mut self, hostname: String) -> Self { 63 | self.hostname = Some(hostname); 64 | self 65 | } 66 | 67 | pub fn user_agent(mut self, user_agent: String) -> Self { 68 | self.user_agent = user_agent; 69 | self 70 | } 71 | 72 | pub fn snappy(mut self, snappy: bool) -> Self { 73 | self.snappy = snappy; 74 | self 75 | } 76 | } -------------------------------------------------------------------------------- /src/producer.rs: -------------------------------------------------------------------------------- 1 | use futures::{Future}; 2 | 3 | use tokio_service::Service; 4 | use tokio_core::reactor::Handle; 5 | 6 | use tokio_proto::{TcpClient}; 7 | use tokio_proto::streaming::{Message}; 8 | use tokio_proto::util::client_proxy::ClientProxy; 9 | 10 | use std::io; 11 | use std::net::SocketAddr; 12 | 13 | use config::Config; 14 | use codec::{NsqMessage, NsqResponseMessage, ClientTypeMap}; 15 | use protocol::{NsqProtocol, RequestMessage}; 16 | 17 | pub struct Producer { 18 | inner: ClientTypeMap>, 19 | } 20 | 21 | impl Producer { 22 | /// Establish a connection and send protocol version. 23 | pub fn connect(addr: &SocketAddr, handle: &Handle, config: Config) -> Box> { 24 | let protocol = NsqProtocol::new(config); 25 | let ret = TcpClient::new(protocol) 26 | .connect(addr, handle) 27 | .map(|client_proxy| { 28 | let type_map = ClientTypeMap { inner: client_proxy }; 29 | Producer { inner: type_map } 30 | }); 31 | 32 | Box::new(ret) 33 | } 34 | 35 | // Publish a message to a topic 36 | pub fn publish(&self, topic: String, message: String) -> Box> { 37 | let mut request = RequestMessage::new(); 38 | request.create_pub_command(topic, message); 39 | 40 | self.handler(request) 41 | } 42 | 43 | // Publish multiple messages to a topic (atomically) 44 | pub fn mpublish(&self, topic: String, messages: Vec) -> Box> { 45 | let mut request = RequestMessage::new(); 46 | request.create_mpub_command(topic, messages); 47 | 48 | self.handler(request) 49 | } 50 | 51 | // Publish a deferred message to a topic 52 | pub fn dpublish(&self, topic: String, message: String, defer_time: i64) -> Box> { 53 | let mut request = RequestMessage::new(); 54 | request.create_dpub_command(topic, message, defer_time); 55 | 56 | self.handler(request) 57 | } 58 | 59 | fn handler(&self, request: RequestMessage) -> Box> { 60 | let service = self.inner.clone(); 61 | let resp = service.inner.call(Message::WithoutBody(request)) 62 | .map_err(|e| { 63 | e.into()} 64 | ) 65 | .and_then(|resp| { 66 | if resp != "OK".into() { 67 | Err(io::Error::new(io::ErrorKind::Other, "expected OK")) 68 | } else { 69 | Ok(resp) 70 | } 71 | }); 72 | 73 | Box::new(resp) 74 | } 75 | } -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | [![Build Status](https://travis-ci.org/wisespace-io/nsqueue.png?branch=master)](https://travis-ci.org/wisespace-io/nsqueue) 2 | [![Crates.io](https://img.shields.io/crates/v/nsqueue.svg)](https://crates.io/crates/nsqueue) 3 | 4 | # nsqueue 5 | A [Tokio](https://tokio.rs/) based client implementation for the [NSQ](https://github.com/bitly/nsq) realtime message processing system 6 | 7 | ## WORK IN PROGRESS 8 | 9 | ### Current features 10 | - [X] PUB 11 | - [X] SUB 12 | - [ ] Discovery 13 | - [ ] Backoff 14 | - [ ] TLS 15 | - [ ] Snappy 16 | - [ ] Auth 17 | 18 | ### Launch NSQ 19 | ``` 20 | $ ./nsqlookupd & 21 | $ ./nsqd --lookupd-tcp-address=127.0.0.1:4160 & 22 | $ ./nsqadmin --lookupd-http-address=127.0.0.1:4161 & 23 | ``` 24 | 25 | ### MPUB 26 | ``` 27 | extern crate futures; 28 | extern crate tokio_core; 29 | extern crate nsqueue; 30 | 31 | use futures::Future; 32 | use tokio_core::reactor::Core; 33 | 34 | use nsqueue::config::*; 35 | use nsqueue::producer::*; 36 | 37 | fn main() { 38 | let mut core = Core::new().unwrap(); 39 | let handle = core.handle(); 40 | let addr = "127.0.0.1:4150".parse().unwrap(); 41 | 42 | let mut messages: Vec = Vec::new(); 43 | messages.push("First message".into()); 44 | messages.push("Second message".into()); 45 | 46 | let res = Producer::connect(&addr, &handle, Config::default()) 47 | .and_then(|conn| { 48 | conn.mpublish("some_topic".into(), messages) 49 | .and_then(move |response| { 50 | println!("Response: {:?}", response); 51 | Ok(()) 52 | }) 53 | }); 54 | core.run(res).unwrap(); 55 | } 56 | ``` 57 | 58 | ### SUB 59 | ``` 60 | extern crate futures; 61 | extern crate tokio_core; 62 | extern crate nsqueue; 63 | 64 | use futures::{Stream, Future}; 65 | use tokio_core::reactor::Core; 66 | use nsqueue::config::*; 67 | use nsqueue::consumer::*; 68 | 69 | fn main() { 70 | let mut core = Core::new().unwrap(); 71 | let handle = core.handle(); 72 | let addr = "127.0.0.1:4150".parse().unwrap(); 73 | 74 | core.run( 75 | Consumer::connect(&addr, &handle, Config::default()) 76 | .and_then(|conn| { 77 | conn.subscribe("some_topic".into(), "some_channel".into()) 78 | .and_then(move |response| { 79 | let ret = response.for_each(move |message| { 80 | if message.message_id == "_heartbeat_" { 81 | conn.nop(); 82 | } else { 83 | println!("Response {:?} {:?}", message.message_id, message.message_body); 84 | conn.fin(message.message_id); // Inform NSQ (Message consumed) 85 | } 86 | Ok(()) 87 | }); 88 | ret 89 | }) 90 | }) 91 | ).unwrap(); 92 | } 93 | ``` 94 | 95 | ## License 96 | 97 | Licensed under either of 98 | 99 | * MIT license (see [LICENSE](LICENSE) or ) 100 | * Apache License, Version 2.0 (see [LICENSE](LICENSE) or ) 101 | -------------------------------------------------------------------------------- /src/consumer.rs: -------------------------------------------------------------------------------- 1 | use futures::{Future, future}; 2 | 3 | use tokio_service::Service; 4 | use tokio_core::reactor::Handle; 5 | use tokio_proto::{TcpClient}; 6 | use tokio_proto::util::client_proxy::ClientProxy; 7 | use tokio_proto::streaming::{Message}; 8 | 9 | use std::io; 10 | use std::net::SocketAddr; 11 | 12 | use config::Config; 13 | use response::{ResponseStream}; 14 | use codec::{NsqMessage, NsqResponseMessage, ClientTypeMap}; 15 | use protocol::{NsqProtocol, RequestMessage}; 16 | 17 | #[derive(Clone)] 18 | pub struct Consumer { 19 | inner: ClientTypeMap>, 20 | } 21 | 22 | impl Consumer { 23 | /// Establish a connection and send protocol version. 24 | pub fn connect(addr: &SocketAddr, handle: &Handle, config: Config) -> Box> { 25 | let protocol = NsqProtocol::new(config); 26 | let ret = TcpClient::new(protocol) 27 | .connect(addr, handle) 28 | .map(|client_proxy| { 29 | let type_map = ClientTypeMap { inner: client_proxy }; 30 | Consumer { inner: type_map } 31 | }); 32 | 33 | Box::new(ret) 34 | } 35 | 36 | #[allow(unused_variables)] 37 | pub fn subscribe(&self, topic: String, channel: String) -> Box> { 38 | let mut request = RequestMessage::new(); 39 | request.create_sub_command(topic, channel); 40 | 41 | let service = self.inner.clone(); 42 | let resp = service.inner.call(Message::WithoutBody(request)) 43 | .map_err(|e| {e.into()}) 44 | .and_then(move |resp| { 45 | let mut request = RequestMessage::new(); 46 | request.create_rdy_command(); 47 | let rdy = service.inner.call(Message::WithoutBody(request)) 48 | .map_err(|e| {e.into()}); 49 | rdy 50 | }) 51 | .map(move |resp| { 52 | match resp { 53 | Message::WithoutBody(str) => { 54 | panic!("Not implemented: {}", str) 55 | }, 56 | Message::WithBody(head, body) => { 57 | ResponseStream { inner: body } 58 | } 59 | } 60 | }); 61 | 62 | Box::new(resp) 63 | } 64 | 65 | #[allow(unused_variables)] 66 | pub fn fin(&self, message_id: String) -> Box> { 67 | let mut request = RequestMessage::new(); 68 | request.create_fin_command(message_id); 69 | 70 | let service = self.inner.clone(); 71 | let resp = service.inner.call(Message::WithoutBody(request)) 72 | .map_err(|e| e.into()) 73 | .and_then(|resp| future::ok(())); 74 | 75 | Box::new(resp) 76 | } 77 | 78 | #[allow(unused_variables)] 79 | pub fn nop(&self) -> Box> { 80 | let mut request = RequestMessage::new(); 81 | request.create_nop_command(); 82 | 83 | let service = self.inner.clone(); 84 | let resp = service.inner.call(Message::WithoutBody(request)) 85 | .map_err(|e| e.into()) 86 | .and_then(|resp| future::ok(())); 87 | 88 | Box::new(resp) 89 | } 90 | } 91 | 92 | impl Service for ClientTypeMap 93 | where T: Service, 94 | T::Future: 'static 95 | { 96 | type Request = RequestMessage; 97 | type Response = NsqResponseMessage; 98 | type Error = io::Error; 99 | type Future = Box>; 100 | 101 | fn call(&self, req: RequestMessage) -> Self::Future { 102 | Box::new(self.inner.call(req)) 103 | } 104 | } -------------------------------------------------------------------------------- /src/protocol.rs: -------------------------------------------------------------------------------- 1 | use std::io; 2 | 3 | use tokio_io::{AsyncRead, AsyncWrite}; 4 | use tokio_io::codec::{Framed}; 5 | use tokio_proto::streaming::pipeline::{Frame, ClientProto}; 6 | 7 | use serde_json::{to_string}; 8 | 9 | use futures::{Future, Stream, Sink}; 10 | 11 | use commands::*; 12 | use response::Message; 13 | use codec::NsqCodec; 14 | use config::Config; 15 | 16 | /// Protocol definition 17 | pub struct NsqProtocol { 18 | pub config: Config, 19 | } 20 | 21 | impl NsqProtocol { 22 | pub fn new(config: Config) -> Self { 23 | NsqProtocol { 24 | config: config, 25 | } 26 | } 27 | } 28 | 29 | #[allow(unused_variables)] 30 | impl ClientProto for NsqProtocol { 31 | type Request = RequestMessage; 32 | type RequestBody = RequestMessage; 33 | type Response = String; 34 | type ResponseBody = Message; 35 | 36 | type Error = io::Error; 37 | type Transport = Framed; 38 | type BindTransport = Box>; 39 | 40 | fn bind_transport(&self, io: T) -> Self::BindTransport { 41 | let config = self.config.clone(); 42 | let mut request = RequestMessage::new(); 43 | request.set_protocol_version(commands::VERSION_2); 44 | 45 | let codec = NsqCodec { 46 | decoding_head: true, 47 | }; 48 | 49 | // Send protocol version 50 | let tst = request.clone(); 51 | let version = Frame::Message {message: tst, body: false }; 52 | let handshake = io.framed(codec).send(version) 53 | .and_then(move |transport| { 54 | let mut request = RequestMessage::new(); 55 | request.create_identify_command(config); 56 | 57 | // Send IDENTIFY 58 | let identify = Frame::Message {message: request.clone(), body: false }; 59 | let ch = transport.send(identify) 60 | .and_then(|transport| transport.into_future().map_err(|(e, _)| e)) 61 | .and_then(|(resp, transport)| { 62 | Ok(transport) 63 | }); 64 | ch 65 | }); 66 | 67 | Box::new(handshake) 68 | } 69 | } 70 | 71 | #[derive(PartialEq, Debug, Clone)] 72 | pub struct RequestMessage { 73 | pub version: Option, 74 | pub header: Option, 75 | pub body: Option, 76 | pub body_messages: Option>, 77 | } 78 | 79 | impl RequestMessage { 80 | pub fn new() -> RequestMessage { 81 | RequestMessage { 82 | version: None, 83 | header: None, 84 | body: None, 85 | body_messages: None, 86 | } 87 | } 88 | 89 | pub fn set_protocol_version(&mut self, version: &str) { 90 | self.version = Some(String::from(version)); 91 | } 92 | 93 | pub fn create_pub_command(&mut self, topic: String, message: String) { 94 | self.header = Some(format!("{} {}\n", commands::PUB, topic)); 95 | self.body = Some(message); 96 | } 97 | 98 | pub fn create_mpub_command(&mut self, topic: String, messages: Vec) { 99 | self.header = Some(format!("{} {}\n", commands::MPUB, topic)); 100 | self.body_messages = Some(messages); 101 | } 102 | 103 | pub fn create_dpub_command(&mut self, topic: String, message: String, defer_time: i64) { 104 | self.header = Some(format!("{} {} {}\n", commands::DPUB, topic, defer_time.to_string())); 105 | self.body = Some(message); 106 | } 107 | 108 | pub fn create_sub_command(&mut self, topic: String, channel: String) { 109 | self.header = Some(format!("{} {} {}\n", commands::SUB, topic, channel)); 110 | } 111 | 112 | pub fn create_identify_command(&mut self, config: Config) { 113 | self.header = Some(format!("{}\n", commands::IDENTIFY)); 114 | // Serialize it to a JSON string. 115 | self.body = Some(to_string(&config).unwrap()); 116 | } 117 | 118 | pub fn create_rdy_command(&mut self) { 119 | self.header = Some(format!("{} 1\n", commands::RDY)); 120 | } 121 | 122 | pub fn create_fin_command(&mut self, message_id: String) { 123 | self.header = Some(format!("{} {}\n", commands::FIN, message_id)); 124 | } 125 | 126 | pub fn create_nop_command(&mut self) { 127 | self.header = Some(format!("{}\n", commands::NOP)); 128 | } 129 | } -------------------------------------------------------------------------------- /src/codec.rs: -------------------------------------------------------------------------------- 1 | use std::io; 2 | use std::io::Cursor; 3 | use std::iter::Iterator; 4 | 5 | use bytes::{Buf, BufMut, BytesMut, BigEndian}; 6 | use tokio_io::codec::{Encoder, Decoder}; 7 | use tokio_proto::streaming::pipeline::Frame; 8 | use tokio_proto::streaming::{Body, Message}; 9 | use std::str; 10 | 11 | use protocol::RequestMessage; 12 | use response::Message as TypeMessage; 13 | 14 | // Header: Size(4-Byte) + FrameType(4-Byte) 15 | const HEADER_LENGTH: usize = 8; 16 | 17 | // Frame Types 18 | const FRAME_TYPE_RESPONSE: i32 = 0x00; 19 | const FRAME_TYPE_ERROR: i32 = 0x01; 20 | const FRAME_TYPE_MESSAGE: i32 = 0x02; 21 | 22 | const HEARTBEAT: &'static str = "_heartbeat_"; 23 | 24 | #[derive(Clone)] 25 | pub struct ClientTypeMap { 26 | pub inner: T, 27 | } 28 | 29 | pub type NsqMessage = Message>; 30 | pub type NsqResponseMessage = Message>; 31 | 32 | /// NSQ codec 33 | pub struct NsqCodec { 34 | pub decoding_head: bool 35 | } 36 | 37 | impl Decoder for NsqCodec { 38 | type Item = Frame; 39 | type Error = io::Error; 40 | 41 | fn decode(&mut self, buf: &mut BytesMut) -> Result, io::Error> { 42 | let length = buf.len(); 43 | 44 | if length < HEADER_LENGTH { 45 | return Ok(None); 46 | } 47 | 48 | let mut cursor = Cursor::new(buf.clone()); 49 | let size: i32 = cursor.get_i32::(); 50 | 51 | if length < size as usize { 52 | return Ok(None); 53 | } 54 | 55 | let frame_type: i32 = cursor.get_i32::(); 56 | 57 | if frame_type == FRAME_TYPE_RESPONSE { 58 | // remove the serialized frame from the buffer. 59 | buf.split_to(HEADER_LENGTH + length); 60 | match str::from_utf8(&cursor.bytes()) { 61 | Ok(s) => { 62 | let decoded_message = s.to_string(); 63 | 64 | // TODO: Implement a proper way to handle the heartbeat 65 | if decoded_message == HEARTBEAT && !self.decoding_head { 66 | Ok(Some(self.heartbeat_message())) 67 | } else if decoded_message == HEARTBEAT { 68 | // toggle streaming 69 | Ok(Some(self.streaming_flag())) 70 | } else { 71 | Ok(Some(Frame::Message { 72 | message: decoded_message, 73 | body: false, 74 | })) 75 | } 76 | } 77 | Err(_) => Err(io::Error::new(io::ErrorKind::Other, "Invalid UTF-8")), 78 | } 79 | } else if frame_type == FRAME_TYPE_ERROR { 80 | Err(io::Error::new(io::ErrorKind::InvalidData, "Invalid packet received")) 81 | } else if frame_type == FRAME_TYPE_MESSAGE { 82 | if self.decoding_head { 83 | // toggle streaming 84 | Ok(Some(self.streaming_flag())) 85 | } else { 86 | let timestamp = cursor.get_i64::(); // timestamp 87 | let _ = cursor.get_u16::(); // attempts 88 | 89 | let data = str::from_utf8(&cursor.bytes()).unwrap().to_string(); 90 | let (id, body) = data.split_at(16); 91 | 92 | let message = TypeMessage{ 93 | timestamp: timestamp, 94 | message_id: id.to_string(), 95 | message_body: body.to_string() 96 | }; 97 | 98 | // remove the serialized frame from the buffer. 99 | buf.split_to(HEADER_LENGTH + length); 100 | 101 | Ok(Some( 102 | Frame::Body { 103 | chunk: Some(message), 104 | } 105 | )) 106 | } 107 | } else { 108 | Ok(None) 109 | } 110 | } 111 | } 112 | 113 | pub type CodecOutputFrame = Frame; 114 | impl Encoder for NsqCodec { 115 | type Item = CodecOutputFrame; 116 | type Error = io::Error; 117 | 118 | fn encode(&mut self, message: Self::Item, buf: &mut BytesMut) -> io::Result<()> { 119 | match message { 120 | Frame::Message { message, .. } => { 121 | if let Some(version) = message.version { 122 | buf.reserve(version.len()); 123 | buf.extend(version.as_bytes()); 124 | } 125 | 126 | if let Some(header) = message.header { 127 | buf.reserve(header.len() + 1); 128 | buf.extend(header.as_bytes()); 129 | } 130 | 131 | if let Some(body) = message.body { 132 | let mut buf_32 = Vec::with_capacity(body.len()); 133 | let body_len = body.len() as u32; 134 | buf_32.put_u32::(body_len); 135 | buf_32.put(&body[..]); 136 | buf.extend(buf_32); 137 | } 138 | 139 | if let Some(body_messages) = message.body_messages { 140 | let total_bytes = body_messages 141 | .iter() 142 | .map(|message| message.len()) 143 | .fold(0, |acc, len| acc + len); 144 | 145 | let mut buf_32 = Vec::with_capacity(total_bytes); 146 | // [4-byte body size] 147 | let body_len = total_bytes as u32; 148 | buf_32.put_u32::(body_len); 149 | // [4-byte num messages] 150 | let messages_len = body_messages.len() as u32; 151 | buf_32.put_u32::(messages_len); 152 | // [ 4-byte message #1 size ][ N-byte binary data ] ... 153 | for message in &body_messages { 154 | let message_len = message.len() as u32; 155 | buf_32.put_u32::(message_len); 156 | buf_32.put(&message[..]); 157 | } 158 | buf.extend(buf_32); 159 | } 160 | Ok(()) 161 | } 162 | Frame::Error { error, .. } => Err(error), 163 | Frame::Body { .. } => panic!("Streaming of Requests is not currently supported"), 164 | } 165 | } 166 | } 167 | 168 | impl NsqCodec { 169 | fn heartbeat_message(&mut self) -> Frame 170 | { 171 | let message = TypeMessage{ 172 | timestamp: 0, 173 | message_id: HEARTBEAT.to_string(), 174 | message_body: HEARTBEAT.to_string() 175 | }; 176 | 177 | Frame::Body { 178 | chunk: Some(message), 179 | } 180 | } 181 | 182 | fn streaming_flag(&mut self) -> Frame 183 | { 184 | self.decoding_head = false; 185 | Frame::Message { 186 | message: "".into(), 187 | body: true, 188 | } 189 | } 190 | } --------------------------------------------------------------------------------