diff --git a/CHANGELOG.adoc b/CHANGELOG.adoc index 543a21f07..ba14f0806 100644 --- a/CHANGELOG.adoc +++ b/CHANGELOG.adoc @@ -16,6 +16,14 @@ endif::[] == 0.5.3.4 +=== Features + +* feature: add options to enable virtual threads + +=== Improvements + +* refactor: migrate synchronized into ReentrantLock + === Fixes * fix: keep instantiated OffsetMapCodecManager so that metrics will not be recreated every commit (#859) diff --git a/README.adoc b/README.adoc index 39f4db1be..b00e0090d 100644 --- a/README.adoc +++ b/README.adoc @@ -1536,6 +1536,14 @@ endif::[] == 0.5.3.4 +=== Features + +* feature: add options to enable virtual threads + +=== Improvements + +* refactor: migrate synchronized into ReentrantLock + === Fixes * fix: keep instantiated OffsetMapCodecManager so that metrics will not be recreated every commit (#859) diff --git a/parallel-consumer-core/src/main/java/io/confluent/csid/utils/SupplierUtils.java b/parallel-consumer-core/src/main/java/io/confluent/csid/utils/SupplierUtils.java index d08616959..a3d1b179e 100644 --- a/parallel-consumer-core/src/main/java/io/confluent/csid/utils/SupplierUtils.java +++ b/parallel-consumer-core/src/main/java/io/confluent/csid/utils/SupplierUtils.java @@ -8,6 +8,7 @@ import java.util.Objects; import java.util.concurrent.atomic.AtomicReference; +import java.util.concurrent.locks.ReentrantLock; import java.util.function.Supplier; @UtilityClass @@ -16,15 +17,19 @@ public class SupplierUtils { public static Supplier memoize(Supplier delegate) { Objects.requireNonNull(delegate); AtomicReference value = new AtomicReference<>(); + ReentrantLock lock = new ReentrantLock(); return () -> { T val = value.get(); if (val == null) { - synchronized (value) { + lock.lock(); + try { val = value.get(); if (val == null) { val = Objects.requireNonNull(delegate.get()); value.set(val); } + } finally { + lock.unlock(); } } return val; diff --git a/parallel-consumer-core/src/main/java/io/confluent/parallelconsumer/ParallelConsumerOptions.java b/parallel-consumer-core/src/main/java/io/confluent/parallelconsumer/ParallelConsumerOptions.java index 05cae8240..6afdfe843 100644 --- a/parallel-consumer-core/src/main/java/io/confluent/parallelconsumer/ParallelConsumerOptions.java +++ b/parallel-consumer-core/src/main/java/io/confluent/parallelconsumer/ParallelConsumerOptions.java @@ -467,6 +467,21 @@ public void validate() { Objects.requireNonNull(consumer, "A consumer must be supplied"); transactionsValidation(); + validateVirtualThreadsSupport(); + } + + private void validateVirtualThreadsSupport() { + if (useVirtualThreads) { + try { + // Thread.ofVirtual(), Executors.newThreadPerTaskExecutor(ThreadFactory) + Class.forName("java.lang.Thread").getMethod("ofVirtual"); + java.util.concurrent.Executors.class.getMethod("newThreadPerTaskExecutor", java.util.concurrent.ThreadFactory.class); + } catch (NoSuchMethodException | ClassNotFoundException e) { + throw new UnsupportedOperationException( + "useVirtualThreads is enabled, but the current JVM does not support Virtual Threads or the required ExecutorService." + ); + } + } } private void transactionsValidation() { @@ -570,4 +585,14 @@ public boolean isProducerSupplied() { */ @Builder.Default public final boolean ignoreReflectiveAccessExceptionsForAutoCommitDisabledCheck = false; + + /** + * Whether to use Virtual Threads (Project Loom) for the worker pool. This feature requires running on JDK 21 or + * higher. If enabled on an older JDK the system will throw an exception during initialization. + *

+ * Note: When enabled, this overrides any custom {@link #managedExecutorService} or {@link #managedThreadFactory} + * settings, as it creates a specific ExecutorService for virtual threads. + */ + @Builder.Default + private final boolean useVirtualThreads = false; } diff --git a/parallel-consumer-core/src/main/java/io/confluent/parallelconsumer/internal/AbstractParallelEoSStreamProcessor.java b/parallel-consumer-core/src/main/java/io/confluent/parallelconsumer/internal/AbstractParallelEoSStreamProcessor.java index fc716650e..87d1575ea 100644 --- a/parallel-consumer-core/src/main/java/io/confluent/parallelconsumer/internal/AbstractParallelEoSStreamProcessor.java +++ b/parallel-consumer-core/src/main/java/io/confluent/parallelconsumer/internal/AbstractParallelEoSStreamProcessor.java @@ -33,6 +33,7 @@ import java.util.*; import java.util.concurrent.*; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.locks.ReentrantLock; import java.util.function.Consumer; import java.util.function.Function; import java.util.function.Supplier; @@ -116,7 +117,7 @@ public Duration getTimeBetweenCommits() { * The pool which is used for running the users' supplied function */ @Getter(PROTECTED) - protected final Supplier workerThreadPool; + protected final Supplier workerThreadPool; private Optional> controlThreadFuture = Optional.empty(); @@ -196,6 +197,8 @@ public static ControllerEventMessage of(WorkContainer work) { */ private final AtomicBoolean commitCommand = new AtomicBoolean(false); + private final ReentrantLock commitLock = new ReentrantLock(); + /** * Multiple of {@link ParallelConsumerOptions#getMaxConcurrency()} to have in our processing queue, in order to make * sure threads always have work to do. @@ -298,7 +301,17 @@ protected AbstractParallelEoSStreamProcessor(ParallelConsumerOptions newOp this.dynamicExtraLoadFactor = module.dynamicExtraLoadFactor(); - workerThreadPool = SupplierUtils.memoize(() -> setupWorkerPool(newOptions.getMaxConcurrency())); + workerThreadPool = SupplierUtils.memoize(() -> { + ExecutorService executor = setupWorkerPool(newOptions.getMaxConcurrency()); + + if (!options.isUseVirtualThreads() && !(executor instanceof ThreadPoolExecutor)) { + throw new IllegalStateException( + "When Virtual Threads are disabled, the worker pool must be an instance of ThreadPoolExecutor. " + + "If you are using a custom ExecutorService, make sure it extends ThreadPoolExecutor." + ); + } + return executor; + }); this.wm = module.workManager(); @@ -345,7 +358,24 @@ private void checkGroupIdConfigured(final org.apache.kafka.clients.consumer.Cons } } - protected ThreadPoolExecutor setupWorkerPool(int poolSize) { + protected ExecutorService setupWorkerPool(int poolSize) { + if (options.isUseVirtualThreads()) { + try { + // Thread.ofVirtual().name("pc-vt-").factory() + Object builder = Class.forName("java.lang.Thread").getMethod("ofVirtual").invoke(null); + Class builderClass = Class.forName("java.lang.Thread$Builder"); + builderClass.getMethod("name", String.class, long.class).invoke(builder, "pc-vt-", 0); + ThreadFactory factory = (ThreadFactory) builderClass.getMethod("factory").invoke(builder); + + // Executors.newThreadPerTaskExecutor(factory) + return (ExecutorService) java.util.concurrent.Executors.class + .getMethod("newThreadPerTaskExecutor", ThreadFactory.class) + .invoke(null, factory); + } catch (Exception e) { + throw new IllegalStateException("Virtual threads not supported on this JVM", e); + } + } + ThreadFactory defaultFactory; try { defaultFactory = InitialContext.doLookup(options.getManagedThreadFactory()); @@ -679,11 +709,11 @@ private void innerDoClose(Duration timeout) throws TimeoutException, ExecutionEx log.debug("Shutting down execution pool..."); //Clear scheduled but not started work in execution pool - workerThreadPool.get().getQueue().clear(); + clearWorkerQueue(); //request graceful shutdown workerThreadPool.get().shutdown(); - if (workerThreadPool.get().getActiveCount() > 0) { - log.info("Inflight work in execution pool: {}, letting to finish on shutdown with timeout: {}", workerThreadPool.get().getActiveCount(), timeout); + if (getPoolActiveCount() > 0) { + log.info("Inflight work in execution pool: {}, letting to finish on shutdown with timeout: {}", getPoolActiveCount(), timeout); } log.debug("Awaiting worker pool termination..."); @@ -708,8 +738,8 @@ private void innerDoClose(Duration timeout) throws TimeoutException, ExecutionEx } awaitingInflightProcessingCompletionOnShutdown.getAndSet(false); - if (workerThreadPool.get().getActiveCount() > 0) { - log.warn("Clean execution pool termination failed - some threads still active despite await and interrupt - is user function swallowing interrupted exception? Threads still not done count: {}", workerThreadPool.get().getActiveCount()); + if (getPoolActiveCount() > 0) { + log.warn("Clean execution pool termination failed - some threads still active despite await and interrupt - is user function swallowing interrupted exception? Threads still not done count: {}", getPoolActiveCount()); } log.debug("Worker pool terminated."); @@ -986,7 +1016,7 @@ private int retrieveAndDistributeNewWork(final Function { int queueSize = getNumberOfUserFunctionsQueued(); log.debug("Stats: \n- pool active: {} queued:{} \n- queue size: {} target: {} loading factor: {}", - workerThreadPool.get().getActiveCount(), queueSize, queueSize, getPoolLoadTarget(), dynamicExtraLoadFactor.getCurrentFactor()); + getPoolActiveCount(), queueSize, queueSize, getPoolLoadTarget(), dynamicExtraLoadFactor.getCurrentFactor()); }); return gotWorkCount; @@ -1105,7 +1135,12 @@ protected int getTargetOutForProcessing() { protected int getQueueTargetLoaded() { //noinspection unchecked - return getPoolLoadTarget() * dynamicExtraLoadFactor.getCurrentFactor(); + ExecutorService executor = workerThreadPool.get(); + if (executor instanceof ThreadPoolExecutor) { + return getPoolLoadTarget() * dynamicExtraLoadFactor.getCurrentFactor(); + } else { + return getPoolLoadTarget(); + } } /** @@ -1140,12 +1175,22 @@ private int getPoolLoadTarget() { } private boolean isPoolQueueLow() { - int queueSize = getNumberOfUserFunctionsQueued(); - int queueTarget = getPoolLoadTarget(); - boolean workAmountBelowTarget = queueSize <= queueTarget; - log.debug("isPoolQueueLow()? workAmountBelowTarget {} {} vs {};", - workAmountBelowTarget, queueSize, queueTarget); - return workAmountBelowTarget; + ExecutorService executor = workerThreadPool.get(); + + // Platform Threads + if (executor instanceof ThreadPoolExecutor) { + int queueSize = getNumberOfUserFunctionsQueued(); + int queueTarget = getPoolLoadTarget(); + boolean workAmountBelowTarget = queueSize <= queueTarget; + log.debug("isPoolQueueLow()? workAmountBelowTarget {} {} vs {};", + workAmountBelowTarget, queueSize, queueTarget); + return workAmountBelowTarget; + } + + // Virtual Threads + int activeCount = wm.getNumberRecordsOutForProcessing(); + int maxConcurrency = options.getTargetAmountOfRecordsInFlight(); + return activeCount <= maxConcurrency; } private void drain() { @@ -1183,7 +1228,7 @@ protected void processWorkCompleteMailBox(final Duration timeToBlockFor) { currentlyPollingWorkCompleteMailBox.getAndSet(true); if (log.isDebugEnabled()) { log.debug("Blocking poll on work until next scheduled offset commit attempt for {}. active threads: {}, queue: {}", - timeToBlockFor, workerThreadPool.get().getActiveCount(), getNumberOfUserFunctionsQueued()); + timeToBlockFor, getPoolActiveCount(), getNumberOfUserFunctionsQueued()); } // wait for work, with a timeToBlockFor for sanity log.trace("Blocking poll {}", timeToBlockFor); @@ -1278,7 +1323,14 @@ protected boolean isTimeToCommitNow() { } private int getNumberOfUserFunctionsQueued() { - return workerThreadPool.get().getQueue().size(); + // Platform Threads + ExecutorService executor = workerThreadPool.get(); + if (executor instanceof ThreadPoolExecutor) { + return ((ThreadPoolExecutor) executor).getQueue().size(); + } + + // Virtual Threads does not queue. + return 0; } @@ -1305,12 +1357,15 @@ private Duration getTimeSinceLastCheck() { * Visible for testing */ protected void commitOffsetsThatAreReady() throws TimeoutException, InterruptedException { - log.trace("Synchronizing on commitCommand..."); - synchronized (commitCommand) { + log.trace("Synchronizing on commitLock..."); + commitLock.lock(); + try { log.debug("Committing offsets that are ready..."); committer.retrieveOffsetsAndCommit(); clearCommitCommand(); this.lastCommitTime = Instant.now(); + } finally { + commitLock.unlock(); } } @@ -1483,8 +1538,11 @@ public void setLongPollTimeout(Duration ofMillis) { */ public void requestCommitAsap() { log.debug("Registering command to commit next chance"); - synchronized (commitCommand) { + commitLock.lock(); + try { this.commitCommand.set(true); + } finally { + commitLock.unlock(); } notifySomethingToDo(); } @@ -1517,18 +1575,45 @@ public void resumeIfPaused() { } private boolean isCommandedToCommit() { - synchronized (commitCommand) { + commitLock.lock(); + try { return this.commitCommand.get(); + } finally { + commitLock.unlock(); } } private void clearCommitCommand() { - synchronized (commitCommand) { + commitLock.lock(); + try { if (commitCommand.get()) { log.debug("Command to commit asap received, clearing"); this.commitCommand.set(false); } + } finally { + commitLock.unlock(); + } + } + + private int getPoolActiveCount() { + ExecutorService executor = workerThreadPool.get(); + + // Platform Threads + if (executor instanceof ThreadPoolExecutor) { + return ((ThreadPoolExecutor) executor).getActiveCount(); + } + + // Virtual Threads + return wm.getNumberRecordsOutForProcessing(); + } + + private void clearWorkerQueue() { + // Platform Threads + ExecutorService executor = workerThreadPool.get(); + if (executor instanceof ThreadPoolExecutor) { + ((ThreadPoolExecutor) executor).getQueue().clear(); } + // Virtual Threads do not have a task queue that can be cleared. } } \ No newline at end of file diff --git a/parallel-consumer-core/src/main/java/io/confluent/parallelconsumer/internal/DynamicLoadFactor.java b/parallel-consumer-core/src/main/java/io/confluent/parallelconsumer/internal/DynamicLoadFactor.java index 38a3034bc..f3d9d06d7 100644 --- a/parallel-consumer-core/src/main/java/io/confluent/parallelconsumer/internal/DynamicLoadFactor.java +++ b/parallel-consumer-core/src/main/java/io/confluent/parallelconsumer/internal/DynamicLoadFactor.java @@ -9,6 +9,7 @@ import java.time.Duration; import java.time.Instant; +import java.util.concurrent.locks.ReentrantLock; /** * Controls a loading factor. Is used to ensure enough messages in multiples of our target concurrency are queued ready @@ -70,6 +71,8 @@ public class DynamicLoadFactor { private long lastSteppedFactor = currentFactor; private Instant lastStepTime = Instant.MIN; + private final ReentrantLock lock = new ReentrantLock(); + public DynamicLoadFactor(int initial, int maximum) { this.currentFactor = initial; this.maxFactor = maximum; @@ -87,19 +90,24 @@ public boolean maybeStepUp() { return false; } - private synchronized boolean doStep() { - if (isMaxReached()) { - return false; - } else { - // compare and set - currentFactor = currentFactor + stepUpFactorBy; - long delta = currentFactor - lastSteppedFactor; - log.debug("Stepped up load factor by {} from {} to {}", delta, lastSteppedFactor, currentFactor); - - // - lastSteppedFactor = currentFactor; - lastStepTime = Instant.now(); - return true; + private boolean doStep() { + lock.lock(); + try { + if (isMaxReached()) { + return false; + } else { + // compare and set + currentFactor = currentFactor + stepUpFactorBy; + long delta = currentFactor - lastSteppedFactor; + log.debug("Stepped up load factor by {} from {} to {}", delta, lastSteppedFactor, currentFactor); + + // + lastSteppedFactor = currentFactor; + lastStepTime = Instant.now(); + return true; + } + } finally { + lock.unlock(); } } diff --git a/parallel-consumer-core/src/main/java/io/confluent/parallelconsumer/internal/ExternalEngine.java b/parallel-consumer-core/src/main/java/io/confluent/parallelconsumer/internal/ExternalEngine.java index 45dade404..bb6c27bc2 100644 --- a/parallel-consumer-core/src/main/java/io/confluent/parallelconsumer/internal/ExternalEngine.java +++ b/parallel-consumer-core/src/main/java/io/confluent/parallelconsumer/internal/ExternalEngine.java @@ -10,6 +10,7 @@ import lombok.extern.slf4j.Slf4j; import java.util.List; +import java.util.concurrent.ExecutorService; import java.util.concurrent.ThreadPoolExecutor; import static io.confluent.csid.utils.StringUtils.msg; @@ -53,7 +54,7 @@ protected void checkPipelinePressure() { * vert.x. */ @Override - protected ThreadPoolExecutor setupWorkerPool(int poolSize) { + protected ExecutorService setupWorkerPool(int poolSize) { return super.setupWorkerPool(1); } diff --git a/parallel-consumer-core/src/main/java/io/confluent/parallelconsumer/internal/ProducerManager.java b/parallel-consumer-core/src/main/java/io/confluent/parallelconsumer/internal/ProducerManager.java index b9c169703..f64c2299e 100644 --- a/parallel-consumer-core/src/main/java/io/confluent/parallelconsumer/internal/ProducerManager.java +++ b/parallel-consumer-core/src/main/java/io/confluent/parallelconsumer/internal/ProducerManager.java @@ -29,6 +29,7 @@ import java.util.Map; import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; +import java.util.concurrent.locks.ReentrantLock; import java.util.concurrent.locks.ReentrantReadWriteLock; import static io.confluent.csid.utils.StringUtils.msg; @@ -69,6 +70,8 @@ public class ProducerManager extends AbstractOffsetCommitter impleme @Getter private ReentrantReadWriteLock producerTransactionLock; + private final ReentrantLock beginTransactionLock = new ReentrantLock(); + public ProducerManager(ProducerWrapper newProducer, ConsumerManager newConsumer, WorkManager wm, @@ -155,14 +158,19 @@ private void lazyMaybeBeginTransaction() { } /** - * Pessimistic lock (synchronized method) on beginning a transaction + * Pessimistic lock (ReentrantLock) on beginning a transaction *

* Thread safe. */ - private synchronized void syncBeginTransaction() { - boolean txNotBegunAlready = !producerWrapper.isTransactionOpen(); - if (txNotBegunAlready) { - beginTransaction(); + private void syncBeginTransaction() { + beginTransactionLock.lock(); + try { + boolean txNotBegunAlready = !producerWrapper.isTransactionOpen(); + if (txNotBegunAlready) { + beginTransaction(); + } + } finally { + beginTransactionLock.unlock(); } } diff --git a/parallel-consumer-core/src/main/java/io/confluent/parallelconsumer/metrics/PCMetrics.java b/parallel-consumer-core/src/main/java/io/confluent/parallelconsumer/metrics/PCMetrics.java index 3725a7af8..630e27a22 100644 --- a/parallel-consumer-core/src/main/java/io/confluent/parallelconsumer/metrics/PCMetrics.java +++ b/parallel-consumer-core/src/main/java/io/confluent/parallelconsumer/metrics/PCMetrics.java @@ -13,6 +13,7 @@ import java.util.*; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.locks.ReentrantLock; import java.util.function.ToDoubleFunction; import static java.util.Collections.singleton; @@ -49,6 +50,8 @@ public class PCMetrics { private final boolean isNoop; + private final ReentrantLock lock = new ReentrantLock(); + /** * @param meterRegistry: meterRegistry to use for meter registration - configured through * {@link io.confluent.parallelconsumer.ParallelConsumerOptions} on PC initialization @@ -200,18 +203,23 @@ public DistributionSummary getDistributionSummaryFromMetricDef( /** * Closes PCMetrics object and cleans up all meters from registry - should be recreated before using it again. */ - public synchronized void close() { - if (this.isClosed.getAndSet(true)) { - //Instance already closed - warn and ignore. - log.warn("Trying to close PCMetrics instance that is already closed."); - return; - } - log.debug("Closing PCMetrics"); - // clean up the instance resources - this.registeredMeters.forEach(this.meterRegistry::remove); - this.registeredMeters.clear(); - if (isNoop) { - this.meterRegistry.close(); + public void close() { + lock.lock(); + try { + if (this.isClosed.getAndSet(true)) { + //Instance already closed - warn and ignore. + log.warn("Trying to close PCMetrics instance that is already closed."); + return; + } + log.debug("Closing PCMetrics"); + // clean up the instance resources + this.registeredMeters.forEach(this.meterRegistry::remove); + this.registeredMeters.clear(); + if (isNoop) { + this.meterRegistry.close(); + } + } finally { + lock.unlock(); } } @@ -223,9 +231,14 @@ public synchronized void close() { * * @param meter to remove. */ - public synchronized void removeMeter(Meter meter) { - if (meter != null) { - removeMeter(meter.getId()); + public void removeMeter(Meter meter) { + lock.lock(); + try { + if (meter != null) { + removeMeter(meter.getId()); + } + } finally { + lock.unlock(); } } diff --git a/parallel-consumer-core/src/test/java/io/confluent/parallelconsumer/AbstractParallelEoSStreamProcessorConfigurationTest.java b/parallel-consumer-core/src/test/java/io/confluent/parallelconsumer/AbstractParallelEoSStreamProcessorConfigurationTest.java index abe2444a8..5cd770cea 100644 --- a/parallel-consumer-core/src/test/java/io/confluent/parallelconsumer/AbstractParallelEoSStreamProcessorConfigurationTest.java +++ b/parallel-consumer-core/src/test/java/io/confluent/parallelconsumer/AbstractParallelEoSStreamProcessorConfigurationTest.java @@ -4,7 +4,7 @@ * Copyright (C) 2020-2024 Confluent, Inc. */ -import io.confluent.parallelconsumer.internal.PCModule; +import io.confluent.parallelconsumer.internal.DynamicLoadFactor; import io.confluent.parallelconsumer.internal.PCModuleTestEnv; import io.confluent.parallelconsumer.internal.TestParallelEoSStreamProcessor; import io.confluent.parallelconsumer.offsets.OffsetMapCodecManager; @@ -12,21 +12,33 @@ import io.confluent.parallelconsumer.state.PartitionState; import io.confluent.parallelconsumer.state.WorkContainer; import io.confluent.parallelconsumer.state.WorkManager; +import lombok.SneakyThrows; import lombok.extern.slf4j.Slf4j; import org.apache.kafka.clients.consumer.ConsumerRecord; -import org.apache.kafka.clients.consumer.OffsetResetStrategy; import org.apache.kafka.clients.consumer.MockConsumer; +import org.apache.kafka.clients.consumer.OffsetResetStrategy; import org.apache.kafka.common.TopicPartition; import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Assumptions; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import pl.tlinkowski.unij.api.UniLists; +import java.lang.reflect.Method; +import java.time.Duration; import java.util.ArrayList; import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Future; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Consumer; import java.util.function.Function; +import java.util.function.Supplier; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; /** * Tests to verify the protected and internal methods of @@ -42,7 +54,8 @@ class AbstractParallelEoSStreamProcessorConfigurationTest { .consumer(consumer) .build(); - ModelUtils mu = new ModelUtils(); + PCModuleTestEnv module; + ModelUtils mu; PartitionState state; WorkManager wm; @@ -50,17 +63,20 @@ class AbstractParallelEoSStreamProcessorConfigurationTest { int partition = 0; TopicPartition tp = new TopicPartition(topic, partition); - PCModule module = new PCModuleTestEnv(); @BeforeEach public void setup() { - state = new PartitionState<>(0, mu.getModule(), tp, OffsetMapCodecManager.HighestOffsetAndIncompletes.of()); - wm = mu.getModule().workManager(); + module = new PCModuleTestEnv(testOptions); + mu = new ModelUtils(module); + + wm = module.workManager(); wm.onPartitionsAssigned(UniLists.of(tp)); + state = new PartitionState<>(0, module, tp, OffsetMapCodecManager.HighestOffsetAndIncompletes.of()); } /** - * Test that the {@link io.confluent.parallelconsumer.internal.AbstractParallelEoSStreamProcessor#getQueueTargetLoaded} + * Test that the + * {@link io.confluent.parallelconsumer.internal.AbstractParallelEoSStreamProcessor#getQueueTargetLoaded} */ @Test void queueTargetLoad() { @@ -118,9 +134,139 @@ void testHandleStaleWorkNoSplit() { testInstance.runUserFunc(dummyFunction, callback, workContainers); - Assertions.assertEquals(testInstance.getMailBoxSuccessCnt(), 2); Assertions.assertEquals(testInstance.getMailBoxFailedCnt(), 0); } } -} + + private boolean isVirtualThreadsSupported() { + try { + Class.forName("java.lang.Thread").getMethod("ofVirtual"); + return true; + } catch (Exception e) { + return false; + } + } + + private boolean isCurrentThreadVirtual() { + try { + Method m = Thread.class.getMethod("isVirtual"); + return (boolean) m.invoke(Thread.currentThread()); + } catch (Exception e) { + return false; + } + } + + @Test + void testVirtualThreadBackpressureStability() { + Assumptions.assumeTrue(isVirtualThreadsSupported(), "JDK 21+ required for Virtual Threads"); + + var vtOptions = ParallelConsumerOptions.builder() + .consumer(consumer) + .useVirtualThreads(true) + .maxConcurrency(100) + .build(); + + try (var testInstance = new TestParallelEoSStreamProcessor(vtOptions) { + @Override + public void checkPipelinePressure() { + super.checkPipelinePressure(); + } + + public DynamicLoadFactor getLoadFactor() { + return this.dynamicExtraLoadFactor; + } + }) { + // Record the initial load factor + int initialLoadFactor = testInstance.getLoadFactor().getCurrentFactor(); + + testInstance.checkPipelinePressure(); + int loadFactorAfterCheck = testInstance.getLoadFactor().getCurrentFactor(); + + // The load factor should remain stable when we're at or near max concurrency + assertThat(loadFactorAfterCheck) + .as("Load factor should not increase excessively when near max concurrency") + .isLessThanOrEqualTo(initialLoadFactor + 1); // Allow for at most 1 step increase + + // Verify the system is stable by checking multiple times + for (int i = 0; i < 5; i++) { + testInstance.checkPipelinePressure(); + } + + int finalLoadFactor = testInstance.getLoadFactor().getCurrentFactor(); + assertThat(finalLoadFactor) + .as("Load factor should remain reasonable after multiple pressure checks") + .isLessThan(DynamicLoadFactor.DEFAULT_MAX_LOADING_FACTOR); + } + } + + @Test + void testVirtualThreadActivationAndConcurrencyLimit() throws Exception { + Assumptions.assumeTrue(isVirtualThreadsSupported(), "JDK 21+ required for Virtual Threads"); + + int maxConcurrency = 10; + var vtOptions = ParallelConsumerOptions.builder() + .consumer(consumer) + .useVirtualThreads(true) + .maxConcurrency(maxConcurrency) + .build(); + + try (var testInstance = new TestParallelEoSStreamProcessor<>(vtOptions) { + @Override + public Supplier getWorkerThreadPool() { + return super.getWorkerThreadPool(); + } + }) { + assertThat(testInstance.getWorkerThreadPool().get()).isNotNull(); + + AtomicInteger virtualThreadCount = new AtomicInteger(); + CountDownLatch latch = new CountDownLatch(1); + + Future future = testInstance.getWorkerThreadPool().get().submit(() -> { + if (isCurrentThreadVirtual()) { + virtualThreadCount.incrementAndGet(); + } + try { + latch.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + }); + + latch.countDown(); + future.get(); + + // Ensure VT is used in the test + assertThat(virtualThreadCount.get()).isEqualTo(1); + } + } + + @Test + @SneakyThrows + void testVirtualThreadCleanShutdown() { + Assumptions.assumeTrue(isVirtualThreadsSupported(), "JDK 21+ required for Virtual Threads"); + + var vtOptions = ParallelConsumerOptions.builder() + .consumer(consumer) + .useVirtualThreads(true) + .build(); + + try (var testInstance = new TestParallelEoSStreamProcessor<>(vtOptions) { + @Override + public Supplier getWorkerThreadPool() { + return super.getWorkerThreadPool(); + } + }) { + testInstance.getWorkerThreadPool().get().submit(() -> { + try { + Thread.sleep(Duration.ofMinutes(1).toMillis()); + } catch (InterruptedException ignored) { + } + }); + + Assertions.assertTimeout(Duration.ofSeconds(5), () -> { + testInstance.close(); + }, "Virtual thread pool should shutdown cleanly and interrupt running tasks"); + } + } +} \ No newline at end of file diff --git a/parallel-consumer-vertx/src/main/java/io/confluent/parallelconsumer/vertx/VertxParallelEoSStreamProcessor.java b/parallel-consumer-vertx/src/main/java/io/confluent/parallelconsumer/vertx/VertxParallelEoSStreamProcessor.java index 6a91218ec..cbb741db5 100644 --- a/parallel-consumer-vertx/src/main/java/io/confluent/parallelconsumer/vertx/VertxParallelEoSStreamProcessor.java +++ b/parallel-consumer-vertx/src/main/java/io/confluent/parallelconsumer/vertx/VertxParallelEoSStreamProcessor.java @@ -31,6 +31,7 @@ import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.concurrent.ExecutorService; import java.util.concurrent.ThreadPoolExecutor; import java.util.concurrent.TimeoutException; import java.util.function.BiFunction; @@ -120,7 +121,7 @@ public VertxParallelEoSStreamProcessor(Vertx vertx, * vert.x. */ @Override - protected ThreadPoolExecutor setupWorkerPool(int poolSize) { + protected ExecutorService setupWorkerPool(int poolSize) { return super.setupWorkerPool(1); }