rubysh 0.0.1

data/lib/rubysh/runner.rb

@@ -0,0 +1,156 @@
+module Rubysh
+  class Runner
+    attr_accessor :command, :targets
+
+    def initialize(command)
+      @runner_state = :initialized
+
+      @command = command
+      @targets = {}
+      @state = {}
+
+      @parallel_io = nil
+
+      prepare!
+    end
+
+    def data(target_name)
+      state = target_state(target_name)
+      raise Rubysh::Error::BaseError.new("Can only access data for readable FDs") unless state[:target_reading?]
+      state[:buffer].join
+    end
+
+    # Ruby's Process::Status. Has fun things like pid and signaled?
+    def full_status(command=nil)
+      command ||= @command
+      @command.status(self)
+    end
+
+    def pid(command=nil)
+      command ||= @command
+      @command.pid(self)
+    end
+
+    # Convenience wrapper
+    def exitstatus(command=nil)
+      if st = full_status(command)
+        st.exitstatus
+      else
+        nil
+      end
+    end
+
+    # API for running/waiting
+    def run_async
+      raise Rubysh::Error::AlreadyRunError.new("You have already run this #{self.class} instance. Cannot run again. You can run its command directly though, which will create a fresh #{self.class} instance.") unless @runner_state == :initialized
+      prepare_io
+      @command.start_async(self)
+      @runner_state = :started
+      self
+    end
+
+    def wait
+      run_io
+      do_wait
+    end
+
+    def run(input={})
+      run_async
+      run_io
+      do_wait
+    end
+
+    def readers
+      readers = {}
+      @targets.each do |target_name, target_state|
+        next unless target_state[:target_reading?]
+        target = target_state[:target]
+        readers[target] = target_name
+      end
+      readers
+    end
+
+    def writers
+      writers = {}
+      @targets.each do |target_name, target_state|
+        next if target_state[:target_reading?]
+        target = target_state[:target]
+        writers[target] = target_name
+      end
+      writers
+    end
+
+    def to_s
+      inspect
+    end
+
+    def inspect
+      extras = []
+      valid_readers = readers.values.map(&:inspect).join(', ')
+      valid_writers = readers.values.map(&:inspect).join(', ')
+
+      extras << "readers: #{valid_readers}" if valid_readers.length > 0
+      extras << "writers: #{valid_writers}" if valid_writers.length > 0
+      if status = exitstatus
+        extras << "exitstatus: #{status}"
+      elsif mypid = pid
+        extras << "pid: #{pid}"
+      end
+      extra_display = extras.length > 0 ? " (#{extras.join(', ')})" : nil
+
+      "#{self.class}: #{command.stringify}#{extra_display}"
+    end
+
+    # Internal helpers
+    def state(object)
+      @state[object] ||= {}
+    end
+
+    # Internal helpers
+    def target_state(target_name)
+      @targets[target_name] || raise(Rubysh::Error::BaseError.new("Invalid target: #{target_name.inspect} (valid targets are: #{@targets.keys.inspect})"))
+    end
+
+    private
+
+    def do_wait
+      raise Rubysh::Error::AlreadyRunError.new("You must run parallel io before waiting. (Perhaps you want to use the 'run' method, which takes care of the plumbing for you?)") unless @runner_state == :parallel_io_ran
+      @command.wait(self)
+      @runner_state = :waited
+      self
+    end
+
+    def run_io
+      raise Rubysh::Error::AlreadyRunError.new("You must start the subprocesses before running parallel io. (Perhaps you want to use the 'run' method, which takes care of the plumbing for you?)") unless @runner_state == :started
+      @parallel_io.run
+      @runner_state = :parallel_io_ran
+      self
+    end
+
+    def prepare!
+      @command.prepare!(self)
+    end
+
+    # Can't build this in the prepare stage because pipes aren't built
+    # there.
+    def prepare_io
+      @parallel_io = Subprocess::ParallelIO.new(readers, writers)
+      @parallel_io.on_read do |target_name, data|
+        if data == Subprocess::ParallelIO::EOF
+          Rubysh.log.debug("EOF reached on #{target_name.inspect}")
+        else
+          Rubysh.log.debug("Just read #{data.inspect} on #{target_name.inspect}")
+          @targets[target_name][:buffer] << data
+        end
+      end
+
+      @parallel_io.on_write do |target_name, written, remaining|
+        if data == Subprocess::ParallelIO::EOF
+          Rubysh.log.debug("EOF reached on #{target_name.inspect}")
+        else
+          Rubysh.log.debug("Just wrote #{written.inspect} on #{target_name.inspect}")
+        end
+      end
+    end
+  end
+end

data/lib/rubysh/subprocess/parallel_io.rb

@@ -0,0 +1,184 @@
+class Rubysh::Subprocess
+  class ParallelIO
+    module EOF; end
+
+    # readers/writers should be hashes mapping {fd => name}
+    def initialize(readers, writers)
+      @finished_readers = Set.new
+      @on_read = nil
+      @readers = readers
+
+      @writers = writers
+      @finished_writers = Set.new
+      @on_write = nil
+      @writer_buffers = {}
+    end
+
+    def on_read(&blk)
+      @on_read = blk
+    end
+
+    def on_write(&blk)
+      @on_write = blk
+    end
+
+    def write(writer_name, data, close_on_complete=true)
+      writer = writer_by_name(writer_name)
+      buffer_state = @writer_buffers[writer] ||= {
+        :data => '',
+        :close_on_complete => nil
+      }
+
+      if buffer_state[:close_on_complete]
+        raise Rubysh::Error::AlreadyClosedError.new("You have already marked #{writer.inspect} as close_on_complete; can't write more data")
+      end
+
+      buffer_state[:close_on_complete] = close_on_complete
+      # XXX: unnecessary copy here
+      buffer_state[:data] += data
+
+      # Note that this leads to a bit of weird semantics if you try
+      # doing a write('') from within an on_write handler, since it'll
+      # call this synchronously. May want to change at some point.
+      finalize_writer_if_done(writer)
+    end
+
+    def close(writer_name)
+      writer = writer_by_name(writer_name)
+      writer.close
+    end
+
+    def available_readers
+      potential = @readers.keys - @finished_readers.to_a
+      potential.select {|reader| !reader.closed?}
+    end
+
+    # Writers with a non-zero number of bytes remaining to write
+    def available_writers
+      potential = @writer_buffers.keys - @finished_writers.to_a
+      potential.select {|writer| !writer.closed? && get_data(writer).length > 0}
+    end
+
+    def run
+      while available_writers.length > 0 || available_readers.length > 0
+        run_once
+      end
+    end
+
+    private
+
+    def run_once
+      potential_readers = available_readers
+      potential_writers = available_writers
+
+      begin
+        ready_readers, ready_writers, _ = IO.select(potential_readers, potential_writers)
+      rescue Errno::EINTR
+        next
+      end
+
+      ready_readers.each do |reader|
+        read_available(reader)
+      end
+
+      ready_writers.each do |writer|
+        write_available(writer)
+      end
+    end
+
+    def read_available(reader)
+      begin
+        data = reader.read_nonblock(4096)
+      rescue EOFError, Errno::EPIPE
+        finalize_reader(reader)
+      rescue Errno::EAGAIN, Errno::EWOULDBLOCK, Errno::EINTR
+      else
+        issue_reader_callback(reader, data)
+      end
+    end
+
+    def finalize_reader(reader)
+      @finished_readers.add(reader)
+      issue_reader_callback(reader, EOF)
+      reader.close
+    end
+
+    def issue_reader_callback(reader, data)
+      if @on_read
+        name = reader_name(reader)
+        @on_read.call(name, data)
+      end
+    end
+
+    def reader_name(reader)
+      @readers.fetch(reader)
+    end
+
+    def write_available(writer)
+      data = get_data(writer)
+      begin
+        count = writer.write_nonblock(data)
+      rescue EOFError, Errno::EPIPE
+        finalize_writer(writer)
+      rescue Errno::EAGAIN, Errno::EWOULDBLOCK, Errno::EINTR
+      else
+        # XXX: This may be a lot of copying. May want to think about
+        # how this scales.
+        written = data[0...count]
+        remaining = data[count..-1]
+        set_data(writer, remaining)
+        issue_writer_callback(writer, written, remaining)
+      end
+      finalize_writer_if_done(writer)
+    end
+
+    # Will only schedule a writer if it has a nonzero number of bytes
+    # left to write, so we need to manually check if we're out after
+    # every run.
+    def finalize_writer_if_done(writer)
+      if !writer.closed? &&
+          buffer_state(writer)[:close_on_complete] &&
+          get_data(writer).length == 0
+        finalize_writer(writer)
+      end
+    end
+
+    def finalize_writer(writer)
+      # TODO: think about how we should deal with errors, maybe
+      remaining = get_data(writer)
+      @finished_writers.add(writer)
+      issue_writer_callback(writer, EOF, remaining)
+      writer.close if buffer_state(writer)[:close_on_complete]
+    end
+
+    def get_data(writer)
+      buffer_state(writer)[:data]
+    end
+
+    def set_data(writer, data)
+      buffer_state(writer)[:data] = data
+    end
+
+    def buffer_state(writer)
+      buffer_state = @writer_buffers[writer]
+      Rubysh.assert(buffer_state, "No buffer state: #{writer.inspect}", true)
+      buffer_state
+    end
+
+    def issue_writer_callback(writer, data, remaining)
+      if @on_write
+        name = writer_name(writer)
+        @on_write.call(name, data, remaining)
+      end
+    end
+
+    def writer_name(writer)
+      @writers.fetch(writer)
+    end
+
+    # Could make this fast, but don't think it matters enough.
+    def writer_by_name(writer_name)
+      @writers.detect {|writer, name| writer_name == name}.first
+    end
+  end
+end

data/lib/rubysh/subprocess/pipe_wrapper.rb

@@ -0,0 +1,61 @@
+class Rubysh::Subprocess
+  class PipeWrapper
+    attr_accessor :reader, :writer
+
+    def initialize(reader_cloexec=true, writer_cloexec=true)
+      @reader, @writer = IO.pipe
+      set_reader_cloexec if reader_cloexec
+      set_writer_cloexec if writer_cloexec
+    end
+
+    def read_only
+      @writer.close
+    end
+
+    def write_only
+      @reader.close
+    end
+
+    def close
+      @writer.close
+      @reader.close
+    end
+
+    def set_reader_cloexec
+      @reader.fcntl(Fcntl::F_SETFD, Fcntl::FD_CLOEXEC)
+    end
+
+    def set_writer_cloexec
+      @writer.fcntl(Fcntl::F_SETFD, Fcntl::FD_CLOEXEC)
+    end
+
+    def nonblock
+      [@reader, @writer].each do |fd|
+        fl = fd.fcntl(Fcntl::F_GETFL)
+        fd.fcntl(Fcntl::F_SETFL, fl | Fcntl::O_NONBLOCK)
+      end
+    end
+
+    def dump_yaml_and_close(msg)
+      begin
+        YAML.dump(msg, @writer)
+      ensure
+        @writer.close
+        Rubysh.assert(@reader.closed?, "Reader should already be closed")
+      end
+    end
+
+    def load_yaml_and_close
+      begin
+        YAML.load(@reader)
+      rescue ArgumentError => e
+        # e.g. ArgumentError: syntax error on line 0, col 2: `' (could
+        # happen if the subprocess was killed while writing a message)
+        raise Rubysh::Error::BaseError.new("Invalid message read from pipe: #{e}")
+      ensure
+        @reader.close
+        Rubysh.assert(@writer.closed?, "Writer should already be closed")
+      end
+    end
+  end
+end

data/lib/rubysh/subprocess.rb

@@ -0,0 +1,154 @@
+# Adapted from https://github.com/ahoward/open4
+require 'fcntl'
+require 'timeout'
+require 'thread'
+
+# Using YAML to avoid the JSON dep. open4 uses Marshal to pass around
+# the exception object, but I'm always a bit sketched by Marshal when
+# it's not needed (i.e. don't want the subprocess to have the ability
+# to execute code in the parent, even if it should lose that ability
+# post-exec.)
+require 'yaml'
+
+require 'rubysh/subprocess/parallel_io'
+require 'rubysh/subprocess/pipe_wrapper'
+
+module Rubysh
+  class Subprocess
+    attr_accessor :command, :args, :directives, :runner
+    attr_accessor :pid, :status, :exec_error
+
+    # TODO: switch directives over to an OrderedHash of some form? Really
+    # want to preserve the semantics here.
+    def initialize(args, directives=[], post_fork=[], runner=nil)
+      raise ArgumentError.new("Must provide an array (#{args.inspect} provided)") unless args.kind_of?(Array)
+      raise ArgumentError.new("No command specified (#{args.inspect} provided)") unless args.length > 0
+      @command = args[0]
+      @args = args[1..-1]
+      @directives = directives
+      @runner = runner
+
+      Rubysh.assert(@directives.length == 0 || @runner, "Directives provided but no runner is", true)
+
+      @exec_status = nil
+      @post_fork = post_fork
+
+      @pid = nil
+      @status = nil
+      @exec_error = nil
+
+      Rubysh.log.debug("Just created: #{self}")
+    end
+
+    def to_s
+      "Subprocess: command=#{@command.inspect} args=#{@args.inspect} directives: #{@directives.inspect}"
+    end
+
+    def run
+      do_run unless @pid
+      @pid
+    end
+
+    def wait(nonblock=false)
+      do_wait(nonblock) unless @status
+      @status
+    end
+
+    private
+
+    def do_run
+      # Create this here so as to not leave an open pipe hanging
+      # around for too long. Not sure what would happen if a child
+      # inherited it.
+      open_exec_status
+      @pid = fork do
+        do_run_child
+      end
+      do_run_parent
+    end
+
+    def open_exec_status
+      @exec_status = PipeWrapper.new
+    end
+
+    def do_run_parent
+      # nil in tests
+      @exec_status.read_only
+      apply_directives_parent
+      handle_exec_error
+    end
+
+    def do_wait(nonblock=false)
+      flags = nonblock ? Process::WNOHANG : nil
+      return nil unless result = Process.waitpid2(@pid, flags)
+
+      pid, @status = result
+      Rubysh.assert(pid == @pid,
+        "Process.waitpid2 returned #{pid} while waiting for #{@pid}",
+        true)
+    end
+
+    def do_run_child
+      # nil in tests
+      @exec_status.write_only
+      run_post_fork
+      apply_directives_child
+      exec_program
+    end
+
+    def run_post_fork
+      @post_fork.each {|blk| blk.call}
+    end
+
+    def apply_directives_parent
+      apply_directives(true)
+    end
+
+    def apply_directives_child
+      apply_directives(false)
+    end
+
+    def apply_directives(is_parent)
+      @directives.each {|directive| apply_directive(directive, is_parent)}
+    end
+
+    def apply_directive(directive, is_parent)
+      if is_parent
+        directive.apply_parent!(runner)
+      else
+        directive.apply!(runner)
+      end
+    end
+
+    def exec_program
+      begin
+        Kernel.exec([command, command], *args)
+        raise Rubysh::Error::UnreachableError.new("This code should be unreachable. If you are seeing this exception, it means someone overrode Kernel.exec. That's not very nice of them.")
+      rescue Exception => e
+        msg = {
+          'message' => e.message,
+          'klass' => e.class.to_s,
+          # TODO: this may need coercion in Ruby1.9
+          'caller' => e.send(:caller)
+        }
+        @exec_status.dump_yaml_and_close(msg)
+        # Note: atexit handlers will fire in this case. May want to do
+        # something about that.
+        raise
+      end
+    end
+
+    def handle_exec_error
+      msg = @exec_status.load_yaml_and_close
+
+      case msg
+      when false
+        # success!
+      when Hash
+        @exec_error = Rubysh::Error::ExecError.new("Failed to exec in subprocess: #{msg['message']}", msg['klass'], msg['caller'])
+      else
+        @exec_error = Rubysh::Error::BaseError.new("Invalid message received over the exec_status pipe: #{msg.inspect}")
+      end
+    end
+  end
+end

data/lib/rubysh/triple_less_than.rb

@@ -0,0 +1,65 @@
+require 'tempfile'
+
+module Rubysh
+  # Looks like bash always buffers <<< to disk
+  class TripleLessThan < BaseDirective
+    class Shell < BaseDirective
+      def initialize(fd)
+        @fd = fd
+      end
+
+      def <(literal)
+        TripleLessThan.new(@fd, literal)
+      end
+
+      def prepare!
+        raise Rubysh::Error::BaseError.new("You have an incorrect usage of <<<, leading to a #{self.class} instance hanging around. Use it as either: Rubysh.<<< 'my string' or Rubysh::FD(3).<<< 'my string'.")
+      end
+
+      def stringify
+        " << #{fd.stringify} (INVALID SYNTAX)"
+      end
+    end
+
+    # TODO: support in-place strings
+    def initialize(fd, literal)
+      @fd = fd
+      @literal = literal
+    end
+
+    def prepare!(runner)
+      tempfile = Tempfile.new('buffer')
+      tempfile.delete
+      tempfile.write(@literal)
+      tempfile.flush
+      tempfile.rewind
+
+      Util.set_cloexec(tempfile)
+
+      state = state(runner)
+      state[:tempfile] = tempfile
+      state[:redirect] = Redirect.new(@fd, '<', tempfile)
+    end
+
+    def stringify
+      fd = Util.to_fileno(@fd)
+      beginning = fd == 0 ? '' : fd.to_s
+      "#{beginning}<<< #{@literal.inspect}"
+    end
+
+    def to_s
+      "TripleLessThan: #{stringify}"
+    end
+
+    def apply_parent!(runner)
+      state = state(runner)
+      state[:tempfile].close
+      state[:redirect].apply_parent!(runner)
+    end
+
+    def apply!(runner)
+      state = state(runner)
+      state[:redirect].apply!(runner)
+    end
+  end
+end

data/lib/rubysh/util.rb

@@ -0,0 +1,55 @@
+require 'set'
+
+module Rubysh
+  module Util
+    def self.to_fileno(file)
+      if file.respond_to?(:fileno)
+        file.fileno
+      else
+        file
+      end
+    end
+
+    # Leaks memory (needed to avoid Ruby 1.8's IO autoclose behavior),
+    # and so you should only use it right before execing.
+    def self.io_without_autoclose(fd_num)
+      fd_num = to_fileno(fd_num)
+      io = IO.new(fd_num)
+      hold(io)
+      io
+    end
+
+    # Should really just shell out to dup2, but looks like we'd need a
+    # C extension to do so. The concurrency story here is a bit off,
+    # and this probably doesn't copy over all FD state
+    # properly. Should be fine for now.
+    def self.dup2(fildes, fildes2)
+      original = io_without_autoclose(fildes)
+
+      begin
+        copy = io_without_autoclose(fildes2)
+      rescue Errno::EBADF
+      else
+        copy.close
+      end
+
+      res = original.fcntl(Fcntl::F_DUPFD, fildes2)
+      Rubysh.assert(res == fildes2, "Tried to open #{fildes2} but ended up with #{res} instead", true)
+    end
+
+    def self.set_cloexec(file, enable=true)
+      file = io_without_autoclose(file) unless file.kind_of?(IO)
+      value = enable ? Fcntl::FD_CLOEXEC : 0
+      file.fcntl(Fcntl::F_SETFD, value)
+    end
+
+    private
+
+    @references = []
+    def self.hold(*references)
+      # Needed for Ruby 1.8, where we can't set IO objects to not
+      # close the underlying FD on destruction
+      @references += references
+    end
+  end
+end

data/lib/rubysh/version.rb

@@ -0,0 +1,3 @@
+module Rubysh
+  VERSION = "0.0.1"
+end