rustmann 0.2.2

Rust/Tokio client for Riemann
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96

use std::io;
use std::net::SocketAddr;
use std::time::Duration;

use futures_legacy::future;
use futures_legacy::stream::SplitSink;
use futures_legacy::sync::mpsc::{self, UnboundedSender};
use futures_legacy::sync::oneshot::{self, Sender};
use futures_legacy::{Future, Sink, Stream};
use protobuf::RepeatedField;
use tokio::codec::Framed;
use tokio::net::TcpStream;
use tokio::prelude::*;

use crate::codec::MsgCodec;
use crate::protos::riemann::{Event, Msg};

#[derive(Debug)]
pub struct Connection {
    sender_queue: UnboundedSender<Sender<Msg>>,
    socket_sender: SplitSink<Framed<TcpStream, MsgCodec>>,
    // callback_queue: VecDeque<Sender<Msg>>,
}

impl Connection {
    pub(crate) fn connect(
        addr: &SocketAddr,
        connect_timeout_ms: u64,
    ) -> impl Future<Item = Connection, Error = io::Error> {
        TcpStream::connect(addr)
            .timeout(Duration::from_millis(connect_timeout_ms))
            .map_err(|e| {
                if e.is_timer() {
                    io::Error::new(io::ErrorKind::TimedOut, e)
                } else {
                    e.into_inner().unwrap()
                }
            })
            .and_then(|socket| {
                socket.set_nodelay(true)?;

                let framed = Framed::new(socket, MsgCodec);
                let (conn_sender, conn_receiver) = framed.split();

                let (cb_queue_tx, cb_queue_rx) = mpsc::unbounded::<Sender<Msg>>();

                let receiver_loop =
                    conn_receiver
                        .zip(cb_queue_rx.map_err(|_| {
                            io::Error::new(io::ErrorKind::Other, "callback queue error")
                        }))
                        .for_each(move |(frame, cb)| {
                            cb.send(frame).map_err(|e| {
                                io::Error::new(
                                    io::ErrorKind::Other,
                                    format!("failed to deliver msg to callback {:?}", e),
                                )
                            })
                        })
                        .map_err(|e| eprintln!("{:?}", e));
                tokio::spawn(receiver_loop);

                Ok(Connection {
                    sender_queue: cb_queue_tx,
                    socket_sender: conn_sender,
                })
            })
    }

    pub(crate) fn send_events(
        &mut self,
        events: &[Event],
        socket_timeout: u64,
    ) -> impl Future<Item = Msg, Error = io::Error> {
        let mut msg = Msg::new();
        msg.set_events(RepeatedField::from_slice(events));

        let (tx, rx) = oneshot::channel::<Msg>();

        future::result(
            self.sender_queue
                .unbounded_send(tx)
                .map_err(|e| io::Error::new(io::ErrorKind::Other, e))
                .and_then(|_| {
                    self.socket_sender
                        .start_send(msg)
                        .and_then(|_| self.socket_sender.poll_complete())
                }),
        )
        .map_err(|e| io::Error::new(io::ErrorKind::UnexpectedEof, e))
        .and_then(move |_| {
            rx.timeout(Duration::from_millis(socket_timeout))
                .map_err(|e| io::Error::new(io::ErrorKind::TimedOut, e))
        })
    }
}