ztk 0.2.4 → 0.2.5

data/lib/ztk/base.rb

@@ -35,26 +35,70 @@ module ZTK
   # and extend functionality as appropriate.
   class Base
 
+    class << self
+
+      # @param [Hash] config Configuration options hash.
+      # @option config [IO] :stdout Instance of IO to be used for STDOUT.
+      # @option config [IO] :stderr Instance of IO to be used for STDERR.
+      # @option config [IO] :stdin Instance of IO to be used for STDIN.
+      # @option config [Logger] :logger Instance of Logger to be used for logging.
+      def build_config(configuration={})
+        if configuration.is_a?(OpenStruct)
+          configuration = configuration.send(:table)
+        end
+
+        rails_logger = Rails.logger if defined?(Rails)
+
+        config = OpenStruct.new({
+          :stdout => $stdout,
+          :stderr => $stderr,
+          :stdin => $stdin,
+          :logger => ($logger || rails_logger || ZTK::Logger.new("/dev/null"))
+        }.merge(configuration))
+
+        (config.stdout && config.stdout.respond_to?(:sync=)) and config.stdout.sync = true
+        (config.stderr && config.stderr.respond_to?(:sync=)) and config.stderr.sync = true
+        (config.stdin  && config.stdin.respond_to?(:sync=))  and config.stdin.sync = true
+        (config.logger && config.logger.respond_to?(:sync=)) and config.logger.sync = true
+
+        config
+      end
+
+      # Removes all key-value pairs which are not core so values do not bleed
+      # into classes they are not meant for.
+      #
+      # This method will leave :stdout, :stderr, :stdin and :logger key-values
+      # intact, while removing all other key-value pairs.
+      def sanitize_config(configuration={})
+        if configuration.is_a?(OpenStruct)
+          configuration = configuration.send(:table)
+        end
+
+        config = configuration.reject do |key,value|
+          !(%w(stdout stderr stdin logger).map(&:to_sym).include?(key))
+        end
+
+        config
+      end
+
+      def log_and_raise(logger, exception, message, shift=1)
+        if logger.is_a?(ZTK::Logger)
+          logger.shift(:fatal, shift) { "EXCEPTION: #{exception.inspect} - #{message.inspect}" }
+        else
+          logger.fatal { "EXCEPTION: #{exception.inspect} - #{message.inspect}" }
+        end
+        raise exception, message
+      end
+
+    end
+
     # @param [Hash] config Configuration options hash.
     # @option config [IO] :stdout Instance of IO to be used for STDOUT.
     # @option config [IO] :stderr Instance of IO to be used for STDERR.
     # @option config [IO] :stdin Instance of IO to be used for STDIN.
     # @option config [Logger] :logger Instance of Logger to be used for logging.
     def initialize(config={})
-      # defined?(Rails) and rails_logger = Rails.logger
-      @config = OpenStruct.new({
-        :stdout => $stdout,
-        :stderr => $stderr,
-        :stdin => $stdin,
-        :logger => $logger
-      }.merge(config))
-
-      @config.stdout.respond_to?(:sync=) and @config.stdout.sync = true
-      @config.stderr.respond_to?(:sync=) and @config.stderr.sync = true
-      @config.stdin.respond_to?(:sync=) and @config.stdin.sync = true
-      @config.logger.respond_to?(:sync=) and @config.logger.sync = true
-
-      log(:debug) { "config(#{@config.inspect})" }
+      @config = Base.build_config(config)
     end
 
     # Configuration OpenStruct accessor method.
@@ -73,7 +117,15 @@ module ZTK
       end
     end
 
-    # Base logging method.
+    def log_and_raise(exception, message)
+      Base.log_and_raise(config.logger, exception, message, 2)
+    end
+
+    # Direct logging method.
+    #
+    # This method provides direct writing of data to the current log device.
+    # This is mainly used for pushing STDOUT and STDERR into the log file in
+    # ZTK::SSH and ZTK::Command, but could easily be used by other classes.
     #
     # The value returned in the block is passed down to the logger specified in
     # the classes configuration.
@@ -81,24 +133,18 @@ module ZTK
     # @param [Symbol] method_name This should be any one of [:debug, :info, :warn, :error, :fatal].
     # @yield No value is passed to the block.
     # @yieldreturn [String] The message to log.
-    def log(method_name, &block)
-      if block_given?
-        @config.logger and @config.logger.method(method_name.to_sym).call { yield }
-      else
-        raise BaseError, "You must supply a block to the log method!"
-      end
-    end
-
     def direct_log(log_level, &blocK)
       @config.logger.nil? and raise BaseError, "You must supply a logger for direct logging support!"
 
       if !block_given?
-        raise BaseError, "You must supply a block to the log method!"
+        log_and_raise(BaseError, "You must supply a block to the log method!")
       elsif (@config.logger.level <= ZTK::Logger.const_get(log_level.to_s.upcase))
         if @config.logger.respond_to?(:logdev)
           @config.logger.logdev.write(yield)
+          @config.logger.logdev.respond_to?(:flush) and @config.logger.logdev.flush
         else
           @config.logger.instance_variable_get(:@logdev).instance_variable_get(:@dev).write(yield)
+          @config.logger.instance_variable_get(:@logdev).instance_variable_get(:@dev).respond_to?(:flush) and @config.logger.instance_variable_get(:@logdev).instance_variable_get(:@dev).flush
         end
       end
     end

data/lib/ztk/benchmark.rb

@@ -34,30 +34,27 @@ module ZTK
     class << self
 
       def bench(options={}, &block)
-        !block_given? and raise BenchmarkError, "You must supply a block!"
+        options = Base.build_config(options)
+        options.logger.debug { "options(#{options.inspect})" }
 
-        options = { :stdout => STDOUT, :logger => $logger, :message => nil, :mark => nil }.merge(options)
+        !block_given? and Base.log_and_raise(options.logger, BenchmarkError, "You must supply a block!")
 
-        stdout = options[:stdout]
-        logger = options[:logger]
-        message = options[:message]
-        mark = options[:mark]
+        check = [options.message, options.mark]
+        (check.any?{ |z| !z.nil? } && !check.all?{ |z| !z.nil? }) and Base.log_and_raise(options.logger, BenchmarkError, "You must supply both a message and a mark!")
 
-        logger and logger.debug { options.inspect }
-
-        (!message.nil? && !mark.nil?) and stdout.print("#{message} ")
+        (options.message && options.mark) and options.stdout.print("#{options.message} ")
         benchmark = ::Benchmark.realtime do
-          if message.nil?
-            yield
-          else
-            ZTK::Spinner.spin do
+          if (options.message && options.mark)
+            ZTK::Spinner.spin(Base.sanitize_config(options)) do
               yield
             end
+          else
+            yield
           end
         end
 
-        (!message.nil? && !mark.nil?) and stdout.print("#{mark}\n" % benchmark)
-        logger and logger.info { "#{message} #{mark}" }
+        (options.message && options.mark) and options.stdout.print("#{options.mark}\n" % benchmark)
+        options.logger.info { "#{options.message} #{options.mark}" }
 
         benchmark
       end

data/lib/ztk/command.rb

@@ -19,6 +19,7 @@
 ################################################################################
 
 require "ostruct"
+require "timeout"
 
 module ZTK
 
@@ -41,8 +42,11 @@ module ZTK
   # @author Zachary Patten <zachary@jovelabs.net>
   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
@@ -56,8 +60,8 @@ module ZTK
     # @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:
     #
@@ -73,16 +77,21 @@ module ZTK
       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
@@ -105,61 +114,71 @@ module ZTK
       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?)
-
-          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 }
-
-          when :stderr then
-            if !stderr_header
-              direct_log(:warn) { log_header("STDERR") }
-              stderr_header = true
-              stdout_header = false
+
+      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[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
+                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

data/lib/ztk/logger.rb

@@ -56,6 +56,11 @@ module ZTK
       self.instance_variable_get(:@logdev).instance_variable_get(:@dev)
     end
 
+    def shift(severity, shift=0, &block)
+      severity = ZTK::Logger.const_get(severity.to_s.upcase)
+      add(severity, nil, nil, shift, &block)
+    end
+
 
   private
 
@@ -81,17 +86,18 @@ module ZTK
     # @param [String] progname Optional name of the program to prefix the log
     #   entry with.
     # @yieldreturn [String] The block should return the desired log message.
-    def add(severity, message = nil, progname = nil, &block)
+    def add(severity, message=nil, progname=nil, shift=0, &block)
       return if (@level > severity)
 
       msg = (block && block.call)
       (msg.nil? || msg.strip.empty?) and return
       @hostname ||= %x(hostname -s).chomp.strip
-      called_by = "#{@hostname}:#{parse_caller(caller[1])}"
+      called_by = "#{@hostname}:#{parse_caller(caller[1+shift])}"
       message = [message, progname, msg].flatten.compact.join(": ")
       message = "%19s.%06d|%05d|%5s|%60s%s\n" % [Time.now.utc.strftime("%Y-%m-%d|%H:%M:%S"), Time.now.utc.usec, Process.pid, SEVERITIES[severity], called_by, message]
 
       @logdev.write(message)
+      @logdev.respond_to?(:flush) and @logdev.flush
 
       true
     end

data/lib/ztk/parallel.rb

@@ -76,12 +76,13 @@ module ZTK
     # @option config [Integer] :max_forks Maximum number of forks to use.
     # @option config [Proc] :before_fork (nil) Proc to call before forking.
     # @option config [Proc] :after_fork (nil) Proc to call after forking.
-    def initialize(config={})
+    def initialize(configuration={})
       super({
         :max_forks => MAX_FORKS
-      }.merge(config))
+      }.merge(configuration))
+      config.logger.debug { "config(#{config.inspect})" }
 
-      raise ParallelError, "max_forks must be equal to or greater than one!" if @config.max_forks < 1
+      (config.max_forks < 1) and log_and_raise(ParallelError, "max_forks must be equal to or greater than one!")
 
       @forks = Array.new
       @results = Array.new
@@ -95,26 +96,26 @@ module ZTK
     #   the parent processes result set.
     # @return [Integer] Returns the pid of the child process forked.
     def process(&block)
-      raise ParallelError, "You must supply a block to the process method!" if !block_given?
+      !block_given? and log_and_raise(ParallelError, "You must supply a block to the process method!")
 
-      log(:debug) { "forks(#{@forks.inspect})" }
+      config.logger.debug { "forks(#{@forks.inspect})" }
 
-      while (@forks.count >= @config.max_forks) do
+      while (@forks.count >= config.max_forks) do
         wait
       end
 
       child_reader, parent_writer = IO.pipe
       parent_reader, child_writer = IO.pipe
 
-      @config.before_fork and @config.before_fork.call(Process.pid)
+      config.before_fork and config.before_fork.call(Process.pid)
       pid = Process.fork do
-        @config.after_fork and @config.after_fork.call(Process.pid)
+        config.after_fork and config.after_fork.call(Process.pid)
 
         parent_writer.close
         parent_reader.close
 
         if !(data = block.call).nil?
-          log(:debug) { "write(#{data.inspect})" }
+          config.logger.debug { "write(#{data.inspect})" }
           child_writer.write(Base64.encode64(Marshal.dump(data)))
         end
 
@@ -122,7 +123,7 @@ module ZTK
         child_writer.close
         Process.exit!(0)
       end
-      @config.after_fork and @config.after_fork.call(Process.pid)
+      config.after_fork and config.after_fork.call(Process.pid)
 
       child_reader.close
       child_writer.close
@@ -142,12 +143,12 @@ module ZTK
     #   status and data returned from the process block.  If wait2() fails nil
     #   is returned.
     def wait
-      log(:debug) { "wait" }
-      log(:debug) { "forks(#{@forks.inspect})" }
+      config.logger.debug { "wait" }
+      config.logger.debug { "forks(#{@forks.inspect})" }
       pid, status = (Process.wait2(-1) rescue nil)
       if !pid.nil? && !status.nil? && !(fork = @forks.select{ |f| f[:pid] == pid }.first).nil?
         data = (Marshal.load(Base64.decode64(fork[:reader].read.to_s)) rescue nil)
-        log(:debug) { "read(#{data.inspect})" }
+        config.logger.debug { "read(#{data.inspect})" }
         !data.nil? and @results.push(data)
         fork[:reader].close
         fork[:writer].close
@@ -162,7 +163,7 @@ module ZTK
     #
     # @return [Array<Object>] The results from all of the *process* blocks.
     def waitall
-      log(:debug) { "waitall" }
+      config.logger.debug { "waitall" }
       while @forks.count > 0
         self.wait
       end
@@ -173,7 +174,7 @@ module ZTK
     #
     # @return [Integer] Current number of active forks.
     def count
-      log(:debug) { "count(#{@forks.count})" }
+      config.logger.debug { "count(#{@forks.count})" }
       @forks.count
     end