From a70976987b01cd0ab0465026a4dc45cf33c64aaa Mon Sep 17 00:00:00 2001 From: Poon Yat Sing Date: Fri, 12 Dec 2025 14:09:59 +0900 Subject: [PATCH] Migrate from Guava Cache to Caffeine --- gateway-ha/pom.xml | 6 ++++ .../gateway/ha/router/BaseRoutingManager.java | 32 ++++++------------- .../gateway/ha/router/TrinoRequestUser.java | 12 +++---- 3 files changed, 20 insertions(+), 30 deletions(-) diff --git a/gateway-ha/pom.xml b/gateway-ha/pom.xml index 36b46ef50..ef5bd58b5 100644 --- a/gateway-ha/pom.xml +++ b/gateway-ha/pom.xml @@ -74,6 +74,12 @@ jackson-dataformat-yaml + + com.github.ben-manes.caffeine + caffeine + 3.2.3 + + com.google.guava guava diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/router/BaseRoutingManager.java b/gateway-ha/src/main/java/io/trino/gateway/ha/router/BaseRoutingManager.java index 996ee158c..73dcca3b9 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/router/BaseRoutingManager.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/router/BaseRoutingManager.java @@ -13,12 +13,11 @@ */ package io.trino.gateway.ha.router; +import com.github.benmanes.caffeine.cache.Caffeine; +import com.github.benmanes.caffeine.cache.LoadingCache; import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Function; import com.google.common.base.Strings; -import com.google.common.cache.CacheBuilder; -import com.google.common.cache.CacheLoader; -import com.google.common.cache.LoadingCache; import io.airlift.log.Logger; import io.trino.gateway.ha.clustermonitor.ClusterStats; import io.trino.gateway.ha.clustermonitor.TrinoStatus; @@ -36,7 +35,6 @@ import java.util.Map; import java.util.Optional; import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.Future; @@ -123,7 +121,7 @@ public String findBackendForQueryId(String queryId) try { backendAddress = queryIdBackendCache.get(queryId); } - catch (ExecutionException e) { + catch (RuntimeException e) { log.warn("Exception while loading queryId from cache %s", e.getLocalizedMessage()); } return backendAddress; @@ -137,7 +135,7 @@ public String findExternalUrlForQueryId(String queryId) try { externalUrl = queryIdExternalUrlCache.get(queryId); } - catch (ExecutionException e) { + catch (RuntimeException e) { log.warn("Exception while loading queryId from cache %s", e.getLocalizedMessage()); } return externalUrl; @@ -155,7 +153,7 @@ public String findRoutingGroupForQueryId(String queryId) try { routingGroup = queryIdRoutingGroupCache.get(queryId); } - catch (ExecutionException e) { + catch (RuntimeException e) { log.warn("Exception while loading queryId from routing group cache %s", e.getLocalizedMessage()); } return routingGroup; @@ -245,9 +243,7 @@ private String searchAllBackendForQuery(String queryId) */ private String findRoutingGroupForUnknownQueryId(String queryId) { - String routingGroup = queryHistoryManager.getRoutingGroupForQueryId(queryId); - setRoutingGroupForQueryId(queryId, routingGroup); - return routingGroup; + return queryHistoryManager.getRoutingGroupForQueryId(queryId); } /** @@ -255,25 +251,15 @@ private String findRoutingGroupForUnknownQueryId(String queryId) */ private String findExternalUrlForUnknownQueryId(String queryId) { - String externalUrl = queryHistoryManager.getExternalUrlForQueryId(queryId); - setExternalUrlForQueryId(queryId, externalUrl); - return externalUrl; + return queryHistoryManager.getExternalUrlForQueryId(queryId); } private LoadingCache buildCache(Function loader) { - return CacheBuilder.newBuilder() + return Caffeine.newBuilder() .maximumSize(10000) .expireAfterAccess(30, TimeUnit.MINUTES) - .build( - new CacheLoader<>() - { - @Override - public String load(String queryId) - { - return loader.apply(queryId); - } - }); + .build(loader::apply); } private boolean isBackendHealthy(String backendId) diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/router/TrinoRequestUser.java b/gateway-ha/src/main/java/io/trino/gateway/ha/router/TrinoRequestUser.java index 449a198f1..6ff32d5c6 100644 --- a/gateway-ha/src/main/java/io/trino/gateway/ha/router/TrinoRequestUser.java +++ b/gateway-ha/src/main/java/io/trino/gateway/ha/router/TrinoRequestUser.java @@ -22,9 +22,8 @@ import com.fasterxml.jackson.databind.SerializerProvider; import com.fasterxml.jackson.databind.annotation.JsonSerialize; import com.fasterxml.jackson.databind.ser.std.StdSerializer; -import com.google.common.cache.CacheBuilder; -import com.google.common.cache.CacheLoader; -import com.google.common.cache.LoadingCache; +import com.github.benmanes.caffeine.cache.Caffeine; +import com.github.benmanes.caffeine.cache.LoadingCache; import com.nimbusds.oauth2.sdk.ParseException; import com.nimbusds.oauth2.sdk.Request; import com.nimbusds.oauth2.sdk.token.BearerAccessToken; @@ -42,7 +41,6 @@ import java.util.Base64; import java.util.Map; import java.util.Optional; -import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; import static com.nimbusds.openid.connect.sdk.UserInfoResponse.parse; @@ -198,7 +196,7 @@ private Optional extractUserFromBearerAuth(String header, String userFie userInfo = Optional.of(userInfoCache.orElseThrow().get(token)); return Optional.of(userInfo.orElseThrow().getSubject().toString()); } - catch (ExecutionException e) { + catch (RuntimeException e) { log.error(e, "Could not get userInfo"); } } @@ -216,10 +214,10 @@ public TrinoRequestUserProvider(RequestAnalyzerConfig config) userField = config.getTokenUserField(); if (config.getOauthTokenInfoUrl() != null) { oauthUserInfoUrl = Optional.of(URI.create(config.getOauthTokenInfoUrl())); - userInfoCache = Optional.of(CacheBuilder.newBuilder() + userInfoCache = Optional.of(Caffeine.newBuilder() .maximumSize(10000) .expireAfterAccess(10, TimeUnit.MINUTES) - .build(CacheLoader.from(this::getUserInfo))); + .build(this::getUserInfo)); } else { oauthUserInfoUrl = Optional.empty();