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"