Pre-register more metrics in the blackholes (#552)

blt · web-flow · commit 225129564f18 · 2023-04-17T19:02:20.000-04:00
This commit pre-registers metrics in our blackholes, a peer to #544. I have made use of OnceCell in situations where a registered metrics needs to cross task boundaries. Fortunately in our blackholes we do not use labels in our metrics so this was a fairly total conversion. Signed-off-by: Brian L. Troutwine <brian.troutwine@datadoghq.com>
diff --git a/src/blackhole/http.rs b/src/blackhole/http.rs
@@ -9,7 +9,8 @@ use hyper::{
     service::{make_service_fn, service_fn},
     Body, Request, Response, Server, StatusCode,
 };
-use once_cell::unsync::OnceCell;
+use metrics::{register_counter, Counter};
+use once_cell::{sync, unsync::OnceCell};
 use serde::{Deserialize, Serialize};
 use tower::ServiceBuilder;
 use tracing::{debug, error, info};
@@ -19,6 +20,18 @@ use crate::signals::Shutdown;
 #[allow(clippy::declare_interior_mutable_const)]
 const RESPONSE: OnceCell<Vec<u8>> = OnceCell::new();
 
+static BYTES_RECEIVED: sync::OnceCell<Counter> = sync::OnceCell::new();
+#[inline]
+fn bytes_received() -> &'static Counter {
+    BYTES_RECEIVED.get_or_init(|| register_counter!("bytes_received"))
+}
+
+static REQUESTS_RECEIVED: sync::OnceCell<Counter> = sync::OnceCell::new();
+#[inline]
+fn requests_received() -> &'static Counter {
+    REQUESTS_RECEIVED.get_or_init(|| register_counter!("requests_received"))
+}
+
 fn default_concurrent_requests_max() -> usize {
     100
 }
@@ -105,7 +118,7 @@ async fn srv(
     req: Request<Body>,
     headers: HeaderMap,
 ) -> Result<Response<Body>, hyper::Error> {
-    metrics::counter!("requests_received", 1);
+    requests_received().increment(1);
 
     let (parts, body) = req.into_parts();
 
@@ -114,7 +127,7 @@ async fn srv(
     match crate::codec::decode(parts.headers.get(hyper::header::CONTENT_ENCODING), bytes) {
         Err(response) => Ok(response),
         Ok(body) => {
-            metrics::counter!("bytes_received", body.len() as u64);
+            bytes_received().increment(body.len() as u64);
 
             let mut okay = Response::default();
             *okay.status_mut() = status;
diff --git a/src/blackhole/sqs.rs b/src/blackhole/sqs.rs
@@ -8,6 +8,8 @@ use hyper::{
     service::{make_service_fn, service_fn},
     Body, Request, Response, Server, StatusCode,
 };
+use metrics::{register_counter, Counter};
+use once_cell::sync;
 use rand::{distributions::Alphanumeric, thread_rng, Rng};
 use serde::Deserialize;
 use tokio::time::Duration;
@@ -23,6 +25,18 @@ pub enum Error {
     Hyper(hyper::Error),
 }
 
+static BYTES_RECEIVED: sync::OnceCell<Counter> = sync::OnceCell::new();
+#[inline]
+fn bytes_received() -> &'static Counter {
+    BYTES_RECEIVED.get_or_init(|| register_counter!("bytes_received"))
+}
+
+static REQUESTS_RECEIVED: sync::OnceCell<Counter> = sync::OnceCell::new();
+#[inline]
+fn requests_received() -> &'static Counter {
+    REQUESTS_RECEIVED.get_or_init(|| register_counter!("requests_received"))
+}
+
 fn default_concurrent_requests_max() -> usize {
     100
 }
@@ -200,10 +214,10 @@ impl DeleteMessageBatch {
 }
 
 async fn srv(req: Request<Body>) -> Result<Response<Body>, hyper::Error> {
-    metrics::counter!("requests_received", 1);
+    requests_received().increment(1);
 
     let bytes = body::to_bytes(req).await?;
-    metrics::counter!("bytes_received", bytes.len() as u64);
+    bytes_received().increment(bytes.len() as u64);
 
     let action: Action = serde_qs::from_bytes(&bytes).unwrap();
 
diff --git a/src/blackhole/tcp.rs b/src/blackhole/tcp.rs
@@ -3,14 +3,27 @@
 use std::{io, net::SocketAddr};
 
 use futures::stream::StreamExt;
-use metrics::counter;
+use metrics::{register_counter, Counter};
+use once_cell::sync;
 use serde::Deserialize;
 use tokio::net::{TcpListener, TcpStream};
 use tokio_util::io::ReaderStream;
 use tracing::info;
 
 use crate::signals::Shutdown;
 
+static BYTES_RECEIVED: sync::OnceCell<Counter> = sync::OnceCell::new();
+#[inline]
+fn bytes_received() -> &'static Counter {
+    BYTES_RECEIVED.get_or_init(|| register_counter!("bytes_received"))
+}
+
+static MESSAGE_RECEIVED: sync::OnceCell<Counter> = sync::OnceCell::new();
+#[inline]
+fn message_received() -> &'static Counter {
+    MESSAGE_RECEIVED.get_or_init(|| register_counter!("message_received"))
+}
+
 #[derive(Debug)]
 /// Errors emitted by [`Tcp`]
 pub enum Error {
@@ -46,9 +59,9 @@ impl Tcp {
         let mut stream = ReaderStream::new(socket);
 
         while let Some(msg) = stream.next().await {
-            counter!("message_received", 1);
+            message_received().increment(1);
             if let Ok(msg) = msg {
-                counter!("bytes_received", msg.len() as u64);
+                bytes_received().increment(msg.len() as u64);
             }
         }
     }
@@ -70,11 +83,13 @@ impl Tcp {
             .await
             .map_err(Error::Io)?;
 
+        let connection_accepted = register_counter!("connection_accepted");
+
         loop {
             tokio::select! {
                 conn = listener.accept() => {
                     let (socket, _) = conn.map_err(Error::Io)?;
-                    counter!("connection_accepted", 1);
+                    connection_accepted.increment(1);
                     tokio::spawn(async move {
                         Self::handle_connection(socket).await;
                     });
diff --git a/src/blackhole/udp.rs b/src/blackhole/udp.rs
@@ -2,7 +2,7 @@
 
 use std::{io, net::SocketAddr};
 
-use metrics::counter;
+use metrics::register_counter;
 use serde::Deserialize;
 use tokio::net::UdpSocket;
 use tracing::info;
@@ -56,14 +56,17 @@ impl Udp {
         let socket = UdpSocket::bind(&self.binding_addr)
             .await
             .map_err(Error::Io)?;
-        let mut buf: Vec<u8> = vec![0; 65536];
+        let mut buf = [0; 65536];
+
+        let bytes_received = register_counter!("bytes_received");
+        let packet_received = register_counter!("packet_received");
 
         loop {
             tokio::select! {
                 packet = socket.recv_from(&mut buf) => {
                     let (bytes, _) = packet.map_err(Error::Io)?;
-                    counter!("packet_received", 1);
-                    counter!("bytes_received", bytes as u64);
+                    packet_received.increment(1);
+                    bytes_received.increment(bytes as u64);
                 }
                 _ = self.shutdown.recv() => {
                     info!("shutdown signal received");
diff --git a/src/blackhole/unix_datagram.rs b/src/blackhole/unix_datagram.rs
@@ -3,7 +3,7 @@
 use std::{io, path::PathBuf};
 
 use futures::TryFutureExt;
-use metrics::counter;
+use metrics::register_counter;
 use serde::Deserialize;
 use tokio::net;
 use tracing::info;
@@ -58,14 +58,15 @@ impl UnixDatagram {
         // socket, ignore any errors.
         let _res = tokio::fs::remove_file(&self.path).map_err(Error::Io);
         let socket = net::UnixDatagram::bind(&self.path).map_err(Error::Io)?;
+        let mut buf = [0; 65536];
 
-        loop {
-            let mut buf = [0; 65536];
+        let bytes_received = register_counter!("bytes_received");
 
+        loop {
             tokio::select! {
                 res = socket.recv(&mut buf) => {
                     let n: usize = res.map_err(Error::Io)?;
-                    counter!("bytes_received", n as u64);
+                    bytes_received.increment(n as u64);
                 }
                 _ = self.shutdown.recv() => {
                     info!("shutdown signal received");
diff --git a/src/blackhole/unix_stream.rs b/src/blackhole/unix_stream.rs
@@ -3,14 +3,27 @@
 use std::{io, path::PathBuf};
 
 use futures::StreamExt;
-use metrics::counter;
+use metrics::{register_counter, Counter};
+use once_cell::sync;
 use serde::Deserialize;
 use tokio::net;
 use tokio_util::io::ReaderStream;
 use tracing::info;
 
 use crate::signals::Shutdown;
 
+static BYTES_RECEIVED: sync::OnceCell<Counter> = sync::OnceCell::new();
+#[inline]
+fn bytes_received() -> &'static Counter {
+    BYTES_RECEIVED.get_or_init(|| register_counter!("bytes_received"))
+}
+
+static MESSAGE_RECEIVED: sync::OnceCell<Counter> = sync::OnceCell::new();
+#[inline]
+fn message_received() -> &'static Counter {
+    MESSAGE_RECEIVED.get_or_init(|| register_counter!("message_received"))
+}
+
 #[derive(Debug)]
 /// Errors produced by [`UnixStream`].
 pub enum Error {
@@ -57,11 +70,13 @@ impl UnixStream {
     pub async fn run(mut self) -> Result<(), Error> {
         let listener = net::UnixListener::bind(&self.path).map_err(Error::Io)?;
 
+        let connection_accepted = register_counter!("connection_accepted");
+
         loop {
             tokio::select! {
                 conn = listener.accept() => {
                     let (socket, _) = conn.map_err(Error::Io)?;
-                    counter!("connection_accepted", 1);
+                    connection_accepted.increment(1);
                     tokio::spawn(async move {
                         Self::handle_connection(socket).await;
                     });
@@ -78,9 +93,9 @@ impl UnixStream {
         let mut stream = ReaderStream::new(socket);
 
         while let Some(msg) = stream.next().await {
-            counter!("message_received", 1);
+            message_received().increment(1);
             if let Ok(msg) = msg {
-                counter!("bytes_received", msg.len() as u64);
+                bytes_received().increment(msg.len() as u64);
             }
         }
     }
