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
use crate::util::data_structures::HashMap;
use futures::{FutureExt, future::LocalBoxFuture, stream::FuturesUnordered};
use std::{hash::Hash, ops::Deref, task::Poll};
/// A local (!Send) adapter for caching and executing an async method
/// from a non-async context.
///
/// The `self_parameter`, `key`, and successful (Ok) results must all be cheap to `clone`.
///
/// Ensures at most one in-flight computation per key. Results are:
/// - cached on success
/// - not retained on error
pub struct LocalPollAdapter<'a, S, K, R> {
pool: FuturesUnordered<LocalBoxFuture<'a, (K, R)>>,
cache: HashMap<K, Poll<R>>,
self_parameter: S,
}
impl<'a, S, K, V, E> LocalPollAdapter<'a, S, K, Result<V, E>>
where
S: Clone + Deref + 'a,
K: Clone + Hash + Eq + 'a,
V: Clone,
{
pub fn new(self_parameter: S) -> Self {
Self {
pool: FuturesUnordered::new(),
cache: HashMap::default(),
self_parameter,
}
}
/// Polls the result for `key`, spawning work if needed.
///
/// If this function returns [`Poll::Pending`], call [`LocalPollAdapter::wait`]
/// to execute the work, then call this function again with the same key
/// to pick up the result.
///
/// Futures that complete immediately are not queued.
pub fn poll<F>(&mut self, f: F, key: K) -> Poll<Result<V, E>>
where
F: AsyncFn(&S::Target, &K) -> Result<V, E> + 'a,
{
match self.cache.get(&key) {
// We have a cached success value, clone it and return.
Some(Poll::Ready(Ok(v))) => return Poll::Ready(Ok(v.clone())),
// We have a cached error value, remove it and return.
// Errors are not Clone, so they are only stored once.
Some(Poll::Ready(Err(_))) => return self.cache.remove(&key).unwrap(),
// This key is already pending.
Some(Poll::Pending) => return Poll::Pending,
// Looks like we have work to do!
None => {}
}
// Created a pinned future that executes the function,
// returning the key and the result.
let mut future = {
let key = key.clone();
let self_parameter = self.self_parameter.clone();
async move {
let v = f(self_parameter.deref(), &key).await;
(key, v)
}
.boxed_local()
};
// Attempt to run the future immediately. If it has no `await` yields,
// it will return here.
if let Some((k, v)) = (&mut future).now_or_never() {
if let Ok(success) = &v {
// Only cache successful results.
self.cache.insert(k, Poll::Ready(Ok(success.clone())));
}
return Poll::Ready(v);
}
// Insert Pending into the cache so we avoid queuing the same future twice.
self.cache.insert(key.clone(), Poll::Pending);
// Add the future to the pending queue.
self.pool.push(future);
Poll::Pending
}
/// Returns the number of pending futures.
pub fn pending_count(&self) -> usize {
self.pool.len()
}
/// Run all pending futures. Returns true if there was no work to do.
pub fn wait(&mut self) -> bool {
let is_empty = self.pool.is_empty();
for (k, v) in crate::util::block_on_stream(&mut self.pool) {
*self
.cache
.get_mut(&k)
.expect("all pending work is in the cache") = Poll::Ready(v);
}
is_empty
}
}
#[cfg(test)]
mod tests {
use super::LocalPollAdapter;
use std::{rc::Rc, task::Poll};
/// Future that yields once.
fn yield_once() -> impl std::future::Future<Output = ()> {
let mut yielded = false;
std::future::poll_fn(move |cx| {
if yielded {
Poll::Ready(())
} else {
yielded = true;
cx.waker().wake_by_ref();
Poll::Pending
}
})
}
struct Thing {}
impl Thing {
async fn widen(&self, i: &i32) -> Result<i64, ()> {
if *i > 10 {
// Big numbers take longer to process (need to test futures that yield).
yield_once().await;
}
if *i % 2 != 0 {
// Odd numbers are not supported (need to test errors).
return Err(());
}
Ok(*i as i64)
}
}
/// Poll wrapper around `Thing`
struct PolledThing<'a> {
poller: LocalPollAdapter<'a, Rc<Thing>, i32, Result<i64, ()>>,
}
impl<'a> PolledThing<'a> {
fn new() -> Self {
Self {
poller: LocalPollAdapter::new(Rc::new(Thing {})),
}
}
// Non-async version of the widen method.
fn widen(&mut self, i: &i32) -> Poll<Result<i64, ()>> {
self.poller.poll(Thing::widen, i.clone())
}
fn wait(&mut self) -> bool {
self.poller.wait()
}
}
#[test]
fn immediate_success() {
let mut p = PolledThing::new();
assert_eq!(p.widen(&2), Poll::Ready(Ok(2)));
assert!(p.wait());
}
#[test]
fn immediate_error() {
let mut p = PolledThing::new();
assert_eq!(p.widen(&1), Poll::Ready(Err(())));
assert!(p.wait());
}
#[test]
fn deferred_error() {
let mut p = PolledThing::new();
assert_eq!(p.widen(&1001), Poll::Pending);
assert!(!p.wait());
assert_eq!(p.widen(&1001), Poll::Ready(Err(())));
assert!(p.wait());
// Errors are not cached
assert_eq!(p.widen(&1001), Poll::Pending);
assert!(!p.wait());
assert_eq!(p.widen(&1001), Poll::Ready(Err(())));
assert!(p.wait());
}
#[test]
fn deferred_success() {
let mut p = PolledThing::new();
assert_eq!(p.widen(&50), Poll::Pending);
assert!(!p.wait());
assert_eq!(p.widen(&50), Poll::Ready(Ok(50)));
assert!(p.wait());
// Success is cached.
assert_eq!(p.widen(&50), Poll::Ready(Ok(50)));
assert!(p.wait());
}
}