This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new 881a1bb49e [python] Keep RESTTokenFileIO usable in a forked child
(#10556)
881a1bb49e is described below
commit 881a1bb49ec136317c43c65c8d0dddcfffedd9e1
Author: Dapeng Sun(孙大鹏) <[email protected]>
AuthorDate: Wed Oct 14 21:49:44 2026 +0800
[python] Keep RESTTokenFileIO usable in a forked child (#10556)
---
.../pypaimon/catalog/rest/rest_token_file_io.py | 36 ++++++++++
.../pypaimon/tests/rest/rest_token_file_io_test.py | 80 ++++++++++++++++++++++
2 files changed, 116 insertions(+)
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 1f83adeb31..3d808e3f33 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 @@ class RESTTokenFileIO(FileIO):
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 @@ class RESTTokenFileIO(FileIO):
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 @@ class RESTTokenFileIO(FileIO):
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 832402f761..91cca6dc8a 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 @@ class RESTTokenFileIOTest(unittest.TestCase):
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
