Fix open-metadata#15503: Cache MultiJWKsProvider response (open-metad…

…ata#15504)
TeddyCr · Mar 11, 2024 · 7d4f127 · 7d4f127
1 parent 1f8e232
commit 7d4f127
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 14 deletions.
diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/events/ChangeEventHandler.java b/openmetadata-service/src/main/java/org/openmetadata/service/events/ChangeEventHandler.java
@@ -68,7 +68,7 @@ public Void process(
         }
         // Always set the Change Event Username as context Principal, the one creating the CE
         changeEvent.setUserName(loggedInUserName);
-        LOG.info(
+        LOG.debug(
             "Recording change event {}:{}:{}:{}",
             changeEvent.getTimestamp(),
             changeEvent.getEntityId(),

diff --git a/...metadata-service/src/main/java/org/openmetadata/service/security/MultiUrlJwkProvider.java b/...metadata-service/src/main/java/org/openmetadata/service/security/MultiUrlJwkProvider.java
@@ -18,30 +18,50 @@
 import com.auth0.jwk.JwkProvider;
 import com.auth0.jwk.SigningKeyNotFoundException;
 import com.auth0.jwk.UrlJwkProvider;
+import com.google.common.cache.CacheBuilder;
+import com.google.common.cache.CacheLoader;
+import com.google.common.cache.LoadingCache;
 import java.net.URL;
 import java.util.List;
+import java.util.concurrent.TimeUnit;
+import org.openmetadata.service.exception.UnhandledServerException;
 
 final class MultiUrlJwkProvider implements JwkProvider {
   private final List<UrlJwkProvider> urlJwkProviders;
+  private LoadingCache<String, Jwk> CACHE =
+      CacheBuilder.newBuilder()
+          .maximumSize(10)
+          .expireAfterWrite(24, TimeUnit.HOURS)
+          .build(
+              new CacheLoader<>() {
+                @Override
+                public Jwk load(String key) throws Exception {
+                  JwkException lastException =
+                      new SigningKeyNotFoundException(
+                          "JWT Token keyID doesn't match the configured keyID. This usually happens if you didn't configure "
+                              + "proper publicKeyUrls under authentication configuration.",
+                          null);
+                  for (UrlJwkProvider jwkProvider : urlJwkProviders) {
+                    try {
+                      return jwkProvider.get(key);
+                    } catch (JwkException e) {
+                      lastException.addSuppressed(e);
+                    }
+                  }
+                  throw lastException;
+                }
+              });
 
   public MultiUrlJwkProvider(List<URL> publicKeyUris) {
     this.urlJwkProviders = publicKeyUris.stream().map(UrlJwkProvider::new).toList();
   }
 
   @Override
-  public Jwk get(String keyId) throws JwkException {
-    JwkException lastException =
-        new SigningKeyNotFoundException(
-            "JWT Token keyID doesn't match the configured keyID. This usually happens if you didn't configure "
-                + "proper publicKeyUrls under authentication configuration.",
-            null);
-    for (UrlJwkProvider jwkProvider : urlJwkProviders) {
-      try {
-        return jwkProvider.get(keyId);
-      } catch (JwkException e) {
-        lastException.addSuppressed(e);
-      }
+  public Jwk get(String keyId) {
+    try {
+      return CACHE.get(keyId);
+    } catch (Exception e) {
+      throw new UnhandledServerException(e.getMessage());
     }
-    throw lastException;
   }
 }