Merge pull request #117 from sschneid/host_snapshots

(QENG-2636) Host snapshots

lib/vmpooler/api/reroute.rb

@@ -50,6 +50,14 @@ module Vmpooler
     put '/vm/:hostname/?' do
       call env.merge("PATH_INFO" => "/api/v#{api_version}/vm/#{params[:hostname]}")
     end
+
+    post '/vm/:hostname/snapshot/?' do
+      call env.merge("PATH_INFO" => "/api/v#{api_version}/vm/#{params[:hostname]}/snapshot")
+    end
+
+    post '/vm/:hostname/snapshot/:snapshot/?' do
+      call env.merge("PATH_INFO" => "/api/v#{api_version}/vm/#{params[:hostname]}/snapshot/#{params[:snapshot]}")
+    end
   end
   end
 end

lib/vmpooler/api/v1.rb

@@ -413,6 +413,11 @@ module Vmpooler
             result[params[:hostname]]['tags'] ||= {}
             result[params[:hostname]]['tags'][$1] = rdata[key]
           end
+
+          if key.match('^snapshot\:(.+?)$')
+            result[params[:hostname]]['snapshots'] ||= []
+            result[params[:hostname]]['snapshots'].push($1)
+          end
         end
 
         if config['domain']
@@ -509,6 +514,51 @@ module Vmpooler
 
       JSON.pretty_generate(result)
     end
+
+    post "#{api_prefix}/vm/:hostname/snapshot/?" do
+      content_type :json
+
+      need_token! if Vmpooler::API.settings.config[:auth]
+
+      status 404
+      result = { 'ok' => false }
+
+      params[:hostname] = hostname_shorten(params[:hostname], config['domain'])
+
+      if backend.exists('vmpooler__vm__' + params[:hostname])
+        result[params[:hostname]] = {}
+
+        o = [('a'..'z'), ('0'..'9')].map(&:to_a).flatten
+        result[params[:hostname]]['snapshot'] = o[rand(25)] + (0...31).map { o[rand(o.length)] }.join
+
+        backend.sadd('vmpooler__tasks__snapshot', params[:hostname] + ':' + result[params[:hostname]]['snapshot'])
+
+        status 202
+        result['ok'] = true
+      end
+
+      JSON.pretty_generate(result)
+    end
+
+    post "#{api_prefix}/vm/:hostname/snapshot/:snapshot/?" do
+      content_type :json
+
+      need_token! if Vmpooler::API.settings.config[:auth]
+
+      status 404
+      result = { 'ok' => false }
+
+      params[:hostname] = hostname_shorten(params[:hostname], config['domain'])
+
+      if backend.exists('vmpooler__vm__' + params[:hostname]) and backend.hget('vmpooler__vm__' + params[:hostname], 'snapshot:' + params[:snapshot])
+        backend.sadd('vmpooler__tasks__snapshot-revert', params[:hostname] + ':' + params[:snapshot])
+
+        status 202
+        result['ok'] = true
+      end
+
+      JSON.pretty_generate(result)
+    end
   end
   end
 end

lib/vmpooler/pool_manager.rb

@@ -292,6 +292,100 @@ module Vmpooler
       end
     end
 
+    def create_vm_snapshot(vm, snapshot_name)
+      Thread.new do
+        _create_vm_snapshot(vm, snapshot_name)
+      end
+    end
+
+    def _create_vm_snapshot(vm, snapshot_name)
+      host = $vsphere['snapshot_manager'].find_vm(vm) ||
+             $vsphere['snapshot_manager'].find_vm_heavy(vm)[vm]
+
+      if (host) && ((! snapshot_name.nil?) || (! snapshot_name.empty?))
+        $logger.log('s', "[ ] [snapshot_manager] '#{vm}' is being snapshotted")
+
+        start = Time.now
+
+        host.CreateSnapshot_Task(
+          name: snapshot_name,
+          description: 'vmpooler',
+          memory: true,
+          quiesce: true
+        ).wait_for_completion
+
+        finish = '%.2f' % (Time.now - start)
+
+        $redis.hset('vmpooler__vm__' + vm, 'snapshot:' + snapshot_name, Time.now.to_s)
+
+        $logger.log('s', "[+] [snapshot_manager] '#{vm}' snapshot created in #{finish} seconds")
+      end
+    end
+
+    def revert_vm_snapshot(vm, snapshot_name)
+      Thread.new do
+        _revert_vm_snapshot(vm, snapshot_name)
+      end
+    end
+
+    def _revert_vm_snapshot(vm, snapshot_name)
+      host = $vsphere['snapshot_manager'].find_vm(vm) ||
+             $vsphere['snapshot_manager'].find_vm_heavy(vm)[vm]
+
+      if host
+        snapshot = $vsphere['snapshot_manager'].find_snapshot(host, snapshot_name)
+
+        if snapshot
+          $logger.log('s', "[ ] [snapshot_manager] '#{vm}' is being reverted to snapshot '#{snapshot_name}'")
+
+          start = Time.now
+
+          snapshot.RevertToSnapshot_Task.wait_for_completion
+
+          finish = '%.2f' % (Time.now - start)
+
+          $logger.log('s', "[<] [snapshot_manager] '#{vm}' reverted to snapshot in #{finish} seconds")
+        end
+      end
+    end
+
+    def check_snapshot_queue
+      $logger.log('d', "[*] [snapshot_manager] starting worker thread")
+
+      $vsphere['snapshot_manager'] ||= Vmpooler::VsphereHelper.new
+
+      $threads['snapshot_manager'] = Thread.new do
+        loop do
+          _check_snapshot_queue
+          sleep(5)
+        end
+      end
+    end
+
+    def _check_snapshot_queue
+      vm = $redis.spop('vmpooler__tasks__snapshot')
+
+      unless vm.nil?
+        begin
+          vm_name, snapshot_name = vm.split(':')
+          create_vm_snapshot(vm_name, snapshot_name)
+        rescue
+          $logger.log('s', "[!] [snapshot_manager] snapshot appears to have failed")
+        end
+      end
+
+      vm = $redis.spop('vmpooler__tasks__snapshot-revert')
+
+      unless vm.nil?
+        begin
+          vm_name, snapshot_name = vm.split(':')
+          revert_vm_snapshot(vm_name, snapshot_name)
+        rescue
+          $logger.log('s', "[!] [snapshot_manager] snapshot revert appears to have failed")
+        end
+      end
+    end
+
     def check_pool(pool)
       $logger.log('d', "[*] [#{pool['name']}] starting worker thread")
 
@@ -443,16 +537,21 @@ module Vmpooler
       $redis.set('vmpooler__tasks__clone', 0)
 
       loop do
+        if ! $threads['snapshot_manager']
+          check_snapshot_queue
+        elsif ! $threads['snapshot_manager'].alive?
+          $logger.log('d', "[!] [snapshot_manager] worker thread died, restarting")
+          check_snapshot_queue
+        end
+
         $config[:pools].each do |pool|
           if ! $threads[pool['name']]
             check_pool(pool)
-          else
+          elsif ! $threads[pool['name']].alive?
-            unless $threads[pool['name']].alive?
             $logger.log('d', "[!] [#{pool['name']}] worker thread died, restarting")
             check_pool(pool)
           end
         end
-        end
 
         sleep(1)
       end

lib/vmpooler/vsphere_helper.rb

@@ -99,6 +99,12 @@ module Vmpooler
       base
     end
 
+    def find_snapshot(vm, snapshotname)
+      if vm.snapshot
+        get_snapshot_list(vm.snapshot.rootSnapshotList, snapshotname)
+      end
+    end
+
     def find_vm(vmname)
       begin
         @connection.serviceInstance.CurrentTime
@@ -178,6 +184,20 @@ module Vmpooler
       )
     end
 
+    def get_snapshot_list(tree, snapshotname)
+      snapshot = nil
+
+      tree.each do |child|
+        if child.name == snapshotname
+          snapshot ||= child.snapshot
+        else
+          snapshot ||= get_snapshot_list(child.childSnapshotList, snapshotname)
+        end
+      end
+
+      snapshot
+    end
+
     def close
       @connection.close
     end

spec/vmpooler/api/v1_spec.rb

@@ -432,6 +432,95 @@ describe Vmpooler::API::V1 do
         end
       end
     end
+
+    describe 'POST /vm/:hostname/snapshot' do
+      context '(auth not configured)' do
+        let(:config) { { auth: false } }
+
+        it 'creates a snapshot' do
+          expect(redis).to receive(:sadd)
+
+          post "#{prefix}/vm/testhost/snapshot"
+
+          expect(last_response.header['Content-Type']).to eq('application/json')
+          expect(JSON.parse(last_response.body)['ok']).to eq(true)
+          expect(JSON.parse(last_response.body)['testhost']['snapshot'].length).to be(32)
+          expect(last_response.status).to eq(202)
+        end
+      end
+
+      context '(auth configured)' do
+        let(:config) { { auth: true } }
+
+        it 'returns a 401 if not authed' do
+          post "#{prefix}/vm/testhost/snapshot"
+
+          expect(last_response).not_to be_ok
+          expect(last_response.header['Content-Type']).to eq('application/json')
+          expect(last_response.body).to eq(JSON.pretty_generate({'ok' => false}))
+          expect(last_response.status).to eq(401)
+        end
+
+        it 'creates a snapshot if authed' do
+          expect(redis).to receive(:sadd)
+
+          post "#{prefix}/vm/testhost/snapshot", "", {
+            'HTTP_X_AUTH_TOKEN' => 'abcdefghijklmnopqrstuvwxyz012345'
+          }
+
+          expect(last_response.header['Content-Type']).to eq('application/json')
+          expect(JSON.parse(last_response.body)['ok']).to eq(true)
+          expect(JSON.parse(last_response.body)['testhost']['snapshot'].length).to be(32)
+          expect(last_response.status).to eq(202)
+        end
+      end
+    end
+
+    describe 'POST /vm/:hostname/snapshot/:snapshot' do
+      context '(auth not configured)' do
+        let(:config) { { auth: false } }
+
+        it 'reverts to a snapshot' do
+          expect(redis).to receive(:exists).with('vmpooler__vm__testhost').and_return(1)
+          expect(redis).to receive(:hget).with('vmpooler__vm__testhost', 'snapshot:testsnapshot').and_return(1)
+          expect(redis).to receive(:sadd)
+
+          post "#{prefix}/vm/testhost/snapshot/testsnapshot"
+
+          expect(last_response.header['Content-Type']).to eq('application/json')
+          expect(last_response.body).to include('"ok": true')
+          expect(last_response.status).to eq(202)
+        end
+      end
+
+      context '(auth configured)' do
+        let(:config) { { auth: true } }
+
+        it 'returns a 401 if not authed' do
+          post "#{prefix}/vm/testhost/snapshot"
+
+          expect(last_response).not_to be_ok
+          expect(last_response.header['Content-Type']).to eq('application/json')
+          expect(last_response.body).to eq(JSON.pretty_generate({'ok' => false}))
+          expect(last_response.status).to eq(401)
+        end
+
+        it 'reverts to a snapshot if authed' do
+          expect(redis).to receive(:exists).with('vmpooler__vm__testhost').and_return(1)
+          expect(redis).to receive(:hget).with('vmpooler__vm__testhost', 'snapshot:testsnapshot').and_return(1)
+          expect(redis).to receive(:sadd)
+
+          post "#{prefix}/vm/testhost/snapshot/testsnapshot", "", {
+            'HTTP_X_AUTH_TOKEN' => 'abcdefghijklmnopqrstuvwxyz012345'
+          }
+
+          expect(last_response.header['Content-Type']).to eq('application/json')
+          expect(last_response.body).to include('"ok": true')
+          expect(last_response.status).to eq(202)
+        end
+      end
+
+    end
   end
 
 end

spec/vmpooler/pool_manager_spec.rb

@@ -248,4 +248,72 @@ describe 'Pool Manager' do
     end
   end
 
+  describe '#_create_vm_snapshot' do
+    let(:snapshot_manager) { 'snapshot_manager' }
+    let(:pool_helper) { double('snapshot_manager') }
+    let(:vsphere) { {snapshot_manager => pool_helper} }
+
+    before do
+      expect(subject).not_to be_nil
+      $vsphere = vsphere
+    end
+
+    context '(valid host)' do
+      let(:vm_host) { double('vmhost') }
+
+      it 'creates a snapshot' do
+        expect(pool_helper).to receive(:find_vm).and_return vm_host
+        expect(logger).to receive(:log)
+        expect(vm_host).to receive_message_chain(:CreateSnapshot_Task, :wait_for_completion)
+        expect(redis).to receive(:hset).with('vmpooler__vm__testvm', 'snapshot:testsnapshot', Time.now.to_s)
+        expect(logger).to receive(:log)
+
+        subject._create_vm_snapshot('testvm', 'testsnapshot')
+      end
+    end
+  end
+
+  describe '#_revert_vm_snapshot' do
+    let(:snapshot_manager) { 'snapshot_manager' }
+    let(:pool_helper) { double('snapshot_manager') }
+    let(:vsphere) { {snapshot_manager => pool_helper} }
+
+    before do
+      expect(subject).not_to be_nil
+      $vsphere = vsphere
+    end
+
+    context '(valid host)' do
+      let(:vm_host) { double('vmhost') }
+      let(:vm_snapshot) { double('vmsnapshot') }
+
+      it 'reverts a snapshot' do
+        expect(pool_helper).to receive(:find_vm).and_return vm_host
+        expect(pool_helper).to receive(:find_snapshot).and_return vm_snapshot
+        expect(logger).to receive(:log)
+        expect(vm_snapshot).to receive_message_chain(:RevertToSnapshot_Task, :wait_for_completion)
+        expect(logger).to receive(:log)
+
+        subject._revert_vm_snapshot('testvm', 'testsnapshot')
+      end
+    end
+  end
+
+  describe '#_check_snapshot_queue' do
+    let(:pool_helper) { double('pool') }
+    let(:vsphere) { {pool => pool_helper} }
+
+    before do
+      expect(subject).not_to be_nil
+      $vsphere = vsphere
+    end
+
+    it 'checks appropriate redis queues' do
+      expect(redis).to receive(:spop).with('vmpooler__tasks__snapshot')
+      expect(redis).to receive(:spop).with('vmpooler__tasks__snapshot-revert')
+
+      subject._check_snapshot_queue
+    end
+  end
+
 end