Skip to main content

slog_telegraf/
telegraf.rs

1use std::{net, io};
2use url::Url;
3use crate::Error;
4use std::io::Write;
5
6/// Telegraf client
7///
8/// ```no_run
9/// use slog_telegraf::{TelegrafDrainBuilder, Client};
10/// let mut client = Client::new("tcp://127.0.0.1:8094".into()).unwrap();
11/// client.write("measurement,tag=value field=10i".as_bytes()).unwrap();
12/// ```
13pub struct Client {
14    connection: Connection
15}
16
17impl Client {
18    pub fn new(url: String) -> Result<Self, Error> {
19        Ok(Client{
20            connection: Connection::new(url)?
21        })
22    }
23
24    pub fn write(&mut self, bytes:&[u8]) -> io::Result<()> {
25        self.connection.write(bytes).map(|_| ())
26    }
27}
28
29enum Connection {
30    Tcp(net::TcpStream),
31    Udp(net::UdpSocket)
32}
33
34impl Connection {
35    pub fn new(url: String) -> Result<Self, Error> {
36        let url = Url::parse(&url)?;
37        let addr = url.socket_addrs(|| None)?;
38
39        match url.scheme() {
40            "tcp" => Ok(Connection::Tcp(net::TcpStream::connect(&*addr)?)),
41            "udp" => {
42                // This will let the OS choose the ip+port
43                let socket = net::UdpSocket::bind(&[net::SocketAddr::from(([0, 0, 0, 0], 0))][..])?;
44                socket.connect(&*addr)?;
45                socket.set_nonblocking(true)?;
46
47                Ok(Connection::Udp(socket))
48            },
49            "" => Err(Error::Custom("Please specify the protocol 'tcp' or 'udp'".to_string())),
50            _ => Err(Error::Custom("Only 'tcp' and 'udp' is currently supported".to_string()))
51        }
52    }
53
54    pub fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
55        match self {
56            Connection::Tcp(tcp_stream) => tcp_stream.write(bytes),
57            Connection::Udp(udp_socket) => udp_socket.send(bytes)
58        }
59    }
60}
61
62#[cfg(test)]
63mod test {
64    use super::*;
65
66    #[test]
67    fn test_connection_new() {
68        assert!(Connection::new("udp://127.0.0.1:12345".into()).is_ok());
69
70        assert!(Connection::new("127.0.0.1:12345".into()).is_err());
71        assert!(Connection::new("http://127.0.0.1:12345".into()).is_err());
72    }
73}