From ed4df21c85813eeec76c4c1d3e4bcd097bbe912f Mon Sep 17 00:00:00 2001 From: Mitchell Hashimoto Date: Wed, 7 Oct 2015 22:52:27 -0400 Subject: [PATCH] commands/snapshot: push and pop --- lib/vagrant/action.rb | 1 + lib/vagrant/action/builtin/is_env_set.rb | 23 +++++++ plugins/commands/snapshot/command/pop.rb | 28 ++++++++ plugins/commands/snapshot/command/push.rb | 33 ++++++++++ .../commands/snapshot/command/push_shared.rb | 65 +++++++++++++++++++ plugins/commands/snapshot/command/root.rb | 10 +++ plugins/providers/virtualbox/action.rb | 7 ++ 7 files changed, 167 insertions(+) create mode 100644 lib/vagrant/action/builtin/is_env_set.rb create mode 100644 plugins/commands/snapshot/command/pop.rb create mode 100644 plugins/commands/snapshot/command/push.rb create mode 100644 plugins/commands/snapshot/command/push_shared.rb diff --git a/lib/vagrant/action.rb b/lib/vagrant/action.rb index fcf8d916d..db4875dbb 100644 --- a/lib/vagrant/action.rb +++ b/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" diff --git a/lib/vagrant/action/builtin/is_env_set.rb b/lib/vagrant/action/builtin/is_env_set.rb new file mode 100644 index 000000000..d08ed9632 --- /dev/null +++ b/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 diff --git a/plugins/commands/snapshot/command/pop.rb b/plugins/commands/snapshot/command/pop.rb new file mode 100644 index 000000000..2c3499f43 --- /dev/null +++ b/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 diff --git a/plugins/commands/snapshot/command/push.rb b/plugins/commands/snapshot/command/push.rb new file mode 100644 index 000000000..c1168bcf2 --- /dev/null +++ b/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 diff --git a/plugins/commands/snapshot/command/push_shared.rb b/plugins/commands/snapshot/command/push_shared.rb new file mode 100644 index 000000000..880f8bcf5 --- /dev/null +++ b/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 diff --git a/plugins/commands/snapshot/command/root.rb b/plugins/commands/snapshot/command/root.rb index 03bcebb80..2ed460fa4 100644 --- a/plugins/commands/snapshot/command/root.rb +++ b/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 diff --git a/plugins/providers/virtualbox/action.rb b/plugins/providers/virtualbox/action.rb index d8b4d7ea1..12773736d 100644 --- a/plugins/providers/virtualbox/action.rb +++ b/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