Bind the closerange-based close_open_fds() on Python 3

billiard/compat.py

@@ -1,4 +1,3 @@
-import errno
 import numbers
 import os
 import subprocess
@@ -118,39 +117,16 @@ def uniq(it):
     return (seen.add(obj) or obj for obj in it if obj not in seen)
 
 
-try:
-    closerange = os.closerange
-except AttributeError:
-
-    def closerange(fd_low, fd_high):  # noqa
-        for fd in reversed(range(fd_low, fd_high)):
-            try:
-                os.close(fd)
-            except OSError as exc:
-                if exc.errno != errno.EBADF:
-                    raise
-
-    def close_open_fds(keep=None):
-        # must make sure this is 0-inclusive (Issue #celery/1882)
-        keep = list(uniq(sorted(
-            f for f in map(maybe_fileno, keep or []) if f is not None
-        )))
-        maxfd = get_fdmax(default=2048)
-        kL, kH = iter([-1] + keep), iter(keep + [maxfd])
-        for low, high in zip_longest(kL, kH):
-            if low + 1 != high:
-                closerange(low + 1, high)
-else:
-    def close_open_fds(keep=None):  # noqa
-        keep = [maybe_fileno(f)
-                for f in (keep or []) if maybe_fileno(f) is not None]
-        for fd in reversed(range(get_fdmax(default=2048))):
-            if fd not in keep:
-                try:
-                    os.close(fd)
-                except OSError as exc:
-                    if exc.errno != errno.EBADF:
-                        raise
+def close_open_fds(keep=None):
+    # must make sure this is 0-inclusive (Issue #celery/1882)
+    keep = list(uniq(sorted(
+        f for f in map(maybe_fileno, keep or []) if f is not None
+    )))
+    maxfd = get_fdmax(default=2048)
+    kL, kH = iter([-1] + keep), iter(keep + [maxfd])
+    for low, high in zip_longest(kL, kH):
+        if low + 1 != high:
+            os.closerange(low + 1, high)
 
 
 def get_errno(exc):

t/unit/test_compat.py

@@ -0,0 +1,46 @@
+from unittest.mock import Mock, call, patch
+
+from billiard.compat import close_open_fds
+
+
+class test_close_open_fds:
+
+    def test_closes_ranges_between_kept_descriptors(self):
+        with patch('billiard.compat.get_fdmax', return_value=10), \
+                patch('os.closerange') as closerange, \
+                patch('os.close') as close:
+            close_open_fds([1, 2, 5])
+        assert closerange.call_args_list == [
+            call(0, 1), call(3, 5), call(6, 10),
+        ]
+        close.assert_not_called()
+
+    def test_accepts_file_objects_and_ignores_none(self):
+        fh = Mock(name='fh')
+        fh.fileno.return_value = 4
+        with patch('billiard.compat.get_fdmax', return_value=6), \
+                patch('os.closerange') as closerange, \
+                patch('os.close'):
+            close_open_fds([None, fh, 4, 0])
+        assert closerange.call_args_list == [call(1, 4), call(5, 6)]
+
+    def test_without_keep_closes_whole_range(self):
+        with patch('billiard.compat.get_fdmax', return_value=3), \
+                patch('os.closerange') as closerange, \
+                patch('os.close'):
+            close_open_fds()
+        closerange.assert_called_once_with(0, 3)
+
+    def test_cost_does_not_grow_with_fdmax(self):
+        """One closerange() call per gap, however large the limit is.
+
+        The previous body called os.close() once per descriptor up to
+        get_fdmax(), which stalls for minutes in containers where the limit
+        is ~1e9 (celery/celery#9886).
+        """
+        with patch('billiard.compat.get_fdmax', return_value=2 ** 20), \
+                patch('os.closerange') as closerange, \
+                patch('os.close') as close:
+            close_open_fds([0, 1, 2])
+        closerange.assert_called_once_with(3, 2 ** 20)
+        close.assert_not_called()