refactor: use cancellation tokens everywhere

Signed-off-by: tjh <[e-mail]>

src/lib.rs

···
 pub mod tap;
-pub mod util;

src/main.rs

···
 use std::{
     collections::{HashSet, VecDeque},
     io,
-    ops::ControlFlow::{self, *},
     time::Duration,
 };
 
 use futures_util::StreamExt;
 use serde_json::Value;
 use sqlx::PgPool;
-use tokio::{
-    sync::{mpsc, watch},
-    task::JoinSet,
-};
+use tokio::{sync::mpsc, task::JoinSet};
+use tokio_util::sync::CancellationToken;
 use tracing::level_filters::LevelFilter;
 use tracing_subscriber::{EnvFilter, layer::SubscriberExt as _, util::SubscriberInitExt as _};
-use trap::{
-    tap::{IdentityEvent, RecordAction, RecordEvent, TapChannel, TapClient, TapEvent},
-    util::with_shutdown,
-};
+use trap::tap::{IdentityEvent, RecordAction, RecordEvent, TapChannel, TapClient, TapEvent};
 
 /// DID provenance.
 #[derive(Debug)]
···
 
     sqlx::migrate!().run(&pool).await?;
 
-    let (shutdown_tx, shutdown_rx) = watch::channel(false);
+    let shutdown = CancellationToken::new();
     let (did_tx, did_rx) = mpsc::unbounded_channel::<(String, DidSource)>();
 
     let tap = TapClient::new(arguments.tap, arguments.tap_password.as_deref())?;
···
         tap_channel,
         pool.clone(),
         did_tx.clone(),
-        shutdown_rx.clone(),
+        shutdown.child_token(),
     ));
 
-    tasks.spawn(did_task(tap, pool, did_rx, shutdown_rx.clone()));
-    tasks.spawn(shutdown_task(shutdown_tx.clone()));
+    tasks.spawn(did_task(tap, pool, did_rx, shutdown.child_token()));
+    tasks.spawn(shutdown_task(shutdown.clone()));
 
     // Submit seed DIDs to the Tap service.
     for did in arguments.seed.into_iter().filter(|s| possible_did(s)) {
···
     for task in tasks.join_all().await {
         if let Err(error) = task {
             tracing::error!(?error, "task failed");
-            shutdown_tx.send(true)?;
+            shutdown.cancel();
         }
     }
 
···
     Ok(row.0)
 }
 
-#[tracing::instrument(skip(channel, pool, tx, rx))]
+#[tracing::instrument(skip(channel, pool, tx, shutdown))]
 async fn event_consumer(
     mut channel: TapChannel,
     pool: PgPool,
     tx: mpsc::UnboundedSender<(String, DidSource)>,
-    mut rx: watch::Receiver<bool>,
+    shutdown: CancellationToken,
 ) -> anyhow::Result<()> {
-    use ControlFlow::Continue;
-
-    while let Continue(Some((event, ack))) = with_shutdown(channel.recv(), &mut rx).await {
+    while let Some(Some((event, ack))) = shutdown.run_until_cancelled(channel.recv()).await {
         let mut transaction = pool.begin().await?;
         match event {
             TapEvent::Record(record) => {
···
     Ok((record_event, parsed_record))
 }
 
-#[tracing::instrument(skip(tap, pool, did_rx, shutdown_rx))]
+#[tracing::instrument(skip(tap, pool, did_rx, shutdown))]
 async fn did_task(
     tap: TapClient,
     pool: PgPool,
     mut did_rx: mpsc::UnboundedReceiver<(String, DidSource)>,
-    mut shutdown_rx: watch::Receiver<bool>,
+    shutdown: CancellationToken,
 ) -> anyhow::Result<()> {
     const BATCH: usize = 64;
 
···
 
     loop {
         tokio::time::sleep(Duration::from_millis(200)).await;
-        match with_shutdown(did_rx.recv_many(&mut dids, BATCH), &mut shutdown_rx).await {
-            Continue(0) | Break(_) => break,
-            Continue(_) => {}
+        match shutdown
+            .run_until_cancelled(did_rx.recv_many(&mut dids, BATCH))
+            .await
+        {
+            Some(0) | None => break,
+            Some(_) => {}
         }
 
         // Convert Vec<Box<Did>> to a Vec<&Did>.
···
     Ok(())
 }
 
-#[tracing::instrument(skip(tx))]
-async fn shutdown_task(tx: watch::Sender<bool>) -> anyhow::Result<()> {
+#[tracing::instrument(skip(shutdown))]
+async fn shutdown_task(shutdown: CancellationToken) -> anyhow::Result<()> {
     tokio::signal::ctrl_c().await?;
     eprintln!();
     tracing::info!("shutdown signal received");
-    tx.send(true)?;
+    shutdown.cancel();
     Ok(())
 }
 

src/tap/channel.rs

···
     let (ack_tx, mut ack_rx) = mpsc::channel(capacity);
     let mut acks: Vec<_> = Default::default();
 
-    'outer: loop {
+    'outer: while !shutdown.is_cancelled() {
         let mut ping_inflight = false;
         let mut timeout = tokio::time::interval(TIMEOUT);
         timeout.tick().await;

src/util.rs

···
-use std::ops;
-
-use tokio::sync::watch;
-
-#[derive(Debug, thiserror::Error)]
-#[error("Shutdown signal received")]
-pub struct Shutdown;
-
-pub async fn with_shutdown<F: Future>(
-    f: F,
-    rx: &mut watch::Receiver<bool>,
-) -> ops::ControlFlow<Shutdown, F::Output> {
-    use ops::ControlFlow::{Break, Continue};
-
-    tokio::select! {
-        result = f => Continue(result),
-        Ok(_) = rx.wait_for(|&v| v) => Break(Shutdown),
-    }
-}