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 c6f958b5ff [python] Drop expired tokens from the token cache of
RESTTokenFileIO (#10573)
c6f958b5ff is described below
commit c6f958b5ff90adc44f1ea419cafb10aa04bd553e
Author: Dapeng Sun(孙大鹏) <[email protected]>
AuthorDate: Sun Oct 18 16:49:55 2026 +0800
[python] Drop expired tokens from the token cache of RESTTokenFileIO
(#10573)
---
.../pypaimon/catalog/rest/rest_token_file_io.py | 22 +++++++
.../pypaimon/tests/rest/rest_token_file_io_test.py | 75 ++++++++++++++++++++++
2 files changed, 97 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 3d808e3f33..fbb39480ed 100644
--- a/paimon-python/pypaimon/catalog/rest/rest_token_file_io.py
+++ b/paimon-python/pypaimon/catalog/rest/rest_token_file_io.py
@@ -53,6 +53,8 @@ class RESTTokenFileIO(FileIO):
_TOKEN_CACHE: dict = {}
_TOKEN_LOCKS: dict = {}
_TOKEN_LOCKS_LOCK = threading.Lock()
+ # Size of either dict at which expired tokens and idle locks are first
dropped; doubles with the live set.
+ _TOKEN_CACHE_PRUNE_SIZE = 1024
# Objects whose lock a forked child resets.
_INSTANCES = weakref.WeakSet()
@@ -318,6 +320,24 @@ class RESTTokenFileIO(FileIO):
def _set_cached_token(self, identifier_str: str, token: RESTToken):
with self._TOKEN_LOCKS_LOCK:
self._TOKEN_CACHE[identifier_str] = token
+ self._prune_if_large()
+
+ @classmethod
+ def _prune_if_large(cls):
+ if max(len(cls._TOKEN_CACHE), len(cls._TOKEN_LOCKS)) >=
RESTTokenFileIO._TOKEN_CACHE_PRUNE_SIZE:
+ cls._prune_expired_tokens()
+
+ @classmethod
+ def _prune_expired_tokens(cls):
+ # Keys embed the caller's identity, which changes whenever STS
credentials rotate.
+ now = int(time.time() * 1000)
+ for key, token in list(cls._TOKEN_CACHE.items()):
+ if token.expire_at_millis <= now:
+ del cls._TOKEN_CACHE[key]
+ for key, lock in list(cls._TOKEN_LOCKS.items()):
+ if key not in cls._TOKEN_CACHE and not lock.locked():
+ del cls._TOKEN_LOCKS[key]
+ RESTTokenFileIO._TOKEN_CACHE_PRUNE_SIZE = max(1024, 2 *
len(cls._TOKEN_CACHE), 2 * len(cls._TOKEN_LOCKS))
def _is_token_expired(
self, token: Optional[RESTToken], minimum_validity_millis: int =
0) -> bool:
@@ -337,6 +357,8 @@ class RESTTokenFileIO(FileIO):
def _get_global_token_lock(self, identifier_str: str) -> threading.Lock:
with self._TOKEN_LOCKS_LOCK:
if identifier_str not in self._TOKEN_LOCKS:
+ # Refreshes that keep failing add locks but no tokens.
+ self._prune_if_large()
self._TOKEN_LOCKS[identifier_str] = threading.Lock()
return self._TOKEN_LOCKS[identifier_str]
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 91cca6dc8a..a7b88220d0 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
@@ -869,6 +869,81 @@ class RESTTokenFileIOTest(unittest.TestCase):
self.assertIsNotNone(cached, "file_io_2 should see file_io_1's cached
token")
self.assertEqual(cached.token.get("fs.oss.accessKeyId"),
"shared-data-ak")
+ def test_token_cache_drops_expired_tokens_of_rotated_identities(self):
+ now = int(time.time() * 1000)
+ live = RESTToken({"fs.oss.accessKeyId": "live"}, now + 7_200_000)
+ expired = RESTToken({"fs.oss.accessKeyId": "expired"}, now - 1)
+ file_io = RESTTokenFileIO(self.identifier, self.warehouse_path,
self.catalog_options)
+ with patch.object(RESTTokenFileIO, '_TOKEN_CACHE', {}), \
+ patch.object(RESTTokenFileIO, '_TOKEN_LOCKS', {}), \
+ patch.object(RESTTokenFileIO, '_TOKEN_CACHE_PRUNE_SIZE', 4):
+ for sts_ak in ("STS.1", "STS.2"):
+ key = sts_ak + ":" + str(self.identifier)
+ file_io._get_global_token_lock(key)
+ file_io._set_cached_token(key, expired)
+ busy = file_io._get_global_token_lock("STS.3:" +
str(self.identifier))
+ busy.acquire()
+ self.addCleanup(busy.release)
+ file_io._set_cached_token("STS.3:other", live)
+ self.assertEqual(3, len(RESTTokenFileIO._TOKEN_CACHE))
+
+ file_io._set_cached_token("STS.4:" + str(self.identifier), live)
+
+ self.assertEqual(["STS.3:other", "STS.4:" + str(self.identifier)],
+ sorted(RESTTokenFileIO._TOKEN_CACHE))
+ # A lock that a refresh still holds stays; idle locks of dropped
tokens go.
+ self.assertEqual(["STS.3:" + str(self.identifier)],
list(RESTTokenFileIO._TOKEN_LOCKS))
+ self.assertEqual(1024, RESTTokenFileIO._TOKEN_CACHE_PRUNE_SIZE)
+
+ def test_token_cache_prune_size_doubles_with_live_tokens(self):
+ now = int(time.time() * 1000)
+ live = RESTToken({"fs.oss.accessKeyId": "live"}, now + 7_200_000)
+ expired = RESTToken({"fs.oss.accessKeyId": "expired"}, now - 1)
+ cache = {"live-%d" % i: live for i in range(1023)}
+ cache["expired-0"] = expired
+ file_io = RESTTokenFileIO(self.identifier, self.warehouse_path,
self.catalog_options)
+ with patch.object(RESTTokenFileIO, '_TOKEN_CACHE', cache), \
+ patch.object(RESTTokenFileIO, '_TOKEN_LOCKS', {}), \
+ patch.object(RESTTokenFileIO, '_TOKEN_CACHE_PRUNE_SIZE', 1024):
+ file_io._set_cached_token("live-1023", live)
+ self.assertNotIn("expired-0", cache)
+ self.assertEqual(2048, RESTTokenFileIO._TOKEN_CACHE_PRUNE_SIZE)
+
+ # Below the doubled size, caching another token scans nothing.
+ file_io._set_cached_token("expired-1", expired)
+ self.assertIn("expired-1", cache)
+ self.assertEqual(1025, len(cache))
+
+ def test_token_cache_prune_size_doubles_with_held_locks(self):
+ held = {"refreshing-%d" % i: threading.Lock() for i in range(1100)}
+ for lock in held.values():
+ lock.acquire()
+ file_io = RESTTokenFileIO(self.identifier, self.warehouse_path,
self.catalog_options)
+ with patch.object(RESTTokenFileIO, '_TOKEN_CACHE', {}), \
+ patch.object(RESTTokenFileIO, '_TOKEN_LOCKS', dict(held)), \
+ patch.object(RESTTokenFileIO, '_TOKEN_CACHE_PRUNE_SIZE', 1024):
+ file_io._get_global_token_lock("new")
+ # Locks of refreshes in flight stay, so the next prune waits until
the lock count doubles.
+ self.assertEqual(1101, len(RESTTokenFileIO._TOKEN_LOCKS))
+ self.assertEqual(2200, RESTTokenFileIO._TOKEN_CACHE_PRUNE_SIZE)
+
+ def test_failed_refreshes_drop_idle_locks(self):
+ keys = ["STS.%d:%s" % (i, self.identifier) for i in range(6)]
+ file_io = RESTTokenFileIO(self.identifier, self.warehouse_path,
self.catalog_options)
+ with patch.object(RESTTokenFileIO, '_TOKEN_CACHE', {}), \
+ patch.object(RESTTokenFileIO, '_TOKEN_LOCKS', {}), \
+ patch.object(RESTTokenFileIO, '_TOKEN_CACHE_PRUNE_SIZE', 4), \
+ patch.object(RESTTokenFileIO, '_build_cache_key',
side_effect=keys), \
+ patch.object(RESTTokenFileIO, 'refresh_token',
side_effect=RuntimeError("forbidden")):
+ for _ in keys:
+ with self.assertRaises(RuntimeError):
+ file_io.try_to_refresh_token()
+
+ # The fifth lock found four idle ones without tokens and dropped
them first.
+ self.assertEqual(keys[4:], sorted(RESTTokenFileIO._TOKEN_LOCKS))
+ self.assertEqual({}, RESTTokenFileIO._TOKEN_CACHE)
+ self.assertEqual(1024, RESTTokenFileIO._TOKEN_CACHE_PRUNE_SIZE)
+
def test_refresh_token_strips_system_table_suffix(self):
"""refresh_token() strips $snapshots suffix before requesting token."""
system_identifier = Identifier.create("db", "my_table",
system_table="snapshots")
