1use std::collections::{BTreeMap, BTreeSet};
2use std::fmt::Debug;
3use std::future::{Future, ready};
4use std::mem;
5use std::pin::Pin;
6
7use fedimint_connectors::ServerResult;
8use fedimint_connectors::error::ServerError;
9use fedimint_core::task::{MaybeSend, MaybeSync};
10use fedimint_core::{NumPeers, PeerId, maybe_add_send, maybe_add_send_sync};
11
12pub trait QueryStrategy<IR, OR = IR> {
19 fn process(
24 &mut self,
25 peer_id: PeerId,
26 response: IR,
27 ) -> impl Future<Output = QueryStep<OR>> + MaybeSend;
28
29 fn process_error(&mut self, _peer_id: PeerId, _error: &ServerError) -> QueryStep<OR> {
36 QueryStep::Continue
37 }
38}
39
40#[derive(Debug)]
46pub enum QueryStep<R> {
47 Retry(BTreeSet<PeerId>),
49 Continue,
51 Success(R),
53 Failure(ServerError),
55}
56
57pub struct FilterMap<R, T> {
60 filter_map: Box<maybe_add_send_sync!(dyn Fn(R) -> ServerResult<T>)>,
61}
62
63impl<R, T> FilterMap<R, T> {
64 pub fn new(
65 filter_map: impl Fn(R) -> ServerResult<T> + MaybeSend + MaybeSync + 'static,
66 ) -> Self {
67 Self {
68 filter_map: Box::new(filter_map),
69 }
70 }
71}
72
73impl<R, T: MaybeSend> QueryStrategy<R, T> for FilterMap<R, T> {
74 fn process(
75 &mut self,
76 _peer: PeerId,
77 response: R,
78 ) -> impl Future<Output = QueryStep<T>> + MaybeSend {
79 ready(match (self.filter_map)(response) {
80 Ok(value) => QueryStep::Success(value),
81 Err(e) => QueryStep::Failure(e),
82 })
83 }
84}
85
86pub struct FilterMapThreshold<R, T> {
89 filter_map: Box<maybe_add_send_sync!(dyn Fn(PeerId, R) -> FilterMapFuture<T>)>,
90 filtered_responses: BTreeMap<PeerId, T>,
91 threshold: usize,
92}
93
94type FilterMapFuture<T> = Pin<Box<maybe_add_send!(dyn Future<Output = ServerResult<T>>)>>;
95
96impl<R, T> FilterMapThreshold<R, T> {
97 pub fn new<Fut>(
100 verifier: impl Fn(PeerId, R) -> Fut + MaybeSend + MaybeSync + 'static,
101 num_peers: NumPeers,
102 ) -> Self
103 where
104 Fut: Future<Output = ServerResult<T>> + MaybeSend + 'static,
105 {
106 Self {
107 filter_map: Box::new(move |peer, response| Box::pin(verifier(peer, response))),
108 filtered_responses: BTreeMap::new(),
109 threshold: num_peers.threshold(),
110 }
111 }
112}
113
114impl<R: MaybeSend, T: MaybeSend> QueryStrategy<R, BTreeMap<PeerId, T>>
115 for FilterMapThreshold<R, T>
116{
117 async fn process(&mut self, peer: PeerId, response: R) -> QueryStep<BTreeMap<PeerId, T>> {
118 match (self.filter_map)(peer, response).await {
119 Ok(response) => {
120 self.filtered_responses.insert(peer, response);
121
122 if self.filtered_responses.len() == self.threshold {
123 QueryStep::Success(mem::take(&mut self.filtered_responses))
124 } else {
125 QueryStep::Continue
126 }
127 }
128 Err(e) => QueryStep::Failure(e),
129 }
130 }
131}
132
133pub struct ThresholdConsensus<R> {
138 responses: BTreeMap<PeerId, R>,
139 retry: BTreeSet<PeerId>,
140 threshold: usize,
141}
142
143impl<R> ThresholdConsensus<R> {
144 pub fn new(num_peers: NumPeers) -> Self {
145 Self {
146 responses: BTreeMap::new(),
147 retry: BTreeSet::new(),
148 threshold: num_peers.threshold(),
149 }
150 }
151}
152
153impl<R: Eq + Clone + MaybeSend> QueryStrategy<R> for ThresholdConsensus<R> {
154 fn process(
155 &mut self,
156 peer: PeerId,
157 response: R,
158 ) -> impl Future<Output = QueryStep<R>> + MaybeSend {
159 self.responses.insert(peer, response.clone());
160
161 if self.responses.values().filter(|r| **r == response).count() == self.threshold {
162 return ready(QueryStep::Success(response));
163 }
164
165 assert!(self.retry.insert(peer));
166
167 ready(if self.retry.len() == self.threshold {
168 QueryStep::Retry(mem::take(&mut self.retry))
169 } else {
170 QueryStep::Continue
171 })
172 }
173}
174
175pub struct ThresholdAgreement<R> {
187 responses: BTreeMap<PeerId, R>,
188 errors: usize,
189 threshold: usize,
190 total: usize,
191}
192
193impl<R> ThresholdAgreement<R> {
194 pub fn new(num_peers: NumPeers) -> Self {
195 Self {
196 responses: BTreeMap::new(),
197 errors: 0,
198 threshold: num_peers.threshold(),
199 total: num_peers.total(),
200 }
201 }
202
203 fn diverged(&mut self) -> Option<QueryStep<Result<R, BTreeMap<PeerId, R>>>> {
206 if self.responses.len() + self.errors < self.total {
207 return None;
208 }
209
210 if self.responses.len() < self.threshold {
214 return None;
215 }
216
217 Some(QueryStep::Success(Err(mem::take(&mut self.responses))))
218 }
219}
220
221impl<R: Eq + Clone + MaybeSend> QueryStrategy<R, Result<R, BTreeMap<PeerId, R>>>
222 for ThresholdAgreement<R>
223{
224 fn process(
225 &mut self,
226 peer: PeerId,
227 response: R,
228 ) -> impl Future<Output = QueryStep<Result<R, BTreeMap<PeerId, R>>>> + MaybeSend {
229 self.responses.insert(peer, response.clone());
230
231 if self.responses.values().filter(|r| **r == response).count() == self.threshold {
232 return ready(QueryStep::Success(Ok(response)));
233 }
234
235 ready(self.diverged().unwrap_or(QueryStep::Continue))
236 }
237
238 fn process_error(
239 &mut self,
240 _peer: PeerId,
241 _error: &ServerError,
242 ) -> QueryStep<Result<R, BTreeMap<PeerId, R>>> {
243 self.errors += 1;
244
245 self.diverged().unwrap_or(QueryStep::Continue)
246 }
247}
248
249#[cfg(test)]
250mod tests;