refactor keepalive to a modulekeepalive-maxcount

2 files changed, 51 insertions, 37 deletions

diff --git a/lib/net/ssh/connection/keepalive.rb b/lib/net/ssh/connection/keepalive.rb
new file mode 100644
index 0000000..6b0daa8
--- /dev/null
+++ b/lib/net/ssh/connection/keepalive.rb
@@ -0,0 +1,47 @@
+module Net; module SSH; module Connection
+
+module Keepalive
+  # Default IO.select timeout threshold
+  DEFAULT_IO_SELECT_TIMEOUT = 300
+
+  def initialize_keepalive
+    @last_keepalive_sent_at = nil
+    @unresponded_keepalive_count = 0
+  end
+
+  def keepalive_enabled?
+    options[:keepalive]
+  end
+
+  def keepalive_interval
+    options[:keepalive_interval] || DEFAULT_IO_SELECT_TIMEOUT
+  end
+
+  def should_send_keepalive?
+    return false unless keepalive_enabled?
+    return true unless @last_keepalive_sent_at
+    Time.now - @last_keepalive_sent_at >= keepalive_interval
+  end
+
+  def keepalive_maxcount
+    (options[:keepalive_maxcount] || 3).to_i
+  end
+
+  def send_keepalive_as_needed(readers, writers)
+    return unless readers.nil? && writers.nil?
+    return unless should_send_keepalive?
+    info { "sending keepalive #{@unresponded_keepalive_count}" }
+
+    @unresponded_keepalive_count += 1
+    send_global_request("keepalive@openssh.com") { |success, response|
+      @unresponded_keepalive_count = 0
+    }
+    if keepalive_maxcount > 0 && @unresponded_keepalive_count > keepalive_maxcount
+      error { "Timeout, server #{host} not responding. Missed #{@unresponded_keepalive_count} timeouts." }
+      raise Net::SSH::Timeout, "Timeout, server #{host} not responding."
+    end
+    @last_keepalive_sent_at = Time.now
+  end
+end
+
+end; end; end
\ No newline at end of file
diff --git a/lib/net/ssh/connection/session.rb b/lib/net/ssh/connection/session.rb
index 12631a1..3137514 100644
--- a/lib/net/ssh/connection/session.rb
+++ b/lib/net/ssh/connection/session.rb
@@ -3,6 +3,7 @@ require 'net/ssh/ruby_compat'
 require 'net/ssh/connection/channel'
 require 'net/ssh/connection/constants'
 require 'net/ssh/service/forward'
+require 'net/ssh/connection/keepalive'
 
 module Net; module SSH; module Connection
 
@@ -23,10 +24,7 @@ module Net; module SSH; module Connection
   #     ssh.exec! "/etc/init.d/some_process start"
   #   end
   class Session
-    include Constants, Loggable
-
-    # Default IO.select timeout threshold
-    DEFAULT_IO_SELECT_TIMEOUT = 300
+    include Constants, Loggable, Keepalive
 
     # The underlying transport layer abstraction (see Net::SSH::Transport::Session).
     attr_reader :transport
@@ -79,8 +77,7 @@ module Net; module SSH; module Connection
       @max_pkt_size = (options.has_key?(:max_pkt_size) ? options[:max_pkt_size] : 0x8000)
       @max_win_size = (options.has_key?(:max_win_size) ? options[:max_win_size] : 0x20000)
 
-      @last_keepalive_sent_at = nil
-      @unresponded_keepalive_count = 0
+      initialize_keepalive
     end
 
     # Retrieves a custom property from this instance. This can be used to
@@ -599,40 +596,10 @@ module Net; module SSH; module Connection
 
       def io_select_wait(wait)
         return wait if wait
-        return wait unless options[:keepalive]
+        return wait unless keepalive_enabled?
         keepalive_interval
       end
 
-      def keepalive_interval
-        options[:keepalive_interval] || DEFAULT_IO_SELECT_TIMEOUT
-      end
-
-      def should_send_keepalive?
-        return false unless options[:keepalive]
-        return true unless @last_keepalive_sent_at
-        Time.now - @last_keepalive_sent_at >= keepalive_interval
-      end
-
-      def keepalive_maxcount
-        (options[:keepalive_maxcount] || 3).to_i
-      end
-
-      def send_keepalive_as_needed(readers, writers)
-        return unless readers.nil? && writers.nil?
-        return unless should_send_keepalive?
-        info { "sending keepalive #{@unresponded_keepalive_count}" }
-
-        @unresponded_keepalive_count += 1
-        send_global_request("keepalive@openssh.com") { |success, response|
-          @unresponded_keepalive_count = 0
-        }
-        if keepalive_maxcount > 0 && @unresponded_keepalive_count > keepalive_maxcount
-          error { "Timeout, server #{host} not responding. Missed #{@unresponded_keepalive_count} timeouts." }
-          raise Net::SSH::Timeout, "Timeout, server #{host} not responding."
-        end
-        @last_keepalive_sent_at = Time.now
-      end
-
       MAP = Constants.constants.inject({}) do |memo, name|
         value = const_get(name)
         next unless Integer === value