diff --git a/catalogs/catalog-lakehouse-generic/src/main/java/org/apache/gravitino/catalog/lakehouse/lance/LanceTableOperations.java b/catalogs/catalog-lakehouse-generic/src/main/java/org/apache/gravitino/catalog/lakehouse/lance/LanceTableOperations.java index 38dbf0e90d2..7edfc468d53 100644 --- a/catalogs/catalog-lakehouse-generic/src/main/java/org/apache/gravitino/catalog/lakehouse/lance/LanceTableOperations.java +++ b/catalogs/catalog-lakehouse-generic/src/main/java/org/apache/gravitino/catalog/lakehouse/lance/LanceTableOperations.java @@ -31,6 +31,8 @@ import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.concurrent.ThreadLocalRandom; +import java.util.function.Function; import java.util.stream.Collectors; import java.util.stream.IntStream; import org.apache.arrow.vector.types.pojo.Field; @@ -63,6 +65,7 @@ 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; @@ -79,6 +82,18 @@ public class LanceTableOperations extends ManagedTableOperations { private static final Logger LOG = LoggerFactory.getLogger(LanceTableOperations.class); + /** + * Max attempts for the optimistic-locked repair-on-load {@code store.update}. Concurrent loads + * can repair the same table at once; the loser of the version CAS re-reads and retries. + */ + private static final int REPAIR_UPDATE_MAX_ATTEMPTS = 5; + + /** Lower bound (inclusive) of the randomized backoff slept between lost-CAS retries. */ + private static final long REPAIR_RETRY_MIN_BACKOFF_MS = 10; + + /** Upper bound (inclusive) of the randomized backoff slept between lost-CAS retries. */ + private static final long REPAIR_RETRY_MAX_BACKOFF_MS = 100; + public enum CreationMode { CREATE, EXIST_OK, @@ -527,10 +542,8 @@ private Column[] extractColumns(Schema arrowSchema) { private Table repairTableMetadata(NameIdentifier ident, Column[] columns, long datasetVersion) { try { TableEntity tableEntity = - store.update( + updateTableWithCasRetry( ident, - TableEntity.class, - Entity.EntityType.TABLE, current -> { if (!needsSchemaRefresh(current, datasetVersion)) { return current; @@ -547,13 +560,76 @@ 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. + */ + private TableEntity updateTableWithCasRetry( + NameIdentifier ident, Function updater) throws IOException { + IOException 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; + } + + lastConflict = e; + LOG.debug( + "Optimistic-lock conflict updating table {} metadata (attempt {}/{}), {}", + ident, + attempt, + REPAIR_UPDATE_MAX_ATTEMPTS, + attempt < REPAIR_UPDATE_MAX_ATTEMPTS ? "retrying" : "retries exhausted", + e); + + if (attempt < REPAIR_UPDATE_MAX_ATTEMPTS) { + backoffBeforeRetry(ident); + } + } + } + throw new IOException( + String.format( + "Failed to update table %s after %d optimistic-lock retries", + ident, REPAIR_UPDATE_MAX_ATTEMPTS), + lastConflict); + } + + /** + * Sleeps a short randomized backoff between two lost-CAS retries so that racing loads de-sync + * instead of immediately colliding again. Package-private so tests can neutralize the sleep. + * + * @param ident the table being updated, used only for the interrupt error message + * @throws IOException if the thread is interrupted while backing off + */ + void backoffBeforeRetry(NameIdentifier ident) throws IOException { + long backoffMillis = + ThreadLocalRandom.current() + .nextLong(REPAIR_RETRY_MIN_BACKOFF_MS, REPAIR_RETRY_MAX_BACKOFF_MS + 1); + try { + Thread.sleep(backoffMillis); + } catch (InterruptedException ie) { + Thread.currentThread().interrupt(); + throw new IOException("Interrupted while retrying metadata update for " + ident, ie); + } + } + private Table recordCheckedEmptyVersion(NameIdentifier ident, long datasetVersion) { try { TableEntity tableEntity = - store.update( + updateTableWithCasRetry( ident, - TableEntity.class, - Entity.EntityType.TABLE, current -> { if (!isDatasetVersionChanged(current.properties(), datasetVersion)) { return current; diff --git a/catalogs/catalog-lakehouse-generic/src/test/java/org/apache/gravitino/catalog/lakehouse/lance/TestLanceConcurrentRepairStress.java b/catalogs/catalog-lakehouse-generic/src/test/java/org/apache/gravitino/catalog/lakehouse/lance/TestLanceConcurrentRepairStress.java new file mode 100644 index 00000000000..4a67be530b5 --- /dev/null +++ b/catalogs/catalog-lakehouse-generic/src/test/java/org/apache/gravitino/catalog/lakehouse/lance/TestLanceConcurrentRepairStress.java @@ -0,0 +1,283 @@ +/* + * 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.lakehouse.lance; + +import static org.apache.gravitino.lance.common.utils.LanceConstants.LANCE_STORAGE_OPTIONS_PREFIX; +import static org.apache.gravitino.lance.common.utils.LanceConstants.LANCE_TABLE_DECLARED; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.when; + +import java.io.IOException; +import java.time.Instant; +import java.util.List; +import java.util.Map; +import java.util.concurrent.Callable; +import java.util.concurrent.CyclicBarrier; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Function; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.Schema; +import org.apache.commons.lang3.tuple.Pair; +import org.apache.gravitino.Config; +import org.apache.gravitino.Entity; +import org.apache.gravitino.Entity.EntityType; +import org.apache.gravitino.EntityStore; +import org.apache.gravitino.HasIdentifier; +import org.apache.gravitino.NameIdentifier; +import org.apache.gravitino.UserPrincipal; +import org.apache.gravitino.catalog.ManagedSchemaOperations; +import org.apache.gravitino.meta.AuditInfo; +import org.apache.gravitino.meta.TableEntity; +import org.apache.gravitino.rel.Table; +import org.apache.gravitino.storage.IdGenerator; +import org.apache.gravitino.utils.Executable; +import org.apache.gravitino.utils.PrincipalUtils; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; +import org.junit.jupiter.api.io.TempDir; +import org.lance.Dataset; +import org.mockito.Mockito; + +/** + * 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 + * 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}). + * + *

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. + */ +public class TestLanceConcurrentRepairStress { + + @TempDir private java.nio.file.Path tempDir; + + private static final NameIdentifier IDENT = NameIdentifier.of("schema", "table"); + + @Test + @Timeout(60) + public void testConcurrentRepairLoadsSurviveCasContention() throws Exception { + // The reported race is "two loads repair the same table at once". Also drive a small herd to + // confirm the bounded CAS retry stays robust beyond the minimal two-load case. Iterations are + // kept modest (and the retry backoff is neutralized in runStress) so this stays a fast, + // non-flaky unit test rather than a soak test — the pre-fix bug already fails on iteration 0. + runStress(2, 500); + runStress(8, 300); + } + + private void runStress(int concurrency, int iterations) throws Exception { + String location = tempDir.resolve("cas-race-" + concurrency).toString(); + Map storageOptions = Map.of("endpoint", "http://endpoint"); + Schema datasetSchema = + new Schema( + List.of( + Field.nullable("id", new ArrowType.Int(32, true)), + Field.nullable("name", new ArrowType.Utf8()))); + + ManagedSchemaOperations schemaOps = mock(ManagedSchemaOperations.class); + IdGenerator idGenerator = mock(IdGenerator.class); + AtomicLong idSeq = new AtomicLong(1000); + when(idGenerator.nextId()).thenAnswer(invocation -> idSeq.incrementAndGet()); + + ExecutorService pool = Executors.newFixedThreadPool(concurrency); + UserPrincipal user = new UserPrincipal("tester"); + int failures = 0; + try { + for (int iter = 0; iter < iterations; iter++) { + // Fresh stale (declared, empty-schema) entity every iteration so each round starts from the + // pre-repair state that triggers repairTableMetadata on every concurrent load. + TableEntity stale = + tableEntity( + List.of(), + Map.of( + Table.PROPERTY_LOCATION, + location, + LANCE_TABLE_DECLARED, + "true", + LANCE_STORAGE_OPTIONS_PREFIX + "endpoint", + "http://endpoint")); + CasEntityStore store = new CasEntityStore(stale); + LanceTableOperations ops = spy(new LanceTableOperations(store, schemaOps, idGenerator)); + // Neutralize the inter-retry backoff: this test validates the CAS re-read/retry recovery, + // not the sleep duration, and real sleeps would make it slow and timing-flaky. + Mockito.doNothing().when(ops).backoffBeforeRetry(Mockito.any()); + Dataset dataset = mock(Dataset.class); + when(dataset.version()).thenReturn(8L); + when(dataset.getSchema()).thenReturn(datasetSchema); + Mockito.doReturn(dataset).when(ops).openDataset(location, storageOptions); + + CyclicBarrier barrier = new CyclicBarrier(concurrency); + List> futures = new java.util.ArrayList<>(concurrency); + for (int t = 0; t < concurrency; t++) { + Callable task = + () -> { + barrier.await(); + return PrincipalUtils.doAs(user, () -> ops.loadTable(IDENT)); + }; + futures.add(pool.submit(task)); + } + + for (Future
future : futures) { + try { + Table loaded = future.get(30, TimeUnit.SECONDS); + Assertions.assertEquals( + 2, loaded.columns().length, "repaired table must expose both columns"); + } catch (ExecutionException e) { + failures++; + // Surface the first real failure with its cause for debugging. + if (failures == 1) { + throw new AssertionError( + "Concurrent repair-on-load failed at iteration " + + iter + + " with concurrency=" + + concurrency, + e.getCause()); + } + } + } + } + } finally { + pool.shutdownNow(); + } + Assertions.assertEquals( + 0, failures, "no concurrent load should fail (concurrency=" + concurrency + ")"); + } + + private static TableEntity tableEntity( + List columns, Map properties) { + return TableEntity.builder() + .withId(1L) + .withName(IDENT.name()) + .withNamespace(IDENT.namespace()) + .withComment("comment") + .withColumns(columns) + .withProperties(properties) + .withAuditInfo( + AuditInfo.builder().withCreator("creator").withCreateTime(Instant.EPOCH).build()) + .build(); + } + + /** + * 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. + */ + private static final class CasEntityStore implements EntityStore { + + private static final class Versioned { + private final long version; + private final TableEntity entity; + + Versioned(long version, TableEntity entity) { + this.version = version; + this.entity = entity; + } + } + + private final AtomicReference ref; + + CasEntityStore(TableEntity initial) { + this.ref = new AtomicReference<>(new Versioned(0L, initial)); + } + + @Override + @SuppressWarnings("unchecked") + public E get( + NameIdentifier ident, EntityType entityType, Class type) { + return (E) ref.get().entity; + } + + @Override + @SuppressWarnings("unchecked") + public E update( + NameIdentifier ident, Class type, EntityType entityType, Function updater) + throws IOException { + Versioned base = ref.get(); + E updated = updater.apply((E) base.entity); + Versioned next = new Versioned(base.version + 1, (TableEntity) updated); + if (ref.compareAndSet(base, next)) { + return updated; + } + throw new IOException("Failed to update the entity: " + ident); + } + + // --- unused surface --------------------------------------------------- + + @Override + public void initialize(Config config) { + throw new UnsupportedOperationException(); + } + + @Override + public boolean exists(NameIdentifier ident, EntityType entityType) { + throw new UnsupportedOperationException(); + } + + @Override + public void put(E e, boolean overwritten) { + throw new UnsupportedOperationException(); + } + + @Override + public List batchGet( + List idents, EntityType entityType, Class clazz) { + throw new UnsupportedOperationException(); + } + + @Override + public boolean delete(NameIdentifier ident, EntityType entityType, boolean cascade) { + throw new UnsupportedOperationException(); + } + + @Override + public int batchDelete( + List> entitiesToDelete, boolean cascade) { + throw new UnsupportedOperationException(); + } + + @Override + public void batchPut(List entities, boolean overwritten) { + throw new UnsupportedOperationException(); + } + + @Override + public R executeInTransaction(Executable executable) { + throw new UnsupportedOperationException(); + } + + @Override + public void close() {} + } +} diff --git a/catalogs/catalog-lakehouse-generic/src/test/java/org/apache/gravitino/catalog/lakehouse/lance/TestLanceTableOperations.java b/catalogs/catalog-lakehouse-generic/src/test/java/org/apache/gravitino/catalog/lakehouse/lance/TestLanceTableOperations.java index 59b2a41ff9e..bf99a64814b 100644 --- a/catalogs/catalog-lakehouse-generic/src/test/java/org/apache/gravitino/catalog/lakehouse/lance/TestLanceTableOperations.java +++ b/catalogs/catalog-lakehouse-generic/src/test/java/org/apache/gravitino/catalog/lakehouse/lance/TestLanceTableOperations.java @@ -34,6 +34,7 @@ import static org.mockito.Mockito.when; import com.google.common.collect.Maps; +import java.io.IOException; import java.time.Instant; import java.util.List; import java.util.Map; @@ -170,6 +171,90 @@ public void testLoadDeclaredTableSchemaFromLocation() throws Exception { Assertions.assertFalse(loadedTable.properties().containsKey(LANCE_TABLE_DECLARED)); } + /** + * 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. + */ + @Test + public void testLoadTableSurvivesConcurrentRepairVersionRace() throws Exception { + NameIdentifier ident = NameIdentifier.of("schema", "table"); + String location = tempDir.resolve("concurrent-repair-table").toString(); + TableEntity tableEntity = + tableEntity( + ident, + List.of(), + Map.of( + Table.PROPERTY_LOCATION, + location, + LANCE_TABLE_DECLARED, + "true", + LANCE_STORAGE_OPTIONS_PREFIX + "endpoint", + "http://endpoint")); + // The winner of the race already repaired the table to the dataset schema and version. + TableEntity alreadyRepairedTableEntity = + tableEntity( + ident, + List.of( + ColumnEntity.builder() + .withId(11L) + .withName("id") + .withDataType(Types.IntegerType.get()) + .withPosition(0) + .withAuditInfo(AuditInfo.EMPTY) + .build(), + ColumnEntity.builder() + .withId(12L) + .withName("name") + .withDataType(Types.StringType.get()) + .withPosition(1) + .withAuditInfo(AuditInfo.EMPTY) + .build()), + Map.of(Table.PROPERTY_LOCATION, location, LANCE_TABLE_VERSION, "8")); + when(store.get(eq(ident), eq(Entity.EntityType.TABLE), eq(TableEntity.class))) + .thenReturn(tableEntity); + 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. + when(store.update(eq(ident), eq(TableEntity.class), eq(Entity.EntityType.TABLE), any())) + .thenThrow(new IOException("Failed to update the entity: " + ident)) + .thenAnswer( + invocation -> { + @SuppressWarnings("unchecked") + Function updater = invocation.getArgument(3); + return updater.apply(alreadyRepairedTableEntity); + }); + + Dataset dataset = mock(Dataset.class); + when(dataset.getSchema()) + .thenReturn( + new Schema( + List.of( + Field.nullable("id", new ArrowType.Int(32, true)), + Field.nullable("name", new ArrowType.Utf8())))); + when(dataset.version()).thenReturn(8L); + Mockito.doReturn(dataset) + .when(lanceTableOps) + .openDataset(location, Map.of("endpoint", "http://endpoint")); + + // A lost repair race must not fail the load: the bounded CAS retry recovers and returns the + // repaired table instead of surfacing the first conflict as RuntimeException("Failed to repair + // table"). + Table loadedTable = + Assertions.assertDoesNotThrow( + () -> + PrincipalUtils.doAs( + new UserPrincipal("tester"), () -> lanceTableOps.loadTable(ident))); + Assertions.assertEquals(2, loadedTable.columns().length); + Assertions.assertEquals("id", loadedTable.columns()[0].name()); + Assertions.assertEquals("name", loadedTable.columns()[1].name()); + } + @Test public void testLoadTableWithStoredColumnsDoesNotReadLocation() throws Exception { NameIdentifier ident = NameIdentifier.of("schema", "table"); diff --git a/core/src/main/java/org/apache/gravitino/storage/relational/service/TableMetaService.java b/core/src/main/java/org/apache/gravitino/storage/relational/service/TableMetaService.java index 72391bc31be..17ee0d5c3b2 100644 --- a/core/src/main/java/org/apache/gravitino/storage/relational/service/TableMetaService.java +++ b/core/src/main/java/org/apache/gravitino/storage/relational/service/TableMetaService.java @@ -56,6 +56,16 @@ /** The service class for table metadata. It provides the basic database operations for table. */ public class TableMetaService { + + /** + * Message prefix of the {@link java.io.IOException} thrown by {@link #updateTable} when the + * optimistic-lock CAS matches zero rows (the stored version advanced under a concurrent update). + * Exposed so callers that retry the lost race (e.g. the Lance repair-on-load path) can recognize + * the conflict without re-declaring the literal. + */ + public static final String UPDATE_ENTITY_CONFLICT_MESSAGE_PREFIX = + "Failed to update the entity: "; + private static final TableMetaService INSTANCE = new TableMetaService(); private BasePOStorageOps ops; @@ -236,7 +246,7 @@ public TableEntity updateTable( if (updateResult.get() > 0) { return newTableEntity; } else { - throw new IOException("Failed to update the entity: " + identifier); + throw new IOException(UPDATE_ENTITY_CONFLICT_MESSAGE_PREFIX + identifier); } }