Don't continually ask the OS for a port, raise an exception (#13)

Don't continually ask the OS for a port, raise an exception instead of
returning None if no port could be found from `pick_unused_port()`.

On very busy machines, we might not be able to get a port from the OS
that is free on both TCP and UDP. Rather than spinning forever, only try
a handful of times.

2 files changed, 27 insertions, 6 deletions

diff --git a/src/portpicker.py b/src/portpicker.py
index c800161..cd92952 100644
--- a/src/portpicker.py
+++ b/src/portpicker.py
@@ -61,6 +61,11 @@ _owned_ports = set()
 _random_ports = set()
 
 
+class NoFreePortFoundException(Exception):
+    """Exception indicating that no free port could be found."""
+    pass
+
+
 def add_reserved_port(port):
     """Add a port that was acquired by means other than the port server."""
     _free_ports.add(port)
@@ -148,6 +153,9 @@ def pick_unused_port(pid=None, portserver_address=None):
 
     Returns:
       A port number that is unused on both TCP and UDP.
+
+    Raises:
+      NoFreePortFoundException: No free port could be found.
     """
     if _free_ports:
         port = _free_ports.pop()
@@ -177,7 +185,10 @@ def _pick_unused_port_without_server():  # Protected. pylint: disable=invalid-na
     should not be called by code outside of this module.
 
     Returns:
-      A port number that is unused on both TCP and UDP.  None on error.
+      A port number that is unused on both TCP and UDP.
+
+    Raises:
+      NoFreePortFoundException: No free port could be found.
     """
     # Try random ports first.
     rng = random.Random()
@@ -187,10 +198,10 @@ def _pick_unused_port_without_server():  # Protected. pylint: disable=invalid-na
             _random_ports.add(port)
             return port
 
-    # Try OS-assigned ports next.
+    # Next, try a few times to get an OS-assigned port.
     # Ambrose discovered that on the 2.6 kernel, calling Bind() on UDP socket
     # returns the same port over and over. So always try TCP first.
-    while True:
+    for _ in range(10):
         # Ask the OS for an unused port.
         port = bind(0, _PROTOS[0][0], _PROTOS[0][1])
         # Check if this port is unused on the other protocol.
@@ -198,6 +209,9 @@ def _pick_unused_port_without_server():  # Protected. pylint: disable=invalid-na
             _random_ports.add(port)
             return port
 
+    # Give up.
+    raise NoFreePortFoundException()
+
 
 def get_port_from_port_server(portserver_address, pid=None):
     """Request a free a port from a system-wide portserver.
diff --git a/src/tests/portpicker_test.py b/src/tests/portpicker_test.py
index b6ac959..8220ca0 100644
--- a/src/tests/portpicker_test.py
+++ b/src/tests/portpicker_test.py
@@ -201,10 +201,17 @@ class PickUnusedPortTest(unittest.TestCase):
                 return None
 
         with mock.patch.object(portpicker, 'bind', bind_with_error):
+            got_at_least_one_port = False
             for _ in range(100):
-                port = portpicker._pick_unused_port_without_server()
-                self.assertTrue(self.IsUnusedTCPPort(port))
-                self.assertTrue(self.IsUnusedUDPPort(port))
+                try:
+                    port = portpicker._pick_unused_port_without_server()
+                except portpicker.NoFreePortFoundException:
+                    continue
+                else:
+                    got_at_least_one_port = True
+                    self.assertTrue(self.IsUnusedTCPPort(port))
+                    self.assertTrue(self.IsUnusedUDPPort(port))
+            self.assertTrue(got_at_least_one_port)
 
     def testIsPortFree(self):
         """This might be flaky unless this test is run with a portserver."""