Skip to main content

wlancfg_lib/client/connection_selection/
fut_manager.rs

1// Copyright 2026 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use 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
17/// Manages the active connection selection futures, allowing for selective cancellation.
18pub struct ConnectionSelectionManager {
19    /// Collection of running connection selection futures, to be polled.
20    #[allow(clippy::type_complexity)]
21    futures: FuturesUnordered<
22        LocalBoxFuture<
23            'static,
24            (SelectionIdentifier, Result<Option<ConnectionSelectionResponse>, anyhow::Error>),
25        >,
26    >,
27    /// Maps a SelectionIdentifier to a sender that can cancel the corresponding future.
28    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    /// Spawns a new connection selection future with an associated SelectionIdentifier.
44    /// Wraps the future in a cancellation logic.
45    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 the manager was previously polled and returned Poll::Pending, it would have stored
69        // a waker, provided by the poller. That waker needs to be awakened to inform the poller
70        // that the manager now has more work to do.
71        if let Some(waker) = &self.waker {
72            waker.wake_by_ref();
73        }
74    }
75
76    /// Cancels the future associated with the given selection_id. Returns true if an active
77    /// selection was cancelled.
78    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    /// Cancels all active connection selections and returns their identifiers.
87    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    /// Returns a list of SelectionIdentifiers for connection selections currently being processed.
96    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
152/// A wrapper future to safely poll the ConnectionSelectionManager within a select! macro.
153pub 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        // Clear any stored waker, in case its obsolete.
166        self.manager.waker = None;
167
168        match Pin::new(&mut self.manager.futures).poll_next(cx) {
169            Poll::Ready(Some((selection_id, result))) => {
170                // Clean up the cancellation map when a future completes.
171                let _ = self.manager.cancellation_handles.remove(&selection_id);
172                Poll::Ready((selection_id, result))
173            }
174            Poll::Ready(None) | Poll::Pending => {
175                // According to the docs, a FuturesUnordered must be polled at least once after a
176                // new future is added in order for that new future to get advanced. We want to
177                // achieve that by having the external caller poll ConnectionSelectionFutures, which
178                // will cause the internal FuturesUnordered to be polled. Therefore, this stores a
179                // copy of the external caller's waker, and use it to wake up that caller when a
180                // new future is spawned. Otherwise, there is no guarantee that the new futures get
181                // polled.
182                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        // This tells select! that futures will continue to be added to the manager, so it should
191        // never be considered terminated.
192        false
193    }
194}
195
196// Identifier for cancellable connection selection futures.
197#[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        // Spawn first
259        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        // Spawn second (duplicate)
266        manager.spawn(
267            selection_id.clone(),
268            Box::pin(async { Ok(ConnectionSelectionResponse::Autoconnect(None)) }),
269        );
270
271        // Should still be 1 because duplicate was ignored
272        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        // Spawn two futures that pending forever
287        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        // Cancel one
293        assert!(manager.cancel(&id_auto));
294
295        // Polling should return cancelled result for id_auto
296        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        // Check active selections - the cancelled handle should be removed immediately.
303        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        // Both should return cancelled
328        let mut futures = ConnectionSelectionFutures::new(&mut manager);
329
330        // We expect two ready results of None
331        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        // Poll empty futures
380        {
381            let mut futures = ConnectionSelectionFutures::new(&mut manager);
382            assert_matches!(Pin::new(&mut futures).poll(&mut cx), Poll::Pending);
383        }
384
385        // Verify the wakers flag is still false.
386        assert!(!flag.0.load(Ordering::SeqCst));
387
388        // Spawn a new future
389        let selection_id = SelectionIdentifier::Automatic;
390        let fut = async move { Ok(ConnectionSelectionResponse::Autoconnect(None)) };
391        manager.spawn(selection_id, Box::pin(fut));
392
393        // Verify waker was triggered
394        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        // Create a waker that increments a counter when awakened.
403        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        // Spawn a future that never completes. This represents an ongoing connection selection future.
417        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        // Spawn a new future.
426        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}