sync: implement PartialEq and Eq for CancellationToken (#8110)

Co-authored-by: Martin Tzvetanov Grigorov <[email protected]>

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<H: core::hash::Hasher>(&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);

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());
+}