mssh 0.0.1

data/LICENSE.md

@@ -0,0 +1,13 @@
+Copyright 2011 Square Inc.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

data/README.md

@@ -0,0 +1,25 @@
+mssh, mcmd
+==========
+
+
+Tools for running multiple commands and ssh jobs in parallel, and easily collecting the result
+
+Usage
+-----
+
+
+<code>mssh -r  host01,host02,host03 "uname -r" -c</code>
+
+BUGS/TODO
+---------
+
+
+ * Optionally Incorporate stderr into -c, with $?
+ * allow commandline manipulation of ssh args
+ * factor out redundancy between bin/mssh and bin/mcmd (cli module?)
+ * incorporate range / foundation lookup syntax for -r
+ * json output mode
+ * to-file output mode
+ * lots of rough spots, not super slick yet
+ * needs testing real bad. 0.1 release
+

data/bin/mcmd

@@ -0,0 +1,97 @@
+#!/usr/bin/env ruby
+
+require 'pp'
+
+$LOAD_PATH.unshift(File.expand_path(File.dirname(__FILE__) + "/../lib"))
+require 'mcmd'
+
+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

@@ -0,0 +1,86 @@
+#!/usr/bin/env ruby
+
+require 'rubygems'
+require 'pp'
+require 'mcmd'
+require 'rangeclient'
+
+require 'optparse'
+options = {
+  :maxflight => 50,
+  :timeout => 60,
+  :global_timeout => 600,
+}
+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('-c', '--collapse', "Collapse similar output ") do |arg|
+    options[:collapse] = arg 
+  end
+  opts.on('-v', '--verbose', "verbose ") 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?
+  # option to ignore as-we-go yield output
+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?
+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 }
+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 "#{command_to_target[p.command.object_id]}: finished" } if options[:verbose]
+
+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 = 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]
+  end
+  # output => [targets ...]
+  stdout_matches.each_pair do |k,v|
+    hosts = range.compress v
+    puts "#{hosts}: '#{k.chomp}'"
+  end
+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"
+  end
+end

data/lib/mcmd.rb

@@ -0,0 +1,281 @@
+#!/usr/bin/ruby
+
+require 'pp'
+
+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
+
+  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
+    
+    pid = fork
+    if not pid.nil?
+      # parent
+      # 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?
+    else
+      # child
+      # setup stdin, out, err
+      STDIN.reopen(stdin_rd)
+      STDOUT.reopen(stdout_wr)
+      STDERR.reopen(stderr_wr)
+      noshell_exec(cmd)
+      raise "can't be reached!!. exec failed!!"
+    end
+  end
+
+  def process_read_fds(read_fds)
+    read_fds.each do |fd|
+      # read available bytes, add to the subproc's read buf
+      if not @subproc_by_fd.has_key?(fd)
+        raise "Select returned a fd which I have not seen! fd: #{fd.inspect}"
+      end
+      subproc = @subproc_by_fd[fd]
+      buf = ""
+      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
+        # 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 run
+    @global_time_start = Time.now.to_i
+    done = false
+    while not done
+      # 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
+      # service running processes
+      self.service_subprocess_io
+      # timeout overdue processes
+      self.process_timeouts
+      # service process cleanup
+      self.wait
+      puts "have #{@subproc_by_pid.length} left to go" if self.debug
+      # 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
+      end
+    end
+    
+    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

@@ -0,0 +1,114 @@
+--- !ruby/object:Gem::Specification 
+name: mssh
+version: !ruby/object:Gem::Version 
+  hash: 29
+  prerelease: 
+  segments: 
+  - 0
+  - 0
+  - 1
+  version: 0.0.1
+platform: ruby
+authors: 
+- Evan Miller
+autorequire: 
+bindir: bin
+cert_chain: []
+
+date: 2012-05-16 00:00:00 Z
+dependencies: 
+- !ruby/object:Gem::Dependency 
+  name: json
+  prerelease: false
+  requirement: &id001 !ruby/object:Gem::Requirement 
+    none: false
+    requirements: 
+    - - ">="
+      - !ruby/object:Gem::Version 
+        hash: 3
+        segments: 
+        - 0
+        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
+executables: 
+- mssh
+- mcmd
+extensions: []
+
+extra_rdoc_files: 
+- LICENSE.md
+files: 
+- lib/mcmd.rb
+- bin/mcmd
+- bin/mssh
+- README.md
+- LICENSE.md
+homepage: http://github.com/square/prodeng
+licenses: []
+
+post_install_message: 
+rdoc_options: 
+- --charset=UTF-8
+require_paths: 
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement 
+  none: false
+  requirements: 
+  - - ">="
+    - !ruby/object:Gem::Version 
+      hash: 3
+      segments: 
+      - 0
+      version: "0"
+required_rubygems_version: !ruby/object:Gem::Requirement 
+  none: false
+  requirements: 
+  - - ">="
+    - !ruby/object:Gem::Version 
+      hash: 23
+      segments: 
+      - 1
+      - 3
+      - 6
+      version: 1.3.6
+requirements: []
+
+rubyforge_project: 
+rubygems_version: 1.8.24
+signing_key: 
+specification_version: 3
+summary: Parallel ssh and command execution.
+test_files: []
+