1use crate::{
2 action::{ActionValue, WireValue},
3 error::{ResourceKind, RuntimeError},
4};
5use hataori_runtime_foundation::protocol::{ActionId, DomainId, LocalityId, RequestId, TraceId};
6use std::{
7 collections::HashMap,
8 future::Future,
9 marker::PhantomData,
10 pin::Pin,
11 sync::{
12 atomic::{AtomicBool, Ordering},
13 Arc, Mutex,
14 },
15 task::{Context, Poll, Waker},
16 time::{Duration, Instant},
17};
18
19struct PromiseState {
20 result: Option<Result<ActionValue, RuntimeError>>,
21 waker: Option<Waker>,
22}
23
24pub(crate) struct Promise {
25 state: Mutex<PromiseState>,
26 done: AtomicBool,
27}
28
29impl Promise {
30 pub(crate) fn new() -> Self {
31 Self {
32 state: Mutex::new(PromiseState {
33 result: None,
34 waker: None,
35 }),
36 done: AtomicBool::new(false),
37 }
38 }
39
40 pub(crate) fn complete(&self, result: Result<ActionValue, RuntimeError>) {
41 let mut state = self.state.lock().unwrap();
42 if state.result.is_none() && !self.done.load(Ordering::Acquire) {
43 state.result = Some(result);
44 self.done.store(true, Ordering::Release);
45 if let Some(waker) = state.waker.take() {
46 waker.wake();
47 }
48 }
49 }
50
51 fn poll(&self, context: &mut Context<'_>) -> Poll<Result<ActionValue, RuntimeError>> {
52 let mut state = self.state.lock().unwrap();
53 if let Some(result) = state.result.take() {
54 Poll::Ready(result)
55 } else {
56 state.waker = Some(context.waker().clone());
57 Poll::Pending
58 }
59 }
60
61 fn done(&self) -> bool {
62 self.done.load(Ordering::Acquire)
63 }
64}
65
66pub(crate) struct PendingEntry {
67 pub promise: Arc<Promise>,
68 pub destination: LocalityId,
69 pub action: ActionId,
70 pub domain: DomainId,
71 pub trace_id: Option<TraceId>,
72 pub expires_at: Instant,
73 pub deadline: Duration,
74 pub cancel_token: Option<Arc<AtomicBool>>,
75}
76
77#[derive(Default)]
78struct PendingState {
79 entries: HashMap<RequestId, PendingEntry>,
80}
81
82pub(crate) struct PendingTable {
83 state: Mutex<PendingState>,
84 limit: usize,
85}
86
87impl PendingTable {
88 pub(crate) fn new(limit: usize) -> Self {
89 Self {
90 state: Mutex::new(PendingState::default()),
91 limit,
92 }
93 }
94
95 pub(crate) fn insert(
96 &self,
97 request: RequestId,
98 entry: PendingEntry,
99 ) -> Result<(), RuntimeError> {
100 let mut state = self.state.lock().unwrap();
101 if state.entries.len() >= self.limit {
102 return Err(RuntimeError::ResourceExhausted {
103 resource: ResourceKind::PendingCalls,
104 limit: self.limit,
105 });
106 }
107 state.entries.insert(request, entry);
108 Ok(())
109 }
110
111 pub(crate) fn take_checked(
112 &self,
113 request: RequestId,
114 source: LocalityId,
115 action: ActionId,
116 ) -> Result<Option<PendingEntry>, RuntimeError> {
117 let mut state = self.state.lock().unwrap();
118 let Some(entry) = state.entries.get(&request) else {
119 return Ok(None);
120 };
121 if entry.destination != source || entry.action != action {
122 return Err(RuntimeError::Protocol(
123 "result source or action does not match pending request".into(),
124 ));
125 }
126 Ok(state.entries.remove(&request))
127 }
128
129 pub(crate) fn remove(&self, request: RequestId) -> Option<PendingEntry> {
130 self.state.lock().unwrap().entries.remove(&request)
131 }
132
133 pub(crate) fn expired(&self, now: Instant) -> Vec<(RequestId, PendingEntry)> {
134 let mut state = self.state.lock().unwrap();
135 let expired: Vec<_> = state
136 .entries
137 .iter()
138 .filter_map(|(request, entry)| (entry.expires_at <= now).then_some(*request))
139 .collect();
140 expired
141 .into_iter()
142 .filter_map(|request| state.entries.remove(&request).map(|entry| (request, entry)))
143 .collect()
144 }
145
146 pub(crate) fn drain(&self) -> Vec<(RequestId, PendingEntry)> {
147 self.state.lock().unwrap().entries.drain().collect()
148 }
149
150 pub(crate) fn len(&self) -> usize {
151 self.state.lock().unwrap().entries.len()
152 }
153}
154
155#[must_use = "dropping a remote future cancels its pending request"]
156pub struct RemoteFuture<T: WireValue> {
157 pub(crate) request: RequestId,
158 promise: Arc<Promise>,
159 cancel: Arc<dyn Fn(RequestId) + Send + Sync>,
160 finished: bool,
161 marker: PhantomData<T>,
162}
163
164impl<T: WireValue> std::fmt::Debug for RemoteFuture<T> {
165 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
166 formatter
167 .debug_struct("RemoteFuture")
168 .field("request", &self.request)
169 .field("finished", &self.finished)
170 .finish()
171 }
172}
173
174impl<T: WireValue> RemoteFuture<T> {
175 pub(crate) fn new(
176 request: RequestId,
177 promise: Arc<Promise>,
178 cancel: Arc<dyn Fn(RequestId) + Send + Sync>,
179 ) -> Self {
180 Self {
181 request,
182 promise,
183 cancel,
184 finished: false,
185 marker: PhantomData,
186 }
187 }
188
189 pub const fn request_id(&self) -> RequestId {
190 self.request
191 }
192}
193
194impl<T: WireValue> Unpin for RemoteFuture<T> {}
195
196impl<T: WireValue> Future for RemoteFuture<T> {
197 type Output = Result<T, RuntimeError>;
198
199 fn poll(self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Self::Output> {
200 let this = self.get_mut();
201 if this.finished {
202 panic!("RemoteFuture polled after completion");
203 }
204 match this.promise.poll(context) {
205 Poll::Ready(Ok(ActionValue::Encoded(segments))) => {
206 this.finished = true;
207 Poll::Ready(T::decode(segments).map_err(|error| {
208 RuntimeError::Protocol(format!("action output decode failed: {error}"))
209 }))
210 }
211 Poll::Ready(Ok(ActionValue::Typed(value))) => {
212 this.finished = true;
213 Poll::Ready(
214 value
215 .downcast::<T>()
216 .map(|value| *value)
217 .map_err(|_| RuntimeError::Protocol("typed action output mismatch".into())),
218 )
219 }
220 Poll::Ready(Err(error)) => {
221 this.finished = true;
222 Poll::Ready(Err(error))
223 }
224 Poll::Pending => Poll::Pending,
225 }
226 }
227}
228
229impl<T: WireValue> Drop for RemoteFuture<T> {
230 fn drop(&mut self) {
231 if !self.finished && !self.promise.done() {
232 (self.cancel)(self.request);
233 }
234 }
235}
236
237#[cfg(test)]
238mod tests;