Adjusted pool size to be minimum of requested pool size or number of hosts. Added test for pool size

Author: Dan

pssh.py

@@ -305,29 +305,33 @@ def __init__(self, hosts,
         """
         :param hosts: Hosts to connect to
         :type hosts: list(str)
-        :param user: (Optional) User to login as. Defaults to logged in user or\
+        :param user: (Optional) User to login as. Defaults to logged in user or \
         user from ~/.ssh/config or /etc/ssh/ssh_config if set
         :type user: str
-        :param password: (Optional) Password to use for login. Defaults to\
+        :param password: (Optional) Password to use for login. Defaults to \
         no password
         :type password: str
-        :param port: (Optional) Port number to use for SSH connection. Defaults\
+        :param port: (Optional) Port number to use for SSH connection. Defaults \
         to None which uses SSH default
         :type port: int
         :param pkey: (Optional) Client's private key to be used to connect with
         :type pkey: :mod:`paramiko.PKey`
-        :param num_retries: (Optional) Number of retries for connection attempts\
+        :param num_retries: (Optional) Number of retries for connection attempts \
         before the client gives up. Defaults to 3.
         :type num_retries: int
-        :param timeout: (Optional) Number of seconds to timout connection attempts\
-        before the client gives up. Defaults to 10.
+        :param timeout: (Optional) Number of seconds to timout connection \
+        attempts before the client gives up. Defaults to 10.
         :type timeout: int
         :param forward_ssh_agent: (Optional) Turn on SSH agent forwarding - \
         equivalent to `ssh -A` from the `ssh` command line utility. \
         Defaults to True if not set.
         :type forward_ssh_agent: bool
         :param pool_size: (Optional) Greenlet pool size. Controls on how many\
-        hosts to execute tasks in parallel. Defaults to 10
+        hosts to execute tasks in parallel. Defaults to number of hosts or 10, \
+        whichever is lower. Pool size will be *equal to* number of hosts if number\
+        of hosts is lower than the pool size specified as that would only \
+        increase overhead with no benefits.
+        
         :type pool_size: int
         
         **Example**
@@ -387,8 +391,8 @@ def __init__(self, hosts,
           
           Connection is terminated.
         """
-        self.pool = gevent.pool.Pool(size=pool_size)
-        self.pool_size = pool_size
+        self.pool_size = len(hosts) if len(hosts) < pool_size else pool_size
+        self.pool = gevent.pool.Pool(size=self.pool_size)
         self.hosts = hosts
         self.user = user
         self.password = password

tests/test_pssh_client.py

@@ -21,9 +21,9 @@
 
 import unittest
 from pssh import ParallelSSHClient, UnknownHostException, \
-    AuthenticationException, ConnectionErrorException
-from fake_server.fake_server import start_server, make_socket, logger as server_logger, \
-    paramiko_logger
+     AuthenticationException, ConnectionErrorException
+from fake_server.fake_server import start_server, make_socket, \
+     logger as server_logger, paramiko_logger
 import random
 import logging
 import gevent
@@ -165,8 +165,9 @@ def test_pssh_client_run_long_command(self):
         output = client.run_command(self.long_running_cmd)
         self.assertTrue('127.0.0.1' in output, msg="Got no output for command")
         stdout = list(output['127.0.0.1']['stdout'])
-        self.assertTrue(len(stdout) == expected_lines, msg="Expected %s lines of response, got %s" %
-                        (expected_lines, len(stdout)))
+        self.assertTrue(len(stdout) == expected_lines,
+                        msg="Expected %s lines of response, got %s" % (
+                            expected_lines, len(stdout)))
         del client
         server.kill()
 
@@ -198,7 +199,8 @@ def test_pssh_client_timeout(self):
             gevent.sleep(0.5)
             cmd.get()
             if not server.exception:
-                raise Exception("Expected gevent.Timeout from socket timeout, got none")
+                raise Exception(
+                    "Expected gevent.Timeout from socket timeout, got none")
             raise server.exception
         except gevent.Timeout:
             pass
@@ -217,7 +219,7 @@ def test_pssh_client_exec_command_password(self):
         output = client.get_stdout(cmd)
         expected = {'127.0.0.1' : {'exit_code' : 0}}
         self.assertEqual(expected, output,
-                         msg = "Got unexpected command output - %s" % (output,))
+                         msg="Got unexpected command output - %s" % (output,))
         del client
         server.join()
 
@@ -232,8 +234,9 @@ def test_pssh_client_long_running_command(self):
         output = client.get_stdout(cmd)
         self.assertTrue('127.0.0.1' in output, msg="Got no output for command")
         stdout = list(output['127.0.0.1']['stdout'])
-        self.assertTrue(len(stdout) == expected_lines, msg="Expected %s lines of response, got %s" %
-                        (expected_lines, len(stdout)))
+        self.assertTrue(len(stdout) == expected_lines,
+                        msg="Expected %s lines of response, got %s" % (
+                            expected_lines, len(stdout)))
         del client
         server.kill()
 
@@ -242,7 +245,8 @@ def test_pssh_client_retries(self):
         expected_num_tries = 2
         with self.assertRaises(ConnectionErrorException) as cm:
             client = ParallelSSHClient(['127.0.0.1'], port=self.listen_port,
-                                       pkey=self.user_key, num_retries=expected_num_tries)
+                                       pkey=self.user_key,
+                                       num_retries=expected_num_tries)
             cmd = client.exec_command('blah')[0]
             cmd.get()
         num_tries = cm.exception.args[-1:][0]
@@ -273,3 +277,24 @@ def test_pssh_copy_file(self):
             os.unlink(filepath)
         del client
         server.join()
+
+    def test_pssh_pool_size(self):
+        """Test pool size logic"""
+        hosts = ['host-%01d' % d for d in xrange(5)]
+        client = ParallelSSHClient(hosts)
+        expected, actual = len(hosts), client.pool.size
+        self.assertEqual(expected, actual,
+                         msg="Expected pool size to be %s, got %s" % (
+                             expected, actual,))
+        hosts = ['host-%01d' % d for d in xrange(15)]
+        client = ParallelSSHClient(hosts)
+        expected, actual = client.pool_size, client.pool.size
+        self.assertEqual(expected, actual,
+                         msg="Expected pool size to be %s, got %s" % (
+                             expected, actual,))
+        hosts = ['host-%01d' % d for d in xrange(15)]
+        client = ParallelSSHClient(hosts, pool_size=len(hosts)+5)
+        expected, actual = len(hosts), client.pool.size
+        self.assertEqual(expected, actual,
+                         msg="Expected pool size to be %s, got %s" % (
+                             expected, actual,))