Skip to content

Commit 633d424

Browse files
feat: add cancellation_token method to RunningService (#218)
1 parent 0d52b7d commit 633d424

File tree

1 file changed

+17
-2
lines changed

1 file changed

+17
-2
lines changed

crates/rmcp/src/service.rs

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -424,14 +424,14 @@ pub struct RunningService<R: ServiceRole, S: Service<R>> {
424424
service: Arc<S>,
425425
peer: Peer<R>,
426426
handle: tokio::task::JoinHandle<QuitReason>,
427-
/// cancellation token with drop guard
427+
cancellation_token: CancellationToken,
428428
dg: DropGuard,
429429
}
430430
impl<R: ServiceRole, S: Service<R>> Deref for RunningService<R, S> {
431431
type Target = Peer<R>;
432432

433433
fn deref(&self) -> &Self::Target {
434-
self.peer()
434+
&self.peer
435435
}
436436
}
437437

@@ -444,6 +444,11 @@ impl<R: ServiceRole, S: Service<R>> RunningService<R, S> {
444444
pub fn service(&self) -> &S {
445445
self.service.as_ref()
446446
}
447+
#[inline]
448+
pub fn cancellation_token(&self) -> RunningServiceCancellationToken {
449+
RunningServiceCancellationToken(self.cancellation_token.clone())
450+
}
451+
#[inline]
447452
pub async fn waiting(self) -> Result<QuitReason, tokio::task::JoinError> {
448453
self.handle.await
449454
}
@@ -454,6 +459,15 @@ impl<R: ServiceRole, S: Service<R>> RunningService<R, S> {
454459
}
455460
}
456461

462+
// use a wrapper type so we can tweak the implementation if needed
463+
pub struct RunningServiceCancellationToken(CancellationToken);
464+
465+
impl RunningServiceCancellationToken {
466+
pub fn cancel(self) {
467+
self.0.cancel();
468+
}
469+
}
470+
457471
#[derive(Debug)]
458472
pub enum QuitReason {
459473
Cancelled,
@@ -801,6 +815,7 @@ where
801815
service,
802816
peer: peer_return,
803817
handle,
818+
cancellation_token: ct.clone(),
804819
dg: ct.drop_guard(),
805820
}
806821
}

0 commit comments

Comments
 (0)