[python] Restore the process-wide FileIO cache of RESTTokenFileIO

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."""

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
 

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."""