[#12650] improvement(core): add OCC for fileset writes

catalogs/catalog-fileset/src/main/java/org/apache/gravitino/catalog/fileset/FilesetCatalogOperations.java

@@ -39,6 +39,7 @@
 import com.google.common.util.concurrent.ThreadFactoryBuilder;
 import java.io.FileNotFoundException;
 import java.io.IOException;
+import java.io.UncheckedIOException;
 import java.time.Instant;
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -709,59 +710,75 @@ public Fileset alterFileset(NameIdentifier ident, FilesetChange... changes)
   @Override
   public boolean dropFileset(NameIdentifier ident) {
     try {
-      FilesetEntity filesetEntity =
-          store.get(ident, Entity.EntityType.FILESET, FilesetEntity.class);
-
-      // For managed fileset, we should delete the related files.
-      if (!disableFSOps && filesetEntity.filesetType() == Fileset.Type.MANAGED) {
-        AtomicReference<IOException> exception = new AtomicReference<>();
-        Map<String, Path> storageLocations =
-            Maps.transformValues(filesetEntity.storageLocations(), Path::new);
-        storageLocations.forEach(
-            (locationName, location) -> {
-              try {
-                Map<String, String> fsConf =
-                    mergeUpLevelConfigurations(ident, filesetEntity.properties(), location);
-                FileSystem fs = getFileSystemWithCache(location, fsConf);
-                if (fs.exists(location)) {
-                  if (!fs.delete(location, true)) {
-                    LOG.warn(
-                        "Failed to delete fileset {} location {} with location name {}",
-                        ident,
-                        location,
-                        locationName);
+      // The relational store runs this cleanup after the metadata CAS wins but before committing
+      // its transaction. The callback therefore sees the exact deleted snapshot, and an I/O
+      // failure can still restore the metadata so the caller may fix permissions and retry.
+      //
+      // The price is that the recursive storage delete runs inside that transaction, holding the
+      // fileset rows and a pooled connection for as long as the filesystem takes. Dropping a
+      // fileset with a very large tree is therefore a slow write for that fileset, and enough
+      // concurrent drops can hold up the connection pool.
+      Optional<FilesetEntity> deletedFileset =
+          store.deleteAndGet(
+              ident,
+              Entity.EntityType.FILESET,
+              FilesetEntity.class,
+              filesetEntity -> {
+                if (!disableFSOps && filesetEntity.filesetType() == Fileset.Type.MANAGED) {
+                  try {
+                    deleteManagedFilesetStorage(ident, filesetEntity);
+                  } catch (IOException ioe) {
+                    throw new UncheckedIOException(ioe);
                   }
-                } else {
-                  LOG.warn(
-                      "Fileset {} location {} with location name {} does not exist",
-                      ident,
-                      location,
-                      locationName);
                 }
-              } catch (IOException ioe) {
-                LOG.warn(
-                    "Failed to delete fileset {} location {} with location name {}",
-                    ident,
-                    location,
-                    locationName,
-                    ioe);
-                exception.set(ioe);
-              }
-            });
-        if (exception.get() != null) {
-          throw exception.get();
-        }
-      }
-
-      return store.delete(ident, Entity.EntityType.FILESET);
+              });
+      return deletedFileset.isPresent();
     } catch (NoSuchEntityException ne) {
       LOG.warn("Fileset {} does not exist", ident);
       return false;
+    } catch (UncheckedIOException uioe) {
+      throw new RuntimeException("Failed to delete fileset " + ident, uioe.getCause());
     } catch (IOException ioe) {
       throw new RuntimeException("Failed to delete fileset " + ident, ioe);
     }
   }
 
+  /**
+   * Removes the storage of a managed fileset while its metadata delete can still be rolled back.
+   *
+   * <p>The first location that cannot be removed stops the loop, so the drop is rejected before it
+   * takes away more data than it already has. The locations removed up to that point are gone for
+   * good, but attempting the remaining ones would only widen that gap.
+   */
+  private void deleteManagedFilesetStorage(NameIdentifier ident, FilesetEntity filesetEntity)
+      throws IOException {
+    Map<String, Path> storageLocations =
+        Maps.transformValues(filesetEntity.storageLocations(), Path::new);
+    for (Map.Entry<String, Path> entry : storageLocations.entrySet()) {
+      String locationName = entry.getKey();
+      Path location = entry.getValue();
+      Map<String, String> fsConf =
+          mergeUpLevelConfigurations(ident, filesetEntity.properties(), location);
+      FileSystem fs = getFileSystemWithCache(location, fsConf);
+      if (!fs.exists(location)) {
+        LOG.warn(
+            "Fileset {} location {} with location name {} does not exist",
+            ident,
+            location,
+            locationName);
+        continue;
+      }
+      if (!fs.delete(location, true) && fs.exists(location)) {
+        // A false return also covers a location that somebody else removed between the check above
+        // and this call. Only a location that is still there is a reason to reject the drop.
+        throw new IOException(
+            String.format(
+                "Failed to delete fileset %s location %s with location name %s",
+                ident, location, locationName));
+      }
+    }
+  }
+
   @Override
   public String getFileLocation(NameIdentifier ident, String subPath, String locationName)
       throws NoSuchFilesetException, NoSuchLocationNameException {

catalogs/catalog-fileset/src/test/java/org/apache/gravitino/catalog/fileset/TestFilesetCatalogOperations.java

@@ -76,6 +76,7 @@
 import org.apache.gravitino.Catalog;
 import org.apache.gravitino.Config;
 import org.apache.gravitino.Configs;
+import org.apache.gravitino.Entity;
 import org.apache.gravitino.EntityStore;
 import org.apache.gravitino.EntityStoreFactory;
 import org.apache.gravitino.GravitinoEnv;
@@ -105,6 +106,7 @@
 import org.apache.gravitino.exceptions.NoSuchFilesetException;
 import org.apache.gravitino.exceptions.NoSuchSchemaException;
 import org.apache.gravitino.exceptions.NonEmptySchemaException;
+import org.apache.gravitino.exceptions.OptimisticLockException;
 import org.apache.gravitino.exceptions.SchemaAlreadyExistsException;
 import org.apache.gravitino.file.FileInfo;
 import org.apache.gravitino.file.Fileset;
@@ -3386,6 +3388,115 @@ private static Stream<Arguments> testRenameArguments() {
             TEST_ROOT_PATH + "/fileset39"));
   }
 
+  @Test
+  public void testDropFilesetKeepsFilesWhenMetadataDropIsRejected() throws IOException {
+    String schemaName = "schema_drop_rejected";
+    String filesetName = "fileset_drop_rejected";
+    String catalogPath = TEST_ROOT_PATH + "/catalog_drop_rejected";
+    createSchema(schemaName, "comment", catalogPath, null);
+    Fileset fileset =
+        createFileset(filesetName, schemaName, "comment", Fileset.Type.MANAGED, catalogPath, null);
+
+    Path filesetPath = new Path(fileset.storageLocation());
+    FileSystem fs = filesetPath.getFileSystem(new Configuration());
+    Assertions.assertTrue(fs.exists(filesetPath));
+
+    NameIdentifier filesetIdent = NameIdentifier.of("m1", "c1", schemaName, filesetName);
+    EntityStore rejectingStore = Mockito.spy(store);
+    Mockito.doThrow(new OptimisticLockException("fileset was modified concurrently"))
+        .when(rejectingStore)
+        .deleteAndGet(
+            Mockito.eq(filesetIdent),
+            Mockito.eq(Entity.EntityType.FILESET),
+            Mockito.eq(FilesetEntity.class),
+            Mockito.any());
+
+    try (FilesetCatalogOperations ops =
+        new FilesetCatalogOperations(rejectingStore, secretManager)) {
+      ops.initialize(
+          ImmutableMap.of(LOCATION, catalogPath),
+          randomCatalogInfo("m1", "c1"),
+          FILESET_PROPERTIES_METADATA);
+      Assertions.assertThrows(OptimisticLockException.class, () -> ops.dropFileset(filesetIdent));
+    }
+
+    // The drop was refused, so the fileset row still advertises this location. Deleting the files
+    // anyway would leave that row pointing at data that is gone.
+    Assertions.assertTrue(fs.exists(filesetPath));
+    Assertions.assertEquals(
+        filesetName,
+        store.get(filesetIdent, Entity.EntityType.FILESET, FilesetEntity.class).name());
+
+    fs.delete(filesetPath, true);
+  }
+
+  @Test
+  public void testDropFilesetRollsBackMetadataWhenStorageDeletionFails() throws IOException {
+    String schemaName = "schema_drop_storage_failure";
+    String filesetName = "fileset_drop_storage_failure";
+    String catalogPath = TEST_ROOT_PATH + "/catalog_drop_storage_failure";
+    createSchema(schemaName, "comment", catalogPath, null);
+    Fileset fileset =
+        createFileset(filesetName, schemaName, "comment", Fileset.Type.MANAGED, catalogPath, null);
+    NameIdentifier filesetIdent = NameIdentifier.of("m1", "c1", schemaName, filesetName);
+
+    FileSystem failingFileSystem = Mockito.mock(FileSystem.class);
+    when(failingFileSystem.exists(any(Path.class))).thenReturn(true);
+    when(failingFileSystem.delete(any(Path.class), Mockito.eq(true)))
+        .thenThrow(new IOException("permission denied"));
+
+    try (FilesetCatalogOperations ops =
+        Mockito.spy(new FilesetCatalogOperations(store, secretManager))) {
+      ops.initialize(
+          ImmutableMap.of(LOCATION, catalogPath),
+          randomCatalogInfo("m1", "c1"),
+          FILESET_PROPERTIES_METADATA);
+      doReturn(failingFileSystem).when(ops).getFileSystemWithCache(any(Path.class), any(Map.class));
+
+      RuntimeException failure =
+          Assertions.assertThrows(RuntimeException.class, () -> ops.dropFileset(filesetIdent));
+      Assertions.assertTrue(failure.getMessage().contains("Failed to delete fileset"));
+    }
+
+    // The failed physical cleanup aborts the outer transaction, so the fileset remains visible and
+    // a caller can repair its filesystem permissions and retry the drop.
+    FilesetEntity survivingFileset =
+        store.get(filesetIdent, Entity.EntityType.FILESET, FilesetEntity.class);
+    Assertions.assertEquals(fileset.storageLocation(), survivingFileset.storageLocation());
+
+    store.delete(filesetIdent, Entity.EntityType.FILESET);
+    new Path(fileset.storageLocation())
+        .getFileSystem(new Configuration())
+        .delete(new Path(fileset.storageLocation()), true);
+  }
+
+  @Test
+  public void testDropFilesetSucceedsWhenTheLocationDisappearsFirst() throws IOException {
+    String schemaName = "schema_drop_vanished";
+    String filesetName = "fileset_drop_vanished";
+    String catalogPath = TEST_ROOT_PATH + "/catalog_drop_vanished";
+    createSchema(schemaName, "comment", catalogPath, null);
+    Fileset fileset =
+        createFileset(filesetName, schemaName, "comment", Fileset.Type.MANAGED, catalogPath, null);
+
+    Path filesetPath = new Path(fileset.storageLocation());
+    FileSystem fs = filesetPath.getFileSystem(new Configuration());
+    Assertions.assertTrue(fs.exists(filesetPath));
+    // Somebody else removed the location already. A drop that finds nothing left to delete has
+    // nothing to complain about.
+    Assertions.assertTrue(fs.delete(filesetPath, true));
+
+    NameIdentifier filesetIdent = NameIdentifier.of("m1", "c1", schemaName, filesetName);
+    try (FilesetCatalogOperations ops = new FilesetCatalogOperations(store, secretManager)) {
+      ops.initialize(
+          ImmutableMap.of(LOCATION, catalogPath),
+          randomCatalogInfo("m1", "c1"),
+          FILESET_PROPERTIES_METADATA);
+      Assertions.assertTrue(ops.dropFileset(filesetIdent));
+      Assertions.assertFalse(ops.dropFileset(filesetIdent), "fileset should be non-existent");
+    }
+  }
+
   private Schema createSchema(String name, String comment, String catalogPath, String schemaPath)
       throws IOException {
     return createSchema(name, comment, catalogPath, schemaPath, false);

core/src/main/java/org/apache/gravitino/EntityStore.java

@@ -23,6 +23,8 @@
 import java.lang.reflect.Array;
 import java.util.Arrays;
 import java.util.List;
+import java.util.Optional;
+import java.util.function.Consumer;
 import java.util.function.Function;
 import org.apache.commons.lang3.tuple.Pair;
 import org.apache.gravitino.Entity.EntityType;
@@ -221,6 +223,83 @@ default boolean delete(NameIdentifier ident, EntityType entityType) throws IOExc
    */
   boolean delete(NameIdentifier ident, EntityType entityType, boolean cascade) throws IOException;
 
+  /**
+   * The only post-delete action an implementation that cannot run it before commit accepts.
+   *
+   * <p>Compared by reference, so a caller that supplies its own action reaches an implementation
+   * that honors the contract or gets told that this one cannot.
+   */
+  Consumer<? extends Entity> NO_POST_DELETE_ACTION = ignored -> {};
+
+  /**
+   * Returns the shared no-op post-delete action.
+   *
+   * @param <E> the entity type
+   * @return an action that does nothing
+   */
+  @SuppressWarnings("unchecked")
+  static <E extends Entity & HasIdentifier> Consumer<E> noPostDeleteAction() {
+    return (Consumer<E>) NO_POST_DELETE_ACTION;
+  }
+
+  /**
+   * Deletes an entity and returns the snapshot chosen by the delete operation.
+   *
+   * <p>The default implementation is intended for stores that serialize operations through {@link
+   * #executeInTransaction(Executable)}. Stores that can read and delete with one native
+   * compare-and-set should override this method so the returned snapshot is exactly the one that
+   * was deleted.
+   *
+   * @param ident the name identifier of the entity
+   * @param entityType the type of the entity
+   * @param clazz the concrete entity class
+   * @param <E> the entity type
+   * @return the deleted entity, or empty when it did not exist
+   * @throws IOException if the delete operation fails
+   */
+  default <E extends Entity & HasIdentifier> Optional<E> deleteAndGet(
+      NameIdentifier ident, EntityType entityType, Class<E> clazz) throws IOException {
+    return deleteAndGet(ident, entityType, clazz, noPostDeleteAction());
+  }
+
+  /**
+   * Deletes an entity, runs an action against the deleted snapshot, and returns that snapshot.
+   *
+   * <p>A transactional store should run the action after its delete has won but before committing.
+   * This lets callers couple non-database cleanup to the metadata transaction: an action failure
+   * can still roll the metadata delete back.
+   *
+   * @param ident the name identifier of the entity
+   * @param entityType the type of the entity
+   * @param clazz the concrete entity class
+   * @param postDeleteAction the action to run after deletion but before commit when supported
+   * @param <E> the entity type
+   * @return the deleted entity, or empty when it did not exist
+   * @throws IOException if the delete operation fails
+   */
+  default <E extends Entity & HasIdentifier> Optional<E> deleteAndGet(
+      NameIdentifier ident, EntityType entityType, Class<E> clazz, Consumer<E> postDeleteAction)
+      throws IOException {
+    if (postDeleteAction != NO_POST_DELETE_ACTION) {
+      // This implementation can only run the action once the delete is committed, which is the
+      // opposite of what the contract promises. Refusing is better than silently leaving the
+      // caller with a committed delete and a failed cleanup.
+      throw new UnsupportedOperationException(
+          "This store cannot run a post-delete action while the delete can still be rolled back");
+    }
+
+    try {
+      E entity = get(ident, entityType, clazz);
+      if (!delete(ident, entityType)) {
+        return Optional.empty();
+      }
+      postDeleteAction.accept(entity);
+      return Optional.of(entity);
+    } catch (NoSuchEntityException e) {
+      return Optional.empty();
+    }
+  }
+
   /**
    * Batch delete entities from the underlying storage by the specified list of {@link
    * org.apache.gravitino.NameIdentifier} and {@link EntityType}.