devices: use PollContext for all virtio deivces

BUG=chromium:816692
TEST=run any VM

Change-Id: I4219050fdb7947ca513f599f1ac57cde6052d397
Reviewed-on: https://chromium-review.googlesource.com/996917
Commit-Ready: Zach Reizner <zachr@chromium.org>
Tested-by: Zach Reizner <zachr@chromium.org>
Reviewed-by: Stephen Barber <smbarber@chromium.org>

5 files changed, 104 insertions, 72 deletions

diff --git a/devices/src/virtio/balloon.rs b/devices/src/virtio/balloon.rs
index 96386d7..2194154 100644
--- a/devices/src/virtio/balloon.rs
+++ b/devices/src/virtio/balloon.rs
@@ -12,7 +12,7 @@ use std::sync::atomic::{AtomicUsize, Ordering};
 use std::thread;
 
 use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
-use sys_util::{self, EventFd, GuestAddress, GuestMemory, Pollable, Poller};
+use sys_util::{self, EventFd, GuestAddress, GuestMemory, PollContext, PollToken};
 
 use super::{VirtioDevice, Queue, DescriptorChain, INTERRUPT_STATUS_CONFIG_CHANGED,
             INTERRUPT_STATUS_USED_RING, TYPE_BALLOON};
@@ -123,49 +123,57 @@ impl Worker {
     }
 
     fn run(&mut self, mut queue_evts: Vec<EventFd>, kill_evt: EventFd) {
-        const POLL_INFLATE: u32 = 0;
-        const POLL_DEFLATE: u32 = 1;
-        const POLL_COMMAND_SOCKET: u32 = 2;
-        const POLL_KILL: u32 = 3;
+        #[derive(PartialEq, PollToken)]
+        enum Token {
+            Inflate,
+            Deflate,
+            CommandSocket,
+            Kill,
+        }
 
         let inflate_queue_evt = queue_evts.remove(0);
         let deflate_queue_evt = queue_evts.remove(0);
 
-        let mut poller = Poller::new(5);
+        let poll_ctx: PollContext<Token> =
+            match PollContext::new()
+                      .and_then(|pc| pc.add(&inflate_queue_evt, Token::Inflate).and(Ok(pc)))
+                      .and_then(|pc| pc.add(&deflate_queue_evt, Token::Deflate).and(Ok(pc)))
+                      .and_then(|pc| pc.add(&self.command_socket, Token::CommandSocket).and(Ok(pc)))
+                      .and_then(|pc| pc.add(&kill_evt, Token::Kill).and(Ok(pc))) {
+                Ok(pc) => pc,
+                Err(e) => {
+                    error!("failed creating PollContext: {:?}", e);
+                    return;
+                }
+            };
+
         'poll: loop {
-            let tokens = match poller.poll(
-                &[
-                    (POLL_INFLATE, &inflate_queue_evt),
-                    (POLL_DEFLATE, &deflate_queue_evt),
-                    (POLL_COMMAND_SOCKET, &self.command_socket as &Pollable),
-                    (POLL_KILL, &kill_evt),
-                ],
-            ) {
+            let events = match poll_ctx.wait() {
                 Ok(v) => v,
                 Err(e) => {
                     error!("failed polling for events: {:?}", e);
-                    break 'poll;
+                    break;
                 }
             };
 
             let mut needs_interrupt = false;
-            'read_tokens: for &token in tokens {
-                match token {
-                    POLL_INFLATE => {
+            for event in events.iter_readable() {
+                match event.token() {
+                    Token::Inflate => {
                         if let Err(e) = inflate_queue_evt.read() {
                             error!("failed reading inflate queue EventFd: {:?}", e);
                             break 'poll;
                         }
                         needs_interrupt |= self.process_inflate_deflate(true);
                     }
-                    POLL_DEFLATE => {
+                    Token::Deflate => {
                         if let Err(e) = deflate_queue_evt.read() {
                             error!("failed reading deflate queue EventFd: {:?}", e);
                             break 'poll;
                         }
                         needs_interrupt |= self.process_inflate_deflate(false);
                     }
-                    POLL_COMMAND_SOCKET => {
+                    Token::CommandSocket => {
                         let mut buf = [0u8; 4];
                         if let Ok(count) = self.command_socket.recv(&mut buf) {
                             if count == 4 {
@@ -174,7 +182,7 @@ impl Worker {
                                 let num_pages = self.config.num_pages.load(Ordering::Relaxed) as
                                     i32;
                                 if increment < 0 && increment.abs() > num_pages {
-                                    continue 'read_tokens;
+                                    continue;
                                 }
                                 self.config.num_pages.fetch_add(
                                     increment as usize,
@@ -184,8 +192,14 @@ impl Worker {
                             }
                         }
                     }
-                    POLL_KILL => break 'poll,
-                    _ => unreachable!(),
+                    Token::Kill => break 'poll,
+                }
+            }
+            for event in events.iter_hungup() {
+                if event.token() == Token::CommandSocket && !event.readable() {
+                    // If this call fails, the command socket was already removed from the
+                    // PollContext.
+                    let _ = poll_ctx.delete(&self.command_socket);
                 }
             }
             if needs_interrupt {
diff --git a/devices/src/virtio/net.rs b/devices/src/virtio/net.rs
index 13482ed..8360158 100644
--- a/devices/src/virtio/net.rs
+++ b/devices/src/virtio/net.rs
@@ -14,7 +14,7 @@ use libc::EAGAIN;
 use net_sys;
 use net_util::{Error as TapError, MacAddress, TapT};
 use sys_util::Error as SysError;
-use sys_util::{EventFd, GuestMemory, Pollable, Poller};
+use sys_util::{EventFd, GuestMemory, PollContext, PollToken};
 use virtio_sys::{vhost, virtio_net};
 use virtio_sys::virtio_net::virtio_net_hdr_v1;
 
@@ -31,6 +31,8 @@ const QUEUE_SIZES: &'static [u16] = &[QUEUE_SIZE, QUEUE_SIZE];
 pub enum NetError {
     /// Creating kill eventfd failed.
     CreateKillEventFd(SysError),
+    /// Creating PollContext failed.
+    CreatePollContext(SysError),
     /// Cloning kill eventfd failed.
     CloneKillEventFd(SysError),
     /// Open tap device failed.
@@ -219,28 +221,30 @@ where
            tx_queue_evt: EventFd,
            kill_evt: EventFd)
            -> Result<(), NetError> {
-        let mut poller = Poller::new(4);
-        // A frame is available for reading from the tap device to receive in the guest.
-        const RX_TAP: u32 = 1;
-        // The guest has made a buffer available to receive a frame into.
-        const RX_QUEUE: u32 = 2;
-        // The transmit queue has a frame that is ready to send from the guest.
-        const TX_QUEUE: u32 = 3;
-        // crosvm has requested the device to shut down.
-        const KILL: u32 = 4;
+        #[derive(PollToken)]
+        enum Token {
+            // A frame is available for reading from the tap device to receive in the guest.
+            RxTap,
+            // The guest has made a buffer available to receive a frame into.
+            RxQueue,
+            // The transmit queue has a frame that is ready to send from the guest.
+            TxQueue,
+            // crosvm has requested the device to shut down.
+            Kill,
+        }
 
-        'poll: loop {
-            let tokens = match poller.poll(&[(RX_TAP, &self.tap),
-                                             (RX_QUEUE, &rx_queue_evt as &Pollable),
-                                             (TX_QUEUE, &tx_queue_evt as &Pollable),
-                                             (KILL, &kill_evt as &Pollable)]) {
-                Ok(v) => v,
-                Err(e) => return Err(NetError::PollError(e)),
-            };
+        let poll_ctx: PollContext<Token> = PollContext::new()
+                      .and_then(|pc| pc.add(&self.tap, Token::RxTap).and(Ok(pc)))
+                      .and_then(|pc| pc.add(&rx_queue_evt, Token::RxQueue).and(Ok(pc)))
+                      .and_then(|pc| pc.add(&tx_queue_evt, Token::TxQueue).and(Ok(pc)))
+                      .and_then(|pc| pc.add(&kill_evt, Token::Kill).and(Ok(pc)))
+                      .map_err(NetError::CreatePollContext)?;
 
-            for &token in tokens {
-                match token {
-                    RX_TAP => {
+        'poll: loop {
+            let events = poll_ctx.wait().map_err(NetError::PollError)?;
+            for event in events.iter_readable() {
+                match event.token() {
+                    Token::RxTap => {
                         // Process a deferred frame first if available. Don't read from tap again
                         // until we manage to receive this deferred frame.
                         if self.deferred_rx {
@@ -252,7 +256,7 @@ where
                         }
                         self.process_rx();
                     }
-                    RX_QUEUE => {
+                    Token::RxQueue => {
                         if let Err(e) = rx_queue_evt.read() {
                             error!("net: error reading rx queue EventFd: {:?}", e);
                             break 'poll;
@@ -262,15 +266,14 @@ where
                             self.deferred_rx = false;
                         }
                     }
-                    TX_QUEUE => {
+                    Token::TxQueue => {
                         if let Err(e) = tx_queue_evt.read() {
                             error!("net: error reading tx queue EventFd: {:?}", e);
                             break 'poll;
                         }
                         self.process_tx();
                     }
-                    KILL => break 'poll,
-                    _ => unreachable!(),
+                    Token::Kill => break 'poll,
                 }
             }
         }
diff --git a/devices/src/virtio/rng.rs b/devices/src/virtio/rng.rs
index 7bf6410..f1cfe0a 100644
--- a/devices/src/virtio/rng.rs
+++ b/devices/src/virtio/rng.rs
@@ -10,7 +10,7 @@ use std::sync::Arc;
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::thread;
 
-use sys_util::{EventFd, GuestMemory, Poller};
+use sys_util::{EventFd, GuestMemory, PollContext, PollToken};
 
 use super::{VirtioDevice, Queue, INTERRUPT_STATUS_USED_RING, TYPE_RNG};
 
@@ -69,12 +69,25 @@ impl Worker {
     }
 
     fn run(&mut self, queue_evt: EventFd, kill_evt: EventFd) {
-        const Q_AVAIL: u32 = 0;
-        const KILL: u32 = 1;
+        #[derive(PollToken)]
+        enum Token {
+            QueueAvailable,
+            Kill,
+        }
+
+        let poll_ctx: PollContext<Token> =
+            match PollContext::new()
+                      .and_then(|pc| pc.add(&queue_evt, Token::QueueAvailable).and(Ok(pc)))
+                      .and_then(|pc| pc.add(&kill_evt, Token::Kill).and(Ok(pc))) {
+                Ok(pc) => pc,
+                Err(e) => {
+                    error!("failed creating PollContext: {:?}", e);
+                    return;
+                }
+            };
 
-        let mut poller = Poller::new(2);
         'poll: loop {
-            let tokens = match poller.poll(&[(Q_AVAIL, &queue_evt), (KILL, &kill_evt)]) {
+            let events = match poll_ctx.wait() {
                 Ok(v) => v,
                 Err(e) => {
                     error!("failed polling for events: {:?}", e);
@@ -83,17 +96,16 @@ impl Worker {
             };
 
             let mut needs_interrupt = false;
-            for &token in tokens {
-                match token {
-                    Q_AVAIL => {
+            for event in events.iter_readable() {
+                match event.token() {
+                    Token::QueueAvailable => {
                         if let Err(e) = queue_evt.read() {
                             error!("failed reading queue EventFd: {:?}", e);
                             break 'poll;
                         }
                         needs_interrupt |= self.process_queue();
                     }
-                    KILL => break 'poll,
-                    _ => unreachable!(),
+                    Token::Kill => break 'poll,
                 }
             }
             if needs_interrupt {
diff --git a/devices/src/virtio/vhost/mod.rs b/devices/src/virtio/vhost/mod.rs
index 0039a5a..0fe15df 100644
--- a/devices/src/virtio/vhost/mod.rs
+++ b/devices/src/virtio/vhost/mod.rs
@@ -21,6 +21,8 @@ pub use self::vsock::Vsock;
 pub enum Error {
     /// Creating kill eventfd failed.
     CreateKillEventFd(SysError),
+    /// Creating poll context failed.
+    CreatePollContext(SysError),
     /// Cloning kill eventfd failed.
     CloneKillEventFd(SysError),
     /// Error while polling for events.
diff --git a/devices/src/virtio/vhost/worker.rs b/devices/src/virtio/vhost/worker.rs
index 12d879a..6d4cab7 100644
--- a/devices/src/virtio/vhost/worker.rs
+++ b/devices/src/virtio/vhost/worker.rs
@@ -6,7 +6,7 @@ use std::os::raw::c_ulonglong;
 use std::sync::Arc;
 use std::sync::atomic::{AtomicUsize, Ordering};
 
-use sys_util::{EventFd, Poller};
+use sys_util::{EventFd, PollContext, PollToken};
 use vhost::Vhost;
 
 use super::{Error, Result};
@@ -100,27 +100,28 @@ impl<T: Vhost> Worker<T> {
 
         activate_vqs(&self.vhost_handle)?;
 
-        const VHOST_IRQ: u32 = 1;
-        const KILL: u32 = 2;
+        #[derive(PollToken)]
+        enum Token {
+            VhostIrq,
+            Kill,
+        }
 
-        let mut poller = Poller::new(2);
+        let poll_ctx: PollContext<Token> = PollContext::new()
+                      .and_then(|pc| pc.add(&self.vhost_interrupt, Token::VhostIrq).and(Ok(pc)))
+                      .and_then(|pc| pc.add(&kill_evt, Token::Kill).and(Ok(pc)))
+                      .map_err(Error::CreatePollContext)?;
 
         'poll: loop {
-            let tokens = match poller.poll(&[(VHOST_IRQ, &self.vhost_interrupt), (KILL, &kill_evt)])
-            {
-                Ok(v) => v,
-                Err(e) => return Err(Error::PollError(e)),
-            };
+            let events = poll_ctx.wait().map_err(Error::PollError)?;
 
             let mut needs_interrupt = false;
-            for &token in tokens {
-                match token {
-                    VHOST_IRQ => {
+            for event in events.iter_readable() {
+                match event.token() {
+                    Token::VhostIrq => {
                         needs_interrupt = true;
                         self.vhost_interrupt.read().map_err(Error::VhostIrqRead)?;
                     },
-                    KILL => break 'poll,
-                    _ => unreachable!(),
+                    Token::Kill => break 'poll,
                 }
             }
             if needs_interrupt {