spawn-ready: Abort background tasks when the service is dropped (#581)

The `SpawnReady` service spawns a background task whenever the inner
service is not ready; but when the `SpawnReady` service is dropped, this
task continues to run until it becomes ready (at which point the ready
service will be dropped). This can cause resource leaks when the inner
service never becomes ready.

This change adds a `Drop` implementation for the `SpawnReady` service
that aborts the background task when one is present.

Co-authored-by: Eliza Weisman <eliza@buoyant.io>

tower/src/spawn_ready/service.rs

@@ -33,6 +33,14 @@ impl<S> SpawnReady<S> {
     }
 }
 
+impl<S> Drop for SpawnReady<S> {
+    fn drop(&mut self) {
+        if let Inner::Future(ref mut task) = self.inner {
+            task.abort();
+        }
+    }
+}
+
 impl<S, Req> Service<Req> for SpawnReady<S>
 where
     Req: 'static,

tower/tests/spawn_ready/main.rs

@@ -5,6 +5,7 @@ mod support;
 use tokio::time;
 use tokio_test::{assert_pending, assert_ready, assert_ready_err, assert_ready_ok};
 use tower::spawn_ready::{SpawnReady, SpawnReadyLayer};
+use tower::util::ServiceExt;
 use tower_test::mock;
 
 #[tokio::test(flavor = "current_thread")]
@@ -46,7 +47,6 @@ async fn when_inner_fails() {
 
 #[tokio::test(flavor = "current_thread")]
 async fn propagates_trace_spans() {
-    use tower::util::ServiceExt;
     use tracing::Instrument;
 
     let _t = support::trace_init();
@@ -59,3 +59,27 @@ async fn propagates_trace_spans() {
 
     result.await.expect("service panicked").expect("failed");
 }
+
+#[cfg(test)]
+#[tokio::test(flavor = "current_thread")]
+async fn abort_on_drop() {
+    let (mock, mut handle) = mock::pair::<(), ()>();
+    let mut svc = SpawnReady::new(mock);
+    handle.allow(0);
+
+    // Drive the service to readiness until we signal a drop.
+    let (drop_tx, drop_rx) = tokio::sync::oneshot::channel();
+    let mut task = tokio_test::task::spawn(async move {
+        tokio::select! {
+            _ = drop_rx => {}
+            _ = svc.ready() => unreachable!("Service must not become ready"),
+        }
+    });
+    assert_pending!(task.poll());
+    assert_pending!(handle.poll_request());
+
+    // End the task and ensure that the inner service has been dropped.
+    assert!(drop_tx.send(()).is_ok());
+    tokio_test::assert_ready!(task.poll());
+    assert!(tokio_test::assert_ready!(handle.poll_request()).is_none());
+}