put DropGuard in its own module

Author: ckampfe

src/channel.rs

@@ -1,4 +1,5 @@
-use crate::{AppState, Done};
+use crate::AppState;
+use crate::drop_guard::DropGuard;
 use axum::Router;
 use axum::body::{Body, BodyDataStream};
 use axum::extract::{Path, Request, State};
@@ -18,8 +19,8 @@ pub(crate) type ChannelClients = Mutex<
         HashMap<
             ChannelName,
             (
-                flume::Sender<(BodyDataStream, HeaderMap, Done)>,
-                flume::Receiver<(BodyDataStream, HeaderMap, Done)>,
+                flume::Sender<(BodyDataStream, HeaderMap, DropGuard)>,
+                flume::Receiver<(BodyDataStream, HeaderMap, DropGuard)>,
             ),
         >,
     >,
@@ -133,13 +134,13 @@ async fn broadcast_to_channel(
 
     let request_body_stream = body.into_data_stream();
 
-    let (done, done_rx) = Done::new();
+    let (drop_guard, drop_guard_rx) = DropGuard::new();
 
-    tx.send_async((request_body_stream, request_headers, done))
+    tx.send_async((request_body_stream, request_headers, drop_guard))
         .await
         .map_err(|_e| StatusCode::INTERNAL_SERVER_ERROR)?;
 
-    done_rx
+    drop_guard_rx
         .await
         .map_err(|_e| StatusCode::INTERNAL_SERVER_ERROR)?;
 
@@ -171,10 +172,10 @@ async fn subscribe_to_channel(
 
     let rx = rx.into_recv_async();
 
-    let (stream, producer_request_headers, _done) =
+    let (request_body_stream, producer_request_headers, _drop_guard) =
         rx.await.map_err(|_e| StatusCode::INTERNAL_SERVER_ERROR)?;
 
-    let body = Body::from_stream(stream);
+    let body = Body::from_stream(request_body_stream);
 
     // we do this because by default, POSTs from curl are `x-www-form-urlencoded`
     let producer_content_type = producer_request_headers
@@ -198,7 +199,8 @@ async fn subscribe_to_channel(
 
 #[cfg(test)]
 mod tests {
-    use crate::{Done, Options, app};
+    use crate::drop_guard::DropGuard;
+    use crate::{Options, app};
     use axum::http::StatusCode;
     use serde::{Deserialize, Serialize};
     use std::{collections::HashSet, sync::atomic::AtomicU16};
@@ -230,11 +232,11 @@ mod tests {
             .await
             .unwrap();
 
-        let (_done, done_rx) = Done::new();
+        let (_drop_guard, drop_guard_rx) = DropGuard::new();
 
         tokio::spawn(async move {
             axum::serve(listener, app(options))
-                .with_graceful_shutdown(async move { done_rx.await.unwrap() })
+                .with_graceful_shutdown(async move { drop_guard_rx.await.unwrap() })
                 .await
                 .unwrap();
         });
@@ -300,11 +302,11 @@ mod tests {
             .await
             .unwrap();
 
-        let (_done, done_rx) = Done::new();
+        let (_drop_guard, drop_guard_rx) = DropGuard::new();
 
         tokio::spawn(async move {
             axum::serve(listener, app(options))
-                .with_graceful_shutdown(async move { done_rx.await.unwrap() })
+                .with_graceful_shutdown(async move { drop_guard_rx.await.unwrap() })
                 .await
                 .unwrap();
         });
@@ -351,11 +353,11 @@ mod tests {
             .await
             .unwrap();
 
-        let (_done, done_rx) = Done::new();
+        let (_drop_guard, drop_guard_rx) = DropGuard::new();
 
         tokio::spawn(async move {
             axum::serve(listener, app(options))
-                .with_graceful_shutdown(async move { done_rx.await.unwrap() })
+                .with_graceful_shutdown(async move { drop_guard_rx.await.unwrap() })
                 .await
                 .unwrap();
         });
@@ -400,11 +402,11 @@ mod tests {
             .await
             .unwrap();
 
-        let (_done, done_rx) = Done::new();
+        let (_drop_guard, drop_guard_rx) = DropGuard::new();
 
         tokio::spawn(async move {
             axum::serve(listener, app(options))
-                .with_graceful_shutdown(async move { done_rx.await.unwrap() })
+                .with_graceful_shutdown(async move { drop_guard_rx.await.unwrap() })
                 .await
                 .unwrap();
         });

src/drop_guard.rs

@@ -0,0 +1,26 @@
+use tokio::sync::oneshot;
+
+/// When a `DropGuard` is dropped, it signals to its paired receiver that it has been dropped.
+/// This allows an external observer, maybe in another task, to know when the
+/// `DropGuard` was dropped.
+pub(crate) struct DropGuard {
+    tx: Option<oneshot::Sender<()>>,
+}
+
+impl DropGuard {
+    pub fn new() -> (DropGuard, oneshot::Receiver<()>) {
+        let (tx, rx) = oneshot::channel();
+        (DropGuard { tx: Some(tx) }, rx)
+    }
+}
+
+impl Drop for DropGuard {
+    fn drop(&mut self) {
+        // needed because we can't move out of a &mut
+        let tx = self
+            .tx
+            .take()
+            .expect("this should never happen, it should always be Some");
+        let _ = tx.send(());
+    }
+}

src/main.rs

@@ -25,32 +25,10 @@ use clap::Parser;
 use std::collections::HashMap;
 use std::error::Error;
 use std::sync::Arc;
-use tokio::sync::{Mutex, oneshot};
+use tokio::sync::Mutex;
 
 mod channel;
-
-/// when this is dropped it signals the oneshot channel
-struct Done {
-    tx: Option<oneshot::Sender<()>>,
-}
-
-impl Done {
-    fn new() -> (Done, oneshot::Receiver<()>) {
-        let (tx, rx) = oneshot::channel();
-        (Done { tx: Some(tx) }, rx)
-    }
-}
-
-impl Drop for Done {
-    fn drop(&mut self) {
-        // needed because we can't move out of a &mut
-        let tx = self
-            .tx
-            .take()
-            .expect("this should never happen, it should always be Some");
-        let _ = tx.send(());
-    }
-}
+mod drop_guard;
 
 async fn app_state(
     State(state): State<Arc<AppState>>,