src_util: add poll module

The poll module adds the Poller object for waiting on mutliple file
descriptors at once. The Pollable trait is introduced so rust objects
can expose a file descriptor useful for polling. An impl for EventFd is
included with this change for testing.

TEST=cargo test
BUG=None

Change-Id: I94fd15a17fe0527c0d29c623badb90668d708689
Reviewed-on: https://chromium-review.googlesource.com/514413
Commit-Ready: Zach Reizner <zachr@chromium.org>
Tested-by: Zach Reizner <zachr@chromium.org>
Reviewed-by: Dylan Reid <dgreid@chromium.org>

1 files changed, 121 insertions, 0 deletions

diff --git a/sys_util/src/poll.rs b/sys_util/src/poll.rs
new file mode 100644
index 0000000..f13faa5
--- /dev/null
+++ b/sys_util/src/poll.rs
@@ -0,0 +1,121 @@
+// Copyright 2017 The Chromium OS Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+use std::os::unix::io::RawFd;
+
+use libc::{nfds_t, pollfd, poll, POLLIN};
+
+use {Result, errno_result};
+
+/// Trait for file descriptors that can be polled for input.
+///
+/// This is marked unsafe because the implementation must promise that the returned RawFd is valid
+/// for polling purposes and that the lifetime of the returned fd is at least that of the trait
+/// object.
+pub unsafe trait Pollable {
+    /// Gets the file descriptor that can be polled for input.
+    fn pollable_fd(&self) -> RawFd;
+}
+
+
+/// Used to poll multiple `Pollable` objects at once.
+///
+/// # Example
+///
+/// ```
+/// # use sys_util::{Result, EventFd, Poller, Pollable};
+/// # fn test() -> Result<()> {
+///     let evt1 = EventFd::new()?;
+///     let evt2 = EventFd::new()?;
+///     evt2.write(1)?;
+///
+///     let pollables: Vec<(u32, &Pollable)> = vec![(1, &evt1), (2, &evt2)];
+///
+///     let mut poller = Poller::new(2);
+///     assert_eq!(poller.poll(&pollables[..]), Ok([2].as_ref()));
+/// #   Ok(())
+/// # }
+/// ```
+pub struct Poller {
+    pollfds: Vec<pollfd>,
+    tokens: Vec<u32>,
+}
+
+impl Poller {
+    /// Constructs a new poller object with the given `capacity` of Pollable objects pre-allocated.
+    pub fn new(capacity: usize) -> Poller {
+        Poller {
+            pollfds: Vec::with_capacity(capacity),
+            tokens: Vec::with_capacity(capacity),
+        }
+    }
+
+    /// Waits for any of the given slice of `token`-`Pollable` tuples to be readable without
+    /// blocking and returns the `token` of each that is readable.
+    ///
+    /// This is guaranteed to not allocate if `pollables.len()` is less than the `capacity` given in
+    /// `Poller::new`.
+    pub fn poll(&mut self, pollables: &[(u32, &Pollable)]) -> Result<&[u32]> {
+        self.pollfds.clear();
+        for pollable in pollables.iter() {
+            self.pollfds
+                .push(pollfd {
+                          fd: pollable.1.pollable_fd(),
+                          events: POLLIN,
+                          revents: 0,
+                      });
+        }
+
+        // Safe because poll is given the correct length of properly initialized pollfds, and we
+        // check the return result.
+        let ret = unsafe {
+            handle_eintr!(poll(self.pollfds.as_mut_ptr(),
+                       self.pollfds.len() as nfds_t,
+                       -1))
+        };
+        if ret < 0 {
+            return errno_result();
+        }
+
+        self.tokens.clear();
+        for (pollfd, pollable) in self.pollfds.iter().zip(pollables.iter()) {
+            if (pollfd.revents & POLLIN) != 0 {
+                self.tokens.push(pollable.0);
+            }
+        }
+
+        Ok(&self.tokens)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use EventFd;
+
+    #[test]
+    fn poller() {
+        let evt1 = EventFd::new().unwrap();
+        let evt2 = EventFd::new().unwrap();
+        evt2.write(1).unwrap();
+
+        let pollables: Vec<(u32, &Pollable)> = vec![(1, &evt1), (2, &evt2)];
+
+        let mut poller = Poller::new(2);
+        assert_eq!(poller.poll(&pollables[..]), Ok([2].as_ref()));
+    }
+
+    #[test]
+    fn poller_multi() {
+        let evt1 = EventFd::new().unwrap();
+        let evt2 = EventFd::new().unwrap();
+        evt1.write(1).unwrap();
+        evt2.write(1).unwrap();
+
+        let pollables: Vec<(u32, &Pollable)> = vec![(1, &evt1), (2, &evt2)];
+
+        let mut poller = Poller::new(2);
+        assert_eq!(poller.poll(&pollables[..]), Ok([1, 2].as_ref()));
+    }
+}