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 8e9fee64f2 [python] Restore the process-wide FileIO cache of
RESTTokenFileIO (#10552)
8e9fee64f2 is described below
commit 8e9fee64f26e7138a2b2e64f499844e88c4b9b09
Author: Dapeng Sun(孙大鹏) <[email protected]>
AuthorDate: Wed Oct 14 13:18:07 2026 +0800
[python] Restore the process-wide FileIO cache of RESTTokenFileIO (#10552)
---
.../pypaimon/catalog/rest/rest_token_file_io.py | 73 +++++----
.../filesystem/jindo_file_system_handler.py | 10 +-
.../pypaimon/tests/jindo_file_system_test.py | 71 ++++++++-
.../pypaimon/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 861e3d3bb1..1f83adeb31 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 @@ class RESTTokenFileIO(FileIO):
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 @@ class RESTTokenFileIO(FileIO):
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 @@ class RESTTokenFileIO(FileIO):
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 c44425e4dc..2219497a80 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 e1ff6dcd10..3b872003a9 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 io
import os
import pickle
import threading
+import time
import types
import unittest
import uuid
@@ -29,10 +30,14 @@ from unittest import mock
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 @@ class _EvictionRecordingNativeStream:
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 @@ class JindoKeepAliveTest(unittest.TestCase):
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 @@ class JindoKeepAliveTest(unittest.TestCase):
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 @@ class JindoKeepAliveTest(unittest.TestCase):
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 8255a78ea9..832402f761 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.common.options.config import CatalogOptions,
OssOptions
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 @@ class RESTTokenFileIOTest(unittest.TestCase):
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 @@ class RESTTokenFileIOTest(unittest.TestCase):
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"
