diff --git a/tokio-util/src/sync/cancellation_token.rs b/tokio-util/src/sync/cancellation_token.rs index 1b397c2bb..46c270a40 100644 --- a/tokio-util/src/sync/cancellation_token.rs +++ b/tokio-util/src/sync/cancellation_token.rs @@ -122,6 +122,26 @@ impl Clone for CancellationToken { } } +impl PartialEq for CancellationToken { + /// Checks if two tokens are equal in terms of their cancellation operation. + /// + /// Two tokens are considered equal if cancelling one will always also cancel the other and vice + /// versa. This is only true for cloned tokens and not for tokens in a parent-child + /// relationship. + fn eq(&self, other: &CancellationToken) -> bool { + Arc::ptr_eq(&self.inner, &other.inner) + } +} + +impl Eq for CancellationToken {} + +impl core::hash::Hash for CancellationToken { + #[inline] + fn hash(&self, state: &mut H) { + Arc::as_ptr(&self.inner).hash(state); + } +} + impl Drop for CancellationToken { fn drop(&mut self) { tree_node::decrease_handle_refcount(&self.inner); diff --git a/tokio-util/tests/sync_cancellation_token.rs b/tokio-util/tests/sync_cancellation_token.rs index 995890171..4a1f1f093 100644 --- a/tokio-util/tests/sync_cancellation_token.rs +++ b/tokio-util/tests/sync_cancellation_token.rs @@ -5,8 +5,10 @@ use tokio::sync::oneshot; use tokio_util::sync::{CancellationToken, WaitForCancellationFuture}; use core::future::Future; +use core::hash::Hash; use core::task::{Context, Poll}; use futures_test::task::new_count_waker; +use std::hash::{DefaultHasher, Hasher}; #[test] fn cancel_token() { @@ -563,3 +565,51 @@ fn run_until_cancelled_owned_test() { ); } } + +#[test] +fn cloned_cancellation_tokens_are_considered_equal() { + let token = CancellationToken::new(); + let token_clone = token.clone(); + + assert_eq!(token, token_clone); +} + +#[test] +fn child_cancellation_tokens_are_not_considered_equal() { + let token = CancellationToken::new(); + let token_clone = token.child_token(); + + assert_ne!(token, token_clone); +} + +#[test] +fn independent_cancellation_tokens_are_not_considered_equal() { + let token1 = CancellationToken::new(); + let token2 = CancellationToken::new(); + + assert_ne!(token1, token2); +} + +#[test] +fn cloned_cancellation_tokens_have_same_hash() { + let token1 = CancellationToken::new(); + let token2 = token1.clone(); + + let mut state1 = DefaultHasher::default(); + token1.hash(&mut state1); + let mut state2 = DefaultHasher::default(); + token2.hash(&mut state2); + assert_eq!(state1.finish(), state2.finish()); +} + +#[test] +fn different_cancellation_tokens_have_different_hash() { + let token1 = CancellationToken::new(); + let token2 = CancellationToken::new(); + + let mut state1 = DefaultHasher::default(); + token1.hash(&mut state1); + let mut state2 = DefaultHasher::default(); + token2.hash(&mut state2); + assert_ne!(state1.finish(), state2.finish()); +}