From d6cb197433bab86e51b922d75946b82d25b8f35d Mon Sep 17 00:00:00 2001 From: Sun Dapeng Date: Wed, 14 Oct 2026 11:17:35 +0800 Subject: [PATCH] [python] Restore the process-wide FileIO cache of RESTTokenFileIO #9715 gave each RESTTokenFileIO its own TTLCache(maxsize=1), so that tables in different buckets that receive the same token stop sharing one bucket-bound OssFileIO. With one entry per instance, a token refresh releases the previous FileIO while streams opened with the still valid previous credentials keep using it, which crashed Jindo readers before #10536, and every RESTTokenFileIO object builds its own FileIO and Jindo client. Share one cache across instances again (1000 entries, 10 hours, like Java's FILE_IO_CACHE). Key it by the token plus the table path and catalog options that FileIO.get builds the FileIO from, which keeps buckets apart, and expire entries 10 hours after their last use like Java's expireAfterAccess. Also align the continuation line flagged by flake8 (E127) in jindo_file_system_test.py, as #10550 does. --- .../catalog/rest/rest_token_file_io.py | 73 +++++--- .../filesystem/jindo_file_system_handler.py | 10 +- .../pypaimon/tests/jindo_file_system_test.py | 71 ++++++- .../tests/rest/rest_token_file_io_test.py | 174 +++++++++++++++++- 4 files changed, 283 insertions(+), 45 deletions(-) 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 861e3d3bb15c..1f83adeb3110 100644 --- a/paimon-python/pypaimon/catalog/rest/rest_token_file_io.py +++ b/paimon-python/pypaimon/catalog/rest/rest_token_file_io.py @@ -40,8 +40,14 @@ class RESTTokenFileIO(FileIO): A FileIO to support getting token from REST Server. """ + _FILE_IO_CACHE_MAXSIZE = 1000 _FILE_IO_CACHE_TTL = 36000 # 10 hours in seconds + # Process-wide like Java's FILE_IO_CACHE, so a token refresh keeps the FileIO + # that streams opened with the previous token still use. + _FILE_IO_CACHE = TTLCache(maxsize=_FILE_IO_CACHE_MAXSIZE, ttl=_FILE_IO_CACHE_TTL) + _FILE_IO_CACHE_LOCK = threading.Lock() + _TOKEN_CACHE: dict = {} _TOKEN_LOCKS: dict = {} _TOKEN_LOCKS_LOCK = threading.Lock() @@ -62,20 +68,16 @@ def __init__(self, identifier: Identifier, path: str, self.api_instance: Optional[RESTApi] = None self.log = logging.getLogger(__name__) self._uri_reader_factory_cache: Optional[UriReaderFactory] = None - self._init_file_io_cache() - - def _init_file_io_cache(self): - # FileIO is bound to this instance's path and catalog options. - self._file_io_cache = TTLCache(maxsize=1, ttl=self._FILE_IO_CACHE_TTL) - self._file_io_cache_lock = threading.Lock() + self._file_io_scope: Optional[str] = None + self._token_lock = threading.Lock() def __getstate__(self): state = self.__dict__.copy() # Remove non-serializable objects state.pop('api_instance', None) state.pop('_uri_reader_factory_cache', None) - state.pop('_file_io_cache', None) - state.pop('_file_io_cache_lock', None) + state.pop('_file_io_scope', None) + state.pop('_token_lock', None) # token can be serialized, but we'll refresh it on deserialization return state @@ -84,13 +86,14 @@ def __setstate__(self, state): self._uri_reader_factory_cache = None # api_instance will be recreated when needed self.api_instance = None - self._init_file_io_cache() + self._file_io_scope = None + self._token_lock = threading.Lock() def file_io(self) -> FileIO: return self._file_io_with_token()[0] def _file_io_with_token(self, minimum_validity_millis: int = 0): - with self._file_io_cache_lock: + with self._token_lock: if minimum_validity_millis: self.try_to_refresh_token(minimum_validity_millis) if not self._has_remaining_lifetime( @@ -101,29 +104,39 @@ def _file_io_with_token(self, minimum_validity_millis: int = 0): else: self.try_to_refresh_token() - if self.token is None: + token = self.token + if token is None: return (FileIO.get( self.path, self.catalog_options or Options({})), None) - cache_key = self.token - cache = self._file_io_cache - file_io = cache.get(cache_key) - if file_io is not None: - return file_io, cache_key - - merged_properties = RESTUtil.merge( - self.catalog_options.to_map() if self.catalog_options else {}, - cache_key.token - ) - if self.catalog_options: - dlf_oss_endpoint = self.catalog_options.get(CatalogOptions.DLF_OSS_ENDPOINT) - if dlf_oss_endpoint and dlf_oss_endpoint.strip(): - merged_properties[OssOptions.OSS_ENDPOINT.key()] = dlf_oss_endpoint - merged_options = Options(merged_properties) - - file_io = FileIO.get(self.path, merged_options) - cache[cache_key] = file_io - return file_io, cache_key + cache_key = (token, self._get_file_io_scope()) + with self._FILE_IO_CACHE_LOCK: + cache = self._FILE_IO_CACHE + file_io = cache.get(cache_key) + if file_io is None: + file_io = FileIO.get(self.path, self._merged_options(token)) + # Re-inserting on every access expires idle entries only, like expireAfterAccess. + cache[cache_key] = file_io + return file_io, token + + def _get_file_io_scope(self) -> str: + # FileIO.get binds its result to the path (e.g. the OSS bucket) and the catalog options. + if self._file_io_scope is None: + options = self.catalog_options.to_map() if self.catalog_options else {} + self._file_io_scope = repr((self.path, sorted( + (str(key), repr(value)) for key, value in options.items()))) + return self._file_io_scope + + def _merged_options(self, token: RESTToken) -> Options: + merged_properties = RESTUtil.merge( + self.catalog_options.to_map() if self.catalog_options else {}, + token.token + ) + if self.catalog_options: + dlf_oss_endpoint = self.catalog_options.get(CatalogOptions.DLF_OSS_ENDPOINT) + if dlf_oss_endpoint and dlf_oss_endpoint.strip(): + merged_properties[OssOptions.OSS_ENDPOINT.key()] = dlf_oss_endpoint + return Options(merged_properties) def _merge_token_with_catalog_options(self, token: dict) -> dict: """Merge token with catalog options, DLF OSS endpoint should override the standard OSS endpoint.""" diff --git a/paimon-python/pypaimon/filesystem/jindo_file_system_handler.py b/paimon-python/pypaimon/filesystem/jindo_file_system_handler.py index c44425e4dc24..2219497a8062 100644 --- a/paimon-python/pypaimon/filesystem/jindo_file_system_handler.py +++ b/paimon-python/pypaimon/filesystem/jindo_file_system_handler.py @@ -134,14 +134,8 @@ class JindoInputFile: def __init__(self, jindo_stream, file_size=None, *, owner): self._stream = jindo_stream self._file_size = file_size - # Keep-alive back-reference to the owning jindo filesystem: the - # stream is a non-owning reference into memory owned by it. Without - # it, dropping the FileIO that holds the handler (e.g. - # RESTTokenFileIO's TTLCache maxsize=1 eviction on token refresh) - # frees the native connection while other threads still read from - # this stream. Keyword-only and required so a new open_* path cannot - # silently skip it; it pins only the filesystem, not the handler, so - # the evicted FileIO's properties do not stay retained. + # The stream does not own its native store: pin the jindo filesystem (not the + # handler) so evicting the handler's FileIO from a cache cannot free it. self._owner = owner self._closed = False diff --git a/paimon-python/pypaimon/tests/jindo_file_system_test.py b/paimon-python/pypaimon/tests/jindo_file_system_test.py index e1ff6dcd1023..3b872003a942 100644 --- a/paimon-python/pypaimon/tests/jindo_file_system_test.py +++ b/paimon-python/pypaimon/tests/jindo_file_system_test.py @@ -20,6 +20,7 @@ import os import pickle import threading +import time import types import unittest import uuid @@ -29,10 +30,14 @@ import pyarrow.fs as pafs from pyarrow.fs import PyFileSystem +from pypaimon.catalog.rest.rest_token import RESTToken +from pypaimon.catalog.rest.rest_token_file_io import RESTTokenFileIO +from pypaimon.common.identifier import Identifier from pypaimon.common.options import Options from pypaimon.common.options.config import OssOptions from pypaimon.filesystem import jindo_file_system_handler as jindo_module from pypaimon.filesystem import oss_user_agent +from pypaimon.filesystem import pyarrow_file_io as pyarrow_file_io_module from pypaimon.filesystem.jindo_file_system_handler import ( JindoFileSystemHandler, JindoInputFile, @@ -387,9 +392,8 @@ def __del__(self): class JindoKeepAliveTest(unittest.TestCase): - """Streams must keep the native store alive across FileIO eviction - (RESTTokenFileIO token refresh) and must close before the store is - released, without relying on attribute dealloc order.""" + """Streams keep the native store alive across a FileIO cache eviction and close + before the store is released, without relying on attribute dealloc order.""" @staticmethod def _handler_with_recording_fs(events): @@ -404,7 +408,7 @@ def _handler_with_recording_fs(events): with mock.patch.object(jindo_module, "JINDO_AVAILABLE", True), \ mock.patch.object(jindo_module, "jfs", fake_jfs), \ mock.patch.object(jindo_module, "jutil", - types.SimpleNamespace(Config=_RecordingConfig)): + types.SimpleNamespace(Config=_RecordingConfig)): handler = JindoFileSystemHandler("oss://bucket/", options) return handler, weakref.ref(fake_fs), weakref.ref(handler) @@ -413,8 +417,8 @@ def test_evicted_filesystem_outlives_open_stream(self): handler, fs_ref, handler_ref = self._handler_with_recording_fs(events) input_file = handler.open_input_file('object') - # TTLCache eviction on token refresh drops the handler chain: the - # handler (and its credentials) must die, the native store must not. + # A FileIO cache eviction drops the handler chain: the handler (and its + # credentials) must die, the native store must not. del handler gc.collect() self.assertIsNone( @@ -488,6 +492,61 @@ def test_owner_argument_is_required(self): JindoOutputFile(_RecordingInputStream()) +class JindoTokenRefreshTest(unittest.TestCase): + """RESTTokenFileIO over the Jindo backend across a data token refresh.""" + + def setUp(self): + RESTTokenFileIO._FILE_IO_CACHE.clear() + self.addCleanup(RESTTokenFileIO._FILE_IO_CACHE.clear) + + @staticmethod + def _token(security_token, expire_at_millis): + return RESTToken({ + OssOptions.OSS_ENDPOINT.key(): "oss-cn-hangzhou.aliyuncs.com", + OssOptions.OSS_ACCESS_KEY_ID.key(): "STS.test-ak", + OssOptions.OSS_ACCESS_KEY_SECRET.key(): "test-sk", + OssOptions.OSS_SECURITY_TOKEN.key(): security_token, + }, expire_at_millis) + + def test_token_refresh_keeps_previous_jindo_handler(self): + events = [] + connected = [] + + def connect(*args, **kwargs): + fs = _EvictionRecordingFileSystem(events) + connected.append(weakref.ref(fs)) + return fs + + table_path = "oss://bucket/db.db/t" + expire_at_millis = int(time.time() * 1000) + 7_200_000 + with mock.patch.object(jindo_module, "JINDO_AVAILABLE", True), \ + mock.patch.object(pyarrow_file_io_module, "JINDO_AVAILABLE", True), \ + mock.patch.object(jindo_module, "jfs", types.SimpleNamespace(connect=connect)), \ + mock.patch.object(jindo_module, "jutil", + types.SimpleNamespace(Config=_RecordingConfig)), \ + mock.patch.object(RESTTokenFileIO, "try_to_refresh_token"): + readers = [RESTTokenFileIO(Identifier.create("db", "t"), table_path, + Options({OssOptions.OSS_IMPL.key(): "jindo"})) + for _ in range(2)] + for reader in readers: + reader.token = self._token("sts-1", expire_at_millis) + stream = readers[0].new_input_stream(table_path + "/bucket-0/data.blob") + handler_ref = weakref.ref(readers[0].filesystem.handler) + self.assertIs(readers[0].file_io(), readers[1].file_io()) + self.assertEqual(1, len(connected)) + + for reader in readers: + reader.token = self._token("sts-2", expire_at_millis + 1) + self.assertIs(readers[0].file_io(), readers[1].file_io()) + gc.collect() + + self.assertEqual(2, len(connected)) + self.assertIsNotNone(handler_ref(), "token refresh released the previous handler") + self.assertEqual(b"x", stream.read(1)) + stream.close() + self.assertEqual(["stream-closed"], events) + + class JindoFileSystemTest(unittest.TestCase): """Test cases for JindoFileSystem.""" 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 8255a78ea977..832402f761fc 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 @@ -15,14 +15,19 @@ # specific language governing permissions and limitations # under the License. +import gc import os import pickle import tempfile +import threading import time import unittest +import weakref from datetime import timedelta from unittest.mock import patch, MagicMock +from cachetools import TTLCache + from pypaimon.catalog.rest.rest_token_file_io import RESTTokenFileIO from pypaimon.catalog.rest.rest_token import RESTToken @@ -32,6 +37,33 @@ from pypaimon.filesystem.local_file_io import LocalFileIO from pypaimon.table.row.blob import BlobDescriptor +_FILE_IO_GET = 'pypaimon.catalog.rest.rest_token_file_io.FileIO.get' + + +class _NonOwningStream: + """Like a pyjindo stream: does not keep the FileIO that opened it alive.""" + + def __init__(self, file_io): + self._file_io = weakref.ref(file_io) + + def read(self): + if self._file_io() is None: + raise RuntimeError('use-after-free: the FileIO of this stream was released') + return b'x' + + +class _FakeFileIO: + + def __init__(self, path, options): + self.path = path + self.options = options + + def new_input_stream(self, path): + return _NonOwningStream(self) + + def exists(self, path): + return True + class RESTTokenFileIOTest(unittest.TestCase): """Test cases for RESTTokenFileIO.""" @@ -42,13 +74,153 @@ def setUp(self): self.warehouse_path = f"file://{self.temp_dir}" self.identifier = Identifier.from_string("default.test_table") self.catalog_options = Options({}) + RESTTokenFileIO._FILE_IO_CACHE.clear() def tearDown(self): """Clean up test fixtures.""" import shutil + RESTTokenFileIO._FILE_IO_CACHE.clear() if os.path.exists(self.temp_dir): shutil.rmtree(self.temp_dir) + @staticmethod + def _valid_token(name): + return RESTToken({OssOptions.OSS_SECURITY_TOKEN.key(): name}, + int(time.time() * 1000) + 7_200_000) + + def _file_io_with_token(self, path, options, token): + file_io = RESTTokenFileIO(self.identifier, path, options) + file_io.token = token + return file_io + + def test_token_refresh_keeps_file_io_of_open_streams(self): + table_path = "oss://bucket/db.db/table" + file_io = RESTTokenFileIO(self.identifier, table_path, self.catalog_options) + self.addCleanup(RESTTokenFileIO._TOKEN_CACHE.pop, 'refresh-test', None) + now = [1700000000] + issued = [] + + def refresh(): + issued.append(RESTToken( + {OssOptions.OSS_SECURITY_TOKEN.key(): "sts-{}".format(len(issued))}, + (now[0] + 2 * 3600) * 1000)) + file_io.token = issued[-1] + + with patch('pypaimon.catalog.rest.rest_token_file_io.time.time', + side_effect=lambda: now[0]), \ + patch.object(file_io, '_build_cache_key', return_value='refresh-test'), \ + patch.object(file_io, 'refresh_token', side_effect=refresh), \ + patch(_FILE_IO_GET, side_effect=_FakeFileIO): + stream = file_io.new_input_stream(table_path + "/bucket-0/data.blob") + # Less than an hour left: the next call refreshes the token. + now[0] += 3600 + 1 + self.assertTrue(file_io.exists(table_path + "/snapshot/LATEST")) + gc.collect() + + self.assertEqual(2, len(issued)) + self.assertEqual(b'x', stream.read()) + + def test_file_io_cache_separates_paths_and_options_with_same_token(self): + token = self._valid_token("shared") + legacy = Options({OssOptions.OSS_IMPL.key(): "legacy"}) + with patch.object(RESTTokenFileIO, 'try_to_refresh_token'), \ + patch(_FILE_IO_GET, side_effect=_FakeFileIO): + bucket_a = self._file_io_with_token( + "oss://bucket-a/db.db/t", Options({}), token).file_io() + bucket_b = self._file_io_with_token( + "oss://bucket-b/db.db/t", Options({}), token).file_io() + bucket_a_legacy = self._file_io_with_token( + "oss://bucket-a/db.db/t", legacy, token).file_io() + + self.assertEqual("oss://bucket-a/db.db/t", bucket_a.path) + self.assertEqual("oss://bucket-b/db.db/t", bucket_b.path) + self.assertIsNot(bucket_a, bucket_a_legacy) + self.assertEqual("legacy", bucket_a_legacy.options.get(OssOptions.OSS_IMPL)) + + def test_same_token_on_other_bucket_keeps_its_own_oss_bucket(self): + token = RESTToken({ + OssOptions.OSS_ENDPOINT.key(): "oss-cn-hangzhou.aliyuncs.com", + OssOptions.OSS_ACCESS_KEY_ID.key(): "ak", + OssOptions.OSS_ACCESS_KEY_SECRET.key(): "sk", + }, int(time.time() * 1000) + 7_200_000) + options = Options({OssOptions.OSS_IMPL.key(): "legacy"}) + with patch.object(RESTTokenFileIO, 'try_to_refresh_token'): + first = self._file_io_with_token("oss://bucket-a/db.db/t", options, token) + second = self._file_io_with_token("oss://bucket-b/db.db/t", options, token) + self.assertEqual("bucket-a", first.file_io()._oss_bucket) + self.assertEqual("bucket-b", second.file_io()._oss_bucket) + + def test_instances_of_one_table_share_file_io(self): + token = self._valid_token("shared") + table_path = "oss://bucket/db.db/t" + with patch.object(RESTTokenFileIO, 'try_to_refresh_token'), \ + patch(_FILE_IO_GET, side_effect=_FakeFileIO) as get_io: + first = self._file_io_with_token(table_path, Options({"k": "v"}), token) + second = self._file_io_with_token(table_path, Options({"k": "v"}), token) + restored = pickle.loads(pickle.dumps(first)) + + self.assertIs(first.file_io(), second.file_io()) + self.assertIs(first.file_io(), restored.file_io()) + self.assertEqual(1, get_io.call_count) + + def test_file_io_cache_expires_idle_entries_only(self): + clock = [0] + hour = 3600 + cache = TTLCache(maxsize=RESTTokenFileIO._FILE_IO_CACHE_MAXSIZE, + ttl=RESTTokenFileIO._FILE_IO_CACHE_TTL, timer=lambda: clock[0]) + with patch.object(RESTTokenFileIO, '_FILE_IO_CACHE', cache), \ + patch.object(RESTTokenFileIO, 'try_to_refresh_token'), \ + patch(_FILE_IO_GET, side_effect=_FakeFileIO): + file_io = self._file_io_with_token( + "oss://bucket/db.db/t", self.catalog_options, self._valid_token("first")) + first = file_io.file_io() + for _ in range(3): + clock[0] += 9 * hour + self.assertIs(first, file_io.file_io()) + + # The previous token's FileIO stays cached until it has been idle for the TTL. + file_io.token = self._valid_token("second") + second = file_io.file_io() + self.assertIsNot(first, second) + first_ref = weakref.ref(first) + del first + clock[0] += 9 * hour + file_io.file_io() + gc.collect() + self.assertIsNotNone(first_ref()) + clock[0] += 2 * hour + self.assertIs(second, file_io.file_io()) + gc.collect() + self.assertIsNone(first_ref()) + + def test_concurrent_lookups_create_one_file_io(self): + token = self._valid_token("shared") + threads_count = 8 + barrier = threading.Barrier(threads_count) + results = [] + + def slow_file_io(path, options): + time.sleep(0.05) + return _FakeFileIO(path, options) + + def lookup(file_io): + barrier.wait() + results.append(file_io.file_io()) + + with patch.object(RESTTokenFileIO, 'try_to_refresh_token'), \ + patch(_FILE_IO_GET, side_effect=slow_file_io) as get_io: + instances = [self._file_io_with_token("oss://bucket/db.db/t", Options({}), token) + for _ in range(threads_count)] + threads = [threading.Thread(target=lookup, args=(i,)) for i in instances] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + self.assertEqual(threads_count, len(results)) + self.assertTrue(all(r is results[0] for r in results)) + self.assertEqual(1, get_io.call_count) + def test_existing_valid_token_does_not_refresh_or_read_shared_cache(self): file_io = RESTTokenFileIO(self.identifier, self.warehouse_path) now = 1700000000 @@ -296,7 +468,7 @@ def test_pickle_serialization(self): self.assertEqual(deserialized_file_io.properties.data, original_file_io.properties.data) self.assertIsNone(deserialized_file_io.api_instance) - self.assertIsNot(deserialized_file_io.file_io(), original_backend) + self.assertIs(deserialized_file_io.file_io(), original_backend) test_file_path = f"file://{self.temp_dir}/pickle_test.txt" test_content = b"pickle test content"