From b1542a6153493d078d49b465c1498deac9141e68 Mon Sep 17 00:00:00 2001
From: Paul Ellenbogen <PaulEllenbogen@gmail.com>
Date: Mon, 28 Apr 2025 02:56:20 -0700
Subject: [PATCH] fix: fix resource leak (#136)

When a running service is dropped, its connections should be closed and its resources should be
cleaned up.

This changes the default from leaving services running to closing them. Library users who want
their services to remain running should call `waiting` on their running service, in a new task
where necessary.

Co-authored-by: Paul Ellenbogen <pe5@cs.princeton.edu>
---
 crates/rmcp/src/service.rs              | 13 +++++++------
 examples/servers/src/generic_service.rs |  2 +-
 2 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/crates/rmcp/src/service.rs b/crates/rmcp/src/service.rs
index ea05edd..54aa081 100644
--- a/crates/rmcp/src/service.rs
+++ b/crates/rmcp/src/service.rs
@@ -21,7 +21,7 @@ mod server;
 pub use server::*;
 #[cfg(feature = "tower")]
 mod tower;
-use tokio_util::sync::CancellationToken;
+use tokio_util::sync::{CancellationToken, DropGuard};
 #[cfg(feature = "tower")]
 pub use tower::*;
 use tracing::instrument;
@@ -429,8 +429,8 @@ pub struct RunningService<R: ServiceRole, S: Service<R>> {
     service: Arc<S>,
     peer: Peer<R>,
     handle: tokio::task::JoinHandle<QuitReason>,
-    /// cancellation token
-    ct: CancellationToken,
+    /// cancellation token with drop guard
+    dg: DropGuard,
 }
 impl<R: ServiceRole, S: Service<R>> Deref for RunningService<R, S> {
     type Target = Peer<R>;
@@ -453,8 +453,9 @@ impl<R: ServiceRole, S: Service<R>> RunningService<R, S> {
         self.handle.await
     }
     pub async fn cancel(self) -> Result<QuitReason, tokio::task::JoinError> {
-        self.ct.cancel();
-        self.waiting().await
+        let RunningService { dg, handle, .. } = self;
+        dg.disarm().cancel();
+        handle.await
     }
 }
 
@@ -759,6 +760,6 @@ where
         service,
         peer: peer_return,
         handle,
-        ct,
+        dg: ct.drop_guard(),
     })
 }
diff --git a/examples/servers/src/generic_service.rs b/examples/servers/src/generic_service.rs
index 546621e..258ed21 100644
--- a/examples/servers/src/generic_service.rs
+++ b/examples/servers/src/generic_service.rs
@@ -13,6 +13,6 @@ async fn main() -> Result<(), Box<dyn Error>> {
 
     let io = (tokio::io::stdin(), tokio::io::stdout());
 
-    serve_server(generic_service, io).await?;
+    serve_server(generic_service, io).await?.waiting().await?;
     Ok(())
 }