From f1ad7234b95a87151b4fbb7b4122c20fce82c443 Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Sun, 18 Jul 2010 08:10:40 -0700
Subject: [PATCH] The new Vagrant::Util::Busy.

---
 lib/vagrant/util/busy.rb       |  59 ++++++++++++++++++
 test/vagrant/util/busy_test.rb | 106 +++++++++++++++++++++++++++++++++
 2 files changed, 165 insertions(+)
 create mode 100644 lib/vagrant/util/busy.rb
 create mode 100644 test/vagrant/util/busy_test.rb

diff --git a/lib/vagrant/util/busy.rb b/lib/vagrant/util/busy.rb
new file mode 100644
index 000000000..5dd6a6ad8
--- /dev/null
+++ b/lib/vagrant/util/busy.rb
@@ -0,0 +1,59 @@
+module Vagrant
+  module Util
+    # Utility class which allows blocks of code to be marked as "busy"
+    # with a specified interrupt handler. During busy areas of code, it
+    # is often undesirable for SIGINTs to immediately kill the application.
+    # This class is a helper to cleanly register callbacks to handle this
+    # situation.
+    class Busy
+      @@registered = []
+      @@mutex = Mutex.new
+
+      class << self
+        # Mark a given block of code as a "busy" block of code, which will
+        # register a SIGINT handler for the duration of the block. When a
+        # SIGINT occurs, the `sig_callback` proc will be called. It is up
+        # to the callback to behave properly and exit the application.
+        def busy(sig_callback)
+          register(sig_callback)
+          yield
+        ensure
+          unregister(sig_callback)
+        end
+
+        # Registers a SIGINT handler. This typically is called from {busy}.
+        # Callbacks are only registered once, so calling this multiple times
+        # with the same callback has no consequence.
+        def register(sig_callback)
+          @@mutex.synchronize do
+            registered << sig_callback
+            registered.uniq!
+
+            # Register the handler if this is our first callback.
+            Signal.trap("INT") { fire_callbacks } if registered.length == 1
+          end
+        end
+
+        # Unregisters a SIGINT handler.
+        def unregister(sig_callback)
+          @@mutex.synchronize do
+            registered.delete(sig_callback)
+
+            # Remove the signal trap if no more registered callbacks exist
+            Signal.trap("INT", "DEFAULT") if registered.empty?
+          end
+        end
+
+        # Fires all the registered callbacks.
+        def fire_callbacks
+          registered.each { |r| r.call }
+        end
+
+        # Helper method to get access to the class variable. This is mostly
+        # exposed for tests. This shouldn't be mucked with directly, since it's
+        # structure may change at any time.
+        def registered; @@registered; end
+      end
+    end
+  end
+end
diff --git a/test/vagrant/util/busy_test.rb b/test/vagrant/util/busy_test.rb
new file mode 100644
index 000000000..2279b6750
--- /dev/null
+++ b/test/vagrant/util/busy_test.rb
@@ -0,0 +1,106 @@
+require File.join(File.dirname(__FILE__), '..', '..', 'test_helper')
+
+class BusyUtilTest < Test::Unit::TestCase
+  setup do
+    @klass = Vagrant::Util::Busy
+  end
+
+  context "registering" do
+    setup do
+      @callback = lambda { puts "FOO" }
+      Signal.stubs(:trap)
+    end
+
+    teardown do
+      @klass.registered.clear
+    end
+
+    should "trap the signal on the first registration" do
+      Signal.expects(:trap).with("INT").once
+      @klass.register(@callback)
+      @klass.register(lambda { puts "BAR" })
+    end
+
+    should "not register the same callback multiple times" do
+      @klass.register(@callback)
+      @klass.register(@callback)
+      @klass.register(@callback)
+      assert_equal 1, @klass.registered.length
+      assert_equal @callback, @klass.registered.first
+    end
+  end
+
+  context "unregistering" do
+    setup do
+      Signal.stubs(:trap)
+
+      @callback = lambda { puts "FOO" }
+    end
+
+    teardown do
+      @klass.registered.clear
+    end
+
+    should "remove the callback and set the trap to DEFAULT when removing final" do
+      @klass.register(@callback)
+      Signal.expects(:trap).with("INT", "DEFAULT").once
+      @klass.unregister(@callback)
+      assert @klass.registered.empty?
+    end
+
+    should "not reset signal trap if not final callback" do
+      @klass.register(@callback)
+      @klass.register(lambda { puts "BAR" })
+      Signal.expects(:trap).never
+      @klass.unregister(@callback)
+    end
+  end
+
+  context "marking for busy" do
+    setup do
+      @callback = lambda { "foo" }
+    end
+
+    should "register, call the block, then unregister" do
+      waiter = mock("waiting")
+      proc = lambda { waiter.ping! }
+
+      seq = sequence('seq')
+      @klass.expects(:register).with(@callback).in_sequence(seq)
+      waiter.expects(:ping!).in_sequence(seq)
+      @klass.expects(:unregister).with(@callback).in_sequence(seq)
+
+      @klass.busy(@callback, &proc)
+    end
+
+    should "unregister callback even if block raises exception" do
+      waiter = mock("waiting")
+      proc = lambda { waiter.ping! }
+
+      seq = sequence('seq')
+      @klass.expects(:register).with(@callback).in_sequence(seq)
+      waiter.expects(:ping!).raises(Exception.new("uh oh!")).in_sequence(seq)
+      @klass.expects(:unregister).with(@callback).in_sequence(seq)
+
+      assert_raises(Exception) { @klass.busy(@callback, &proc) }
+    end
+  end
+
+  context "firing callbacks" do
+    setup do
+      Signal.stubs(:trap)
+    end
+
+    teardown do
+      @klass.registered.clear
+    end
+
+    should "just call the registered callbacks" do
+      waiting = mock("waiting")
+      waiting.expects(:ping!).once
+
+      @klass.register(lambda { waiting.ping! })
+      @klass.fire_callbacks
+    end
+  end
+end