From 70c5e93084ed2e786e62c48bbfeaf8bad48534a2 Mon Sep 17 00:00:00 2001
From: Jamis Buck <jamis@37signals.com>
Date: Mon, 7 Apr 2008 23:07:25 -0600
Subject: :on_error support

---
 lib/net/ssh/multi/server.rb          | 13 +++++++++++++
 lib/net/ssh/multi/session.rb         | 22 +++++++++++++++++++++-
 lib/net/ssh/multi/session_actions.rb |  2 +-
 3 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/lib/net/ssh/multi/server.rb b/lib/net/ssh/multi/server.rb
index 999e788..3354b0a 100644
--- a/lib/net/ssh/multi/server.rb
+++ b/lib/net/ssh/multi/server.rb
@@ -36,6 +36,7 @@ module Net; module SSH; module Multi
       @user = user
       @options = options.dup
       @gateway = @options.delete(:via)
+      @failed = false
     end
 
     # Returns the value of the server property with the given +key+. Server
@@ -82,6 +83,18 @@ module Net; module SSH; module Multi
       @inspect ||= "#<%s:0x%x %s>" % [self.class.name, object_id, to_s]
     end
 
+    # Returns +true+ if this server has ever failed a connection attempt.
+    def failed?
+      @failed
+    end
+
+    # Indicates (by default) that this server has just failed a connection
+    # attempt. If +flag+ is false, this can be used to reset the failed flag
+    # so that a retry may be attempted.
+    def fail!(flag=true)
+      @failed = flag
+    end
+
     # Returns the Net::SSH session object for this server. If +require_session+
     # is false and the session has not previously been created, this will
     # return +nil+. If +require_session+ is true, the session will be instantiated
diff --git a/lib/net/ssh/multi/session.rb b/lib/net/ssh/multi/session.rb
index 561f9d1..6c31dea 100644
--- a/lib/net/ssh/multi/session.rb
+++ b/lib/net/ssh/multi/session.rb
@@ -60,6 +60,11 @@ module Net; module SSH; module Multi
     # of sessions will be open at any given time.
     attr_accessor :concurrent_connections
 
+    # How connection errors should be handled. This defaults to :fail, but
+    # may be set to :ignore if connection errors should be ignored, or
+    # :warn if connection errors should cause a warning.
+    attr_accessor :on_error
+
     # The number of connections that are currently open.
     attr_reader :open_connections #:nodoc:
 
@@ -86,6 +91,7 @@ module Net; module SSH; module Multi
       @gateway = nil
       @open_groups = []
       @connect_threads = []
+      @on_error = :fail
 
       @open_connections = 0
       @pending_sessions = []
@@ -351,6 +357,9 @@ module Net; module SSH; module Multi
     # If +force+ is true, the concurrent_connections check is skipped and a real
     # connection is always returned.
     def next_session(server, force=false) #:nodoc:
+      # don't retry a failed attempt
+      return nil if server.failed?
+
       @session_mutex.synchronize do
         if !force && concurrent_connections && concurrent_connections <= open_connections
           connection = PendingConnection.new(server)
@@ -364,8 +373,19 @@ module Net; module SSH; module Multi
       begin
         server.new_session
       rescue Exception => e
+        server.fail!
         @session_mutex.synchronize { @open_connections -= 1 }
-        raise
+
+        case on_error
+        when :ignore then
+          # do nothing
+        when :warn then
+          warn("error connecting to #{server}: #{e.class} (#{e.message})")
+        else
+          raise
+        end
+
+        return nil
       end
     end
 
diff --git a/lib/net/ssh/multi/session_actions.rb b/lib/net/ssh/multi/session_actions.rb
index 35c69ed..2d87392 100644
--- a/lib/net/ssh/multi/session_actions.rb
+++ b/lib/net/ssh/multi/session_actions.rb
@@ -35,7 +35,7 @@ module Net; module SSH; module Multi
     def sessions
       threads = servers.map { |server| Thread.new { server.session(true) } if server.session.nil? }
       threads.each { |thread| thread.join if thread }
-      servers.map { |server| server.session }
+      servers.map { |server| server.session }.compact
     end
 
     # Sends a global request to the sessions for all contained servers
-- 
cgit v1.2.1