From cfae7ffaf5e728348a127d54c72ec6f04d4dd4c3 Mon Sep 17 00:00:00 2001
From: Olof Johansson <olof@lixom.net>
Date: Tue, 20 Sep 2022 16:41:42 -0700
Subject: [PATCH] Add `available` traits to find out how many client tokens are
 unused

Returns amount of tokens in the read-side pipe (number of jobs that could be started).

Signed-off-by: Olof Johansson <olof@lixom.net>
---
 src/lib.rs      | 13 +++++++++++++
 src/unix.rs     |  7 +++++++
 src/wasm.rs     |  5 +++++
 src/windows.rs  | 20 ++++++++++++++++++++
 tests/server.rs | 24 ++++++++++++++++++++++++
 5 files changed, 69 insertions(+)

diff --git a/src/lib.rs b/src/lib.rs
index 1272eaa..6d07884 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -273,6 +273,19 @@ impl Client {
         })
     }
 
+    /// Returns amount of tokens in the read-side pipe.
+    ///
+    /// # Return value
+    ///
+    /// Number of bytes available to be read from the jobserver pipe
+    ///
+    /// # Errors
+    ///
+    /// Underlying errors from the ioctl will be passed up.
+    pub fn available(&self) -> io::Result<usize> {
+        self.inner.available()
+    }
+
     /// Configures a child process to have access to this client's jobserver as
     /// well.
     ///
diff --git a/src/unix.rs b/src/unix.rs
index 2373fb7..2474dee 100644
--- a/src/unix.rs
+++ b/src/unix.rs
@@ -3,6 +3,7 @@ use libc::c_int;
 use std::fs::File;
 use std::io::{self, Read, Write};
 use std::mem;
+use std::mem::MaybeUninit;
 use std::os::unix::prelude::*;
 use std::process::Command;
 use std::ptr;
@@ -204,6 +205,12 @@ impl Client {
         format!("{},{}", self.read.as_raw_fd(), self.write.as_raw_fd())
     }
 
+    pub fn available(&self) -> io::Result<usize> {
+        let mut len = MaybeUninit::<c_int>::uninit();
+        cvt(unsafe { libc::ioctl(self.read.as_raw_fd(), libc::FIONREAD, len.as_mut_ptr()) })?;
+        Ok(unsafe { len.assume_init() } as usize)
+    }
+
     pub fn configure(&self, cmd: &mut Command) {
         // Here we basically just want to say that in the child process
         // we'll configure the read/write file descriptors to *not* be
diff --git a/src/wasm.rs b/src/wasm.rs
index b88a9d9..3793bd6 100644
--- a/src/wasm.rs
+++ b/src/wasm.rs
@@ -59,6 +59,11 @@ impl Client {
         );
     }
 
+    pub fn available(&self) -> io::Result<usize> {
+        let lock = self.inner.count.lock().unwrap_or_else(|e| e.into_inner());
+        Ok(*lock)
+    }
+
     pub fn configure(&self, _cmd: &mut Command) {
         unreachable!();
     }
diff --git a/src/windows.rs b/src/windows.rs
index d795c1c..6791efe 100644
--- a/src/windows.rs
+++ b/src/windows.rs
@@ -170,6 +170,26 @@ impl Client {
         self.name.clone()
     }
 
+    pub fn available(&self) -> io::Result<usize> {
+        // Can't read value of a semaphore on Windows, so
+        // try to acquire without sleeping, since we can find out the
+        // old value on release. If acquisiton fails, then available is 0.
+        unsafe {
+            let r = WaitForSingleObject(self.sem.0, 0);
+            if r != WAIT_OBJECT_0 {
+                Ok(0)
+            } else {
+                let mut prev: LONG = 0;
+                let r = ReleaseSemaphore(self.sem.0, 1, &mut prev);
+                if r != 0 {
+                    Ok(prev as usize + 1)
+                } else {
+                    Err(io::Error::last_os_error())
+                }
+            }
+        }
+    }
+
     pub fn configure(&self, _cmd: &mut Command) {
         // nothing to do here, we gave the name of our semaphore to the
         // child above
diff --git a/tests/server.rs b/tests/server.rs
index c92676e..70ea218 100644
--- a/tests/server.rs
+++ b/tests/server.rs
@@ -33,6 +33,30 @@ fn server_multiple() {
     drop((a, b));
 }
 
+#[test]
+fn server_available() {
+    let c = t!(Client::new(10));
+    assert_eq!(c.available().unwrap(), 10);
+    let a = c.acquire().unwrap();
+    assert_eq!(c.available().unwrap(), 9);
+    drop(a);
+    assert_eq!(c.available().unwrap(), 10);
+}
+
+#[test]
+fn server_none_available() {
+    let c = t!(Client::new(2));
+    assert_eq!(c.available().unwrap(), 2);
+    let a = c.acquire().unwrap();
+    assert_eq!(c.available().unwrap(), 1);
+    let b = c.acquire().unwrap();
+    assert_eq!(c.available().unwrap(), 0);
+    drop(a);
+    assert_eq!(c.available().unwrap(), 1);
+    drop(b);
+    assert_eq!(c.available().unwrap(), 2);
+}
+
 #[test]
 fn server_blocks() {
     let c = t!(Client::new(1));