[python] Keep RESTTokenFileIO usable in a forked child

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)

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