feat(new source): add initial websocket source

src/common/backoff.rs

@@ -0,0 +1,81 @@
+use std::time::Duration;
+
+// `tokio-retry` crate
+// MIT License
+// Copyright (c) 2017 Sam Rijs
+//
+/// A retry strategy driven by exponential back-off.
+///
+/// The power corresponds to the number of past attempts.
+#[derive(Debug, Clone)]
+pub struct ExponentialBackoff {
+    current: u64,
+    base: u64,
+    factor: u64,
+    max_delay: Option<Duration>,
+}
+
+impl ExponentialBackoff {
+    /// Constructs a new exponential back-off strategy,
+    /// given a base duration in milliseconds.
+    ///
+    /// The resulting duration is calculated by taking the base to the `n`-th power,
+    /// where `n` denotes the number of past attempts.
+    pub const fn from_millis(base: u64) -> ExponentialBackoff {
+        ExponentialBackoff {
+            current: base,
+            base,
+            factor: 1u64,
+            max_delay: None,
+        }
+    }
+
+    /// A multiplicative factor that will be applied to the retry delay.
+    ///
+    /// For example, using a factor of `1000` will make each delay in units of seconds.
+    ///
+    /// Default factor is `1`.
+    pub const fn factor(mut self, factor: u64) -> ExponentialBackoff {
+        self.factor = factor;
+        self
+    }
+
+    /// Apply a maximum delay. No retry delay will be longer than this `Duration`.
+    pub const fn max_delay(mut self, duration: Duration) -> ExponentialBackoff {
+        self.max_delay = Some(duration);
+        self
+    }
+
+    /// Resents the exponential back-off strategy to its initial state.
+    pub fn reset(&mut self) {
+        self.current = self.base;
+    }
+}
+
+impl Iterator for ExponentialBackoff {
+    type Item = Duration;
+
+    fn next(&mut self) -> Option<Duration> {
+        // set delay duration by applying factor
+        let duration = if let Some(duration) = self.current.checked_mul(self.factor) {
+            Duration::from_millis(duration)
+        } else {
+            Duration::from_millis(std::u64::MAX)
+        };
+
+        // check if we reached max delay
+        if let Some(ref max_delay) = self.max_delay {
+            if duration > *max_delay {
+                return Some(*max_delay);
+            }
+        }
+
+        if let Some(next) = self.current.checked_mul(self.base) {
+            self.current = next;
+        } else {
+            self.current = std::u64::MAX;
+        }
+
+        Some(duration)
+    }
+}

src/common/mod.rs

@@ -17,3 +17,8 @@ pub(crate) mod sqs;
 
 #[cfg(any(feature = "sources-aws_s3", feature = "sinks-aws_s3"))]
 pub(crate) mod s3;
+
+pub mod websocket;
+
+pub(crate) mod backoff;
+pub(crate) mod ping;

src/common/ping.rs

@@ -0,0 +1,29 @@
+use std::{
+    task::{Context, Poll},
+    time::Duration,
+};
+
+use tokio::time;
+
+pub struct PingInterval {
+    interval: Option<time::Interval>,
+}
+
+impl PingInterval {
+    pub fn new(period: Option<u64>) -> Self {
+        Self {
+            interval: period.map(|period| time::interval(Duration::from_secs(period))),
+        }
+    }
+
+    pub fn poll_tick(&mut self, cx: &mut Context<'_>) -> Poll<time::Instant> {
+        match self.interval.as_mut() {
+            Some(interval) => interval.poll_tick(cx),
+            None => Poll::Pending,
+        }
+    }
+
+    pub async fn tick(&mut self) -> time::Instant {
+        std::future::poll_fn(|cx| self.poll_tick(cx)).await
+    }
+}

src/common/websocket.rs

@@ -0,0 +1,154 @@
+use std::{fmt::Debug, net::SocketAddr, time::Duration};
+
+use snafu::{ResultExt, Snafu};
+use tokio::{net::TcpStream, time};
+use tokio_tungstenite::{
+    client_async_with_config,
+    tungstenite::{
+        client::{uri_mode, IntoClientRequest},
+        error::{Error as WsError, ProtocolError, UrlError},
+        handshake::client::Request as WsRequest,
+        protocol::WebSocketConfig,
+        stream::Mode as UriMode,
+    },
+    WebSocketStream as WsStream,
+};
+
+use crate::{
+    common::backoff::ExponentialBackoff,
+    dns,
+    http::Auth,
+    internal_events::{WsConnectionEstablished, WsConnectionFailedError},
+    tls::{MaybeTlsSettings, MaybeTlsStream, TlsError},
+};
+
+#[derive(Debug, Snafu)]
+#[snafu(visibility(pub))]
+pub enum WebSocketError {
+    #[snafu(display("Creating WebSocket client failed: {}", source))]
+    CreateFailed { source: WsError },
+    #[snafu(display("Connect error: {}", source))]
+    ConnectError { source: TlsError },
+    #[snafu(display("Unable to resolve DNS: {}", source))]
+    DnsError { source: dns::DnsError },
+    #[snafu(display("No addresses returned."))]
+    NoAddresses,
+}
+
+#[derive(Clone)]
+pub(crate) struct WebSocketConnector {
+    uri: String,
+    host: String,
+    port: u16,
+    tls: MaybeTlsSettings,
+    auth: Option<Auth>,
+}
+
+impl WebSocketConnector {
+    pub(crate) fn new(
+        uri: String,
+        tls: MaybeTlsSettings,
+        auth: Option<Auth>,
+    ) -> Result<Self, WebSocketError> {
+        let request = (&uri).into_client_request().context(CreateFailedSnafu)?;
+        let (host, port) = Self::extract_host_and_port(&request).context(CreateFailedSnafu)?;
+
+        Ok(Self {
+            uri,
+            host,
+            port,
+            tls,
+            auth,
+        })
+    }
+
+    fn extract_host_and_port(request: &WsRequest) -> Result<(String, u16), WsError> {
+        let host = request
+            .uri()
+            .host()
+            .ok_or(WsError::Url(UrlError::NoHostName))?
+            .to_string();
+        let mode = uri_mode(request.uri())?;
+        let port = request.uri().port_u16().unwrap_or(match mode {
+            UriMode::Tls => 443,
+            UriMode::Plain => 80,
+        });
+
+        Ok((host, port))
+    }
+
+    const fn fresh_backoff() -> ExponentialBackoff {
+        ExponentialBackoff::from_millis(2)
+            .factor(250)
+            .max_delay(Duration::from_secs(60))
+    }
+
+    async fn tls_connect(&self) -> Result<MaybeTlsStream<TcpStream>, WebSocketError> {
+        let ip = dns::Resolver
+            .lookup_ip(self.host.clone())
+            .await
+            .context(DnsSnafu)?
+            .next()
+            .ok_or(WebSocketError::NoAddresses)?;
+
+        let addr = SocketAddr::new(ip, self.port);
+        self.tls
+            .connect(&self.host, &addr)
+            .await
+            .context(ConnectSnafu)
+    }
+
+    async fn connect(&self) -> Result<WsStream<MaybeTlsStream<TcpStream>>, WebSocketError> {
+        let mut request = (&self.uri)
+            .into_client_request()
+            .context(CreateFailedSnafu)?;
+
+        if let Some(auth) = &self.auth {
+            auth.apply(&mut request);
+        }
+
+        let maybe_tls = self.tls_connect().await?;
+
+        let ws_config = WebSocketConfig {
+            max_send_queue: None, // don't buffer messages
+            ..Default::default()
+        };
+
+        let (ws_stream, _response) = client_async_with_config(request, maybe_tls, Some(ws_config))
+            .await
+            .context(CreateFailedSnafu)?;
+
+        Ok(ws_stream)
+    }
+
+    pub(crate) async fn connect_backoff(&self) -> WsStream<MaybeTlsStream<TcpStream>> {
+        let mut backoff = Self::fresh_backoff();
+        loop {
+            match self.connect().await {
+                Ok(ws_stream) => {
+                    emit!(WsConnectionEstablished {});
+                    return ws_stream;
+                }
+                Err(error) => {
+                    emit!(WsConnectionFailedError {
+                        error: Box::new(error)
+                    });
+                    time::sleep(backoff.next().unwrap()).await;
+                }
+            }
+        }
+    }
+
+    pub(crate) async fn healthcheck(&self) -> crate::Result<()> {
+        self.connect().await.map(|_| ()).map_err(Into::into)
+    }
+}
+
+pub(crate) const fn is_closed(error: &WsError) -> bool {
+    matches!(
+        error,
+        WsError::ConnectionClosed
+            | WsError::AlreadyClosed
+            | WsError::Protocol(ProtocolError::ResetWithoutClosingHandshake)
+    )
+}

src/config/enterprise.rs

@@ -21,6 +21,7 @@ use super::{
     SourceOuter, TransformOuter,
 };
 use crate::{
+    common::backoff::ExponentialBackoff,
     common::datadog::{get_api_base_endpoint, get_base_domain_region, Region},
     conditions::AnyCondition,
     http::{HttpClient, HttpError},
@@ -29,7 +30,7 @@ use crate::{
             default_site, logs::DatadogLogsConfig, metrics::DatadogMetricsConfig,
             DatadogCommonConfig,
         },
-        util::{http::RequestConfig, retries::ExponentialBackoff},
+        util::http::RequestConfig,
     },
     sources::{
         host_metrics::{Collector, HostMetricsConfig},

src/internal_events/websocket.rs

@@ -89,3 +89,24 @@ impl InternalEvent for WsConnectionError {
         Some("WsConnectionError")
     }
 }
+
+pub struct WsMessageReceived {
+    pub url: String,
+}
+
+impl InternalEvent for WsMessageReceived {
+    fn emit(self) {
+        trace!(
+            message = "Event received.",
+            url =  %self.url,
+        );
+        counter!(
+            "component_received_events_total", 1,
+            "uri" => self.url.clone(),
+        );
+    }
+
+    fn name(&self) -> Option<&'static str> {
+        Some("WsMessageReceived")
+    }
+}

src/lib.rs

@@ -66,6 +66,7 @@ pub mod async_read;
 pub mod aws;
 #[allow(unreachable_pub)]
 pub mod codecs;
+#[allow(unreachable_pub)]
 pub(crate) mod common;
 pub mod encoding_transcode;
 pub mod enrichment_tables;