diff --git a/billiard/compat.py b/billiard/compat.py index b75455a..1324166 100644 --- a/billiard/compat.py +++ b/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): diff --git a/t/unit/test_compat.py b/t/unit/test_compat.py index 2def862..040553a 100644 --- a/t/unit/test_compat.py +++ b/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()