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/CapabilityHelpers.java b/core/src/main/java/org/apache/gravitino/catalog/CapabilityHelpers.java index eae018c2c0d..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,11 +52,16 @@ public class CapabilityHelpers { public static Capability getCapability(NameIdentifier ident, CatalogManager catalogManager) { NameIdentifier catalogIdent = getCatalogIdentifier(ident); - CatalogManager.CatalogWrapper c = catalogManager.loadCatalogAndWrap(catalogIdent); + // Acquire the lease outside the try so a missing catalog keeps propagating its + // NoSuchCatalogException (a 404) instead of being wrapped into a plain RuntimeException (a + // 500); only the capability lookup itself is wrapped. + CatalogLease lease = catalogManager.acquireCatalogLease(catalogIdent); try { - return c.capabilities(); + return lease.wrapper().capabilities(); } catch (Exception e) { throw new RuntimeException("Failed to get capabilities for catalog: " + catalogIdent, e); + } finally { + lease.close(); } } 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..5cc7e80e1a7 --- /dev/null +++ b/core/src/main/java/org/apache/gravitino/catalog/CatalogLease.java @@ -0,0 +1,78 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.gravitino.catalog; + +import java.util.concurrent.atomic.AtomicBoolean; +import org.apache.gravitino.NameIdentifier; +import org.apache.gravitino.catalog.CatalogManager.CatalogWrapper; +import org.apache.gravitino.connector.BaseCatalog; + +/** + * A lease on a {@link CatalogWrapper} held for the duration of one catalog operation. + * + *

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(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; + } + + /** + * 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 72c0139cf43..c6eb0699e9f 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; @@ -124,6 +125,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,18 +162,48 @@ 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; + // 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; - 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) { this.classLoader = classLoader; + this.pool = null; + this.poolEntry = null; } /** Pooled constructor: ClassLoader is managed by the pool with reference counting. */ @@ -182,6 +218,79 @@ 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. + */ + 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. + * + *

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. + */ + 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()}. + */ + 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. + */ + 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,20 +394,43 @@ 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() { + // 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) { @@ -306,7 +438,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; @@ -368,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). @@ -396,12 +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(); } - ((CatalogWrapper) v).close(); }) .scheduler( Scheduler.forScheduledExecutorService( @@ -430,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(); } /** @@ -560,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 @@ -569,12 +719,38 @@ 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 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. * @@ -583,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. */ @@ -628,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; @@ -748,14 +925,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 +966,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, @@ -828,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. */ @@ -841,31 +1018,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; }); @@ -911,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); @@ -937,8 +1115,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 +1263,50 @@ 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 { + 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); + } + + // 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 +1314,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 +1513,27 @@ 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 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) { 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/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 new file mode 100644 index 00000000000..0e1e51d63ef --- /dev/null +++ b/core/src/test/java/org/apache/gravitino/catalog/CatalogTestUtils.java @@ -0,0 +1,61 @@ +/* + * 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; +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 { + + 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); + } + + /** + * 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/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(); + } } 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..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() @@ -964,6 +966,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()); @@ -1019,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() @@ -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()); @@ -1070,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() @@ -1094,6 +1094,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()); @@ -1122,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()) @@ -1141,6 +1142,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(); @@ -1164,17 +1166,16 @@ 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); 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()); @@ -1201,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(); @@ -1222,20 +1220,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..84e3aa936ca --- /dev/null +++ b/core/src/test/java/org/apache/gravitino/catalog/TestCatalogWrapperLease.java @@ -0,0 +1,439 @@ +/* + * 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.connector.BaseCatalog; +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; +import org.mockito.Mockito; + +/** + * 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(expiringConfig)); + 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 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"); + CatalogWrapper wrapper = catalogManager.getCatalogCache().getIfPresent(ident); + Assertions.assertNotNull(wrapper); + + 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); + 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..cd508ce04af 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,12 @@ public void setUp() { when(catalogManager.loadCatalogAndWrap(NameIdentifier.of(METALAKE, ICEBERG_CATALOG))) .thenReturn(icebergWrapper); + when(catalogManager.acquireCatalogLease(NameIdentifier.of(METALAKE, ICEBERG_CATALOG))) + .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 -> 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 84c9678a668..35d02c89462 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,8 @@ 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 -> 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 85649cdea18..1fc71d4f5cf 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,8 @@ 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 -> CatalogTestUtils.unmanagedLease(mockWrapper)); TableNormalizeDispatcher dispatcher = new TableNormalizeDispatcher(mockDispatcher, mockCatalogManager); 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 fe0447fa750..347787b84a1 100644 --- a/core/src/test/java/org/apache/gravitino/hook/TestFilesetHookDispatcher.java +++ b/core/src/test/java/org/apache/gravitino/hook/TestFilesetHookDispatcher.java @@ -55,6 +55,7 @@ import org.apache.gravitino.authorization.Owner; import org.apache.gravitino.authorization.OwnerDispatcher; 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; @@ -98,8 +99,10 @@ 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)); authorizationPlugin = Mockito.mock(AuthorizationPlugin.class); Mockito.when(catalog.getAuthorizationPlugin()).thenReturn(authorizationPlugin); } @@ -115,6 +118,8 @@ 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 -> 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 faba932cd6a..a8dfde594cd 100644 --- a/core/src/test/java/org/apache/gravitino/hook/TestFunctionHookDispatcher.java +++ b/core/src/test/java/org/apache/gravitino/hook/TestFunctionHookDispatcher.java @@ -32,6 +32,7 @@ import org.apache.gravitino.authorization.Owner; import org.apache.gravitino.authorization.OwnerDispatcher; 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; @@ -64,6 +65,8 @@ 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 -> CatalogTestUtils.unmanagedLease(catalogWrapper)); Mockito.when( dispatcher.registerFunction( @@ -147,6 +150,8 @@ 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 -> CatalogTestUtils.unmanagedLease(mockWrapper)); OwnerDispatcher mockOwnerDispatcher = Mockito.mock(OwnerDispatcher.class); FunctionDispatcher mockFunctionDispatcher = Mockito.mock(FunctionDispatcher.class); @@ -200,6 +205,8 @@ 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 -> 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 c6777b58919..97d9bf263a6 100644 --- a/core/src/test/java/org/apache/gravitino/hook/TestModelHookDispatcher.java +++ b/core/src/test/java/org/apache/gravitino/hook/TestModelHookDispatcher.java @@ -33,6 +33,7 @@ import org.apache.gravitino.authorization.Owner; import org.apache.gravitino.authorization.OwnerDispatcher; 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; @@ -62,6 +63,8 @@ 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 -> 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 3e05c3c83ce..539e9276cd0 100644 --- a/core/src/test/java/org/apache/gravitino/hook/TestSchemaHookDispatcher.java +++ b/core/src/test/java/org/apache/gravitino/hook/TestSchemaHookDispatcher.java @@ -46,6 +46,7 @@ import org.apache.gravitino.authorization.Owner; import org.apache.gravitino.authorization.OwnerDispatcher; 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; @@ -78,6 +79,8 @@ 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 -> 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 bf9431a504c..7d0867ae342 100644 --- a/core/src/test/java/org/apache/gravitino/hook/TestTableHookDispatcher.java +++ b/core/src/test/java/org/apache/gravitino/hook/TestTableHookDispatcher.java @@ -35,6 +35,7 @@ import org.apache.gravitino.authorization.Owner; import org.apache.gravitino.authorization.OwnerDispatcher; 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; @@ -85,6 +86,8 @@ 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 -> CatalogTestUtils.unmanagedLease(wrapper)); OwnerDispatcher ownerDispatcher = Mockito.mock(OwnerDispatcher.class); TableDispatcher dispatcher = Mockito.mock(TableDispatcher.class); @@ -151,6 +154,8 @@ 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 -> 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 539af87f5ab..35fcb8ca7f6 100644 --- a/core/src/test/java/org/apache/gravitino/hook/TestTopicHookDispatcher.java +++ b/core/src/test/java/org/apache/gravitino/hook/TestTopicHookDispatcher.java @@ -32,6 +32,7 @@ import org.apache.gravitino.authorization.Owner; import org.apache.gravitino.authorization.OwnerDispatcher; 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; @@ -72,8 +73,10 @@ 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)); authorizationPlugin = Mockito.mock(AuthorizationPlugin.class); Mockito.when(catalog.getAuthorizationPlugin()).thenReturn(authorizationPlugin); } @@ -89,6 +92,8 @@ 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 -> CatalogTestUtils.unmanagedLease(mockWrapper)); 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/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/dispatcher/TestIcebergAsyncPurge.java b/iceberg/iceberg-rest-server/src/test/java/org/apache/gravitino/iceberg/service/dispatcher/TestIcebergAsyncPurge.java index 3ee6bf69071..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 @@ -33,6 +33,7 @@ import org.apache.gravitino.GravitinoEnv; import org.apache.gravitino.auth.AuthConstants; 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; @@ -241,6 +242,8 @@ 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 -> CatalogTestUtils.unmanagedLease(wrapper)); when(wrapper.catalog()).thenReturn(catalog); when(catalog.entity()).thenReturn(entity); when(entity.id()).thenReturn(CATALOG_ID); 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()); + } } 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"));