1 files changed, 22 insertions, 10 deletions
diff --git a/devices/src/virtio/block.rs b/devices/src/virtio/block.rs
index ca0b222..2faa45d 100644
--- a/devices/src/virtio/block.rs
+++ b/devices/src/virtio/block.rs
@@ -11,7 +11,7 @@ use std::sync::atomic::{AtomicUsize, Ordering};
 use std::thread;
 
 use sys_util::Result as SysResult;
-use sys_util::{EventFd, GuestAddress, GuestMemory, GuestMemoryError, Poller};
+use sys_util::{EventFd, GuestAddress, GuestMemory, GuestMemoryError, PollContext, PollToken};
 
 use super::{VirtioDevice, Queue, DescriptorChain, INTERRUPT_STATUS_USED_RING, TYPE_BLOCK};
 
@@ -232,12 +232,25 @@ impl<T: DiskFile> Worker<T> {
     }
 
     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);
@@ -246,17 +259,16 @@ impl<T: DiskFile> Worker<T> {
             };
 
             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(0);
                     }
-                    KILL => break 'poll,
-                    _ => unreachable!(),
+                    Token::Kill => break 'poll,
                 }
             }
             if needs_interrupt {