floatyhelper 2.0.6 → 2.1.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 12632be13fdd03beaa837d72ea01a0ba7fff08494a79ed4f839070debfc35e89
-  data.tar.gz: d9ca928fc876f292aca393aa4b1071508797b4883e3e64797a2ff75fb436f12e
+  metadata.gz: a0a23b077fb6f7d707267d06af856f7eb2c5854bd1ae65b325fd8098a57c5623
+  data.tar.gz: 812e17bf9526dc862df4095119ee388d6e4e8113773a7311600fe6e35ccde743
 SHA512:
-  metadata.gz: ec8bc4eedba90e8457588265add5cdf55924f4583f65b9da4e5e1404771d009b1cfb5cbf09abcb8376ff7eeeb5cdf65b8f96d63ad5643fcaf11dcae9663489d1
-  data.tar.gz: 3fa22d4d9a06edae1e6bf5db430b9b383f36549666c59056dfb5a87adc754b652da76af3ed0c1dd35f1222e8ca4f486d6298ce6951fff2cf7f580a122f2d8a8d
+  metadata.gz: 70b52c8b29fcb0dd530bd39efd69961da64784383705c57fc315e6c7bd013f9acc5339c46d66aeca3bdcb549bdc3b11f0b40e7ed0a987e93867ebd596c6afcb2
+  data.tar.gz: fe9a5d8fae49c366df1c6a125c1600d9d11b6c9c759bcc25e4aec7f2e9fce3a6a28785a61e42c9fae79efc943307694f0e71b39114b2e796e71637749040e8b8

data/lib/floatyhelper/floaty.rb

@@ -3,6 +3,7 @@ require 'json'
 require 'open3'
 require 'colorize'
 require 'pty'
+require 'thread'
 
 class Floaty
   FLOATY_SERVICES = {
@@ -137,10 +138,16 @@ class Floaty
         while !Process.waitpid(io.pid, Process::WNOHANG)
           line = io.gets
           unless line.nil?
-            puts line
+            # Always print initial job_id request, VM results, or if we aren't running in a separate thread.
+            puts line if line =~ /^\s*-\s*[\w-]+\./ || line =~ /Requesting VMs/ || Thread.current['print'].nil? || Thread.current['print'] == true
             output += line
           end
         end
+        # If the process finishes before we get all the output, get the rest and handle it.
+        while line = io.gets
+          puts line if line =~ /^\s*-\s*[\w-]+\./ || Thread.current['print'].nil? || Thread.current['print'] == true
+          output += line
+        end
       end
       # Need to check exit status
       output

data/lib/floatyhelper/version.rb

@@ -1,3 +1,3 @@
 module FloatyhelperVersion
-  VERSION = '2.0.6'.freeze
+  VERSION = '2.1.0'.freeze
 end

data/lib/floatyhelper/vm.rb

@@ -7,8 +7,10 @@ require 'floatyhelper/config'
 require 'colorize'
 require 'net/http'
 require 'json'
+require 'thread'
 
 class VM
+  # Assume we only have 1 of each if we can't find the info
   def self.pooled_platforms_default
     [
       'centos-7-x86_64',
@@ -17,10 +19,13 @@ class VM
       'redhat-7-x86_64',
       'redhat-8-x86_64',
       'redhat-fips-7-x86_64',
+      'redhat-fips-8-x86_64',
       'scientific-7-x86_64',
       'sles-12-x86_64',
+      'sles-15-x86_64',
       'ubuntu-1804-x86_64',
-    ]
+      'ubuntu-2004-x86_64',
+    ].map { |x| [x,1] }.to_h
   end
 
   def self.find_pooled_platforms
@@ -33,7 +38,7 @@ class VM
       result = JSON.parse(result.body)
       # Techinally, 'max > 0' tells you if it's a pooled platform, but if
       # the pool is empty, we'll want to fall back to ABS anyway.
-      result['pools'].select { |_pool, info| info['ready'].positive? }.map { |pool, _info| pool.gsub('-pixa4','') }
+      result['pools'].select { |_pool, info| info['ready'].positive? }.map { |pool, info| [pool.gsub('-pixa4','').gsub('-pooled',''), info['ready']] }.to_h
     rescue StandardError
       # Not a great practice to swallow all errors, but this list is probably
       # pretty stable, so let's just pass along the default.
@@ -48,7 +53,7 @@ class VM
     Groups.delete_tag(id) if Groups.tag?(id)
     Groups.delete_all if id == 'all'
     hosts = hosts.select { |host| alive(host) }
-    puts Floaty.floaty_cmd("delete #{hosts.join(',')} --service vmpooler") unless hosts.empty?
+    hosts.each_slice(10) { |s| puts Floaty.floaty_cmd("delete #{s.join(',')} --service vmpooler") } unless hosts.empty?
   end
 
   def self.query(host)
@@ -149,28 +154,98 @@ class VM
   end
 
   def self.revert(id, snaptag)
+    clr = Config.get_config_setting('vertical_snapshot_status')
+    clr = clr.to_s.downcase == 'true'
     snaptag ||= 'Blank Snaptag'
     shas = VM.getsnapshot(id, snaptag)
+    lastlog = {}
+    puts `tput clear` if clr
     shas.each do |host, sha|
+      # Force a session to show up in lastlog
+      lastlog[host] = `ssh -t root@#{host} lastlog -u root 2>/dev/null`.gsub("\r",'').gsub("\n",'')
       print "#{host}: #{sha} --- "
       puts Floaty.floaty_cmd("revert #{host} #{sha} --service vmpooler")
     end
-    puts 'Waiting 10 seconds for revert to take effect'
-    sleep(10)
+    puts
+    puts 'Waiting for all VMs to finish reverting...'
+    alldone = false
+    until alldone
+      puts `tput cup #{shas.count + 2}` if clr
+      alldone = true
+      print "\r" unless clr
+      lastlog.each do |host, value|
+        begin
+          result =`ssh root@#{host} lastlog -u root 2>/dev/null`.gsub("\r",'').gsub("\n",'')
+          # Original will have some garbage about pseudo-terminal that we won't have here
+          # Also, can't do this on Windows yet, so pretend we're done for those hosts.
+          done = !lastlog[host].include?(result) || result.empty?
+        rescue
+          done = false
+        end
+        status = done ? 'Done'.green : 'Wait'.yellow
+        if clr
+          puts "* %s #{status} *" % "#{host}:".ljust(16)
+        else
+          print "* %s #{status} *" % "#{host}:".ljust(16)
+        end
+        alldone &= done
+      end
+      sleep(1)
+    end
+    puts ''
+  end
+
+  # Response should be something like:
+  # - vm1.delivery.puppetlabs.net (platform-tag)\n
+  # - vm2.delivery.puppetlabs.net (platform-tag)\n
+  def self.parse_response(response)
+    response.scan(/- ([\w-]+)(?:.vmpooler-prod.puppet.net|.delivery.puppetlabs.net) \(/).flatten
   end
 
   ### Get a VM from floaty ###
-  def self.get_vm(platform: 'centos-7-x86_64', force_abs: false)
-    if !find_pooled_platforms.include?(platform.gsub('-pixa4','')) || force_abs
-      response = Floaty.floaty_cmd("get #{platform} --service abs --priority 1", use_pty: true)
-      vmlinesplit = response.chomp.split('- ')
+  def self.get_vm(platform: 'centos-7-x86_64', force_abs: false, number: 1)
+    pools = find_pooled_platforms
+    platform = platform.gsub('-pixa4','')
+    vms_per_thread = 10
+    if !pools.keys.include?(platform) || force_abs || pools[platform] < number
+      numthreads = number / vms_per_thread
+      numthreads += 1 if number % vms_per_thread != 0
+      threads = []
+      responses = []
+      mutex = Mutex.new
+      numleft = number
+      numthreads.times do
+        num = numleft < vms_per_thread ? numleft : vms_per_thread
+        numleft -= num
+        threads << Thread.new do
+          Thread.current['print'] = false
+          myresponse = Floaty.floaty_cmd("get #{platform}=#{num} --service abs --priority 1 --force", use_pty: true)
+          mutex.synchronize { responses << myresponse }
+        end
+      end
+      # Setting print to true then joining causes the thread not
+      # to pick up the variable change sometimes. So this polls
+      # the thread status and sets the print variable in a loop.
+      puts "Fetching VMs using #{threads.count} thread#{threads.count > 1 ? 's' : ''}. Output shown will be from one of the currently alive threads, as well as the result of any completed threads, until all threads are finished.".cyan unless threads.count == 1
+      while threads.any? { |t| t.alive? } do
+        # If we have no living threads with print set to true,
+        # pick one and turn it on.
+        if threads.select { |t| t.alive? && t['print'] == true }.empty?
+          t = threads.find{ |t| t.alive? }
+          t['print'] = true if !t.nil?
+        end
+      end
+      # Just in case
+      threads.each { |t| t.join }
+      vms = []
+      responses.each { |r| vms += parse_response(r) }
       # When use_pty is true, we've already printed stderr/stdout, so no need to do so again.
-      raise "Error obtaining a VM" if vmlinesplit.count == 1
-      vmlinesplit[1].split[0].split('.')[0]
+      raise "Error obtaining a VM" if vms.empty?
+      vms
     else
-      response = Floaty.floaty_cmd("get #{platform} --service vmpooler")
+      response = Floaty.floaty_cmd("get #{platform}=#{number} --service vmpooler")
       raise "Error obtaining a VM: #{response}" if response.include?('error')
-      response.split[1].split('.')[0]
+      parse_response(response)
     end
   end
 end

data/lib/floatyhelper.rb

@@ -10,6 +10,8 @@ require 'floatyhelper/groups'
 require 'floatyhelper/vm'
 require 'floatyhelper/version'
 require 'colorize'
+require 'open3'
+require 'thread'
 
 class Floatyhelper
   include Commander::Methods
@@ -371,15 +373,24 @@ class Floatyhelper
         ABS with the --abs flag.
       EOT
       c.option '--abs', 'Force use of ABS to get platform, even if platform is pooled.'
+      c.option '--number INTEGER', Integer, 'Number of VMs of this flavor to check out. Default 1.'
+      c.option '-y', "Don't ask if we're sure we want to request more than 10 VMs."
       c.action do |args, options|
         platform = args[0] || 'centos-7-x86_64'
         tag = args.length > 1 ? args[1] : 'Unassigned'
+        num = options.number || 1
+        if num > 10 && !options.y
+          answer = ask("Checking out more than 10 VMs at once can have an impact on the Vmpooler infrastructure. Please do not use this option unless you know what you're doing and clean them up when you're finished. Proceed? [Y/n]".yellow)
+          if !answer.empty? && answer.capitalize != 'Y'
+            return
+          end
+        end
         begin
-          host = VM.get_vm(platform: platform, force_abs: options.abs)
-          Groups.addhosts([host], tag)
-          puts "Added #{host} to tag #{tag}".green
+          hosts = VM.get_vm(platform: platform, force_abs: options.abs, number: num)
+          Groups.addhosts(hosts, tag)
+          puts "Added #{hosts.join(', ')} to tag #{tag}".green
           amount = Config.get_config_setting('increaselife').to_i
-          VM.increaselife(host, amount)
+          VM.increaselife(hosts, amount)
         rescue StandardError => e
           puts e.message.red
         end
@@ -395,6 +406,101 @@ class Floatyhelper
       end
     end
 
+    command :runon do |c|
+      c.syntax = 'floatyhelper runon <tag> <command>'
+      c.summary = 'Run a command on all hosts under the given tag'
+      c.description = 'Run a command on all hosts under the given tag using SSH. Requires that the key to SSH into the hosts is in ssh-agent already.'
+      c.option '--user', 'Run command as a user other than root'
+      c.action do |args, options|
+        if args.length == 0
+          puts 'Must specify a tag and command to run'.red
+          return
+        end
+        if args.length == 1
+          puts 'Must specify a command to run'.red
+          return
+        end
+
+        tag = args[0]
+        command = args[1..].join(' ')
+        user = options.user || 'root'
+        
+        if !Groups.tag?(tag)
+          puts "#{tag} does not appear to be a tag managed by floatyhelper.".red
+          return
+        end
+        hosts = Groups.get_hosts(tag)
+        threads = []
+        responses = []
+        mutex = Mutex.new
+        hosts.each do |host|
+          threads << Thread.new do
+            output, status = Open3.capture2e("ssh -A -o StrictHostKeyChecking=no #{user}@#{host}.delivery.puppetlabs.net \"#{command}\"")
+            mutex.synchronize do
+              responses << [host,output,status]
+            end
+          end
+        end
+        puts
+        (1..threads.count).each do |i|
+          print "\rWaiting for thread #{i}/#{threads.count} to finish..."
+          threads[i-1].join
+        end
+        print "\rDone\n"
+        responses.each do |r|
+          color = r[2].exitstatus.zero? ? :white : :red
+          puts "*** #{r[0]}: Exit code #{r[2].exitstatus} ***".colorize(color)
+          puts r[1].colorize(color)
+        end
+      end
+    end
+
+    command :scpto do |c|
+      c.syntax = 'floatyhelper scp <tag> <path_to_local_file> <path_to_remote_file>'
+      c.summary = 'SCP a file from this host to all hosts under the given tag at the given path'
+      c.description = 'SCP a file from this host to all hosts under the given tag at the given path. This assumes the necessary SSH key is already loaded into ssh-agent.'
+      c.option '--user', 'Log in as a user other than root'
+      c.action do |args, options|
+        if args.length < 3
+          puts 'Must specify a tag, local path, and remote path'.red
+          return
+        end
+
+        tag = args[0]
+        local = args[1]
+        remote = args[2]
+        user = options.user || 'root'
+        
+        if !Groups.tag?(tag)
+          puts "#{tag} does not appear to be a tag managed by floatyhelper.".red
+          return
+        end
+        hosts = Groups.get_hosts(tag)
+        threads = []
+        responses = []
+        mutex = Mutex.new
+        hosts.each do |host|
+          threads << Thread.new do
+            output, status = Open3.capture2e("scp #{local} #{user}@#{host}.delivery.puppetlabs.net:#{remote}")
+            mutex.synchronize do
+              responses << [host,output,status]
+            end
+          end
+        end
+        puts
+        (1..threads.count).each do |i|
+          print "\rWaiting for thread #{i}/#{threads.count} to finish..."
+          threads[i-1].join
+        end
+        print "\rDone\n"
+        responses.each do |r|
+          color = r[2].exitstatus.zero? ? :white : :red
+          puts "*** #{r[0]}: Exit code #{r[2].exitstatus} ***".colorize(color)
+          puts r[1].colorize(color)
+        end
+      end
+    end
+
     run!
   end
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: floatyhelper
 version: !ruby/object:Gem::Version
-  version: 2.0.6
+  version: 2.1.0
 platform: ruby
 authors:
 - Nick Burgan-Illig
 autorequire: 
 bindir: bin
 cert_chain: []
-date: 2021-09-21 00:00:00.000000000 Z
+date: 2023-08-15 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: bundler