From 0bbee7df4876b2701f781f1fb7351661ade83f46 Mon Sep 17 00:00:00 2001
From: Loic Nageleisen <loic.nageleisen@gmail.com>
Date: Thu, 17 Dec 2015 08:54:10 +0100
Subject: [PATCH] improve waitgroup

---
 lib/channel/waitgroup.rb | 17 ++++++++++-------
 test/test_waitgroup.rb   | 29 +++++++++++++++--------------
 2 files changed, 25 insertions(+), 21 deletions(-)

diff --git a/lib/channel/waitgroup.rb b/lib/channel/waitgroup.rb
index 38f2c50..a025e49 100644
--- a/lib/channel/waitgroup.rb
+++ b/lib/channel/waitgroup.rb
@@ -5,16 +5,19 @@ class WaitGroup
     @waiting = []
   end
 
-  def add(count)
-    sync! { @count += count }
+  def add(delta)
+    sync! do
+      @count += delta
+      fail 'negative WaitGroup counter' if @count < 0
+      if @waiting.any? && delta > 0 && @count == delta
+        fail 'misuse: add called concurrently with wait'
+      end
+      wake!
+    end
   end
 
   def done
-    sync! do
-      fail 'negative count' if done?
-      @count -= 1
-      wake!
-    end
+    add(-1)
   end
 
   private def done?
diff --git a/test/test_waitgroup.rb b/test/test_waitgroup.rb
index 867436c..b05922f 100644
--- a/test/test_waitgroup.rb
+++ b/test/test_waitgroup.rb
@@ -2,6 +2,9 @@ require 'test/unit'
 require 'thread'
 require 'channel/waitgroup'
 
+# rubocop:disable Metrics/AbcSize
+# rubocop:disable Metrics/MethodLength
+
 class TestWaitGroup < Test::Unit::TestCase
   module Util
     def meanwhile(*procs, &blk)
@@ -24,27 +27,25 @@ class TestWaitGroup < Test::Unit::TestCase
   def test_waitgroup
     Time.now.tap do |start|
       wg = WaitGroup.new
-      wg.add(5)
+      wg.add(2)
       ok1 = false
       ok2 = false
-      ok3 = false
-      ok4 = false
-      ok5 = false
-      go -> { sleep 0.1; ok1 = true; wg.done }
-      go -> { sleep 0.2; ok2 = true; wg.done }
-      go -> { sleep 0.3; ok3 = true; wg.done }
-      go -> { sleep 0.4; ok4 = true; wg.done }
-      go -> { sleep 0.5; ok5 = true; wg.done }
+      go -> { sleep 0.3; ok1 = true; wg.done }
+      go -> { sleep 0.5; ok2 = true; wg.done }
       wg.wait
       duration = Time.now - start
-      assert_equal(true, duration > 0.45)
-      assert_equal(true, duration < 0.70)
+      assert_equal(true, duration > 0.48)
+      assert_equal(true, duration < 0.52)
       assert_true(ok1)
       assert_true(ok2)
-      assert_true(ok3)
-      assert_true(ok4)
-      assert_true(ok5)
       assert_raises(RuntimeError) { wg.done }
     end
   end
+
+  def test_waitgroup_concurrent_add
+    wg = WaitGroup.new
+    go -> { wg.wait }
+    sleep 0.1
+    assert_raises(RuntimeError) { wg.add(1) }
+  end
 end