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
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126

#[cfg(test)]
#[macro_use]
extern crate serde_derive;

use futures::Future;
use pulsar::{Connection, Producer};

pub struct ConnectionManager {
    addr: String,
    executor: tokio::runtime::TaskExecutor
}

impl r2d2::ManageConnection for ConnectionManager {
    type Connection = Connection;
    type Error = pulsar::ConnectionError;

    fn connect(&self) -> Result<Self::Connection, Self::Error> {
        Future::wait(Connection::new(self.addr.clone(), None, None, self.executor.clone()))
    }

    fn is_valid(&self, conn: &mut Self::Connection) -> Result<(), Self::Error> {
        Future::wait(conn.sender().lookup_topic(String::from("test"), false))
            .map(|_| ())
    }

    fn has_broken(&self, conn: &mut Self::Connection) -> bool {
        !conn.is_valid()
    }
}

pub struct ProducerConnectionManager {
    addr: String,
    executor: tokio::runtime::TaskExecutor,
}

impl ProducerConnectionManager {
    pub fn new<S: Into<String>>(addr: S, executor: tokio::runtime::TaskExecutor) -> ProducerConnectionManager {
        ProducerConnectionManager {
            addr: addr.into(),
            executor,
        }
    }
}

impl r2d2::ManageConnection for ProducerConnectionManager {
    type Connection = Producer;
    type Error = pulsar::ProducerError;

    fn connect(&self) -> Result<Producer, Self::Error> {
        Producer::new(self.addr.as_str(), None, None, None, self.executor.clone())
            .wait()
    }

    fn is_valid(&self, conn: &mut Producer) -> Result<(), Self::Error> {
        conn.check_connection().wait().map_err(|e| pulsar::ProducerError::Connection(e))
    }

    fn has_broken(&self, conn: &mut Producer) -> bool {
        !conn.is_valid()
    }
}


#[cfg(test)]
mod tests {
    use super::*;
    use r2d2;
    use futures::Stream;
    use pulsar::{Consumer, ConsumerBuilder, SubType, ConsumerError, ConnectionError};

    #[derive(Debug, Clone, Serialize, Deserialize)]
    struct TestData {
        pub data: String
    }

    #[test]
    #[ignore]
    fn it_works() {
        let addr = "127.0.0.1:6650";
        let runtime = tokio::runtime::Runtime::new().unwrap();
        let pool = r2d2::Pool::new(ProducerConnectionManager::new(
            addr,
            runtime.executor()
        )).unwrap();

        let mut a = pool.get().unwrap();
        let mut b = pool.get().unwrap();

        let dataz = TestData { data: String::from("dataz") };

        let send_1 = a.send_json("r2d2_test", &dataz.clone());
        let send_2 = b.send_json("r2d2_test", &dataz.clone());

        send_1.join(send_2).wait().unwrap();

        let consumer: Consumer<TestData> = ConsumerBuilder::new(addr, runtime.executor())
            .with_subscription("r2d2_test")
            .with_topic("r2d2_test")
            .with_subscription_type(SubType::Exclusive)
            .build()
            .wait()
            .unwrap();

        let mut consumed_messages = 0;
        let consumed = consumer.for_each(move |(msg, ack)| {
            consumed_messages += 1;
            ack.ack();
            if let Err(e) = msg {
                println!("Consumer error: {}", e);
            }
            if consumed_messages == 2 {
                //err here to shutdown
                Err(ConsumerError::Connection(ConnectionError::Unexpected("Done!".to_owned())))
            } else {
                Ok(())
            }
        }).wait();

        match &consumed {
            Err(ConsumerError::Connection(ConnectionError::Unexpected(msg))) if msg.as_str() == "Done!" => {},
            other => panic!("Unexpected consumer shutdown. Found: {:?}", other)
        }

        runtime.shutdown_now();
    }
}