[#12345] improvement(core): add OCC for table writes

catalogs/catalog-lakehouse-generic/src/main/java/org/apache/gravitino/catalog/lakehouse/lance/LanceTableOperations.java

@@ -50,6 +50,7 @@
 import org.apache.gravitino.exceptions.NoSuchEntityException;
 import org.apache.gravitino.exceptions.NoSuchSchemaException;
 import org.apache.gravitino.exceptions.NoSuchTableException;
+import org.apache.gravitino.exceptions.OptimisticLockException;
 import org.apache.gravitino.exceptions.TableAlreadyExistsException;
 import org.apache.gravitino.lance.common.ops.gravitino.LanceDataTypeConverter;
 import org.apache.gravitino.lance.common.utils.LanceConstants;
@@ -65,7 +66,6 @@
 import org.apache.gravitino.rel.expressions.transforms.Transform;
 import org.apache.gravitino.rel.indexes.Index;
 import org.apache.gravitino.storage.IdGenerator;
-import org.apache.gravitino.storage.relational.service.TableMetaService;
 import org.apache.gravitino.utils.PrincipalUtils;
 import org.lance.Dataset;
 import org.lance.ReadOptions;
@@ -561,30 +561,20 @@ private Table repairTableMetadata(NameIdentifier ident, Column[] columns, long d
   }
 
   /**
-   * Applies an idempotent update to the stored table, retrying when the optimistic-lock CAS is lost
-   * to a concurrent update. The repair-on-load path runs on every {@code loadTable}, so concurrent
-   * loads of the same table race on the version CAS; {@code store.update} surfaces the lost race as
-   * an {@link IOException} whose message starts with {@link
-   * TableMetaService#UPDATE_ENTITY_CONFLICT_MESSAGE_PREFIX}. Because the updater is idempotent, the
-   * loser sleeps a short randomized backoff (to avoid re-colliding), re-reads the latest (already
-   * repaired) entity, and retries instead of failing the whole load with a fatal error. Other IO
-   * failures (DB outage, serialization errors, etc.) are not conflicts and fail fast.
+   * Repairs stored table metadata and retries only when another repair wins the same race.
+   *
+   * <p>Two concurrent loads can both read the old metadata. The first update wins; the second gets
+   * {@link OptimisticLockException}. Repair is safe to run more than once, so the loser waits
+   * briefly, reads the winner's latest row, and tries again. Ordinary IO failures are not safe to
+   * retry here and are returned immediately.
    */
   private TableEntity updateTableWithCasRetry(
       NameIdentifier ident, Function<TableEntity, TableEntity> updater) throws IOException {
-    IOException lastConflict = null;
+    OptimisticLockException lastConflict = null;
     for (int attempt = 1; attempt <= REPAIR_UPDATE_MAX_ATTEMPTS; attempt++) {
       try {
         return store.update(ident, TableEntity.class, Entity.EntityType.TABLE, updater);
-      } catch (IOException e) {
-        // Only retry when the update matched 0 rows (lost optimistic-lock CAS). Other IO failures
-        // (DB outage, serialization errors, etc.) should fail fast.
-        String message = e.getMessage();
-        if (message == null
-            || !message.startsWith(TableMetaService.UPDATE_ENTITY_CONFLICT_MESSAGE_PREFIX)) {
-          throw e;
-        }
-
+      } catch (OptimisticLockException e) {
         lastConflict = e;
         LOG.debug(
             "Optimistic-lock conflict updating table {} metadata (attempt {}/{}), {}",
@@ -599,11 +589,13 @@ private TableEntity updateTableWithCasRetry(
         }
       }
     }
-    throw new IOException(
-        String.format(
-            "Failed to update table %s after %d optimistic-lock retries",
-            ident, REPAIR_UPDATE_MAX_ATTEMPTS),
-        lastConflict);
+    // Reaching here means every attempt lost to another writer. Keep the exception type so the
+    // caller still knows this is an OCC conflict, and add the attempt count for diagnosis.
+    throw new OptimisticLockException(
+        lastConflict,
+        "Failed to repair table %s after %d optimistic-lock attempts",
+        ident,
+        REPAIR_UPDATE_MAX_ATTEMPTS);
   }
 
   /**

catalogs/catalog-lakehouse-generic/src/test/java/org/apache/gravitino/catalog/lakehouse/lance/TestLanceConcurrentRepairStress.java

@@ -50,6 +50,7 @@
 import org.apache.gravitino.NameIdentifier;
 import org.apache.gravitino.UserPrincipal;
 import org.apache.gravitino.catalog.ManagedSchemaOperations;
+import org.apache.gravitino.exceptions.OptimisticLockException;
 import org.apache.gravitino.meta.AuditInfo;
 import org.apache.gravitino.meta.TableEntity;
 import org.apache.gravitino.rel.Table;
@@ -66,16 +67,16 @@
 /**
  * Real multi-threaded reproduction of the repair-on-load optimistic-lock race behind #11891. Unlike
  * {@code TestLanceTableOperations#testLoadTableSurvivesConcurrentRepairVersionRace} (a
- * deterministic mock that throws one scripted {@code IOException}), this test drives {@link
+ * deterministic mock that throws one scripted conflict), this test drives {@link
  * LanceTableOperations#loadTable} from several threads at once against a {@link CasEntityStore}
  * that models the production relational store's compare-and-set semantics faithfully: every {@code
  * update} bumps a version guarded by the base version, so concurrent updates from the same base
  * conflict and exactly one wins per generation — just like {@code TableMetaService.updateTable}
  * ({@code UPDATE ... WHERE current_version = old}).
  *
- * <p>Before the CAS retry, the loser of the race got {@code IOException("Failed to update the
- * entity")}, rethrown as a fatal {@code RuntimeException} (HTTP 500). This test asserts every
- * concurrent load returns the repaired table instead.
+ * <p>Before the CAS retry, the loser of the race got an {@link OptimisticLockException}, rethrown
+ * as a fatal error (HTTP 500). This test asserts every concurrent load returns the repaired table
+ * instead.
  */
 public class TestLanceConcurrentRepairStress {
 
@@ -189,10 +190,10 @@ private static TableEntity tableEntity(
   /**
    * In-memory {@link EntityStore} that reproduces the relational store's optimistic-lock CAS:
    * {@code update} reads a versioned snapshot, applies the (idempotent) updater, and commits only
-   * if the version has not advanced since the read — otherwise it throws {@code IOException("Failed
-   * to update the entity")}, exactly as {@code TableMetaService.updateTable} does when {@code
-   * UPDATE ... WHERE current_version = old} matches zero rows. Every commit bumps the version, so
-   * even a no-op update invalidates a concurrent update from the same base, matching production.
+   * if the version has not advanced since the read — otherwise it throws an {@link
+   * OptimisticLockException}, exactly as the relational table service does when {@code UPDATE ...
+   * WHERE current_version = old} matches zero rows. Every commit bumps the version, so even a no-op
+   * update invalidates a concurrent update from the same base, matching production.
    */
   private static final class CasEntityStore implements EntityStore {
 
@@ -230,7 +231,7 @@ public <E extends Entity & HasIdentifier> E update(
       if (ref.compareAndSet(base, next)) {
         return updated;
       }
-      throw new IOException("Failed to update the entity: " + ident);
+      throw new OptimisticLockException("mock conflict for %s", ident);
     }
 
     // --- unused surface ---------------------------------------------------

catalogs/catalog-lakehouse-generic/src/test/java/org/apache/gravitino/catalog/lakehouse/lance/TestLanceTableOperations.java

@@ -48,6 +48,7 @@
 import org.apache.gravitino.NameIdentifier;
 import org.apache.gravitino.UserPrincipal;
 import org.apache.gravitino.catalog.ManagedSchemaOperations;
+import org.apache.gravitino.exceptions.OptimisticLockException;
 import org.apache.gravitino.meta.AuditInfo;
 import org.apache.gravitino.meta.ColumnEntity;
 import org.apache.gravitino.meta.TableEntity;
@@ -174,10 +175,10 @@ public void testLoadDeclaredTableSchemaFromLocation() throws Exception {
   /**
    * Reproduces the concurrent repair-on-load race seen in {@code LanceSparkRESTServiceIT}. When two
    * loads repair the same table at once, the optimistic-locked {@code store.update} of the slower
-   * one matches zero rows and {@code TableMetaService} surfaces it as {@code IOException("Failed to
-   * update the entity")}. Before the CAS retry, {@code repairTableMetadata} rethrew it as a fatal
-   * {@code RuntimeException} (HTTP 500) instead of tolerating the concurrent update. This test
-   * asserts that the lost race is benign and load returns a usable table.
+   * one matches zero rows and the store surfaces an {@link OptimisticLockException}. Before the CAS
+   * retry, {@code repairTableMetadata} rethrew it as a fatal error (HTTP 500) instead of tolerating
+   * the concurrent update. This test asserts that the lost race is benign and load returns a usable
+   * table.
    */
   @Test
   public void testLoadTableSurvivesConcurrentRepairVersionRace() throws Exception {
@@ -219,10 +220,10 @@ public void testLoadTableSurvivesConcurrentRepairVersionRace() throws Exception
     when(idGenerator.nextId()).thenReturn(10L, 11L);
 
     // First repair attempt loses the optimistic-lock CAS (a concurrent load already bumped the
-    // version): TableMetaService surfaces exactly this IOException. The retry re-reads the winner's
-    // already-repaired entity, against which the idempotent updater succeeds.
+    // version). The retry re-reads the winner's already-repaired entity, against which the
+    // idempotent updater succeeds.
     when(store.update(eq(ident), eq(TableEntity.class), eq(Entity.EntityType.TABLE), any()))
-        .thenThrow(new IOException("Failed to update the entity: " + ident))
+        .thenThrow(new OptimisticLockException("mock conflict"))
         .thenAnswer(
             invocation -> {
               @SuppressWarnings("unchecked")
@@ -255,6 +256,58 @@ public void testLoadTableSurvivesConcurrentRepairVersionRace() throws Exception
     Assertions.assertEquals("name", loadedTable.columns()[1].name());
   }
 
+  @Test
+  public void testRepairStopsAfterBoundedOptimisticLockRetries() throws Exception {
+    NameIdentifier ident = prepareDeclaredTableForRepair("repair-conflict-exhausted");
+    when(store.update(eq(ident), eq(TableEntity.class), eq(Entity.EntityType.TABLE), any()))
+        .thenThrow(new OptimisticLockException("mock conflict"));
+    // Remove the real sleep so this test checks the retry bound without becoming timing-sensitive.
+    Mockito.doNothing().when(lanceTableOps).backoffBeforeRetry(any());
+
+    OptimisticLockException failure =
+        Assertions.assertThrows(
+            OptimisticLockException.class,
+            () ->
+                PrincipalUtils.doAs(
+                    new UserPrincipal("tester"), () -> lanceTableOps.loadTable(ident)));
+    Assertions.assertInstanceOf(OptimisticLockException.class, failure.getCause());
+    verify(store, Mockito.times(5))
+        .update(eq(ident), eq(TableEntity.class), eq(Entity.EntityType.TABLE), any());
+  }
+
+  @Test
+  public void testRepairDoesNotRetryOrdinaryIoFailure() throws Exception {
+    NameIdentifier ident = prepareDeclaredTableForRepair("repair-io-failure");
+    when(store.update(eq(ident), eq(TableEntity.class), eq(Entity.EntityType.TABLE), any()))
+        .thenThrow(new IOException("database unavailable"));
+
+    RuntimeException failure =
+        Assertions.assertThrows(
+            RuntimeException.class,
+            () ->
+                PrincipalUtils.doAs(
+                    new UserPrincipal("tester"), () -> lanceTableOps.loadTable(ident)));
+    Assertions.assertInstanceOf(IOException.class, failure.getCause());
+    verify(store, Mockito.times(1))
+        .update(eq(ident), eq(TableEntity.class), eq(Entity.EntityType.TABLE), any());
+  }
+
+  @Test
+  public void testRepairBackoffPreservesThreadInterrupt() {
+    NameIdentifier ident = NameIdentifier.of("schema", "table");
+    Thread.currentThread().interrupt();
+    try {
+      IOException failure =
+          Assertions.assertThrows(IOException.class, () -> lanceTableOps.backoffBeforeRetry(ident));
+
+      Assertions.assertInstanceOf(InterruptedException.class, failure.getCause());
+      Assertions.assertTrue(Thread.currentThread().isInterrupted());
+    } finally {
+      // JUnit reuses worker threads, so do not leak this test's interrupt flag into another test.
+      Thread.interrupted();
+    }
+  }
+
   @Test
   public void testLoadTableWithStoredColumnsDoesNotReadLocation() throws Exception {
     NameIdentifier ident = NameIdentifier.of("schema", "table");
@@ -954,4 +1007,31 @@ private static TableEntity tableEntity(
             AuditInfo.builder().withCreator("creator").withCreateTime(Instant.EPOCH).build())
         .build();
   }
+
+  private NameIdentifier prepareDeclaredTableForRepair(String directoryName) throws Exception {
+    NameIdentifier ident = NameIdentifier.of("schema", "table");
+    String location = tempDir.resolve(directoryName).toString();
+    TableEntity tableEntity =
+        tableEntity(
+            ident,
+            List.of(),
+            Map.of(
+                Table.PROPERTY_LOCATION,
+                location,
+                LANCE_TABLE_DECLARED,
+                "true",
+                LANCE_STORAGE_OPTIONS_PREFIX + "endpoint",
+                "http://endpoint"));
+    when(store.get(eq(ident), eq(Entity.EntityType.TABLE), eq(TableEntity.class)))
+        .thenReturn(tableEntity);
+
+    Dataset dataset = mock(Dataset.class);
+    when(dataset.getSchema())
+        .thenReturn(new Schema(List.of(Field.nullable("id", new ArrowType.Int(32, true)))));
+    when(dataset.version()).thenReturn(8L);
+    Mockito.doReturn(dataset)
+        .when(lanceTableOps)
+        .openDataset(location, Map.of("endpoint", "http://endpoint"));
+    return ident;
+  }
 }

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

@@ -35,6 +35,7 @@
 import org.apache.gravitino.connector.PropertiesMetadata;
 import org.apache.gravitino.connector.capability.Capability;
 import org.apache.gravitino.exceptions.NoSuchEntityException;
+import org.apache.gravitino.exceptions.OptimisticLockException;
 import org.apache.gravitino.file.FilesetChange;
 import org.apache.gravitino.messaging.TopicChange;
 import org.apache.gravitino.rel.SupportsPartitions;
@@ -219,6 +220,9 @@ protected <R extends HasIdentifier> R operateOnEntity(
     R ret = null;
     try {
       ret = fn.apply(ident);
+    } catch (OptimisticLockException e) {
+      // A version conflict is actionable: the caller must retry against the latest entity.
+      throw e;
     } catch (NoSuchEntityException e) {
       // Case 2: The table is created by Gravitino, but has no corresponding entity in Gravitino.
       LOG.error(FormattedErrorMessages.ENTITY_NOT_FOUND, ident);

core/src/main/java/org/apache/gravitino/storage/relational/mapper/TableMetaMapper.java

@@ -73,6 +73,15 @@ TablePO selectTableByFullQualifiedName(
       @Param("schemaName") String schemaName,
       @Param("tableName") String tableName);
 
+  /**
+   * Selects and exclusively locks an active table metadata row.
+   *
+   * @param tableId the table ID
+   * @return the active table metadata, or {@code null} when it no longer exists
+   */
+  @SelectProvider(type = TableMetaSQLProviderFactory.class, method = "selectTableMetaByIdForUpdate")
+  TablePO selectTableMetaByIdForUpdate(@Param("tableId") Long tableId);
+
   @InsertProvider(type = TableMetaSQLProviderFactory.class, method = "insertTableMeta")
   void insertTableMeta(@Param("tableMeta") TablePO tablePO);
 
@@ -87,10 +96,18 @@ Integer updateTableMeta(
       @Param("oldTableMeta") TablePO oldTablePO,
       @Param("newSchemaId") Long newSchemaId);
 
+  /**
+   * Soft-deletes a table only if its version has not changed since the caller read it.
+   *
+   * @param tableId the table ID
+   * @param currentVersion the version observed by the caller
+   * @return the number of deleted rows; zero means the table changed or disappeared
+   */
   @UpdateProvider(
       type = TableMetaSQLProviderFactory.class,
       method = "softDeleteTableMetasByTableId")
-  Integer softDeleteTableMetasByTableId(@Param("tableId") Long tableId);
+  Integer softDeleteTableMetasByTableId(
+      @Param("tableId") Long tableId, @Param("currentVersion") Long currentVersion);
 
   @UpdateProvider(
       type = TableMetaSQLProviderFactory.class,

core/src/main/java/org/apache/gravitino/storage/relational/mapper/TableMetaSQLProviderFactory.java

@@ -89,6 +89,16 @@ public static String selectTableMetaById(@Param("tableId") Long tableId) {
     return getProvider().selectTableMetaById(tableId);
   }
 
+  /**
+   * Returns SQL that selects and exclusively locks an active table metadata row.
+   *
+   * @param tableId the table ID
+   * @return the locking select SQL
+   */
+  public static String selectTableMetaByIdForUpdate(@Param("tableId") Long tableId) {
+    return getProvider().selectTableMetaByIdForUpdate(tableId);
+  }
+
   public static String insertTableMeta(@Param("tableMeta") TablePO tablePO) {
     return getProvider().insertTableMeta(tablePO);
   }
@@ -104,8 +114,16 @@ public static String updateTableMeta(
     return getProvider().updateTableMeta(newTablePO, oldTablePO, newSchemaId);
   }
 
-  public static String softDeleteTableMetasByTableId(@Param("tableId") Long tableId) {
-    return getProvider().softDeleteTableMetasByTableId(tableId);
+  /**
+   * Returns SQL that soft-deletes a table with a version check.
+   *
+   * @param tableId the table ID
+   * @param currentVersion the version observed by the caller
+   * @return the version-checked delete SQL
+   */
+  public static String softDeleteTableMetasByTableId(
+      @Param("tableId") Long tableId, @Param("currentVersion") Long currentVersion) {
+    return getProvider().softDeleteTableMetasByTableId(tableId, currentVersion);
   }
 
   public static String softDeleteTableMetasByMetalakeId(@Param("metalakeId") Long metalakeId) {