fix(prometheus_exporter sink): close listener before the sink finishes

src/sinks/prometheus/exporter.rs

@@ -3,13 +3,15 @@ use std::{
     hash::Hash,
     mem::{Discriminant, discriminant},
     net::{IpAddr, Ipv4Addr, SocketAddr},
+    pin::Pin,
     sync::{Arc, RwLock},
+    task::{Context, Poll},
     time::{Duration, Instant},
 };
 
 use async_trait::async_trait;
 use base64::prelude::{BASE64_STANDARD, Engine as _};
-use futures::{FutureExt, StreamExt, future, stream::BoxStream};
+use futures::{FutureExt, Stream, StreamExt, future, stream::BoxStream};
 use hyper::{
     Body, Method, Request, Response, Server, StatusCode,
     body::HttpBody,
@@ -19,6 +21,7 @@ use hyper::{
 use indexmap::{IndexMap, map::Entry};
 use serde_with::serde_as;
 use stream_cancel::{Trigger, Tripwire};
+use tokio::sync::oneshot;
 use tower::ServiceBuilder;
 use tower_http::compression::CompressionLayer;
 use tracing::{Instrument, Span};
@@ -228,11 +231,28 @@ impl ValidatedSink for PrometheusExporterConfig {
 }
 
 struct PrometheusExporter {
-    server_shutdown_trigger: Option<Trigger>,
+    /// Shutdown trigger of the running HTTP server, and a receiver that completes when the server
+    /// closes its listener.
+    server: Option<(Trigger, oneshot::Receiver<()>)>,
     config: PrometheusExporterConfig,
     metrics: Arc<RwLock<IndexMap<MetricRef, (Metric, MetricMetadata)>>>,
 }
 
+/// Accept stream that drops `_closed` after it closes its listener.
+struct NotifyOnClose<S> {
+    // Rust drops fields in declaration order, so the listener is closed before `_closed` is dropped.
+    incoming: S,
+    _closed: oneshot::Sender<()>,
+}
+
+impl<S: Stream + Unpin> Stream for NotifyOnClose<S> {
+    type Item = S::Item;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<S::Item>> {
+        self.incoming.poll_next_unpin(cx)
+    }
+}
+
 /// Expiration metadata for a metric.
 #[derive(Clone, Copy, Debug)]
 struct MetricMetadata {
@@ -430,14 +450,14 @@ impl Handler {
 impl PrometheusExporter {
     fn new(config: PrometheusExporterConfig) -> Self {
         Self {
-            server_shutdown_trigger: None,
+            server: None,
             config,
             metrics: Arc::new(RwLock::new(IndexMap::new())),
         }
     }
 
     async fn start_server_if_needed(&mut self) -> crate::Result<()> {
-        if self.server_shutdown_trigger.is_some() {
+        if self.server.is_some() {
             return Ok(());
         }
 
@@ -479,21 +499,28 @@ impl PrometheusExporter {
 
         let tls = MaybeTlsSettings::from_config(tls.as_ref(), true)?;
         let listener = tls.bind(&address).await?;
+        // On shutdown, the server drops its accept stream at once and then drains open
+        // connections. The stream drops `closed` with the listener.
+        let (closed, listener_closed) = oneshot::channel();
+        let incoming = NotifyOnClose {
+            incoming: listener.accept_stream(),
+            _closed: closed,
+        };
 
         crate::spawn_in_current_span(async move {
             info!(message = "Building HTTP server.", address = %address);
 
-            Server::builder(hyper::server::accept::from_stream(listener.accept_stream()))
+            if let Err(error) = Server::builder(hyper::server::accept::from_stream(incoming))
                 .serve(new_service)
                 .with_graceful_shutdown(tripwire.then(crate::shutdown::tripwire_handler))
                 .instrument(span)
                 .await
-                .map_err(|error| error!("Server error: {error}."))?;
-
-            Ok::<(), ()>(())
+            {
+                error!("Server error: {error}.");
+            }
         });
 
-        self.server_shutdown_trigger = Some(trigger);
+        self.server = Some((trigger, listener_closed));
         Ok(())
     }
 
@@ -614,6 +641,15 @@ impl StreamSink<Event, SinkError> for PrometheusExporter {
             finalizers.update_status(EventStatus::Delivered);
         }
 
+        // Wait for the server to close its listener before the sink finishes, because a reload
+        // that binds the same address starts only after the old sink finishes. Open connections
+        // continue to drain in the server task, so a stalled client cannot block the sink.
+        if let Some((trigger, listener_closed)) = self.server.take() {
+            drop(trigger);
+            // The sender never sends. The receive completes when the sender is dropped.
+            _ = listener_closed.await;
+        }
+
         Ok(())
     }
 }
@@ -642,7 +678,7 @@ mod tests {
         test_util::{
             addr::next_addr,
             components::{SINK_TAGS, run_and_assert_sink_compliance},
-            random_string, trace_init,
+            random_string, trace_init, wait_for_tcp,
         },
         tls::MaybeTlsSettings,
     };
@@ -669,6 +705,61 @@ mod tests {
             .expect("default config should validate");
     }
 
+    #[tokio::test]
+    async fn finished_sink_releases_address() {
+        let (_guard, address) = next_addr();
+        let config = PrometheusExporterConfig {
+            address,
+            ..Default::default()
+        };
+        let (sink, _) = SinkConfig::build(&config, SinkContext::default())
+            .await
+            .unwrap();
+
+        sink.run(stream::empty::<Event>().map(Into::into))
+            .await
+            .unwrap();
+
+        // A reload binds the address again as soon as the old sink finishes.
+        std::net::TcpListener::bind(address).expect("finished sink still holds its address");
+    }
+
+    #[tokio::test]
+    async fn stalled_scrape_does_not_block_sink_finish() {
+        use tokio::{io::AsyncWriteExt, net::TcpStream, sync::mpsc};
+        use tokio_stream::wrappers::UnboundedReceiverStream;
+
+        let (_guard, address) = next_addr();
+        let config = PrometheusExporterConfig {
+            address,
+            ..Default::default()
+        };
+        let (sink, _) = SinkConfig::build(&config, SinkContext::default())
+            .await
+            .unwrap();
+        let (tx, rx) = mpsc::unbounded_channel::<Event>();
+        let sink_handle = tokio::spawn(sink.run(UnboundedReceiverStream::new(rx).map(Into::into)));
+        wait_for_tcp(address).await;
+
+        // A client that never completes its request keeps its connection busy, so a graceful
+        // drain of the server waits for it.
+        let mut client = TcpStream::connect(address).await.unwrap();
+        client
+            .write_all(b"GET /metrics HTTP/1.1\r\n")
+            .await
+            .unwrap();
+        time::sleep(std::time::Duration::from_millis(100)).await;
+
+        drop(tx);
+        time::timeout(std::time::Duration::from_secs(5), sink_handle)
+            .await
+            .expect("stalled scrape blocked the sink")
+            .unwrap()
+            .unwrap();
+        std::net::TcpListener::bind(address).expect("finished sink still holds its address");
+        drop(client);
+    }
+
     #[tokio::test]
     async fn prometheus_notls() {
         export_and_fetch_simple(None).await;