diem_logger/
struct_log.rs

1// Copyright (c) The Diem Core Contributors
2// SPDX-License-Identifier: Apache-2.0
3
4// Copyright 2021 Conflux Foundation. All rights reserved.
5// Conflux is free software and distributed under GNU General Public License.
6// See http://www.gnu.org/licenses/
7
8//! Implementations for sending logs to external log processes e.g. Logstash
9//!
10//! Handles sending logs under disconnects, and retries.  Tries to continue to
11//! make progress on a log but eventually drops older logs to continue to make
12//! progress on newer logs.
13
14use crate::counters::{
15    STRUCT_LOG_CONNECT_ERROR_COUNT, STRUCT_LOG_TCP_CONNECT_COUNT,
16};
17use std::{
18    io::{self, Write},
19    net::{TcpStream, ToSocketAddrs},
20    time::Duration,
21};
22
23const WRITE_TIMEOUT_MS: u64 = 2000;
24const CONNECTION_TIMEOUT_MS: u64 = 5000;
25
26/// A wrapper for `TcpStream` that handles reconnecting to the endpoint
27/// automatically
28///
29/// `TcpWriter::write()` will block on the message until it is connected.
30pub(crate) struct TcpWriter {
31    /// The DNS name or IP address logs are being sent to
32    endpoint: String,
33    /// The `TCPStream` to write to, which will be `None` when disconnected
34    stream: Option<TcpStream>,
35}
36
37impl TcpWriter {
38    pub fn new(endpoint: String) -> Self {
39        Self {
40            endpoint,
41            stream: None,
42        }
43    }
44
45    pub fn endpoint(&self) -> &str { &self.endpoint }
46
47    /// Ensure that we get a connection, no matter how long it takes
48    /// This will block until there is a connection
49    fn refresh_connection(&mut self) {
50        loop {
51            match self.connect() {
52                Ok(stream) => {
53                    self.stream = Some(stream);
54                    return;
55                }
56                Err(e) => {
57                    eprintln!("[Logging] Failed to connect: {}", e);
58                    STRUCT_LOG_CONNECT_ERROR_COUNT.inc();
59                }
60            }
61
62            // Sleep a second so this doesn't just spin as fast as possible
63            std::thread::sleep(Duration::from_millis(1000));
64        }
65    }
66
67    /// Connect and ensure the write timeout is set
68    fn connect(&mut self) -> io::Result<TcpStream> {
69        STRUCT_LOG_TCP_CONNECT_COUNT.inc();
70
71        let mut last_error = io::Error::new(
72            io::ErrorKind::Other,
73            format!("Unable to resolve and connect to {}", self.endpoint),
74        );
75
76        // resolve addresses to handle DNS names
77        for addr in self.endpoint.to_socket_addrs()? {
78            match TcpStream::connect_timeout(
79                &addr,
80                Duration::from_millis(CONNECTION_TIMEOUT_MS),
81            ) {
82                Ok(stream) => {
83                    // Set the write timeout
84                    if let Err(err) = stream.set_write_timeout(Some(
85                        Duration::from_millis(WRITE_TIMEOUT_MS),
86                    )) {
87                        STRUCT_LOG_CONNECT_ERROR_COUNT.inc();
88                        eprintln!(
89                            "[Logging] Failed to set write timeout: {}",
90                            err
91                        );
92                        continue;
93                    }
94                    return Ok(stream);
95                }
96                Err(err) => last_error = err,
97            }
98        }
99
100        Err(last_error)
101    }
102}
103
104impl Write for TcpWriter {
105    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
106        // Refresh the connection if it's missing
107        if self.stream.is_none() {
108            self.refresh_connection();
109        }
110
111        // Attempt to write, and if it fails clear underlying stream
112        // This doesn't guarantee a message cut off mid send will work, but it
113        // does guarantee that we will connect first
114        self.stream
115            .as_mut()
116            .ok_or_else(|| {
117                io::Error::new(io::ErrorKind::NotConnected, "No stream")
118            })
119            .and_then(|stream| stream.write(buf))
120            .map_err(|e| {
121                self.stream = None;
122                e
123            })
124    }
125
126    fn flush(&mut self) -> io::Result<()> {
127        if let Some(mut stream) = self.stream.as_ref() {
128            stream.flush()
129        } else {
130            Err(io::Error::new(
131                io::ErrorKind::NotConnected,
132                "Can't flush, not connected",
133            ))
134        }
135    }
136}
diem_logger/struct_log.rs

diem_logger/
struct_log.rs
