From 5bf78d77ada685826b33fd62e69df179f3de8a1c Mon Sep 17 00:00:00 2001
From: Maarten de Vries <maarten@de-vri.es>
Date: Tue, 21 Jan 2020 01:50:31 +0100
Subject: [PATCH] Add a method to test if split streams come from the same
 stream. (#1762)

* Add a method to test if split streams come from the same stream.

The exposed stream ID can also be used as key in associative containers.

* Document the fact that split stream IDs can dangle.
---
 tokio/src/io/split.rs   | 42 ++++++++++++++++++++++++++++++++++++++++-
 tokio/tests/io_split.rs | 10 ++++++++++
 2 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/tokio/src/io/split.rs b/tokio/src/io/split.rs
index 7f21533ee4e..e56e6ee95dc 100644
--- a/tokio/src/io/split.rs
+++ b/tokio/src/io/split.rs
@@ -51,6 +51,18 @@ cfg_io_util! {
     }
 }
 
+/// An opaque ID for the parent stream of a split half.
+///
+/// If you keep a `SplitStreamId` around after both halves have been dropped or reunited,
+/// the stream ID is dangling.
+/// The same ID may then be used for other split streams.
+/// To avoid this, do not keep `SplitStreamId` around after both half have been dropped.
+///
+/// Note that it is still impossible to unsplit two halves from a different stream,
+/// since at-least one half has not been dropped in that scenario.
+#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
+pub struct SplitStreamId(usize);
+
 struct Inner<T> {
     locked: AtomicBool,
     stream: UnsafeCell<T>,
@@ -61,14 +73,28 @@ struct Guard<'a, T> {
 }
 
 impl<T> ReadHalf<T> {
+    /// Get an opaque ID for the parent stream.
+    ///
+    /// This can be used to check if two halves have been split from the
+    /// same stream.
+    /// The stream ID can also be used as key in associative containers.
+    ///
+    /// Note that stream IDs may dangle when both halves are dropped.
+    /// See [`SplitStreamId`] for more information.
+    pub fn stream_id(&self) -> SplitStreamId {
+        SplitStreamId(&*self.inner as *const Inner<T> as usize)
+    }
+
     /// Reunite with a previously split `WriteHalf`.
     ///
     /// # Panics
     ///
     /// If this `ReadHalf` and the given `WriteHalf` do not originate from the
     /// same `split` operation this method will panic.
+    /// This can be checked ahead of time by comparing the stream ID
+    /// of the two halves.
     pub fn unsplit(self, wr: WriteHalf<T>) -> T {
-        if Arc::ptr_eq(&self.inner, &wr.inner) {
+        if self.stream_id() == wr.stream_id() {
             drop(wr);
 
             let inner = Arc::try_unwrap(self.inner)
@@ -82,6 +108,20 @@ impl<T> ReadHalf<T> {
     }
 }
 
+impl<T> WriteHalf<T> {
+    /// Get an opaque ID for the parent stream.
+    ///
+    /// This can be used to check if two halves have been split from the
+    /// same stream.
+    /// The stream ID can also be used as key in associative containers.
+    ///
+    /// Note that stream IDs may dangle when both halves are dropped.
+    /// See [`SplitStreamId`] for more information.
+    pub fn stream_id(&self) -> SplitStreamId {
+        SplitStreamId(&*self.inner as *const Inner<T> as usize)
+    }
+}
+
 impl<T: AsyncRead> AsyncRead for ReadHalf<T> {
     fn poll_read(
         self: Pin<&mut Self>,
diff --git a/tokio/tests/io_split.rs b/tokio/tests/io_split.rs
index 0080f46af59..f6326ae9b7e 100644
--- a/tokio/tests/io_split.rs
+++ b/tokio/tests/io_split.rs
@@ -45,6 +45,16 @@ fn is_send_and_sync() {
     assert_bound::<WriteHalf<RW>>();
 }
 
+#[test]
+fn split_stream_id() {
+    let (r1, w1) = split(RW);
+    let (r2, w2) = split(RW);
+    assert_eq!(r1.stream_id(), w1.stream_id());
+    assert_eq!(r1.stream_id(), w1.stream_id());
+    assert_ne!(r1.stream_id(), w2.stream_id());
+    assert_ne!(r2.stream_id(), w1.stream_id());
+}
+
 #[test]
 fn unsplit_ok() {
     let (r, w) = split(RW);