From 95db76b2b46e00684bb9cd621775240891d7dad4 Mon Sep 17 00:00:00 2001 From: yuqi Date: Mon, 10 Aug 2026 18:26:15 +0800 Subject: [PATCH 1/6] [#12403] fix(core): lease cached catalog wrappers during operations CatalogManager closed CatalogWrapper synchronously from the catalog cache removal listener. Caffeine runs that listener asynchronously and outside the local TreeLock, so an expiry, a remote change-log invalidation or a drop could close the catalog, clear its reference and release the pooled ClassLoader while another thread was still running an operation on that wrapper. CatalogWrapper now counts active operations: tryAcquire() takes a lease, release() returns it and retire() (called from the removal listener) only marks the wrapper unusable for new leases. The catalog and the ClassLoader are cleaned up exactly once, when the wrapper is retired and its last lease is released. Operations obtain a CatalogLease from CatalogManager.acquireCatalogLease(), which reloads a fresh wrapper when the cached one has already retired, and all production uses of the wrapper were migrated to it. --- .../gravitino/catalog/CapabilityHelpers.java | 5 +- .../gravitino/catalog/CatalogLease.java | 93 +++++ .../gravitino/catalog/CatalogManager.java | 273 +++++++++++---- .../catalog/OperationDispatcher.java | 15 +- .../gravitino/catalog/TestCatalogManager.java | 18 +- .../catalog/TestCatalogWrapperLease.java | 326 ++++++++++++++++++ .../TestFunctionOperationDispatcher.java | 6 + .../TestPartitionNormalizeDispatcher.java | 3 + .../catalog/TestTableNormalizeDispatcher.java | 3 + .../hook/TestFilesetHookDispatcher.java | 7 + .../hook/TestFunctionHookDispatcher.java | 10 + .../hook/TestModelHookDispatcher.java | 4 + .../hook/TestSchemaHookDispatcher.java | 4 + .../hook/TestTableHookDispatcher.java | 7 + .../hook/TestTopicHookDispatcher.java | 7 + .../dispatcher/IcebergCleanupHelper.java | 13 +- .../dispatcher/TestIcebergAsyncPurge.java | 4 + 17 files changed, 717 insertions(+), 81 deletions(-) create mode 100644 core/src/main/java/org/apache/gravitino/catalog/CatalogLease.java create mode 100644 core/src/test/java/org/apache/gravitino/catalog/TestCatalogWrapperLease.java diff --git a/core/src/main/java/org/apache/gravitino/catalog/CapabilityHelpers.java b/core/src/main/java/org/apache/gravitino/catalog/CapabilityHelpers.java index eae018c2c0d..125bbb021bf 100644 --- a/core/src/main/java/org/apache/gravitino/catalog/CapabilityHelpers.java +++ b/core/src/main/java/org/apache/gravitino/catalog/CapabilityHelpers.java @@ -52,9 +52,8 @@ public class CapabilityHelpers { public static Capability getCapability(NameIdentifier ident, CatalogManager catalogManager) { NameIdentifier catalogIdent = getCatalogIdentifier(ident); - CatalogManager.CatalogWrapper c = catalogManager.loadCatalogAndWrap(catalogIdent); - try { - return c.capabilities(); + try (CatalogLease lease = catalogManager.acquireCatalogLease(catalogIdent)) { + return lease.wrapper().capabilities(); } catch (Exception e) { throw new RuntimeException("Failed to get capabilities for catalog: " + catalogIdent, e); } diff --git a/core/src/main/java/org/apache/gravitino/catalog/CatalogLease.java b/core/src/main/java/org/apache/gravitino/catalog/CatalogLease.java new file mode 100644 index 00000000000..ad0e687bd9b --- /dev/null +++ b/core/src/main/java/org/apache/gravitino/catalog/CatalogLease.java @@ -0,0 +1,93 @@ +/* + * 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 com.google.common.base.Preconditions; +import java.util.concurrent.atomic.AtomicBoolean; +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. + * + *

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. + * + *

Leases are obtained from {@link CatalogManager#acquireCatalogLease(org.apache.gravitino + * .NameIdentifier)} and must be closed exactly once, ideally with try-with-resources: + * + * 

{@code
+ * try (CatalogLease lease = catalogManager.acquireCatalogLease(ident)) {
+ *   return lease.wrapper().doWithTableOps(ops -> ops.loadTable(tableIdent));
+ * }
+ * }
+ */ +public final class CatalogLease implements AutoCloseable { + + private final CatalogWrapper wrapper; + private final AtomicBoolean released = new AtomicBoolean(false); + + CatalogLease(CatalogWrapper wrapper) { + this.wrapper = wrapper; + } + + /** + * Takes a lease on a wrapper the caller already holds. Prefer {@link + * CatalogManager#acquireCatalogLease(org.apache.gravitino.NameIdentifier)}, which reloads a fresh + * wrapper when the cached one has been retired. + * + * @param wrapper the wrapper to lease. + * @return a lease on the given wrapper. + * @throws IllegalStateException if the wrapper has already been retired. + */ + public static CatalogLease of(CatalogWrapper wrapper) { + Preconditions.checkState( + wrapper.tryAcquire(), "Catalog wrapper has already been retired, cannot lease it"); + return new CatalogLease(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(); + } + } +} diff --git a/core/src/main/java/org/apache/gravitino/catalog/CatalogManager.java b/core/src/main/java/org/apache/gravitino/catalog/CatalogManager.java index 3bafc882b87..2bd78315359 100644 --- a/core/src/main/java/org/apache/gravitino/catalog/CatalogManager.java +++ b/core/src/main/java/org/apache/gravitino/catalog/CatalogManager.java @@ -124,6 +124,11 @@ public class CatalogManager implements CatalogDispatcher, Closeable { private static final String CATALOG_DOES_NOT_EXIST_MSG = "Catalog %s does not exist"; + // Bounds the retry loop in acquireCatalogLease() when the wrapper read from the cache is retired + // by a concurrent eviction before the lease can be taken. Each retry reloads a fresh wrapper, so + // exhausting the attempts means the catalog is being evicted continuously. + private static final int MAX_LEASE_ATTEMPTS = 5; + private static final Logger LOG = LoggerFactory.getLogger(CatalogManager.class); private static final Set CONTRIB_CATALOGS_TYPES = @@ -156,14 +161,35 @@ public class CatalogManager implements CatalogDispatcher, Closeable { "uri", "fs.defaultFS"); - /** Wrapper class for a catalog instance and its class loader. */ + /** + * Wrapper class for a catalog instance and its class loader. + * + *

A wrapper is shared by all threads that read it from the catalog cache, while cache eviction + * (expiry, explicit invalidation, or remote change-log invalidation) happens outside the tree + * lock. To keep an eviction from tearing down a catalog that an in-flight operation is still + * using, the wrapper counts active operations: {@link #tryAcquire()} takes a lease, {@link + * #release()} returns it, and {@link #retire()} (called from the cache removal listener) only + * marks the wrapper unusable for new leases. The catalog and the ClassLoader are cleaned up + * exactly once, when the wrapper is retired and the last lease has been released. + */ public static class CatalogWrapper { private BaseCatalog catalog; private IsolatedClassLoader classLoader; private ClassLoaderPool pool; private PooledClassLoaderEntry poolEntry; - private boolean closed = false; + + /** Guards {@link #activeOps}, {@link #retired} and {@link #cleanupStarted}. */ + private final Object leaseLock = new Object(); + + /** Number of leases currently held by in-flight operations. */ + private int activeOps = 0; + + /** Set when the wrapper leaves the cache; no new lease can be acquired afterwards. */ + private boolean retired = false; + + /** Set by the thread that claims the (exactly-once) resource cleanup. */ + private boolean cleanupStarted = false; /** Non-pooled constructor: each catalog owns its ClassLoader exclusively. */ CatalogWrapper(IsolatedClassLoader classLoader) { @@ -182,6 +208,74 @@ public BaseCatalog catalog() { return catalog; } + /** + * Tries to take a lease on this wrapper, keeping its catalog and ClassLoader alive until the + * lease is released. + * + * @return true if the lease was taken, false if the wrapper has already been retired and the + * caller must load a fresh wrapper. + */ + public boolean tryAcquire() { + synchronized (leaseLock) { + if (retired) { + return false; + } + activeOps++; + return true; + } + } + + /** + * Releases a lease taken by {@link #tryAcquire()}. Cleans up the catalog and the ClassLoader if + * this was the last lease on an already retired wrapper. + */ + public void release() { + boolean shouldCleanup; + synchronized (leaseLock) { + Preconditions.checkState(activeOps > 0, "Releasing a lease that was never acquired"); + activeOps--; + shouldCleanup = claimCleanupIfIdle(); + } + + if (shouldCleanup) { + cleanup(); + } + } + + /** + * Retires this wrapper: no new lease can be taken. The catalog and the ClassLoader are cleaned + * up immediately if no operation is in flight, otherwise by the last {@link #release()}. + */ + public void retire() { + boolean shouldCleanup; + synchronized (leaseLock) { + retired = true; + shouldCleanup = claimCleanupIfIdle(); + } + + if (shouldCleanup) { + cleanup(); + } + } + + /** + * Returns whether this wrapper has been retired and can no longer serve new operations. + * + * @return true if the wrapper has been retired. + */ + public boolean isRetired() { + synchronized (leaseLock) { + return retired; + } + } + + @VisibleForTesting + int activeOperations() { + synchronized (leaseLock) { + return activeOps; + } + } + public R doWithSchemaOps(ThrowableFunction fn) throws Exception { return classLoader.withClassLoader( cl -> { @@ -285,13 +379,30 @@ public Capability capabilities() throws Exception { return classLoader.withClassLoader(cl -> catalog.capability()); } - public synchronized void close() { - if (closed) { - // Idempotent: a second close() must not re-run pool release or classloader cleanup. - return; + /** + * Retires the wrapper and, once no operation is in flight anymore, releases its resources. Kept + * as an alias of {@link #retire()} so callers that own a wrapper exclusively (for example + * {@link CatalogManager#testConnection}) can keep using the {@link java.io.Closeable}-style + * API. + */ + public void close() { + retire(); + } + + /** + * Claims the exactly-once resource cleanup when the wrapper is retired and idle. Must be called + * while holding {@link #leaseLock}; the caller runs {@link #cleanup()} outside the lock so a + * slow catalog close does not block {@link #tryAcquire()}. + */ + private boolean claimCleanupIfIdle() { + if (!retired || activeOps > 0 || cleanupStarted) { + return false; } - closed = true; + cleanupStarted = true; + return true; + } + private void cleanup() { try { classLoader.withClassLoader( cl -> { @@ -306,7 +417,7 @@ public synchronized void close() { } finally { // Release the pool reference (or clean up the dedicated ClassLoader) in a finally so a // failure while closing the catalog cannot permanently leak the pooled ClassLoader - // reference (close() is idempotent, so a retry would otherwise skip this). + // reference (cleanup() runs exactly once, so a retry would otherwise skip this). if (poolEntry != null) { pool.release(poolEntry); poolEntry = null; @@ -401,7 +512,10 @@ public CatalogManager( listener.accept((NameIdentifier) k); } } - ((CatalogWrapper) v).close(); + // Retire rather than close: an operation that already leased this wrapper keeps + // it alive, and the actual catalog/ClassLoader cleanup runs when the last lease + // is released. + ((CatalogWrapper) v).retire(); }) .scheduler( Scheduler.forScheduledExecutorService( @@ -569,9 +683,11 @@ public Catalog loadCatalog(NameIdentifier ident) throws NoSuchCatalogException { ident, LockType.READ, () -> { - BaseCatalog baseCatalog = loadCatalogAndWrap(ident).catalog(); - baseCatalog.checkMetalakeInUse(); - return baseCatalog; + try (CatalogLease lease = acquireCatalogLease(ident)) { + BaseCatalog baseCatalog = lease.catalog(); + baseCatalog.checkMetalakeInUse(); + return baseCatalog; + } }); } @@ -748,14 +864,14 @@ public void enableCatalog(NameIdentifier ident) metalakeIdent, LockType.WRITE, () -> { - BaseCatalog baseCatalog = loadCatalogAndWrap(ident).catalog(); - baseCatalog.checkMetalakeInUse(); + try (CatalogLease lease = acquireCatalogLease(ident)) { + BaseCatalog baseCatalog = lease.catalog(); + baseCatalog.checkMetalakeInUse(); - if (baseCatalog.catalogInUse()) { - return null; - } + if (baseCatalog.catalogInUse()) { + return null; + } - try { store.update( ident, CatalogEntity.class, @@ -789,14 +905,14 @@ public void disableCatalog(NameIdentifier ident) throws NoSuchCatalogException { metalakeIdent, LockType.WRITE, () -> { - BaseCatalog baseCatalog = loadCatalogAndWrap(ident).catalog(); - baseCatalog.checkMetalakeInUse(); + try (CatalogLease lease = acquireCatalogLease(ident)) { + BaseCatalog baseCatalog = lease.catalog(); + baseCatalog.checkMetalakeInUse(); - if (!baseCatalog.catalogInUse()) { - return null; - } + if (!baseCatalog.catalogInUse()) { + return null; + } - try { store.update( ident, CatalogEntity.class, @@ -841,31 +957,30 @@ public Catalog alterCatalog(NameIdentifier ident, CatalogChange... changes) LockType.READ, () -> { // There could be a race issue that someone is using the catalog from cache while we are - // updating it. - CatalogWrapper catalogWrapper = loadCatalogAndWrap(ident); - if (catalogWrapper == null) { - throw new NoSuchCatalogException(CATALOG_DOES_NOT_EXIST_MSG, ident); - } - - BaseCatalog catalog = catalogWrapper.catalog(); - catalog.checkMetalakeAndCatalogInUse(); - - try { - catalogWrapper.doWithPropertiesMeta( - f -> { - Pair, Map> alterProperty = - getCatalogAlterProperty(changes); - validatePropertyForAlter( - f.catalogPropertiesMetadata(), - alterProperty.getLeft(), - alterProperty.getRight()); - return null; - }); - } catch (IllegalArgumentException e1) { - throw e1; - } catch (Exception e) { - LOG.error("Failed to alter catalog {}", ident, e); - throw new RuntimeException(e); + // updating it. The lease keeps the wrapper alive for the whole validation. + try (CatalogLease lease = acquireCatalogLease(ident)) { + BaseCatalog catalog = lease.catalog(); + catalog.checkMetalakeAndCatalogInUse(); + + try { + lease + .wrapper() + .doWithPropertiesMeta( + f -> { + Pair, Map> alterProperty = + getCatalogAlterProperty(changes); + validatePropertyForAlter( + f.catalogPropertiesMetadata(), + alterProperty.getLeft(), + alterProperty.getRight()); + return null; + }); + } catch (IllegalArgumentException e1) { + throw e1; + } catch (Exception e) { + LOG.error("Failed to alter catalog {}", ident, e); + throw new RuntimeException(e); + } } return null; }); @@ -937,8 +1052,8 @@ public boolean dropCatalog(NameIdentifier ident, boolean force) metalakeIdent, LockType.WRITE, () -> { - try { - CatalogWrapper catalogWrapper = loadCatalogAndWrap(ident); + try (CatalogLease lease = acquireCatalogLease(ident)) { + CatalogWrapper catalogWrapper = lease.wrapper(); catalogWrapper.catalog().checkMetalakeInUse(); boolean catalogInUse = catalogWrapper.catalog().catalogInUse(); @@ -1085,10 +1200,44 @@ private boolean containsUserCreatedSchemas( return false; } + /** + * Loads the catalog with the specified identifier, wraps it in a CatalogWrapper, caches the + * wrapper for reuse, and takes a lease on it. The lease keeps the catalog and its ClassLoader + * alive for the duration of the operation even if the cache evicts the wrapper concurrently, so + * the caller must close the lease when the operation is done, ideally with try-with-resources. + * + *

If the cached wrapper has already been retired (by an eviction, an invalidation or a drop), + * the stale entry is evicted and a fresh wrapper is loaded and cached. + * + * @param ident The identifier of the catalog to load. + * @return A lease on the CatalogWrapper containing the loaded catalog. + * @throws NoSuchCatalogException If the specified catalog does not exist. + */ + public CatalogLease acquireCatalogLease(NameIdentifier ident) throws NoSuchCatalogException { + for (int attempt = 0; attempt < MAX_LEASE_ATTEMPTS; attempt++) { + CatalogWrapper wrapper = loadCatalogAndWrap(ident); + if (wrapper.tryAcquire()) { + return new CatalogLease(wrapper); + } + + // The cached wrapper was retired between the cache lookup and the lease attempt. Evict the + // stale entry and reload a fresh one. Use a conditional remove so we do not clobber a + // wrapper that another thread may have concurrently reloaded into the cache. + catalogCache.asMap().remove(ident, wrapper); + } + + throw new GravitinoRuntimeException( + "Failed to acquire a lease on catalog %s after %d attempts", ident, MAX_LEASE_ATTEMPTS); + } + /** * Loads the catalog with the specified identifier, wraps it in a CatalogWrapper, and caches the - * wrapper for reuse. If the cached wrapper has already been closed (its underlying catalog is - * null), the stale entry is evicted and a fresh wrapper is loaded and cached. + * wrapper for reuse. If the cached wrapper has already been retired, the stale entry is evicted + * and a fresh wrapper is loaded and cached. + * + *

The returned wrapper is not leased, so a concurrent cache eviction may retire and close it + * while the caller is still using it. Prefer {@link #acquireCatalogLease(NameIdentifier)}, which + * keeps the wrapper alive for the duration of the operation. * * @param ident The identifier of the catalog to load. * @return The wrapped CatalogWrapper containing the loaded catalog. @@ -1096,14 +1245,14 @@ private boolean containsUserCreatedSchemas( */ public CatalogWrapper loadCatalogAndWrap(NameIdentifier ident) throws NoSuchCatalogException { CatalogWrapper wrapper = catalogCache.get(ident, this::loadCatalogInternal); - if (wrapper.catalog() != null) { + if (!wrapper.isRetired()) { return wrapper; } - // The cached wrapper has already been closed (catalog() == null), e.g. by a prior - // dropCatalog or cache eviction. Evict the stale entry and reload a fresh one. - // Use a conditional remove so we do not clobber a wrapper that another thread may - // have concurrently reloaded into the cache between our initial get and this remove. + // The cached wrapper has already been retired, e.g. by a prior dropCatalog or cache eviction. + // Evict the stale entry and reload a fresh one. Use a conditional remove so we do not clobber + // a wrapper that another thread may have concurrently reloaded into the cache between our + // initial get and this remove. catalogCache.asMap().remove(ident, wrapper); return catalogCache.get(ident, this::loadCatalogInternal); } @@ -1295,9 +1444,11 @@ private Map getResolvedProperties(CatalogEntity entity) { // down a throwaway BaseCatalog (and leaking its authorizationPlugin) on every listCatalogsInfo // call, and keeps the classLoaderSharingEnabled branching in a single place // (createCatalogWrapper). - CatalogWrapper catalogWrapper = loadCatalogAndWrap(entity.nameIdentifier()); - return catalogWrapper.classLoader.withClassLoader( - cl -> catalogWrapper.catalog.properties(), RuntimeException.class); + try (CatalogLease lease = acquireCatalogLease(entity.nameIdentifier())) { + CatalogWrapper catalogWrapper = lease.wrapper(); + return catalogWrapper.classLoader.withClassLoader( + cl -> catalogWrapper.catalog.properties(), RuntimeException.class); + } } private BaseCatalog createBaseCatalog(IsolatedClassLoader classLoader, CatalogEntity entity) { diff --git a/core/src/main/java/org/apache/gravitino/catalog/OperationDispatcher.java b/core/src/main/java/org/apache/gravitino/catalog/OperationDispatcher.java index 4fa1808fc69..40ae30807da 100644 --- a/core/src/main/java/org/apache/gravitino/catalog/OperationDispatcher.java +++ b/core/src/main/java/org/apache/gravitino/catalog/OperationDispatcher.java @@ -86,8 +86,9 @@ protected R doWithTable( throws E { try { NameIdentifier catalogIdent = getCatalogIdentifier(tableIdent); - CatalogManager.CatalogWrapper c = catalogManager.loadCatalogAndWrap(catalogIdent); - return c.doWithPartitionOps(tableIdent, fn); + try (CatalogLease lease = catalogManager.acquireCatalogLease(catalogIdent)) { + return lease.wrapper().doWithPartitionOps(tableIdent, fn); + } } catch (Exception exception) { if (ex.isInstance(exception)) { throw ex.cast(exception); @@ -103,8 +104,9 @@ protected R doWithCatalog( NameIdentifier ident, ThrowableFunction fn, Class ex) throws E { try { - CatalogManager.CatalogWrapper c = catalogManager.loadCatalogAndWrap(ident); - return fn.apply(c); + try (CatalogLease lease = catalogManager.acquireCatalogLease(ident)) { + return fn.apply(lease.wrapper()); + } } catch (Exception exception) { if (ex.isInstance(exception)) { throw ex.cast(exception); @@ -123,8 +125,9 @@ protected R doWithCatalog( Class ex2) throws E1, E2 { try { - CatalogManager.CatalogWrapper c = catalogManager.loadCatalogAndWrap(ident); - return fn.apply(c); + try (CatalogLease lease = catalogManager.acquireCatalogLease(ident)) { + return fn.apply(lease.wrapper()); + } } catch (Exception exception) { if (ex1.isInstance(exception)) { throw ex1.cast(exception); diff --git a/core/src/test/java/org/apache/gravitino/catalog/TestCatalogManager.java b/core/src/test/java/org/apache/gravitino/catalog/TestCatalogManager.java index c55c4044692..3dbbbf65ad3 100644 --- a/core/src/test/java/org/apache/gravitino/catalog/TestCatalogManager.java +++ b/core/src/test/java/org/apache/gravitino/catalog/TestCatalogManager.java @@ -964,6 +964,7 @@ public void testDropCatalogSkipsImportedSchemas() throws Exception { Capability capability = Mockito.mock(Capability.class); CapabilityResult unsupportedResult = CapabilityResult.unsupported("Not managed"); Mockito.doReturn(wrapper).when(catalogManager).loadCatalogAndWrap(ident); + Mockito.when(wrapper.tryAcquire()).thenReturn(true); Mockito.doReturn(catalog).when(wrapper).catalog(); Mockito.doReturn(capability).when(wrapper).capabilities(); Mockito.doReturn(unsupportedResult).when(capability).managedStorage(any()); @@ -1043,6 +1044,7 @@ public void testDropCatalogIgnoresMissingSchema() throws Exception { Capability capability = Mockito.mock(Capability.class); CapabilityResult unsupportedResult = CapabilityResult.unsupported("Not managed"); Mockito.doReturn(wrapper).when(catalogManager).loadCatalogAndWrap(ident); + Mockito.when(wrapper.tryAcquire()).thenReturn(true); Mockito.doReturn(catalog).when(wrapper).catalog(); Mockito.doReturn(capability).when(wrapper).capabilities(); Mockito.doReturn(unsupportedResult).when(capability).managedStorage(any()); @@ -1094,6 +1096,7 @@ public void testDropCatalogFailsOnSchemaClassificationError() throws Exception { Capability capability = Mockito.mock(Capability.class); CapabilityResult unsupportedResult = CapabilityResult.unsupported("Not managed"); Mockito.doReturn(wrapper).when(catalogManager).loadCatalogAndWrap(ident); + Mockito.when(wrapper.tryAcquire()).thenReturn(true); Mockito.doReturn(catalog).when(wrapper).catalog(); Mockito.doReturn(capability).when(wrapper).capabilities(); Mockito.doReturn(unsupportedResult).when(capability).managedStorage(any()); @@ -1141,6 +1144,7 @@ public void testForceDropCatalog() throws Exception { Capability capability = Mockito.mock(Capability.class); CapabilityResult unsupportedResult = CapabilityResult.unsupported("Not managed"); Mockito.doReturn(catalogWrapper).when(catalogManager).loadCatalogAndWrap(ident); + Mockito.when(catalogWrapper.tryAcquire()).thenReturn(true); Mockito.doReturn(capability).when(catalogWrapper).capabilities(); Mockito.doReturn(unsupportedResult).when(capability).managedStorage(any()); Mockito.doReturn(catalog).when(catalogWrapper).catalog(); @@ -1175,6 +1179,7 @@ void testDropCatalogInvalidatesCacheAfterStoreDelete() throws Exception { Capability capability = Mockito.mock(Capability.class); CapabilityResult unsupportedResult = CapabilityResult.unsupported("Not managed"); Mockito.doReturn(catalogWrapper).when(catalogManager).loadCatalogAndWrap(ident); + Mockito.when(catalogWrapper.tryAcquire()).thenReturn(true); Mockito.doReturn(catalog).when(catalogWrapper).catalog(); Mockito.doReturn(capability).when(catalogWrapper).capabilities(); Mockito.doReturn(unsupportedResult).when(capability).managedStorage(any()); @@ -1222,20 +1227,23 @@ void testDropCatalogReloadsClosedCachedWrapper() throws Exception { void testLoadCatalogAndWrapDoesNotInvalidateConcurrentlyReloadedWrapper() { NameIdentifier ident = NameIdentifier.of("metalake", "concurrent_cache_reload_test"); - CatalogManager.CatalogWrapper closedWrapper = Mockito.mock(CatalogManager.CatalogWrapper.class); + CatalogManager.CatalogWrapper retiredWrapper = + Mockito.mock(CatalogManager.CatalogWrapper.class); CatalogManager.CatalogWrapper freshWrapper = Mockito.mock(CatalogManager.CatalogWrapper.class); BaseCatalog freshCatalog = Mockito.mock(BaseCatalog.class); Mockito.doReturn(freshCatalog).when(freshWrapper).catalog(); + // The retired wrapper reports its state only after another thread has already reloaded a fresh + // wrapper into the cache, so the stale-entry eviction must not clobber that fresh wrapper. Mockito.doAnswer( invocation -> { catalogManager.getCatalogCache().put(ident, freshWrapper); - return null; + return true; }) - .when(closedWrapper) - .catalog(); + .when(retiredWrapper) + .isRetired(); try { - catalogManager.getCatalogCache().put(ident, closedWrapper); + catalogManager.getCatalogCache().put(ident, retiredWrapper); CatalogManager.CatalogWrapper loadedWrapper = catalogManager.loadCatalogAndWrap(ident); diff --git a/core/src/test/java/org/apache/gravitino/catalog/TestCatalogWrapperLease.java b/core/src/test/java/org/apache/gravitino/catalog/TestCatalogWrapperLease.java new file mode 100644 index 00000000000..e40f20ee99d --- /dev/null +++ b/core/src/test/java/org/apache/gravitino/catalog/TestCatalogWrapperLease.java @@ -0,0 +1,326 @@ +/* + * 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 static org.awaitility.Awaitility.await; + +import com.google.common.collect.ImmutableMap; +import java.io.IOException; +import java.time.Duration; +import java.time.Instant; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import org.apache.commons.lang3.reflect.FieldUtils; +import org.apache.gravitino.Catalog; +import org.apache.gravitino.Config; +import org.apache.gravitino.Configs; +import org.apache.gravitino.GravitinoEnv; +import org.apache.gravitino.NameIdentifier; +import org.apache.gravitino.Namespace; +import org.apache.gravitino.catalog.CatalogManager.CatalogWrapper; +import org.apache.gravitino.lock.LockManager; +import org.apache.gravitino.meta.AuditInfo; +import org.apache.gravitino.meta.BaseMetalake; +import org.apache.gravitino.meta.SchemaVersion; +import org.apache.gravitino.secret.SecretManager; +import org.apache.gravitino.storage.RandomIdGenerator; +import org.apache.gravitino.storage.memory.TestMemoryEntityStore; +import org.apache.gravitino.storage.memory.TestMemoryEntityStore.InMemoryEntityStore; +import org.apache.gravitino.storage.relational.po.cache.EntityChangeRecord; +import org.apache.gravitino.storage.relational.po.cache.OperateType; +import org.apache.gravitino.utils.ClassLoaderPool; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +/** + * Tests that a catalog wrapper evicted from the catalog cache is not torn down while an operation + * is still using it. Cache eviction (expiry, explicit invalidation, remote change-log invalidation, + * and drop) must only retire the wrapper; the catalog and the ClassLoader are cleaned up when the + * last lease is released, exactly once. + */ +public class TestCatalogWrapperLease { + + private static final String METALAKE = "metalake"; + private static final String PROVIDER = "test"; + private static final Map PROPS = + ImmutableMap.of("key1", "value1", "key2", "value2", "key5-1", "value3"); + + private static Config config; + private static InMemoryEntityStore entityStore; + + private CatalogManager catalogManager; + + private static final BaseMetalake METALAKE_ENTITY = + BaseMetalake.builder() + .withId(1L) + .withName(METALAKE) + .withAuditInfo( + AuditInfo.builder().withCreator("test").withCreateTime(Instant.now()).build()) + .withVersion(SchemaVersion.V_0_1) + .build(); + + @BeforeAll + public static void setUp() throws IOException, IllegalAccessException { + config = new Config(false) {}; + config.set(Configs.CATALOG_LOAD_ISOLATED, false); + + entityStore = new TestMemoryEntityStore.InMemoryEntityStore(); + entityStore.initialize(config); + entityStore.put(METALAKE_ENTITY, true); + + FieldUtils.writeField(GravitinoEnv.getInstance(), "lockManager", new LockManager(config), true); + } + + @AfterAll + public static void tearDown() throws IOException { + if (entityStore != null) { + entityStore.close(); + entityStore = null; + } + } + + @BeforeEach + public void beforeEach() { + catalogManager = + new CatalogManager(config, entityStore, new RandomIdGenerator(), new SecretManager(config)); + } + + @AfterEach + public void afterEach() throws IOException { + if (catalogManager != null) { + catalogManager.close(); + catalogManager = null; + } + entityStore.clear(); + entityStore.put(METALAKE_ENTITY, true); + } + + @Test + public void testCacheInvalidationDefersCleanupUntilLeaseIsReleased() throws Exception { + NameIdentifier ident = createCatalog("invalidate_with_lease"); + + CatalogLease lease = catalogManager.acquireCatalogLease(ident); + CatalogWrapper wrapper = lease.wrapper(); + Assertions.assertEquals(1, wrapper.activeOperations()); + + catalogManager.getCatalogCache().invalidate(ident); + // Caffeine runs the removal listener asynchronously, so wait for the retirement to land. + await().atMost(Duration.ofSeconds(10)).until(wrapper::isRetired); + + Assertions.assertTrue(wrapper.isRetired(), "eviction must retire the wrapper"); + Assertions.assertNotNull( + wrapper.catalog(), "a leased wrapper must not be closed by cache eviction"); + // The leased wrapper is still fully usable: this is the operation that used to fail with an + // NPE (or NoClassDefFoundError) once the removal listener closed the wrapper underneath it. + Assertions.assertDoesNotThrow( + () -> + wrapper.doWithSchemaOps(ops -> ops.listSchemas(Namespace.of(METALAKE, ident.name())))); + + lease.close(); + // Closing the same lease twice must not double-release the active-operation count. + lease.close(); + + Assertions.assertEquals(0, wrapper.activeOperations()); + Assertions.assertNull(wrapper.catalog(), "the last lease release must clean up the catalog"); + } + + @Test + public void testCacheExpiryDefersCleanupUntilLeaseIsReleased() throws Exception { + Config expiringConfig = new Config(false) {}; + expiringConfig.set(Configs.CATALOG_LOAD_ISOLATED, false); + expiringConfig.set(Configs.CATALOG_CACHE_EVICTION_INTERVAL_MS, 1L); + + CatalogManager expiringManager = + new CatalogManager( + expiringConfig, entityStore, new RandomIdGenerator(), new SecretManager(config)); + try { + NameIdentifier ident = NameIdentifier.of(METALAKE, "expiring_catalog"); + expiringManager.createCatalog(ident, Catalog.Type.RELATIONAL, PROVIDER, "comment", PROPS); + + try (CatalogLease lease = expiringManager.acquireCatalogLease(ident)) { + CatalogWrapper wrapper = lease.wrapper(); + + await().atMost(Duration.ofSeconds(10)).until(wrapper::isRetired); + + Assertions.assertNull(expiringManager.getCatalogCache().getIfPresent(ident)); + Assertions.assertNotNull( + wrapper.catalog(), "a leased wrapper must survive cache expiration"); + Assertions.assertDoesNotThrow( + () -> + wrapper.doWithSchemaOps( + ops -> ops.listSchemas(Namespace.of(METALAKE, ident.name())))); + } + } finally { + expiringManager.close(); + } + } + + @Test + public void testRemoteChangeLogInvalidationDefersCleanupUntilLeaseIsReleased() throws Exception { + NameIdentifier ident = createCatalog("remote_invalidation"); + + try (CatalogLease lease = catalogManager.acquireCatalogLease(ident)) { + CatalogWrapper wrapper = lease.wrapper(); + + new CatalogChangeLogListener(catalogManager) + .onEntityChange( + List.of( + new EntityChangeRecord( + 1L, + METALAKE, + "CATALOG", + METALAKE + "." + ident.name(), + OperateType.ALTER, + 0L))); + + Assertions.assertNull(catalogManager.getCatalogCache().getIfPresent(ident)); + await().atMost(Duration.ofSeconds(10)).until(wrapper::isRetired); + Assertions.assertTrue(wrapper.isRetired()); + Assertions.assertNotNull( + wrapper.catalog(), "a leased wrapper must survive remote change-log invalidation"); + Assertions.assertDoesNotThrow( + () -> + wrapper.doWithSchemaOps( + ops -> ops.listSchemas(Namespace.of(METALAKE, ident.name())))); + } + } + + @Test + public void testDropCatalogDefersCleanupUntilLeaseIsReleased() throws Exception { + NameIdentifier ident = createCatalog("drop_with_lease"); + + CatalogLease lease = catalogManager.acquireCatalogLease(ident); + CatalogWrapper wrapper = lease.wrapper(); + + catalogManager.disableCatalog(ident); + Assertions.assertTrue(catalogManager.dropCatalog(ident)); + + Assertions.assertNotNull(wrapper.catalog(), "a leased wrapper must survive a concurrent drop"); + // Caffeine runs the removal listener asynchronously, so wait for the retirement to land before + // releasing the lease that defers the cleanup. + await().atMost(Duration.ofSeconds(10)).until(wrapper::isRetired); + lease.close(); + Assertions.assertNull(wrapper.catalog()); + } + + @Test + public void testAcquireLeaseReloadsRetiredWrapper() throws Exception { + NameIdentifier ident = createCatalog("retired_reload"); + + CatalogWrapper retiredWrapper = catalogManager.getCatalogCache().getIfPresent(ident); + Assertions.assertNotNull(retiredWrapper); + retiredWrapper.retire(); + + try (CatalogLease lease = catalogManager.acquireCatalogLease(ident)) { + Assertions.assertNotSame( + retiredWrapper, lease.wrapper(), "a retired wrapper must not be leased again"); + Assertions.assertFalse(lease.wrapper().isRetired()); + Assertions.assertNotNull(lease.catalog()); + } + } + + @Test + public void testCleanupRunsExactlyOnceForRepeatedRetireAndRelease() throws Exception { + // Two catalogs of the same provider share one pooled ClassLoader, so the pool entry survives + // as long as exactly one reference is released per wrapper. A double cleanup would drop the + // reference count to zero and destroy the ClassLoader the second catalog is still using. + NameIdentifier ident1 = createCatalog("exactly_once_1"); + createCatalog("exactly_once_2"); + + ClassLoaderPool pool = + (ClassLoaderPool) FieldUtils.readField(catalogManager, "classLoaderPool", true); + Assertions.assertEquals(1, pool.size(), "same-provider catalogs share one pooled entry"); + + CatalogLease lease = catalogManager.acquireCatalogLease(ident1); + CatalogWrapper wrapper = lease.wrapper(); + + // Repeated retirements (eviction + explicit invalidation + close) and a lease release must + // together release the pooled ClassLoader reference exactly once. + wrapper.retire(); + wrapper.retire(); + wrapper.close(); + lease.close(); + wrapper.retire(); + + Assertions.assertNull(wrapper.catalog()); + Assertions.assertEquals( + 1, pool.size(), "the pooled ClassLoader of the second catalog must stay alive"); + } + + @Test + public void testConcurrentEvictionDoesNotCloseCatalogUnderRunningOperation() throws Exception { + NameIdentifier ident = createCatalog("concurrent_eviction"); + + CountDownLatch leaseAcquired = new CountDownLatch(1); + CountDownLatch evicted = new CountDownLatch(1); + ExecutorService executor = Executors.newSingleThreadExecutor(); + try { + Future operation = + executor.submit( + () -> { + try (CatalogLease lease = catalogManager.acquireCatalogLease(ident)) { + leaseAcquired.countDown(); + Assertions.assertTrue(evicted.await(10, TimeUnit.SECONDS)); + // Runs after the wrapper has been evicted from the cache; without a lease the + // wrapper's catalog would already be closed here. + lease + .wrapper() + .doWithSchemaOps( + ops -> ops.listSchemas(Namespace.of(METALAKE, ident.name()))); + return true; + } + }); + + Assertions.assertTrue(leaseAcquired.await(10, TimeUnit.SECONDS)); + CatalogWrapper leasedWrapper = catalogManager.getCatalogCache().getIfPresent(ident); + Assertions.assertNotNull(leasedWrapper); + catalogManager.getCatalogCache().invalidate(ident); + await().atMost(Duration.ofSeconds(10)).until(leasedWrapper::isRetired); + evicted.countDown(); + + Assertions.assertTrue(operation.get(10, TimeUnit.SECONDS)); + } finally { + executor.shutdownNow(); + } + } + + @Test + public void testReleaseWithoutAcquireIsRejected() throws Exception { + NameIdentifier ident = createCatalog("release_without_acquire"); + CatalogWrapper wrapper = catalogManager.getCatalogCache().getIfPresent(ident); + Assertions.assertNotNull(wrapper); + + Assertions.assertThrows(IllegalStateException.class, wrapper::release); + } + + private NameIdentifier createCatalog(String name) { + NameIdentifier ident = NameIdentifier.of(METALAKE, name); + catalogManager.createCatalog(ident, Catalog.Type.RELATIONAL, PROVIDER, "comment", PROPS); + return ident; + } +} diff --git a/core/src/test/java/org/apache/gravitino/catalog/TestFunctionOperationDispatcher.java b/core/src/test/java/org/apache/gravitino/catalog/TestFunctionOperationDispatcher.java index 41caa1303f2..21909901e16 100644 --- a/core/src/test/java/org/apache/gravitino/catalog/TestFunctionOperationDispatcher.java +++ b/core/src/test/java/org/apache/gravitino/catalog/TestFunctionOperationDispatcher.java @@ -59,8 +59,14 @@ public void setUp() { when(catalogManager.loadCatalogAndWrap(NameIdentifier.of(METALAKE, ICEBERG_CATALOG))) .thenReturn(icebergWrapper); + when(catalogManager.acquireCatalogLease(NameIdentifier.of(METALAKE, ICEBERG_CATALOG))) + .thenAnswer(invocation -> CatalogLease.of(icebergWrapper)); + when(icebergWrapper.tryAcquire()).thenReturn(true); when(catalogManager.loadCatalogAndWrap(NameIdentifier.of(METALAKE, HIVE_CATALOG))) .thenReturn(hiveWrapper); + when(catalogManager.acquireCatalogLease(NameIdentifier.of(METALAKE, HIVE_CATALOG))) + .thenAnswer(invocation -> CatalogLease.of(hiveWrapper)); + when(hiveWrapper.tryAcquire()).thenReturn(true); dispatcher = new FunctionOperationDispatcher( diff --git a/core/src/test/java/org/apache/gravitino/catalog/TestPartitionNormalizeDispatcher.java b/core/src/test/java/org/apache/gravitino/catalog/TestPartitionNormalizeDispatcher.java index 84c9678a668..6a9a3ebc3e1 100644 --- a/core/src/test/java/org/apache/gravitino/catalog/TestPartitionNormalizeDispatcher.java +++ b/core/src/test/java/org/apache/gravitino/catalog/TestPartitionNormalizeDispatcher.java @@ -113,6 +113,9 @@ public void testAddPartitionListPartitionsGetPartitionRoundTrip() throws Excepti .thenReturn(TestCapabilityHelpers.QUOTE_AWARE_CAPABILITY); Mockito.when(mockCatalogManager.loadCatalogAndWrap(Mockito.any(NameIdentifier.class))) .thenReturn(mockWrapper); + Mockito.when(mockCatalogManager.acquireCatalogLease(Mockito.any(NameIdentifier.class))) + .thenAnswer(invocation -> CatalogLease.of(mockWrapper)); + Mockito.when(mockWrapper.tryAcquire()).thenReturn(true); PartitionNormalizeDispatcher dispatcher = new PartitionNormalizeDispatcher(mockDispatcher, mockCatalogManager); diff --git a/core/src/test/java/org/apache/gravitino/catalog/TestTableNormalizeDispatcher.java b/core/src/test/java/org/apache/gravitino/catalog/TestTableNormalizeDispatcher.java index 85649cdea18..8f36c283f4b 100644 --- a/core/src/test/java/org/apache/gravitino/catalog/TestTableNormalizeDispatcher.java +++ b/core/src/test/java/org/apache/gravitino/catalog/TestTableNormalizeDispatcher.java @@ -228,6 +228,9 @@ public void testCreateTableListTablesLoadTableRoundTrip() throws Exception { .thenReturn(TestCapabilityHelpers.QUOTE_AWARE_CAPABILITY); Mockito.when(mockCatalogManager.loadCatalogAndWrap(Mockito.any(NameIdentifier.class))) .thenReturn(mockWrapper); + Mockito.when(mockCatalogManager.acquireCatalogLease(Mockito.any(NameIdentifier.class))) + .thenAnswer(invocation -> CatalogLease.of(mockWrapper)); + Mockito.when(mockWrapper.tryAcquire()).thenReturn(true); TableNormalizeDispatcher dispatcher = new TableNormalizeDispatcher(mockDispatcher, mockCatalogManager); diff --git a/core/src/test/java/org/apache/gravitino/hook/TestFilesetHookDispatcher.java b/core/src/test/java/org/apache/gravitino/hook/TestFilesetHookDispatcher.java index 943444019e7..45f8b8d0640 100644 --- a/core/src/test/java/org/apache/gravitino/hook/TestFilesetHookDispatcher.java +++ b/core/src/test/java/org/apache/gravitino/hook/TestFilesetHookDispatcher.java @@ -54,6 +54,7 @@ import org.apache.gravitino.authorization.AccessControlManager; import org.apache.gravitino.authorization.Owner; import org.apache.gravitino.authorization.OwnerDispatcher; +import org.apache.gravitino.catalog.CatalogLease; import org.apache.gravitino.catalog.CatalogManager; import org.apache.gravitino.catalog.FilesetDispatcher; import org.apache.gravitino.catalog.TestFilesetOperationDispatcher; @@ -100,6 +101,9 @@ public static void initialize() throws Exception { Mockito.when(catalogWrapper.capabilities()).thenReturn(Capability.DEFAULT); Mockito.when(catalogManager.loadCatalog(any())).thenReturn(catalog); Mockito.when(catalogManager.loadCatalogAndWrap(any())).thenReturn(catalogWrapper); + Mockito.when(catalogManager.acquireCatalogLease(any())) + .thenAnswer(invocation -> CatalogLease.of(catalogWrapper)); + Mockito.when(catalogWrapper.tryAcquire()).thenReturn(true); authorizationPlugin = Mockito.mock(AuthorizationPlugin.class); Mockito.when(catalog.getAuthorizationPlugin()).thenReturn(authorizationPlugin); } @@ -115,6 +119,9 @@ public void testCreateFilesetSetsOwnerWithNormalizedIdentifier() throws Exceptio CatalogManager.CatalogWrapper mockWrapper = Mockito.mock(CatalogManager.CatalogWrapper.class); Mockito.when(mockWrapper.capabilities()).thenReturn(new CaseInsensitiveCapability()); Mockito.when(mockCatalogManager.loadCatalogAndWrap(any())).thenReturn(mockWrapper); + Mockito.when(mockCatalogManager.acquireCatalogLease(any())) + .thenAnswer(invocation -> CatalogLease.of(mockWrapper)); + Mockito.when(mockWrapper.tryAcquire()).thenReturn(true); OwnerDispatcher mockOwnerDispatcher = Mockito.mock(OwnerDispatcher.class); FilesetDispatcher mockFilesetDispatcher = Mockito.mock(FilesetDispatcher.class); diff --git a/core/src/test/java/org/apache/gravitino/hook/TestFunctionHookDispatcher.java b/core/src/test/java/org/apache/gravitino/hook/TestFunctionHookDispatcher.java index faba932cd6a..ef378fa14ff 100644 --- a/core/src/test/java/org/apache/gravitino/hook/TestFunctionHookDispatcher.java +++ b/core/src/test/java/org/apache/gravitino/hook/TestFunctionHookDispatcher.java @@ -31,6 +31,7 @@ import org.apache.gravitino.auth.AuthConstants; import org.apache.gravitino.authorization.Owner; import org.apache.gravitino.authorization.OwnerDispatcher; +import org.apache.gravitino.catalog.CatalogLease; import org.apache.gravitino.catalog.CatalogManager; import org.apache.gravitino.catalog.FunctionDispatcher; import org.apache.gravitino.connector.capability.Capability; @@ -64,6 +65,9 @@ public void testRegisterFunctionSetOwnerAfterRegister() throws Exception { Mockito.mock(CatalogManager.CatalogWrapper.class); Mockito.when(catalogWrapper.capabilities()).thenReturn(Capability.DEFAULT); Mockito.when(catalogManager.loadCatalogAndWrap(any())).thenReturn(catalogWrapper); + Mockito.when(catalogManager.acquireCatalogLease(any())) + .thenAnswer(invocation -> CatalogLease.of(catalogWrapper)); + Mockito.when(catalogWrapper.tryAcquire()).thenReturn(true); Mockito.when( dispatcher.registerFunction( @@ -147,6 +151,9 @@ public void testRegisterFunctionSetsOwnerWithNormalizedIdentifier() throws Excep CatalogManager.CatalogWrapper mockWrapper = Mockito.mock(CatalogManager.CatalogWrapper.class); Mockito.when(mockWrapper.capabilities()).thenReturn(new CaseInsensitiveCapability()); Mockito.when(mockCatalogManager.loadCatalogAndWrap(any())).thenReturn(mockWrapper); + Mockito.when(mockCatalogManager.acquireCatalogLease(any())) + .thenAnswer(invocation -> CatalogLease.of(mockWrapper)); + Mockito.when(mockWrapper.tryAcquire()).thenReturn(true); OwnerDispatcher mockOwnerDispatcher = Mockito.mock(OwnerDispatcher.class); FunctionDispatcher mockFunctionDispatcher = Mockito.mock(FunctionDispatcher.class); @@ -200,6 +207,9 @@ public void testRegisterFunctionThrowsWhenSetOwnerFails() throws Exception { Mockito.mock(CatalogManager.CatalogWrapper.class); Mockito.when(catalogWrapper.capabilities()).thenReturn(Capability.DEFAULT); Mockito.when(catalogManager.loadCatalogAndWrap(any())).thenReturn(catalogWrapper); + Mockito.when(catalogManager.acquireCatalogLease(any())) + .thenAnswer(invocation -> CatalogLease.of(catalogWrapper)); + Mockito.when(catalogWrapper.tryAcquire()).thenReturn(true); FunctionDispatcher mockFunctionDispatcher = Mockito.mock(FunctionDispatcher.class); Function mockFunction = Mockito.mock(Function.class); diff --git a/core/src/test/java/org/apache/gravitino/hook/TestModelHookDispatcher.java b/core/src/test/java/org/apache/gravitino/hook/TestModelHookDispatcher.java index c6777b58919..02c4fc7ac48 100644 --- a/core/src/test/java/org/apache/gravitino/hook/TestModelHookDispatcher.java +++ b/core/src/test/java/org/apache/gravitino/hook/TestModelHookDispatcher.java @@ -32,6 +32,7 @@ import org.apache.gravitino.NameIdentifier; import org.apache.gravitino.authorization.Owner; import org.apache.gravitino.authorization.OwnerDispatcher; +import org.apache.gravitino.catalog.CatalogLease; import org.apache.gravitino.catalog.CatalogManager; import org.apache.gravitino.catalog.ModelDispatcher; import org.apache.gravitino.connector.capability.Capability; @@ -62,6 +63,9 @@ public void setUp() throws Exception { mockCatalogManager = mock(CatalogManager.class); mockCatalogWrapper = mock(CatalogManager.CatalogWrapper.class); when(mockCatalogManager.loadCatalogAndWrap(any())).thenReturn(mockCatalogWrapper); + when(mockCatalogManager.acquireCatalogLease(any())) + .thenAnswer(invocation -> CatalogLease.of(mockCatalogWrapper)); + when(mockCatalogWrapper.tryAcquire()).thenReturn(true); when(mockCatalogWrapper.capabilities()).thenReturn(Capability.DEFAULT); savedOwnerDispatcher = GravitinoEnv.getInstance().ownerDispatcher(); // Read the catalogManager field directly via reflection because the public accessor diff --git a/core/src/test/java/org/apache/gravitino/hook/TestSchemaHookDispatcher.java b/core/src/test/java/org/apache/gravitino/hook/TestSchemaHookDispatcher.java index 3e05c3c83ce..f515d977732 100644 --- a/core/src/test/java/org/apache/gravitino/hook/TestSchemaHookDispatcher.java +++ b/core/src/test/java/org/apache/gravitino/hook/TestSchemaHookDispatcher.java @@ -45,6 +45,7 @@ import org.apache.gravitino.authorization.AuthorizationUtils; import org.apache.gravitino.authorization.Owner; import org.apache.gravitino.authorization.OwnerDispatcher; +import org.apache.gravitino.catalog.CatalogLease; import org.apache.gravitino.catalog.CatalogManager; import org.apache.gravitino.catalog.SchemaDispatcher; import org.apache.gravitino.connector.capability.Capability; @@ -78,6 +79,9 @@ public void setUp() throws Exception { mockCatalogManager = mock(CatalogManager.class); mockCatalogWrapper = mock(CatalogManager.CatalogWrapper.class); when(mockCatalogManager.loadCatalogAndWrap(any())).thenReturn(mockCatalogWrapper); + when(mockCatalogManager.acquireCatalogLease(any())) + .thenAnswer(invocation -> CatalogLease.of(mockCatalogWrapper)); + when(mockCatalogWrapper.tryAcquire()).thenReturn(true); when(mockCatalogWrapper.capabilities()).thenReturn(Capability.DEFAULT); savedOwnerDispatcher = GravitinoEnv.getInstance().ownerDispatcher(); // Tests in this class that rely on the singleton catalogManager always go through diff --git a/core/src/test/java/org/apache/gravitino/hook/TestTableHookDispatcher.java b/core/src/test/java/org/apache/gravitino/hook/TestTableHookDispatcher.java index bf9431a504c..50ffd6bc7b9 100644 --- a/core/src/test/java/org/apache/gravitino/hook/TestTableHookDispatcher.java +++ b/core/src/test/java/org/apache/gravitino/hook/TestTableHookDispatcher.java @@ -34,6 +34,7 @@ import org.apache.gravitino.authorization.AuthorizationUtils; import org.apache.gravitino.authorization.Owner; import org.apache.gravitino.authorization.OwnerDispatcher; +import org.apache.gravitino.catalog.CatalogLease; import org.apache.gravitino.catalog.CatalogManager; import org.apache.gravitino.catalog.TableDispatcher; import org.apache.gravitino.connector.capability.Capability; @@ -85,6 +86,9 @@ public void testCreateTableSetsOwnerWithNormalizedIdentifier() throws Exception CatalogManager.CatalogWrapper wrapper = Mockito.mock(CatalogManager.CatalogWrapper.class); Mockito.when(wrapper.capabilities()).thenReturn(new CaseInsensitiveCapability()); Mockito.when(catalogManager.loadCatalogAndWrap(any())).thenReturn(wrapper); + Mockito.when(catalogManager.acquireCatalogLease(any())) + .thenAnswer(invocation -> CatalogLease.of(wrapper)); + Mockito.when(wrapper.tryAcquire()).thenReturn(true); OwnerDispatcher ownerDispatcher = Mockito.mock(OwnerDispatcher.class); TableDispatcher dispatcher = Mockito.mock(TableDispatcher.class); @@ -151,6 +155,9 @@ public void testCreateTableThrowsWhenSetOwnerFails() throws Exception { CatalogManager.CatalogWrapper wrapper = Mockito.mock(CatalogManager.CatalogWrapper.class); Mockito.when(wrapper.capabilities()).thenReturn(Capability.DEFAULT); Mockito.when(catalogManager.loadCatalogAndWrap(any())).thenReturn(wrapper); + Mockito.when(catalogManager.acquireCatalogLease(any())) + .thenAnswer(invocation -> CatalogLease.of(wrapper)); + Mockito.when(wrapper.tryAcquire()).thenReturn(true); TableHookDispatcher hook = new TableHookDispatcher(dispatcher, () -> ownerDispatcher, catalogManager); diff --git a/core/src/test/java/org/apache/gravitino/hook/TestTopicHookDispatcher.java b/core/src/test/java/org/apache/gravitino/hook/TestTopicHookDispatcher.java index 539af87f5ab..dfb16ecd61b 100644 --- a/core/src/test/java/org/apache/gravitino/hook/TestTopicHookDispatcher.java +++ b/core/src/test/java/org/apache/gravitino/hook/TestTopicHookDispatcher.java @@ -31,6 +31,7 @@ import org.apache.gravitino.authorization.AccessControlManager; import org.apache.gravitino.authorization.Owner; import org.apache.gravitino.authorization.OwnerDispatcher; +import org.apache.gravitino.catalog.CatalogLease; import org.apache.gravitino.catalog.CatalogManager; import org.apache.gravitino.catalog.TestOperationDispatcher; import org.apache.gravitino.catalog.TestTopicOperationDispatcher; @@ -74,6 +75,9 @@ public static void initialize() throws Exception { Mockito.when(catalogWrapper.capabilities()).thenReturn(Capability.DEFAULT); Mockito.when(catalogManager.loadCatalog(any())).thenReturn(catalog); Mockito.when(catalogManager.loadCatalogAndWrap(any())).thenReturn(catalogWrapper); + Mockito.when(catalogManager.acquireCatalogLease(any())) + .thenAnswer(invocation -> CatalogLease.of(catalogWrapper)); + Mockito.when(catalogWrapper.tryAcquire()).thenReturn(true); authorizationPlugin = Mockito.mock(AuthorizationPlugin.class); Mockito.when(catalog.getAuthorizationPlugin()).thenReturn(authorizationPlugin); } @@ -89,6 +93,9 @@ public void testCreateTopicSetsOwnerWithNormalizedIdentifier() throws Exception CatalogManager.CatalogWrapper mockWrapper = Mockito.mock(CatalogManager.CatalogWrapper.class); Mockito.when(mockWrapper.capabilities()).thenReturn(new CaseInsensitiveCapability()); Mockito.when(mockCatalogManager.loadCatalogAndWrap(any())).thenReturn(mockWrapper); + Mockito.when(mockCatalogManager.acquireCatalogLease(any())) + .thenAnswer(invocation -> CatalogLease.of(mockWrapper)); + Mockito.when(mockWrapper.tryAcquire()).thenReturn(true); OwnerDispatcher mockOwnerDispatcher = Mockito.mock(OwnerDispatcher.class); TopicDispatcher mockTopicDispatcher = Mockito.mock(TopicDispatcher.class); diff --git a/iceberg/iceberg-rest-server/src/main/java/org/apache/gravitino/iceberg/service/dispatcher/IcebergCleanupHelper.java b/iceberg/iceberg-rest-server/src/main/java/org/apache/gravitino/iceberg/service/dispatcher/IcebergCleanupHelper.java index b01f30eca09..3bcb41b7e2d 100644 --- a/iceberg/iceberg-rest-server/src/main/java/org/apache/gravitino/iceberg/service/dispatcher/IcebergCleanupHelper.java +++ b/iceberg/iceberg-rest-server/src/main/java/org/apache/gravitino/iceberg/service/dispatcher/IcebergCleanupHelper.java @@ -22,6 +22,7 @@ import java.util.Optional; import org.apache.gravitino.GravitinoEnv; import org.apache.gravitino.NameIdentifier; +import org.apache.gravitino.catalog.CatalogLease; import org.apache.gravitino.iceberg.service.authorization.IcebergRESTServerContext; import org.apache.gravitino.iceberg.service.cleanup.IcebergCleanupManager; import org.apache.iceberg.catalog.Namespace; @@ -42,12 +43,12 @@ private IcebergCleanupHelper() {} */ static long catalogId(String catalogName) { String metalake = IcebergRESTServerContext.getInstance().metalakeName(); - return GravitinoEnv.getInstance() - .catalogManager() - .loadCatalogAndWrap(NameIdentifier.of(metalake, catalogName)) - .catalog() - .entity() - .id(); + try (CatalogLease lease = + GravitinoEnv.getInstance() + .catalogManager() + .acquireCatalogLease(NameIdentifier.of(metalake, catalogName))) { + return lease.catalog().entity().id(); + } } /** diff --git a/iceberg/iceberg-rest-server/src/test/java/org/apache/gravitino/iceberg/service/dispatcher/TestIcebergAsyncPurge.java b/iceberg/iceberg-rest-server/src/test/java/org/apache/gravitino/iceberg/service/dispatcher/TestIcebergAsyncPurge.java index 3ee6bf69071..0b64ed2c175 100644 --- a/iceberg/iceberg-rest-server/src/test/java/org/apache/gravitino/iceberg/service/dispatcher/TestIcebergAsyncPurge.java +++ b/iceberg/iceberg-rest-server/src/test/java/org/apache/gravitino/iceberg/service/dispatcher/TestIcebergAsyncPurge.java @@ -32,6 +32,7 @@ import java.util.Optional; import org.apache.gravitino.GravitinoEnv; import org.apache.gravitino.auth.AuthConstants; +import org.apache.gravitino.catalog.CatalogLease; import org.apache.gravitino.catalog.CatalogManager; import org.apache.gravitino.connector.BaseCatalog; import org.apache.gravitino.iceberg.service.CatalogWrapperForREST; @@ -241,6 +242,9 @@ private static MockedStatic mockCatalogId() { envStatic.when(GravitinoEnv::getInstance).thenReturn(env); when(env.catalogManager()).thenReturn(catalogManager); when(catalogManager.loadCatalogAndWrap(any())).thenReturn(wrapper); + when(catalogManager.acquireCatalogLease(any())) + .thenAnswer(invocation -> CatalogLease.of(wrapper)); + when(wrapper.tryAcquire()).thenReturn(true); when(wrapper.catalog()).thenReturn(catalog); when(catalog.entity()).thenReturn(entity); when(entity.id()).thenReturn(CATALOG_ID); From 5223b66f1b25952fa69c0d1dc65749a931ee92f7 Mon Sep 17 00:00:00 2001 From: yuqi Date: Mon, 10 Aug 2026 18:49:18 +0800 Subject: [PATCH 2/6] [#12403] refactor(core): keep the catalog lease API package-private Drop CatalogLease.of(), which existed only so tests with a mocked CatalogManager could hand out a lease, and narrow tryAcquire/release/retire/ isRetired to package-private: every production caller lives in the catalog package. Tests outside that package now build leases through a test-only CatalogTestUtils helper. Also document that the deferred cleanup runs on the thread that releases the last lease, which is a request thread when the wrapper was evicted while an operation was in flight. --- .../gravitino/catalog/CatalogLease.java | 16 -------- .../gravitino/catalog/CatalogManager.java | 13 ++++-- .../gravitino/catalog/CatalogTestUtils.java | 40 +++++++++++++++++++ .../TestFunctionOperationDispatcher.java | 6 +-- .../TestPartitionNormalizeDispatcher.java | 3 +- .../catalog/TestTableNormalizeDispatcher.java | 3 +- .../hook/TestFilesetHookDispatcher.java | 8 ++-- .../hook/TestFunctionHookDispatcher.java | 11 ++--- .../hook/TestModelHookDispatcher.java | 5 +-- .../hook/TestSchemaHookDispatcher.java | 5 +-- .../hook/TestTableHookDispatcher.java | 8 ++-- .../hook/TestTopicHookDispatcher.java | 8 ++-- .../dispatcher/TestIcebergAsyncPurge.java | 5 +-- 13 files changed, 72 insertions(+), 59 deletions(-) create mode 100644 core/src/test/java/org/apache/gravitino/catalog/CatalogTestUtils.java diff --git a/core/src/main/java/org/apache/gravitino/catalog/CatalogLease.java b/core/src/main/java/org/apache/gravitino/catalog/CatalogLease.java index ad0e687bd9b..93af7d692c1 100644 --- a/core/src/main/java/org/apache/gravitino/catalog/CatalogLease.java +++ b/core/src/main/java/org/apache/gravitino/catalog/CatalogLease.java @@ -18,7 +18,6 @@ */ package org.apache.gravitino.catalog; -import com.google.common.base.Preconditions; import java.util.concurrent.atomic.AtomicBoolean; import org.apache.gravitino.catalog.CatalogManager.CatalogWrapper; import org.apache.gravitino.connector.BaseCatalog; @@ -50,21 +49,6 @@ public final class CatalogLease implements AutoCloseable { this.wrapper = wrapper; } - /** - * Takes a lease on a wrapper the caller already holds. Prefer {@link - * CatalogManager#acquireCatalogLease(org.apache.gravitino.NameIdentifier)}, which reloads a fresh - * wrapper when the cached one has been retired. - * - * @param wrapper the wrapper to lease. - * @return a lease on the given wrapper. - * @throws IllegalStateException if the wrapper has already been retired. - */ - public static CatalogLease of(CatalogWrapper wrapper) { - Preconditions.checkState( - wrapper.tryAcquire(), "Catalog wrapper has already been retired, cannot lease it"); - return new CatalogLease(wrapper); - } - /** * Returns the leased catalog wrapper. * diff --git a/core/src/main/java/org/apache/gravitino/catalog/CatalogManager.java b/core/src/main/java/org/apache/gravitino/catalog/CatalogManager.java index 2bd78315359..ed841b94bb3 100644 --- a/core/src/main/java/org/apache/gravitino/catalog/CatalogManager.java +++ b/core/src/main/java/org/apache/gravitino/catalog/CatalogManager.java @@ -215,7 +215,7 @@ public BaseCatalog catalog() { * @return true if the lease was taken, false if the wrapper has already been retired and the * caller must load a fresh wrapper. */ - public boolean tryAcquire() { + boolean tryAcquire() { synchronized (leaseLock) { if (retired) { return false; @@ -228,8 +228,13 @@ public boolean tryAcquire() { /** * Releases a lease taken by {@link #tryAcquire()}. Cleans up the catalog and the ClassLoader if * this was the last lease on an already retired wrapper. + * + *

Note that in that case the cleanup runs on the releasing thread, which is usually a + * request thread, so a slow catalog close is charged to that request. This only happens when + * the wrapper was evicted while the operation was in flight; the common case is that eviction + * finds no lease and cleans up on the cache's own thread. */ - public void release() { + void release() { boolean shouldCleanup; synchronized (leaseLock) { Preconditions.checkState(activeOps > 0, "Releasing a lease that was never acquired"); @@ -246,7 +251,7 @@ public void release() { * Retires this wrapper: no new lease can be taken. The catalog and the ClassLoader are cleaned * up immediately if no operation is in flight, otherwise by the last {@link #release()}. */ - public void retire() { + void retire() { boolean shouldCleanup; synchronized (leaseLock) { retired = true; @@ -263,7 +268,7 @@ public void retire() { * * @return true if the wrapper has been retired. */ - public boolean isRetired() { + boolean isRetired() { synchronized (leaseLock) { return retired; } diff --git a/core/src/test/java/org/apache/gravitino/catalog/CatalogTestUtils.java b/core/src/test/java/org/apache/gravitino/catalog/CatalogTestUtils.java new file mode 100644 index 00000000000..c52a9ee4711 --- /dev/null +++ b/core/src/test/java/org/apache/gravitino/catalog/CatalogTestUtils.java @@ -0,0 +1,40 @@ +/* + * 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 org.apache.gravitino.catalog.CatalogManager.CatalogWrapper; + +/** Test-only helpers for catalog internals that are package-private in production code. */ +public final class CatalogTestUtils { + + private CatalogTestUtils() {} + + /** + * Wraps a mocked {@link CatalogWrapper} into a {@link CatalogLease} without going through the + * wrapper's active-operation counting, for tests that stub {@link + * CatalogManager#acquireCatalogLease} on a mocked manager. Production code must obtain leases + * from {@link CatalogManager#acquireCatalogLease}. + * + * @param wrapper the (usually mocked) wrapper to hand out. + * @return a lease over the given wrapper. + */ + public static CatalogLease unmanagedLease(CatalogWrapper wrapper) { + return new CatalogLease(wrapper); + } +} diff --git a/core/src/test/java/org/apache/gravitino/catalog/TestFunctionOperationDispatcher.java b/core/src/test/java/org/apache/gravitino/catalog/TestFunctionOperationDispatcher.java index 21909901e16..cd508ce04af 100644 --- a/core/src/test/java/org/apache/gravitino/catalog/TestFunctionOperationDispatcher.java +++ b/core/src/test/java/org/apache/gravitino/catalog/TestFunctionOperationDispatcher.java @@ -60,13 +60,11 @@ public void setUp() { when(catalogManager.loadCatalogAndWrap(NameIdentifier.of(METALAKE, ICEBERG_CATALOG))) .thenReturn(icebergWrapper); when(catalogManager.acquireCatalogLease(NameIdentifier.of(METALAKE, ICEBERG_CATALOG))) - .thenAnswer(invocation -> CatalogLease.of(icebergWrapper)); - when(icebergWrapper.tryAcquire()).thenReturn(true); + .thenAnswer(invocation -> CatalogTestUtils.unmanagedLease(icebergWrapper)); when(catalogManager.loadCatalogAndWrap(NameIdentifier.of(METALAKE, HIVE_CATALOG))) .thenReturn(hiveWrapper); when(catalogManager.acquireCatalogLease(NameIdentifier.of(METALAKE, HIVE_CATALOG))) - .thenAnswer(invocation -> CatalogLease.of(hiveWrapper)); - when(hiveWrapper.tryAcquire()).thenReturn(true); + .thenAnswer(invocation -> CatalogTestUtils.unmanagedLease(hiveWrapper)); dispatcher = new FunctionOperationDispatcher( diff --git a/core/src/test/java/org/apache/gravitino/catalog/TestPartitionNormalizeDispatcher.java b/core/src/test/java/org/apache/gravitino/catalog/TestPartitionNormalizeDispatcher.java index 6a9a3ebc3e1..35d02c89462 100644 --- a/core/src/test/java/org/apache/gravitino/catalog/TestPartitionNormalizeDispatcher.java +++ b/core/src/test/java/org/apache/gravitino/catalog/TestPartitionNormalizeDispatcher.java @@ -114,8 +114,7 @@ public void testAddPartitionListPartitionsGetPartitionRoundTrip() throws Excepti Mockito.when(mockCatalogManager.loadCatalogAndWrap(Mockito.any(NameIdentifier.class))) .thenReturn(mockWrapper); Mockito.when(mockCatalogManager.acquireCatalogLease(Mockito.any(NameIdentifier.class))) - .thenAnswer(invocation -> CatalogLease.of(mockWrapper)); - Mockito.when(mockWrapper.tryAcquire()).thenReturn(true); + .thenAnswer(invocation -> CatalogTestUtils.unmanagedLease(mockWrapper)); PartitionNormalizeDispatcher dispatcher = new PartitionNormalizeDispatcher(mockDispatcher, mockCatalogManager); diff --git a/core/src/test/java/org/apache/gravitino/catalog/TestTableNormalizeDispatcher.java b/core/src/test/java/org/apache/gravitino/catalog/TestTableNormalizeDispatcher.java index 8f36c283f4b..1fc71d4f5cf 100644 --- a/core/src/test/java/org/apache/gravitino/catalog/TestTableNormalizeDispatcher.java +++ b/core/src/test/java/org/apache/gravitino/catalog/TestTableNormalizeDispatcher.java @@ -229,8 +229,7 @@ public void testCreateTableListTablesLoadTableRoundTrip() throws Exception { Mockito.when(mockCatalogManager.loadCatalogAndWrap(Mockito.any(NameIdentifier.class))) .thenReturn(mockWrapper); Mockito.when(mockCatalogManager.acquireCatalogLease(Mockito.any(NameIdentifier.class))) - .thenAnswer(invocation -> CatalogLease.of(mockWrapper)); - Mockito.when(mockWrapper.tryAcquire()).thenReturn(true); + .thenAnswer(invocation -> CatalogTestUtils.unmanagedLease(mockWrapper)); TableNormalizeDispatcher dispatcher = new TableNormalizeDispatcher(mockDispatcher, mockCatalogManager); diff --git a/core/src/test/java/org/apache/gravitino/hook/TestFilesetHookDispatcher.java b/core/src/test/java/org/apache/gravitino/hook/TestFilesetHookDispatcher.java index 45f8b8d0640..c8d876e60d2 100644 --- a/core/src/test/java/org/apache/gravitino/hook/TestFilesetHookDispatcher.java +++ b/core/src/test/java/org/apache/gravitino/hook/TestFilesetHookDispatcher.java @@ -54,8 +54,8 @@ import org.apache.gravitino.authorization.AccessControlManager; import org.apache.gravitino.authorization.Owner; import org.apache.gravitino.authorization.OwnerDispatcher; -import org.apache.gravitino.catalog.CatalogLease; import org.apache.gravitino.catalog.CatalogManager; +import org.apache.gravitino.catalog.CatalogTestUtils; import org.apache.gravitino.catalog.FilesetDispatcher; import org.apache.gravitino.catalog.TestFilesetOperationDispatcher; import org.apache.gravitino.catalog.TestOperationDispatcher; @@ -102,8 +102,7 @@ public static void initialize() throws Exception { Mockito.when(catalogManager.loadCatalog(any())).thenReturn(catalog); Mockito.when(catalogManager.loadCatalogAndWrap(any())).thenReturn(catalogWrapper); Mockito.when(catalogManager.acquireCatalogLease(any())) - .thenAnswer(invocation -> CatalogLease.of(catalogWrapper)); - Mockito.when(catalogWrapper.tryAcquire()).thenReturn(true); + .thenAnswer(invocation -> CatalogTestUtils.unmanagedLease(catalogWrapper)); authorizationPlugin = Mockito.mock(AuthorizationPlugin.class); Mockito.when(catalog.getAuthorizationPlugin()).thenReturn(authorizationPlugin); } @@ -120,8 +119,7 @@ public void testCreateFilesetSetsOwnerWithNormalizedIdentifier() throws Exceptio Mockito.when(mockWrapper.capabilities()).thenReturn(new CaseInsensitiveCapability()); Mockito.when(mockCatalogManager.loadCatalogAndWrap(any())).thenReturn(mockWrapper); Mockito.when(mockCatalogManager.acquireCatalogLease(any())) - .thenAnswer(invocation -> CatalogLease.of(mockWrapper)); - Mockito.when(mockWrapper.tryAcquire()).thenReturn(true); + .thenAnswer(invocation -> CatalogTestUtils.unmanagedLease(mockWrapper)); OwnerDispatcher mockOwnerDispatcher = Mockito.mock(OwnerDispatcher.class); FilesetDispatcher mockFilesetDispatcher = Mockito.mock(FilesetDispatcher.class); diff --git a/core/src/test/java/org/apache/gravitino/hook/TestFunctionHookDispatcher.java b/core/src/test/java/org/apache/gravitino/hook/TestFunctionHookDispatcher.java index ef378fa14ff..a8dfde594cd 100644 --- a/core/src/test/java/org/apache/gravitino/hook/TestFunctionHookDispatcher.java +++ b/core/src/test/java/org/apache/gravitino/hook/TestFunctionHookDispatcher.java @@ -31,8 +31,8 @@ import org.apache.gravitino.auth.AuthConstants; import org.apache.gravitino.authorization.Owner; import org.apache.gravitino.authorization.OwnerDispatcher; -import org.apache.gravitino.catalog.CatalogLease; import org.apache.gravitino.catalog.CatalogManager; +import org.apache.gravitino.catalog.CatalogTestUtils; import org.apache.gravitino.catalog.FunctionDispatcher; import org.apache.gravitino.connector.capability.Capability; import org.apache.gravitino.connector.capability.CapabilityResult; @@ -66,8 +66,7 @@ public void testRegisterFunctionSetOwnerAfterRegister() throws Exception { Mockito.when(catalogWrapper.capabilities()).thenReturn(Capability.DEFAULT); Mockito.when(catalogManager.loadCatalogAndWrap(any())).thenReturn(catalogWrapper); Mockito.when(catalogManager.acquireCatalogLease(any())) - .thenAnswer(invocation -> CatalogLease.of(catalogWrapper)); - Mockito.when(catalogWrapper.tryAcquire()).thenReturn(true); + .thenAnswer(invocation -> CatalogTestUtils.unmanagedLease(catalogWrapper)); Mockito.when( dispatcher.registerFunction( @@ -152,8 +151,7 @@ public void testRegisterFunctionSetsOwnerWithNormalizedIdentifier() throws Excep Mockito.when(mockWrapper.capabilities()).thenReturn(new CaseInsensitiveCapability()); Mockito.when(mockCatalogManager.loadCatalogAndWrap(any())).thenReturn(mockWrapper); Mockito.when(mockCatalogManager.acquireCatalogLease(any())) - .thenAnswer(invocation -> CatalogLease.of(mockWrapper)); - Mockito.when(mockWrapper.tryAcquire()).thenReturn(true); + .thenAnswer(invocation -> CatalogTestUtils.unmanagedLease(mockWrapper)); OwnerDispatcher mockOwnerDispatcher = Mockito.mock(OwnerDispatcher.class); FunctionDispatcher mockFunctionDispatcher = Mockito.mock(FunctionDispatcher.class); @@ -208,8 +206,7 @@ public void testRegisterFunctionThrowsWhenSetOwnerFails() throws Exception { Mockito.when(catalogWrapper.capabilities()).thenReturn(Capability.DEFAULT); Mockito.when(catalogManager.loadCatalogAndWrap(any())).thenReturn(catalogWrapper); Mockito.when(catalogManager.acquireCatalogLease(any())) - .thenAnswer(invocation -> CatalogLease.of(catalogWrapper)); - Mockito.when(catalogWrapper.tryAcquire()).thenReturn(true); + .thenAnswer(invocation -> CatalogTestUtils.unmanagedLease(catalogWrapper)); FunctionDispatcher mockFunctionDispatcher = Mockito.mock(FunctionDispatcher.class); Function mockFunction = Mockito.mock(Function.class); diff --git a/core/src/test/java/org/apache/gravitino/hook/TestModelHookDispatcher.java b/core/src/test/java/org/apache/gravitino/hook/TestModelHookDispatcher.java index 02c4fc7ac48..97d9bf263a6 100644 --- a/core/src/test/java/org/apache/gravitino/hook/TestModelHookDispatcher.java +++ b/core/src/test/java/org/apache/gravitino/hook/TestModelHookDispatcher.java @@ -32,8 +32,8 @@ import org.apache.gravitino.NameIdentifier; import org.apache.gravitino.authorization.Owner; import org.apache.gravitino.authorization.OwnerDispatcher; -import org.apache.gravitino.catalog.CatalogLease; import org.apache.gravitino.catalog.CatalogManager; +import org.apache.gravitino.catalog.CatalogTestUtils; import org.apache.gravitino.catalog.ModelDispatcher; import org.apache.gravitino.connector.capability.Capability; import org.apache.gravitino.connector.capability.CapabilityResult; @@ -64,8 +64,7 @@ public void setUp() throws Exception { mockCatalogWrapper = mock(CatalogManager.CatalogWrapper.class); when(mockCatalogManager.loadCatalogAndWrap(any())).thenReturn(mockCatalogWrapper); when(mockCatalogManager.acquireCatalogLease(any())) - .thenAnswer(invocation -> CatalogLease.of(mockCatalogWrapper)); - when(mockCatalogWrapper.tryAcquire()).thenReturn(true); + .thenAnswer(invocation -> CatalogTestUtils.unmanagedLease(mockCatalogWrapper)); when(mockCatalogWrapper.capabilities()).thenReturn(Capability.DEFAULT); savedOwnerDispatcher = GravitinoEnv.getInstance().ownerDispatcher(); // Read the catalogManager field directly via reflection because the public accessor diff --git a/core/src/test/java/org/apache/gravitino/hook/TestSchemaHookDispatcher.java b/core/src/test/java/org/apache/gravitino/hook/TestSchemaHookDispatcher.java index f515d977732..539e9276cd0 100644 --- a/core/src/test/java/org/apache/gravitino/hook/TestSchemaHookDispatcher.java +++ b/core/src/test/java/org/apache/gravitino/hook/TestSchemaHookDispatcher.java @@ -45,8 +45,8 @@ import org.apache.gravitino.authorization.AuthorizationUtils; import org.apache.gravitino.authorization.Owner; import org.apache.gravitino.authorization.OwnerDispatcher; -import org.apache.gravitino.catalog.CatalogLease; import org.apache.gravitino.catalog.CatalogManager; +import org.apache.gravitino.catalog.CatalogTestUtils; import org.apache.gravitino.catalog.SchemaDispatcher; import org.apache.gravitino.connector.capability.Capability; import org.apache.gravitino.connector.capability.CapabilityResult; @@ -80,8 +80,7 @@ public void setUp() throws Exception { mockCatalogWrapper = mock(CatalogManager.CatalogWrapper.class); when(mockCatalogManager.loadCatalogAndWrap(any())).thenReturn(mockCatalogWrapper); when(mockCatalogManager.acquireCatalogLease(any())) - .thenAnswer(invocation -> CatalogLease.of(mockCatalogWrapper)); - when(mockCatalogWrapper.tryAcquire()).thenReturn(true); + .thenAnswer(invocation -> CatalogTestUtils.unmanagedLease(mockCatalogWrapper)); when(mockCatalogWrapper.capabilities()).thenReturn(Capability.DEFAULT); savedOwnerDispatcher = GravitinoEnv.getInstance().ownerDispatcher(); // Tests in this class that rely on the singleton catalogManager always go through diff --git a/core/src/test/java/org/apache/gravitino/hook/TestTableHookDispatcher.java b/core/src/test/java/org/apache/gravitino/hook/TestTableHookDispatcher.java index 50ffd6bc7b9..7d0867ae342 100644 --- a/core/src/test/java/org/apache/gravitino/hook/TestTableHookDispatcher.java +++ b/core/src/test/java/org/apache/gravitino/hook/TestTableHookDispatcher.java @@ -34,8 +34,8 @@ import org.apache.gravitino.authorization.AuthorizationUtils; import org.apache.gravitino.authorization.Owner; import org.apache.gravitino.authorization.OwnerDispatcher; -import org.apache.gravitino.catalog.CatalogLease; import org.apache.gravitino.catalog.CatalogManager; +import org.apache.gravitino.catalog.CatalogTestUtils; import org.apache.gravitino.catalog.TableDispatcher; import org.apache.gravitino.connector.capability.Capability; import org.apache.gravitino.connector.capability.CapabilityResult; @@ -87,8 +87,7 @@ public void testCreateTableSetsOwnerWithNormalizedIdentifier() throws Exception Mockito.when(wrapper.capabilities()).thenReturn(new CaseInsensitiveCapability()); Mockito.when(catalogManager.loadCatalogAndWrap(any())).thenReturn(wrapper); Mockito.when(catalogManager.acquireCatalogLease(any())) - .thenAnswer(invocation -> CatalogLease.of(wrapper)); - Mockito.when(wrapper.tryAcquire()).thenReturn(true); + .thenAnswer(invocation -> CatalogTestUtils.unmanagedLease(wrapper)); OwnerDispatcher ownerDispatcher = Mockito.mock(OwnerDispatcher.class); TableDispatcher dispatcher = Mockito.mock(TableDispatcher.class); @@ -156,8 +155,7 @@ public void testCreateTableThrowsWhenSetOwnerFails() throws Exception { Mockito.when(wrapper.capabilities()).thenReturn(Capability.DEFAULT); Mockito.when(catalogManager.loadCatalogAndWrap(any())).thenReturn(wrapper); Mockito.when(catalogManager.acquireCatalogLease(any())) - .thenAnswer(invocation -> CatalogLease.of(wrapper)); - Mockito.when(wrapper.tryAcquire()).thenReturn(true); + .thenAnswer(invocation -> CatalogTestUtils.unmanagedLease(wrapper)); TableHookDispatcher hook = new TableHookDispatcher(dispatcher, () -> ownerDispatcher, catalogManager); diff --git a/core/src/test/java/org/apache/gravitino/hook/TestTopicHookDispatcher.java b/core/src/test/java/org/apache/gravitino/hook/TestTopicHookDispatcher.java index dfb16ecd61b..d5bb231eb8f 100644 --- a/core/src/test/java/org/apache/gravitino/hook/TestTopicHookDispatcher.java +++ b/core/src/test/java/org/apache/gravitino/hook/TestTopicHookDispatcher.java @@ -31,8 +31,8 @@ import org.apache.gravitino.authorization.AccessControlManager; import org.apache.gravitino.authorization.Owner; import org.apache.gravitino.authorization.OwnerDispatcher; -import org.apache.gravitino.catalog.CatalogLease; import org.apache.gravitino.catalog.CatalogManager; +import org.apache.gravitino.catalog.CatalogTestUtils; import org.apache.gravitino.catalog.TestOperationDispatcher; import org.apache.gravitino.catalog.TestTopicOperationDispatcher; import org.apache.gravitino.catalog.TopicDispatcher; @@ -76,8 +76,7 @@ public static void initialize() throws Exception { Mockito.when(catalogManager.loadCatalog(any())).thenReturn(catalog); Mockito.when(catalogManager.loadCatalogAndWrap(any())).thenReturn(catalogWrapper); Mockito.when(catalogManager.acquireCatalogLease(any())) - .thenAnswer(invocation -> CatalogLease.of(catalogWrapper)); - Mockito.when(catalogWrapper.tryAcquire()).thenReturn(true); + .thenAnswer(invocation -> CatalogTestUtils.unmanagedLease(catalogWrapper)); authorizationPlugin = Mockito.mock(AuthorizationPlugin.class); Mockito.when(catalog.getAuthorizationPlugin()).thenReturn(authorizationPlugin); } @@ -94,8 +93,7 @@ public void testCreateTopicSetsOwnerWithNormalizedIdentifier() throws Exception Mockito.when(mockWrapper.capabilities()).thenReturn(new CaseInsensitiveCapability()); Mockito.when(mockCatalogManager.loadCatalogAndWrap(any())).thenReturn(mockWrapper); Mockito.when(mockCatalogManager.acquireCatalogLease(any())) - .thenAnswer(invocation -> CatalogLease.of(mockWrapper)); - Mockito.when(mockWrapper.tryAcquire()).thenReturn(true); + .thenAnswer(invocation -> CatalogTestUtils.unmanagedLease(mockWrapper)); OwnerDispatcher mockOwnerDispatcher = Mockito.mock(OwnerDispatcher.class); TopicDispatcher mockTopicDispatcher = Mockito.mock(TopicDispatcher.class); diff --git a/iceberg/iceberg-rest-server/src/test/java/org/apache/gravitino/iceberg/service/dispatcher/TestIcebergAsyncPurge.java b/iceberg/iceberg-rest-server/src/test/java/org/apache/gravitino/iceberg/service/dispatcher/TestIcebergAsyncPurge.java index 0b64ed2c175..ae63c101454 100644 --- a/iceberg/iceberg-rest-server/src/test/java/org/apache/gravitino/iceberg/service/dispatcher/TestIcebergAsyncPurge.java +++ b/iceberg/iceberg-rest-server/src/test/java/org/apache/gravitino/iceberg/service/dispatcher/TestIcebergAsyncPurge.java @@ -32,8 +32,8 @@ import java.util.Optional; import org.apache.gravitino.GravitinoEnv; import org.apache.gravitino.auth.AuthConstants; -import org.apache.gravitino.catalog.CatalogLease; import org.apache.gravitino.catalog.CatalogManager; +import org.apache.gravitino.catalog.CatalogTestUtils; import org.apache.gravitino.connector.BaseCatalog; import org.apache.gravitino.iceberg.service.CatalogWrapperForREST; import org.apache.gravitino.iceberg.service.IcebergCatalogWrapperManager; @@ -243,8 +243,7 @@ private static MockedStatic mockCatalogId() { when(env.catalogManager()).thenReturn(catalogManager); when(catalogManager.loadCatalogAndWrap(any())).thenReturn(wrapper); when(catalogManager.acquireCatalogLease(any())) - .thenAnswer(invocation -> CatalogLease.of(wrapper)); - when(wrapper.tryAcquire()).thenReturn(true); + .thenAnswer(invocation -> CatalogTestUtils.unmanagedLease(wrapper)); when(wrapper.catalog()).thenReturn(catalog); when(catalog.entity()).thenReturn(entity); when(entity.id()).thenReturn(CATALOG_ID); From b8eda2855d4e63d3b77182fe1f5e292d67a062e9 Mon Sep 17 00:00:00 2001 From: yuqi Date: Mon, 10 Aug 2026 19:02:28 +0800 Subject: [PATCH 3/6] [#12403] fix(core): keep NoSuchCatalogException unwrapped in getCapability CapabilityHelpers.getCapability() acquired the catalog lease inside the try block that wraps failures into a RuntimeException, so a missing catalog turned a 404 into a 500 on the normalize paths. Acquire the lease outside the try, as the catalog load was before, and wrap only the capability lookup. --- .../gravitino/catalog/CapabilityHelpers.java | 8 ++++- .../catalog/TestCapabilityHelpers.java | 36 +++++++++++++++++++ 2 files changed, 43 insertions(+), 1 deletion(-) diff --git a/core/src/main/java/org/apache/gravitino/catalog/CapabilityHelpers.java b/core/src/main/java/org/apache/gravitino/catalog/CapabilityHelpers.java index 125bbb021bf..56a20852ca3 100644 --- a/core/src/main/java/org/apache/gravitino/catalog/CapabilityHelpers.java +++ b/core/src/main/java/org/apache/gravitino/catalog/CapabilityHelpers.java @@ -52,10 +52,16 @@ public class CapabilityHelpers { public static Capability getCapability(NameIdentifier ident, CatalogManager catalogManager) { NameIdentifier catalogIdent = getCatalogIdentifier(ident); - try (CatalogLease lease = catalogManager.acquireCatalogLease(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 lease.wrapper().capabilities(); } catch (Exception e) { throw new RuntimeException("Failed to get capabilities for catalog: " + catalogIdent, e); + } finally { + lease.close(); } } diff --git a/core/src/test/java/org/apache/gravitino/catalog/TestCapabilityHelpers.java b/core/src/test/java/org/apache/gravitino/catalog/TestCapabilityHelpers.java index 529ee4992c7..df1bb9648c4 100644 --- a/core/src/test/java/org/apache/gravitino/catalog/TestCapabilityHelpers.java +++ b/core/src/test/java/org/apache/gravitino/catalog/TestCapabilityHelpers.java @@ -23,6 +23,7 @@ import org.apache.gravitino.Namespace; import org.apache.gravitino.connector.capability.Capability; import org.apache.gravitino.connector.capability.CapabilityResult; +import org.apache.gravitino.exceptions.NoSuchCatalogException; import org.apache.gravitino.rel.expressions.literals.Literal; import org.apache.gravitino.rel.expressions.literals.Literals; import org.apache.gravitino.rel.partitions.IdentityPartition; @@ -30,6 +31,7 @@ import org.apache.gravitino.rel.partitions.Partitions; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; +import org.mockito.Mockito; public class TestCapabilityHelpers { @@ -142,4 +144,38 @@ void testApplyCapabilitiesValidatesNameBeforeNormalizing() { Assertions.assertEquals("My Table", result.name()); } + + @Test + void testGetCapabilityPropagatesNoSuchCatalogException() { + NameIdentifier tableIdent = + NameIdentifier.of(Namespace.of("metalake", "catalog", "schema"), "table"); + CatalogManager catalogManager = Mockito.mock(CatalogManager.class); + Mockito.when(catalogManager.acquireCatalogLease(Mockito.any())) + .thenThrow(new NoSuchCatalogException("Catalog %s does not exist", tableIdent)); + + // A missing catalog must stay a NoSuchCatalogException (mapped to 404 by the REST layer) + // instead of being wrapped into a plain RuntimeException (a 500). + Assertions.assertThrows( + NoSuchCatalogException.class, + () -> CapabilityHelpers.getCapability(tableIdent, catalogManager)); + } + + @Test + void testGetCapabilityWrapsCapabilityFailureAndReleasesLease() throws Exception { + NameIdentifier tableIdent = + NameIdentifier.of(Namespace.of("metalake", "catalog", "schema"), "table"); + CatalogManager catalogManager = Mockito.mock(CatalogManager.class); + CatalogManager.CatalogWrapper wrapper = Mockito.mock(CatalogManager.CatalogWrapper.class); + Mockito.when(catalogManager.acquireCatalogLease(Mockito.any())) + .thenAnswer(invocation -> CatalogTestUtils.unmanagedLease(wrapper)); + Mockito.when(wrapper.capabilities()).thenThrow(new IllegalStateException("boom")); + + RuntimeException e = + Assertions.assertThrows( + RuntimeException.class, + () -> CapabilityHelpers.getCapability(tableIdent, catalogManager)); + + Assertions.assertInstanceOf(IllegalStateException.class, e.getCause()); + Mockito.verify(wrapper).release(); + } } From 56837650194acd4098129664dfc99a1a659de54c Mon Sep 17 00:00:00 2001 From: yuqi Date: Mon, 10 Aug 2026 20:04:16 +0800 Subject: [PATCH 4/6] [#12403] fix(core): address review comments on the catalog lease - Make CatalogWrapper#catalog volatile so the cleanup that nulls it outside leaseLock is properly published to unleased readers (loadCatalogAndWrap callers); keep the nulling out of the lock so a slow catalog close cannot stall tryAcquire. Make classLoader and pool final, they never change after construction. - Keep the Javadoc @link to acquireCatalogLease on one line by importing NameIdentifier, so the link renders. - Build the expiring test's SecretManager from the same config as its CatalogManager. --- .../apache/gravitino/catalog/CatalogLease.java | 5 +++-- .../apache/gravitino/catalog/CatalogManager.java | 15 ++++++++++++--- .../catalog/TestCatalogWrapperLease.java | 5 ++++- 3 files changed, 19 insertions(+), 6 deletions(-) diff --git a/core/src/main/java/org/apache/gravitino/catalog/CatalogLease.java b/core/src/main/java/org/apache/gravitino/catalog/CatalogLease.java index 93af7d692c1..5cc7e80e1a7 100644 --- a/core/src/main/java/org/apache/gravitino/catalog/CatalogLease.java +++ b/core/src/main/java/org/apache/gravitino/catalog/CatalogLease.java @@ -19,6 +19,7 @@ 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; @@ -31,8 +32,8 @@ * resources are released once the wrapper is retired and its last lease is closed, so an operation * can never observe a half-closed catalog. * - *

Leases are obtained from {@link CatalogManager#acquireCatalogLease(org.apache.gravitino - * .NameIdentifier)} and must be closed exactly once, ideally with try-with-resources: + *

Leases are obtained from {@link CatalogManager#acquireCatalogLease(NameIdentifier)} and must + * be closed exactly once, ideally with try-with-resources: * * 

{@code
  * try (CatalogLease lease = catalogManager.acquireCatalogLease(ident)) {
diff --git a/core/src/main/java/org/apache/gravitino/catalog/CatalogManager.java b/core/src/main/java/org/apache/gravitino/catalog/CatalogManager.java
index ed841b94bb3..d153cee2bdb 100644
--- a/core/src/main/java/org/apache/gravitino/catalog/CatalogManager.java
+++ b/core/src/main/java/org/apache/gravitino/catalog/CatalogManager.java
@@ -174,9 +174,16 @@ public class CatalogManager implements CatalogDispatcher, Closeable {
    */
   public static class CatalogWrapper {
 
-    private BaseCatalog catalog;
-    private IsolatedClassLoader classLoader;
-    private ClassLoaderPool pool;
+    // Volatile because cleanup() nulls it outside leaseLock (holding the lock across a catalog
+    // close would stall tryAcquire), while unleased readers such as callers of
+    // loadCatalogAndWrap() may read it from another thread. Leased readers cannot race with
+    // cleanup at all: cleanup is only claimed once the wrapper is retired and no lease is held.
+    private volatile BaseCatalog catalog;
+
+    private final IsolatedClassLoader classLoader;
+    private final ClassLoaderPool pool;
+
+    // Only written by the single thread that claims the cleanup, and read nowhere else.
     private PooledClassLoaderEntry poolEntry;
 
     /** Guards {@link #activeOps}, {@link #retired} and {@link #cleanupStarted}. */
@@ -194,6 +201,8 @@ public static class CatalogWrapper {
     /** Non-pooled constructor: each catalog owns its ClassLoader exclusively. */
     CatalogWrapper(IsolatedClassLoader classLoader) {
       this.classLoader = classLoader;
+      this.pool = null;
+      this.poolEntry = null;
     }
 
     /** Pooled constructor: ClassLoader is managed by the pool with reference counting. */
diff --git a/core/src/test/java/org/apache/gravitino/catalog/TestCatalogWrapperLease.java b/core/src/test/java/org/apache/gravitino/catalog/TestCatalogWrapperLease.java
index e40f20ee99d..c055e1bd9fd 100644
--- a/core/src/test/java/org/apache/gravitino/catalog/TestCatalogWrapperLease.java
+++ b/core/src/test/java/org/apache/gravitino/catalog/TestCatalogWrapperLease.java
@@ -157,7 +157,10 @@ public void testCacheExpiryDefersCleanupUntilLeaseIsReleased() throws Exception
 
     CatalogManager expiringManager =
         new CatalogManager(
-            expiringConfig, entityStore, new RandomIdGenerator(), new SecretManager(config));
+            expiringConfig,
+            entityStore,
+            new RandomIdGenerator(),
+            new SecretManager(expiringConfig));
     try {
       NameIdentifier ident = NameIdentifier.of(METALAKE, "expiring_catalog");
       expiringManager.createCatalog(ident, Catalog.Type.RELATIONAL, PROVIDER, "comment", PROPS);

From 06b0aaa63a53c3cf3ae62bfe49343230f831e4a0 Mon Sep 17 00:00:00 2001
From: yuqi 
Date: Wed, 12 Aug 2026 13:31:51 +0800
Subject: [PATCH 5/6] [#12403] fix(core): address catalog lease lifecycle gaps

---
 .../authorization/AuthorizationUtils.java     |  91 +++++++-------
 .../gravitino/catalog/CatalogManager.java     | 113 ++++++++++++++----
 .../gravitino/hook/CatalogHookDispatcher.java |  24 ++--
 .../gravitino/utils/ClassLoaderPool.java      |  91 ++++++++++----
 .../TestAccessControlManager.java             |   3 +-
 ...estAccessControlManagerForPermissions.java |   3 +-
 .../authorization/TestAuthorizationUtils.java |   9 +-
 .../authorization/TestOwnerManager.java       |   4 +-
 .../gravitino/catalog/CatalogTestUtils.java   |  21 ++++
 .../gravitino/catalog/TestCatalogManager.java |  47 ++++----
 .../catalog/TestCatalogWrapperLease.java      | 110 +++++++++++++++++
 .../catalog/TestTableOperationDispatcher.java |  15 ++-
 .../catalog/TestTopicOperationDispatcher.java |   3 +-
 .../catalog/TestViewOperationDispatcher.java  |  24 ++--
 .../hook/TestFilesetHookDispatcher.java       |   2 +-
 .../hook/TestTopicHookDispatcher.java         |   2 +-
 .../DynamicIcebergConfigProvider.java         | 105 +++++++++-------
 .../TestDynamicIcebergConfigProvider.java     |  98 +++++++--------
 18 files changed, 506 insertions(+), 259 deletions(-)

diff --git a/core/src/main/java/org/apache/gravitino/authorization/AuthorizationUtils.java b/core/src/main/java/org/apache/gravitino/authorization/AuthorizationUtils.java
index 76f77770e85..3298ab5aefd 100644
--- a/core/src/main/java/org/apache/gravitino/authorization/AuthorizationUtils.java
+++ b/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,20 +367,16 @@ 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);
         }
       }
     }
@@ -388,9 +384,10 @@ public static void callAuthorizationPluginForSecurableObjects(
 
   public static void callAuthorizationPluginForMetadataObject(
       String metalake, MetadataObject metadataObject, Consumer consumer) {
-    List loadedCatalogs = loadMetadataObjectCatalog(metalake, metadataObject);
-    for (Catalog catalog : loadedCatalogs) {
-      callAuthorizationPluginImpl(consumer, catalog);
+    CatalogManager catalogManager = GravitinoEnv.getInstance().catalogManager();
+    List 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 locations) {
+  public static void removeCatalogPrivileges(NameIdentifier catalogIdent, List 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 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 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 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 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 loadMetadataObjectCatalog(
+  private static List getMetadataObjectCatalogs(
       String metalake, MetadataObject metadataObject) {
     CatalogManager catalogManager = GravitinoEnv.getInstance().catalogManager();
-    List loadedCatalogs = Lists.newArrayList();
+    List 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
diff --git a/core/src/main/java/org/apache/gravitino/catalog/CatalogManager.java b/core/src/main/java/org/apache/gravitino/catalog/CatalogManager.java
index d153cee2bdb..67fbd178541 100644
--- a/core/src/main/java/org/apache/gravitino/catalog/CatalogManager.java
+++ b/core/src/main/java/org/apache/gravitino/catalog/CatalogManager.java
@@ -55,6 +55,7 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ScheduledThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Consumer;
 import java.util.stream.Collectors;
@@ -417,13 +418,19 @@ private boolean claimCleanupIfIdle() {
     }
 
     private void cleanup() {
+      // Drop the reference before closing so a failing close() cannot leave a half-closed catalog
+      // reachable: cleanup() runs exactly once, so a null assignment after close() would be skipped
+      // on that path. Unleased readers then see null and fail fast instead of using a closed
+      // catalog; leased readers cannot race with cleanup at all.
+      BaseCatalog toClose = catalog;
+      catalog = null;
+
       try {
         classLoader.withClassLoader(
             cl -> {
-              if (catalog != null) {
-                catalog.close();
+              if (toClose != null) {
+                toClose.close();
               }
-              catalog = null;
               return null;
             });
       } catch (Exception e) {
@@ -493,6 +500,8 @@ private ModelCatalog asModels() {
   private final ConcurrentHashMap localMutationCounts =
       new ConcurrentHashMap<>();
 
+  private final AtomicBoolean closed = new AtomicBoolean(false);
+
   // Set to true when a CatalogChangeLogListener is active. markLocalMutation() is a no-op
   // unless this flag is set, preventing unbounded growth of localMutationCounts in deployments
   // that do not use a relational entity store (where the poller never runs).
@@ -521,15 +530,19 @@ public CatalogManager(
             .removalListener(
                 (k, v, c) -> {
                   LOG.debug("Removed catalog cache entry, identifier={}, cause={}", k, c);
-                  for (Consumer listener : removalListeners) {
-                    if (k != null) {
-                      listener.accept((NameIdentifier) k);
+                  try {
+                    for (Consumer listener : removalListeners) {
+                      if (k != null) {
+                        listener.accept((NameIdentifier) k);
+                      }
                     }
+                  } finally {
+                    // Retire rather than close: an operation that already leased this wrapper
+                    // keeps it alive, and the actual catalog/ClassLoader cleanup runs when the
+                    // last lease is released. Keep this in finally so a faulty external listener
+                    // cannot skip resource cleanup.
+                    ((CatalogWrapper) v).retire();
                   }
-                  // Retire rather than close: an operation that already leased this wrapper keeps
-                  // it alive, and the actual catalog/ClassLoader cleanup runs when the last lease
-                  // is released.
-                  ((CatalogWrapper) v).retire();
                 })
             .scheduler(
                 Scheduler.forScheduledExecutorService(
@@ -558,18 +571,27 @@ public CatalogManager(
   }
 
   /**
-   * Closes the CatalogManager and releases any resources associated with it. This method
-   * invalidates all cached catalog instances and clears the cache.
+   * Closes the CatalogManager and invalidates all cached catalog instances. Idle resources are
+   * released immediately; resources protected by active leases are released when their last lease
+   * is closed.
    */
   @Override
   public void close() {
+    if (!closed.compareAndSet(false, true)) {
+      return;
+    }
     if (catalogChangeLogListener != null) {
       ((SupportsEntityChangeLog) store).unregisterEntityChangeLogListener(catalogChangeLogListener);
       trackLocalMutations = false;
       localMutationCounts.clear();
     }
+    List wrappers = new ArrayList<>(catalogCache.asMap().values());
     catalogCache.invalidateAll();
-    classLoaderPool.close();
+    // The removal listener is asynchronous, so retire the wrappers synchronously before closing
+    // the pool. Active leases defer wrapper cleanup and keep their pooled ClassLoader reference
+    // alive until the last operation releases it.
+    wrappers.forEach(CatalogWrapper::retire);
+    classLoaderPool.closeWhenIdle();
   }
 
   /**
@@ -688,7 +710,7 @@ public Catalog[] listCatalogsInfo(Namespace namespace) throws NoSuchMetalakeExce
    * Loads the catalog with the specified identifier.
    *
    * @param ident The identifier of the catalog to load.
-   * @return The loaded catalog.
+   * @return A metadata snapshot of the loaded catalog. Connector resources are not exposed.
    * @throws NoSuchCatalogException If the specified catalog does not exist.
    */
   @Override
@@ -700,11 +722,35 @@ public Catalog loadCatalog(NameIdentifier ident) throws NoSuchCatalogException {
           try (CatalogLease lease = acquireCatalogLease(ident)) {
             BaseCatalog baseCatalog = lease.catalog();
             baseCatalog.checkMetalakeInUse();
-            return baseCatalog;
+            return toCatalogInfo(lease.wrapper());
           }
         });
   }
 
+  /**
+   * Runs an operation against a catalog while keeping its catalog instance and ClassLoader alive.
+   *
+   * 
Callers that need connector-only state, such as the authorization plugin or raw catalog
+   * properties, must use this method instead of casting the metadata snapshot returned by {@link
+   * #loadCatalog(NameIdentifier)}.
+   *
+   * @param ident The identifier of the catalog to use.
+   * @param operation The operation to run against the live catalog instance.
+   * @return The value returned by the operation.
+   * @param  The result type of the operation.
+   * @throws NoSuchCatalogException If the specified catalog does not exist.
+   */
+  public  R doWithCatalog(NameIdentifier ident, ThrowableFunction operation)
+      throws NoSuchCatalogException {
+    try (CatalogLease lease = acquireCatalogLease(ident)) {
+      return lease.wrapper().doWithCredentialOps(operation);
+    } catch (RuntimeException e) {
+      throw e;
+    } catch (Exception e) {
+      throw new RuntimeException("Failed to operate on catalog: " + ident, e);
+    }
+  }
+
   /**
    * Creates a new catalog with the provided details.
    *
@@ -713,7 +759,7 @@ public Catalog loadCatalog(NameIdentifier ident) throws NoSuchCatalogException {
    * @param provider The provider of the new catalog.
    * @param comment The comment for the new catalog.
    * @param properties The properties of the new catalog.
-   * @return The created catalog.
+   * @return A metadata snapshot of the created catalog.
    * @throws NoSuchMetalakeException If the specified metalake does not exist.
    * @throws CatalogAlreadyExistsException If a catalog with the same identifier already exists.
    */
@@ -758,11 +804,12 @@ public Catalog createCatalog(
           boolean needClean = true;
           try {
             store.put(e, false /* overwrite */);
-            CatalogWrapper wrapper =
-                catalogCache.get(ident, id -> createCatalogWrapper(e, mergedConfig));
+            catalogCache.get(ident, id -> createCatalogWrapper(e, mergedConfig));
 
             needClean = false;
-            return wrapper.catalog;
+            try (CatalogLease lease = acquireCatalogLease(ident)) {
+              return toCatalogInfo(lease.wrapper());
+            }
 
           } catch (EntityAlreadyExistsException e1) {
             needClean = false;
@@ -958,7 +1005,7 @@ public void disableCatalog(NameIdentifier ident) throws NoSuchCatalogException {
    *
    * @param ident The identifier of the catalog to alter.
    * @param changes The changes to apply to the catalog.
-   * @return The altered catalog.
+   * @return A metadata snapshot of the altered catalog.
    * @throws NoSuchCatalogException If the specified catalog does not exist.
    * @throws IllegalArgumentException If an unsupported catalog change is provided.
    */
@@ -1040,7 +1087,9 @@ public Catalog alterCatalog(NameIdentifier ident, CatalogChange... changes)
             // a background thread from overwriting it with stale data between invalidate and put.
             CatalogWrapper newWrapper = createCatalogWrapper(convertedCatalog, null);
             catalogCache.put(convertedCatalog.nameIdentifier(), newWrapper);
-            return newWrapper.catalog();
+            try (CatalogLease lease = acquireCatalogLease(convertedCatalog.nameIdentifier())) {
+              return toCatalogInfo(lease.wrapper());
+            }
 
           } catch (NoSuchEntityException ne) {
             LOG.warn("Catalog {} does not exist", ident, ne);
@@ -1228,8 +1277,14 @@ private boolean containsUserCreatedSchemas(
    * @throws NoSuchCatalogException If the specified catalog does not exist.
    */
   public CatalogLease acquireCatalogLease(NameIdentifier ident) throws NoSuchCatalogException {
+    checkOpen();
     for (int attempt = 0; attempt < MAX_LEASE_ATTEMPTS; attempt++) {
       CatalogWrapper wrapper = loadCatalogAndWrap(ident);
+      if (closed.get()) {
+        catalogCache.asMap().remove(ident, wrapper);
+        wrapper.retire();
+        checkOpen();
+      }
       if (wrapper.tryAcquire()) {
         return new CatalogLease(wrapper);
       }
@@ -1465,6 +1520,22 @@ private Map getResolvedProperties(CatalogEntity entity) {
     }
   }
 
+  private Catalog toCatalogInfo(CatalogWrapper wrapper) {
+    try {
+      return wrapper.doWithCredentialOps(
+          catalog ->
+              catalog.entity().toCatalogInfoWithResolvedProps(new HashMap<>(catalog.properties())));
+    } catch (RuntimeException e) {
+      throw e;
+    } catch (Exception e) {
+      throw new RuntimeException("Failed to create catalog metadata snapshot", e);
+    }
+  }
+
+  private void checkOpen() {
+    Preconditions.checkState(!closed.get(), "CatalogManager is already closed");
+  }
+
   private BaseCatalog createBaseCatalog(IsolatedClassLoader classLoader, CatalogEntity entity) {
     // Load Catalog class instance
     BaseCatalog catalog = createCatalogInstance(classLoader, entity.getProvider());
diff --git a/core/src/main/java/org/apache/gravitino/hook/CatalogHookDispatcher.java b/core/src/main/java/org/apache/gravitino/hook/CatalogHookDispatcher.java
index 81d324856ea..272ecfb80ee 100644
--- a/core/src/main/java/org/apache/gravitino/hook/CatalogHookDispatcher.java
+++ b/core/src/main/java/org/apache/gravitino/hook/CatalogHookDispatcher.java
@@ -31,7 +31,6 @@
 import org.apache.gravitino.authorization.Owner;
 import org.apache.gravitino.authorization.OwnerDispatcher;
 import org.apache.gravitino.catalog.CatalogDispatcher;
-import org.apache.gravitino.connector.BaseCatalog;
 import org.apache.gravitino.exceptions.CatalogAlreadyExistsException;
 import org.apache.gravitino.exceptions.CatalogInUseException;
 import org.apache.gravitino.exceptions.CatalogNotInUseException;
@@ -94,9 +93,16 @@ public Catalog createCatalog(
 
       // Apply the metalake securable object privileges to authorization plugin
       FutureGrantManager futureGrantManager = GravitinoEnv.getInstance().futureGrantManager();
-      if (futureGrantManager != null && catalog instanceof BaseCatalog) {
-        futureGrantManager.grantNewlyCreatedCatalog(
-            ident.namespace().level(0), (BaseCatalog) catalog);
+      if (futureGrantManager != null) {
+        GravitinoEnv.getInstance()
+            .catalogManager()
+            .doWithCatalog(
+                ident,
+                leasedCatalog -> {
+                  futureGrantManager.grantNewlyCreatedCatalog(
+                      ident.namespace().level(0), leasedCatalog);
+                  return null;
+                });
       }
     } catch (Exception postHookException) {
       LOG.warn(
@@ -143,13 +149,9 @@ public boolean dropCatalog(NameIdentifier ident, boolean force)
       return false;
     }
 
-    Catalog catalog = dispatcher.loadCatalog(ident);
-
-    if (catalog != null) {
-      List locations =
-          AuthorizationUtils.getMetadataObjectLocation(ident, Entity.EntityType.CATALOG);
-      AuthorizationUtils.removeCatalogPrivileges(catalog, locations);
-    }
+    List locations =
+        AuthorizationUtils.getMetadataObjectLocation(ident, Entity.EntityType.CATALOG);
+    AuthorizationUtils.removeCatalogPrivileges(ident, locations);
 
     // We should call the authorization plugin before dropping the catalog, because the dropping
     // catalog will close the authorization plugin.
diff --git a/core/src/main/java/org/apache/gravitino/utils/ClassLoaderPool.java b/core/src/main/java/org/apache/gravitino/utils/ClassLoaderPool.java
index fd376fb0e1d..f5c9f61f418 100644
--- a/core/src/main/java/org/apache/gravitino/utils/ClassLoaderPool.java
+++ b/core/src/main/java/org/apache/gravitino/utils/ClassLoaderPool.java
@@ -25,6 +25,7 @@
 import java.util.Enumeration;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
 import java.util.function.Supplier;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -54,6 +55,8 @@ public class ClassLoaderPool implements Closeable {
 
   private final AtomicBoolean closed = new AtomicBoolean(false);
 
+  private final ReentrantReadWriteLock lifecycleLock = new ReentrantReadWriteLock();
+
   /**
    * Acquires a ClassLoader entry for the given key. If an entry already exists, increments the
    * reference count. Otherwise, creates a new entry using the provided factory.
@@ -64,24 +67,29 @@ public class ClassLoaderPool implements Closeable {
    * @throws IllegalStateException if the pool has been closed.
    */
   public PooledClassLoaderEntry acquire(ClassLoaderKey key, Supplier factory) {
-    return pool.compute(
-        key,
-        (k, existing) -> {
-          if (closed.get()) {
-            throw new IllegalStateException("ClassLoaderPool is already closed");
-          }
-          if (existing != null) {
-            existing.incrementRefCount();
-            LOG.debug("Reusing ClassLoader for key {}, refCount={}.", key, existing.refCount());
-            return existing;
-          }
-          // If the factory throws (e.g., invalid classpath), the exception propagates to the
-          // caller and ConcurrentHashMap leaves the key unmapped.
-          IsolatedClassLoader classLoader = factory.get();
-          PooledClassLoaderEntry newEntry = new PooledClassLoaderEntry(k, classLoader);
-          LOG.info("Created new ClassLoader for key {}, refCount=1.", key);
-          return newEntry;
-        });
+    lifecycleLock.readLock().lock();
+    try {
+      return pool.compute(
+          key,
+          (k, existing) -> {
+            if (closed.get()) {
+              throw new IllegalStateException("ClassLoaderPool is already closed");
+            }
+            if (existing != null) {
+              existing.incrementRefCount();
+              LOG.debug("Reusing ClassLoader for key {}, refCount={}.", key, existing.refCount());
+              return existing;
+            }
+            // If the factory throws (e.g., invalid classpath), the exception propagates to the
+            // caller and ConcurrentHashMap leaves the key unmapped.
+            IsolatedClassLoader classLoader = factory.get();
+            PooledClassLoaderEntry newEntry = new PooledClassLoaderEntry(k, classLoader);
+            LOG.info("Created new ClassLoader for key {}, refCount=1.", key);
+            return newEntry;
+          });
+    } finally {
+      lifecycleLock.readLock().unlock();
+    }
   }
 
   /**
@@ -124,12 +132,30 @@ public int size() {
    */
   @Override
   public void close() {
-    closed.set(true);
-    // Drain with a loop to catch entries inserted by concurrent acquire() calls that were
-    // already past the closed check when we set the flag. Since closed=true prevents any new
-    // entries from being created, this loop is guaranteed to terminate.
-    while (!pool.isEmpty()) {
+    lifecycleLock.writeLock().lock();
+    try {
+      closed.set(true);
       pool.keySet().forEach(this::removeAndCleanup);
+    } finally {
+      lifecycleLock.writeLock().unlock();
+    }
+  }
+
+  /**
+   * Prevents new acquisitions and cleans up idle entries while deferring entries that still have
+   * owners. A deferred entry is cleaned up by {@link #release(PooledClassLoaderEntry)} when its
+   * reference count reaches zero.
+   *
+   * 
This is intended for graceful owner shutdown: unlike {@link #close()}, it never closes a
+   * ClassLoader that an active owner may still be using.
+   */
+  public void closeWhenIdle() {
+    lifecycleLock.writeLock().lock();
+    try {
+      closed.set(true);
+      pool.keySet().forEach(this::removeAndCleanupIfIdle);
+    } finally {
+      lifecycleLock.writeLock().unlock();
     }
   }
 
@@ -153,6 +179,25 @@ private void removeAndCleanup(ClassLoaderKey key) {
         });
   }
 
+  private void removeAndCleanupIfIdle(ClassLoaderKey key) {
+    pool.compute(
+        key,
+        (k, existing) -> {
+          if (existing == null) {
+            return null;
+          }
+          if (existing.refCount() > 0) {
+            LOG.info(
+                "Deferring ClassLoader cleanup for key {} with {} active reference(s).",
+                k,
+                existing.refCount());
+            return existing;
+          }
+          doFinalCleanup(existing);
+          return null;
+        });
+  }
+
   /**
    * Performs final cleanup when a ClassLoader's reference count reaches zero. This includes:
    *
diff --git a/core/src/test/java/org/apache/gravitino/authorization/TestAccessControlManager.java b/core/src/test/java/org/apache/gravitino/authorization/TestAccessControlManager.java
index fbc54b29030..68d2ca42544 100644
--- a/core/src/test/java/org/apache/gravitino/authorization/TestAccessControlManager.java
+++ b/core/src/test/java/org/apache/gravitino/authorization/TestAccessControlManager.java
@@ -66,6 +66,7 @@
 import org.apache.gravitino.Namespace;
 import org.apache.gravitino.StringIdentifier;
 import org.apache.gravitino.catalog.CatalogManager;
+import org.apache.gravitino.catalog.CatalogTestUtils;
 import org.apache.gravitino.connector.BaseCatalog;
 import org.apache.gravitino.connector.authorization.AuthorizationPlugin;
 import org.apache.gravitino.exceptions.GroupAlreadyExistsException;
@@ -194,7 +195,7 @@ public static void setUp() throws Exception {
         GravitinoEnv.getInstance(), "accessControlDispatcher", accessControlManager, true);
     FieldUtils.writeField(GravitinoEnv.getInstance(), "catalogManager", catalogManager, true);
     BaseCatalog catalog = mock(BaseCatalog.class);
-    when(catalogManager.loadCatalog(any())).thenReturn(catalog);
+    CatalogTestUtils.mockDoWithCatalog(catalogManager, catalog);
     authorizationPlugin = mock(AuthorizationPlugin.class);
     when(catalog.getAuthorizationPlugin()).thenReturn(authorizationPlugin);
   }
diff --git a/core/src/test/java/org/apache/gravitino/authorization/TestAccessControlManagerForPermissions.java b/core/src/test/java/org/apache/gravitino/authorization/TestAccessControlManagerForPermissions.java
index 7a6564270b9..228ae3b1acf 100644
--- a/core/src/test/java/org/apache/gravitino/authorization/TestAccessControlManagerForPermissions.java
+++ b/core/src/test/java/org/apache/gravitino/authorization/TestAccessControlManagerForPermissions.java
@@ -42,6 +42,7 @@
 import org.apache.gravitino.NameIdentifier;
 import org.apache.gravitino.Namespace;
 import org.apache.gravitino.catalog.CatalogManager;
+import org.apache.gravitino.catalog.CatalogTestUtils;
 import org.apache.gravitino.connector.BaseCatalog;
 import org.apache.gravitino.connector.authorization.AuthorizationPlugin;
 import org.apache.gravitino.exceptions.IllegalRoleException;
@@ -176,7 +177,7 @@ public static void setUp() throws Exception {
         GravitinoEnv.getInstance(), "accessControlDispatcher", accessControlManager, true);
     FieldUtils.writeField(GravitinoEnv.getInstance(), "catalogManager", catalogManager, true);
     BaseCatalog catalog = Mockito.mock(BaseCatalog.class);
-    Mockito.when(catalogManager.loadCatalog(any())).thenReturn(catalog);
+    CatalogTestUtils.mockDoWithCatalog(catalogManager, catalog);
     Mockito.when(catalogManager.listCatalogs(Mockito.any()))
         .thenReturn(new NameIdentifier[] {NameIdentifier.of("metalake", "catalog")});
     authorizationPlugin = Mockito.mock(AuthorizationPlugin.class);
diff --git a/core/src/test/java/org/apache/gravitino/authorization/TestAuthorizationUtils.java b/core/src/test/java/org/apache/gravitino/authorization/TestAuthorizationUtils.java
index 48bae76cbd8..9d44687b9fc 100644
--- a/core/src/test/java/org/apache/gravitino/authorization/TestAuthorizationUtils.java
+++ b/core/src/test/java/org/apache/gravitino/authorization/TestAuthorizationUtils.java
@@ -33,6 +33,7 @@
 import org.apache.gravitino.Schema;
 import org.apache.gravitino.catalog.CatalogDispatcher;
 import org.apache.gravitino.catalog.CatalogManager;
+import org.apache.gravitino.catalog.CatalogTestUtils;
 import org.apache.gravitino.catalog.SchemaDispatcher;
 import org.apache.gravitino.catalog.TableDispatcher;
 import org.apache.gravitino.connector.BaseCatalog;
@@ -357,7 +358,7 @@ void testRenamePrivilegesNotifiesOldEntityNameIdMappingChange() {
     AccessControlDispatcher accessControlDispatcher = Mockito.mock(AccessControlDispatcher.class);
     CatalogManager catalogManager = Mockito.mock(CatalogManager.class);
     BaseCatalog baseCatalog = Mockito.mock(BaseCatalog.class);
-    Mockito.when(catalogManager.loadCatalog(Mockito.any())).thenReturn(baseCatalog);
+    CatalogTestUtils.mockDoWithCatalog(catalogManager, baseCatalog);
 
     GravitinoEnv envMock = Mockito.mock(GravitinoEnv.class);
     Mockito.when(envMock.gravitinoAuthorizer()).thenReturn(authorizer);
@@ -379,14 +380,13 @@ void testRenamePrivilegesNotifiesOldEntityNameIdMappingChange() {
   @Test
   void testRenameTablePrivilegesNotifiesAuthorizationPluginWithExpectedChange() {
     NameIdentifier ident = NameIdentifier.of("metalake", "catalog", "schema", "table");
-    NameIdentifier catalogIdent = NameIdentifier.of("metalake", "catalog");
     List locations = Lists.newArrayList("/warehouse/schema/table");
 
     AccessControlDispatcher accessControlDispatcher = Mockito.mock(AccessControlDispatcher.class);
     CatalogManager catalogManager = Mockito.mock(CatalogManager.class);
     BaseCatalog baseCatalog = Mockito.mock(BaseCatalog.class);
     AuthorizationPlugin authorizationPlugin = Mockito.mock(AuthorizationPlugin.class);
-    Mockito.when(catalogManager.loadCatalog(catalogIdent)).thenReturn(baseCatalog);
+    CatalogTestUtils.mockDoWithCatalog(catalogManager, baseCatalog);
     Mockito.when(baseCatalog.getAuthorizationPlugin()).thenReturn(authorizationPlugin);
 
     GravitinoEnv envMock = Mockito.mock(GravitinoEnv.class);
@@ -419,14 +419,13 @@ void testRenameTablePrivilegesNotifiesAuthorizationPluginWithExpectedChange() {
   @Test
   void testRemoveTablePrivilegesNotifiesAuthorizationPluginWithExpectedChange() {
     NameIdentifier ident = NameIdentifier.of("metalake", "catalog", "schema", "table");
-    NameIdentifier catalogIdent = NameIdentifier.of("metalake", "catalog");
     List locations = Lists.newArrayList("/warehouse/schema/table");
 
     AccessControlDispatcher accessControlDispatcher = Mockito.mock(AccessControlDispatcher.class);
     CatalogManager catalogManager = Mockito.mock(CatalogManager.class);
     BaseCatalog baseCatalog = Mockito.mock(BaseCatalog.class);
     AuthorizationPlugin authorizationPlugin = Mockito.mock(AuthorizationPlugin.class);
-    Mockito.when(catalogManager.loadCatalog(catalogIdent)).thenReturn(baseCatalog);
+    CatalogTestUtils.mockDoWithCatalog(catalogManager, baseCatalog);
     Mockito.when(baseCatalog.getAuthorizationPlugin()).thenReturn(authorizationPlugin);
 
     GravitinoEnv envMock = Mockito.mock(GravitinoEnv.class);
diff --git a/core/src/test/java/org/apache/gravitino/authorization/TestOwnerManager.java b/core/src/test/java/org/apache/gravitino/authorization/TestOwnerManager.java
index 564000a7ca2..8bc9234d9fd 100644
--- a/core/src/test/java/org/apache/gravitino/authorization/TestOwnerManager.java
+++ b/core/src/test/java/org/apache/gravitino/authorization/TestOwnerManager.java
@@ -38,7 +38,6 @@
 import static org.apache.gravitino.Configs.TREE_LOCK_MAX_NODE_IN_MEMORY;
 import static org.apache.gravitino.Configs.TREE_LOCK_MIN_NODE_IN_MEMORY;
 import static org.apache.gravitino.Configs.VERSION_RETENTION_COUNT;
-import static org.mockito.ArgumentMatchers.any;
 
 import com.google.common.collect.Lists;
 import java.io.File;
@@ -60,6 +59,7 @@
 import org.apache.gravitino.NameIdentifier;
 import org.apache.gravitino.Namespace;
 import org.apache.gravitino.catalog.CatalogManager;
+import org.apache.gravitino.catalog.CatalogTestUtils;
 import org.apache.gravitino.connector.BaseCatalog;
 import org.apache.gravitino.connector.authorization.AuthorizationPlugin;
 import org.apache.gravitino.exceptions.NoSuchMetadataObjectException;
@@ -177,7 +177,7 @@ public static void setUp() throws IOException, IllegalAccessException {
 
     ownerManager = new OwnerManager(entityStore);
     BaseCatalog catalog = Mockito.mock(BaseCatalog.class);
-    Mockito.when(catalogManager.loadCatalog(any())).thenReturn(catalog);
+    CatalogTestUtils.mockDoWithCatalog(catalogManager, catalog);
     Mockito.when(catalogManager.listCatalogs(Mockito.any()))
         .thenReturn(new NameIdentifier[] {NameIdentifier.of("metalake", "catalog")});
     Mockito.when(catalog.getAuthorizationPlugin()).thenReturn(authorizationPlugin);
diff --git a/core/src/test/java/org/apache/gravitino/catalog/CatalogTestUtils.java b/core/src/test/java/org/apache/gravitino/catalog/CatalogTestUtils.java
index c52a9ee4711..0e1e51d63ef 100644
--- a/core/src/test/java/org/apache/gravitino/catalog/CatalogTestUtils.java
+++ b/core/src/test/java/org/apache/gravitino/catalog/CatalogTestUtils.java
@@ -19,6 +19,9 @@
 package org.apache.gravitino.catalog;
 
 import org.apache.gravitino.catalog.CatalogManager.CatalogWrapper;
+import org.apache.gravitino.connector.BaseCatalog;
+import org.apache.gravitino.utils.ThrowableFunction;
+import org.mockito.Mockito;
 
 /** Test-only helpers for catalog internals that are package-private in production code. */
 public final class CatalogTestUtils {
@@ -37,4 +40,22 @@ private CatalogTestUtils() {}
   public static CatalogLease unmanagedLease(CatalogWrapper wrapper) {
     return new CatalogLease(wrapper);
   }
+
+  /**
+   * Stubs a mocked manager so {@link CatalogManager#doWithCatalog} invokes its callback with the
+   * supplied catalog.
+   *
+   * @param catalogManager the mocked catalog manager.
+   * @param catalog the live catalog to pass to callbacks.
+   */
+  @SuppressWarnings("unchecked")
+  public static void mockDoWithCatalog(CatalogManager catalogManager, BaseCatalog catalog) {
+    Mockito.doAnswer(
+            invocation -> {
+              ThrowableFunction operation = invocation.getArgument(1);
+              return operation.apply(catalog);
+            })
+        .when(catalogManager)
+        .doWithCatalog(Mockito.any(), Mockito.any());
+  }
 }
diff --git a/core/src/test/java/org/apache/gravitino/catalog/TestCatalogManager.java b/core/src/test/java/org/apache/gravitino/catalog/TestCatalogManager.java
index 3dbbbf65ad3..ab98bf4ee64 100644
--- a/core/src/test/java/org/apache/gravitino/catalog/TestCatalogManager.java
+++ b/core/src/test/java/org/apache/gravitino/catalog/TestCatalogManager.java
@@ -566,6 +566,9 @@ public void testLoadCatalog() {
     catalogManager.createCatalog(ident, Catalog.Type.RELATIONAL, provider, "comment", props);
 
     Catalog catalog = catalogManager.loadCatalog(ident);
+    Assertions.assertFalse(
+        catalog instanceof BaseCatalog,
+        "loadCatalog must not expose a live catalog after its lease is released");
     Assertions.assertEquals("test21", catalog.name());
     Assertions.assertEquals("comment", catalog.comment());
     testProperties(props, catalog.properties());
@@ -652,10 +655,8 @@ void testAlterCatalogRefreshesCacheAfterStoreUpdate() throws Exception {
             PROPERTY_KEY5_PREFIX + "1",
             "value3");
 
-    Catalog catalog =
-        catalogManager.createCatalog(ident, Catalog.Type.RELATIONAL, provider, "comment", props);
-    CatalogEntity originalEntity = entityStore.get(ident, EntityType.CATALOG, CatalogEntity.class);
-    FieldUtils.writeField(catalog, "entity", originalEntity, true);
+    catalogManager.createCatalog(ident, Catalog.Type.RELATIONAL, provider, "comment", props);
+    BaseCatalog catalog = catalogManager.loadCatalogAndWrap(ident).catalog();
 
     CatalogManager.CatalogWrapper staleWrapper =
         Mockito.mock(CatalogManager.CatalogWrapper.class, Mockito.RETURNS_DEEP_STUBS);
@@ -664,8 +665,11 @@ void testAlterCatalogRefreshesCacheAfterStoreUpdate() throws Exception {
     CatalogManager.CatalogWrapper freshWrapper =
         Mockito.mock(CatalogManager.CatalogWrapper.class, Mockito.RETURNS_DEEP_STUBS);
     BaseCatalog freshCatalog = Mockito.mock(BaseCatalog.class);
-    Mockito.doReturn("cache_race_test_renamed").when(freshCatalog).name();
+    Catalog freshCatalogInfo = Mockito.mock(Catalog.class);
+    Mockito.doReturn("cache_race_test_renamed").when(freshCatalogInfo).name();
     Mockito.doReturn(freshCatalog).when(freshWrapper).catalog();
+    Mockito.doReturn(true).when(freshWrapper).tryAcquire();
+    Mockito.doReturn(freshCatalogInfo).when(freshWrapper).doWithCredentialOps(any());
 
     AtomicBoolean staleInserted = new AtomicBoolean(false);
     Answer insertStaleWrapper =
@@ -935,12 +939,10 @@ public void testDropCatalogSkipsImportedSchemas() throws Exception {
             "value3");
     String comment = "comment";
 
-    Catalog catalog =
-        catalogManager.createCatalog(ident, Catalog.Type.RELATIONAL, provider, comment, props);
+    catalogManager.createCatalog(ident, Catalog.Type.RELATIONAL, provider, comment, props);
     Mockito.doCallRealMethod().when(catalogManager).loadCatalogAndWrap(ident);
     Assertions.assertDoesNotThrow(() -> catalogManager.disableCatalog(ident));
-    CatalogEntity catalogEntity = entityStore.get(ident, EntityType.CATALOG, CatalogEntity.class);
-    FieldUtils.writeField(catalog, "entity", catalogEntity, true);
+    BaseCatalog catalog = catalogManager.loadCatalogAndWrap(ident).catalog();
 
     SchemaEntity importedSchemaEntity =
         SchemaEntity.builder()
@@ -1020,12 +1022,10 @@ public void testDropCatalogIgnoresMissingSchema() throws Exception {
             "value3");
     String comment = "comment";
 
-    Catalog catalog =
-        catalogManager.createCatalog(ident, Catalog.Type.RELATIONAL, provider, comment, props);
+    catalogManager.createCatalog(ident, Catalog.Type.RELATIONAL, provider, comment, props);
     Mockito.doCallRealMethod().when(catalogManager).loadCatalogAndWrap(ident);
     Assertions.assertDoesNotThrow(() -> catalogManager.disableCatalog(ident));
-    CatalogEntity catalogEntity = entityStore.get(ident, EntityType.CATALOG, CatalogEntity.class);
-    FieldUtils.writeField(catalog, "entity", catalogEntity, true);
+    BaseCatalog catalog = catalogManager.loadCatalogAndWrap(ident).catalog();
 
     SchemaEntity schemaEntity =
         SchemaEntity.builder()
@@ -1072,12 +1072,10 @@ public void testDropCatalogFailsOnSchemaClassificationError() throws Exception {
             "value3");
     String comment = "comment";
 
-    Catalog catalog =
-        catalogManager.createCatalog(ident, Catalog.Type.RELATIONAL, provider, comment, props);
+    catalogManager.createCatalog(ident, Catalog.Type.RELATIONAL, provider, comment, props);
     Mockito.doCallRealMethod().when(catalogManager).loadCatalogAndWrap(ident);
     Assertions.assertDoesNotThrow(() -> catalogManager.disableCatalog(ident));
-    CatalogEntity catalogEntity = entityStore.get(ident, EntityType.CATALOG, CatalogEntity.class);
-    FieldUtils.writeField(catalog, "entity", catalogEntity, true);
+    BaseCatalog catalog = catalogManager.loadCatalogAndWrap(ident).catalog();
 
     SchemaEntity schemaEntity =
         SchemaEntity.builder()
@@ -1125,8 +1123,8 @@ public void testForceDropCatalog() throws Exception {
             PROPERTY_KEY5_PREFIX + "1",
             "value3");
     String comment = "comment";
-    Catalog catalog =
-        catalogManager.createCatalog(ident, Catalog.Type.RELATIONAL, provider, comment, props);
+    catalogManager.createCatalog(ident, Catalog.Type.RELATIONAL, provider, comment, props);
+    BaseCatalog catalog = catalogManager.loadCatalogAndWrap(ident).catalog();
     SchemaEntity schemaEntity =
         SchemaEntity.builder()
             .withId(RandomIdGenerator.INSTANCE.nextId())
@@ -1168,11 +1166,9 @@ void testDropCatalogInvalidatesCacheAfterStoreDelete() throws Exception {
             PROPERTY_KEY5_PREFIX + "1",
             "value3");
 
-    Catalog catalog =
-        catalogManager.createCatalog(ident, Catalog.Type.RELATIONAL, provider, "comment", props);
+    catalogManager.createCatalog(ident, Catalog.Type.RELATIONAL, provider, "comment", props);
     Assertions.assertDoesNotThrow(() -> catalogManager.disableCatalog(ident));
-    CatalogEntity entity = entityStore.get(ident, EntityType.CATALOG, CatalogEntity.class);
-    FieldUtils.writeField(catalog, "entity", entity, true);
+    BaseCatalog catalog = catalogManager.loadCatalogAndWrap(ident).catalog();
 
     CatalogManager.CatalogWrapper catalogWrapper =
         Mockito.mock(CatalogManager.CatalogWrapper.class, Mockito.RETURNS_DEEP_STUBS);
@@ -1206,11 +1202,8 @@ void testDropCatalogReloadsClosedCachedWrapper() throws Exception {
             PROPERTY_KEY5_PREFIX + "1",
             "value3");
 
-    Catalog catalog =
-        catalogManager.createCatalog(ident, Catalog.Type.RELATIONAL, provider, "comment", props);
+    catalogManager.createCatalog(ident, Catalog.Type.RELATIONAL, provider, "comment", props);
     Assertions.assertDoesNotThrow(() -> catalogManager.disableCatalog(ident));
-    CatalogEntity entity = entityStore.get(ident, EntityType.CATALOG, CatalogEntity.class);
-    FieldUtils.writeField(catalog, "entity", entity, true);
 
     CatalogManager.CatalogWrapper closedWrapper = catalogManager.loadCatalogAndWrap(ident);
     closedWrapper.close();
diff --git a/core/src/test/java/org/apache/gravitino/catalog/TestCatalogWrapperLease.java b/core/src/test/java/org/apache/gravitino/catalog/TestCatalogWrapperLease.java
index c055e1bd9fd..84e3aa936ca 100644
--- a/core/src/test/java/org/apache/gravitino/catalog/TestCatalogWrapperLease.java
+++ b/core/src/test/java/org/apache/gravitino/catalog/TestCatalogWrapperLease.java
@@ -39,6 +39,7 @@
 import org.apache.gravitino.NameIdentifier;
 import org.apache.gravitino.Namespace;
 import org.apache.gravitino.catalog.CatalogManager.CatalogWrapper;
+import org.apache.gravitino.connector.BaseCatalog;
 import org.apache.gravitino.lock.LockManager;
 import org.apache.gravitino.meta.AuditInfo;
 import org.apache.gravitino.meta.BaseMetalake;
@@ -56,6 +57,7 @@
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.mockito.Mockito;
 
 /**
  * Tests that a catalog wrapper evicted from the catalog cache is not torn down while an operation
@@ -312,6 +314,88 @@ public void testConcurrentEvictionDoesNotCloseCatalogUnderRunningOperation() thr
     }
   }
 
+  @Test
+  public void testDoWithCatalogKeepsLeaseForEntireCallback() throws Exception {
+    NameIdentifier ident = createCatalog("callback_with_lease");
+    CountDownLatch callbackStarted = new CountDownLatch(1);
+    CountDownLatch continueCallback = new CountDownLatch(1);
+    ExecutorService executor = Executors.newSingleThreadExecutor();
+    try {
+      Future operation =
+          executor.submit(
+              () ->
+                  catalogManager.doWithCatalog(
+                      ident,
+                      catalog -> {
+                        callbackStarted.countDown();
+                        Assertions.assertTrue(continueCallback.await(10, TimeUnit.SECONDS));
+                        catalog.ops();
+                        return true;
+                      }));
+
+      Assertions.assertTrue(callbackStarted.await(10, TimeUnit.SECONDS));
+      CatalogWrapper wrapper = catalogManager.getCatalogCache().getIfPresent(ident);
+      Assertions.assertNotNull(wrapper);
+
+      catalogManager.getCatalogCache().invalidate(ident);
+      await().atMost(Duration.ofSeconds(10)).until(wrapper::isRetired);
+
+      Assertions.assertNotNull(
+          wrapper.catalog(), "the callback lease must survive a concurrent invalidation");
+      continueCallback.countDown();
+      Assertions.assertTrue(operation.get(10, TimeUnit.SECONDS));
+      Assertions.assertNull(wrapper.catalog());
+    } finally {
+      continueCallback.countDown();
+      executor.shutdownNow();
+    }
+  }
+
+  @Test
+  public void testManagerCloseDefersCleanupUntilLeaseIsReleased() throws Exception {
+    NameIdentifier ident = createCatalog("manager_close_with_lease");
+
+    CatalogLease lease = catalogManager.acquireCatalogLease(ident);
+    CatalogWrapper wrapper = lease.wrapper();
+    ClassLoaderPool pool =
+        (ClassLoaderPool) FieldUtils.readField(catalogManager, "classLoaderPool", true);
+
+    catalogManager.close();
+
+    Assertions.assertTrue(wrapper.isRetired());
+    Assertions.assertNotNull(
+        wrapper.catalog(), "manager shutdown must not close a catalog with an active lease");
+    Assertions.assertEquals(
+        1, pool.size(), "manager shutdown must retain an actively leased ClassLoader");
+    Assertions.assertDoesNotThrow(
+        () ->
+            wrapper.doWithSchemaOps(ops -> ops.listSchemas(Namespace.of(METALAKE, ident.name()))));
+    Assertions.assertThrows(
+        IllegalStateException.class, () -> catalogManager.acquireCatalogLease(ident));
+
+    lease.close();
+
+    Assertions.assertNull(wrapper.catalog());
+    Assertions.assertEquals(0, pool.size());
+  }
+
+  @Test
+  public void testRemovalListenerFailureDoesNotSkipWrapperRetirement() {
+    NameIdentifier ident = createCatalog("failing_removal_listener");
+    CatalogWrapper wrapper = catalogManager.getCatalogCache().getIfPresent(ident);
+    Assertions.assertNotNull(wrapper);
+    catalogManager.addCatalogCacheRemoveListener(
+        ignored -> {
+          throw new RuntimeException("listener failed");
+        });
+
+    catalogManager.getCatalogCache().invalidate(ident);
+
+    await().atMost(Duration.ofSeconds(10)).until(wrapper::isRetired);
+    Assertions.assertNull(
+        wrapper.catalog(), "listener failures must not prevent wrapper resource cleanup");
+  }
+
   @Test
   public void testReleaseWithoutAcquireIsRejected() throws Exception {
     NameIdentifier ident = createCatalog("release_without_acquire");
@@ -321,6 +405,32 @@ public void testReleaseWithoutAcquireIsRejected() throws Exception {
     Assertions.assertThrows(IllegalStateException.class, wrapper::release);
   }
 
+  @Test
+  public void testCleanupClearsCatalogEvenWhenCloseFails() throws Exception {
+    NameIdentifier ident = createCatalog("failing_close");
+
+    ClassLoaderPool pool =
+        (ClassLoaderPool) FieldUtils.readField(catalogManager, "classLoaderPool", true);
+    Assertions.assertEquals(1, pool.size());
+
+    CatalogWrapper wrapper = catalogManager.getCatalogCache().getIfPresent(ident);
+    Assertions.assertNotNull(wrapper);
+
+    BaseCatalog failingCatalog = Mockito.mock(BaseCatalog.class);
+    Mockito.doThrow(new IOException("close failed")).when(failingCatalog).close();
+    FieldUtils.writeField(wrapper, "catalog", failingCatalog, true);
+
+    // Cleanup runs exactly once, so a close() failure must not leave the reference behind: there
+    // is no second chance to clear it.
+    wrapper.retire();
+
+    Mockito.verify(failingCatalog).close();
+    Assertions.assertNull(
+        wrapper.catalog(), "a failing close must still drop the catalog reference");
+    Assertions.assertEquals(
+        0, pool.size(), "a failing close must still release the pooled ClassLoader");
+  }
+
   private NameIdentifier createCatalog(String name) {
     NameIdentifier ident = NameIdentifier.of(METALAKE, name);
     catalogManager.createCatalog(ident, Catalog.Type.RELATIONAL, PROVIDER, "comment", PROPS);
diff --git a/core/src/test/java/org/apache/gravitino/catalog/TestTableOperationDispatcher.java b/core/src/test/java/org/apache/gravitino/catalog/TestTableOperationDispatcher.java
index a834398c71f..ce53de28a7f 100644
--- a/core/src/test/java/org/apache/gravitino/catalog/TestTableOperationDispatcher.java
+++ b/core/src/test/java/org/apache/gravitino/catalog/TestTableOperationDispatcher.java
@@ -565,7 +565,8 @@ public void testDropTableCleansUpAutoDroppedParentSchemas() throws Exception {
     tableOperationDispatcher.createTable(tableIdent, columns, "comment", props, new Transform[0]);
 
     TestCatalog testCatalog =
-        (TestCatalog) catalogManager.loadCatalog(NameIdentifier.of(metalake, catalog));
+        (TestCatalog)
+            catalogManager.loadCatalogAndWrap(NameIdentifier.of(metalake, catalog)).catalog();
     TestCatalogOperations testCatalogOperations = (TestCatalogOperations) testCatalog.ops();
     Assertions.assertTrue(testCatalogOperations.dropSchema(schemaIdent, false));
     Assertions.assertFalse(testCatalogOperations.schemaExists(schemaIdent));
@@ -604,7 +605,8 @@ public void testDropMissingTableCleansUpSchemas() throws Exception {
     // now-empty namespaces, so the catalog no longer knows the table (dropTable returns false),
     // while Gravitino still holds the orphaned schema entities.
     TestCatalog testCatalog =
-        (TestCatalog) catalogManager.loadCatalog(NameIdentifier.of(metalake, catalog));
+        (TestCatalog)
+            catalogManager.loadCatalogAndWrap(NameIdentifier.of(metalake, catalog)).catalog();
     TestCatalogOperations testCatalogOperations = (TestCatalogOperations) testCatalog.ops();
     Assertions.assertTrue(testCatalogOperations.dropTable(tableIdent));
     Assertions.assertTrue(testCatalogOperations.dropSchema(schemaIdent, false));
@@ -646,7 +648,8 @@ public void testPurgeMissingTableCleansUpSchemas() throws Exception {
     // now-empty namespaces, so the catalog no longer knows the table (purgeTable returns false),
     // while Gravitino still holds the orphaned schema entities.
     TestCatalog testCatalog =
-        (TestCatalog) catalogManager.loadCatalog(NameIdentifier.of(metalake, catalog));
+        (TestCatalog)
+            catalogManager.loadCatalogAndWrap(NameIdentifier.of(metalake, catalog)).catalog();
     TestCatalogOperations testCatalogOperations = (TestCatalogOperations) testCatalog.ops();
     Assertions.assertTrue(testCatalogOperations.purgeTable(tableIdent));
     Assertions.assertTrue(testCatalogOperations.dropSchema(schemaIdent, false));
@@ -667,7 +670,8 @@ public void testCreateTableNeedImportingSchema() throws IOException {
     NameIdentifier tableIdent = NameIdentifier.of(tableNs, "topic81");
     Map props = ImmutableMap.of("k1", "v1", "k2", "v2");
     TestCatalog testCatalog =
-        (TestCatalog) catalogManager.loadCatalog(NameIdentifier.of(metalake, catalog));
+        (TestCatalog)
+            catalogManager.loadCatalogAndWrap(NameIdentifier.of(metalake, catalog)).catalog();
     TestCatalogOperations testCatalogOperations = (TestCatalogOperations) testCatalog.ops();
     testCatalogOperations.createSchema(
         NameIdentifier.of(tableNs.levels()), "", Collections.emptyMap());
@@ -736,7 +740,8 @@ public void testCreateAndLoadTableWithColumn() throws IOException {
 
     // Test if the column from table is not matched with the column from table entity
     TestCatalog testCatalog =
-        (TestCatalog) catalogManager.loadCatalog(NameIdentifier.of(metalake, catalog));
+        (TestCatalog)
+            catalogManager.loadCatalogAndWrap(NameIdentifier.of(metalake, catalog)).catalog();
     TestCatalogOperations testCatalogOperations = (TestCatalogOperations) testCatalog.ops();
 
     // 1. Update the existing column
diff --git a/core/src/test/java/org/apache/gravitino/catalog/TestTopicOperationDispatcher.java b/core/src/test/java/org/apache/gravitino/catalog/TestTopicOperationDispatcher.java
index 084f12132b8..83f5b9a23a2 100644
--- a/core/src/test/java/org/apache/gravitino/catalog/TestTopicOperationDispatcher.java
+++ b/core/src/test/java/org/apache/gravitino/catalog/TestTopicOperationDispatcher.java
@@ -265,7 +265,8 @@ public void testCreateTopicNeedImportingSchema() throws IOException {
     NameIdentifier topicIdent = NameIdentifier.of(topicNs, "topic61");
     Map props = ImmutableMap.of("k1", "v1", "k2", "v2");
     TestCatalog testCatalog =
-        (TestCatalog) catalogManager.loadCatalog(NameIdentifier.of(metalake, catalog));
+        (TestCatalog)
+            catalogManager.loadCatalogAndWrap(NameIdentifier.of(metalake, catalog)).catalog();
     TestCatalogOperations testCatalogOperations = (TestCatalogOperations) testCatalog.ops();
     testCatalogOperations.createSchema(
         NameIdentifier.of(topicNs.levels()), "", Collections.emptyMap());
diff --git a/core/src/test/java/org/apache/gravitino/catalog/TestViewOperationDispatcher.java b/core/src/test/java/org/apache/gravitino/catalog/TestViewOperationDispatcher.java
index 26372f662e6..43eadc5a0a1 100644
--- a/core/src/test/java/org/apache/gravitino/catalog/TestViewOperationDispatcher.java
+++ b/core/src/test/java/org/apache/gravitino/catalog/TestViewOperationDispatcher.java
@@ -156,7 +156,8 @@ public void testLoadView() throws IOException {
 
     // Mock the catalog operations to return the view
     TestCatalog testCatalog =
-        (TestCatalog) catalogManager.loadCatalog(NameIdentifier.of(metalake, catalog));
+        (TestCatalog)
+            catalogManager.loadCatalogAndWrap(NameIdentifier.of(metalake, catalog)).catalog();
     TestCatalogOperations testCatalogOperations = (TestCatalogOperations) testCatalog.ops();
     testCatalogOperations.views.put(viewIdent1, mockView);
 
@@ -188,7 +189,8 @@ public void testLoadViewWithMultipleViews() throws IOException {
 
     // Create multiple views
     TestCatalog testCatalog =
-        (TestCatalog) catalogManager.loadCatalog(NameIdentifier.of(metalake, catalog));
+        (TestCatalog)
+            catalogManager.loadCatalogAndWrap(NameIdentifier.of(metalake, catalog)).catalog();
     TestCatalogOperations testCatalogOperations = (TestCatalogOperations) testCatalog.ops();
 
     for (int i = 1; i <= 3; i++) {
@@ -258,7 +260,8 @@ public void testLoadViewAutoImportsIntoEntityStore() throws IOException {
     View mockView = createMockView("auto_import_view", props, auditInfo);
 
     TestCatalog testCatalog =
-        (TestCatalog) catalogManager.loadCatalog(NameIdentifier.of(metalake, catalog));
+        (TestCatalog)
+            catalogManager.loadCatalogAndWrap(NameIdentifier.of(metalake, catalog)).catalog();
     TestCatalogOperations testCatalogOperations = (TestCatalogOperations) testCatalog.ops();
     testCatalogOperations.views.put(viewIdent, mockView);
 
@@ -307,7 +310,8 @@ public void testLoadViewSkipsImportWhenAlreadyInEntityStore() throws IOException
     View mockView = createMockView("already_imported_view", props, auditInfo);
 
     TestCatalog testCatalog =
-        (TestCatalog) catalogManager.loadCatalog(NameIdentifier.of(metalake, catalog));
+        (TestCatalog)
+            catalogManager.loadCatalogAndWrap(NameIdentifier.of(metalake, catalog)).catalog();
     TestCatalogOperations testCatalogOperations = (TestCatalogOperations) testCatalog.ops();
     testCatalogOperations.views.put(viewIdent, mockView);
 
@@ -341,7 +345,8 @@ public void testLoadViewAutoImportWithMultipleConcurrentLoads() throws Exception
     View mockView = createMockView("concurrent_view", props, auditInfo);
 
     TestCatalog testCatalog =
-        (TestCatalog) catalogManager.loadCatalog(NameIdentifier.of(metalake, catalog));
+        (TestCatalog)
+            catalogManager.loadCatalogAndWrap(NameIdentifier.of(metalake, catalog)).catalog();
     TestCatalogOperations testCatalogOperations = (TestCatalogOperations) testCatalog.ops();
     testCatalogOperations.views.put(viewIdent, mockView);
 
@@ -404,7 +409,8 @@ public void testLoadViewAfterManualDelete() throws IOException {
     View mockView = createMockView("deleted_view", props, auditInfo);
 
     TestCatalog testCatalog =
-        (TestCatalog) catalogManager.loadCatalog(NameIdentifier.of(metalake, catalog));
+        (TestCatalog)
+            catalogManager.loadCatalogAndWrap(NameIdentifier.of(metalake, catalog)).catalog();
     TestCatalogOperations testCatalogOperations = (TestCatalogOperations) testCatalog.ops();
     testCatalogOperations.views.put(viewIdent, mockView);
 
@@ -529,7 +535,8 @@ public void testDropViewCleansUpAutoDroppedParentSchemas() throws Exception {
         viewIdent, null, new Column[0], representations, null, null, ImmutableMap.of("k1", "v1"));
 
     TestCatalog testCatalog =
-        (TestCatalog) catalogManager.loadCatalog(NameIdentifier.of(metalake, catalog));
+        (TestCatalog)
+            catalogManager.loadCatalogAndWrap(NameIdentifier.of(metalake, catalog)).catalog();
     TestCatalogOperations testCatalogOperations = (TestCatalogOperations) testCatalog.ops();
     Assertions.assertTrue(testCatalogOperations.dropSchema(schemaIdent, false));
     Assertions.assertFalse(testCatalogOperations.schemaExists(schemaIdent));
@@ -563,7 +570,8 @@ public void testDropMissingViewCleansUpSchemas() throws Exception {
     // now-empty namespaces, so the catalog no longer knows the view (dropView returns false),
     // while Gravitino still holds the orphaned schema entities.
     TestCatalog testCatalog =
-        (TestCatalog) catalogManager.loadCatalog(NameIdentifier.of(metalake, catalog));
+        (TestCatalog)
+            catalogManager.loadCatalogAndWrap(NameIdentifier.of(metalake, catalog)).catalog();
     TestCatalogOperations testCatalogOperations = (TestCatalogOperations) testCatalog.ops();
     Assertions.assertTrue(testCatalogOperations.dropView(viewIdent));
     Assertions.assertTrue(testCatalogOperations.dropSchema(schemaIdent, false));
diff --git a/core/src/test/java/org/apache/gravitino/hook/TestFilesetHookDispatcher.java b/core/src/test/java/org/apache/gravitino/hook/TestFilesetHookDispatcher.java
index c8d876e60d2..928a5d21b33 100644
--- a/core/src/test/java/org/apache/gravitino/hook/TestFilesetHookDispatcher.java
+++ b/core/src/test/java/org/apache/gravitino/hook/TestFilesetHookDispatcher.java
@@ -99,7 +99,7 @@ public static void initialize() throws Exception {
         Mockito.mock(CatalogManager.CatalogWrapper.class);
     Mockito.when(catalogWrapper.catalog()).thenReturn(catalog);
     Mockito.when(catalogWrapper.capabilities()).thenReturn(Capability.DEFAULT);
-    Mockito.when(catalogManager.loadCatalog(any())).thenReturn(catalog);
+    CatalogTestUtils.mockDoWithCatalog(catalogManager, catalog);
     Mockito.when(catalogManager.loadCatalogAndWrap(any())).thenReturn(catalogWrapper);
     Mockito.when(catalogManager.acquireCatalogLease(any()))
         .thenAnswer(invocation -> CatalogTestUtils.unmanagedLease(catalogWrapper));
diff --git a/core/src/test/java/org/apache/gravitino/hook/TestTopicHookDispatcher.java b/core/src/test/java/org/apache/gravitino/hook/TestTopicHookDispatcher.java
index d5bb231eb8f..35fcb8ca7f6 100644
--- a/core/src/test/java/org/apache/gravitino/hook/TestTopicHookDispatcher.java
+++ b/core/src/test/java/org/apache/gravitino/hook/TestTopicHookDispatcher.java
@@ -73,7 +73,7 @@ public static void initialize() throws Exception {
         Mockito.mock(CatalogManager.CatalogWrapper.class);
     Mockito.when(catalogWrapper.catalog()).thenReturn(catalog);
     Mockito.when(catalogWrapper.capabilities()).thenReturn(Capability.DEFAULT);
-    Mockito.when(catalogManager.loadCatalog(any())).thenReturn(catalog);
+    CatalogTestUtils.mockDoWithCatalog(catalogManager, catalog);
     Mockito.when(catalogManager.loadCatalogAndWrap(any())).thenReturn(catalogWrapper);
     Mockito.when(catalogManager.acquireCatalogLease(any()))
         .thenAnswer(invocation -> CatalogTestUtils.unmanagedLease(catalogWrapper));
diff --git a/iceberg/iceberg-rest-server/src/main/java/org/apache/gravitino/iceberg/service/provider/DynamicIcebergConfigProvider.java b/iceberg/iceberg-rest-server/src/main/java/org/apache/gravitino/iceberg/service/provider/DynamicIcebergConfigProvider.java
index 43ec66812e2..49d12df45f3 100644
--- a/iceberg/iceberg-rest-server/src/main/java/org/apache/gravitino/iceberg/service/provider/DynamicIcebergConfigProvider.java
+++ b/iceberg/iceberg-rest-server/src/main/java/org/apache/gravitino/iceberg/service/provider/DynamicIcebergConfigProvider.java
@@ -32,12 +32,11 @@
 import org.apache.gravitino.GravitinoEnv;
 import org.apache.gravitino.NameIdentifier;
 import org.apache.gravitino.auth.AuthProperties;
-import org.apache.gravitino.catalog.CatalogDispatcher;
+import org.apache.gravitino.catalog.CatalogManager;
 import org.apache.gravitino.catalog.lakehouse.iceberg.IcebergConstants;
 import org.apache.gravitino.client.DefaultOAuth2TokenProvider;
 import org.apache.gravitino.client.GravitinoClient;
 import org.apache.gravitino.client.GravitinoClient.ClientBuilder;
-import org.apache.gravitino.connector.BaseCatalog;
 import org.apache.gravitino.credential.JdbcCredential;
 import org.apache.gravitino.credential.SupportsCredentials;
 import org.apache.gravitino.exceptions.NoSuchCatalogException;
@@ -93,46 +92,12 @@ public Optional getIcebergCatalogConfig(String catalogName) {
                           IcebergConfig.ICEBERG_CONFIG_PREFIX
                               + IcebergConstants.ICEBERG_REST_DEFAULT_DYNAMIC_CATALOG_NAME)));
     }
-    Catalog catalog;
+    Map catalogProperties;
     try {
-      catalog = getCatalogFetcher().loadCatalog(catalogName);
+      catalogProperties = getCatalogFetcher().loadCatalogProperties(catalogName);
     } catch (NoSuchCatalogException e) {
       return Optional.empty();
     }
-
-    Preconditions.checkArgument(
-        "lakehouse-iceberg".equals(catalog.provider()),
-        String.format("%s.%s is not iceberg catalog", gravitinoMetalake, catalogName));
-
-    // Sensitive credentials (e.g. jdbc-password) are marked hidden in PropertiesMetadata and
-    // filtered out of catalog.properties(). We need two different strategies to recover them:
-    //
-    // Auxiliary mode: the catalog is a BaseCatalog running in the same JVM as the Gravitino
-    // server. Call propertiesWithCredentialProviders() which returns the raw entity properties
-    // including all hidden fields.
-    //
-    // Standalone mode: the catalog is a client-side object obtained via the Gravitino REST API.
-    // Call getCredentials() to retrieve vended credentials, then inject any JdbcCredential
-    // fields into the properties map so the JDBC backend can connect.
-    Map catalogProperties;
-    if (catalog instanceof BaseCatalog) {
-      catalogProperties = ((BaseCatalog) catalog).propertiesWithCredentialProviders();
-    } else {
-      catalogProperties = new HashMap<>(catalog.properties());
-      if (catalog instanceof SupportsCredentials) {
-        Arrays.stream(((SupportsCredentials) catalog).getCredentials())
-            .filter(c -> c instanceof JdbcCredential)
-            .map(c -> (JdbcCredential) c)
-            .findFirst()
-            .ifPresent(
-                jdbc -> {
-                  catalogProperties.putIfAbsent(
-                      IcebergConstants.GRAVITINO_JDBC_USER, jdbc.jdbcUser());
-                  catalogProperties.putIfAbsent(
-                      IcebergConstants.GRAVITINO_JDBC_PASSWORD, jdbc.jdbcPassword());
-                });
-      }
-    }
     return Optional.of(getIcebergConfigFromCatalogProperties(catalogProperties));
   }
 
@@ -266,13 +231,37 @@ void setCatalogFetcher(CatalogFetcher catalogFetcher) {
   interface CatalogFetcher extends Closeable {
     Catalog loadCatalog(String catalogName) throws NoSuchCatalogException;
 
+    default Map loadCatalogProperties(String catalogName)
+        throws NoSuchCatalogException {
+      Catalog catalog = loadCatalog(catalogName);
+      Preconditions.checkArgument(
+          "lakehouse-iceberg".equals(catalog.provider()),
+          String.format("Catalog %s is not an Iceberg catalog", catalogName));
+
+      Map catalogProperties = new HashMap<>(catalog.properties());
+      if (catalog instanceof SupportsCredentials) {
+        Arrays.stream(((SupportsCredentials) catalog).getCredentials())
+            .filter(c -> c instanceof JdbcCredential)
+            .map(c -> (JdbcCredential) c)
+            .findFirst()
+            .ifPresent(
+                jdbc -> {
+                  catalogProperties.putIfAbsent(
+                      IcebergConstants.GRAVITINO_JDBC_USER, jdbc.jdbcUser());
+                  catalogProperties.putIfAbsent(
+                      IcebergConstants.GRAVITINO_JDBC_PASSWORD, jdbc.jdbcPassword());
+                });
+      }
+      return catalogProperties;
+    }
+
     @Override
     default void close() {}
   }
 
   /**
-   * Internal catalog fetcher that uses CatalogDispatcher directly. This bypasses the HTTP layer and
-   * is used when running in auxiliary mode (embedded in Gravitino server).
+   * Internal catalog fetcher that uses the lease-aware CatalogManager API. This bypasses the HTTP
+   * layer and is used when running in auxiliary mode (embedded in Gravitino server).
    *
    * 
Note: When authorization is enabled (which requires auxiliary mode),
    * IcebergCatalogWrapperManager bypasses its cache to avoid consistency issues between the
@@ -280,22 +269,46 @@ default void close() {}
    */
   private static class InternalCatalogFetcher implements CatalogFetcher {
     private final String metalake;
-    private final CatalogDispatcher catalogDispatcher;
+    private final CatalogManager catalogManager;
 
     InternalCatalogFetcher(String metalake) {
       this.metalake = metalake;
-      CatalogDispatcher dispatcher = GravitinoEnv.getInstance().internalCatalogDispatcher();
+      CatalogManager manager;
+      try {
+        manager = GravitinoEnv.getInstance().catalogManager();
+      } catch (IllegalArgumentException e) {
+        throw new IllegalStateException(
+            "Internal CatalogManager is not available. "
+                + "Internal catalog fetcher requires running within Gravitino server.",
+            e);
+      }
       Preconditions.checkState(
-          dispatcher != null,
-          "Internal CatalogDispatcher is not available. "
+          manager != null,
+          "Internal CatalogManager is not available. "
               + "Internal catalog fetcher requires running within Gravitino server.");
-      this.catalogDispatcher = dispatcher;
+      this.catalogManager = manager;
     }
 
     @Override
     public Catalog loadCatalog(String catalogName) throws NoSuchCatalogException {
       NameIdentifier catalogIdent = NameIdentifierUtil.ofCatalog(metalake, catalogName);
-      return catalogDispatcher.loadCatalog(catalogIdent);
+      return catalogManager.loadCatalog(catalogIdent);
+    }
+
+    @Override
+    public Map loadCatalogProperties(String catalogName)
+        throws NoSuchCatalogException {
+      NameIdentifier catalogIdent = NameIdentifierUtil.ofCatalog(metalake, catalogName);
+      return catalogManager.doWithCatalog(
+          catalogIdent,
+          catalog -> {
+            Preconditions.checkArgument(
+                "lakehouse-iceberg".equals(catalog.provider()),
+                String.format("Catalog %s is not an Iceberg catalog", catalogName));
+            // Auxiliary mode needs raw properties, including hidden credentials. Copy them while
+            // the lease is held so no live BaseCatalog escapes into the caller.
+            return new HashMap<>(catalog.propertiesWithCredentialProviders());
+          });
     }
   }
 
diff --git a/iceberg/iceberg-rest-server/src/test/java/org/apache/gravitino/iceberg/service/provider/TestDynamicIcebergConfigProvider.java b/iceberg/iceberg-rest-server/src/test/java/org/apache/gravitino/iceberg/service/provider/TestDynamicIcebergConfigProvider.java
index d5b2024152f..231f2329185 100644
--- a/iceberg/iceberg-rest-server/src/test/java/org/apache/gravitino/iceberg/service/provider/TestDynamicIcebergConfigProvider.java
+++ b/iceberg/iceberg-rest-server/src/test/java/org/apache/gravitino/iceberg/service/provider/TestDynamicIcebergConfigProvider.java
@@ -33,13 +33,15 @@
 import org.apache.gravitino.Catalog;
 import org.apache.gravitino.GravitinoEnv;
 import org.apache.gravitino.NameIdentifier;
-import org.apache.gravitino.catalog.CatalogDispatcher;
+import org.apache.gravitino.catalog.CatalogManager;
 import org.apache.gravitino.catalog.lakehouse.iceberg.IcebergConstants;
+import org.apache.gravitino.connector.BaseCatalog;
 import org.apache.gravitino.exceptions.NoSuchCatalogException;
 import org.apache.gravitino.iceberg.common.IcebergConfig;
 import org.apache.gravitino.iceberg.common.ops.IcebergCatalogWrapper;
 import org.apache.gravitino.iceberg.service.authorization.IcebergRESTServerContext;
 import org.apache.gravitino.utils.NameIdentifierUtil;
+import org.apache.gravitino.utils.ThrowableFunction;
 import org.apache.iceberg.hive.HiveCatalog;
 import org.apache.iceberg.jdbc.JdbcCatalog;
 import org.junit.jupiter.api.AfterEach;
@@ -63,6 +65,7 @@ public void setUp() throws IllegalAccessException {
   public void tearDown() throws IllegalAccessException {
     // Clean up GravitinoEnv and IcebergRESTServerContext state after each test
     FieldUtils.writeField(GravitinoEnv.getInstance(), "internalCatalogDispatcher", null, true);
+    FieldUtils.writeField(GravitinoEnv.getInstance(), "catalogManager", null, true);
     resetServerContext();
   }
 
@@ -282,15 +285,13 @@ public void testInternalCatalogFetcher() throws IllegalAccessException {
     // Enable authorization to use internal fetcher
     createMockServerContext(true);
 
-    // Mock CatalogDispatchers
-    CatalogDispatcher mockCatalogDispatcher = Mockito.mock(CatalogDispatcher.class);
-    CatalogDispatcher mockInternalCatalogDispatcher = Mockito.mock(CatalogDispatcher.class);
-    Catalog mockCatalog = Mockito.mock(Catalog.class);
+    CatalogManager mockCatalogManager = Mockito.mock(CatalogManager.class);
+    BaseCatalog mockCatalog = Mockito.mock(BaseCatalog.class);
 
     NameIdentifier catalogIdent = NameIdentifierUtil.ofCatalog(metalakeName, catalogName);
-    Mockito.when(mockInternalCatalogDispatcher.loadCatalog(catalogIdent)).thenReturn(mockCatalog);
+    mockDoWithCatalog(mockCatalogManager, mockCatalog);
     Mockito.when(mockCatalog.provider()).thenReturn("lakehouse-iceberg");
-    Mockito.when(mockCatalog.properties())
+    Mockito.when(mockCatalog.propertiesWithCredentialProviders())
         .thenReturn(
             new HashMap() {
               {
@@ -299,14 +300,7 @@ public void testInternalCatalogFetcher() throws IllegalAccessException {
               }
             });
 
-    // Set the mock CatalogDispatchers to GravitinoEnv
-    FieldUtils.writeField(
-        GravitinoEnv.getInstance(), "catalogDispatcher", mockCatalogDispatcher, true);
-    FieldUtils.writeField(
-        GravitinoEnv.getInstance(),
-        "internalCatalogDispatcher",
-        mockInternalCatalogDispatcher,
-        true);
+    FieldUtils.writeField(GravitinoEnv.getInstance(), "catalogManager", mockCatalogManager, true);
 
     // Initialize provider with required properties
     Map properties = new HashMap<>();
@@ -315,12 +309,11 @@ public void testInternalCatalogFetcher() throws IllegalAccessException {
     DynamicIcebergConfigProvider provider = new DynamicIcebergConfigProvider();
     provider.initialize(properties);
 
-    // Test that internal interface is used (internal CatalogDispatcher should be called)
+    // Test that the internal lease-aware interface is used.
     Optional icebergConfig = provider.getIcebergCatalogConfig(catalogName);
 
     Assertions.assertTrue(icebergConfig.isPresent());
-    Mockito.verify(mockInternalCatalogDispatcher).loadCatalog(catalogIdent);
-    Mockito.verify(mockCatalogDispatcher, Mockito.never()).loadCatalog(catalogIdent);
+    Mockito.verify(mockCatalogManager).doWithCatalog(Mockito.eq(catalogIdent), Mockito.any());
   }
 
   @Test
@@ -362,15 +355,15 @@ public void testHttpCatalogFetcherUsedWhenAuthorizationDisabled() throws Illegal
   }
 
   @Test
-  public void testInternalCatalogFetcherWithNullCatalogDispatcher() throws IllegalAccessException {
+  public void testInternalCatalogFetcherWithNullCatalogManager() throws IllegalAccessException {
     String metalakeName = "test_metalake";
     String catalogName = "internal_catalog";
 
     // Enable authorization to use internal fetcher
     createMockServerContext(true);
 
-    // Ensure internal CatalogDispatcher is null (simulating GravitinoEnv not initialized)
-    FieldUtils.writeField(GravitinoEnv.getInstance(), "internalCatalogDispatcher", null, true);
+    // Ensure CatalogManager is null (simulating GravitinoEnv not initialized)
+    FieldUtils.writeField(GravitinoEnv.getInstance(), "catalogManager", null, true);
 
     // Initialize provider with required properties
     Map properties = new HashMap<>();
@@ -383,7 +376,7 @@ public void testInternalCatalogFetcherWithNullCatalogDispatcher() throws Illegal
         Assertions.assertThrows(
             IllegalStateException.class, () -> provider.getIcebergCatalogConfig(catalogName));
     Assertions.assertEquals(
-        "Internal CatalogDispatcher is not available. "
+        "Internal CatalogManager is not available. "
             + "Internal catalog fetcher requires running within Gravitino server.",
         exception.getMessage());
   }
@@ -396,22 +389,15 @@ public void testInternalCatalogFetcherNoSuchCatalogException() throws IllegalAcc
     // Enable authorization to use internal fetcher
     createMockServerContext(true);
 
-    // Mock internal CatalogDispatcher to throw NoSuchCatalogException
-    CatalogDispatcher mockCatalogDispatcher = Mockito.mock(CatalogDispatcher.class);
-    CatalogDispatcher mockInternalCatalogDispatcher = Mockito.mock(CatalogDispatcher.class);
+    // Mock the lease-aware CatalogManager to throw NoSuchCatalogException.
+    CatalogManager mockCatalogManager = Mockito.mock(CatalogManager.class);
     NameIdentifier catalogIdent =
         NameIdentifierUtil.ofCatalog(metalakeName, nonExistentCatalogName);
-    Mockito.when(mockInternalCatalogDispatcher.loadCatalog(catalogIdent))
-        .thenThrow(new NoSuchCatalogException("Catalog not found: %s", nonExistentCatalogName));
-
-    // Set the mock CatalogDispatchers to GravitinoEnv
-    FieldUtils.writeField(
-        GravitinoEnv.getInstance(), "catalogDispatcher", mockCatalogDispatcher, true);
-    FieldUtils.writeField(
-        GravitinoEnv.getInstance(),
-        "internalCatalogDispatcher",
-        mockInternalCatalogDispatcher,
-        true);
+    Mockito.doThrow(new NoSuchCatalogException("Catalog not found: %s", nonExistentCatalogName))
+        .when(mockCatalogManager)
+        .doWithCatalog(Mockito.eq(catalogIdent), Mockito.any());
+
+    FieldUtils.writeField(GravitinoEnv.getInstance(), "catalogManager", mockCatalogManager, true);
 
     // Initialize provider with required properties
     Map properties = new HashMap<>();
@@ -424,8 +410,7 @@ public void testInternalCatalogFetcherNoSuchCatalogException() throws IllegalAcc
     Optional result = provider.getIcebergCatalogConfig(nonExistentCatalogName);
 
     Assertions.assertFalse(result.isPresent());
-    Mockito.verify(mockInternalCatalogDispatcher).loadCatalog(catalogIdent);
-    Mockito.verify(mockCatalogDispatcher, Mockito.never()).loadCatalog(catalogIdent);
+    Mockito.verify(mockCatalogManager).doWithCatalog(Mockito.eq(catalogIdent), Mockito.any());
   }
 
   @Test
@@ -509,15 +494,13 @@ public void testConcurrentAccessWithInternalFetcher() throws Exception {
     // Enable authorization to use internal fetcher
     createMockServerContext(true);
 
-    // Mock CatalogDispatchers
-    CatalogDispatcher mockCatalogDispatcher = Mockito.mock(CatalogDispatcher.class);
-    CatalogDispatcher mockInternalCatalogDispatcher = Mockito.mock(CatalogDispatcher.class);
-    Catalog mockCatalog = Mockito.mock(Catalog.class);
+    CatalogManager mockCatalogManager = Mockito.mock(CatalogManager.class);
+    BaseCatalog mockCatalog = Mockito.mock(BaseCatalog.class);
 
     NameIdentifier catalogIdent = NameIdentifierUtil.ofCatalog(metalakeName, catalogName);
-    Mockito.when(mockInternalCatalogDispatcher.loadCatalog(catalogIdent)).thenReturn(mockCatalog);
+    mockDoWithCatalog(mockCatalogManager, mockCatalog);
     Mockito.when(mockCatalog.provider()).thenReturn("lakehouse-iceberg");
-    Mockito.when(mockCatalog.properties())
+    Mockito.when(mockCatalog.propertiesWithCredentialProviders())
         .thenReturn(
             new HashMap() {
               {
@@ -526,14 +509,7 @@ public void testConcurrentAccessWithInternalFetcher() throws Exception {
               }
             });
 
-    // Set the mock CatalogDispatchers to GravitinoEnv
-    FieldUtils.writeField(
-        GravitinoEnv.getInstance(), "catalogDispatcher", mockCatalogDispatcher, true);
-    FieldUtils.writeField(
-        GravitinoEnv.getInstance(),
-        "internalCatalogDispatcher",
-        mockInternalCatalogDispatcher,
-        true);
+    FieldUtils.writeField(GravitinoEnv.getInstance(), "catalogManager", mockCatalogManager, true);
 
     // Initialize provider with required properties
     Map properties = new HashMap<>();
@@ -576,12 +552,22 @@ public void testConcurrentAccessWithInternalFetcher() throws Exception {
       Assertions.assertTrue(result.isPresent(), "Each thread should get a valid config");
     }
 
-    // Verify internal CatalogDispatcher was called (at least once, possibly more due to
-    // concurrency)
-    Mockito.verify(mockInternalCatalogDispatcher, Mockito.atLeastOnce()).loadCatalog(catalogIdent);
-    Mockito.verify(mockCatalogDispatcher, Mockito.never()).loadCatalog(catalogIdent);
+    // Verify the internal lease-aware path was called (possibly more than once due to concurrency).
+    Mockito.verify(mockCatalogManager, Mockito.atLeastOnce())
+        .doWithCatalog(Mockito.eq(catalogIdent), Mockito.any());
 
     executor.shutdown();
     executor.awaitTermination(5, TimeUnit.SECONDS);
   }
+
+  @SuppressWarnings("unchecked")
+  private static void mockDoWithCatalog(CatalogManager catalogManager, BaseCatalog baseCatalog) {
+    Mockito.doAnswer(
+            invocation -> {
+              ThrowableFunction operation = invocation.getArgument(1);
+              return operation.apply(baseCatalog);
+            })
+        .when(catalogManager)
+        .doWithCatalog(Mockito.any(), Mockito.any());
+  }
 }

From 50f663ce69cc5541b7fe134363b103812afdf1ed Mon Sep 17 00:00:00 2001
From: yuqi 
Date: Wed, 12 Aug 2026 14:48:18 +0800
Subject: [PATCH 6/6] [#12403] test(server-common): update catalog lease mock

---
 .../gravitino/server/authorization/TestMetadataIdConverter.java | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/server-common/src/test/java/org/apache/gravitino/server/authorization/TestMetadataIdConverter.java b/server-common/src/test/java/org/apache/gravitino/server/authorization/TestMetadataIdConverter.java
index 3adfa55f021..8fe345f4ce9 100644
--- a/server-common/src/test/java/org/apache/gravitino/server/authorization/TestMetadataIdConverter.java
+++ b/server-common/src/test/java/org/apache/gravitino/server/authorization/TestMetadataIdConverter.java
@@ -216,7 +216,7 @@ void testConvertReturnsEmptyWhenParentCatalogDoesNotExist() throws IllegalAccess
     MetadataObject fileset =
         MetadataObjects.of(
             ImmutableList.of("missing_catalog", "schema", "fileset"), MetadataObject.Type.FILESET);
-    when(mockCatalogManager.loadCatalogAndWrap(NameIdentifier.of("metalake", "missing_catalog")))
+    when(mockCatalogManager.acquireCatalogLease(NameIdentifier.of("metalake", "missing_catalog")))
         .thenThrow(
             new NoSuchCatalogException("Catalog %s does not exist", "metalake.missing_catalog"));