JingsongLi commented on code in PR #10178: URL: https://github.com/apache/paimon/pull/10178#discussion_r4253798440
########## paimon-common/src/main/java/org/apache/paimon/rest/RESTTokenRefresher.java: ########## @@ -0,0 +1,268 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.paimon.rest; + +import org.apache.paimon.annotation.VisibleForTesting; +import org.apache.paimon.catalog.Identifier; +import org.apache.paimon.options.Options; +import org.apache.paimon.rest.responses.GetTableTokenResponse; + +import org.apache.paimon.shade.guava30.com.google.common.collect.ImmutableMap; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import javax.annotation.Nullable; + +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.concurrent.locks.ReentrantLock; +import java.util.function.LongSupplier; + +import static org.apache.paimon.rest.RESTApi.TOKEN_EXPIRATION_SAFE_TIME_MILLIS; +import static org.apache.paimon.rest.RESTCatalogOptions.DLF_OSS_ENDPOINT; +import static org.apache.paimon.rest.RESTCatalogOptions.IO_CACHE_ENABLED; + +/** + * Loads the data token of a table from the REST server and reloads it before it expires, for {@link + * RESTTokenFileIO} and for storage credentials providers created from its options. + */ +public class RESTTokenRefresher { + + private static final Logger LOG = LoggerFactory.getLogger(RESTTokenRefresher.class); + + /** Database of the table whose token is refreshed, set by {@link RESTTokenFileIO}. */ + public static final String DATABASE = "data-token.database"; + + /** Object name of the table whose token is refreshed, including any branch. */ + public static final String OBJECT = "data-token.object"; + + /** Expiration of the token already present in the options. */ + public static final String EXPIRES_AT_MILLIS = "data-token.expires-at-millis"; + + /** Keys of the token already present in the options, replaced as a whole on reload. */ + public static final String TOKEN_KEYS = "data-token.keys"; + + // A reloaded token is kept at least this long, and a failed reload is retried after it. + static final long RELOAD_INTERVAL_MILLIS = 10_000L; + + // Shifts the clock of refreshers that use the system clock, for tests only. + @VisibleForTesting static volatile long clockOffsetMillis; + + private final Options catalogOptions; + private final Identifier identifier; + private final LongSupplier clock; + private final ReentrantLock lock = new ReentrantLock(); + + @Nullable private volatile CachedToken cached; + + // Guarded by lock. + @Nullable private RESTApi api; + private long nextAttemptMillis; + @Nullable private RuntimeException lastFailure; + + RESTTokenRefresher( + Options catalogOptions, + Identifier identifier, + @Nullable RESTApi api, + @Nullable RESTToken token, + LongSupplier clock) { + this.catalogOptions = catalogOptions; + this.identifier = identifier; + this.api = api; + this.clock = clock; + this.cached = token == null ? null : new CachedToken(token, currentTimeMillis()); + } + + /** Describes the table and the token merged into the options, so a refresher can be created. */ + public static void configure(Options options, Identifier identifier, RESTToken token) { + options.set(DATABASE, identifier.getDatabaseName()); + options.set(OBJECT, identifier.getObjectName()); + options.set(EXPIRES_AT_MILLIS, String.valueOf(token.expireAtMillis())); + options.set(TOKEN_KEYS, String.join(",", token.token().keySet())); + } + + /** Whether the options name a table, see {@link #configure}. */ + public static boolean isConfigured(Options options) { + return options.containsKey(DATABASE) && options.containsKey(OBJECT); + } + + /** + * Creates a refresher from catalog options that also carry the current token and the keys set + * by {@link #configure}. + */ + public static RESTTokenRefresher fromOptions(Options options) { + return fromOptions(options, null); + } + + @VisibleForTesting + static RESTTokenRefresher fromOptions(Options options, @Nullable RESTApi api) { + Identifier identifier = new Identifier(options.get(DATABASE), options.get(OBJECT)); + String expiresAt = options.get(EXPIRES_AT_MILLIS); + String tokenKeys = options.get(TOKEN_KEYS); + RESTToken token = null; + if (expiresAt != null && tokenKeys != null) { + Map<String, String> values = new LinkedHashMap<>(); + for (String key : tokenKeys.split(",")) { + String value = options.get(key); + if (value != null) { + values.put(key, value); + } + } + token = new RESTToken(values, Long.parseLong(expiresAt)); + } + return new RESTTokenRefresher( + options, identifier, api, token, RESTTokenRefresher::systemTimeMillis); + } + + /** Returns the current token, reloading it from the catalog when it is about to expire. */ + public RESTToken token() { + CachedToken current = cached; + long now = currentTimeMillis(); + if (current != null && now < current.reloadAtMillis) { + return current.token; + } + // While the token is still valid, one caller reloads it and the others keep using it. + if (current != null && now < current.token.expireAtMillis()) { + if (!lock.tryLock()) { + return current.token; + } + } else { + lock.lock(); + } + try { + return reload(); + } finally { + lock.unlock(); + } + } + + /** Returns the current token, or reloads it when it has less than the given time left. */ + public RESTToken tokenValidFor(long minimumValidityMillis) { + RESTToken token = token(); + if (remainingMillis(token) >= minimumValidityMillis) { + return token; + } + lock.lock(); + try { + CachedToken current = cached; + if (current != null && remainingMillis(current.token) >= minimumValidityMillis) { + return current.token; + } + RESTToken loaded = load(); + cached = new CachedToken(loaded, currentTimeMillis()); + lastFailure = null; + return loaded; + } finally { + lock.unlock(); + } + } + + /** Milliseconds left until the token expires. */ + public long remainingMillis(RESTToken token) { + return token.expireAtMillis() - currentTimeMillis(); + } + + private RESTToken reload() { + CachedToken current = cached; + long now = currentTimeMillis(); + if (current != null && now < current.reloadAtMillis) { + return current.token; + } + boolean valid = current != null && now < current.token.expireAtMillis(); + if (now < nextAttemptMillis) { + if (valid) { + return current.token; + } + throw new IllegalStateException( + "The data token of " + identifier + " expired and reloading it failed.", + lastFailure); + } + try { + RESTToken loaded = load(); + cached = new CachedToken(loaded, now); + lastFailure = null; + return loaded; + } catch (RuntimeException e) { + lastFailure = e; + nextAttemptMillis = now + RELOAD_INTERVAL_MILLIS; Review Comment: [P2] Start the retry cooldown when the refresh attempt finishes `now` is sampled before load(). If a failed REST request takes more than the 10-second reload interval, this deadline is already in the past when the call returns. The next stream request immediately calls REST again instead of reusing the still-valid token during the promised cooldown. This is reachable with the production client: HttpClientUtils allows three-minute connection-request/response timeouts and HTTP retries. A clock-controlled production-refresher probe starts with a two-minute token, advances 61 seconds into its refresh window, and has loadTableToken consume 20 seconds before failing; two immediate token() calls perform two loads although the old token is still valid (expected one). Under a slow catalog outage, successive reads can repeatedly pay the REST timeout and keep sending refresh traffic. Please sample the completion time in the failure handler and set nextAttemptMillis relative to it. Recheck the old token expiry at that time too: the existing `valid` flag w as computed before the same slow request and can otherwise return an already-expired bundle. Cover slow failure both while the old token remains valid and when it expires during the attempt. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
