[#12403] fix(core): defer catalog wrapper cleanup with an operation lease

core/src/main/java/org/apache/gravitino/authorization/AuthorizationUtils.java

@@ -23,6 +23,7 @@
 import com.google.common.collect.Lists;
 import com.google.common.collect.Sets;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.List;
 import java.util.Set;
@@ -40,7 +41,6 @@
 import org.apache.gravitino.catalog.CatalogManager;
 import org.apache.gravitino.catalog.FilesetDispatcher;
 import org.apache.gravitino.catalog.hive.HiveConstants;
-import org.apache.gravitino.connector.BaseCatalog;
 import org.apache.gravitino.connector.authorization.AuthorizationPlugin;
 import org.apache.gravitino.dto.authorization.PrivilegeDTO;
 import org.apache.gravitino.dto.util.DTOConverters;
@@ -367,30 +367,27 @@ public static void callAuthorizationPluginForSecurableObjects(
     for (SecurableObject securableObject : securableObjects) {
       if (needApplyAuthorizationPluginAllCatalogs(securableObject)) {
         NameIdentifier[] catalogs = catalogManager.listCatalogs(Namespace.of(metalake));
-        // ListCatalogsInfo return `CatalogInfo` instead of `BaseCatalog`, we need `BaseCatalog` to
-        // call authorization plugin method.
         for (NameIdentifier catalog : catalogs) {
-          callAuthorizationPluginImpl(consumer, catalogManager.loadCatalog(catalog));
+          callAuthorizationPluginImpl(consumer, catalogManager, catalog);
         }
 
       } else if (needApplyAuthorization(securableObject.type())) {
         NameIdentifier catalogIdent =
             NameIdentifierUtil.getCatalogIdentifier(
                 MetadataObjectUtil.toEntityIdent(metalake, securableObject));
-        Catalog catalog = catalogManager.loadCatalog(catalogIdent);
-        if (!catalogsAlreadySet.contains(catalog.name())) {
-          catalogsAlreadySet.add(catalog.name());
-          callAuthorizationPluginImpl(consumer, catalog);
+        if (catalogsAlreadySet.add(catalogIdent.name())) {
+          callAuthorizationPluginImpl(consumer, catalogManager, catalogIdent);
         }
       }
     }
   }
 
   public static void callAuthorizationPluginForMetadataObject(
       String metalake, MetadataObject metadataObject, Consumer<AuthorizationPlugin> consumer) {
-    List<Catalog> loadedCatalogs = loadMetadataObjectCatalog(metalake, metadataObject);
-    for (Catalog catalog : loadedCatalogs) {
-      callAuthorizationPluginImpl(consumer, catalog);
+    CatalogManager catalogManager = GravitinoEnv.getInstance().catalogManager();
+    List<NameIdentifier> catalogIdents = getMetadataObjectCatalogs(metalake, metadataObject);
+    for (NameIdentifier catalogIdent : catalogIdents) {
+      callAuthorizationPluginImpl(consumer, catalogManager, catalogIdent);
     }
   }
 
@@ -504,18 +501,19 @@ public static void authorizationPluginRemovePrivileges(
     }
   }
 
-  public static void removeCatalogPrivileges(Catalog catalog, List<String> locations) {
+  public static void removeCatalogPrivileges(NameIdentifier catalogIdent, List<String> locations) {
     // If we enable authorization, we should remove the privileges about the entity in the
     // authorization plugin.
     MetadataObject metadataObject =
-        MetadataObjects.of(null, catalog.name(), MetadataObject.Type.CATALOG);
+        MetadataObjects.of(null, catalogIdent.name(), MetadataObject.Type.CATALOG);
     MetadataObjectChange removeObject = MetadataObjectChange.remove(metadataObject, locations);
 
     callAuthorizationPluginImpl(
         authorizationPlugin -> {
           authorizationPlugin.onMetadataUpdated(removeObject);
         },
-        catalog);
+        GravitinoEnv.getInstance().catalogManager(),
+        catalogIdent);
   }
 
   public static void authorizationPluginRenamePrivileges(
@@ -600,35 +598,33 @@ private static boolean needApplyAuthorization(MetadataObject.Type type) {
   }
 
   private static void callAuthorizationPluginImpl(
-      BiConsumer<AuthorizationPlugin, String> consumer, Catalog catalog) {
-
-    if (catalog instanceof BaseCatalog) {
-      BaseCatalog baseCatalog = (BaseCatalog) catalog;
-      if (baseCatalog.getAuthorizationPlugin() != null) {
-        consumer.accept(baseCatalog.getAuthorizationPlugin(), catalog.name());
-      }
-    } else {
-      throw new IllegalArgumentException(
-          String.format(
-              "Catalog %s is not a BaseCatalog, we don't support authorization plugin for it",
-              catalog.type()));
-    }
+      BiConsumer<AuthorizationPlugin, String> consumer,
+      CatalogManager catalogManager,
+      NameIdentifier catalogIdent) {
+    catalogManager.doWithCatalog(
+        catalogIdent,
+        catalog -> {
+          AuthorizationPlugin authorizationPlugin = catalog.getAuthorizationPlugin();
+          if (authorizationPlugin != null) {
+            consumer.accept(authorizationPlugin, catalog.name());
+          }
+          return null;
+        });
   }
 
   private static void callAuthorizationPluginImpl(
-      Consumer<AuthorizationPlugin> consumer, Catalog catalog) {
-
-    if (catalog instanceof BaseCatalog) {
-      BaseCatalog baseCatalog = (BaseCatalog) catalog;
-      if (baseCatalog.getAuthorizationPlugin() != null) {
-        consumer.accept(baseCatalog.getAuthorizationPlugin());
-      }
-    } else {
-      throw new IllegalArgumentException(
-          String.format(
-              "Catalog %s is not a BaseCatalog, we don't support authorization plugin for it",
-              catalog.type()));
-    }
+      Consumer<AuthorizationPlugin> consumer,
+      CatalogManager catalogManager,
+      NameIdentifier catalogIdent) {
+    catalogManager.doWithCatalog(
+        catalogIdent,
+        catalog -> {
+          AuthorizationPlugin authorizationPlugin = catalog.getAuthorizationPlugin();
+          if (authorizationPlugin != null) {
+            consumer.accept(authorizationPlugin);
+          }
+          return null;
+        });
   }
 
   private static void checkCatalogType(
@@ -642,26 +638,21 @@ private static void checkCatalogType(
     }
   }
 
-  private static List<Catalog> loadMetadataObjectCatalog(
+  private static List<NameIdentifier> getMetadataObjectCatalogs(
       String metalake, MetadataObject metadataObject) {
     CatalogManager catalogManager = GravitinoEnv.getInstance().catalogManager();
-    List<Catalog> loadedCatalogs = Lists.newArrayList();
+    List<NameIdentifier> catalogIdents = Lists.newArrayList();
     if (needApplyAuthorizationPluginAllCatalogs(metadataObject.type())) {
       NameIdentifier[] catalogs = catalogManager.listCatalogs(Namespace.of(metalake));
-      // ListCatalogsInfo return `CatalogInfo` instead of `BaseCatalog`, we need `BaseCatalog` to
-      // call authorization plugin method.
-      for (NameIdentifier catalog : catalogs) {
-        loadedCatalogs.add(catalogManager.loadCatalog(catalog));
-      }
+      catalogIdents.addAll(Arrays.asList(catalogs));
     } else if (needApplyAuthorization(metadataObject.type())) {
       NameIdentifier catalogIdent =
           NameIdentifierUtil.getCatalogIdentifier(
               MetadataObjectUtil.toEntityIdent(metalake, metadataObject));
-      Catalog catalog = catalogManager.loadCatalog(catalogIdent);
-      loadedCatalogs.add(catalog);
+      catalogIdents.add(catalogIdent);
     }
 
-    return loadedCatalogs;
+    return catalogIdents;
   }
 
   // The Hive default schema location is Hive warehouse directory

core/src/main/java/org/apache/gravitino/catalog/CapabilityHelpers.java

@@ -52,11 +52,16 @@ public class CapabilityHelpers {
 
   public static Capability getCapability(NameIdentifier ident, CatalogManager catalogManager) {
     NameIdentifier catalogIdent = getCatalogIdentifier(ident);
-    CatalogManager.CatalogWrapper c = catalogManager.loadCatalogAndWrap(catalogIdent);
+    // Acquire the lease outside the try so a missing catalog keeps propagating its
+    // NoSuchCatalogException (a 404) instead of being wrapped into a plain RuntimeException (a
+    // 500); only the capability lookup itself is wrapped.
+    CatalogLease lease = catalogManager.acquireCatalogLease(catalogIdent);
     try {
-      return c.capabilities();
+      return lease.wrapper().capabilities();
     } catch (Exception e) {
       throw new RuntimeException("Failed to get capabilities for catalog: " + catalogIdent, e);
+    } finally {
+      lease.close();
     }
   }
 

core/src/main/java/org/apache/gravitino/catalog/CatalogLease.java

@@ -0,0 +1,78 @@
+/*
+ * 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.gravitino.catalog;
+
+import java.util.concurrent.atomic.AtomicBoolean;
+import org.apache.gravitino.NameIdentifier;
+import org.apache.gravitino.catalog.CatalogManager.CatalogWrapper;
+import org.apache.gravitino.connector.BaseCatalog;
+
+/**
+ * A lease on a {@link CatalogWrapper} held for the duration of one catalog operation.
+ *
+ * <p>While the lease is held, the wrapper's catalog instance and its {@link
+ * org.apache.gravitino.utils.IsolatedClassLoader} stay alive even if the catalog cache evicts the
+ * wrapper concurrently (expiry, explicit invalidation, or remote change-log invalidation). The
+ * resources are released once the wrapper is retired and its last lease is closed, so an operation
+ * can never observe a half-closed catalog.
+ *
+ * <p>Leases are obtained from {@link CatalogManager#acquireCatalogLease(NameIdentifier)} and must
+ * be closed exactly once, ideally with try-with-resources:
+ *
+ * <pre>{@code
+ * try (CatalogLease lease = catalogManager.acquireCatalogLease(ident)) {
+ *   return lease.wrapper().doWithTableOps(ops -> ops.loadTable(tableIdent));
+ * }
+ * }</pre>
+ */
+public final class CatalogLease implements AutoCloseable {
+
+  private final CatalogWrapper wrapper;
+  private final AtomicBoolean released = new AtomicBoolean(false);
+
+  CatalogLease(CatalogWrapper wrapper) {
+    this.wrapper = wrapper;
+  }
+
+  /**
+   * Returns the leased catalog wrapper.
+   *
+   * @return the leased catalog wrapper, guaranteed to stay usable until this lease is closed.
+   */
+  public CatalogWrapper wrapper() {
+    return wrapper;
+  }
+
+  /**
+   * Returns the catalog of the leased wrapper.
+   *
+   * @return the leased catalog, guaranteed to stay usable until this lease is closed.
+   */
+  public BaseCatalog catalog() {
+    return wrapper.catalog();
+  }
+
+  /** Releases the lease. Closing an already closed lease is a no-op. */
+  @Override
+  public void close() {
+    if (released.compareAndSet(false, true)) {
+      wrapper.release();
+    }
+  }
+}