From 8813d1598ec4d78b579c9f417ca7b5fb825bcba0 Mon Sep 17 00:00:00 2001
From: Xiaojian Sun <sunxiaojian926@163.com>
Date: Thu, 23 Jan 2025 20:24:49 +0800
Subject: [PATCH] [#6270] (Improvement) unified cache framework (#6271)

### What changes were proposed in this pull request?

Unified cache framework

### Why are the changes needed?

Fix: https://github.com/apache/gravitino/issues/6270

### Does this PR introduce _any_ user-facing change?

N/A

### How was this patch tested?

N/A
---
 spark-connector/spark-common/build.gradle.kts         |  1 +
 .../connector/catalog/GravitinoCatalogManager.java    | 11 +++++------
 spark-connector/v3.3/spark-runtime/build.gradle.kts   |  1 +
 spark-connector/v3.4/spark-runtime/build.gradle.kts   |  1 +
 spark-connector/v3.5/spark-runtime/build.gradle.kts   |  1 +
 5 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/spark-connector/spark-common/build.gradle.kts b/spark-connector/spark-common/build.gradle.kts
index 06e0077d21e..dc6bbd258e7 100644
--- a/spark-connector/spark-common/build.gradle.kts
+++ b/spark-connector/spark-common/build.gradle.kts
@@ -40,6 +40,7 @@ val scalaCollectionCompatVersion: String = libs.versions.scala.collection.compat
 dependencies {
   implementation(project(":catalogs:catalog-common"))
   implementation(libs.guava)
+  implementation(libs.caffeine)
 
   compileOnly(project(":clients:client-java-runtime", configuration = "shadow"))
   compileOnly("org.apache.iceberg:iceberg-spark-runtime-${sparkMajorVersion}_$scalaVersion:$icebergVersion")
diff --git a/spark-connector/spark-common/src/main/java/org/apache/gravitino/spark/connector/catalog/GravitinoCatalogManager.java b/spark-connector/spark-common/src/main/java/org/apache/gravitino/spark/connector/catalog/GravitinoCatalogManager.java
index f04193f33a3..a484972656a 100644
--- a/spark-connector/spark-common/src/main/java/org/apache/gravitino/spark/connector/catalog/GravitinoCatalogManager.java
+++ b/spark-connector/spark-common/src/main/java/org/apache/gravitino/spark/connector/catalog/GravitinoCatalogManager.java
@@ -18,13 +18,12 @@
  */
 package org.apache.gravitino.spark.connector.catalog;
 
+import com.github.benmanes.caffeine.cache.Cache;
+import com.github.benmanes.caffeine.cache.Caffeine;
 import com.google.common.base.Preconditions;
 import com.google.common.base.Supplier;
-import com.google.common.cache.Cache;
-import com.google.common.cache.CacheBuilder;
 import java.util.Arrays;
 import java.util.Map;
-import java.util.concurrent.ExecutionException;
 import org.apache.gravitino.Catalog;
 import org.apache.gravitino.client.GravitinoClient;
 import org.slf4j.Logger;
@@ -42,7 +41,7 @@ public class GravitinoCatalogManager {
   private GravitinoCatalogManager(Supplier<GravitinoClient> clientBuilder) {
     this.gravitinoClient = clientBuilder.get();
     // Will not evict catalog by default
-    this.gravitinoCatalogs = CacheBuilder.newBuilder().build();
+    this.gravitinoCatalogs = Caffeine.newBuilder().build();
   }
 
   public static GravitinoCatalogManager create(Supplier<GravitinoClient> clientBuilder) {
@@ -69,8 +68,8 @@ public void close() {
 
   public Catalog getGravitinoCatalogInfo(String name) {
     try {
-      return gravitinoCatalogs.get(name, () -> loadCatalog(name));
-    } catch (ExecutionException e) {
+      return gravitinoCatalogs.get(name, catalogName -> loadCatalog(catalogName));
+    } catch (Exception e) {
       LOG.error(String.format("Load catalog %s failed", name), e);
       throw new RuntimeException(e);
     }
diff --git a/spark-connector/v3.3/spark-runtime/build.gradle.kts b/spark-connector/v3.3/spark-runtime/build.gradle.kts
index 851473f4bf7..7f4b7f2edcf 100644
--- a/spark-connector/v3.3/spark-runtime/build.gradle.kts
+++ b/spark-connector/v3.3/spark-runtime/build.gradle.kts
@@ -51,6 +51,7 @@ tasks.withType<ShadowJar>(ShadowJar::class.java) {
   relocate("com.google", "org.apache.gravitino.shaded.com.google")
   relocate("google", "org.apache.gravitino.shaded.google")
   relocate("org.apache.hc", "org.apache.gravitino.shaded.org.apache.hc")
+  relocate("com.github.benmanes.caffeine", "org.apache.gravitino.shaded.com.github.benmanes.caffeine")
 }
 
 publishing {
diff --git a/spark-connector/v3.4/spark-runtime/build.gradle.kts b/spark-connector/v3.4/spark-runtime/build.gradle.kts
index 0b766fef85b..13de294c055 100644
--- a/spark-connector/v3.4/spark-runtime/build.gradle.kts
+++ b/spark-connector/v3.4/spark-runtime/build.gradle.kts
@@ -51,6 +51,7 @@ tasks.withType<ShadowJar>(ShadowJar::class.java) {
   relocate("com.google", "org.apache.gravitino.shaded.com.google")
   relocate("google", "org.apache.gravitino.shaded.google")
   relocate("org.apache.hc", "org.apache.gravitino.shaded.org.apache.hc")
+  relocate("com.github.benmanes.caffeine", "org.apache.gravitino.shaded.com.github.benmanes.caffeine")
 }
 
 publishing {
diff --git a/spark-connector/v3.5/spark-runtime/build.gradle.kts b/spark-connector/v3.5/spark-runtime/build.gradle.kts
index 2e21e5bd755..c1cdae6c69d 100644
--- a/spark-connector/v3.5/spark-runtime/build.gradle.kts
+++ b/spark-connector/v3.5/spark-runtime/build.gradle.kts
@@ -51,6 +51,7 @@ tasks.withType<ShadowJar>(ShadowJar::class.java) {
   relocate("com.google", "org.apache.gravitino.shaded.com.google")
   relocate("google", "org.apache.gravitino.shaded.google")
   relocate("org.apache.hc", "org.apache.gravitino.shaded.org.apache.hc")
+  relocate("com.github.benmanes.caffeine", "org.apache.gravitino.shaded.com.github.benmanes.caffeine")
 }
 
 publishing {