lib/munin-ruby/connection.rb in munin-ruby-0.2.2 vs lib/munin-ruby/connection.rb in munin-ruby-0.2.3

@@ -1,7 +1,5 @@
-require 'socket'
-
 module Munin
   class Connection
     include Munin::Parser
     
     attr_reader :host, :port
@@ -27,17 +25,23 @@
     
     # Establish connection to the server
     #
     def open
       begin
-        @socket = TCPSocket.new(@host, @port)
-        @socket.sync = true
-        welcome = @socket.gets
-        unless welcome =~ /^# munin node at/
-          raise Munin::AccessDenied
+        begin
+          with_timeout do
+            @socket = TCPSocket.new(@host, @port)
+            @socket.sync = true
+            welcome = @socket.gets
+            unless welcome =~ /^# munin node at/
+              raise Munin::AccessDenied
+            end
+            @connected = true
+          end
+        rescue Timeout::Error
+          raise Munin::ConnectionError, "Timed out talking to #{@host}"
         end
-        @connected = true
       rescue Errno::ETIMEDOUT, Errno::ECONNREFUSED, Errno::ECONNRESET => ex
         raise Munin::ConnectionError, ex.message
       rescue EOFError
         raise Munin::AccessDenied
       rescue Exception => ex
@@ -64,35 +68,59 @@
           raise Munin::ConnectionError, "Not connected."
         else
           open
         end
       end
-      @socket.puts("#{str.strip}\n")
+
+      begin
+        with_timeout { @socket.puts("#{str.strip}\n") }
+      rescue Timeout::Error
+        raise Munin::ConnectionError, "Timed out on #{@host} trying to send."
+      end
     end
     
     # Reads a single line from socket
     #
     def read_line
       begin
-        @socket.gets.to_s.strip
+        with_timeout { @socket.gets.to_s.strip }
       rescue Errno::ETIMEDOUT, Errno::ECONNREFUSED, Errno::ECONNRESET, EOFError => ex
         raise Munin::ConnectionError, ex.message
+      rescue Timeout::Error
+        raise Munin::ConnectionError, "Timed out reading from #{@host}."
       end
     end
     
     # Reads a packet of data until '.' reached
     #
     def read_packet
       begin
-        lines = []
-        while(str = @socket.readline.to_s) do
-          break if str.strip == '.'
-          lines << str.strip
+        with_timeout do
+          lines = []
+          while(str = @socket.readline.to_s) do
+            break if str.strip == '.'
+            lines << str.strip
+          end
+          parse_error(lines)
+          lines
         end
-        parse_error(lines)
-        lines
       rescue Errno::ETIMEDOUT, Errno::ECONNREFUSED, Errno::ECONNRESET, EOFError => ex
         raise Munin::ConnectionError, ex.message
+      rescue Timeout::Error
+        raise Munin::ConnectionError, "Timed out reading from #{@host}."
       end
     end
+
+    private
+
+    # Execute operation with timeout
+    # @param [Block] block Block to execute
+    def with_timeout(time=Munin::TIMEOUT_TIME)
+      raise ArgumentError, "Block required" if !block_given?
+      if Munin::TIMEOUT_CLASS.respond_to?(:timeout_after)
+        Munin::TIMEOUT_CLASS.timeout_after(time) { yield }
+      else
+        Munin::TIMEOUT_CLASS.timeout(time) { yield }
+      end
+    end
   end
-end
\ No newline at end of file
+end