1 files changed, 49 insertions, 24 deletions

diff --git a/crates/teleia/src/net/fig.rs b/crates/teleia/src/net/fig.rs
index 0d5e16d..d49e498 100644
--- a/crates/teleia/src/net/fig.rs
+++ b/crates/teleia/src/net/fig.rs
@@ -50,7 +50,6 @@ impl BinaryClient {
         self.out_buf.write_u32::<byteorder::LE>(buf.len() as u32).wrap_err("failed to send message")?;
         self.out_buf.write_all(buf).wrap_err("failed to send message")?;
         Ok(())
-        
     }
     pub fn publish(&mut self, ev: &[u8], data: &[u8]) -> Erm<()> {
         write!(self.out_buf, "p").wrap_err("failed to send publish message to bus")?;
@@ -58,7 +57,7 @@ impl BinaryClient {
         self.write_length_prefixed(data)?;
         Ok(())
     }
-    fn pop_incoming_message(&mut self) -> Option<BinaryMessage> {
+    pub fn pop_incoming_message(&mut self) -> Option<BinaryMessage> {
         let mut reader = std::io::Cursor::new(&self.in_buf);
         let event_len = reader.read_u32::<byteorder::LE>().ok()?;
         let mut event = vec![0; event_len as usize];
@@ -70,31 +69,57 @@ impl BinaryClient {
         self.in_buf.drain(..len);
         Some(BinaryMessage { event, data })
     }
-    pub fn pump(&mut self) -> Erm<Option<BinaryMessage>> {
-        let mut events = polling::Events::new();
-        events.clear();
-        self.poller.wait(&mut events, Some(std::time::Duration::from_secs(0)))
-            .wrap_err("failed to poll message bus")?;
+    fn try_read(&mut self) -> Erm<()> {
+        let mut buf = [0; 1024];
+        match self.socket.read(&mut buf) {
+            Ok(sz) => self.in_buf.extend_from_slice(&buf[..sz]),
+            Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {},
+            e => { e.wrap_err("failed to read from bus socket")?; },
+        }
+        Ok(())
+    }
+    fn try_write(&mut self) -> Erm<()> {
+        if !self.out_buf.is_empty() {
+            match self.socket.write(&self.out_buf) {
+                Ok(sz) => { self.out_buf.drain(..sz); },
+                Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {},
+                e => { e.wrap_err("failed to write to bus socket")?; },
+            }
+        }
+        Ok(())
+    }
+    fn resubscribe(&mut self) -> Erm<()> {
+        self.poller.modify(&self.socket, polling::Event::readable(KEY)).wrap_err("failed to update event to poll")?;
+        if !self.out_buf.is_empty() {
+            self.poller.modify(&self.socket, polling::Event::writable(KEY)).wrap_err("failed to update event to poll")?;
+        }
+        Ok(())
+    }
+    fn process_events(&mut self, events: &mut polling::Events) -> Erm<()> {
         for ev in events.iter() {
             if ev.key == KEY {
-                if ev.readable {
-                    let mut buf = [0; 1024];
-                    match self.socket.read(&mut buf) {
-                        Ok(sz) => self.in_buf.extend_from_slice(&buf[..sz]),
-                        Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {},
-                        e => { e.wrap_err("failed to read from bus socket")?; },
-                    }
-                }
-                if ev.writable && !self.out_buf.is_empty() {
-                    match self.socket.write(&self.out_buf) {
-                        Ok(sz) => { self.out_buf.drain(..sz); },
-                        Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {},
-                        e => { e.wrap_err("failed to write to bus socket")?; },
-                    }
-                }
-                self.poller.modify(&self.socket, polling::Event::all(KEY)).wrap_err("failed to update event to poll")?;
+                if ev.readable { self.try_read()?; }
+                if ev.writable { self.try_write()?; }
             }
         }
-        Ok(self.pop_incoming_message())
+        Ok(())
+    }
+    pub fn pump(&mut self) -> Erm<()> {
+        let mut events = polling::Events::new();
+        events.clear();
+        self.resubscribe()?;
+        self.poller.wait(&mut events, Some(std::time::Duration::from_secs(0)))
+            .wrap_err("failed to poll message bus")?;
+        self.process_events(&mut events)?;
+        Ok(())
+    }
+    pub fn wait(&mut self) -> Erm<()> {
+        let mut events = polling::Events::new();
+        events.clear();
+        self.resubscribe()?;
+        self.poller.wait(&mut events, None)
+            .wrap_err("failed to poll message bus")?;
+        self.process_events(&mut events)?;
+        Ok(())
     }
 }