diff --git a/hudi-sync/hudi-hive-sync/src/main/java/org/apache/hudi/hive/HiveSyncConfigHolder.java b/hudi-sync/hudi-hive-sync/src/main/java/org/apache/hudi/hive/HiveSyncConfigHolder.java index 14afa81ea30f7..84e2f7bae0e8d 100644 --- a/hudi-sync/hudi-hive-sync/src/main/java/org/apache/hudi/hive/HiveSyncConfigHolder.java +++ b/hudi-sync/hudi-hive-sync/src/main/java/org/apache/hudi/hive/HiveSyncConfigHolder.java @@ -132,18 +132,24 @@ public class HiveSyncConfigHolder { + "Hive Driver workers, with ADD and TOUCH additionally split into batches of " + "`hoodie.datasource.hive_sync.batch_num` partitions per statement (ADD was already batched " + "before this flag existed; only its dispatch becomes parallel here). SET_LOCATION remains one " - + "statement per partition, as Hive SQL has no multi-partition form. DROP remains serial. " - + "Table-level statements (create/alter table, last commit time, writer version) continue to run " - + "on the single session Driver. Default off; the default HiveQL path is unchanged unless " + + "statement per partition, as Hive SQL has no multi-partition form. DROP is also parallelized, " + + "but over a pool of metastore (Thrift) clients rather than Hive Driver workers, since it is " + + "issued as dropPartition calls rather than SQL; drops are split into batches of " + + "`hoodie.datasource.hive_sync.batch_num` partitions and fanned across those clients. DROP falls " + + "back to sequential execution on the single session client when " + + "`hoodie.datasource.hive_sync.use_spark_catalog` is true, as the Spark catalog client cannot be " + + "pooled. Table-level statements (create/alter table, last commit time, writer version) continue " + + "to run on the single session Driver. Default off; the default HiveQL path is unchanged unless " + "explicitly opted in."); public static final ConfigProperty HIVE_SYNC_BATCHING_THREADS = ConfigProperty .key("hoodie.datasource.hive_sync.batching.threads") .defaultValue(4) .markAdvanced() .sinceVersion("1.3.0") - .withDocumentation("Pool size (number of Hive Driver workers) and worker-thread count for parallel " - + "HiveQL partition dispatch when `hoodie.datasource.hive_sync.batching.enabled` is true. " - + "Ignored otherwise."); + .withDocumentation("Number of worker threads used for parallel HiveQL partition dispatch when " + + "`hoodie.datasource.hive_sync.batching.enabled` is true. The same value sizes both pools: the " + + "Hive Driver workers used for ADD/TOUCH/SET_LOCATION and the metastore (Thrift) clients used " + + "for DROP. Ignored otherwise."); public static final ConfigProperty HIVE_SYNC_MODE = ConfigProperty .key("hoodie.datasource.hive_sync.mode") .noDefaultValue() diff --git a/hudi-sync/hudi-hive-sync/src/main/java/org/apache/hudi/hive/HoodieHiveSyncClient.java b/hudi-sync/hudi-hive-sync/src/main/java/org/apache/hudi/hive/HoodieHiveSyncClient.java index accc60cb90664..fe7d5be219b47 100644 --- a/hudi-sync/hudi-hive-sync/src/main/java/org/apache/hudi/hive/HoodieHiveSyncClient.java +++ b/hudi-sync/hudi-hive-sync/src/main/java/org/apache/hudi/hive/HoodieHiveSyncClient.java @@ -38,6 +38,7 @@ import org.apache.hudi.hive.ddl.JDBCBasedMetadataOperator; import org.apache.hudi.hive.ddl.JDBCExecutor; import org.apache.hudi.hive.util.HiveDriverPool; +import org.apache.hudi.hive.util.IMetaStoreClientPool; import org.apache.hudi.hive.util.IMetaStoreClientUtil; import org.apache.hudi.hive.util.PartitionFilterGenerator; import org.apache.hudi.sync.common.HoodieSyncClient; @@ -89,6 +90,11 @@ public class HoodieHiveSyncClient extends HoodieSyncClient { private final Map initialTableByName = new HashMap<>(); DDLExecutor ddlExecutor; private IMetaStoreClient client; + // Present only when HIVE_SYNC_BATCHING_ENABLED and sync mode is HIVEQL. Owned by + // this class; closed in close() before Hive.closeCurrent(). HiveQueryDDLExecutor + // uses it only for DROP (Hive Thrift, not Hive Driver) — see IMetaStoreClientPool + // javadoc. + private Option partitionClientPool = Option.empty(); // Present only when HIVE_SYNC_BATCHING_ENABLED and sync mode is HIVEQL (explicit // or legacy default). Owned by HiveQueryDDLExecutor; this field is kept for // reference only — close() is delegated through ddlExecutor.close(). @@ -131,8 +137,7 @@ public HoodieHiveSyncClient(HiveSyncConfig config, HoodieTableMetaClient metaCli ddlExecutor = new HMSDDLExecutor(config, this.client); break; case HIVEQL: - this.partitionDriverPool = maybeBuildHiveDriverPool(config); - ddlExecutor = new HiveQueryDDLExecutor(config, this.client, this.partitionDriverPool); + ddlExecutor = buildHiveQueryDDLExecutor(config); break; case JDBC: JDBCExecutor jdbcExecutor = new JDBCExecutor(config); @@ -150,23 +155,23 @@ public HoodieHiveSyncClient(HiveSyncConfig config, HoodieTableMetaClient metaCli jdbcMetadataOperator = new JDBCBasedMetadataOperator( jdbcExecutor.getConnection(), databaseName); } else { - this.partitionDriverPool = maybeBuildHiveDriverPool(config); - ddlExecutor = new HiveQueryDDLExecutor(config, this.client, this.partitionDriverPool); + ddlExecutor = buildHiveQueryDDLExecutor(config); } } } catch (Exception e) { - // The pool owns live daemon threads and Hive Drivers, and is built before the - // executor that would otherwise own its lifecycle. Any throw between those two - // points would leak it -- notably QueryBasedDDLExecutor's super(config), which - // runs the PartitionValueExtractor reflection before HiveQueryDDLExecutor's own - // try block is even entered. Closing here covers every such window; the pool's - // close() is idempotent, so overlapping with the executor's cleanup is harmless. - closePartitionDriverPoolQuietly(); + // The pools own live daemon threads, Hive Drivers, and Thrift sockets, and are + // built before the executor that would otherwise own their lifecycle. Any throw + // between those two points would leak them -- notably QueryBasedDDLExecutor's + // super(config), which runs the PartitionValueExtractor reflection before + // HiveQueryDDLExecutor's own try block is even entered. Closing here covers every + // such window; both close() methods are idempotent, so overlapping with the + // executor's cleanup (or buildHiveQueryDDLExecutor's rollback) is harmless. + closePartitionPoolsQuietly(); throw new HoodieHiveSyncException("Failed to create HiveMetaStoreClient", e); } } - private void closePartitionDriverPoolQuietly() { + private void closePartitionPoolsQuietly() { partitionDriverPool.ifPresent(pool -> { try { pool.close(); @@ -174,6 +179,15 @@ private void closePartitionDriverPoolQuietly() { log.warn("Error closing HiveDriverPool during failed sync client construction", e); } }); + partitionDriverPool = Option.empty(); + partitionClientPool.ifPresent(pool -> { + try { + pool.close(); + } catch (Exception e) { + log.warn("Error closing IMetaStoreClient pool during failed sync client construction", e); + } + }); + partitionClientPool = Option.empty(); } /** @@ -227,6 +241,22 @@ private IMetaStoreClient createMetaStoreClient(HiveSyncConfig config) { } } + private Option maybeBuildPartitionClientPool(HiveSyncConfig config) { + if (!config.getBooleanOrDefault(HIVE_SYNC_BATCHING_ENABLED)) { + return Option.empty(); + } + if (config.getBooleanOrDefault(HIVE_SYNC_USE_SPARK_CATALOG)) { + // The Spark catalog client is constructed via reflection against a Spark-side + // class and isn't compatible with the direct RetryingMetaStoreClient pool path. + // Fall back to single-client sequential behavior rather than failing the sync. + log.warn("{}=true is not supported with {}=true; falling back to sequential partition drop.", + HIVE_SYNC_BATCHING_ENABLED.key(), HIVE_SYNC_USE_SPARK_CATALOG.key()); + return Option.empty(); + } + int size = config.getIntOrDefault(HIVE_SYNC_BATCHING_THREADS); + return Option.of(new IMetaStoreClientPool(config, size)); + } + private Option maybeBuildHiveDriverPool(HiveSyncConfig config) { if (!config.getBooleanOrDefault(HIVE_SYNC_BATCHING_ENABLED)) { return Option.empty(); @@ -235,6 +265,25 @@ private Option maybeBuildHiveDriverPool(HiveSyncConfig config) { return Option.of(new HiveDriverPool(config, size)); } + /** + * Builds the (optional) partition-phase pools and the {@link HiveQueryDDLExecutor} + * that uses them, rolling back whichever pool(s) already got built if a later step + * in this sequence throws. Without this, a failure in {@code maybeBuildPartitionClientPool} + * (after {@code partitionDriverPool} was already built) or in the executor's own + * constructor would leak the already-built pool's worker threads and Thrift/Driver + * connections, since this constructor's outer catch just rethrows. + */ + private HiveQueryDDLExecutor buildHiveQueryDDLExecutor(HiveSyncConfig config) { + try { + this.partitionDriverPool = maybeBuildHiveDriverPool(config); + this.partitionClientPool = maybeBuildPartitionClientPool(config); + return new HiveQueryDDLExecutor(config, this.client, this.partitionDriverPool, this.partitionClientPool); + } catch (Exception e) { + closePartitionPoolsQuietly(); + throw e; + } + } + private Table getInitialTable(String table) { return initialTableByName.computeIfAbsent(table, t -> { try { @@ -630,22 +679,38 @@ public void deleteLastReplicatedTimeStamp(String tableName) { @Override public void close() { try { - ddlExecutor.close(); - if (client != null) { - // Close the proxied IMetaStoreClient directly before Hive.closeCurrent(). - // When RetryingMetaStoreClient rebuilds the underlying client on a transient - // TException, the fresh MSC is reachable only through this proxy, while the - // thread-local Hive singleton still references the older instance. So - // Hive.closeCurrent() alone closes the stale MSC and orphans the retry-created - // one, leaking a connection per sync cycle. client.close() releases the live - // MSC by identity; Hive.closeCurrent() remains a fallback for the singleton path. - try { - client.close(); - } catch (Exception e) { - log.warn("Failed to close IMetaStoreClient directly; Hive.closeCurrent() will run anyway", e); + try { + ddlExecutor.close(); + } finally { + // Close the partition client pool before Hive.closeCurrent() so the + // RetryingMetaStoreClient instances held by the pool release their Thrift + // sockets without racing the ThreadLocal Hive cleanup. Runs even if + // ddlExecutor.close() above threw, so the pool's Thrift sockets and + // worker threads aren't leaked. + if (partitionClientPool.isPresent()) { + try { + partitionClientPool.get().close(); + } catch (Exception e) { + log.warn("Error closing IMetaStoreClient pool", e); + } + partitionClientPool = Option.empty(); + } + if (client != null) { + // Close the proxied IMetaStoreClient directly before Hive.closeCurrent(). + // When RetryingMetaStoreClient rebuilds the underlying client on a transient + // TException, the fresh MSC is reachable only through this proxy, while the + // thread-local Hive singleton still references the older instance. So + // Hive.closeCurrent() alone closes the stale MSC and orphans the retry-created + // one, leaking a connection per sync cycle. client.close() releases the live + // MSC by identity; Hive.closeCurrent() remains a fallback for the singleton path. + try { + client.close(); + } catch (Exception e) { + log.warn("Failed to close IMetaStoreClient directly; Hive.closeCurrent() will run anyway", e); + } + Hive.closeCurrent(); + client = null; } - Hive.closeCurrent(); - client = null; } } catch (Exception e) { log.error("Could not close connection ", e); diff --git a/hudi-sync/hudi-hive-sync/src/main/java/org/apache/hudi/hive/ddl/HiveQueryDDLExecutor.java b/hudi-sync/hudi-hive-sync/src/main/java/org/apache/hudi/hive/ddl/HiveQueryDDLExecutor.java index c853313182ed1..70ff8f5e760f3 100644 --- a/hudi-sync/hudi-hive-sync/src/main/java/org/apache/hudi/hive/ddl/HiveQueryDDLExecutor.java +++ b/hudi-sync/hudi-hive-sync/src/main/java/org/apache/hudi/hive/ddl/HiveQueryDDLExecutor.java @@ -18,12 +18,14 @@ package org.apache.hudi.hive.ddl; +import org.apache.hudi.common.util.CollectionUtils; import org.apache.hudi.common.util.HoodieTimer; import org.apache.hudi.common.util.Option; import org.apache.hudi.hive.HiveSyncConfig; import org.apache.hudi.hive.HoodieHiveSyncException; import org.apache.hudi.hive.util.HiveDriverPool; import org.apache.hudi.hive.util.HivePartitionUtil; +import org.apache.hudi.hive.util.IMetaStoreClientPool; import lombok.extern.slf4j.Slf4j; import org.apache.hadoop.hive.metastore.IMetaStoreClient; @@ -59,16 +61,22 @@ public class HiveQueryDDLExecutor extends QueryBasedDDLExecutor { // (createTable, schema evolution, single-statement runSQL callers) always uses the // session `hiveDriver` above. See HiveDriverPool javadoc. private final Option driverPool; + // When present, dropPartitionsToTable fans batches across this Thrift client pool. + // Owned by HoodieHiveSyncClient; close() is delegated through there. See + // IMetaStoreClientPool javadoc for the usage contract (partition-row ops only). + private final Option metaStoreClientPool; public HiveQueryDDLExecutor(HiveSyncConfig config, IMetaStoreClient metaStoreClient) { - this(config, metaStoreClient, Option.empty()); + this(config, metaStoreClient, Option.empty(), Option.empty()); } public HiveQueryDDLExecutor(HiveSyncConfig config, IMetaStoreClient metaStoreClient, - Option driverPool) { + Option driverPool, + Option metaStoreClientPool) { super(config); this.metaStoreClient = metaStoreClient; this.driverPool = driverPool; + this.metaStoreClientPool = metaStoreClientPool; try { this.sessionState = new SessionState(config.getHiveConf(), UserGroupInformation.getCurrentUser().getShortUserName()); @@ -209,21 +217,57 @@ public void dropPartitionsToTable(String tableName, List partitionsToDro log.info("Drop partitions {} on {}", partitionsToDrop.size(), tableName); try { - for (String dropPartition : partitionsToDrop) { - if (HivePartitionUtil.partitionExists(metaStoreClient, tableName, dropPartition, partitionValueExtractor, - config)) { - String partitionClause = - HivePartitionUtil.getPartitionClauseForDrop(dropPartition, partitionValueExtractor, config); - metaStoreClient.dropPartition(databaseName, tableName, partitionClause, false); - } - log.info("Drop partition {} on {}", dropPartition, tableName); - } + int batchSyncPartitionNum = config.getIntOrDefault(HIVE_BATCH_SYNC_PARTITION_NUM); + List> batches = CollectionUtils.batches(partitionsToDrop, batchSyncPartitionNum); + runDropBatches(tableName, batches); } catch (Exception e) { log.error("{} drop partition failed", tableId(databaseName, tableName), e); throw new HoodieHiveSyncException(tableId(databaseName, tableName) + " drop partition failed", e); } } + /** + * Drops partitions one batch at a time. When {@link #metaStoreClientPool} is present, + * batches fan out across the pool's worker threads (each borrowing an independent + * IMetaStoreClient); otherwise batches are dispatched sequentially against the + * session client. Hive has no batch-drop primitive that matches dropPartition's + * semantics, so each worker still iterates its chunk one partition at a time — the + * win is fanning chunks across independent Thrift clients. + * + *

First-error semantics come from {@link ParallelDispatch}, shared with + * {@code HiveDriverPool}: the first failure is rethrown, batches that have not started + * are stopped via the task-side abort flag, and later failures are logged at WARN. + */ + private void runDropBatches(String tableName, List> batches) throws Exception { + if (!metaStoreClientPool.isPresent()) { + for (List batch : batches) { + applyDropBatch(metaStoreClient, tableName, batch); + } + return; + } + IMetaStoreClientPool pool = metaStoreClientPool.get(); + pool.awaitAll( + pool.dispatchAll(batches, (client, batch) -> applyDropBatch(client, tableName, batch)), + "drop partition"); + } + + private void applyDropBatch(IMetaStoreClient client, String tableName, List batch) throws Exception { + int dropped = 0; + for (String dropPartition : batch) { + if (HivePartitionUtil.partitionExists(client, tableName, dropPartition, + partitionValueExtractor, config)) { + String partitionClause = + HivePartitionUtil.getPartitionClauseForDrop(dropPartition, partitionValueExtractor, config); + client.dropPartition(databaseName, tableName, partitionClause, false); + dropped++; + } + // Per-partition detail stays at debug: a batch can hold thousands of partitions + // and N workers log concurrently, so INFO carries the per-batch summary instead. + log.debug("Dropped partition {} on {}", dropPartition, tableName); + } + log.info("Dropped {} of {} partitions in batch on {}", dropped, batch.size(), tableName); + } + @Override public void close() { // Close the pool first so the worker threads stop dispatching against their diff --git a/hudi-sync/hudi-hive-sync/src/main/java/org/apache/hudi/hive/util/HiveDriverPool.java b/hudi-sync/hudi-hive-sync/src/main/java/org/apache/hudi/hive/util/HiveDriverPool.java index 5950bf442313d..b9f4db106ac1f 100644 --- a/hudi-sync/hudi-hive-sync/src/main/java/org/apache/hudi/hive/util/HiveDriverPool.java +++ b/hudi-sync/hudi-hive-sync/src/main/java/org/apache/hudi/hive/util/HiveDriverPool.java @@ -18,7 +18,6 @@ package org.apache.hudi.hive.util; -import org.apache.hudi.common.util.VisibleForTesting; import org.apache.hudi.exception.HoodieException; import org.apache.hudi.hive.HiveSyncConfig; import org.apache.hudi.hive.HoodieHiveSyncException; @@ -32,15 +31,10 @@ import java.util.ArrayList; import java.util.List; -import java.util.concurrent.CancellationException; -import java.util.concurrent.CountDownLatch; -import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; -import java.util.concurrent.Future; import java.util.concurrent.ThreadFactory; import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import static org.apache.hudi.sync.common.HoodieSyncConfig.META_SYNC_DATABASE_NAME; @@ -124,24 +118,14 @@ public void runOnEachWorker(List setupSqls) { if (setupSqls.isEmpty()) { return; } - Dispatch dispatch = new Dispatch(workers.size()); + ParallelDispatch dispatch = new ParallelDispatch(workers.size()); for (Worker worker : workers) { - dispatch.add(worker.executor.submit(() -> { - if (dispatch.aborted()) { - throw new CancellationException("Skipped after an earlier setup statement failed"); - } - try { - for (String sql : setupSqls) { - worker.driver.run(sql); - } - } catch (Throwable t) { - dispatch.abort(); - throw t; - } finally { - dispatch.taskSettled(); + dispatch.add(worker.executor.submit(dispatch.guard(() -> { + for (String sql : setupSqls) { + worker.driver.run(sql); } return null; - })); + }, "Skipped after an earlier setup statement failed"))); } dispatch.sealed(); awaitAll(dispatch); @@ -150,42 +134,28 @@ public void runOnEachWorker(List setupSqls) { /** * Dispatches each SQL string to a worker (round-robin) and returns a handle to the * in-flight batch — this method does not block. The caller is responsible for - * awaiting completion via {@link #awaitAll(Dispatch)} and collecting errors. SQL text + * awaiting completion via {@link #awaitAll(ParallelDispatch)} and collecting errors. SQL text * is intentionally not logged per-statement here: batched TOUCH/ADD statements can * be many kilobytes, and N parallel workers would multiply the log volume. See - * {@link #awaitAll(Dispatch)} for the per-call summary log. + * {@link #awaitAll(ParallelDispatch)} for the per-call summary log. * *

Statements are spread round-robin across workers, so worker w owns * indices {@code w, w + N, w + 2N, ...}. Each worker drains its own queue * independently, which is why abort has to be observed by the tasks themselves - * rather than by the awaiting thread — see {@link Dispatch}. + * rather than by the awaiting thread — see {@link ParallelDispatch}. */ - public Dispatch dispatchAll(List sqls) { + public ParallelDispatch dispatchAll(List sqls) { if (closed) { throw new IllegalStateException("Cannot dispatch to a closed HiveDriverPool"); } - Dispatch dispatch = new Dispatch(sqls.size()); + ParallelDispatch dispatch = new ParallelDispatch(sqls.size()); for (int i = 0; i < sqls.size(); i++) { String sql = sqls.get(i); Worker worker = workers.get(i % workers.size()); - dispatch.add(worker.executor.submit(() -> { - // Abort check inside the task: a worker can dequeue this statement while the - // awaiting thread is still parked on some other worker's slower statement, so - // Future.cancel() alone cannot stop it in time. Checking here means no - // statement starts after a sibling has already failed. - if (dispatch.aborted()) { - throw new CancellationException("Skipped after an earlier statement failed"); - } - try { - worker.driver.run(sql); - } catch (Throwable t) { - dispatch.abort(); - throw t; - } finally { - dispatch.taskSettled(); - } + dispatch.add(worker.executor.submit(dispatch.guard(() -> { + worker.driver.run(sql); return null; - })); + }, "Skipped after an earlier statement failed"))); } dispatch.sealed(); return dispatch; @@ -200,144 +170,17 @@ public Dispatch dispatchAll(List sqls) { * Callers do not need per-statement results (Hive's Driver.run side-effects the * metastore), so this method is void. */ - public void awaitAll(Dispatch dispatch) { + public void awaitAll(ParallelDispatch dispatch) { long start = System.currentTimeMillis(); - // Block until either every task settled or one of them aborted the batch. Only - // then walk the futures — by that point no un-started task can still begin, so - // the walk order no longer affects how much extra DDL gets applied. - dispatch.awaitSettledOrAborted(); - int cancelled = dispatch.cancelPending(); + ParallelDispatch.Outcome outcome = dispatch.awaitOutcome(); + outcome.suppressed().forEach(e -> + LOG.warn("Additional SQL batch failed (suppressed in favor of first error)", e)); - Exception firstError = null; - int completed = 0; - for (Future f : dispatch.futures()) { - try { - f.get(); - completed++; - } catch (CancellationException ce) { - // Either we cancelled it before it started, or the task itself observed the - // abort flag and bailed. Not a new failure; just note it for the summary. - cancelled++; - } catch (InterruptedException ie) { - Thread.currentThread().interrupt(); - if (firstError == null) { - firstError = ie; - } - } catch (ExecutionException ee) { - Exception cause = unwrap(ee); - if (cause instanceof CancellationException) { - cancelled++; - } else if (firstError == null) { - firstError = cause; - } else { - LOG.warn("Additional SQL batch failed (suppressed in favor of first error)", cause); - } - } - } - if (firstError != null) { - throw new HoodieHiveSyncException("Failed in executing SQL", firstError); + if (outcome.failed()) { + throw new HoodieHiveSyncException("Failed in executing SQL", outcome.firstError()); } LOG.info("Completed {} SQL statements ({} cancelled) in {} ms across {} workers", - completed, cancelled, System.currentTimeMillis() - start, size); - } - - /** - * Handle to one {@link #dispatchAll(List)} batch: the submitted futures plus the - * shared abort flag the tasks consult before running. - * - *

The flag exists because the futures belong to N independent single-thread - * executors. Cancelling from the awaiting thread is inherently late — a worker can - * pull its next statement off its own queue at any moment — so each task also - * re-checks {@link #aborted()} on entry. That is what actually bounds how much extra - * partition DDL a failed sync can apply. - */ - public static final class Dispatch { - private final List> futures; - private final int total; - private final AtomicInteger settled = new AtomicInteger(0); - private final AtomicBoolean aborted = new AtomicBoolean(false); - private final CountDownLatch done = new CountDownLatch(1); - private volatile boolean sealed; - - private Dispatch(int total) { - this.total = total; - this.futures = new ArrayList<>(total); - } - - private void add(Future future) { - futures.add(future); - } - - // Called once submission finishes. A task that settles before the last submit - // would otherwise see settled < total and never trip the latch, so re-check here. - private void sealed() { - sealed = true; - signalIfComplete(); - } - - private boolean aborted() { - return aborted.get(); - } - - private void abort() { - aborted.set(true); - done.countDown(); - } - - private void taskSettled() { - settled.incrementAndGet(); - signalIfComplete(); - } - - private void signalIfComplete() { - if (sealed && settled.get() >= total) { - done.countDown(); - } - } - - private void awaitSettledOrAborted() { - if (total == 0) { - return; - } - try { - done.await(); - } catch (InterruptedException ie) { - Thread.currentThread().interrupt(); - aborted.set(true); - } - } - - // mayInterruptIfRunning=false: the worker thread is bound to a Hive Driver whose - // state we don't want to corrupt mid-statement. Cancel only tasks that haven't - // started; in-flight statements run to completion. - private int cancelPending() { - int cancelled = 0; - for (Future f : futures) { - if (f.cancel(false)) { - cancelled++; - } - } - return cancelled; - } - - private List> futures() { - return futures; - } - - @VisibleForTesting - public int size() { - return futures.size(); - } - - @VisibleForTesting - public Future futureAt(int index) { - return futures.get(index); - } - } - - private static Exception unwrap(ExecutionException ee) { - Throwable cause = ee.getCause(); - return (cause instanceof Exception) ? (Exception) cause : ee; + outcome.completed(), outcome.cancelled(), System.currentTimeMillis() - start, size); } public int size() { diff --git a/hudi-sync/hudi-hive-sync/src/main/java/org/apache/hudi/hive/util/IMetaStoreClientPool.java b/hudi-sync/hudi-hive-sync/src/main/java/org/apache/hudi/hive/util/IMetaStoreClientPool.java new file mode 100644 index 0000000000000..da1c7a8aac8d9 --- /dev/null +++ b/hudi-sync/hudi-hive-sync/src/main/java/org/apache/hudi/hive/util/IMetaStoreClientPool.java @@ -0,0 +1,262 @@ +/* + * 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.hudi.hive.util; + +import org.apache.hudi.exception.HoodieException; +import org.apache.hudi.hive.HiveSyncConfig; + +import org.apache.hadoop.hive.conf.HiveConf; +import org.apache.hadoop.hive.metastore.IMetaStoreClient; +import org.apache.hadoop.hive.metastore.RetryingMetaStoreClient; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.ArrayBlockingQueue; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.ThreadFactory; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; + +/** + * Pool of {@link IMetaStoreClient} instances for parallel partition sync. + * + *

Each pooled client wraps an independent Thrift connection to the Hive Metastore. + * Callers borrow a client via {@link #run(ClientAction)}, which blocks until a client + * is available, executes the action, and returns the client to the pool. A worker + * thread pool of the same size is exposed via {@link #executor()} so callers can fan + * out their batches to match the number of available clients. + * + *

Usage contract: pool clients must be used only for partition-row + * operations — {@code add_partitions}, {@code alter_partitions}, {@code dropPartition}, + * {@code getPartition}. Table-row operations ({@code createTable}, {@code alter_table}, + * {@code getTable} used as the read half of a read-modify-write of table parameters) + * must continue to go through the session client held by + * {@code HoodieHiveSyncClient.client} on the sync driver thread. Mixing the two would + * risk lost updates on table parameters such as the last-commit-time-synced marker. + * + *

The pool is gated behind {@code hoodie.datasource.hive_sync.batching.enabled} and + * is constructed for sync mode HIVEQL, where it backs the DROP path only. DROP goes + * through {@code IMetaStoreClient.dropPartition} (Thrift), whereas ADD/UPDATE/TOUCH go + * through the thread-bound Hive {@code Driver} and use {@code HiveDriverPool} instead. + */ +public class IMetaStoreClientPool implements AutoCloseable { + + private static final Logger LOG = LoggerFactory.getLogger(IMetaStoreClientPool.class); + + private final ArrayBlockingQueue available; + private final List all; + private final ExecutorService executor; + private final int size; + private volatile boolean closed; + + public IMetaStoreClientPool(HiveSyncConfig config, int size) { + this(buildClients(config, size), size); + } + + // Package-private for tests: accepts a pre-built list of clients so we can + // exercise borrow/return/close semantics without a live metastore. + IMetaStoreClientPool(List clients, int size) { + if (size < 1) { + throw new IllegalArgumentException("Pool size must be >= 1, got " + size); + } + if (clients.size() != size) { + throw new IllegalArgumentException("Expected " + size + " clients, got " + clients.size()); + } + this.size = size; + this.available = new ArrayBlockingQueue<>(size); + this.all = new ArrayList<>(clients); + this.available.addAll(clients); + this.executor = Executors.newFixedThreadPool(size, new PoolThreadFactory()); + LOG.info("Initialized IMetaStoreClient pool with {} clients", size); + } + + private static List buildClients(HiveSyncConfig config, int size) { + if (size < 1) { + throw new IllegalArgumentException("Pool size must be >= 1, got " + size); + } + HiveConf hiveConf = config.getHiveConf(); + List clients = new ArrayList<>(size); + try { + for (int i = 0; i < size; i++) { + clients.add(newClient(hiveConf)); + } + return clients; + } catch (Exception e) { + // Construction failed mid-way; close any clients we already built before + // surfacing the error so we don't leak Thrift sockets. + for (IMetaStoreClient c : clients) { + try { + c.close(); + } catch (Exception ignore) { + // intentional: best-effort cleanup during failure + } + } + throw new HoodieException("Failed to construct IMetaStoreClient pool of size " + size, e); + } + } + + private static IMetaStoreClient newClient(HiveConf hiveConf) { + try { + // RetryingMetaStoreClient.getProxy returns an independent IMetaStoreClient + // (one Thrift socket per call), bypassing the Hive thread-local cache that + // Hive.get(conf) would use. This is what gives us N truly independent clients. + return RetryingMetaStoreClient.getProxy(hiveConf, true); + } catch (Exception e) { + throw new HoodieException("Failed to create IMetaStoreClient for pool", e); + } + } + + /** + * Borrows a client, runs the action, and returns the client to the pool. + * Blocks if all clients are in use until one becomes available. + */ + public T run(ClientAction action) throws Exception { + if (closed) { + throw new IllegalStateException("Cannot borrow from a closed IMetaStoreClient pool"); + } + IMetaStoreClient client = available.take(); + try { + return action.apply(client); + } finally { + // Always return the client to the pool, even on failure. Thrift clients + // recover transparently from transactional errors at the HMS layer; + // RetryingMetaStoreClient handles transient socket failures internally. + if (!closed) { + available.offer(client); + } + } + } + + /** + * Submits one task per item, each borrowing a pooled client for the duration of its + * call, and returns a handle to the in-flight batch — this method does not block. The + * caller awaits completion via {@link #awaitAll(ParallelDispatch, String)}. + * + *

Tasks observe a shared abort flag, so a failure on any item stops items that have + * not started yet even while a slower sibling is still mid-call. See + * {@link ParallelDispatch} for why waiting on futures alone does not achieve that. + */ + public ParallelDispatch dispatchAll(List items, ClientConsumer action) { + if (closed) { + throw new IllegalStateException("Cannot dispatch to a closed IMetaStoreClient pool"); + } + ParallelDispatch dispatch = new ParallelDispatch(items.size()); + for (T item : items) { + dispatch.add(executor.submit(dispatch.guard(() -> { + run(client -> { + action.accept(client, item); + return null; + }); + return null; + }, "Skipped after an earlier batch failed"))); + } + dispatch.sealed(); + return dispatch; + } + + /** + * Awaits a batch from {@link #dispatchAll(List, ClientConsumer)}, cancels whatever had + * not started, and rethrows the first real failure. Later failures are logged at WARN. + */ + public void awaitAll(ParallelDispatch dispatch, String description) throws Exception { + long start = System.currentTimeMillis(); + ParallelDispatch.Outcome outcome = dispatch.awaitOutcome(); + outcome.suppressed().forEach(e -> + LOG.warn("Additional {} batch failed (suppressed in favor of first error)", description, e)); + + if (outcome.failed()) { + LOG.error("{} dispatch aborted after first failure ({} batches cancelled)", + description, outcome.cancelled()); + throw outcome.firstError(); + } + LOG.info("Completed {} {} batches ({} cancelled) in {} ms across {} clients", + outcome.completed(), description, outcome.cancelled(), + System.currentTimeMillis() - start, size); + } + + /** + * Worker thread pool sized to match the client pool. Use this to fan out + * batches so the number of in-flight Thrift calls cannot exceed the + * number of pooled clients. + */ + public ExecutorService executor() { + return executor; + } + + public int size() { + return size; + } + + @Override + public void close() { + if (closed) { + return; + } + closed = true; + executor.shutdown(); + try { + if (!executor.awaitTermination(30, TimeUnit.SECONDS)) { + executor.shutdownNow(); + } + } catch (InterruptedException ie) { + executor.shutdownNow(); + Thread.currentThread().interrupt(); + } + closeQuietly(); + } + + private void closeQuietly() { + for (IMetaStoreClient client : all) { + try { + client.close(); + } catch (Exception e) { + LOG.warn("Error closing pooled IMetaStoreClient", e); + } + } + available.clear(); + all.clear(); + } + + @FunctionalInterface + public interface ClientAction { + T apply(IMetaStoreClient client) throws Exception; + } + + /** Work applied to one fanned-out item using a borrowed client. */ + @FunctionalInterface + public interface ClientConsumer { + void accept(IMetaStoreClient client, T item) throws Exception; + } + + private static final class PoolThreadFactory implements ThreadFactory { + private static final AtomicInteger POOL_ID = new AtomicInteger(0); + private final AtomicInteger threadId = new AtomicInteger(0); + private final String namePrefix = "hudi-hive-sync-pool-" + POOL_ID.incrementAndGet() + "-"; + + @Override + public Thread newThread(Runnable r) { + Thread t = new Thread(r, namePrefix + threadId.incrementAndGet()); + t.setDaemon(true); + return t; + } + } +} diff --git a/hudi-sync/hudi-hive-sync/src/main/java/org/apache/hudi/hive/util/ParallelDispatch.java b/hudi-sync/hudi-hive-sync/src/main/java/org/apache/hudi/hive/util/ParallelDispatch.java new file mode 100644 index 0000000000000..255fabbce4c96 --- /dev/null +++ b/hudi-sync/hudi-hive-sync/src/main/java/org/apache/hudi/hive/util/ParallelDispatch.java @@ -0,0 +1,240 @@ +/* + * 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.hudi.hive.util; + +import org.apache.hudi.common.util.VisibleForTesting; + +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.Callable; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Future; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; + +/** + * Handle to one fan-out batch of partition work: the submitted futures plus the shared + * abort flag the tasks consult before running. + * + *

The abort flag exists because waiting on futures in submission order is not + * enough to stop queued work. If a later task fails quickly while an earlier one is slow, + * the awaiting thread is still parked on the earlier {@code Future.get()}, and the + * executor happily keeps starting every queued task in the meantime. By the time the + * failure is observed, most of the "not-yet-started" work has already run. + * + *

Two mechanisms fix that, and both are needed: + *

    + *
  • a {@link CountDownLatch} tripped by the first abort, so the awaiting + * thread wakes on failure rather than on its turn in submission order;
    • + *
  • a task-side {@link #aborted()} check on entry, because cancelling from the + * awaiting thread is inherently late — a worker can pull its next task off the + * queue at any moment.
    • + *
+ * + *

Shared by {@link HiveDriverPool} (Hive {@code Driver} statements) and + * {@link IMetaStoreClientPool} (Thrift {@code dropPartition} batches), which fan out over + * different execution models but need identical abort-on-first-error semantics. + */ +public final class ParallelDispatch { + + private final List> futures; + private final int total; + private final AtomicInteger settled = new AtomicInteger(0); + private final AtomicBoolean aborted = new AtomicBoolean(false); + private final CountDownLatch done = new CountDownLatch(1); + private volatile boolean sealed; + + ParallelDispatch(int total) { + this.total = total; + this.futures = new ArrayList<>(total); + } + + void add(Future future) { + futures.add(future); + } + + // Called once submission finishes. A task that settles before the last submit + // would otherwise see settled < total and never trip the latch, so re-check here. + void sealed() { + sealed = true; + signalIfComplete(); + } + + boolean aborted() { + return aborted.get(); + } + + void abort() { + aborted.set(true); + done.countDown(); + } + + void taskSettled() { + settled.incrementAndGet(); + signalIfComplete(); + } + + private void signalIfComplete() { + if (sealed && settled.get() >= total) { + done.countDown(); + } + } + + void awaitSettledOrAborted() { + if (total == 0) { + return; + } + try { + done.await(); + } catch (InterruptedException ie) { + Thread.currentThread().interrupt(); + aborted.set(true); + } + } + + // mayInterruptIfRunning=false: a worker may be mid-statement against a Hive Driver or + // a Thrift client, and we don't want to tear that down partway. Cancel only tasks that + // haven't started; in-flight work runs to completion. + int cancelPending() { + int cancelled = 0; + for (Future f : futures) { + if (f.cancel(false)) { + cancelled++; + } + } + return cancelled; + } + + List> futures() { + return futures; + } + + /** + * Wraps {@code body} so it observes this batch's abort flag: it skips itself if a + * sibling has already failed, trips the flag if it fails, and always records that it + * settled so the awaiting thread can be released. + */ + Callable guard(Callable body, String skipMessage) { + return () -> { + if (aborted()) { + throw new CancellationException(skipMessage); + } + try { + return body.call(); + } catch (Throwable t) { + abort(); + throw t; + } finally { + taskSettled(); + } + }; + } + + /** + * Waits for the batch to settle (or abort), cancels whatever had not started, and + * returns the outcome. Errors are observed in completion order, not submission + * order, so a failure on a fast worker stops the other queues even while a slow worker + * is still mid-statement. + */ + Outcome awaitOutcome() { + awaitSettledOrAborted(); + int cancelled = cancelPending(); + + Exception firstError = null; + int completed = 0; + List suppressed = new ArrayList<>(); + for (Future f : futures) { + try { + f.get(); + completed++; + } catch (CancellationException ce) { + // Either we cancelled it before it started, or the task itself observed the + // abort flag and bailed. Not a new failure; just note it for the summary. + cancelled++; + } catch (InterruptedException ie) { + Thread.currentThread().interrupt(); + if (firstError == null) { + firstError = ie; + } + } catch (ExecutionException ee) { + Exception cause = unwrap(ee); + if (cause instanceof CancellationException) { + cancelled++; + } else if (firstError == null) { + firstError = cause; + } else { + suppressed.add(cause); + } + } + } + return new Outcome(firstError, completed, cancelled, suppressed); + } + + private static Exception unwrap(ExecutionException ee) { + Throwable cause = ee.getCause(); + return (cause instanceof Exception) ? (Exception) cause : ee; + } + + @VisibleForTesting + public int size() { + return futures.size(); + } + + @VisibleForTesting + public Future futureAt(int index) { + return futures.get(index); + } + + /** Result of awaiting a batch: the first real failure, if any, plus counts for logging. */ + static final class Outcome { + private final Exception firstError; + private final int completed; + private final int cancelled; + private final List suppressed; + + private Outcome(Exception firstError, int completed, int cancelled, List suppressed) { + this.firstError = firstError; + this.completed = completed; + this.cancelled = cancelled; + this.suppressed = suppressed; + } + + Exception firstError() { + return firstError; + } + + boolean failed() { + return firstError != null; + } + + int completed() { + return completed; + } + + int cancelled() { + return cancelled; + } + + List suppressed() { + return suppressed; + } + } +} diff --git a/hudi-sync/hudi-hive-sync/src/test/java/org/apache/hudi/hive/TestHiveSyncTool.java b/hudi-sync/hudi-hive-sync/src/test/java/org/apache/hudi/hive/TestHiveSyncTool.java index 4f109c9cdd64e..7e2ebbe2bc21a 100644 --- a/hudi-sync/hudi-hive-sync/src/test/java/org/apache/hudi/hive/TestHiveSyncTool.java +++ b/hudi-sync/hudi-hive-sync/src/test/java/org/apache/hudi/hive/TestHiveSyncTool.java @@ -365,6 +365,46 @@ public void testHiveQLSyncWithBatchingEnabled() throws Exception { "Incremental add via parallel HiveQL batching should sync the new partitions"); } + /** + * Exercises the DROP path in HiveQL mode with batching on. DROP routes through + * IMetaStoreClient.dropPartition (Thrift, not Hive Driver), so when batching is + * enabled it fans out across IMetaStoreClientPool. Verifies the partition set + * shrinks as expected when batches drop in parallel. + */ + @Test + public void testHiveQLDropPartitionsWithBatching() throws Exception { + hiveSyncProps.setProperty(HIVE_SYNC_MODE.key(), HiveSyncMode.HIVEQL.name()); + hiveSyncProps.setProperty(HIVE_SYNC_BATCHING_ENABLED.key(), "true"); + hiveSyncProps.setProperty(HIVE_SYNC_BATCHING_THREADS.key(), "3"); + // Small batch_num so we get multiple drop batches dispatched in parallel. + hiveSyncProps.setProperty(HIVE_BATCH_SYNC_PARTITION_NUM.key(), "2"); + + int partitionCount = 8; + HiveTestUtil.createCOWTable("100", partitionCount, true); + reInitHiveSyncClient(); + reSyncHiveTable(); + assertEquals(partitionCount, hiveClient.getAllPartitions(HiveTestUtil.TABLE_NAME).size(), + "All partitions should be added before drop test"); + + // Drop half the partitions through the parallel pool path. + List existing = hiveClient.getAllPartitions(HiveTestUtil.TABLE_NAME).stream() + .map(p -> getRelativePartitionPath(new Path(basePath), new Path(p.getStorageLocation()))) + .collect(Collectors.toList()); + List toDrop = existing.subList(0, partitionCount / 2); + hiveClient.dropPartitions(HiveTestUtil.TABLE_NAME, toDrop); + + List remaining = hiveClient.getAllPartitions(HiveTestUtil.TABLE_NAME); + assertEquals(partitionCount - toDrop.size(), remaining.size(), + "Parallel DROP should remove exactly the requested partitions"); + Set remainingPaths = remaining.stream() + .map(p -> getRelativePartitionPath(new Path(basePath), new Path(p.getStorageLocation()))) + .collect(Collectors.toSet()); + for (String dropped : toDrop) { + assertFalse(remainingPaths.contains(dropped), + "Dropped partition " + dropped + " must not appear in remaining set"); + } + } + /** * Exercises the SET_LOCATION path in HiveQL mode with batching on. SET_LOCATION * emits one ALTER PARTITION ... SET LOCATION statement per partition (Hive SQL diff --git a/hudi-sync/hudi-hive-sync/src/test/java/org/apache/hudi/hive/util/TestHiveDriverPool.java b/hudi-sync/hudi-hive-sync/src/test/java/org/apache/hudi/hive/util/TestHiveDriverPool.java index fd58b7023aa19..414400e36796b 100644 --- a/hudi-sync/hudi-hive-sync/src/test/java/org/apache/hudi/hive/util/TestHiveDriverPool.java +++ b/hudi-sync/hudi-hive-sync/src/test/java/org/apache/hudi/hive/util/TestHiveDriverPool.java @@ -108,7 +108,7 @@ void runAllDispatchesEachSqlAcrossWorkers() throws Exception { }; try (HiveDriverPool pool = new HiveDriverPool(config, 2, factory)) { List sqls = Arrays.asList("SELECT 1", "SELECT 2", "SELECT 3", "SELECT 4"); - HiveDriverPool.Dispatch futures = pool.dispatchAll(sqls); + ParallelDispatch futures = pool.dispatchAll(sqls); pool.awaitAll(futures); assertEquals(2, seenThreadsByDriver.size(), "Expected exactly 2 worker Drivers"); int totalCalls = seenThreadsByDriver.values().stream().mapToInt(Set::size).sum(); @@ -135,7 +135,7 @@ void awaitAllThrowsFirstError() throws Exception { return d; }; try (HiveDriverPool pool = new HiveDriverPool(config, 2, factory)) { - HiveDriverPool.Dispatch futures = pool.dispatchAll(Arrays.asList("OK", "FAIL", "OK")); + ParallelDispatch futures = pool.dispatchAll(Arrays.asList("OK", "FAIL", "OK")); HoodieHiveSyncException ex = assertThrows(HoodieHiveSyncException.class, () -> pool.awaitAll(futures)); assertTrue(ex.getCause() != null && ex.getCause().getMessage().contains("boom")); @@ -161,7 +161,7 @@ void concurrentDispatchBoundedByPoolSize() throws Exception { }; try (HiveDriverPool pool = new HiveDriverPool(config, 2, factory)) { // 5 SQLs against pool of size 2 → max in-flight should be 2. - HiveDriverPool.Dispatch futures = pool.dispatchAll(Arrays.asList("a", "b", "c", "d", "e")); + ParallelDispatch futures = pool.dispatchAll(Arrays.asList("a", "b", "c", "d", "e")); // Release after a short wait so all SQLs progress. Thread.sleep(150); hold.countDown(); @@ -212,7 +212,7 @@ void runOnEachWorkerRunsSetupOnEveryWorker() throws Exception { }; try (HiveDriverPool pool = new HiveDriverPool(config, 3, factory)) { pool.runOnEachWorker(Arrays.asList("USE `db1`")); - HiveDriverPool.Dispatch futures = pool.dispatchAll(Arrays.asList("ALTER 1", "ALTER 2", "ALTER 3")); + ParallelDispatch futures = pool.dispatchAll(Arrays.asList("ALTER 1", "ALTER 2", "ALTER 3")); pool.awaitAll(futures); assertEquals(3, sqlsByDriver.size(), "Expected one Driver per worker"); @@ -249,7 +249,7 @@ void awaitAllCancelsPendingFuturesOnFirstError() throws Exception { return d; }; try (HiveDriverPool pool = new HiveDriverPool(config, 1, factory)) { - HiveDriverPool.Dispatch dispatch = pool.dispatchAll(Arrays.asList("FAIL", "PENDING_A", "PENDING_B")); + ParallelDispatch dispatch = pool.dispatchAll(Arrays.asList("FAIL", "PENDING_A", "PENDING_B")); HoodieHiveSyncException ex = assertThrows(HoodieHiveSyncException.class, () -> pool.awaitAll(dispatch)); @@ -299,7 +299,7 @@ void awaitAllStopsLaterWorkerWhenEarlierFutureIsSlow() throws Exception { }; try (HiveDriverPool pool = new HiveDriverPool(config, 2, factory)) { // Round-robin over 2 workers: index 0 -> worker 0, indices 1 and 2 -> worker 1. - HiveDriverPool.Dispatch dispatch = + ParallelDispatch dispatch = pool.dispatchAll(Arrays.asList("SLOW", "FAIL", "AFTER_FAIL")); assertTrue(failed.await(5, TimeUnit.SECONDS), "FAIL must have run"); releaseSlow.countDown(); diff --git a/hudi-sync/hudi-hive-sync/src/test/java/org/apache/hudi/hive/util/TestIMetaStoreClientPool.java b/hudi-sync/hudi-hive-sync/src/test/java/org/apache/hudi/hive/util/TestIMetaStoreClientPool.java new file mode 100644 index 0000000000000..85f2297da0a71 --- /dev/null +++ b/hudi-sync/hudi-hive-sync/src/test/java/org/apache/hudi/hive/util/TestIMetaStoreClientPool.java @@ -0,0 +1,251 @@ +/* + * 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.hudi.hive.util; + +import org.apache.hadoop.hive.metastore.IMetaStoreClient; +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.stream.Collectors; +import java.util.stream.IntStream; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; + +/** + * Unit tests for {@link IMetaStoreClientPool}: borrow/return bounding, batch fan-out, + * abort-on-first-error, and close semantics. Uses mock clients so no metastore is needed. + */ +class TestIMetaStoreClientPool { + + private static List mockClients(int n) { + return IntStream.range(0, n) + .mapToObj(i -> mock(IMetaStoreClient.class)) + .collect(Collectors.toList()); + } + + @Test + void runBorrowsAndReturnsTheSameClient() throws Exception { + List clients = mockClients(1); + try (IMetaStoreClientPool pool = new IMetaStoreClientPool(clients, 1)) { + // Given a single-client pool, two sequential borrows must both succeed -- + // which can only happen if the first borrow returned its client. + IMetaStoreClient first = pool.run(c -> c); + IMetaStoreClient second = pool.run(c -> c); + + assertEquals(clients.get(0), first); + assertEquals(first, second, "Client must be returned to the pool after each run"); + } + } + + @Test + void runReturnsClientEvenWhenActionThrows() throws Exception { + List clients = mockClients(1); + try (IMetaStoreClientPool pool = new IMetaStoreClientPool(clients, 1)) { + assertThrows(IllegalStateException.class, () -> pool.run(c -> { + throw new IllegalStateException("boom"); + })); + + // If the failed borrow had leaked the client, this would block forever. + IMetaStoreClient reused = pool.run(c -> c); + assertEquals(clients.get(0), reused, "A failed action must still return its client"); + } + } + + @Test + void dispatchAllRunsEveryBatch() throws Exception { + List clients = mockClients(3); + List batches = Arrays.asList("b0", "b1", "b2", "b3", "b4"); + List applied = Collections.synchronizedList(new ArrayList<>()); + try (IMetaStoreClientPool pool = new IMetaStoreClientPool(clients, 3)) { + pool.awaitAll(pool.dispatchAll(batches, (client, batch) -> applied.add(batch)), "test"); + } + + assertEquals(5, applied.size()); + assertTrue(applied.containsAll(batches), "Every batch must be applied exactly once"); + } + + @Test + void concurrentBatchesBoundedByPoolSize() throws Exception { + int poolSize = 2; + List clients = mockClients(poolSize); + AtomicInteger inFlight = new AtomicInteger(); + AtomicInteger maxInFlight = new AtomicInteger(); + try (IMetaStoreClientPool pool = new IMetaStoreClientPool(clients, poolSize)) { + pool.awaitAll(pool.dispatchAll(Arrays.asList("a", "b", "c", "d", "e"), (client, batch) -> { + int now = inFlight.incrementAndGet(); + maxInFlight.accumulateAndGet(now, Math::max); + Thread.sleep(20); + inFlight.decrementAndGet(); + }), "test"); + } + + assertTrue(maxInFlight.get() <= poolSize, + "In-flight Thrift calls must never exceed the client count, saw " + maxInFlight.get()); + } + + @Test + void eachConcurrentBatchGetsADistinctClient() throws Exception { + int poolSize = 3; + List clients = mockClients(poolSize); + Set seenConcurrently = ConcurrentHashMap.newKeySet(); + CountDownLatch allBorrowed = new CountDownLatch(poolSize); + try (IMetaStoreClientPool pool = new IMetaStoreClientPool(clients, poolSize)) { + pool.awaitAll(pool.dispatchAll(Arrays.asList("a", "b", "c"), (client, batch) -> { + seenConcurrently.add(client); + // Hold every client at once so none can be recycled to another batch. + allBorrowed.countDown(); + assertTrue(allBorrowed.await(5, TimeUnit.SECONDS)); + }), "test"); + } + + assertEquals(poolSize, seenConcurrently.size(), + "Concurrent batches must not share a Thrift client"); + } + + @Test + void awaitAllThrowsFirstError() { + List clients = mockClients(2); + IMetaStoreClientPool pool = new IMetaStoreClientPool(clients, 2); + try { + Exception ex = assertThrows(Exception.class, () -> + pool.awaitAll(pool.dispatchAll(Arrays.asList("ok", "boom"), (client, batch) -> { + if (batch.equals("boom")) { + throw new IllegalStateException("drop failed"); + } + }), "test")); + + assertEquals("drop failed", ex.getMessage(), + "The original cause must surface unwrapped, not as an ExecutionException"); + } finally { + pool.close(); + } + } + + /** + * Regression for the in-order-await bug: waiting on futures in submission order lets + * the executor keep starting queued batches after a sibling has already failed. + * + *

Given a single-client pool so batches run strictly in order, when the first batch + * fails, then no later batch may reach the metastore -- the task-side abort flag has to + * stop them, since {@code Future.cancel} from the awaiting thread is inherently late. + */ + @Test + void abortStopsQueuedBatchesAfterFirstFailure() { + List clients = mockClients(1); + List applied = Collections.synchronizedList(new ArrayList<>()); + IMetaStoreClientPool pool = new IMetaStoreClientPool(clients, 1); + try { + assertThrows(Exception.class, () -> + pool.awaitAll(pool.dispatchAll(Arrays.asList("FAIL", "AFTER_A", "AFTER_B"), + (client, batch) -> { + applied.add(batch); + if (batch.equals("FAIL")) { + throw new IllegalStateException("drop failed"); + } + }), "test")); + + assertEquals(Collections.singletonList("FAIL"), applied, + "Batches queued behind the failure must never reach the metastore"); + } finally { + pool.close(); + } + } + + /** + * The same abort guarantee, but with the failure landing on a later future than + * a still-running one. This is the interleaving Future-order waiting cannot handle: the + * awaiting thread is parked on SLOW's future while the failing batch races ahead. + */ + @Test + void abortStopsQueuedBatchesWhenEarlierBatchIsSlow() throws Exception { + List clients = mockClients(2); + List applied = Collections.synchronizedList(new ArrayList<>()); + CountDownLatch failed = new CountDownLatch(1); + CountDownLatch releaseSlow = new CountDownLatch(1); + IMetaStoreClientPool pool = new IMetaStoreClientPool(clients, 2); + try { + // SLOW occupies one client and blocks until FAIL has thrown, pinning the + // interleaving where awaitAll is still parked on future 0. + List batches = Arrays.asList("SLOW", "FAIL", "AFTER_FAIL", "AFTER_FAIL_2"); + ParallelDispatch dispatch = pool.dispatchAll(batches, (client, batch) -> { + applied.add(batch); + if (batch.equals("FAIL")) { + failed.countDown(); + throw new IllegalStateException("drop failed"); + } + if (batch.equals("SLOW")) { + releaseSlow.await(5, TimeUnit.SECONDS); + } + }); + assertTrue(failed.await(5, TimeUnit.SECONDS), "FAIL must have run"); + releaseSlow.countDown(); + + assertThrows(Exception.class, () -> pool.awaitAll(dispatch, "test")); + + assertFalse(applied.contains("AFTER_FAIL_2"), + "Batches queued behind a failure must not be applied while an earlier " + + "batch on another client is still running"); + } finally { + pool.close(); + } + } + + @Test + void closeIsIdempotentAndClosesEveryClient() throws Exception { + List clients = mockClients(2); + IMetaStoreClientPool pool = new IMetaStoreClientPool(clients, 2); + + pool.close(); + pool.close(); + + for (IMetaStoreClient client : clients) { + verify(client).close(); + } + assertThrows(IllegalStateException.class, () -> pool.run(c -> c), + "Borrowing from a closed pool must fail fast"); + } + + @Test + void dispatchOnClosedPoolFailsFast() { + IMetaStoreClientPool pool = new IMetaStoreClientPool(mockClients(1), 1); + pool.close(); + + assertThrows(IllegalStateException.class, + () -> pool.dispatchAll(Collections.singletonList("a"), (client, batch) -> { })); + } + + @Test + void rejectsSizeMismatchAndNonPositiveSize() { + assertThrows(IllegalArgumentException.class, () -> new IMetaStoreClientPool(mockClients(1), 2)); + assertThrows(IllegalArgumentException.class, () -> new IMetaStoreClientPool(mockClients(0), 0)); + } +}