Add support for :max_hosts option in task definition or run() (closes…

… #10264)

git-svn-id: http://svn.rubyonrails.org/rails/tools/capistrano@8924 5ecf4fe2-1ee6-0310-87b1-e25e094e27de

CHANGELOG

@@ -1,5 +1,7 @@
 *SVN*
 
+* Add support for :max_hosts option in task definition or run() [Rob Holland <rob@inversepath.com>]
+
 * Distributed git support for better operability with remote_cache strategy [voidlock]
 
 * Use a default line length in help text if line length is otherwise too small [Jamis Buck]

lib/capistrano/configuration/connections.rb

@@ -1,3 +1,5 @@
+
+require 'enumerator'
 require 'capistrano/gateway'
 require 'capistrano/ssh'
 
@@ -95,6 +97,14 @@ def establish_connections_to(servers)
         end
       end
 
+      # Destroys sessions for each server in the list.
+      def teardown_connections_to(servers)
+        servers.each do |server|
+           @sessions[server].close
+           @sessions.delete(server)
+         end
+      end
+
       # Determines the set of servers within the current task's scope and
       # establishes connections to them, and then yields that list of
       # servers.
@@ -120,22 +130,32 @@ def execute_on_servers(options={})
         servers = [servers.first] if options[:once]
         logger.trace "servers: #{servers.map { |s| s.host }.inspect}"
 
-        # establish connections to those servers, as necessary
-        begin
-          establish_connections_to(servers)
-        rescue ConnectionError => error
-          raise error unless task && task.continue_on_error?
-          error.hosts.each do |h|
-            servers.delete(h)
-            failed!(h)
+        max_hosts = (options[:max_hosts] || (task && task.max_hosts) || servers.size).to_i
+        is_subset = max_hosts < servers.size
+
+        # establish connections to those servers in groups of max_hosts, as necessary
+        servers.each_slice(max_hosts) do |servers_slice|
+          begin
+            establish_connections_to(servers_slice)
+          rescue ConnectionError => error
+            raise error unless task && task.continue_on_error?
+            error.hosts.each do |h|
+              servers_slice.delete(h)
+              failed!(h)
+            end
+          end
+
+          begin
+            yield servers_slice
+          rescue RemoteError => error
+            raise error unless task && task.continue_on_error?
+            error.hosts.each { |h| failed!(h) }
           end
-        end
 
-        begin
-          yield servers
-        rescue RemoteError => error
-          raise error unless task && task.continue_on_error?
-          error.hosts.each { |h| failed!(h) }
+          # if dealing with a subset (e.g., :max_hosts is less than the
+          # number of servers available) teardown the subset of connections
+          # that were just made, so that we can make room for the next subset.
+          teardown_connections_to(servers_slice) if is_subset
         end
       end
 

lib/capistrano/task_definition.rb

@@ -3,12 +3,13 @@
 module Capistrano
   # Represents the definition of a single task.
   class TaskDefinition
-    attr_reader :name, :namespace, :options, :body, :desc, :on_error
+    attr_reader :name, :namespace, :options, :body, :desc, :on_error, :max_hosts
 
     def initialize(name, namespace, options={}, &block)
       @name, @namespace, @options = name, namespace, options
       @desc = @options.delete(:desc)
       @on_error = options.delete(:on_error)
+      @max_hosts = options[:max_hosts] && options[:max_hosts].to_i
       @body = block or raise ArgumentError, "a task requires a block"
       @servers = nil
     end

test/configuration/connections_test.rb

@@ -217,7 +217,6 @@ def test_execute_on_servers_should_not_try_to_connect_to_hosts_with_connection_e
     @config.current_task = mock_task(:on_error => :continue)
     @config.expects(:find_servers_for_task).with(@config.current_task, {}).returns(list)
     Capistrano::SSH.expects(:connect).times(2).raises(Exception).then.returns(:success)
-    @config.expects(:failed!).with(server("cap1"))
     @config.execute_on_servers do |servers|
       assert_equal %w(cap2), servers.map { |s| s.host }
     end
@@ -260,7 +259,7 @@ def test_execute_on_servers_should_not_try_to_connect_to_hosts_with_upload_error
       assert_equal %w(cap2), servers.map { |s| s.host }
     end
   end
-
+  
   def test_connect_should_establish_connections_to_all_servers_in_scope
     assert @config.sessions.empty?
     @config.current_task = mock_task
@@ -269,7 +268,43 @@ def test_connect_should_establish_connections_to_all_servers_in_scope
     @config.connect!
     assert_equal %w(cap1 cap2 cap3), @config.sessions.keys.sort.map { |s| s.host }
   end
-
+
+  def test_execute_on_servers_should_only_run_on_tasks_max_hosts_hosts_at_once
+    cap1 = server("cap1")
+    cap2 = server("cap2")
+    connection1 = mock()
+    connection2 = mock()
+    connection1.expects(:close)
+    connection2.expects(:close)
+    @config.current_task = mock_task(:max_hosts => 1)
+    @config.expects(:find_servers_for_task).with(@config.current_task, {}).returns([cap1, cap2])
+    Capistrano::SSH.expects(:connect).times(2).returns(connection1).then.returns(connection2)
+    block_called = 0
+    @config.execute_on_servers do |servers|
+      block_called += 1
+      assert_equal 1, servers.size
+    end
+    assert_equal 2, block_called
+  end
+
+  def test_execute_on_servers_should_only_run_on_max_hosts_hosts_at_once
+    cap1 = server("cap1")
+    cap2 = server("cap2")
+    connection1 = mock()
+    connection2 = mock()
+    connection1.expects(:close)
+    connection2.expects(:close)
+    @config.current_task = mock_task(:max_hosts => 1)
+    @config.expects(:find_servers_for_task).with(@config.current_task, {}).returns([cap1, cap2])
+    Capistrano::SSH.expects(:connect).times(2).returns(connection1).then.returns(connection2)
+    block_called = 0
+    @config.execute_on_servers do |servers|
+      block_called += 1
+      assert_equal 1, servers.size
+    end
+    assert_equal 2, block_called
+  end
+
   def test_connect_should_honor_once_option
     assert @config.sessions.empty?
     @config.current_task = mock_task
@@ -283,6 +318,11 @@ def test_connect_should_honor_once_option
 
     def mock_task(options={})
       continue_on_error = options[:on_error] == :continue
-      stub("task", :fully_qualified_name => "name", :options => options, :continue_on_error? => continue_on_error)
+      stub("task",
+        :fully_qualified_name => "name",
+        :options => options,
+        :continue_on_error? => continue_on_error,
+        :max_hosts => options[:max_hosts]
+      )
     end
 end