lib/heroku/executor.rb in heroku-commander-0.1.0 vs lib/heroku/executor.rb in heroku-commander-0.2.0

@@ -1,81 +1,111 @@
 module Heroku
   class Executor
 
     class Terminate < StandardError
+
+      attr_accessor :timeout
+
+      def initialize(timeout = 0)
+        @timeout = timeout
+      end
     end
 
     class << self
 
       # Executes a command and yields output line-by-line.
       def run(cmd, options = {}, &block)
         lines = []
+        running_pid = nil
         logger = options[:logger]
         logger.debug "Running: #{cmd}" if logger
         PTY.spawn(cmd) do |r, w, pid|
+          running_pid = pid
           logger.debug "Started: #{pid}" if logger
           terminated = false
           begin
             r.sync = true
-            r.each do |line|
-              line.strip! if line
-              logger.debug "#{pid}: #{line}" if logger
-              if block_given?
-                yield line
-              end
-              lines << line
-            end
-          rescue Heroku::Executor::Terminate
-            logger.debug "Terminating #{pid}." if logger
-            Process.kill("TERM", pid)
+            read_from(r, pid, options, lines, &block)
+          rescue Heroku::Executor::Terminate => e
+            logger.debug "Waiting: #{e.timeout} second(s) to terminate #{pid}" if logger
+            terminate_process! pid, e.timeout
             terminated = true
           rescue Errno::EIO, IOError => e
-            logger.debug "#{e.class}: #{e.message}" if logger
+            # ignore
           rescue PTY::ChildExited => e
             logger.debug "Terminated: #{pid}" if logger
-            terminted = true
+            terminated = true
             raise e
           ensure
             unless terminated
+              # wait for process
               logger.debug "Waiting: #{pid}" if logger
-              Process.wait(pid) rescue Errno::ECHILD
+              ::Process.wait(pid)
             end
           end
         end
-        check_exit_status! cmd, $?.exitstatus, lines
+        check_exit_status! cmd, running_pid, $?.exitstatus, lines
         lines
       rescue Errno::ECHILD => e
-        logger.debug "#{e.class}: #{e.message}" if logger
-        check_exit_status! cmd, $?.exitstatus, lines
+        check_exit_status! cmd, running_pid, $?.exitstatus, lines
         lines
       rescue PTY::ChildExited => e
-        logger.debug "#{e.class}: #{e.message}" if logger
-        check_exit_status! cmd, $!.status.exitstatus, lines
+        check_exit_status! cmd, running_pid, $!.status.exitstatus, lines
         lines
       rescue Heroku::Commander::Errors::Base => e
         logger.debug "Error: #{e.problem}" if logger
         raise
       rescue Exception => e
         logger.debug "#{e.class}: #{e.respond_to?(:problem) ? e.problem : e.message}" if logger
         raise Heroku::Commander::Errors::CommandError.new({
           :cmd => cmd,
+          :pid => running_pid,
           :status => $?.exitstatus,
           :message => e.message,
           :inner_exception => e,
           :lines => lines
         })
       end
 
       private
 
-        def check_exit_status!(cmd, status, lines = nil)
+        def read_from(r, pid, options, lines, &block)
+          logger = options[:logger]
+          while ! r.eof do
+            line = r.readline
+            line.strip! if line
+            logger.debug "#{pid}: #{line}" if logger
+            if block_given?
+              yield line
+            end
+            lines << line
+          end
+        end
+
+        def check_exit_status!(cmd, pid, status, lines = nil)
           return if ! status || status == 0
           raise Heroku::Commander::Errors::CommandError.new({
             :cmd => cmd,
+            :pid => pid,
             :status => status,
             :message => "The command #{cmd} failed with exit status #{status}.",
             :lines => lines
           })
+        end
+
+        def terminate_process!(pid, timeout)
+          if timeout
+            # Delay terminating of the process, usually to let the output flush.
+            Thread.new(pid, timeout) do |pid, timeout|
+              begin
+                sleep(timeout) if timeout
+                ::Process.kill("TERM", pid)
+              rescue
+              end
+            end
+          else
+            ::Process.kill("TERM", pid)
+          end
         end
 
     end
   end
 end