1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
use crate::channel::error::Result;
use std::future::Future;
use std::pin::Pin;
use std::sync::atomic::{Ordering, fence};
use std::task::{Context, Poll};
use crate::channel::CHANNEL_STORE_LOAD_ORDER;
/// Future implementation for zero-cost async receive
pub struct RecvFuture<'a, T> {
pub(super) receiver: &'a super::recv::HybridReceiver<T>,
pub(super) id: Option<u64>,
}
impl<T: Send> Future for RecvFuture<'_, T> {
type Output = Result<T>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
// Fast path check
if let Some(value) = self.receiver.ring.try_consume() {
return Poll::Ready(Ok(value));
}
if self.receiver.closed.load(Ordering::Acquire) {
return Poll::Ready(Err(crate::channel::error::ChannelError::Closed));
}
let this = self.get_mut();
let mut wakers = this
.receiver
.async_wakers
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
// Register FIRST, then re-check. The sender's counter gate runs
// without this lock, so registering after the re-check leaves a plain
// interleaving hole (no reordering required): re-check sees an empty
// ring, the sender produces and reads a zero `waker_count`, then the
// registration lands — nobody ever wakes the task. Registration
// before the fenced re-check closes it: either the re-check observes
// the message, or the registration is ordered before the sender's
// gate load and the sender drains-and-wakes this waker.
let waker = cx.waker();
let id = if let Some(id) = this.id {
if let Some(pos) = wakers.iter().position(|(w_id, _)| *w_id == id) {
wakers[pos].1.clone_from(waker);
} else {
// A previous send drained this future's entry; re-register
// under the same ID and re-count it.
wakers.push((id, waker.clone()));
this.receiver
.waker_count
.fetch_add(1, CHANNEL_STORE_LOAD_ORDER);
}
id
} else {
let id = this.receiver.next_id.fetch_add(1, Ordering::Relaxed);
wakers.push((id, waker.clone()));
this.id = Some(id);
this.receiver
.waker_count
.fetch_add(1, CHANNEL_STORE_LOAD_ORDER);
id
};
// Dekker fence between the registration above and the re-checks
// below; pairs with the fence in `notify_consumers` between the
// sender's publication and its counter gate loads (the same
// store-buffer pair `moirai-sync`'s `FutexMutex` documents).
fence(CHANNEL_STORE_LOAD_ORDER);
let ready = match this.receiver.ring.try_consume() {
Some(value) => Some(Ok(value)),
_ => {
if this.receiver.closed.load(Ordering::Acquire) {
Some(Err(crate::channel::error::ChannelError::Closed))
} else {
None
}
}
};
if let Some(output) = ready {
// Deregister while the lock is still held: a resolved future must
// not leave a counted waker behind for senders to wake.
if let Some(pos) = wakers.iter().position(|(w_id, _)| *w_id == id) {
wakers.remove(pos);
this.receiver
.waker_count
.fetch_sub(1, CHANNEL_STORE_LOAD_ORDER);
}
return Poll::Ready(output);
}
Poll::Pending
}
}
impl<T> Drop for RecvFuture<'_, T> {
fn drop(&mut self) {
if let Some(id) = self.id {
let mut wakers = self
.receiver
.async_wakers
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if let Some(pos) = wakers.iter().position(|(w_id, _)| *w_id == id) {
wakers.remove(pos);
self.receiver
.waker_count
.fetch_sub(1, CHANNEL_STORE_LOAD_ORDER);
}
}
}
}