tree-wide: use PollContext::build_with where possible

The old method of creating a PollContext and calling `add` inside of
`and_then` chains was an ugly way handle the Results that can crop up
after each call. The `build_with` function is equivalent but operates on
a slice which has way less boilerplate.

TEST=./build_test
BUG=None

Change-Id: I8b0d6532680e04c501187397bd211014a2363c25
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/crosvm/+/1715581
Tested-by: Zach Reizner <zachr@chromium.org>
Tested-by: kokoro <noreply+kokoro@google.com>
Auto-Submit: Zach Reizner <zachr@chromium.org>
Reviewed-by: Daniel Verkamp <dverkamp@chromium.org>
Commit-Queue: Zach Reizner <zachr@chromium.org>

14 files changed, 82 insertions, 132 deletions

diff --git a/devices/src/pit.rs b/devices/src/pit.rs
index 63f31f5..86ad4c5 100644
--- a/devices/src/pit.rs
+++ b/devices/src/pit.rs
@@ -745,10 +745,9 @@ impl Worker {
             Kill,
         }
 
-        let poll_ctx: PollContext<Token> = PollContext::new()
-            .and_then(|pc| pc.add(&self.fd, Token::TimerExpire).and(Ok(pc)))
-            .and_then(|pc| pc.add(&kill_evt, Token::Kill).and(Ok(pc)))
-            .map_err(PitError::CreatePollContext)?;
+        let poll_ctx: PollContext<Token> =
+            PollContext::build_with(&[(&self.fd, Token::TimerExpire), (&kill_evt, Token::Kill)])
+                .map_err(PitError::CreatePollContext)?;
 
         loop {
             let events = poll_ctx.wait().map_err(PitError::PollError)?;
diff --git a/devices/src/virtio/balloon.rs b/devices/src/virtio/balloon.rs
index 451421b..5dc0967 100644
--- a/devices/src/virtio/balloon.rs
+++ b/devices/src/virtio/balloon.rs
@@ -143,19 +143,13 @@ impl Worker {
         let inflate_queue_evt = queue_evts.remove(0);
         let deflate_queue_evt = queue_evts.remove(0);
 
-        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(&self.interrupt_resample_evt, Token::InterruptResample)
-                    .and(Ok(pc))
-            })
-            .and_then(|pc| pc.add(&kill_evt, Token::Kill).and(Ok(pc)))
-        {
+        let poll_ctx: PollContext<Token> = match PollContext::build_with(&[
+            (&inflate_queue_evt, Token::Inflate),
+            (&deflate_queue_evt, Token::Deflate),
+            (&self.command_socket, Token::CommandSocket),
+            (&self.interrupt_resample_evt, Token::InterruptResample),
+            (&kill_evt, Token::Kill),
+        ]) {
             Ok(pc) => pc,
             Err(e) => {
                 error!("failed creating PollContext: {}", e);
diff --git a/devices/src/virtio/block.rs b/devices/src/virtio/block.rs
index f638a0e..07e2ca4 100644
--- a/devices/src/virtio/block.rs
+++ b/devices/src/virtio/block.rs
@@ -772,16 +772,13 @@ impl<T: DiskFile> Worker<T> {
         };
         let mut flush_timer_armed = false;
 
-        let poll_ctx: PollContext<Token> = match PollContext::new()
-            .and_then(|pc| pc.add(&flush_timer, Token::FlushTimer).and(Ok(pc)))
-            .and_then(|pc| pc.add(&queue_evt, Token::QueueAvailable).and(Ok(pc)))
-            .and_then(|pc| pc.add(&control_socket, Token::ControlRequest).and(Ok(pc)))
-            .and_then(|pc| {
-                pc.add(&self.interrupt_resample_evt, Token::InterruptResample)
-                    .and(Ok(pc))
-            })
-            .and_then(|pc| pc.add(&kill_evt, Token::Kill).and(Ok(pc)))
-        {
+        let poll_ctx: PollContext<Token> = match PollContext::build_with(&[
+            (&flush_timer, Token::FlushTimer),
+            (&queue_evt, Token::QueueAvailable),
+            (&control_socket, Token::ControlRequest),
+            (&self.interrupt_resample_evt, Token::InterruptResample),
+            (&kill_evt, Token::Kill),
+        ]) {
             Ok(pc) => pc,
             Err(e) => {
                 error!("failed creating PollContext: {}", e);
diff --git a/devices/src/virtio/gpu/mod.rs b/devices/src/virtio/gpu/mod.rs
index 0ab83fd..6cf60d9 100644
--- a/devices/src/virtio/gpu/mod.rs
+++ b/devices/src/virtio/gpu/mod.rs
@@ -497,19 +497,13 @@ impl Worker {
             ResourceBridge { index: usize },
         }
 
-        let poll_ctx: PollContext<Token> = match PollContext::new()
-            .and_then(|pc| pc.add(&self.ctrl_evt, Token::CtrlQueue).and(Ok(pc)))
-            .and_then(|pc| pc.add(&self.cursor_evt, Token::CursorQueue).and(Ok(pc)))
-            .and_then(|pc| {
-                pc.add(&*self.state.display().borrow(), Token::Display)
-                    .and(Ok(pc))
-            })
-            .and_then(|pc| {
-                pc.add(&self.interrupt_resample_evt, Token::InterruptResample)
-                    .and(Ok(pc))
-            })
-            .and_then(|pc| pc.add(&self.kill_evt, Token::Kill).and(Ok(pc)))
-        {
+        let poll_ctx: PollContext<Token> = match PollContext::build_with(&[
+            (&self.ctrl_evt, Token::CtrlQueue),
+            (&self.cursor_evt, Token::CursorQueue),
+            (&*self.state.display().borrow(), Token::Display),
+            (&self.interrupt_resample_evt, Token::InterruptResample),
+            (&self.kill_evt, Token::Kill),
+        ]) {
             Ok(pc) => pc,
             Err(e) => {
                 error!("failed creating PollContext: {}", e);
diff --git a/devices/src/virtio/input/mod.rs b/devices/src/virtio/input/mod.rs
index 195a2e0..3bed5ef 100644
--- a/devices/src/virtio/input/mod.rs
+++ b/devices/src/virtio/input/mod.rs
@@ -566,25 +566,13 @@ impl<T: EventSource> Worker<T> {
             InterruptResample,
             Kill,
         }
-        let poll_ctx: PollContext<Token> = match PollContext::new()
-            .and_then(|pc| {
-                pc.add(&event_queue_evt_fd, Token::EventQAvailable)
-                    .and(Ok(pc))
-            })
-            .and_then(|pc| {
-                pc.add(&status_queue_evt_fd, Token::StatusQAvailable)
-                    .and(Ok(pc))
-            })
-            .and_then(|pc| {
-                pc.add(&self.event_source, Token::InputEventsAvailable)
-                    .and(Ok(pc))
-            })
-            .and_then(|pc| {
-                pc.add(&self.interrupt_resample_evt, Token::InterruptResample)
-                    .and(Ok(pc))
-            })
-            .and_then(|pc| pc.add(&kill_evt, Token::Kill).and(Ok(pc)))
-        {
+        let poll_ctx: PollContext<Token> = match PollContext::build_with(&[
+            (&event_queue_evt_fd, Token::EventQAvailable),
+            (&status_queue_evt_fd, Token::StatusQAvailable),
+            (&self.event_source, Token::InputEventsAvailable),
+            (&self.interrupt_resample_evt, Token::InterruptResample),
+            (&kill_evt, Token::Kill),
+        ]) {
             Ok(poll_ctx) => poll_ctx,
             Err(e) => {
                 error!("failed creating PollContext: {}", e);
diff --git a/devices/src/virtio/net.rs b/devices/src/virtio/net.rs
index cfb43b9..73a8499 100644
--- a/devices/src/virtio/net.rs
+++ b/devices/src/virtio/net.rs
@@ -209,16 +209,14 @@ where
             Kill,
         }
 
-        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(&self.interrupt_resample_evt, Token::InterruptResample)
-                    .and(Ok(pc))
-            })
-            .and_then(|pc| pc.add(&kill_evt, Token::Kill).and(Ok(pc)))
-            .map_err(NetError::CreatePollContext)?;
+        let poll_ctx: PollContext<Token> = PollContext::build_with(&[
+            (&self.tap, Token::RxTap),
+            (&rx_queue_evt, Token::RxQueue),
+            (&tx_queue_evt, Token::TxQueue),
+            (&self.interrupt_resample_evt, Token::InterruptResample),
+            (&kill_evt, Token::Kill),
+        ])
+        .map_err(NetError::CreatePollContext)?;
 
         'poll: loop {
             let events = poll_ctx.wait().map_err(NetError::PollError)?;
diff --git a/devices/src/virtio/p9.rs b/devices/src/virtio/p9.rs
index 9785eb7..a3cb007 100644
--- a/devices/src/virtio/p9.rs
+++ b/devices/src/virtio/p9.rs
@@ -124,14 +124,12 @@ impl Worker {
             Kill,
         }
 
-        let poll_ctx: PollContext<Token> = PollContext::new()
-            .and_then(|pc| pc.add(&queue_evt, Token::QueueReady).and(Ok(pc)))
-            .and_then(|pc| {
-                pc.add(&self.interrupt_resample_evt, Token::InterruptResample)
-                    .and(Ok(pc))
-            })
-            .and_then(|pc| pc.add(&kill_evt, Token::Kill).and(Ok(pc)))
-            .map_err(P9Error::CreatePollContext)?;
+        let poll_ctx: PollContext<Token> = PollContext::build_with(&[
+            (&queue_evt, Token::QueueReady),
+            (&self.interrupt_resample_evt, Token::InterruptResample),
+            (&kill_evt, Token::Kill),
+        ])
+        .map_err(P9Error::CreatePollContext)?;
 
         loop {
             let events = poll_ctx.wait().map_err(P9Error::PollError)?;
diff --git a/devices/src/virtio/pmem.rs b/devices/src/virtio/pmem.rs
index 5a4fa32..607ce58 100644
--- a/devices/src/virtio/pmem.rs
+++ b/devices/src/virtio/pmem.rs
@@ -198,14 +198,11 @@ impl Worker {
             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(&self.interrupt_resample_event, Token::InterruptResample)
-                    .and(Ok(pc))
-            })
-            .and_then(|pc| pc.add(&kill_evt, Token::Kill).and(Ok(pc)))
-        {
+        let poll_ctx: PollContext<Token> = match PollContext::build_with(&[
+            (&queue_evt, Token::QueueAvailable),
+            (&self.interrupt_resample_event, Token::InterruptResample),
+            (&kill_evt, Token::Kill),
+        ]) {
             Ok(pc) => pc,
             Err(e) => {
                 error!("failed creating PollContext: {}", e);
diff --git a/devices/src/virtio/rng.rs b/devices/src/virtio/rng.rs
index 19f83bc..21d5c28 100644
--- a/devices/src/virtio/rng.rs
+++ b/devices/src/virtio/rng.rs
@@ -85,14 +85,11 @@ impl Worker {
             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(&self.interrupt_resample_evt, Token::InterruptResample)
-                    .and(Ok(pc))
-            })
-            .and_then(|pc| pc.add(&kill_evt, Token::Kill).and(Ok(pc)))
-        {
+        let poll_ctx: PollContext<Token> = match PollContext::build_with(&[
+            (&queue_evt, Token::QueueAvailable),
+            (&self.interrupt_resample_evt, Token::InterruptResample),
+            (&kill_evt, Token::Kill),
+        ]) {
             Ok(pc) => pc,
             Err(e) => {
                 error!("failed creating PollContext: {}", e);
diff --git a/devices/src/virtio/tpm.rs b/devices/src/virtio/tpm.rs
index 8715abe..e8c50f2 100644
--- a/devices/src/virtio/tpm.rs
+++ b/devices/src/virtio/tpm.rs
@@ -139,14 +139,11 @@ impl Worker {
             Kill,
         }
 
-        let poll_ctx = match PollContext::new()
-            .and_then(|pc| pc.add(&self.queue_evt, Token::QueueAvailable).and(Ok(pc)))
-            .and_then(|pc| {
-                pc.add(&self.interrupt_resample_evt, Token::InterruptResample)
-                    .and(Ok(pc))
-            })
-            .and_then(|pc| pc.add(&self.kill_evt, Token::Kill).and(Ok(pc)))
-        {
+        let poll_ctx = match PollContext::build_with(&[
+            (&self.queue_evt, Token::QueueAvailable),
+            (&self.interrupt_resample_evt, Token::InterruptResample),
+            (&self.kill_evt, Token::Kill),
+        ]) {
             Ok(pc) => pc,
             Err(e) => {
                 error!("vtpm failed creating PollContext: {}", e);
diff --git a/devices/src/virtio/vhost/worker.rs b/devices/src/virtio/vhost/worker.rs
index 9389f82..54cbb63 100644
--- a/devices/src/virtio/vhost/worker.rs
+++ b/devices/src/virtio/vhost/worker.rs
@@ -119,14 +119,12 @@ impl<T: Vhost> Worker<T> {
             Kill,
         }
 
-        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(&self.interrupt_resample_evt, Token::InterruptResample)
-                    .and(Ok(pc))
-            })
-            .and_then(|pc| pc.add(&kill_evt, Token::Kill).and(Ok(pc)))
-            .map_err(Error::CreatePollContext)?;
+        let poll_ctx: PollContext<Token> = PollContext::build_with(&[
+            (&self.vhost_interrupt, Token::VhostIrq),
+            (&self.interrupt_resample_evt, Token::InterruptResample),
+            (&kill_evt, Token::Kill),
+        ])
+        .map_err(Error::CreatePollContext)?;
 
         'poll: loop {
             let events = poll_ctx.wait().map_err(Error::PollError)?;
diff --git a/devices/src/virtio/wl.rs b/devices/src/virtio/wl.rs
index dd8ce86..d40089b 100644
--- a/devices/src/virtio/wl.rs
+++ b/devices/src/virtio/wl.rs
@@ -1531,15 +1531,13 @@ impl Worker {
             InterruptResample,
         }
 
-        let poll_ctx: PollContext<Token> = match PollContext::new()
-            .and_then(|pc| pc.add(&in_queue_evt, Token::InQueue).and(Ok(pc)))
-            .and_then(|pc| pc.add(&out_queue_evt, Token::OutQueue).and(Ok(pc)))
-            .and_then(|pc| pc.add(&kill_evt, Token::Kill).and(Ok(pc)))
-            .and_then(|pc| pc.add(&self.state.poll_ctx, Token::State).and(Ok(pc)))
-            .and_then(|pc| {
-                pc.add(&self.interrupt_resample_evt, Token::InterruptResample)
-                    .and(Ok(pc))
-            }) {
+        let poll_ctx: PollContext<Token> = match PollContext::build_with(&[
+            (&in_queue_evt, Token::InQueue),
+            (&out_queue_evt, Token::OutQueue),
+            (&kill_evt, Token::Kill),
+            (&self.state.poll_ctx, Token::State),
+            (&self.interrupt_resample_evt, Token::InterruptResample),
+        ]) {
             Ok(pc) => pc,
             Err(e) => {
                 error!("failed creating PollContext: {}", e);
diff --git a/src/linux.rs b/src/linux.rs
index 8d4a4ae..58a0e10 100644
--- a/src/linux.rs
+++ b/src/linux.rs
@@ -1432,16 +1432,15 @@ fn run_control(
         .set_raw_mode()
         .expect("failed to set terminal raw mode");
 
-    let poll_ctx = PollContext::new().map_err(Error::CreatePollContext)?;
-    poll_ctx
-        .add(&linux.exit_evt, Token::Exit)
-        .map_err(Error::PollContextAdd)?;
+    let poll_ctx = PollContext::build_with(&[
+        (&linux.exit_evt, Token::Exit),
+        (&sigchld_fd, Token::ChildSignal),
+    ])
+    .map_err(Error::PollContextAdd)?;
+
     if let Err(e) = poll_ctx.add(&stdin_handle, Token::Stdin) {
         warn!("failed to add stdin to poll context: {}", e);
     }
-    poll_ctx
-        .add(&sigchld_fd, Token::ChildSignal)
-        .map_err(Error::PollContextAdd)?;
 
     if let Some(socket_server) = &control_server_socket {
         poll_ctx
diff --git a/src/plugin/mod.rs b/src/plugin/mod.rs
index cf8332a..3b6d82b 100644
--- a/src/plugin/mod.rs
+++ b/src/plugin/mod.rs
@@ -624,13 +624,9 @@ pub fn run_config(cfg: Config) -> Result<()> {
     let kill_signaled = Arc::new(AtomicBool::new(false));
     let mut vcpu_handles = Vec::with_capacity(vcpu_count as usize);
 
-    let poll_ctx = PollContext::new().map_err(Error::CreatePollContext)?;
-    poll_ctx
-        .add(&exit_evt, Token::Exit)
-        .map_err(Error::PollContextAdd)?;
-    poll_ctx
-        .add(&sigchld_fd, Token::ChildSignal)
-        .map_err(Error::PollContextAdd)?;
+    let poll_ctx =
+        PollContext::build_with(&[(&exit_evt, Token::Exit), (&sigchld_fd, Token::ChildSignal)])
+            .map_err(Error::PollContextAdd)?;
 
     let mut sockets_to_drop = Vec::new();
     let mut redo_poll_ctx_sockets = true;