Fix races in collector_plugin test

Summary:
There are a few races in the test that show up in stress test. We need
to use the barrier a few more times.

Reviewed By: dschatzberg

Differential Revision: D34598821

fbshipit-source-id: d97cf5f425b1fd58e4bbbd708a0ba77ed5f690a2

1 files changed, 12 insertions, 3 deletions

diff --git a/below/model/src/collector_plugin.rs b/below/model/src/collector_plugin.rs
index efcd0345..4f55ddc3 100644
--- a/below/model/src/collector_plugin.rs
+++ b/below/model/src/collector_plugin.rs
@@ -151,22 +151,31 @@ mod test {
             // Test overwriting sample
             futures::executor::block_on(collector.collect_and_update()).unwrap();
             c.wait(); // <-- 1
+            // Consumer checking overwritten sample
+            c.wait(); // <-- 2
             // Test sending None
             futures::executor::block_on(collector.collect_and_update()).unwrap();
-            c.wait(); // <-- 2
+            c.wait(); // <-- 3
+            // Consumer checking None
+            c.wait(); // <-- 4
             // Test sending error. Will fail on both collector and consumer threads.
             let is_error = matches!(
                 futures::executor::block_on(collector.collect_and_update()),
                 Err(_)
             );
-            c.wait(); // <-- 3
+            c.wait(); // <-- 5
             assert!(is_error, "Collector did not return an error");
         });
+        // Collector overwriting sample
         barrier.wait(); // <-- 1
         assert_eq!(Some(2), consumer.try_take().unwrap());
         barrier.wait(); // <-- 2
-        assert_eq!(None, consumer.try_take().unwrap());
+        // Collector sending None
         barrier.wait(); // <-- 3
+        assert_eq!(None, consumer.try_take().unwrap());
+        barrier.wait(); // <-- 4
+        // Collector sending error
+        barrier.wait(); // <-- 5
         assert!(matches!(consumer.try_take(), Err(_)));
 
         handle.join().unwrap();