GracefulHalt builtin

lib/vagrant/action.rb

@@ -14,6 +14,7 @@ module Vagrant
       autoload :Confirm, "vagrant/action/builtin/confirm"
       autoload :ConfigValidate, "vagrant/action/builtin/config_validate"
       autoload :EnvSet,  "vagrant/action/builtin/env_set"
+      autoload :GracefulHalt, "vagrant/action/builtin/graceful_halt"
       autoload :Provision, "vagrant/action/builtin/provision"
       autoload :SSHExec, "vagrant/action/builtin/ssh_exec"
       autoload :SSHRun,  "vagrant/action/builtin/ssh_run"

lib/vagrant/action/builtin/graceful_halt.rb

@@ -0,0 +1,67 @@
+require "log4r"
+
+module Vagrant
+  module Action
+    module Builtin
+      # This middleware class will attempt to perform a graceful shutdown
+      # of the machine using the guest implementation. This middleware is
+      # compatible with the {Call} middleware so you can branch based on
+      # the result, which is true if the halt succeeded and false otherwise.
+      class GracefulHalt
+        # Note: Any of the arguments can be arrays as well.
+        #
+        # @param [Symbol] target_state The target state ID that means that
+        #   the machine was properly shut down.
+        # @param [Symbol] source_state The source state ID that the machine
+        #   must be in to be shut down.
+        def initialize(app, env, target_state, source_state=nil)
+          @app          = app
+          @logger       = Log4r::Logger.new("vagrant::action::builtin::graceful_halt")
+          @source_state = source_state
+          @target_state = target_state
+        end
+
+        def call(env)
+          graceful = true
+          graceful = !env[:force_halt] if env.has_key?(:force_halt)
+
+          # By default, we didn't succeed.
+          env[:result] = false
+
+          if graceful && @source_state
+            @logger.info("Verifying source state of machine: #{@source_state.inspect}")
+
+            # If we're not in the proper source state, then we don't
+            # attempt to halt the machine
+            current_state = env[:machine].state.id
+            if current_state != @source_state
+              @logger.info("Invalid source state, not halting: #{current_state}")
+              graceful = false
+            end
+          end
+
+          # Only attempt to perform graceful shutdown under certain cases
+          # checked above.
+          if graceful
+            env[:ui].info I18n.t("vagrant.actions.vm.halt.graceful")
+            env[:machine].guest.halt
+
+            @logger.debug("Waiting for target graceful halt state: #{@target_state}")
+            count = 0
+            while env[:machine].state.id != @target_state
+              count += 1
+              return if count >= env[:machine].config.vm.graceful_halt_retry_count
+              sleep env[:machine].config.vm.graceful_halt_retry_interval
+            end
+
+            # The result of this matters on whether we reached our
+            # proper target state or not.
+            env[:result] = env[:machine].state.id == @target_state
+          end
+
+          @app.call(env)
+        end
+      end
+    end
+  end
+end

test/unit/vagrant/action/builtin/graceful_halt_test.rb

@@ -0,0 +1,61 @@
+require File.expand_path("../../../../base", __FILE__)
+
+describe Vagrant::Action::Builtin::GracefulHalt do
+  let(:app) { lambda { |env| } }
+  let(:env) { { :machine => machine, :ui => ui } }
+  let(:machine) do
+    result = double("machine")
+    result.stub(:config).and_return(machine_config)
+    result.stub(:guest).and_return(machine_guest)
+    result.stub(:state).and_return(machine_state)
+    result
+  end
+  let(:machine_config) do
+    double("machine_config").tap do |top_config|
+      vm_config = double("machien_vm_config")
+      vm_config.stub(:graceful_halt_retry_count => 2)
+      vm_config.stub(:graceful_halt_retry_interval => 0)
+      top_config.stub(:vm => vm_config)
+    end
+  end
+  let(:machine_guest) { double("machine_guest") }
+  let(:machine_state) do
+    double("machine_state").tap do |result|
+      result.stub(:id).and_return(:unknown)
+    end
+  end
+  let(:target_state) { :target }
+  let(:ui) do
+    double("ui").tap do |result|
+      result.stub(:info)
+    end
+  end
+
+  it "should do nothing if force is specified" do
+    env[:force_halt] = true
+
+    machine_guest.should_not_receive(:halt)
+
+    described_class.new(app, env, target_state).call(env)
+
+    env[:result].should == false
+  end
+
+  it "should do nothing if there is an invalid source state" do
+    machine_state.stub(:id).and_return(:invalid_source)
+    machine_guest.should_not_receive(:halt)
+
+    described_class.new(app, env, target_state, :target_source).call(env)
+
+    env[:result].should == false
+  end
+
+  it "should gracefully halt and wait for the target state" do
+    machine_guest.should_receive(:halt).once
+    machine_state.stub(:id).and_return(target_state)
+
+    described_class.new(app, env, target_state).call(env)
+
+    env[:result].should == true
+  end
+end