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
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
//! Observation of work that stays on its original worker.
//!
//! Local polls access the worker directly. The first foreign poll installs a
//! completion channel, which all subsequent polls use. The original identity
//! remains available for cancellation until observation ends.
use super::{
mailbox::{Cancel, Forward, Mailbox, Message},
runtime::Local,
};
use crate::Error;
use commonware_utils::channel::oneshot;
use std::{
future::Future,
mem,
pin::Pin,
sync::Weak,
task::{Context, Poll, Waker},
};
/// Inspection without cloning or invoking the observer's waker.
pub enum Observation<T> {
/// The registration has completed and relinquished its output.
Ready(T),
/// The installed waker already matches the current task.
Pending,
/// A new waker must be cloned outside the worker borrow.
Refresh,
}
/// Worker-owned operations needed by a movable observer.
pub trait Key: Copy + Unpin {
/// Terminal value produced by the owning subsystem.
type Output;
/// Inspect completion without running callbacks under the worker borrow.
fn observe(self, local: &mut Local, waker: &Waker) -> Observation<Self::Output>;
/// Replace the observer and return its displaced waker for deferred destruction.
fn refresh(self, local: &mut Local, waker: Waker) -> Option<Waker>;
/// Transfer observation to a completion channel on the owning worker.
fn forward(self, sender: oneshot::Sender<Result<Self::Output, Error>>) -> Forward;
/// Release observation, leaving retirement policy to the owning subsystem.
fn cancel(self) -> Cancel;
}
/// Registration ownership retained between polls.
enum State<K: Key> {
/// Registration awaiting observation, directly or through a forwarding channel.
Waiting {
/// Owning worker, without extending its lifetime.
mailbox: Weak<Mailbox>,
/// Identity retained until completion or cancellation.
key: K,
/// Permanent observation path once forwarding is requested.
receiver: Option<oneshot::Receiver<Result<K::Output, Error>>>,
},
/// Registration rejected by a closing worker, with closure not yet observed.
Closed,
/// Observation completed or released, with no cancellation identity retained.
Done,
}
/// A movable observer of one registration on its original worker.
pub struct Registration<K: Key> {
/// Retains the cancellation identity until observation completes or is released.
state: State<K>,
}
impl<K: Key> Registration<K> {
/// Observe work already registered with this worker.
pub const fn new(mailbox: Weak<Mailbox>, key: K) -> Self {
Self {
state: State::Waiting {
mailbox,
key,
receiver: None,
},
}
}
/// Report that the worker rejected registration during shutdown.
pub const fn closed() -> Self {
Self {
state: State::Closed,
}
}
/// Return the identity of a waiting registration.
#[cfg(test)]
pub fn key(&self) -> K {
let State::Waiting { key, .. } = self.state else {
panic!("registration already completed");
};
key
}
/// Release the registration before destroying a channel's observer waker.
fn release(&mut self) {
// Consume the identity before cleanup can unwind. Publish cancellation
// before dropping the receiver, whose waker destructor may run user code.
if let State::Waiting {
mailbox,
key,
receiver,
} = mem::replace(&mut self.state, State::Done)
{
Local::cancel(&mailbox, key.cancel());
drop(receiver);
}
}
}
impl<K: Key> Future for Registration<K> {
type Output = Result<K::Output, Error>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
// Pending polls retain the owner and key so Drop can still release the
// registration if polling unwinds.
let this = self.get_mut();
let (mailbox, key, receiver) = match &mut this.state {
State::Waiting {
mailbox,
key,
receiver,
} => (mailbox, *key, receiver),
State::Closed => {
this.state = State::Done;
return Poll::Ready(Err(Error::Closed));
}
State::Done => panic!("io_uring registration polled after completion"),
};
// Direct access requires the owning worker and an unforwarded registration.
// Once installed, the channel remains the observation path even on that worker.
let owner = receiver.is_none().then(|| Local::owner(mailbox)).flatten();
let Some(owner) = owner else {
// A missing or closed worker drops the sender, allowing the receiver
// to report closure without retaining the worker.
let receiver = receiver.get_or_insert_with(|| {
let (sender, receiver) = oneshot::channel();
if let Some(mailbox) = mailbox.upgrade() {
let _ = mailbox.send(Message::Forward(key.forward(sender)));
}
receiver
});
let result = std::task::ready!(Pin::new(receiver).poll(cx));
this.state = State::Done;
return Poll::Ready(result.unwrap_or(Err(Error::Closed)));
};
// A closing worker owns cleanup. Release observation after ending this
// borrow because cancellation reenters the worker.
let mut local = owner.borrow_mut();
if local.closing {
drop(local);
this.release();
return Poll::Ready(Err(Error::Closed));
}
// Only a changed pending waker needs a refresh. Ready results and matching
// wakers return without invoking observer callbacks.
match key.observe(&mut local, cx.waker()) {
Observation::Ready(output) => {
this.state = State::Done;
return Poll::Ready(Ok(output));
}
Observation::Pending => return Poll::Pending,
Observation::Refresh => {}
}
// Waker callbacks may reenter the runtime, so clone outside the borrow.
// A clone panic leaves the registration identity available for
// cancellation. A clone can also admit a durability request, whose
// admission polls the notified permit waits and can complete this
// registration if its permit was granted since the last service, so
// observe again before installing the clone. Defer destruction of
// whichever waker is left over.
drop(local);
let waker = cx.waker().clone();
let mut local = owner.borrow_mut();
let (poll, unused) = match key.observe(&mut local, cx.waker()) {
Observation::Ready(output) => {
this.state = State::Done;
(Poll::Ready(Ok(output)), Some(waker))
}
Observation::Pending => (Poll::Pending, Some(waker)),
Observation::Refresh => (Poll::Pending, key.refresh(&mut local, waker)),
};
local.deferred.drops.extend(unused);
poll
}
}
impl<K: Key> Drop for Registration<K> {
fn drop(&mut self) {
self.release();
}
}