sundapeng commented on code in PR #10178: URL: https://github.com/apache/paimon/pull/10178#discussion_r4253860944
########## 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: Agreed, fixed in 3d9f3c2f0. `reload()` now reads the clock again when the request finishes. A failure starts the retry interval from that time and rethrows when the current token has expired by then. A successful reload derives the new token's reload time from its arrival as well, since a slow request returning a short-lived token had the same flaw and could be reloaded on the next call. `RESTTokenRefresherTest` covers, with a 20-second request on the test clock: - a failure while the old token stays valid: two immediate calls make one request, and the next one follows after the retry interval; - a failure during which the old token expires: the REST error is thrown, then the retry-interval error without another request; - a successful load returning a 16-second token: it is kept for the reload interval instead of being reloaded on the next call. All three fail on the previous code. -- 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]
