List the fd directory in close_open_fds() instead of walking the limit

billiard/compat.py

@@ -117,11 +117,62 @@ def uniq(it):
     return (seen.add(obj) or obj for obj in it if obj not in seen)
 
 
+# Directory listing the descriptors open in the calling process.  Same
+# choice as CPython's FD_DIR (Modules/_posixsubprocess.c, Cygwin added in
+# gh-148575): /dev/fd on macOS, Cygwin, FreeBSD and DragonFly, /proc/self/fd
+# elsewhere.  On FreeBSD and DragonFly it is only trusted when fdescfs is
+# mounted, see _open_fds().
+_BSD_WITH_FDESCFS = ('freebsd', 'dragonfly')
+if sys.platform.startswith(('cygwin', 'darwin') + _BSD_WITH_FDESCFS):
+    _FD_DIR = '/dev/fd'
+else:
+    _FD_DIR = '/proc/self/fd'
+
+
+def _dev_fd_is_fdescfs():
+    # devfs alone creates only /dev/fd/0-2, while fdescfs creates entries for
+    # every descriptor the process has open.  Same check as CPython's
+    # _is_fdescfs_mounted_on_dev_fd().
+    try:
+        return os.stat('/dev').st_dev != os.stat(_FD_DIR).st_dev
+    except OSError:
+        return False
+
+
+def _open_fds():
+    """Return the descriptors open in this process, or None.
+
+    None means this platform has no directory listing them, and the caller
+    has to fall back to a numeric range.
+    """
+    if sys.platform.startswith(_BSD_WITH_FDESCFS) and not _dev_fd_is_fdescfs():
+        return None
+    try:
+        names = os.listdir(_FD_DIR)
+    except OSError:
+        return None
+    return sorted(int(name) for name in names if name.isdigit())
+
+
 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
     )))
+    fds = _open_fds()
+    if fds is not None:
+        # Only the descriptors actually open are touched, so the cost does
+        # not depend on RLIMIT_NOFILE (celery/celery#9886).
+        keep = set(keep)
+        for fd in fds:
+            if fd in keep:
+                continue
+            try:
+                os.close(fd)
+            except OSError:
+                # Same as os.closerange() below: closing is best effort.
+                pass
+        return
     maxfd = get_fdmax(default=2048)
     kL, kH = iter([-1] + keep), iter(keep + [maxfd])
     for low, high in zip_longest(kL, kH):

t/unit/test_compat.py

@@ -1,12 +1,61 @@
+import errno
+import sys
+import tempfile
 from unittest.mock import Mock, call, patch
 
+import pytest
+
+from billiard import compat
 from billiard.compat import close_open_fds
 
 
 class test_close_open_fds:
 
+    def test_closes_only_listed_descriptors(self):
+        with patch('billiard.compat._open_fds',
+                   return_value=[0, 1, 2, 7, 9]), \
+                patch('os.close') as close, \
+                patch('os.closerange') as closerange:
+            close_open_fds([1, 2])
+        assert close.call_args_list == [call(0), call(7), call(9)]
+        closerange.assert_not_called()
+
+    def test_accepts_file_objects_and_ignores_none(self):
+        fh = Mock(name='fh')
+        fh.fileno.return_value = 3
+        with patch('billiard.compat._open_fds', return_value=[3, 4]), \
+                patch('os.close') as close, \
+                patch('os.closerange'):
+            close_open_fds([None, fh])
+        close.assert_called_once_with(4)
+
+    @pytest.mark.parametrize('code', [errno.EBADF, errno.EINTR, errno.EIO])
+    def test_ignores_close_errors_like_closerange(self, code):
+        exc = OSError()
+        exc.errno = code
+        with patch('billiard.compat._open_fds', return_value=[5, 6]), \
+                patch('os.close', side_effect=exc) as close, \
+                patch('os.closerange'):
+            close_open_fds()
+        assert close.call_args_list == [call(5), call(6)]
+
+    def test_never_calls_get_fdmax_with_fd_dir(self):
+        """The limit is irrelevant when the open descriptors can be listed."""
+        with patch('billiard.compat._open_fds', return_value=[0, 1, 2]), \
+                patch('billiard.compat.get_fdmax') as get_fdmax, \
+                patch('os.close'), \
+                patch('os.closerange') as closerange:
+            close_open_fds([0, 1, 2])
+        get_fdmax.assert_not_called()
+        closerange.assert_not_called()
+
+
+class test_close_open_fds_fallback:
+    """Without an fd directory the gaps between kept descriptors are closed."""
+
     def test_closes_ranges_between_kept_descriptors(self):
-        with patch('billiard.compat.get_fdmax', return_value=10), \
+        with patch('billiard.compat._open_fds', return_value=None), \
+                patch('billiard.compat.get_fdmax', return_value=10), \
                 patch('os.closerange') as closerange, \
                 patch('os.close') as close:
             close_open_fds([1, 2, 5])
@@ -18,14 +67,16 @@ def test_closes_ranges_between_kept_descriptors(self):
     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), \
+        with patch('billiard.compat._open_fds', return_value=None), \
+                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), \
+        with patch('billiard.compat._open_fds', return_value=None), \
+                patch('billiard.compat.get_fdmax', return_value=3), \
                 patch('os.closerange') as closerange, \
                 patch('os.close'):
             close_open_fds()
@@ -34,13 +85,54 @@ def test_without_keep_closes_whole_range(self):
     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).
+        The body before celery/billiard#455 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), \
+        with patch('billiard.compat._open_fds', return_value=None), \
+                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()
+
+
+class test_open_fds:
+
+    def test_fd_dir_matches_this_platform(self):
+        if sys.platform == 'darwin':
+            assert compat._FD_DIR == '/dev/fd'
+        elif sys.platform.startswith(('cygwin', 'freebsd', 'dragonfly')):
+            assert compat._FD_DIR == '/dev/fd'
+        else:
+            assert compat._FD_DIR == '/proc/self/fd'
+
+    @pytest.mark.skipif(sys.platform == 'win32', reason='no fd directory')
+    def test_lists_a_descriptor_this_process_opened(self):
+        with tempfile.TemporaryFile() as fh:
+            fds = compat._open_fds()
+            if fds is None:
+                pytest.skip(
+                    f'fd directory {compat._FD_DIR!r} not available here')
+            assert fh.fileno() in fds
+
+    def test_returns_none_when_fd_dir_is_missing(self):
+        with patch('os.listdir', side_effect=OSError()):
+            assert compat._open_fds() is None
+
+    def test_returns_none_on_freebsd_without_fdescfs(self):
+        with patch.object(sys, 'platform', 'freebsd14'), \
+                patch('billiard.compat._dev_fd_is_fdescfs',
+                      return_value=False), \
+                patch('os.listdir') as listdir:
+            assert compat._open_fds() is None
+        listdir.assert_not_called()
+
+    def test_dev_fd_is_fdescfs(self):
+        with patch('os.stat', side_effect=[Mock(st_dev=1), Mock(st_dev=2)]):
+            assert compat._dev_fd_is_fdescfs()
+        with patch('os.stat', side_effect=[Mock(st_dev=1), Mock(st_dev=1)]):
+            assert not compat._dev_fd_is_fdescfs()
+        with patch('os.stat', side_effect=OSError()):
+            assert not compat._dev_fd_is_fdescfs()