commands/snapshot: push and pop

lib/vagrant/action.rb

@@ -20,6 +20,7 @@ module Vagrant
       autoload :HandleBox, "vagrant/action/builtin/handle_box"
       autoload :HandleBoxUrl, "vagrant/action/builtin/handle_box_url"
       autoload :HandleForwardedPortCollisions, "vagrant/action/builtin/handle_forwarded_port_collisions"
+      autoload :IsEnvSet, "vagrant/action/builtin/is_env_set"
       autoload :IsState, "vagrant/action/builtin/is_state"
       autoload :Lock, "vagrant/action/builtin/lock"
       autoload :Message, "vagrant/action/builtin/message"

lib/vagrant/action/builtin/is_env_set.rb

@@ -0,0 +1,23 @@
+module Vagrant
+  module Action
+    module Builtin
+      # This middleware is meant to be used with Call and can check if
+      # a variable in env is set.
+      class IsEnvSet
+        def initialize(app, env, key, **opts)
+          @app    = app
+          @logger = Log4r::Logger.new("vagrant::action::builtin::is_env_set")
+          @key    = key
+          @invert = !!opts[:invert]
+        end
+
+        def call(env)
+          @logger.debug("Checking if env is set: '#{@key}'")
+          env[:result] = !!env[@key]
+          @logger.debug(" - Result: #{env[:result].inspect}")
+          @app.call(env)
+        end
+      end
+    end
+  end
+end

plugins/commands/snapshot/command/pop.rb

@@ -0,0 +1,28 @@
+require 'json'
+require 'optparse'
+
+require_relative "push_shared"
+
+module VagrantPlugins
+  module CommandSnapshot
+    module Command
+      class Pop < Vagrant.plugin("2", :command)
+        include PushShared
+
+        def execute
+          opts = OptionParser.new do |o|
+            o.banner = "Usage: vagrant snapshot pop [options] [vm-name]"
+            o.separator ""
+            o.separator "Restore state that was pushed with `vagrant snapshot push`."
+          end
+
+          # Parse the options
+          argv = parse_options(opts)
+          return if !argv
+
+          return shared_exec(argv, method(:pop))
+        end
+      end
+    end
+  end
+end

plugins/commands/snapshot/command/push.rb

@@ -0,0 +1,33 @@
+require 'json'
+require 'optparse'
+
+require_relative "push_shared"
+
+module VagrantPlugins
+  module CommandSnapshot
+    module Command
+      class Push < Vagrant.plugin("2", :command)
+        include PushShared
+
+        def execute
+          opts = OptionParser.new do |o|
+            o.banner = "Usage: vagrant snapshot push [options] [vm-name]"
+            o.separator ""
+            o.separator "Take a snapshot of the current state of the machine and 'push'"
+            o.separator "it onto the stack of states. You can use `vagrant snapshot pop`"
+            o.separator "to restore back to this state at any time."
+            o.separator ""
+            o.separator "If you use `vagrant snapshot save` or restore at any point after"
+            o.separator "a push, pop will still bring you back to this pushed state."
+          end
+
+          # Parse the options
+          argv = parse_options(opts)
+          return if !argv
+
+          return shared_exec(argv, method(:push))
+        end
+      end
+    end
+  end
+end

plugins/commands/snapshot/command/push_shared.rb

@@ -0,0 +1,65 @@
+require 'json'
+
+module VagrantPlugins
+  module CommandSnapshot
+    module Command
+      module PushShared
+        def shared_exec(argv, m)
+          with_target_vms(argv) do |vm|
+            if !vm.id
+              vm.ui.info("Not created. Cannot push snapshot state.")
+              next
+            end
+
+            vm.env.lock("machine-snapshot-stack") do
+              m.call(vm)
+            end
+          end
+
+          0
+        end
+
+        def push(machine)
+          snapshot_name = "push_#{Time.now.to_i}_#{rand(10000)}"
+
+          # Save the snapshot. This will raise an exception if it fails.
+          machine.action(:snapshot_save, snapshot_name: snapshot_name)
+
+          # Success! Write the resulting stack out
+          modify_snapshot_stack(machine) do |stack|
+            stack << snapshot_name
+          end
+        end
+
+        def pop(machine)
+          modify_snapshot_stack(machine) do |stack|
+            name = stack.pop
+
+            # Restore the snapshot and tell the provider to delete it as well.
+            machine.action(
+              :snapshot_restore,
+              snapshot_name: name,
+              snapshot_delete: true)
+          end
+        end
+
+        protected
+
+        def modify_snapshot_stack(machine)
+          # Get the stack
+          snapshot_stack = []
+          snapshot_file = machine.data_dir.join("snapshot_stack")
+          snapshot_stack = JSON.parse(snapshot_file.read) if snapshot_file.file?
+
+          # Yield it so it can be modified
+          yield snapshot_stack
+
+          # Write it out
+          snapshot_file.open("w+") do |f|
+            f.write(JSON.dump(snapshot_stack))
+          end
+        end
+      end
+    end
+  end
+end

plugins/commands/snapshot/command/root.rb

@@ -33,6 +33,16 @@ module VagrantPlugins
             require File.expand_path("../list", __FILE__)
             List
           end
+
+          @subcommands.register(:push) do
+            require File.expand_path("../push", __FILE__)
+            Push
+          end
+
+          @subcommands.register(:pop) do
+            require File.expand_path("../pop", __FILE__)
+            Pop
+          end
         end
 
         def execute

plugins/providers/virtualbox/action.rb

@@ -266,6 +266,13 @@ module VagrantPlugins
             b2.use EnvSet, force_halt: true
             b2.use action_halt
             b2.use SnapshotRestore
+
+            b2.use Call, IsEnvSet, :snapshot_delete do |env2, b3|
+              if env2[:result]
+                b3.use action_snapshot_delete
+              end
+            end
+
             b2.use action_start
           end
         end