mssh 0.0.1 → 0.0.4

data/bin/mcmd-ev

@@ -0,0 +1,97 @@
+#!/usr/bin/ruby
+
+require 'pp'
+
+$LOAD_PATH.unshift(File.expand_path(File.dirname(__FILE__) + "/../lib"))
+require 'mcmd-ev'
+
+require 'optparse'
+options = {
+  :maxflight => 200,
+  :timeout => 60,
+  :global_timeout => 0,
+}
+
+optparse = OptionParser.new do |opts|
+  opts.on('-r', '--range RANGE', 'currently takes a CSV list') do |arg|
+    options[:range] = arg 
+  end
+  opts.on('-m', '--maxflight 50', 'How many subprocesses? 50 by default') do |arg|
+    options[:maxflight] = arg 
+  end
+  opts.on('-t', '--timeout 60', 'How many seconds may each individual process take? 0 for no timeout') do |arg|
+    options[:timeout] = arg 
+  end
+  opts.on('-g', '--global_timeout 600', 'How many seconds for the whole shebang 0 for no timeout') do |arg|
+    options[:global_timeout] = arg 
+  end
+  opts.on('--noshell', "Don't invoke a shell. Args will be passed to exec verbatim ") do |arg|
+    options[:noshell] = arg 
+  end
+  opts.on('-c', '--collapse', "Collapse similar output ") do |arg|
+    options[:collapse] = arg 
+  end
+  opts.on('-v', '--verbose', "Verbose output") do |arg|
+    options[:verbose] = arg 
+  end
+  opts.on('-d', '--debug', "Debug output") do |arg|
+    options[:debug] = arg 
+  end
+  # option to merge stdin/stdout into one buf? how should this work?
+  # option to ignore as-we-go yield output - this is off by default now except for success/fail
+end
+optparse.parse!
+
+raise "Error, need -r argument" if options[:range].nil? or options[:range].empty?
+raise "Error, need command to run" if ARGV.size.zero?
+
+m = MultipleCmd.new
+
+targets = options[:range].split ","
+
+m.commands = targets.map { |t| ["/bin/sh", "-c"].push ARGV.map { |arg| arg.gsub('HOSTNAME', t)}.join " " }
+command_to_target = Hash.new
+targets.size.times do |i|
+  command_to_target[m.commands[i].object_id] = targets[i]
+end
+m.yield_startcmd = lambda { |p| puts "#{command_to_target[p.command.object_id]}: starting" } if options[:verbose]
+m.yield_wait = lambda { |p| puts "#{p.success? ? 'SUCCESS' : 'FAILURE'} #{command_to_target[p.command.object_id]}: '#{p.stdout_buf}'" }
+# m.yield_proc_timeout = lambda { |p| puts "am killing #{p.inspect}"}
+
+m.perchild_timeout = options[:timeout].to_i
+m.global_timeout = options[:global_timeout].to_i
+m.maxflight = options[:maxflight].to_i
+m.verbose = options[:verbose]
+m.debug = options[:debug]
+
+result = m.run
+
+if options[:collapse]
+  # print a collapsed summary
+  stdout_matches_success = Hash.new
+  stdout_matches_failure = Hash.new
+  result.each do |r|
+    if r[:retval].success?
+      stdout_matches_success[r[:stdout_buf]] = [] if stdout_matches_success[r[:stdout_buf]].nil?
+      stdout_matches_success[r[:stdout_buf]] << command_to_target[r[:command].object_id]
+    else
+      stdout_matches_failure[r[:stdout_buf]] = [] if stdout_matches_failure[r[:stdout_buf]].nil?
+      stdout_matches_failure[r[:stdout_buf]] << command_to_target[r[:command].object_id]
+    end
+  end
+  # output => [targets ...]
+  stdout_matches_success.each_pair do |k,v|
+    puts "SUCCESS: #{v.join ','}: #{k}"
+  end
+  stdout_matches_failure.each_pair do |k,v|
+    puts "FAILURE: #{v.join ','}: #{k}"
+  end
+else
+  # we already printed while in-flight; do nothing
+    # not collapse, print one per host
+#  result.each do |r|
+#    target = command_to_target[r[:command].object_id]
+#    puts "#{target}: '#{r[:stdout_buf].chomp}'\n"
+#  end
+end
+

data/bin/mssh

@@ -3,7 +3,6 @@
 require 'rubygems'
 require 'pp'
 require 'mcmd'
-require 'rangeclient'
 
 require 'optparse'
 options = {
@@ -15,6 +14,10 @@ optparse = OptionParser.new do |opts|
   opts.on('-r', '--range RANGE', 'currently takes a CSV list') do |arg|
     options[:range] = arg 
   end
+  opts.on('-f', '--file FILE', 'List of hostnames in a FILE \
+          use (/dev/stdin) for reading from stdin') do |arg|
+    options[:file] = arg 
+  end
   opts.on('-m', '--maxflight 50', 'How many subprocesses? 50 by default') do |arg|
     options[:maxflight] = arg 
   end
@@ -38,18 +41,26 @@ optparse = OptionParser.new do |opts|
 end
 optparse.parse!
 
-raise "Error, need -r argument" if options[:range].nil? or options[:range].empty?
+targets = []
+if (!options[:range].nil?)
+    require 'rangeclient'
+    range = Range::Client.new
+    targets = range.expand options[:range]
+elsif (!options[:file].nil?)
+    targets_fd = File.open(options[:file])
+    targets_fd.read.each { |x| targets << x.chomp }
+else 
+    raise "Error, need -r or -f option" 
+end
+
+
 raise "Error, need command to run" if ARGV.size.zero?
 raise "Error, too many arguments" if ARGV.size != 1
 
 
-range = Range::Client.new
 
 m = MultipleCmd.new
-
-targets = range.expand options[:range]
-
-m.commands = targets.map { |t| ["/usr/bin/ssh", "-2", "-oPasswordAuthentication=no", "-A", t].push ARGV.first }
+m.commands = targets.map { |t| ["/usr/bin/ssh","-2", "-oBatchMode=yes", "-A", t].push ARGV.first }
 command_to_target = Hash.new
 targets.size.times do |i|
   command_to_target[m.commands[i].object_id] = targets[i]
@@ -65,15 +76,18 @@ m.debug = options[:debug]
 
 result = m.run
 
-if options[:collapse]
+if (options[:collapse] and options[:range])
   # print a collapsed summary
-  stdout_matches = Hash.new
+  out_matches = Hash.new
   result.each do |r|
-    stdout_matches[r[:stdout_buf]] = [] if stdout_matches[r[:stdout_buf]].nil?
-    stdout_matches[r[:stdout_buf]] << command_to_target[r[:command].object_id]
+    out = "" 
+    out += r[:stdout_buf].chomp if(!r[:stdout_buf].nil?)
+    out += r[:stderr_buf].chomp if(!r[:stderr_buf].nil?)
+    out_matches[out] = [] if out_matches[out].nil?
+    out_matches[out] << command_to_target[r[:command].object_id]
   end
   # output => [targets ...]
-  stdout_matches.each_pair do |k,v|
+  out_matches.each_pair do |k,v|
     hosts = range.compress v
     puts "#{hosts}: '#{k.chomp}'"
   end
@@ -81,6 +95,14 @@ else
   # not collapse, print one per host
   result.each do |r|
     target = command_to_target[r[:command].object_id]
-    puts "#{target}: #{r[:retval] == 0 ? 'SUCCESS:':'FAILURE:'} '#{r[:stdout_buf].chomp}'\n"
+    out = "" 
+    out += r[:stdout_buf].chomp if(!r[:stdout_buf].nil?)
+    out += r[:stderr_buf].chomp if(!r[:stderr_buf].nil?)
+    if (r[:retval] == 0 )
+        puts "#{target}:SUCCESS: #{out}\n"
+    else
+        exit_code = r[:retval].exitstatus.to_s if(!r[:retval].nil?)
+        puts "#{target}:FAILURE[#{exit_code}]: #{out}\n"
+    end
   end
 end

data/lib/mcmd-ev.rb

@@ -0,0 +1,390 @@
+#!/usr/bin/ruby
+
+require 'pp'
+require 'rubygems'
+require 'eventmachine'
+
+
+# TODO - figure out why IO.sysread isn't raising on EOF
+module FdWatcher
+  def notify_readable
+    puts "readable: '#{@io.inspect}'"
+    puts "mcmd obj: '#{@mcmd}'"
+    # read available bytes, add to the subproc's read buf
+    if not @mcmd.subproc_by_fd.has_key?(@io)
+      raise "Select returned a fd which I have not seen! fd: #{@io.inspect}"
+    end
+    subproc = @mcmd.subproc_by_fd[@io]
+    buf = ""
+    begin
+      buf = @io.sysread(4096)
+      puts "saw #{buf}"
+      if buf.nil?
+        raise " Impossible result from sysread()"
+      end
+      # no exception? bytes were read. append them.                                                                                   
+      if @io == subproc.stdout_fd
+        subproc.stdout_buf << buf
+        # FIXME if we've read > maxbuf, allow closing/ignoring the fd instead of hard kill                                            
+        if subproc.stdout_buf.bytesize > @mcmd.max_read_size
+          # self.kill_process(subproc) # can't kill this here, need a way to mark-to-kill                                             
+        end
+      elsif @io == subproc.stderr_fd
+        subproc.stderr_buf << buf
+        # FIXME if we've read > maxbuf, allow closing/ignoring the fd instead of hard kill                                            
+        if subproc.stderr_buf.bytesize > @mcmd.max_read_size
+          # self.kill_process(subproc) # "" above                                                                                     
+        end
+      end
+    rescue SystemCallError, EOFError => ex
+      puts "DEBUG: saw read exception #{ex}"
+      @mcmd.periodic
+    rescue => ex
+      puts "else #{ex.inspect}"
+    end
+  end
+
+  def unbind
+    puts "GOT UNBIND"
+    @mcmd.periodic
+    # move periodic() bits under here, to be called as necessary
+  end
+  def notify_writable
+#   puts "writable: '#{@io.inspect}'"
+  end
+end
+
+class FdClass < EventMachine::Connection
+  include FdWatcher
+end
+
+module ProcessWatcher
+  def process_exited
+    puts 'the forked child died!'
+    EM.stop_event_loop
+    puts "stopped"
+  end
+end
+
+
+class MultipleCmd
+
+  attr_accessor :global_timeout, :maxflight, :perchild_timeout, :commands
+  attr_accessor :yield_wait, :yield_startcmd, :debug, :yield_proc_timeout
+  attr_accessor :verbose, :poll_period, :max_read_size
+  attr_accessor :subproc_by_fd
+
+  def initialize
+    # these are re-initialized after every run
+    @subproc_by_pid = Hash.new
+    @subproc_by_fd = Hash.new
+    @processed_commands = []
+    # end items which are re-initialized
+
+    self.commands = []
+    self.perchild_timeout = 60
+    self.global_timeout = 0
+    self.maxflight = 200
+    self.debug = false
+    self.poll_period = 0.5 # shouldn't need adjusting
+    self.max_read_size = 2 ** 19 # 512k
+  end
+
+  def noshell_exec(cmd)
+    if cmd.length == 1
+      Kernel.exec([cmd[0], cmd[0]])
+    else 
+      Kernel.exec([cmd[0], cmd[0]], *cmd[1..-1])
+    end
+  end
+  
+  # I should probably move this whole method
+  # into SubProc and make the subproc_by_* into
+  # class variables
+  def add_subprocess(cmd)
+    stdin_rd, stdin_wr = IO.pipe
+    stdout_rd, stdout_wr = IO.pipe
+    stderr_rd, stderr_wr = IO.pipe
+    subproc = MultipleCmd::SubProc.new
+    subproc.stdin_fd = stdin_wr
+    subproc.stdout_fd = stdout_rd
+    subproc.stderr_fd = stderr_rd
+    subproc.command = cmd
+    
+    puts "Before fork"
+    pid = fork
+    if not pid.nil?
+      # parent
+      stdin_rd.close rescue true
+      stdout_wr.close rescue true
+      stderr_wr.close rescue true
+      # for mapping to subproc by pid
+      subproc.pid = pid
+      @subproc_by_pid[pid] = subproc
+      # for mapping to subproc by i/o handle (returned from select)
+      @subproc_by_fd[stdin_rd] = subproc
+      @subproc_by_fd[stdin_wr] = subproc
+      @subproc_by_fd[stdout_rd] = subproc
+      @subproc_by_fd[stdout_wr] = subproc
+      @subproc_by_fd[stderr_rd] = subproc
+      @subproc_by_fd[stderr_wr] = subproc
+      
+      self.yield_startcmd.call(subproc) unless self.yield_startcmd.nil?
+      temp_self = self
+      # now add fds to EM
+      EM.watch(stdin_wr, FdClass) { |c|
+        c.notify_writable = true
+        c.instance_eval do
+          @mcmd = temp_self
+        end
+      }
+      EM.watch(stdout_rd, FdClass) { |c|
+        c.notify_readable = true
+        c.instance_eval do
+          @mcmd = temp_self
+        end
+      }
+      EM.watch(stderr_rd, FdClass) { |c|
+        c.notify_readable = true
+        c.instance_eval do
+          @mcmd = temp_self
+        end
+      }
+      # EM.watch_process(pid, ProcessWatcher) # is this kqueue only? If so, probably don't need it FIXME
+    else
+      # child
+      # setup stdin, out, err
+      stdin_wr.close rescue true
+      stdout_rd.close rescue true
+      stderr_rd.close rescue true
+      STDIN.reopen(stdin_rd)
+      STDOUT.reopen(stdout_wr)
+      STDERR.reopen(stderr_wr)
+      puts "EXEC: #{cmd}"
+      noshell_exec(cmd)
+      raise "can't be reached!!. exec failed!!"
+    end
+  end
+
+  def process_read_fds(read_fds)
+    read_fds.each do |fd|
+      begin
+        buf = fd.sysread(4096)
+        
+        if buf.nil?
+          raise " Impossible result from sysread()"
+        end
+        # no exception? bytes were read. append them.
+        if fd == subproc.stdout_fd
+          subproc.stdout_buf << buf
+          # FIXME if we've read > maxbuf, allow closing/ignoring the fd instead of hard kill
+          if subproc.stdout_buf.bytesize > self.max_read_size
+            # self.kill_process(subproc) # can't kill this here, need a way to mark-to-kill
+          end
+        elsif fd == subproc.stderr_fd
+          subproc.stderr_buf << buf
+          # FIXME if we've read > maxbuf, allow closing/ignoring the fd instead of hard kill
+          if subproc.stderr_buf.bytesize > self.max_read_size
+            # self.kill_process(subproc) # "" above
+          end
+        end
+      rescue SystemCallError, EOFError => ex
+        puts "DEBUG: saw read exception #{ex}" if self.debug
+        detach
+        # clear out the read fd for this subproc
+        # finalize read i/o
+        # if we're reading, it was the process's stdout or stderr
+        if fd == subproc.stdout_fd
+          subproc.stdout_fd = nil 
+        elsif fd == subproc.stderr_fd
+          subproc.stderr_fd = nil
+        else
+          raise "impossible: operating on a subproc where the fd isn't found, even though it's mapped"
+        end
+        fd.close rescue true
+      end
+    end
+  end # process_read_fds()
+  def process_write_fds(write_fds)
+    write_fds.each do |fd|
+      raise "working on an unknown fd #{fd}" unless @subproc_by_fd.has_key?(fd)
+      subproc = @subproc_by_fd[fd]
+      buf = ""
+      # add writing here, todo. not core feature
+    end
+  end
+  def process_err_fds(err_fds)
+  end
+  
+  # iterate and service fds in child procs, collect data and status
+  def service_subprocess_io
+    write_fds = @subproc_by_pid.values.select {|x| not x.stdin_fd.nil? and not x.terminated}.map {|x| x.stdin_fd}
+    read_fds = @subproc_by_pid.values.select {|x| not x.terminated}.map {|x| [x.stdout_fd, x.stderr_fd].select {|x| not x.nil? } }.flatten
+
+    read_fds, write_fds, err_fds = IO.select(read_fds, write_fds, nil, self.poll_period)
+
+    self.process_read_fds(read_fds) unless read_fds.nil?
+    self.process_write_fds(write_fds) unless write_fds.nil?
+    self.process_err_fds(err_fds) unless err_fds.nil?
+    # errors? 
+  end
+
+  def process_timeouts
+    now = Time.now.to_i
+    @subproc_by_pid.values.each do |p|
+      if (now - p.time_start) > self.perchild_timeout
+        # expire this child process
+        
+        self.yield_proc_timeout.call(p) unless self.yield_proc_timeout.nil?
+        self.kill_process(p)
+      end
+    end
+  end
+
+  def kill_process(p)
+    # do not remove from pid list until waited on
+    @subproc_by_fd.delete(p.stdin_fd)
+    @subproc_by_fd.delete(p.stdout_fd)
+    @subproc_by_fd.delete(p.stderr_fd)
+    # must kill after deleting from maps
+    # kill closes fds
+    p.kill
+  end
+
+  def periodic
+    puts "in periodic"
+    self.wait
+    # This could probably move into the child-cleanup callback
+    # start up as many as maxflight processes
+    while @subproc_by_pid.length < self.maxflight and not @commands.empty?
+      # take one from @commands and start it
+      commands = @commands.shift
+      self.add_subprocess(commands)
+    end
+    
+    puts "have #{@subproc_by_pid.length} left to go"
+    # if we have nothing in flight (active pid)
+    # and nothing pending on the input list
+    # then we're done
+    if @subproc_by_pid.length.zero? and @commands.empty?
+      done = true
+      puts "stopping ev loop"
+      EM.stop_event_loop
+    end
+  end
+
+  def run
+    @global_time_start = Time.now.to_i
+    done = false
+    EM.run do
+      EM.schedule do
+        puts "schedule!"
+        self.periodic
+      end
+#      EM.add_periodic_timer(0.5) do
+#        puts "TIMER!!"
+#        self.periodic
+#      end # end periodic timer
+    end
+
+    ## These things get broken up into individual proc/fd callbacks
+      # service running processes
+#      self.service_subprocess_io
+      # timeout overdue processes
+#      self.process_timeouts
+      # service process cleanup
+    
+    ## this is still the same, post-event-loop
+    data = self.return_rundata
+    # these are re-initialized after every run
+    @subproc_by_pid = Hash.new
+    @subproc_by_fd = Hash.new
+    @processed_commands = []
+    # end items which are re-initialized
+    return data
+  end
+  
+  def return_rundata
+    data = []
+    @processed_commands.each do |c|
+      #FIXME pass through the process object
+      data << {
+        :pid => c.pid,
+        :write_buf_position => c.write_buf_position,
+        :stdout_buf => c.stdout_buf,
+        :stderr_buf => c.stderr_buf,
+        :command => c.command,
+        :time_start => c.time_start,
+        :time_end => c.time_end,
+        :retval => c.retval,
+      }
+    end
+    return data
+  end
+  
+  def wait
+    possible_children = true
+    just_reaped = Array.new
+    while possible_children
+      begin
+        pid = Process::waitpid(-1, Process::WNOHANG)
+        if pid.nil?
+          possible_children = false
+        else
+          # pid is now gone. remove from subproc_by_pid and
+          # add to the processed commands list
+          p = @subproc_by_pid[pid]
+          p.time_end = Time.now.to_i
+          p.retval = $?
+          @subproc_by_pid.delete(pid)
+          @processed_commands << p
+          just_reaped << p
+        end
+      rescue Errno::ECHILD => ex
+        # ECHILD. ignore.
+        possible_children = false
+      end
+    end
+    # We may have waited on a child before reading all its output. Collect those missing bits. No blocking.
+    if not just_reaped.empty?
+      read_fds = just_reaped.select {|x| not x.terminated}.map {|x| [x.stdout_fd, x.stderr_fd].select {|x| not x.nil? } }.flatten
+      read_fds, write_fds, err_fds = IO.select(read_fds, nil, nil, 0)
+      self.process_read_fds(read_fds) unless read_fds.nil?
+    end
+    just_reaped.each do |p|
+      self.yield_wait.call(p) unless self.yield_wait.nil?
+    end
+  end
+
+end
+
+class MultipleCmd::SubProc
+  attr_accessor :stdin_fd, :stdout_fd, :stderr_fd, :write_buf_position
+  attr_accessor :time_start, :time_end, :pid, :retval, :stdout_buf, :stderr_buf, :command, :terminated
+
+  def initialize
+    self.write_buf_position = 0
+    self.time_start = Time.now.to_i
+    self.stdout_buf = ""
+    self.stderr_buf = ""
+    self.terminated = false
+  end
+
+  # when a process has out-stayed its welcome
+  def kill
+    self.stdin_fd.close rescue true
+    self.stdout_fd.close rescue true
+    self.stderr_fd.close rescue true
+    #TODO configurable sig?
+    Process::kill("KILL", self.pid)
+    self.terminated = true
+  end
+
+
+  # some heuristic to determine if this job was successful
+  # for now, trust retval. Also check stderr?
+  def success?
+    self.retval.success?
+  end
+end
+

metadata

@@ -1,13 +1,13 @@
 --- !ruby/object:Gem::Specification 
 name: mssh
 version: !ruby/object:Gem::Version 
-  hash: 29
+  hash: 23
   prerelease: 
   segments: 
   - 0
   - 0
-  - 1
-  version: 0.0.1
+  - 4
+  version: 0.0.4
 platform: ruby
 authors: 
 - Evan Miller
@@ -15,7 +15,7 @@ autorequire:
 bindir: bin
 cert_chain: []
 
-date: 2012-05-16 00:00:00 Z
+date: 2012-09-04 00:00:00 Z
 dependencies: 
 - !ruby/object:Gem::Dependency 
   name: json
@@ -31,34 +31,6 @@ dependencies:
         version: "0"
   type: :runtime
   version_requirements: *id001
-- !ruby/object:Gem::Dependency 
-  name: rangeclient
-  prerelease: false
-  requirement: &id002 !ruby/object:Gem::Requirement 
-    none: false
-    requirements: 
-    - - ">="
-      - !ruby/object:Gem::Version 
-        hash: 3
-        segments: 
-        - 0
-        version: "0"
-  type: :runtime
-  version_requirements: *id002
-- !ruby/object:Gem::Dependency 
-  name: rdoc
-  prerelease: false
-  requirement: &id003 !ruby/object:Gem::Requirement 
-    none: false
-    requirements: 
-    - - ">="
-      - !ruby/object:Gem::Version 
-        hash: 3
-        segments: 
-        - 0
-        version: "0"
-  type: :runtime
-  version_requirements: *id003
 description: Simple library for running jobs and sshing to many hosts at once.
 email: 
 - github@squareup.com
@@ -70,8 +42,10 @@ extensions: []
 extra_rdoc_files: 
 - LICENSE.md
 files: 
+- lib/mcmd-ev.rb
 - lib/mcmd.rb
 - bin/mcmd
+- bin/mcmd-ev
 - bin/mssh
 - README.md
 - LICENSE.md