Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 7 additions & 8 deletions crates/datadog-serverless-compat/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -104,14 +104,12 @@ pub async fn main() {
None
}
};
let dd_dogstatsd_port: u16 = if dd_dogstatsd_windows_pipe_name.is_some() {
0 // Override to 0 when using Windows named pipe
} else {
env::var("DD_DOGSTATSD_PORT")
.ok()
.and_then(|port| port.parse::<u16>().ok())
.unwrap_or(DEFAULT_DOGSTATSD_PORT)
};
// UDP listener also runs when a named pipe is set, so clients without
// named-pipe support can still send metrics.
let dd_dogstatsd_port: u16 = env::var("DD_DOGSTATSD_PORT")
.ok()
.and_then(|port| port.parse::<u16>().ok())
.unwrap_or(DEFAULT_DOGSTATSD_PORT);
let dd_site = env::var("DD_SITE").unwrap_or_else(|_| "datadoghq.com".to_string());
let dd_use_dogstatsd = env::var("DD_USE_DOGSTATSD")
.map(|val| val.to_lowercase() != "false")
Expand Down Expand Up @@ -252,6 +250,7 @@ pub async fn main() {
)
.await;
if let Some(ref windows_pipe_name) = dd_dogstatsd_windows_pipe_name {
// The dogstatsd crate logs whether the best-effort UDP bind succeeded.
info!("dogstatsd-pipe: starting to listen on pipe {windows_pipe_name}");
} else {
info!("dogstatsd-udp: starting to listen on port {dd_dogstatsd_port}");
Expand Down
105 changes: 83 additions & 22 deletions crates/dogstatsd/src/dogstatsd.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,12 @@ use tracing::{debug, error, trace};

// Windows-specific imports
#[cfg(all(windows, feature = "windows-pipes"))]
use {std::sync::Arc, tokio::io::AsyncReadExt, tokio::net::windows::named_pipe::ServerOptions};
use {
std::sync::Arc,
tokio::io::AsyncReadExt,
tokio::net::windows::named_pipe::ServerOptions,
tracing::{info, warn},
};

// Default buffer size for receiving DogStatsD packets (one read call).
// Used for both UDP recv_from and Windows named pipe reads.
Expand All @@ -31,7 +36,8 @@ const DEFAULT_QUEUE_SIZE: usize = 1024;
pub struct DogStatsDConfig {
/// Host to bind UDP socket to (e.g., "127.0.0.1")
pub host: String,
/// Port to bind UDP socket to (e.g., 8125), will be 0 if we're using a Named Pipe
/// Port to bind UDP socket to (e.g., 8125). The UDP socket is also bound when a named pipe
/// is configured, so clients without named-pipe support can still send metrics.
pub port: u16,
/// Optional namespace to prepend to all metric names (e.g., "myapp")
pub metric_namespace: Option<String>,
Expand Down Expand Up @@ -94,11 +100,14 @@ enum BufferReader {
#[allow(dead_code)]
MirrorTest(Vec<u8>, SocketAddr),

/// Windows named pipe reader (Windows-only transport)
/// Windows named pipe reader (Windows-only transport), read alongside the UDP socket.
/// `udp_socket` is `None` when the UDP bind failed, e.g. because another agent on a
/// shared host already owns the port.
#[cfg(all(windows, feature = "windows-pipes"))]
NamedPipe {
pipe_name: Arc<String>,
receiver: Arc<tokio::sync::Mutex<tokio::sync::mpsc::UnboundedReceiver<Vec<u8>>>>,
udp_socket: Option<tokio::net::UdpSocket>,
},
}

Expand All @@ -121,17 +130,34 @@ impl BufferReader {
BufferReader::NamedPipe {
pipe_name,
receiver,
} => match receiver.lock().await.recv().await {
Some(data) => {
let len = data.len().min(buf.len());
buf[..len].copy_from_slice(&data[..len]);
Ok((len, MessageSource::NamedPipe(pipe_name.clone())))
udp_socket,
} => {
// Wait on whichever transport delivers first. The UDP arm stays pending
// when the UDP bind failed.
let pipe_data = tokio::select! {
udp_result = async {
match udp_socket.as_ref() {
Some(socket) => socket.recv_from(buf).await,
None => std::future::pending().await,
}
} => {
let (amt, src) = udp_result?;
return Ok((amt, MessageSource::Network(src)));
}
pipe_data = async { receiver.lock().await.recv().await } => pipe_data,
};
match pipe_data {
Some(data) => {
let len = data.len().min(buf.len());
buf[..len].copy_from_slice(&data[..len]);
Ok((len, MessageSource::NamedPipe(pipe_name.clone())))
}
None => Err(std::io::Error::new(
std::io::ErrorKind::ConnectionReset,
"named pipe channel closed",
)),
}
None => Err(std::io::Error::new(
std::io::ErrorKind::ConnectionReset,
"named pipe channel closed",
)),
},
}
}
}

Expand All @@ -141,17 +167,30 @@ impl BufferReader {
/// kernel buffer without re-entering tokio's event loop.
fn try_read_into(&mut self, buf: &mut [u8]) -> std::io::Result<Option<(usize, MessageSource)>> {
match self {
BufferReader::UdpSocket(socket) => match socket.try_recv_from(buf) {
Ok((amt, src)) => Ok(Some((amt, MessageSource::Network(src)))),
Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => Ok(None),
Err(e) => Err(e),
},
BufferReader::UdpSocket(socket) => try_recv_udp(socket, buf),
// Drain only the UDP side; the named pipe channel is read by `read_into`.
#[cfg(all(windows, feature = "windows-pipes"))]
BufferReader::NamedPipe {
udp_socket: Some(socket),
..
} => try_recv_udp(socket, buf),
// Non-UDP transports don't support non-blocking reads
Comment on lines +171 to 177
_ => Ok(None),
}
}
}

fn try_recv_udp(
socket: &tokio::net::UdpSocket,
buf: &mut [u8],
) -> std::io::Result<Option<(usize, MessageSource)>> {
match socket.try_recv_from(buf) {
Ok((amt, src)) => Ok(Some((amt, MessageSource::Network(src)))),
Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => Ok(None),
Err(e) => Err(e),
}
}

/// DogStatsD server to receive, parse, and forward metrics.
pub struct DogStatsD {
cancel_token: tokio_util::sync::CancellationToken,
Expand All @@ -165,7 +204,8 @@ pub struct DogStatsD {
impl DogStatsD {
/// Creates a new DogStatsD server instance.
///
/// The server will bind to either a UDP socket or Windows named pipe based on the config.
/// The server always binds a UDP socket and also listens on a Windows named pipe when one
/// is configured.
/// Metrics received will be forwarded to the provided aggregator_handle.
#[must_use]
pub async fn new(
Expand All @@ -191,10 +231,30 @@ impl DogStatsD {
None => DEFAULT_BUFFER_SIZE,
};

let addr = format!("{}:{}", config.host, config.port);

let buffer_reader = if let Some(pipe_name_ref) = pipe_name_opt {
// Windows named pipe transport
// Windows named pipe transport, plus UDP for clients without named-pipe support
#[cfg(all(windows, feature = "windows-pipes"))]
{
// UDP bind is best-effort when a named pipe is configured. Multiple Azure apps
// share the same VM on certain Windows Azure plans, and only one agent can own
// the port. The named pipe is unique per app, so a UDP bind failure is not fatal.
let udp_socket = match create_udp_socket(&addr, config.so_rcvbuf).await {
Ok(socket) => {
info!("dogstatsd-udp: starting to listen on {}", addr);
Some(socket)
}
Err(e) => {
warn!(
"DogStatsD could not bind UDP {}: {}. Named-pipe transport is active; \
metrics sent over UDP will not be received by this agent.",
addr, e
);
None
}
};

let pipe_name = Arc::new(pipe_name_ref.clone());

// Create channel for receiving data from client handlers
Expand All @@ -213,6 +273,7 @@ impl DogStatsD {
BufferReader::NamedPipe {
pipe_name,
receiver,
udp_socket,
}
}
#[cfg(not(all(windows, feature = "windows-pipes")))]
Expand All @@ -225,8 +286,8 @@ impl DogStatsD {
);
}
} else {
// UDP socket for all platforms
let addr = format!("{}:{}", config.host, config.port);
// UDP socket for all platforms. Without a named pipe, UDP is the only transport, so a
// bind failure is fatal.
// TODO (UDS socket)
#[allow(clippy::expect_used)]
let socket = create_udp_socket(&addr, config.so_rcvbuf)
Expand Down
72 changes: 72 additions & 0 deletions crates/dogstatsd/tests/integration_test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -548,6 +548,78 @@ async fn test_named_pipe_basic_communication() {
handle.shutdown().expect("shutdown failed");
}

#[cfg(test)]
#[cfg(all(windows, feature = "windows-pipes"))]
#[tokio::test]
async fn test_named_pipe_and_udp_dual_listen() {
let pipe_name = r"\\.\pipe\test_dogstatsd_dual_listen";
let udp_port: u16 = 18130;
let (service, handle) = AggregatorService::new(SortedTags::parse("test:value").unwrap(), 1_024)
.expect("aggregator service creation failed");
tokio::spawn(service.run());

let cancel_token = CancellationToken::new();

// Both transports are deliberately set on the same server: a non-zero UDP
// port AND a pipe name. They must come up concurrently.
let dogstatsd_task = {
let handle = handle.clone();
let cancel_token = cancel_token.clone();
tokio::spawn(async move {
let dogstatsd = DogStatsD::new(
&DogStatsDConfig {
host: "127.0.0.1".to_string(),
port: udp_port,
metric_namespace: None,
windows_pipe_name: Some(pipe_name.to_string()),
so_rcvbuf: None,
buffer_size: None,
queue_size: None,
},
handle,
cancel_token,
)
.await;
dogstatsd.spin().await;
})
};

sleep(Duration::from_millis(100)).await;

// One metric per transport, distinguishable by metric name.
let socket = UdpSocket::bind("0.0.0.0:0").await.expect("bind failed");
socket
.send_to(b"dual.udp.metric:1|c", ("127.0.0.1", udp_port))
.await
.expect("UDP send failed");

let mut client = ClientOptions::new().open(pipe_name).expect("client open");
client
.write_all(b"dual.pipe.metric:1|c\n")
.await
.expect("pipe write failed");
client.flush().await.expect("pipe flush failed");

sleep(Duration::from_millis(100)).await;

let response = handle.flush().await.expect("flush failed");
let series_json = serde_json::to_string(&response.series).expect("serialize series");
assert!(
series_json.contains("dual.udp.metric"),
"UDP metric was not received when a pipe name is also configured: {series_json}"
);
assert!(
series_json.contains("dual.pipe.metric"),
"pipe metric was not received when UDP is also listening: {series_json}"
);

// Cleanup
cancel_token.cancel();
let result = timeout(Duration::from_millis(500), dogstatsd_task).await;
assert!(result.is_ok(), "task should complete after cancellation");
handle.shutdown().expect("shutdown failed");
}

#[cfg(test)]
#[cfg(all(windows, feature = "windows-pipes"))]
#[tokio::test]
Expand Down
Loading