lib/ztk/command.rb in ztk-0.2.4 vs lib/ztk/command.rb in ztk-0.2.5

@@ -17,10 +17,11 @@
 #   limitations under the License.
 #
 ################################################################################
 
 require "ostruct"
+require "timeout"
 
 module ZTK
 
   # ZTK::Command Error Class
   #
@@ -39,12 +40,15 @@
   #     cmd = ZTK::Command.new(:stdout => std_combo, :stderr => std_combo)
   #
   # @author Zachary Patten <[e-mail]>
   class Command < ZTK::Base
 
-    def initialize(config={})
-      super(config)
+    def initialize(configuration={})
+      super({
+        :timeout => 600
+      }.merge(configuration))
+      config.logger.debug { "config(#{config.inspect})" }
     end
 
     def inspect
       @hostname ||= %x(hostname -f).chomp
       "#{ENV['USER']}@#{@hostname}"
@@ -54,12 +58,12 @@
     #
     # @param [String] command The command to execute.
     # @param [Hash] options The options hash for executing the command.
     #
     # @return [OpenStruct#output] The output of the command, both STDOUT and
-    #   STDERR.
-    # @return [OpenStruct#exit] The exit status (i.e. $?).
+    #   STDERR combined.
+    # @return [OpenStruct#exit_code] The exit code of the process.
     #
     # @example Execute a command:
     #
     #   cmd = ZTK::Command.new
     #   puts cmd.exec("hostname -f").inspect
@@ -71,20 +75,25 @@
         header = [sep, "[ #{tag} ]", sep, "[ #{self.inspect} ]", sep, "[ #{tag} ]", sep].join
         "#{header}\n"
       end
 
       options = OpenStruct.new({ :exit_code => 0, :silence => false }.merge(options))
-      log(:debug) { "options(#{options.inspect})" }
-      log(:debug) { "command(#{command.inspect})" }
 
+      config.logger.debug { "config(#{config.inspect})" }
+      config.logger.debug { "options(#{options.inspect})" }
+      config.logger.info { "command(#{command.inspect})" }
+
       output = ""
+      exit_code = -1
       stdout_header = false
       stderr_header = false
 
       parent_stdout_reader, child_stdout_writer = IO.pipe
       parent_stderr_reader, child_stderr_writer = IO.pipe
 
+      start_time = Time.now.utc
+
       pid = Process.fork do
         parent_stdout_reader.close
         parent_stderr_reader.close
 
         STDOUT.reopen(child_stdout_writer)
@@ -103,64 +112,74 @@
       reader_writer_map = {parent_stdout_reader => @config.stdout, parent_stderr_reader => @config.stderr}
 
       direct_log(:debug) { log_header("COMMAND") }
       direct_log(:debug) { "#{command}\n" }
       direct_log(:debug) { log_header("STARTED") }
-      loop do
-        break if reader_writer_map.keys.all?{ |reader| reader.eof? }
 
-        sockets = IO.select(reader_writer_map.keys).first
-        sockets.each do |socket|
-          data = socket.read
-          next if (data.nil? || data.empty?)
+      begin
+        Timeout.timeout(config.timeout) do
+          loop do
+            pipes = IO.select(reader_writer_map.keys, [], reader_writer_map.keys).first
+            pipes.each do |pipe|
+              data = pipe.read
+              next if (data.nil? || data.empty?)
 
-          case reader_writer_key[socket]
-          when :stdout then
-            if !stdout_header
-              direct_log(:debug) { log_header("STDOUT") }
-              stdout_header = true
-              stderr_header = false
-            end
-            reader_writer_map[socket].write(data) unless options.silence
-            direct_log(:debug) { data }
+              case reader_writer_key[pipe]
+              when :stdout then
+                if !stdout_header
+                  direct_log(:debug) { log_header("STDOUT") }
+                  stdout_header = true
+                  stderr_header = false
+                end
+                reader_writer_map[pipe].write(data) unless options.silence
+                direct_log(:debug) { data }
 
-          when :stderr then
-            if !stderr_header
-              direct_log(:warn) { log_header("STDERR") }
-              stderr_header = true
-              stdout_header = false
+              when :stderr then
+                if !stderr_header
+                  direct_log(:warn) { log_header("STDERR") }
+                  stderr_header = true
+                  stdout_header = false
+                end
+                reader_writer_map[pipe].write(data) unless options.silence
+                direct_log(:warn) { data }
+              end
+
+              output += data
             end
-            reader_writer_map[socket].write(data) unless options.silence
-            direct_log(:warn) { data }
+            break if reader_writer_map.keys.all?{ |reader| reader.eof? }
           end
-
-          output += data
         end
+      rescue Timeout::Error => e
+        direct_log(:debug) { log_header("TIMEOUT") }
+        log_and_raise(CommandError, "Process timed out after #{config.timeout} seconds!")
       end
-      direct_log(:debug) { log_header("STOPPED") }
 
       Process.waitpid(pid)
+      direct_log(:debug) { log_header("STOPPED") }
 
       parent_stdout_reader.close
       parent_stderr_reader.close
 
-      log(:debug) { "exit_code(#{$?.inspect})" }
-
-      if ($? != options.exit_code)
-        message = "exec(#{command.inspect}, #{options.inspect}) failed! [#{$?.inspect}]"
-        log(:fatal) { message }
-        raise CommandError, message
+      if RUBY_VERSION >= "1.9"
+        exit_code = $?.exitstatus
+      else
+        exit_code = $?
       end
 
-      OpenStruct.new(:output => output, :exit => $?)
+      config.logger.debug { "exit_code(#{exit_code})" }
+
+      if (exit_code != options.exit_code)
+        log_and_raise(CommandError, "exec(#{command.inspect}, #{options.inspect}) failed! [#{exit_code}]")
+      end
+      OpenStruct.new(:output => output, :exit_code => exit_code)
     end
 
     def upload(*args)
-      raise CommandError, "Not Implemented"
+      log_and_raise(CommandError, "Not Supported")
     end
 
     def download(*args)
-      raise CommandError, "Not Implemented"
+      log_and_raise(CommandError, "Not Supported")
     end
 
   end
 
 end