wlancfg_lib/client/connection_selection/
fut_manager.rs1use crate::client::connection_selection::ConnectionSelectionRequester;
6use crate::client::types;
7use crate::mode_management::iface_manager_api::ConnectAttemptRequest;
8use anyhow::format_err;
9use futures::channel::oneshot;
10use futures::future::{FusedFuture, Future, FutureExt, LocalBoxFuture};
11use futures::stream::{FuturesUnordered, Stream};
12use log::warn;
13use std::collections::HashMap;
14use std::pin::Pin;
15use std::task::{Context, Poll, Waker};
16
17pub struct ConnectionSelectionManager {
19 #[allow(clippy::type_complexity)]
21 futures: FuturesUnordered<
22 LocalBoxFuture<
23 'static,
24 (SelectionIdentifier, Result<Option<ConnectionSelectionResponse>, anyhow::Error>),
25 >,
26 >,
27 cancellation_handles: HashMap<SelectionIdentifier, oneshot::Sender<()>>,
29 requester: ConnectionSelectionRequester,
30 waker: Option<Waker>,
31}
32
33impl ConnectionSelectionManager {
34 pub fn new(requester: ConnectionSelectionRequester) -> Self {
35 Self {
36 futures: FuturesUnordered::new(),
37 cancellation_handles: HashMap::new(),
38 requester,
39 waker: None,
40 }
41 }
42
43 pub fn spawn(
46 &mut self,
47 selection_id: SelectionIdentifier,
48 selection_future: LocalBoxFuture<
49 'static,
50 Result<ConnectionSelectionResponse, anyhow::Error>,
51 >,
52 ) {
53 if self.cancellation_handles.contains_key(&selection_id) {
54 warn!("{} already in progress. Ignoring new request.", selection_id);
55 return;
56 }
57
58 let (sender, receiver) = oneshot::channel();
59 let _ = self.cancellation_handles.insert(selection_id.clone(), sender);
60 let cancellable_future = async move {
61 futures::select! {
62 result = selection_future.fuse() => (selection_id.clone(), result.map(Some)),
63 _ = receiver.fuse() => (selection_id, Ok(None)),
64 }
65 };
66
67 self.futures.push(cancellable_future.boxed_local());
68 if let Some(waker) = &self.waker {
72 waker.wake_by_ref();
73 }
74 }
75
76 pub fn cancel(&mut self, selection_id: &SelectionIdentifier) -> bool {
79 if let Some(cancel_tx) = self.cancellation_handles.remove(selection_id) {
80 let _ = cancel_tx.send(());
81 return true;
82 }
83 false
84 }
85
86 pub fn cancel_all(&mut self) -> Vec<SelectionIdentifier> {
88 let active = self.active_selections();
89 for selection_id in &active {
90 let _ = self.cancel(selection_id);
91 }
92 active
93 }
94
95 pub fn active_selections(&self) -> Vec<SelectionIdentifier> {
97 self.cancellation_handles.keys().cloned().collect()
98 }
99
100 pub fn initiate_automatic_connection_selection(&mut self) {
101 let mut requester = self.requester.clone();
102 let fut = async move {
103 requester
104 .do_connection_selection(None, types::ConnectReason::IdleInterfaceAutoconnect)
105 .await
106 .map_err(|e| format_err!("Error sending connection selection request: {}.", e))
107 .map(ConnectionSelectionResponse::Autoconnect)
108 };
109 self.spawn(SelectionIdentifier::Automatic, Box::pin(fut).boxed_local());
110 }
111
112 pub fn initiate_connection_selection_for_connect_request(
113 &mut self,
114 request: ConnectAttemptRequest,
115 ) {
116 let mut requester = self.requester.clone();
117 let network = request.network.clone();
118 let fut = async move {
119 requester
120 .do_connection_selection(Some(request.network.clone()), request.reason)
121 .await
122 .map_err(|e| format_err!("Error sending connection selection request: {}.", e))
123 .map(move |candidate| ConnectionSelectionResponse::ConnectRequest {
124 candidate,
125 request,
126 })
127 };
128 self.spawn(SelectionIdentifier::ConnectRequest(network), Box::pin(fut).boxed_local());
129 }
130
131 #[cfg(test)]
132 #[allow(clippy::type_complexity)]
133 pub fn get_futures(
134 &mut self,
135 ) -> &mut FuturesUnordered<
136 LocalBoxFuture<
137 'static,
138 (SelectionIdentifier, Result<Option<ConnectionSelectionResponse>, anyhow::Error>),
139 >,
140 > {
141 &mut self.futures
142 }
143
144 #[cfg(test)]
145 pub fn get_cancellation_handles(
146 &mut self,
147 ) -> &mut HashMap<SelectionIdentifier, oneshot::Sender<()>> {
148 &mut self.cancellation_handles
149 }
150}
151
152pub struct ConnectionSelectionFutures<'a> {
154 manager: &'a mut ConnectionSelectionManager,
155}
156impl<'a> ConnectionSelectionFutures<'a> {
157 pub fn new(manager: &'a mut ConnectionSelectionManager) -> Self {
158 Self { manager }
159 }
160}
161impl<'a> Future for ConnectionSelectionFutures<'a> {
162 type Output = (SelectionIdentifier, Result<Option<ConnectionSelectionResponse>, anyhow::Error>);
163
164 fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
165 self.manager.waker = None;
167
168 match Pin::new(&mut self.manager.futures).poll_next(cx) {
169 Poll::Ready(Some((selection_id, result))) => {
170 let _ = self.manager.cancellation_handles.remove(&selection_id);
172 Poll::Ready((selection_id, result))
173 }
174 Poll::Ready(None) | Poll::Pending => {
175 self.manager.waker = Some(cx.waker().clone());
183 Poll::Pending
184 }
185 }
186 }
187}
188impl<'a> FusedFuture for ConnectionSelectionFutures<'a> {
189 fn is_terminated(&self) -> bool {
190 false
193 }
194}
195
196#[derive(Clone, PartialEq, Eq, Hash)]
198#[cfg_attr(test, derive(Debug))]
199pub enum SelectionIdentifier {
200 ConnectRequest(types::NetworkIdentifier),
201 Automatic,
202}
203impl std::fmt::Display for SelectionIdentifier {
204 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
205 match self {
206 SelectionIdentifier::ConnectRequest(id) => {
207 write!(f, "Connection selection for network: {}", id)
208 }
209 SelectionIdentifier::Automatic => write!(f, "Automatic connection selection"),
210 }
211 }
212}
213
214#[cfg_attr(test, derive(Debug))]
215pub enum ConnectionSelectionResponse {
216 ConnectRequest { candidate: Option<types::ScannedCandidate>, request: ConnectAttemptRequest },
217 Autoconnect(Option<types::ScannedCandidate>),
218}
219
220#[cfg(test)]
221mod tests {
222 use super::*;
223 use crate::util::testing::generate_random_network_identifier;
224 use assert_matches::assert_matches;
225 use fuchsia_async::TestExecutor;
226 use futures::task::Poll;
227 use std::sync::Arc;
228 use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
229 use std::task::{Wake, Waker};
230
231 fn fake_requester() -> ConnectionSelectionRequester {
232 let (sender, _receiver) = futures::channel::mpsc::channel(1);
233 ConnectionSelectionRequester::new(sender)
234 }
235
236 #[test]
237 fn test_spawn_and_poll() {
238 let mut exec = TestExecutor::new();
239 let mut manager = ConnectionSelectionManager::new(fake_requester());
240
241 let selection_id = SelectionIdentifier::Automatic;
242 let fut = async move { Ok(ConnectionSelectionResponse::Autoconnect(None)) };
243 manager.spawn(selection_id.clone(), Box::pin(fut));
244
245 let mut futures = ConnectionSelectionFutures::new(&mut manager);
246 assert_matches!(
247 exec.run_until_stalled(&mut futures),
248 Poll::Ready((id, Ok(Some(ConnectionSelectionResponse::Autoconnect(None))))) if id == selection_id
249 );
250 }
251
252 #[test]
253 fn test_spawn_duplicate() {
254 let mut manager = ConnectionSelectionManager::new(fake_requester());
255
256 let selection_id = SelectionIdentifier::Automatic;
257
258 manager.spawn(
260 selection_id.clone(),
261 Box::pin(async { Ok(ConnectionSelectionResponse::Autoconnect(None)) }),
262 );
263 assert_eq!(manager.active_selections().len(), 1);
264
265 manager.spawn(
267 selection_id.clone(),
268 Box::pin(async { Ok(ConnectionSelectionResponse::Autoconnect(None)) }),
269 );
270
271 assert_eq!(manager.active_selections().len(), 1);
273 }
274
275 #[test]
276 fn test_cancel_specific() {
277 let mut exec = TestExecutor::new();
278 let mut manager = ConnectionSelectionManager::new(fake_requester());
279
280 let id_auto = SelectionIdentifier::Automatic;
281 let id_req = SelectionIdentifier::ConnectRequest(types::NetworkIdentifier {
282 ssid: types::Ssid::from_bytes_unchecked(b"test".to_vec()),
283 security_type: types::SecurityType::Wpa2,
284 });
285
286 manager.spawn(id_auto.clone(), Box::pin(futures::future::pending()));
288 manager.spawn(id_req.clone(), Box::pin(futures::future::pending()));
289
290 assert_eq!(manager.active_selections().len(), 2);
291
292 assert!(manager.cancel(&id_auto));
294
295 let mut futures = ConnectionSelectionFutures::new(&mut manager);
297 assert_matches!(
298 exec.run_until_stalled(&mut futures),
299 Poll::Ready((id, Ok(None))) if id == id_auto
300 );
301
302 assert_eq!(manager.active_selections().len(), 1);
304 assert!(manager.active_selections().contains(&id_req));
305 }
306
307 #[test]
308 fn test_cancel_all() {
309 let mut exec = TestExecutor::new();
310 let mut manager = ConnectionSelectionManager::new(fake_requester());
311
312 let id1 = SelectionIdentifier::Automatic;
313 let id2 = SelectionIdentifier::ConnectRequest(types::NetworkIdentifier {
314 ssid: types::Ssid::from_bytes_unchecked(b"test".to_vec()),
315 security_type: types::SecurityType::Wpa2,
316 });
317
318 manager.spawn(id1.clone(), Box::pin(futures::future::pending()));
319 manager.spawn(id2.clone(), Box::pin(futures::future::pending()));
320
321 assert_eq!(manager.active_selections().len(), 2);
322
323 assert_eq!(manager.cancel_all().len(), 2);
324
325 assert_eq!(manager.active_selections().len(), 0);
326
327 let mut futures = ConnectionSelectionFutures::new(&mut manager);
329
330 let mut cancelled_count = 0;
332
333 loop {
334 match exec.run_until_stalled(&mut futures) {
335 Poll::Ready((_, Ok(None))) => cancelled_count += 1,
336 Poll::Pending => break,
337 _ => panic!("Unexpected result"),
338 }
339 }
340
341 assert_eq!(cancelled_count, 2);
342 }
343
344 #[test]
345 fn test_internal_state_access() {
346 let mut manager = ConnectionSelectionManager::new(fake_requester());
347 let id = SelectionIdentifier::Automatic;
348 manager.spawn(id.clone(), Box::pin(futures::future::pending()));
349
350 assert_eq!(manager.get_cancellation_handles().len(), 1);
351 assert!(!manager.get_futures().is_empty());
352 }
353
354 #[test]
355 fn test_poll_empty_twice_does_not_panic() {
356 let mut exec = TestExecutor::new();
357 let mut manager = ConnectionSelectionManager::new(fake_requester());
358
359 let mut futures = ConnectionSelectionFutures::new(&mut manager);
360 assert_matches!(exec.run_until_stalled(&mut futures), Poll::Pending);
361 assert_matches!(exec.run_until_stalled(&mut futures), Poll::Pending);
362 }
363
364 #[test]
365 fn test_spawn_wakes_after_empty() {
366 struct FlagWaker(AtomicBool);
367 impl Wake for FlagWaker {
368 fn wake(self: Arc<Self>) {
369 self.0.store(true, Ordering::SeqCst);
370 }
371 }
372
373 let mut manager = ConnectionSelectionManager::new(fake_requester());
374
375 let flag = Arc::new(FlagWaker(AtomicBool::new(false)));
376 let waker = Waker::from(flag.clone());
377 let mut cx = Context::from_waker(&waker);
378
379 {
381 let mut futures = ConnectionSelectionFutures::new(&mut manager);
382 assert_matches!(Pin::new(&mut futures).poll(&mut cx), Poll::Pending);
383 }
384
385 assert!(!flag.0.load(Ordering::SeqCst));
387
388 let selection_id = SelectionIdentifier::Automatic;
390 let fut = async move { Ok(ConnectionSelectionResponse::Autoconnect(None)) };
391 manager.spawn(selection_id, Box::pin(fut));
392
393 assert!(
395 flag.0.load(Ordering::SeqCst),
396 "Waker should have been awakened by spawn() after empty poll"
397 );
398 }
399
400 #[test]
401 fn test_spawn_wakes_while_busy() {
402 struct CounterWaker(AtomicUsize);
404 impl Wake for CounterWaker {
405 fn wake(self: Arc<Self>) {
406 let _ = self.0.fetch_add(1, Ordering::SeqCst);
407 }
408 }
409
410 let mut manager = ConnectionSelectionManager::new(fake_requester());
411
412 let flag = Arc::new(CounterWaker(AtomicUsize::new(0)));
413 let waker = Waker::from(flag.clone());
414 let mut cx = Context::from_waker(&waker);
415
416 manager.spawn(SelectionIdentifier::Automatic, Box::pin(futures::future::pending()));
418
419 {
420 let mut futures = ConnectionSelectionFutures::new(&mut manager);
421 assert_matches!(Pin::new(&mut futures).poll(&mut cx), Poll::Pending);
422 }
423 assert_eq!(flag.0.load(Ordering::SeqCst), 1);
424
425 let selection_id =
427 SelectionIdentifier::ConnectRequest(generate_random_network_identifier());
428 manager.spawn(
429 selection_id.clone(),
430 Box::pin(futures::future::ready(Ok(ConnectionSelectionResponse::Autoconnect(None)))),
431 );
432 assert_eq!(
433 flag.0.load(Ordering::SeqCst),
434 2,
435 "Manager should wake stored waker on spawn even when busy"
436 );
437 }
438}