Skip to content

Commit 4910799

Browse files
tones111taiki-e
authored andcommitted
provide a non-destructive mechanism to determine if a sink/stream are paired
1 parent 40eed5c commit 4910799

File tree

2 files changed

+86
-1
lines changed

2 files changed

+86
-1
lines changed

futures-util/src/lock/bilock.rs

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -149,14 +149,19 @@ impl<T> BiLock<T> {
149149
BiLockAcquire { bilock: self }
150150
}
151151

152+
/// Returns `true` only if the other `BiLock<T>` originated from the same call to `BiLock::new`.
153+
pub fn is_pair_of(&self, other: &Self) -> bool {
154+
Arc::ptr_eq(&self.arc, &other.arc)
155+
}
156+
152157
/// Attempts to put the two "halves" of a `BiLock<T>` back together and
153158
/// recover the original value. Succeeds only if the two `BiLock<T>`s
154159
/// originated from the same call to `BiLock::new`.
155160
pub fn reunite(self, other: Self) -> Result<T, ReuniteError<T>>
156161
where
157162
T: Unpin,
158163
{
159-
if Arc::ptr_eq(&self.arc, &other.arc) {
164+
if self.is_pair_of(&other) {
160165
drop(other);
161166
let inner = Arc::try_unwrap(self.arc)
162167
.ok()

futures-util/src/stream/stream/split.rs

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,13 @@ pub struct SplitStream<S>(BiLock<S>);
1515

1616
impl<S> Unpin for SplitStream<S> {}
1717

18+
impl<S> SplitStream<S> {
19+
/// Returns `true` if the `SplitStream<S>` and `SplitSink<S>` originate from the same call to `StreamExt::split`.
20+
pub fn is_pair_of<Item>(&self, other: &SplitSink<S, Item>) -> bool {
21+
other.is_pair_of(&self)
22+
}
23+
}
24+
1825
impl<S: Unpin> SplitStream<S> {
1926
/// Attempts to put the two "halves" of a split `Stream + Sink` back
2027
/// together. Succeeds only if the `SplitStream<S>` and `SplitSink<S>` are
@@ -60,6 +67,13 @@ impl<S: Sink<Item> + Unpin, Item> SplitSink<S, Item> {
6067
}
6168
}
6269

70+
impl<S, Item> SplitSink<S, Item> {
71+
/// Returns `true` if the `SplitStream<S>` and `SplitSink<S>` originate from the same call to `StreamExt::split`.
72+
pub fn is_pair_of(&self, other: &SplitStream<S>) -> bool {
73+
self.lock.is_pair_of(&other.0)
74+
}
75+
}
76+
6377
impl<S: Sink<Item>, Item> SplitSink<S, Item> {
6478
fn poll_flush_slot(
6579
mut inner: Pin<&mut S>,
@@ -142,3 +156,69 @@ impl<T, Item> fmt::Display for ReuniteError<T, Item> {
142156

143157
#[cfg(feature = "std")]
144158
impl<T: core::any::Any, Item> std::error::Error for ReuniteError<T, Item> {}
159+
160+
#[cfg(test)]
161+
mod tests {
162+
use super::*;
163+
use crate::{sink::Sink, stream::StreamExt};
164+
use core::marker::PhantomData;
165+
166+
struct NopStream<Item> {
167+
phantom: PhantomData<Item>,
168+
}
169+
170+
impl<Item> Stream for NopStream<Item> {
171+
type Item = Item;
172+
173+
fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
174+
todo!()
175+
}
176+
}
177+
178+
impl<Item> Sink<Item> for NopStream<Item> {
179+
type Error = ();
180+
181+
fn poll_ready(
182+
self: Pin<&mut Self>,
183+
_cx: &mut Context<'_>,
184+
) -> Poll<Result<(), Self::Error>> {
185+
todo!()
186+
}
187+
188+
fn start_send(self: Pin<&mut Self>, _item: Item) -> Result<(), Self::Error> {
189+
todo!()
190+
}
191+
192+
fn poll_flush(
193+
self: Pin<&mut Self>,
194+
_cx: &mut Context<'_>,
195+
) -> Poll<Result<(), Self::Error>> {
196+
todo!()
197+
}
198+
199+
fn poll_close(
200+
self: Pin<&mut Self>,
201+
_cx: &mut Context<'_>,
202+
) -> Poll<Result<(), Self::Error>> {
203+
todo!()
204+
}
205+
}
206+
207+
#[test]
208+
fn test_pairing() {
209+
let s1 = NopStream::<()> { phantom: PhantomData };
210+
let (sink1, stream1) = s1.split();
211+
assert!(sink1.is_pair_of(&stream1));
212+
assert!(stream1.is_pair_of(&sink1));
213+
214+
let s2 = NopStream::<()> { phantom: PhantomData };
215+
let (sink2, stream2) = s2.split();
216+
assert!(sink2.is_pair_of(&stream2));
217+
assert!(stream2.is_pair_of(&sink2));
218+
219+
assert!(!sink1.is_pair_of(&stream2));
220+
assert!(!stream1.is_pair_of(&sink2));
221+
assert!(!sink2.is_pair_of(&stream1));
222+
assert!(!stream2.is_pair_of(&sink1));
223+
}
224+
}

0 commit comments

Comments
 (0)