diff --git a/lib/channel/waitgroup.rb b/lib/channel/waitgroup.rb new file mode 100644 index 0000000..38f2c50 --- /dev/null +++ b/lib/channel/waitgroup.rb @@ -0,0 +1,38 @@ +class WaitGroup + def initialize + @mutex = Mutex.new + @count = 0 + @waiting = [] + end + + def add(count) + sync! { @count += count } + end + + def done + sync! do + fail 'negative count' if done? + @count -= 1 + wake! + end + end + + private def done? + @count == 0 + end + + def wait + sync! do + @waiting << Thread.current + @mutex.sleep until done? + end + end + + private def wake! + @waiting.each { |t| t.wakeup if t.alive? } + end + + private def sync!(&block) + @mutex.synchronize(&block) + end +end diff --git a/test/test_waitgroup.rb b/test/test_waitgroup.rb new file mode 100644 index 0000000..867436c --- /dev/null +++ b/test/test_waitgroup.rb @@ -0,0 +1,50 @@ +require 'test/unit' +require 'thread' +require 'channel/waitgroup' + +class TestWaitGroup < Test::Unit::TestCase + module Util + def meanwhile(*procs, &blk) + threads = procs.map { |p| Thread.new(&p) } + blk.call + threads.each(&:join) + end + end + + module Assert + def assert_raise_with_message(exc, msg, &block) + e = assert_raise(exc, &block) + assert_match(msg, e.message) + end + end + + include Util + include Assert + + def test_waitgroup + Time.now.tap do |start| + wg = WaitGroup.new + wg.add(5) + 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 } + wg.wait + duration = Time.now - start + assert_equal(true, duration > 0.45) + assert_equal(true, duration < 0.70) + assert_true(ok1) + assert_true(ok2) + assert_true(ok3) + assert_true(ok4) + assert_true(ok5) + assert_raises(RuntimeError) { wg.done } + end + end +end