From 8152adf1913fcad03f16c951cb82353cb47c4949 Mon Sep 17 00:00:00 2001
From: Rob David <rob.j.david@googlemail.com>
Date: Wed, 16 Jun 2021 22:36:07 +0100
Subject: [PATCH] Reduce CPU usage of wait group by blocking

---
 spec/wait_group_spec.cr   | 50 +++++++++++++++++++++++++++++++++++++++
 src/xmpp/event_manager.cr | 22 ++++++++++++-----
 2 files changed, 66 insertions(+), 6 deletions(-)
 create mode 100644 spec/wait_group_spec.cr

diff --git a/spec/wait_group_spec.cr b/spec/wait_group_spec.cr
new file mode 100644
index 0000000..898ffc9
--- /dev/null
+++ b/spec/wait_group_spec.cr
@@ -0,0 +1,50 @@
+require "./spec_helper"
+
+module XMPP
+  describe WaitGroup do
+
+    it "Should not block on creation" do
+      wg = WaitGroup.new
+      wg.wait
+    end
+
+    it "Should block whilst not done" do
+      wg = WaitGroup.new
+      wg.add.should eq 1
+      progress = Atomic(Int32).new(0)
+      spawn { progress.add(1); wg.wait; progress.add(1) }
+      sleep 5.milliseconds
+      # Should get to wait but not beyond it
+      progress.get.should eq 1
+      wg.done.should eq 0
+      while progress.get != 2
+         sleep 1.milliseconds
+      end
+    end
+
+    it "Should support multiple waiters" do
+      wg = WaitGroup.new
+      wg.add.should eq 1
+      unblocked = Atomic(Int32).new(0)
+      nwaiters = 10
+      (1..nwaiters).each { spawn { wg.wait; unblocked.add(1) } }
+      sleep 5.milliseconds
+      unblocked.get.should eq 0
+      wg.done.should eq 0
+      while unblocked.get != nwaiters
+        sleep 1.milliseconds
+      end
+    end
+
+    it "Won't block again after done" do
+      wg = WaitGroup.new
+      wg.add
+      wg.done
+      wg.done?.should be_true
+      wg.wait
+      wg.add
+      wg.done?.should be_true
+    end
+
+  end
+end
diff --git a/src/xmpp/event_manager.cr b/src/xmpp/event_manager.cr
index 5b0ea23..9ebe326 100644
--- a/src/xmpp/event_manager.cr
+++ b/src/xmpp/event_manager.cr
@@ -105,22 +105,32 @@ module XMPP
   private class WaitGroup
     def initialize
       @count = Atomic(Int32).new(0)
-      @span = Time::Span.new(nanoseconds: 5000)
+      @chan = Channel(Nil).new
     end
 
     def add(n = 1)
-      @count.add n
+      return @count.get if n == 0
+      count = @count.add(n) + n # New value
+      @chan.close if count <= 0
+      count
+    end
+
+    def count
+      @count.get
     end
 
     def done
       add(-1)
     end
 
+    def done?
+      @chan.closed?
+    end
+
     def wait
-      loop do
-        return if @count.get == 0
-        sleep(@span)
-      end
+      return if @count.get == 0 # Don't block when first constructed
+      @chan.receive
+    rescue Channel::ClosedError
     end
   end
 end