cros_async: Fix double cancellation of fd wakers

When converting from VecDeque to Slab, the remove function return
changed from returning None when missing a token a panic.
Instead, of relying on that explicitly store an option to the waker so
there can be no mistake, log an error if there is ever an unexpected
token or if a waker fires twice.

Change-Id: Icfbb3124ca4ad3cd721a3722a299f28358e8547b
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/crosvm/+/2545127
Tested-by: kokoro <noreply+kokoro@google.com>
Commit-Queue: Dylan Reid <dgreid@chromium.org>
Reviewed-by: Chirantan Ekbote <chirantan@chromium.org>

1 files changed, 13 insertions, 6 deletions

diff --git a/cros_async/src/fd_executor.rs b/cros_async/src/fd_executor.rs
index 515391e90..08ca2f4fb 100644
--- a/cros_async/src/fd_executor.rs
+++ b/cros_async/src/fd_executor.rs
@@ -21,7 +21,7 @@ use std::task::Waker;
 
 use slab::Slab;
 
-use sys_util::{PollContext, WatchingEvents};
+use sys_util::{error, PollContext, WatchingEvents};
 
 use crate::executor::{ExecutableFuture, Executor, FutureList};
 use crate::WakerToken;
@@ -150,7 +150,7 @@ pub(crate) fn add_future(future: Pin<Box<dyn Future<Output = ()>>>) -> Result<()
 // Tracks active wakers and associates wakers with the futures that registered them.
 struct FdWakerState {
     poll_ctx: PollContext<usize>,
-    tokens: Slab<(File, Waker)>,
+    tokens: Slab<(File, Option<Waker>)>,
     new_futures: VecDeque<ExecutableFuture<()>>,
 }
 
@@ -175,7 +175,7 @@ impl FdWakerState {
         self.poll_ctx
             .add_fd_with_events(&duped_fd, events, next_token)
             .map_err(Error::SubmittingWaker)?;
-        entry.insert((duped_fd, waker));
+        entry.insert((duped_fd, Some(waker)));
         Ok(WakerToken(next_token))
     }
 
@@ -184,9 +184,16 @@ impl FdWakerState {
         let events = self.poll_ctx.wait().map_err(Error::PollContextError)?;
         for e in events.iter() {
             let token = e.token();
-            let (fd, waker) = self.tokens.remove(token);
-            self.poll_ctx.delete(&fd).map_err(Error::PollContextError)?;
-            waker.wake();
+            if let Some((fd, waker)) = self.tokens.get_mut(token) {
+                self.poll_ctx.delete(fd).map_err(Error::PollContextError)?;
+                if let Some(waker) = waker.take() {
+                    waker.wake();
+                } else {
+                    error!("Woken twice");
+                }
+            } else {
+                error!("Unknown waker");
+            }
         }
         Ok(())
     }