diff --git a/paimon-python/pypaimon/catalog/rest/rest_token_file_io.py b/paimon-python/pypaimon/catalog/rest/rest_token_file_io.py index 1f83adeb3110..3d808e3f3307 100644 --- a/paimon-python/pypaimon/catalog/rest/rest_token_file_io.py +++ b/paimon-python/pypaimon/catalog/rest/rest_token_file_io.py @@ -16,8 +16,10 @@ # under the License. import logging +import os import threading import time +import weakref from datetime import timedelta from typing import Optional, Union @@ -52,6 +54,9 @@ class RESTTokenFileIO(FileIO): _TOKEN_LOCKS: dict = {} _TOKEN_LOCKS_LOCK = threading.Lock() + # Objects whose lock a forked child resets. + _INSTANCES = weakref.WeakSet() + def __init__(self, identifier: Identifier, path: str, catalog_options: Optional[Union[dict, Options]] = None): self.identifier = identifier @@ -70,6 +75,7 @@ def __init__(self, identifier: Identifier, path: str, self._uri_reader_factory_cache: Optional[UriReaderFactory] = None self._file_io_scope: Optional[str] = None self._token_lock = threading.Lock() + self._register_for_fork() def __getstate__(self): state = self.__dict__.copy() @@ -88,6 +94,29 @@ def __setstate__(self, state): self.api_instance = None self._file_io_scope = None self._token_lock = threading.Lock() + self._register_for_fork() + + def _register_for_fork(self): + # The cache lock is held across fork, so a child never sees the set mid-update. + with self._FILE_IO_CACHE_LOCK: + RESTTokenFileIO._INSTANCES.add(self) + + @classmethod + def _before_fork(cls): + cls._FILE_IO_CACHE_LOCK.acquire() + + @classmethod + def _after_fork_in_parent(cls): + cls._FILE_IO_CACHE_LOCK.release() + + @classmethod + def _after_fork_in_child(cls): + # Parent threads that held these locks do not exist in the child. + cls._FILE_IO_CACHE_LOCK = threading.Lock() + cls._TOKEN_LOCKS_LOCK = threading.Lock() + cls._TOKEN_LOCKS = {} + for file_io in list(cls._INSTANCES): + file_io._token_lock = threading.Lock() def file_io(self) -> FileIO: return self._file_io_with_token()[0] @@ -346,3 +375,10 @@ def _existing_valid_token(self) -> Optional[RESTToken]: def close(self): pass + + +if hasattr(os, "register_at_fork"): + # Like logging: a fork waits for the cache lock, and the child starts with free locks. + os.register_at_fork(before=RESTTokenFileIO._before_fork, + after_in_parent=RESTTokenFileIO._after_fork_in_parent, + after_in_child=RESTTokenFileIO._after_fork_in_child) diff --git a/paimon-python/pypaimon/tests/rest/rest_token_file_io_test.py b/paimon-python/pypaimon/tests/rest/rest_token_file_io_test.py index 832402f761fc..91cca6dc8ac3 100644 --- a/paimon-python/pypaimon/tests/rest/rest_token_file_io_test.py +++ b/paimon-python/pypaimon/tests/rest/rest_token_file_io_test.py @@ -18,10 +18,12 @@ import gc import os import pickle +import signal import tempfile import threading import time import unittest +import warnings import weakref from datetime import timedelta from unittest.mock import patch, MagicMock @@ -221,6 +223,84 @@ def lookup(file_io): self.assertTrue(all(r is results[0] for r in results)) self.assertEqual(1, get_io.call_count) + @staticmethod + def _run_in_forked_child(action): + with warnings.catch_warnings(): + # Python 3.12+ warns about forking with other threads alive, which these tests do on purpose. + warnings.simplefilter("ignore", DeprecationWarning) + pid = os.fork() + if pid == 0: + code = 1 + try: + signal.alarm(10) + action() + code = 0 + finally: + os._exit(code) + return os.waitpid(pid, 0)[1] + + def _assert_child_succeeded(self, status): + self.assertFalse(os.WIFSIGNALED(status) and os.WTERMSIG(status) == signal.SIGALRM, + "the forked child blocked on a lock held by a parent thread") + self.assertTrue(os.WIFEXITED(status) and os.WEXITSTATUS(status) == 0, status) + + @unittest.skipUnless(hasattr(os, "register_at_fork"), "needs os.register_at_fork") + def test_fork_while_another_thread_creates_file_io(self): + token = self._valid_token("shared") + creating = threading.Event() + + def slow_file_io(path, options): + if path.endswith("/slow"): + creating.set() + time.sleep(0.3) + return _FakeFileIO(path, options) + + with patch.object(RESTTokenFileIO, 'try_to_refresh_token'), \ + patch(_FILE_IO_GET, side_effect=slow_file_io) as get_io: + other = self._file_io_with_token("oss://bucket/db.db/other", Options({}), token) + slow = self._file_io_with_token("oss://bucket/db.db/slow", Options({}), token) + + def in_child(): + other.file_io() + # The fork waited for the creation, so the child finds that FileIO cached. + slow.file_io() + assert get_io.call_count == 2, get_io.call_count + + creator = threading.Thread(target=slow.file_io) + creator.start() + self.assertTrue(creating.wait(5)) + status = self._run_in_forked_child(in_child) + creator.join() + self._assert_child_succeeded(status) + + @unittest.skipUnless(hasattr(os, "register_at_fork"), "needs os.register_at_fork") + def test_fork_while_another_thread_refreshes_token(self): + table_path = "oss://bucket/db.db/t" + refreshing, release = threading.Event(), threading.Event() + holders = [] + + def refresh(file_io): + if threading.current_thread() in holders: + refreshing.set() + release.wait(10) + file_io.token = self._valid_token("refreshed") + + self.addCleanup(RESTTokenFileIO._TOKEN_CACHE.pop, 'fork-test', None) + with patch.object(RESTTokenFileIO, '_build_cache_key', return_value='fork-test'), \ + patch.object(RESTTokenFileIO, 'refresh_token', autospec=True, side_effect=refresh), \ + patch(_FILE_IO_GET, side_effect=_FakeFileIO): + file_io = RESTTokenFileIO(self.identifier, table_path, Options({})) + holders.append(threading.Thread(target=file_io.file_io)) + holders[0].start() + self.assertTrue(refreshing.wait(5)) + # The child uses the object whose locks are held, and another object of that table. + status = self._run_in_forked_child(lambda: ( + file_io.file_io(), + RESTTokenFileIO(self.identifier, table_path, Options({})).file_io())) + release.set() + holders[0].join() + self._assert_child_succeeded(status) + def test_existing_valid_token_does_not_refresh_or_read_shared_cache(self): file_io = RESTTokenFileIO(self.identifier, self.warehouse_path) now = 1700000000