Make Input threadsafe (optionally) (hydro-project#8)

Author: justinj

covid_tracing/src/main.rs

@@ -1,12 +1,10 @@
 #![feature(never_type)]
 
-use std::sync::mpsc;
 use std::time::Duration;
 
 use hydroflow::compiled::pull::SymmetricHashJoin;
 use hydroflow::compiled::{ForEach, Pivot, Tee};
 use hydroflow::scheduled::collections::Iter;
-use hydroflow::scheduled::ctx::SendCtx;
 use hydroflow::scheduled::handoff::VecHandoff;
 use hydroflow::scheduled::Hydroflow;
 use hydroflow::{tl, tlt};
@@ -23,26 +21,11 @@ fn main() {
     type Phone = &'static str;
     type DateTime = usize;
 
-    let (contacts_send, contacts_recv) = mpsc::channel::<(Pid, Pid, DateTime)>();
-    let (diagnosed_send, diagnosed_recv) = mpsc::channel::<(Pid, (DateTime, DateTime))>();
-    let (people_send, people_recv) = mpsc::channel::<(Pid, (Name, Phone))>();
-
     let mut df = Hydroflow::new();
 
-    let contacts_out = df.add_source(move |send: &mut SendCtx<VecHandoff<_>>| {
-        send.give(Iter(contacts_recv.try_iter()));
-    });
-    let contacts_op_id = contacts_out.op_id();
-
-    let diagnosed_out = df.add_source(move |send: &mut SendCtx<VecHandoff<_>>| {
-        send.give(Iter(diagnosed_recv.try_iter()));
-    });
-    let diagnosed_op_id = diagnosed_out.op_id();
-
-    let people_out = df.add_source(move |send: &mut SendCtx<VecHandoff<_>>| {
-        send.give(Iter(people_recv.try_iter()));
-    });
-    let people_op_id = people_out.op_id();
+    let (contacts_send, contacts_out) = df.add_channel_input();
+    let (diagnosed_send, diagnosed_out) = df.add_channel_input();
+    let (people_send, people_out) = df.add_channel_input();
 
     let mut exposed_contacts = Default::default();
 
@@ -127,17 +110,12 @@ fn main() {
     df.add_edge(people_out, people_in);
     df.add_edge(notifs_out, notifs_in);
 
-    let reactor = df.reactor();
-
     let all_people = people::get_people();
 
     let inner = all_people.clone();
-    let inner_reactor = reactor.clone();
     std::thread::spawn(move || {
-        for person in inner {
-            people_send.send(person).unwrap();
-        }
-        inner_reactor.trigger(people_op_id).unwrap();
+        people_send.give(Iter(inner.into_iter()));
+        people_send.flush();
     });
 
     std::thread::spawn(move || {
@@ -152,10 +130,8 @@ fn main() {
                         let p1 = rng.gen_range(0..all_people.len());
                         let p2 = rng.gen_range(0..all_people.len());
                         if p1 != p2 {
-                            contacts_send
-                                .send((all_people[p1].0, all_people[p2].0, t))
-                                .unwrap();
-                            reactor.trigger(contacts_op_id).unwrap();
+                            contacts_send.give(Some((all_people[p1].0, all_people[p2].0, t)));
+                            contacts_send.flush();
                         }
                     }
                 }
@@ -164,9 +140,8 @@ fn main() {
                     if !all_people.is_empty() {
                         let p = rng.gen_range(0..all_people.len());
                         diagnosed_send
-                            .send((all_people[p].0, (t, t + TRANSMISSIBLE_DURATION)))
-                            .unwrap();
-                        reactor.trigger(diagnosed_op_id).unwrap();
+                            .give(Some((all_people[p].0, (t, t + TRANSMISSIBLE_DURATION))));
+                        diagnosed_send.flush();
                     }
                 }
                 _ => unreachable!(),

hydroflow/src/scheduled/input.rs

@@ -1,31 +1,67 @@
-use std::{cell::RefCell, rc::Rc};
+use std::{cell::RefCell, marker::PhantomData, rc::Rc, sync::mpsc::SyncSender};
 
 use super::{OpId, Reactor};
 
-pub struct Input<T> {
+pub trait Give<T> {
+    fn give(&self, t: T) -> bool;
+}
+
+pub struct Buffer<T>(pub(crate) Rc<RefCell<Vec<T>>>);
+impl<T> Give<T> for Buffer<T> {
+    fn give(&self, t: T) -> bool {
+        (*self.0).borrow_mut().push(t);
+        true
+    }
+}
+
+impl<T> Default for Buffer<T> {
+    fn default() -> Self {
+        Buffer(Rc::new(RefCell::new(Vec::new())))
+    }
+}
+
+impl<T> Clone for Buffer<T> {
+    fn clone(&self) -> Self {
+        Buffer(self.0.clone())
+    }
+}
+
+impl<T> Give<T> for SyncSender<T> {
+    fn give(&self, t: T) -> bool {
+        matches!(self.send(t), Ok(_))
+    }
+}
+
+// TODO(justin): this thing should probably give Vecs to the Givable, and buffer
+// stuff up and automatically flush, but postponing that until we have occasion
+// to benchmark it.
+pub struct Input<T, G>
+where
+    G: Give<T>,
+{
     reactor: Reactor,
     op_id: OpId,
-    data: Rc<RefCell<Vec<T>>>,
+    givable: G,
+    _marker: PhantomData<T>,
 }
-impl<T> Input<T> {
-    pub fn new(reactor: Reactor, op_id: OpId, data: Rc<RefCell<Vec<T>>>) -> Self {
+impl<T, G> Input<T, G>
+where
+    G: Give<T>,
+{
+    pub fn new(reactor: Reactor, op_id: OpId, givable: G) -> Self {
         Input {
             reactor,
             op_id,
-            data,
+            givable,
+            _marker: PhantomData,
         }
     }
 
     pub fn give(&self, t: T) {
-        (*self.data).borrow_mut().push(t);
+        self.givable.give(t);
     }
 
     pub fn flush(&self) {
         self.reactor.trigger(self.op_id).unwrap(/* TODO(justin) */);
     }
-
-    pub fn give_vec(&self, t: &mut Vec<T>) {
-        (*self.data).borrow_mut().extend(t.drain(..));
-        self.flush();
-    }
 }

hydroflow/src/scheduled/mod.rs

@@ -17,7 +17,7 @@ use handoff::{Handoff, HandoffList, HandoffMeta, NullHandoff, VecHandoff};
 use subgraph::{NtoMClosureSubgraph, Subgraph, VariadicClosureSubgraph};
 
 use self::handoff::CanReceive;
-use self::input::Input;
+use self::input::{Buffer, Input};
 
 pub type OpId = usize;
 pub type HandoffId = usize;
@@ -286,25 +286,41 @@ impl Hydroflow {
 
     /**
      * Adds an "input" operator, along with a handle to insert data into it.
-     * TODO(justin): this thing currently cannot be passed across threads, but I
-     * think if we replaced the Vec with a channel it could.
      */
-    pub fn add_input<T, W>(&mut self) -> (Input<T>, OutputPort<W>)
+    pub fn add_input<T, W>(&mut self) -> (Input<T, Buffer<T>>, OutputPort<W>)
     where
         T: 'static,
         W: 'static + Handoff + CanReceive<T>,
     {
-        let input = Rc::new(RefCell::new(Vec::new()));
+        let input = Buffer::default();
         let inner_input = input.clone();
         let output_port = self.add_source::<_, W>(move |send| {
-            for x in (*inner_input).borrow_mut().drain(..) {
+            for x in (*inner_input.0).borrow_mut().drain(..) {
                 send.give(x);
             }
         });
         let id = output_port.op_id;
         (Input::new(self.reactor(), id, input), output_port)
     }
 
+    /**
+     * Adds a threadsafe "input" operator, along with a handle to insert data into it.
+     */
+    pub fn add_channel_input<T, W>(&mut self) -> (Input<T, SyncSender<T>>, OutputPort<W>)
+    where
+        T: 'static,
+        W: 'static + Handoff + CanReceive<T>,
+    {
+        let (sender, receiver) = mpsc::sync_channel(8000);
+        let output_port = self.add_source::<_, W>(move |send| {
+            for x in receiver.try_iter() {
+                send.give(x);
+            }
+        });
+        let id = output_port.op_id;
+        (Input::new(self.reactor(), id, sender), output_port)
+    }
+
     /**
      * Adds a new compiled subgraph with no inputs and one output.
      */
@@ -351,7 +367,7 @@ impl Hydroflow {
 }
 
 /**
- * A handle into a specific [Hydroflow] instance for triggering operators to run, possibly from another thread.Default
+ * A handle into a specific [Hydroflow] instance for triggering operators to run, possibly from another thread.
  */
 #[derive(Clone)]
 pub struct Reactor {
@@ -650,3 +666,36 @@ fn test_input_handle() {
 
     assert_eq!((*vec).borrow().clone(), vec![1, 2, 3, 4, 5, 6]);
 }
+
+#[test]
+fn test_input_handle_thread() {
+    let mut df = Hydroflow::new();
+
+    let (input, output_port) = df.add_channel_input();
+
+    let vec = Rc::new(RefCell::new(Vec::new()));
+    let inner_vec = vec.clone();
+    let input_port = df.add_sink(move |recv: &mut RecvCtx<VecHandoff<usize>>| {
+        for v in recv.take_inner() {
+            (*inner_vec).borrow_mut().push(v);
+        }
+    });
+
+    df.add_edge(output_port, input_port);
+
+    let (done, wait) = mpsc::channel();
+
+    std::thread::spawn(move || {
+        input.give(Some(1));
+        input.give(Some(2));
+        input.give(Some(3));
+        input.flush();
+        done.send(()).unwrap();
+    });
+
+    wait.recv().unwrap();
+
+    df.tick();
+
+    assert_eq!((*vec).borrow().clone(), vec![1, 2, 3]);
+}