-
Notifications
You must be signed in to change notification settings - Fork 79
Add unix domain socket for telemetry server #128
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,71 @@ | ||
//! Network address types that support both TCP and Unix domain sockets. | ||
//! | ||
//! This module provides the [`ListenAddr`] enum, a flexible address type that can represent | ||
//! either TCP socket addresses or Unix domain socket paths. | ||
|
||
#[cfg(feature = "settings")] | ||
use crate::settings::Settings; | ||
#[cfg(any(feature = "telemetry-server", feature = "settings"))] | ||
use serde::Deserialize; | ||
#[cfg(feature = "settings")] | ||
use serde::Serialize; | ||
use std::fmt; | ||
use std::net::{Ipv4Addr, SocketAddr}; | ||
|
||
/// Address that can be either TCP socket or Unix domain socket endpoint | ||
#[derive(Clone, Debug)] | ||
#[cfg_attr( | ||
any(feature = "telemetry-server", feature = "settings"), | ||
derive(Deserialize) | ||
)] | ||
#[cfg_attr(feature = "settings", derive(Serialize))] | ||
#[cfg_attr( | ||
any(feature = "telemetry-server", feature = "settings"), | ||
serde(untagged) | ||
)] | ||
pub enum ListenAddr { | ||
/// TCP network socket address | ||
Tcp(std::net::SocketAddr), | ||
/// Unix domain socket path | ||
#[cfg(unix)] | ||
Unix(std::path::PathBuf), | ||
} | ||
|
||
impl Default for ListenAddr { | ||
fn default() -> Self { | ||
ListenAddr::Tcp((Ipv4Addr::LOCALHOST, 0).into()) | ||
} | ||
} | ||
|
||
#[cfg(feature = "settings")] | ||
impl From<crate::settings::net::SocketAddr> for ListenAddr { | ||
fn from(addr: crate::settings::net::SocketAddr) -> Self { | ||
ListenAddr::Tcp(addr.into()) | ||
} | ||
} | ||
|
||
impl From<SocketAddr> for ListenAddr { | ||
fn from(addr: SocketAddr) -> Self { | ||
ListenAddr::Tcp(addr) | ||
} | ||
} | ||
|
||
#[cfg(unix)] | ||
impl From<std::path::PathBuf> for ListenAddr { | ||
fn from(path: std::path::PathBuf) -> Self { | ||
ListenAddr::Unix(path) | ||
} | ||
} | ||
|
||
impl fmt::Display for ListenAddr { | ||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { | ||
match self { | ||
ListenAddr::Tcp(addr) => write!(f, "{addr}"), | ||
#[cfg(unix)] | ||
ListenAddr::Unix(path) => write!(f, "{}", path.display()), | ||
} | ||
} | ||
} | ||
|
||
#[cfg(feature = "settings")] | ||
impl Settings for ListenAddr {} |
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -66,6 +66,8 @@ | |
|
||
mod utils; | ||
|
||
pub mod addr; | ||
|
||
#[cfg(feature = "cli")] | ||
pub mod cli; | ||
|
||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,6 +1,7 @@ | ||
#[cfg(feature = "metrics")] | ||
use super::metrics; | ||
use super::settings::TelemetrySettings; | ||
use crate::addr::ListenAddr; | ||
use crate::telemetry::log; | ||
use crate::BootstrapResult; | ||
use anyhow::Context as _; | ||
|
@@ -14,18 +15,127 @@ use std::net::SocketAddr; | |
use std::pin::Pin; | ||
use std::sync::Arc; | ||
use std::task::{Context, Poll}; | ||
use tokio::io::{AsyncRead, AsyncWrite}; | ||
use tokio::net::TcpListener; | ||
#[cfg(unix)] | ||
use tokio::net::{TcpStream, UnixListener, UnixStream}; | ||
use tokio::sync::watch; | ||
|
||
mod router; | ||
|
||
use router::Router; | ||
|
||
enum TelemetryStream { | ||
Tcp(TcpStream), | ||
#[cfg(unix)] | ||
Unix(UnixStream), | ||
} | ||
|
||
impl AsyncRead for TelemetryStream { | ||
fn poll_read( | ||
self: Pin<&mut Self>, | ||
cx: &mut Context<'_>, | ||
buf: &mut tokio::io::ReadBuf<'_>, | ||
) -> Poll<std::io::Result<()>> { | ||
match self.get_mut() { | ||
TelemetryStream::Tcp(stream) => Pin::new(stream).poll_read(cx, buf), | ||
#[cfg(unix)] | ||
TelemetryStream::Unix(stream) => Pin::new(stream).poll_read(cx, buf), | ||
} | ||
} | ||
} | ||
|
||
impl AsyncWrite for TelemetryStream { | ||
fn poll_write( | ||
self: Pin<&mut Self>, | ||
cx: &mut Context<'_>, | ||
buf: &[u8], | ||
) -> Poll<Result<usize, std::io::Error>> { | ||
match self.get_mut() { | ||
TelemetryStream::Tcp(stream) => Pin::new(stream).poll_write(cx, buf), | ||
#[cfg(unix)] | ||
TelemetryStream::Unix(stream) => Pin::new(stream).poll_write(cx, buf), | ||
} | ||
} | ||
|
||
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), std::io::Error>> { | ||
match self.get_mut() { | ||
TelemetryStream::Tcp(stream) => Pin::new(stream).poll_flush(cx), | ||
#[cfg(unix)] | ||
TelemetryStream::Unix(stream) => Pin::new(stream).poll_flush(cx), | ||
} | ||
} | ||
|
||
fn poll_shutdown( | ||
self: Pin<&mut Self>, | ||
cx: &mut Context<'_>, | ||
) -> Poll<Result<(), std::io::Error>> { | ||
match self.get_mut() { | ||
TelemetryStream::Tcp(stream) => Pin::new(stream).poll_shutdown(cx), | ||
#[cfg(unix)] | ||
TelemetryStream::Unix(stream) => Pin::new(stream).poll_shutdown(cx), | ||
} | ||
} | ||
} | ||
|
||
enum TelemetryListener { | ||
Tcp(TcpListener), | ||
#[cfg(unix)] | ||
Unix(UnixListener), | ||
} | ||
|
||
impl TelemetryListener { | ||
pub(crate) fn local_addr(&self) -> BootstrapResult<ListenAddr> { | ||
match self { | ||
TelemetryListener::Tcp(listener) => Ok(listener.local_addr()?.into()), | ||
#[cfg(unix)] | ||
TelemetryListener::Unix(listener) => match listener.local_addr()?.as_pathname() { | ||
Some(path) => Ok(path.to_path_buf().into()), | ||
None => Err(anyhow::anyhow!("unix socket listener has no pathname")), | ||
}, | ||
} | ||
} | ||
|
||
pub(crate) async fn accept(&self) -> std::io::Result<TelemetryStream> { | ||
match self { | ||
TelemetryListener::Tcp(listener) => listener | ||
.accept() | ||
.await | ||
.map(|(conn, _)| TelemetryStream::Tcp(conn)), | ||
#[cfg(unix)] | ||
TelemetryListener::Unix(listener) => listener | ||
.accept() | ||
.await | ||
.map(|(conn, _)| TelemetryStream::Unix(conn)), | ||
} | ||
} | ||
|
||
pub(crate) fn poll_accept( | ||
&mut self, | ||
cx: &mut std::task::Context<'_>, | ||
) -> std::task::Poll<std::io::Result<TelemetryStream>> { | ||
match self { | ||
TelemetryListener::Tcp(listener) => match std::task::ready!(listener.poll_accept(cx)) { | ||
Ok((conn, _)) => std::task::Poll::Ready(Ok(TelemetryStream::Tcp(conn))), | ||
Err(e) => std::task::Poll::Ready(Err(e)), | ||
}, | ||
#[cfg(unix)] | ||
TelemetryListener::Unix(listener) => { | ||
match std::task::ready!(listener.poll_accept(cx)) { | ||
Ok((conn, _)) => std::task::Poll::Ready(Ok(TelemetryStream::Unix(conn))), | ||
Err(e) => std::task::Poll::Ready(Err(e)), | ||
} | ||
} | ||
} | ||
} | ||
} | ||
|
||
pub use router::{ | ||
BoxError, TelemetryRouteHandler, TelemetryRouteHandlerFuture, TelemetryServerRoute, | ||
}; | ||
|
||
pub(super) struct TelemetryServerFuture { | ||
listener: TcpListener, | ||
listener: TelemetryListener, | ||
router: Router, | ||
} | ||
|
||
|
@@ -47,27 +157,38 @@ impl TelemetryServerFuture { | |
.map_err(|err| anyhow::anyhow!(err))?; | ||
} | ||
|
||
let addr = settings.server.addr; | ||
|
||
#[cfg(feature = "settings")] | ||
let addr = SocketAddr::from(addr); | ||
|
||
let router = Router::new(custom_routes, settings); | ||
|
||
let listener = { | ||
let std_listener = std::net::TcpListener::from( | ||
bind_socket(addr).with_context(|| format!("binding to socket {addr:?}"))?, | ||
); | ||
|
||
std_listener.set_nonblocking(true)?; | ||
let router = Router::new(custom_routes, Arc::clone(&settings)); | ||
|
||
let listener = match &settings.server.addr { | ||
ListenAddr::Tcp(addr) => { | ||
let std_listener = std::net::TcpListener::from( | ||
bind_socket(*addr) | ||
.with_context(|| format!("binding to TCP socket {addr:?}"))?, | ||
); | ||
std_listener.set_nonblocking(true)?; | ||
let tokio_listener = tokio::net::TcpListener::from_std(std_listener)?; | ||
TelemetryListener::Tcp(tokio_listener) | ||
} | ||
#[cfg(unix)] | ||
ListenAddr::Unix(path) => { | ||
// Remove existing socket file if it exists to avoid bind errors | ||
if path.exists() { | ||
if let Err(e) = std::fs::remove_file(path) { | ||
log::warn!("failed to remove existing Unix socket file"; "path" => %path.display(), "error" => e); | ||
} | ||
} | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Do we really want to silently delete existing sockets? I have a gut feeling that this could hide subtle problems There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. It would be easy to get into a bad state during a crash and a restart because we can't reuse old sockets. |
||
|
||
tokio::net::TcpListener::from_std(std_listener)? | ||
let unix_listener = UnixListener::bind(path) | ||
.with_context(|| format!("binding to Unix socket {path:?}"))?; | ||
TelemetryListener::Unix(unix_listener) | ||
} | ||
}; | ||
|
||
Ok(Some(TelemetryServerFuture { listener, router })) | ||
} | ||
pub(super) fn local_addr(&self) -> SocketAddr { | ||
self.listener.local_addr().unwrap() | ||
|
||
pub(super) fn local_addr(&self) -> BootstrapResult<ListenAddr> { | ||
self.listener.local_addr() | ||
} | ||
|
||
// Adapted from Hyper 0.14 Server stuff and axum::serve::serve. | ||
|
@@ -87,15 +208,12 @@ impl TelemetryServerFuture { | |
let (close_tx, close_rx) = watch::channel(()); | ||
let listener = self.listener; | ||
|
||
pin_mut!(listener); | ||
|
||
loop { | ||
let socket = tokio::select! { | ||
conn = listener.accept() => match conn { | ||
Ok((conn, _)) => TokioIo::new(conn), | ||
Ok(conn) => TokioIo::new(conn), | ||
Err(e) => { | ||
log::warn!("failed to accept connection"; "error" => e); | ||
|
||
continue; | ||
} | ||
}, | ||
|
@@ -140,11 +258,10 @@ impl Future for TelemetryServerFuture { | |
let this = &mut *self; | ||
|
||
loop { | ||
let socket = match ready!(Pin::new(&mut this.listener).poll_accept(cx)) { | ||
Ok((conn, _)) => TokioIo::new(conn), | ||
let socket = match ready!(this.listener.poll_accept(cx)) { | ||
Ok(conn) => TokioIo::new(conn), | ||
Err(e) => { | ||
log::warn!("failed to accept connection"; "error" => e); | ||
|
||
continue; | ||
} | ||
}; | ||
|
Uh oh!
There was an error while loading. Please reload this page.