This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/dev/pr-12316-cec3276740d989c8cb67c7b8c37e0a817e728059 in repository https://gitbox.apache.org/repos/asf/seatunnel.git
commit f867a5410f039e2706a450424ffd9ea8323d5104 Author: Daniel <[email protected]> AuthorDate: Wed Oct 28 07:20:45 2026 +0000 [Fix][Zeta] Hand the source checkpoint lock to barrier injectors instead of racing them (#12316) Co-authored-by: DanielLeens <[email protected]> Co-authored-by: Claude Fable 5.1 <[email protected]> Co-authored-by: David Zollo <[email protected]> Co-authored-by: David Zollo <[email protected]> --- .../engine/e2e/BackpressureSlowSinkIT.java | 110 +++- .../task/flow/SourceCheckpointLockHandoff.java | 110 ++++ .../server/task/flow/SourceFlowLifeCycle.java | 84 ++- .../task/flow/SourceCheckpointLockHandoffTest.java | 579 +++++++++++++++++++++ 4 files changed, 839 insertions(+), 44 deletions(-) diff --git a/seatunnel-e2e/seatunnel-engine-e2e/connector-seatunnel-e2e-base/src/test/java/org/apache/seatunnel/engine/e2e/BackpressureSlowSinkIT.java b/seatunnel-e2e/seatunnel-engine-e2e/connector-seatunnel-e2e-base/src/test/java/org/apache/seatunnel/engine/e2e/BackpressureSlowSinkIT.java index decca6e307..ebbaa4b5a1 100644 --- a/seatunnel-e2e/seatunnel-engine-e2e/connector-seatunnel-e2e-base/src/test/java/org/apache/seatunnel/engine/e2e/BackpressureSlowSinkIT.java +++ b/seatunnel-e2e/seatunnel-engine-e2e/connector-seatunnel-e2e-base/src/test/java/org/apache/seatunnel/engine/e2e/BackpressureSlowSinkIT.java @@ -17,6 +17,9 @@ package org.apache.seatunnel.engine.e2e; +import org.apache.seatunnel.api.common.metrics.Counter; +import org.apache.seatunnel.api.common.metrics.MetricNames; +import org.apache.seatunnel.api.options.EnvCommonOptions; import org.apache.seatunnel.engine.client.SeaTunnelClient; import org.apache.seatunnel.engine.client.job.ClientJobExecutionEnvironment; import org.apache.seatunnel.engine.client.job.ClientJobProxy; @@ -24,14 +27,20 @@ import org.apache.seatunnel.engine.common.config.ConfigProvider; import org.apache.seatunnel.engine.common.config.JobConfig; import org.apache.seatunnel.engine.common.config.SeaTunnelConfig; import org.apache.seatunnel.engine.common.job.JobStatus; +import org.apache.seatunnel.engine.server.SeaTunnelServer; import org.apache.seatunnel.engine.server.SeaTunnelServerStarter; +import org.apache.seatunnel.engine.server.dag.physical.PhysicalVertex; +import org.apache.seatunnel.engine.server.dag.physical.SubPlan; +import org.apache.seatunnel.engine.server.execution.Task; import org.apache.seatunnel.engine.server.rest.RestConstant; +import org.apache.seatunnel.engine.server.task.SourceSeaTunnelTask; import org.awaitility.Awaitility; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; import com.hazelcast.client.config.ClientConfig; import com.hazelcast.instance.impl.HazelcastInstanceImpl; @@ -121,22 +130,6 @@ public class BackpressureSlowSinkIT { config.getEngineConfig().getSlotServiceConfig().setDynamicSlot(false); config.getEngineConfig().getSlotServiceConfig().setSlotNum(20); node = SeaTunnelServerStarter.createHazelcastInstance(config); - - String filePath = TestUtils.getResource(CONF_FILE); - JobConfig jobConfig = new JobConfig(); - jobConfig.setName("backpressure_slow_sink_it"); - - ClientConfig clientConfig = ConfigProvider.locateAndGetClientConfig(); - clientConfig.setClusterName(testClusterName); - engineClient = new SeaTunnelClient(clientConfig); - ClientJobExecutionEnvironment jobExecutionEnv = - engineClient.createExecutionContext(filePath, jobConfig, config); - - jobProxy = jobExecutionEnv.execute(); - Awaitility.await() - .atMost(2, TimeUnit.MINUTES) - .untilAsserted( - () -> Assertions.assertEquals(JobStatus.RUNNING, jobProxy.getJobStatus())); } @AfterEach @@ -167,8 +160,11 @@ public class BackpressureSlowSinkIT { * * @throws Exception if polling the cluster's REST endpoints fails unexpectedly */ - @Test - public void testCheckpointsKeepCompletingUnderSustainedBackpressure() throws Exception { + @ParameterizedTest + @ValueSource(longs = {0L, 1000L}) + public void testCheckpointsKeepCompletingUnderSustainedBackpressure(long flushIntervalMs) + throws Exception { + startJob(flushIntervalMs); String checkpointBaseUrl = HOST + config.getEngineConfig().getHttpConfig().getPort() @@ -193,6 +189,8 @@ public class BackpressureSlowSinkIT { Assertions.assertTrue(getLong(counts, "completed") >= 1L); }); + Counter flushSignals = sourceFlushCounter(jobId); + List<Long> flushSamples = new ArrayList<>(); List<Long> completedSamples = new ArrayList<>(); List<Long> queueCapacitySamples = new ArrayList<>(); List<Long> blockedNsSamples = new ArrayList<>(); @@ -215,6 +213,7 @@ public class BackpressureSlowSinkIT { getLong(counts, "failed"), "no checkpoint should fail while the slow sink is applying backpressure"); completedSamples.add(getLong(counts, "completed")); + flushSamples.add(flushSignals.getCount()); Map<String, Object> edgesResp = getEdgesMetrics(jobId, metricsBaseUrl); List<Map<String, Object>> edges = castList(edgesResp.get("edges")); @@ -279,10 +278,83 @@ public class BackpressureSlowSinkIT { + "be > 0 for at least one sample); got " + blockedNsSamples); + // A timer tick must also cross the busy reader monitor repeatedly; checkpoint progress + // alone would not detect an omitted timer registration. Keep the same sustained workload + // and conservative progress floor rather than avoiding long polls with smaller splits. + if (flushIntervalMs > 0) { + long additionalFlushes = + flushSamples.get(flushSamples.size() - 1) - flushSamples.get(0); + Assertions.assertTrue( + additionalFlushes >= 3L, + "timer flush must keep progressing under sustained backpressure: " + + flushSamples); + } else { + Assertions.assertEquals( + 0L, flushSignals.getCount(), "timer flushing is disabled by default"); + } + // 4. Final health check after sustaining backpressure for the whole window. Assertions.assertEquals(JobStatus.RUNNING, jobProxy.getJobStatus()); } + /** + * Starts the unchanged single-split long-poll workload with optional timer-flush contention. + * + * <p>The zero-interval variant retains the original barrier-only regression; the positive + * interval also exercises the shared timer worker against the same busy source monitor. + */ + private void startJob(long flushIntervalMs) throws Exception { + String filePath = TestUtils.getResource(CONF_FILE); + JobConfig jobConfig = new JobConfig(); + jobConfig.setName("backpressure_slow_sink_it_" + flushIntervalMs); + if (flushIntervalMs > 0) { + jobConfig + .getEnvOptions() + .put(EnvCommonOptions.SINK_FLUSH_INTERVAL.key(), flushIntervalMs); + } + + ClientConfig clientConfig = ConfigProvider.locateAndGetClientConfig(); + clientConfig.setClusterName(config.getHazelcastConfig().getClusterName()); + engineClient = new SeaTunnelClient(clientConfig); + ClientJobExecutionEnvironment jobExecutionEnv = + engineClient.createExecutionContext(filePath, jobConfig, config); + + jobProxy = jobExecutionEnv.execute(); + Awaitility.await() + .atMost(2, TimeUnit.MINUTES) + .untilAsserted( + () -> Assertions.assertEquals(JobStatus.RUNNING, jobProxy.getJobStatus())); + } + + /** + * Reads the active worker source counter rather than a submitted or serialized task template. + * + * <p>The single-node test uses the physical plan only to resolve the stable source location; + * the TaskExecutionService supplies the actual running task and its live metric instance. + */ + private Counter sourceFlushCounter(long jobId) { + SeaTunnelServer server = node.node.getNodeEngine().getService(SeaTunnelServer.SERVICE_NAME); + for (SubPlan pipeline : + server.getCoordinatorService() + .getJobMaster(jobId) + .getPhysicalPlan() + .getPipelineList()) { + for (PhysicalVertex vertex : pipeline.getPhysicalVertexList()) { + for (Task template : vertex.getTaskGroup().getTasks()) { + if (template instanceof SourceSeaTunnelTask) { + SourceSeaTunnelTask<?, ?> source = + server.getTaskExecutionService() + .getTask( + ((SourceSeaTunnelTask<?, ?>) template) + .getTaskLocation()); + return source.getMetricsContext().counter(MetricNames.FLUSH_SIGNAL_TOTAL); + } + } + } + } + throw new AssertionError("No active source task found for backpressure job " + jobId); + } + private Map<String, Object> getCheckpointOverview(long jobId, String baseUrl) { return given().get(baseUrl + RestConstant.REST_URL_CHECKPOINT_OVERVIEW + "/" + jobId) .then() diff --git a/seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/task/flow/SourceCheckpointLockHandoff.java b/seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/task/flow/SourceCheckpointLockHandoff.java new file mode 100644 index 0000000000..7a324266b6 --- /dev/null +++ b/seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/task/flow/SourceCheckpointLockHandoff.java @@ -0,0 +1,110 @@ +/* + * 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.seatunnel.engine.server.task.flow; + +import java.util.concurrent.atomic.AtomicInteger; + +/** + * Cooperative handoff of the source checkpoint lock to barrier and timer-flush injectors. + * + * <p>A source reader can hold the checkpoint lock for the whole duration of {@code + * SourceReader#pollNext} and, under backpressure, every record it emits inside that call can block + * on the bounded intermediate queue; one poll can therefore hold the lock for many seconds. A + * barrier injector ({@link SourceFlowLifeCycle#triggerBarrier}) and the timer-flush callback need + * the same monitor. Java monitors are not fair: when the reader releases the lock and immediately + * re-enters the next poll, it wins the race against the parked injector almost every time, so a + * barrier can wait an unbounded number of poll cycles. {@code Thread.sleep(0)} between polls does + * not change that. + * + * <p>This class turns the unfair race into an explicit handoff without touching the public {@code + * Collector#getCheckpointLock()} contract: an injector registers itself before contending for the + * lock and deregisters after releasing it, and the reader thread, at a point where it holds no + * lock, waits until no injector is registered before starting its next poll. An announcement can + * race with a reader that already passed the check, but the reader cannot repeatedly overtake the + * injector across subsequent poll boundaries. Both barriers and timer-flush ticks register here. + * + * <p>The reader waits only while holding no checkpoint monitor. Injectors never wait on this + * handoff primitive, but can still wait for the poll already in progress, other injectors, and + * downstream delivery. A source that never returns from a poll or a downstream that never drains + * still prevents progress; there is no wall-clock checkpoint-latency guarantee. Registration must + * always be released in a {@code finally} block, including when injection fails. + */ +final class SourceCheckpointLockHandoff { + + /** + * Sleep slice while an injector is registered, avoiding a busy spin without adding another + * monitor or a notification protocol to the existing checkpoint monitor. Reader resumption + * follows OS timer granularity, not a precise millisecond deadline; downstream backpressure can + * keep the injector registered for much longer than one slice. + */ + static final long WAIT_SLICE_MS = 1L; + + /** + * Number of injectors that have announced themselves and not yet finished. More than one is + * possible when a checkpoint and a savepoint barrier are triggered close together. + */ + private final AtomicInteger pendingInjectors = new AtomicInteger(); + + /** + * Announces an injector that is about to contend for the checkpoint lock. Must be paired with + * {@link #injectorFinished()} in a {@code finally} block. + */ + void injectorArriving() { + pendingInjectors.incrementAndGet(); + } + + /** + * Withdraws an injector announcement after the injector released the checkpoint lock, whether + * the injection succeeded or failed. + */ + void injectorFinished() { + pendingInjectors.decrementAndGet(); + } + + /** + * Returns whether at least one injector is registered. + * + * @return true when an injector is waiting for or holding the checkpoint lock + */ + boolean hasPendingInjector() { + return pendingInjectors.get() > 0; + } + + /** + * Blocks the calling reader thread until no injector is registered. The caller must not hold + * the checkpoint lock, otherwise the injector it waits for can never acquire it. + * + * @return nanoseconds spent waiting, so the caller can account the time as idle + * @throws InterruptedException if the reader thread is interrupted, e.g. by task cancellation + */ + long awaitInjectors() throws InterruptedException { + // Preserve the interruption check of the Thread.sleep(0) replaced at the call site, + // including when no contender is pending and the reader takes the fast path. + if (Thread.interrupted()) { + throw new InterruptedException("Source checkpoint handoff interrupted"); + } + if (!hasPendingInjector()) { + return 0L; + } + long startNs = System.nanoTime(); + while (hasPendingInjector()) { + Thread.sleep(WAIT_SLICE_MS); + } + return System.nanoTime() - startNs; + } +} diff --git a/seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/task/flow/SourceFlowLifeCycle.java b/seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/task/flow/SourceFlowLifeCycle.java index 40a2cea117..fb51040064 100644 --- a/seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/task/flow/SourceFlowLifeCycle.java +++ b/seatunnel-engine/seatunnel-engine-server/src/main/java/org/apache/seatunnel/engine/server/task/flow/SourceFlowLifeCycle.java @@ -123,6 +123,15 @@ public class SourceFlowLifeCycle<T, SplitT extends SourceSplit> extends ActionFl private final AtomicReference<SchemaChangePhase> schemaChangePhase = new AtomicReference<>(); + /** + * Hands the checkpoint lock from the reader thread to barrier and timer-flush injectors. The + * reader holds the lock for the whole of {@code pollNext} and immediately re-enters the next + * poll, so without this explicit handoff a barrier thread parked on the unfair monitor can lose + * the race for many consecutive polls; under backpressure each poll can last seconds. + */ + private final SourceCheckpointLockHandoff checkpointLockHandoff = + new SourceCheckpointLockHandoff(); + private final long flushIntervalMs; private transient volatile ScheduledFuture<?> flushFuture; @@ -205,11 +214,21 @@ public class SourceFlowLifeCycle<T, SplitT extends SourceSplit> extends ActionFl if (prepareClose) { return; } + // Timer ticks contend on the same monitor as barriers. Register before contending so a + // busy reader cannot repeatedly reacquire the monitor ahead of the shared timer worker. + checkpointLockHandoff.injectorArriving(); try { - collector.sendFlushSignal( - currentTaskLocation.getJobId(), currentTaskLocation.getTaskID()); + synchronized (collector.getCheckpointLock()) { + // A closing barrier may have acquired the monitor while this tick was waiting. + if (!prepareClose) { + collector.sendFlushSignal( + currentTaskLocation.getJobId(), currentTaskLocation.getTaskID()); + } + } } catch (Exception e) { log.warn("Failed to broadcast FlushSignal from task {}", currentTaskLocation, e); + } finally { + checkpointLockHandoff.injectorFinished(); } } @@ -255,9 +274,11 @@ public class SourceFlowLifeCycle<T, SplitT extends SourceSplit> extends ActionFl * </ol> * * <p><b>Checkpoint lock interaction:</b> The reader holds the checkpoint lock during {@code - * pollNext}. A brief {@code Thread.sleep(0L)} after a non-empty poll gives the checkpoint - * thread a chance to acquire the lock via {@link #triggerBarrier(Barrier)}, preventing - * checkpoint starvation under high CPU load. + * pollNext}. After a non-empty poll, and before polling again, the reader waits on {@link + * SourceCheckpointLockHandoff} while a barrier injector ({@link #triggerBarrier(Barrier)}) or + * timer-flush tick has announced itself, so the reader cannot repeatedly overtake that + * contender at poll boundaries. The poll already in progress and downstream delivery must still + * finish; this is not a wall-clock bound. The wait is accounted as source idle time. * * @throws Exception if polling or schema-change triggering fails */ @@ -289,13 +310,18 @@ public class SourceFlowLifeCycle<T, SplitT extends SourceSplit> extends ActionFl } collector.resetEmptyThisPollNext(); /* - * The current thread obtain a checkpoint lock in the method {@link - * SourceReader#pollNext(Collector)}. When trigger the checkpoint or savepoint, - * other threads try to obtain the lock in the method {@link - * SourceFlowLifeCycle#triggerBarrier(Barrier)}. When high CPU load, checkpoint - * process may be blocked as long time. So we need sleep to free the CPU. + * The reader owned the checkpoint lock inside SourceReader#pollNext(Collector) + * and has released it here. A checkpoint or savepoint thread waiting in + * SourceFlowLifeCycle#triggerBarrier(Barrier) contends for the same unfair + * monitor and would usually lose to the reader's immediate re-entry, so the + * reader explicitly yields until every announced barrier or timer-flush + * injector has finished. The former Thread.sleep(0L) only yielded the CPU and + * did not prevent that race. */ - Thread.sleep(0L); + long handoffWaitNs = checkpointLockHandoff.awaitInjectors(); + if (metricsEnabled) { + sourceIdleNs.inc(handoffWaitNs); + } } if (collector.captureSchemaChangeBeforeCheckpointSignal()) { @@ -500,21 +526,29 @@ public class SourceFlowLifeCycle<T, SplitT extends SourceSplit> extends ActionFl long startTime = System.currentTimeMillis(); - // Block the reader from adding barrier to the collector. - synchronized (collector.getCheckpointLock()) { - if (barrier.prepareClose(this.currentTaskLocation)) { - this.prepareClose = true; - } - if (barrier.snapshot()) { - List<byte[]> states = - serializeStates(splitSerializer, reader.snapshotState(barrier.getId())); - runningTask.addState(barrier, ActionStateKey.of(sourceAction), states); + // Announce before contending for the lock so the reader thread, which re-enters the lock + // for every poll, stops polling until this barrier has been injected; withdraw in finally + // so a failed injection can never leave the reader parked. + checkpointLockHandoff.injectorArriving(); + try { + // Block the reader from adding barrier to the collector. + synchronized (collector.getCheckpointLock()) { + if (barrier.prepareClose(this.currentTaskLocation)) { + this.prepareClose = true; + } + if (barrier.snapshot()) { + List<byte[]> states = + serializeStates(splitSerializer, reader.snapshotState(barrier.getId())); + runningTask.addState(barrier, ActionStateKey.of(sourceAction), states); + } + // ack after #addState + runningTask.ack(barrier); + log.debug("source ack barrier finished, taskId: [{}]", runningTask.getTaskID()); + collector.sendRecordToNext(new Record<>(barrier)); + log.debug("send record to next finished, taskId: [{}]", runningTask.getTaskID()); } - // ack after #addState - runningTask.ack(barrier); - log.debug("source ack barrier finished, taskId: [{}]", runningTask.getTaskID()); - collector.sendRecordToNext(new Record<>(barrier)); - log.debug("send record to next finished, taskId: [{}]", runningTask.getTaskID()); + } finally { + checkpointLockHandoff.injectorFinished(); } log.debug( diff --git a/seatunnel-engine/seatunnel-engine-server/src/test/java/org/apache/seatunnel/engine/server/task/flow/SourceCheckpointLockHandoffTest.java b/seatunnel-engine/seatunnel-engine-server/src/test/java/org/apache/seatunnel/engine/server/task/flow/SourceCheckpointLockHandoffTest.java new file mode 100644 index 0000000000..11b395e369 --- /dev/null +++ b/seatunnel-engine/seatunnel-engine-server/src/test/java/org/apache/seatunnel/engine/server/task/flow/SourceCheckpointLockHandoffTest.java @@ -0,0 +1,579 @@ +/* + * 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.seatunnel.engine.server.task.flow; + +import org.apache.seatunnel.api.common.metrics.MetricNames; +import org.apache.seatunnel.api.serialization.Serializer; +import org.apache.seatunnel.api.signal.FlushSignal; +import org.apache.seatunnel.api.source.Boundedness; +import org.apache.seatunnel.api.source.Collector; +import org.apache.seatunnel.api.source.SeaTunnelSource; +import org.apache.seatunnel.api.source.SourceReader; +import org.apache.seatunnel.api.source.SourceSplit; +import org.apache.seatunnel.api.table.type.BasicType; +import org.apache.seatunnel.api.table.type.Record; +import org.apache.seatunnel.api.table.type.SeaTunnelDataType; +import org.apache.seatunnel.api.table.type.SeaTunnelRow; +import org.apache.seatunnel.api.table.type.SeaTunnelRowType; +import org.apache.seatunnel.core.starter.flowcontrol.FlowControlStrategy; +import org.apache.seatunnel.engine.common.config.EngineConfig; +import org.apache.seatunnel.engine.common.utils.concurrent.CompletableFuture; +import org.apache.seatunnel.engine.core.checkpoint.CheckpointType; +import org.apache.seatunnel.engine.core.dag.actions.SourceAction; +import org.apache.seatunnel.engine.server.TaskExecutionService; +import org.apache.seatunnel.engine.server.checkpoint.CheckpointBarrier; +import org.apache.seatunnel.engine.server.execution.TaskExecutionContext; +import org.apache.seatunnel.engine.server.execution.TaskGroupLocation; +import org.apache.seatunnel.engine.server.execution.TaskLocation; +import org.apache.seatunnel.engine.server.metrics.SeaTunnelMetricsContext; +import org.apache.seatunnel.engine.server.task.SeaTunnelSourceCollector; +import org.apache.seatunnel.engine.server.task.SeaTunnelTask; +import org.apache.seatunnel.engine.server.task.record.Barrier; + +import org.awaitility.Awaitility; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.mockito.Mockito; + +import com.hazelcast.cluster.Address; +import com.hazelcast.spi.impl.operationservice.impl.InvocationFuture; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ScheduledFuture; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; + +/** + * Protects the handoff counter and its real source lifecycle barrier/timer wiring. + * + * <p>Tests use actual collector monitors and controlled poll/injection boundaries so starvation, + * premature reader resumption, unbalanced failure cleanup, and cancellation each fail directly. + */ +public class SourceCheckpointLockHandoffTest { + + private static final long JOIN_TIMEOUT_MS = 10_000L; + + @Test + public void awaitReturnsImmediatelyWhenNoInjectorIsPending() throws InterruptedException { + SourceCheckpointLockHandoff handoff = new SourceCheckpointLockHandoff(); + Assertions.assertFalse(handoff.hasPendingInjector()); + Assertions.assertEquals(0L, handoff.awaitInjectors()); + } + + @Test + public void fastPathHonoursInterruptWithoutPendingInjector() { + SourceCheckpointLockHandoff handoff = new SourceCheckpointLockHandoff(); + try { + Thread.currentThread().interrupt(); + Assertions.assertThrows(InterruptedException.class, handoff::awaitInjectors); + Assertions.assertFalse(Thread.currentThread().isInterrupted()); + } finally { + // Keep the test runner healthy even when the assertion detects a regression. + Thread.interrupted(); + } + } + + @Test + public void pendingStateTracksEveryInjectorIndependently() { + SourceCheckpointLockHandoff handoff = new SourceCheckpointLockHandoff(); + handoff.injectorArriving(); + handoff.injectorArriving(); + Assertions.assertTrue(handoff.hasPendingInjector()); + handoff.injectorFinished(); + Assertions.assertTrue(handoff.hasPendingInjector()); + handoff.injectorFinished(); + Assertions.assertFalse(handoff.hasPendingInjector()); + } + + @Test + public void injectorAcquiresLockWithinOnePollCycle() throws Exception { + SourceCheckpointLockHandoff handoff = new SourceCheckpointLockHandoff(); + Object checkpointLock = new Object(); + AtomicInteger completedPolls = new AtomicInteger(); + AtomicBoolean readerRunning = new AtomicBoolean(true); + CountDownLatch readerHoldsLock = new CountDownLatch(1); + CountDownLatch releasePoll = new CountDownLatch(1); + AtomicInteger pollsBeforeInjector = new AtomicInteger(-1); + AtomicInteger pollsWhenInjectorAcquired = new AtomicInteger(-1); + AtomicReference<Throwable> failure = new AtomicReference<>(); + Thread reader = + thread( + "handoff-reader", + failure, + () -> { + while (readerRunning.get()) { + synchronized (checkpointLock) { + readerHoldsLock.countDown(); + Assertions.assertTrue( + releasePoll.await( + JOIN_TIMEOUT_MS, TimeUnit.MILLISECONDS)); + } + completedPolls.incrementAndGet(); + handoff.awaitInjectors(); + } + }); + Thread injector = + thread( + "handoff-injector", + failure, + () -> { + handoff.injectorArriving(); + try { + // Announce before sampling. Preemption before registration must not + // charge + // unrelated reader polls against the handoff's cross-poll bound. + pollsBeforeInjector.set(completedPolls.get()); + synchronized (checkpointLock) { + pollsWhenInjectorAcquired.set(completedPolls.get()); + } + } finally { + handoff.injectorFinished(); + } + }); + try { + reader.start(); + Assertions.assertTrue(readerHoldsLock.await(JOIN_TIMEOUT_MS, TimeUnit.MILLISECONDS)); + injector.start(); + awaitThreadState(injector, Thread.State.BLOCKED, failure); + releasePoll.countDown(); + injector.join(JOIN_TIMEOUT_MS); + Assertions.assertFalse(injector.isAlive(), "injector did not acquire the monitor"); + Assertions.assertNull(failure.get()); + int pollsWaited = pollsWhenInjectorAcquired.get() - pollsBeforeInjector.get(); + Assertions.assertTrue( + pollsWaited >= 0 && pollsWaited <= 1, + "injector waited " + pollsWaited + " poll cycles"); + Assertions.assertFalse(handoff.hasPendingInjector()); + } finally { + readerRunning.set(false); + releasePoll.countDown(); + stopThreads(reader, injector); + } + } + + @Test + public void waitingReaderResumesWhenInjectorFinishes() throws Exception { + SourceCheckpointLockHandoff handoff = new SourceCheckpointLockHandoff(); + handoff.injectorArriving(); + AtomicBoolean resumed = new AtomicBoolean(); + AtomicReference<Throwable> failure = new AtomicReference<>(); + Thread reader = + thread( + "waiting-reader", + failure, + () -> { + handoff.awaitInjectors(); + resumed.set(true); + }); + boolean registered = true; + try { + reader.start(); + awaitThreadState(reader, Thread.State.TIMED_WAITING, failure); + Assertions.assertFalse(resumed.get()); + handoff.injectorFinished(); + registered = false; + reader.join(JOIN_TIMEOUT_MS); + Assertions.assertFalse(reader.isAlive()); + Assertions.assertTrue(resumed.get()); + Assertions.assertNull(failure.get()); + } finally { + if (registered) { + handoff.injectorFinished(); + } + stopThreads(reader); + } + } + + @Test + public void waitingReaderHonoursInterrupt() throws Exception { + SourceCheckpointLockHandoff handoff = new SourceCheckpointLockHandoff(); + handoff.injectorArriving(); + AtomicBoolean interrupted = new AtomicBoolean(); + AtomicReference<Throwable> failure = new AtomicReference<>(); + Thread reader = + thread( + "interrupted-reader", + failure, + () -> { + try { + handoff.awaitInjectors(); + } catch (InterruptedException e) { + interrupted.set(true); + } + }); + try { + reader.start(); + awaitThreadState(reader, Thread.State.TIMED_WAITING, failure); + reader.interrupt(); + reader.join(JOIN_TIMEOUT_MS); + Assertions.assertFalse(reader.isAlive()); + Assertions.assertTrue(interrupted.get()); + Assertions.assertNull(failure.get()); + } finally { + handoff.injectorFinished(); + stopThreads(reader); + } + } + + @Test + public void lifecycleCollectHandsOffToBarrierBeforeReturning() throws Exception { + assertLifecycleHandoff(false); + } + + @Test + public void lifecycleCollectHandsOffToTimerBeforeReturning() throws Exception { + assertLifecycleHandoff(true); + } + + @Test + public void lifecycleFailuresReleaseBarrierAndTimerRegistrations() throws Exception { + for (String failurePoint : Arrays.asList("snapshot", "barrier-forward", "flush-forward")) { + LifecycleFixture fixture = new LifecycleFixture(); + AtomicReference<Throwable> failure = new AtomicReference<>(); + Thread reader = thread("reader-after-failed-injection", failure, fixture.flow::collect); + try { + fixture.releasePoll.countDown(); + fixture.releaseInjection.countDown(); + if ("snapshot".equals(failurePoint)) { + Mockito.when(fixture.reader.snapshotState(11L)) + .thenThrow(new IOException("snapshot failure")); + Assertions.assertThrows( + IOException.class, () -> fixture.flow.triggerBarrier(fixture.barrier)); + } else { + fixture.failForward.set(true); + if ("flush-forward".equals(failurePoint)) { + fixture.timerTick.get().run(); + } else { + Assertions.assertThrows( + IOException.class, + () -> fixture.flow.triggerBarrier(fixture.barrier)); + } + } + fixture.failForward.set(false); + reader.start(); + reader.join(JOIN_TIMEOUT_MS); + Assertions.assertFalse(reader.isAlive(), failurePoint + " left the reader parked"); + Assertions.assertNull(failure.get()); + Assertions.assertTrue( + fixture.records.get(fixture.records.size() - 1).getData() + instanceof SeaTunnelRow); + } finally { + fixture.releasePoll.countDown(); + fixture.releaseInjection.countDown(); + try { + stopThreads(reader); + } finally { + fixture.flow.close(); + } + } + } + } + + @Test + public void queuedTimerDoesNotCrossClosingBarrier() throws Exception { + LifecycleFixture fixture = new LifecycleFixture(); + AtomicReference<Throwable> failure = new AtomicReference<>(); + Thread timer = thread("timer-behind-close", failure, () -> fixture.timerTick.get().run()); + try { + fixture.releaseInjection.countDown(); + synchronized (fixture.collector.getCheckpointLock()) { + timer.start(); + awaitThreadState(timer, Thread.State.BLOCKED, failure); + fixture.flow.triggerBarrier( + new CheckpointBarrier(12L, 0L, CheckpointType.COMPLETED_POINT_TYPE)); + } + timer.join(JOIN_TIMEOUT_MS); + Assertions.assertFalse(timer.isAlive()); + Assertions.assertNull(failure.get()); + Assertions.assertEquals(1, fixture.records.size()); + Assertions.assertTrue(fixture.records.get(0).getData() instanceof Barrier); + Assertions.assertEquals( + 0L, fixture.metrics.counter(MetricNames.FLUSH_SIGNAL_TOTAL).getCount()); + } finally { + try { + stopThreads(timer); + } finally { + fixture.flow.close(); + } + } + } + + private static void assertLifecycleHandoff(boolean timer) throws Exception { + LifecycleFixture fixture = new LifecycleFixture(); + AtomicReference<Throwable> failure = new AtomicReference<>(); + Thread reader = thread("lifecycle-reader", failure, fixture.flow::collect); + Thread injector = + thread( + "lifecycle-injector", + failure, + () -> { + if (timer) { + fixture.timerTick.get().run(); + } else { + fixture.flow.triggerBarrier(fixture.barrier); + } + }); + try { + reader.start(); + Assertions.assertTrue( + fixture.pollEntered.await(JOIN_TIMEOUT_MS, TimeUnit.MILLISECONDS)); + injector.start(); + awaitThreadState(injector, Thread.State.BLOCKED, failure); + fixture.releasePoll.countDown(); + Assertions.assertTrue( + fixture.injectionEntered.await(JOIN_TIMEOUT_MS, TimeUnit.MILLISECONDS)); + Assertions.assertTrue( + fixture.pollReturned.await(JOIN_TIMEOUT_MS, TimeUnit.MILLISECONDS)); + // The real reader poll has returned. The real collector is forwarding the control + // record on another thread, so collect must now wait on the handoff outside its lock. + awaitThreadState(reader, Thread.State.TIMED_WAITING, failure); + fixture.releaseInjection.countDown(); + injector.join(JOIN_TIMEOUT_MS); + reader.join(JOIN_TIMEOUT_MS); + Assertions.assertFalse(injector.isAlive()); + Assertions.assertFalse(reader.isAlive()); + Assertions.assertNull(failure.get()); + fixture.flow.collect(); + Assertions.assertEquals(3, fixture.records.size()); + Assertions.assertTrue(fixture.records.get(0).getData() instanceof SeaTunnelRow); + Assertions.assertTrue(fixture.records.get(2).getData() instanceof SeaTunnelRow); + Assertions.assertTrue(fixture.records.get(1).getData() instanceof FlushSignal == timer); + if (!timer) { + Assertions.assertSame(fixture.barrier, fixture.records.get(1).getData()); + Assertions.assertEquals( + Arrays.asList("snapshot", "state", "ack", "forward"), fixture.barrierOrder); + Assertions.assertEquals(1, fixture.savedStates.get().size()); + Assertions.assertArrayEquals( + new byte[] {1, 2, 3}, fixture.savedStates.get().get(0)); + } + Assertions.assertTrue( + fixture.metrics.counter(MetricNames.SOURCE_IDLE_NANOS + "#1").getCount() > 0); + } finally { + fixture.releasePoll.countDown(); + fixture.releaseInjection.countDown(); + try { + stopThreads(reader, injector); + } finally { + fixture.flow.close(); + } + } + } + + private static void awaitThreadState( + Thread thread, Thread.State expected, AtomicReference<Throwable> failure) { + Awaitility.await() + .atMost(JOIN_TIMEOUT_MS, TimeUnit.MILLISECONDS) + .untilAsserted( + () -> { + Assertions.assertNull(failure.get()); + Assertions.assertEquals(expected, thread.getState()); + }); + } + + /** + * Captures background failures on the asserting thread and preserves cancellation interrupts. + * + * <p>Every caller owns unconditional bounded shutdown through {@link #stopThreads(Thread...)}. + */ + private static Thread thread( + String name, AtomicReference<Throwable> failure, CheckedRunnable action) { + return new Thread( + () -> { + try { + action.run(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } catch (Throwable e) { + failure.compareAndSet(null, e); + } + }, + name); + } + + private static void stopThreads(Thread... threads) throws InterruptedException { + for (Thread thread : threads) { + thread.interrupt(); + } + for (Thread thread : threads) { + thread.join(JOIN_TIMEOUT_MS); + } + for (Thread thread : threads) { + Assertions.assertFalse( + thread.isAlive(), "test-owned thread did not terminate: " + thread.getName()); + } + } + + /** + * A test action whose checked failures must be observed by the owning assertion thread. + * + * <p>This is only thread ownership plumbing; source/collector lifecycle behavior runs normally. + */ + private interface CheckedRunnable { + void run() throws Exception; + } + + /** + * Drives public lifecycle methods with a real collector and controlled reader/output edges. + * + * <p>No private lifecycle method is reflected or replaced: hook registers the timer callback, + * init creates the reader normally, and actual checkpoint serialization/ack/forwarding runs. + */ + private static final class LifecycleFixture { + private final CountDownLatch pollEntered = new CountDownLatch(1); + private final CountDownLatch releasePoll = new CountDownLatch(1); + private final CountDownLatch pollReturned = new CountDownLatch(1); + private final CountDownLatch injectionEntered = new CountDownLatch(1); + private final CountDownLatch releaseInjection = new CountDownLatch(1); + private final AtomicBoolean failForward = new AtomicBoolean(); + private final List<Record<?>> records = new ArrayList<>(); + private final List<String> barrierOrder = new ArrayList<>(); + private final AtomicReference<List<byte[]>> savedStates = new AtomicReference<>(); + private final AtomicReference<Runnable> timerTick = new AtomicReference<>(); + private final SeaTunnelMetricsContext metrics = new SeaTunnelMetricsContext(); + private final SourceReader<SeaTunnelRow, SourceSplit> reader; + private final SeaTunnelSourceCollector<SeaTunnelRow> collector; + private final SourceFlowLifeCycle<SeaTunnelRow, SourceSplit> flow; + private final CheckpointBarrier barrier = + new CheckpointBarrier(11L, 0L, CheckpointType.CHECKPOINT_TYPE); + + @SuppressWarnings("unchecked") + private LifecycleFixture() throws Exception { + TaskLocation location = new TaskLocation(new TaskGroupLocation(1L, 1, 1L), 1L, 0); + SeaTunnelTask task = Mockito.mock(SeaTunnelTask.class); + TaskExecutionContext execution = Mockito.mock(TaskExecutionContext.class); + TaskExecutionService service = Mockito.mock(TaskExecutionService.class); + Mockito.when(task.getTaskLocation()).thenReturn(location); + Mockito.when(task.getTaskID()).thenReturn(location.getTaskID()); + Mockito.when(task.getExecutionContext()).thenReturn(execution); + Mockito.when(task.isObservabilityEnabled()).thenReturn(true); + Mockito.when(execution.getTaskExecutionService()).thenReturn(service); + InvocationFuture<Object> address = Mockito.mock(InvocationFuture.class); + Mockito.when(address.get()).thenReturn(new Address("localhost", 5801)); + Mockito.doReturn(address).when(execution).sendToMaster(Mockito.any()); + Mockito.doAnswer( + call -> { + timerTick.set(call.getArgument(1)); + return Mockito.mock(ScheduledFuture.class); + }) + .when(service) + .registerTimerFlushTask(Mockito.eq(location), Mockito.any(), Mockito.eq(1000L)); + reader = Mockito.mock(SourceReader.class); + SeaTunnelSource<SeaTunnelRow, SourceSplit, ?> source = + Mockito.mock(SeaTunnelSource.class); + Mockito.doReturn(reader).when(source).createReader(Mockito.any()); + Mockito.when(source.getBoundedness()).thenReturn(Boundedness.UNBOUNDED); + Serializer<SourceSplit> serializer = Mockito.mock(Serializer.class); + SourceSplit split = Mockito.mock(SourceSplit.class); + Mockito.when(serializer.serialize(split)).thenReturn(new byte[] {1, 2, 3}); + Mockito.when(source.getSplitSerializer()).thenReturn(serializer); + Mockito.when(reader.snapshotState(11L)) + .thenAnswer( + call -> { + barrierOrder.add("snapshot"); + return Collections.singletonList(split); + }); + Mockito.doAnswer( + call -> { + barrierOrder.add("state"); + savedStates.set(call.getArgument(2)); + return null; + }) + .when(task) + .addState(Mockito.any(), Mockito.any(), Mockito.anyList()); + Mockito.doAnswer( + call -> { + barrierOrder.add("ack"); + return null; + }) + .when(task) + .ack(Mockito.any()); + OneInputFlowLifeCycle<Record<?>> output = + record -> { + if (!(record.getData() instanceof SeaTunnelRow)) { + if (failForward.get()) { + throw new IOException("control forwarding failure"); + } + injectionEntered.countDown(); + try { + Assertions.assertTrue( + releaseInjection.await( + JOIN_TIMEOUT_MS, TimeUnit.MILLISECONDS)); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IOException(e); + } + if (record.getData() instanceof Barrier) { + barrierOrder.add("forward"); + } + } + records.add(record); + }; + collector = + new SeaTunnelSourceCollector<>( + new Object(), + Collections.singletonList(output), + metrics, + FlowControlStrategy.builder().build(), + new SeaTunnelRowType( + new String[] {"id"}, + new SeaTunnelDataType[] {BasicType.INT_TYPE}), + Collections.emptyList(), + task, + new EngineConfig(), + Collections.emptyMap()); + Mockito.doAnswer( + call -> { + Collector<SeaTunnelRow> target = call.getArgument(0); + synchronized (target.getCheckpointLock()) { + target.collect(new SeaTunnelRow(new Object[] {1})); + pollEntered.countDown(); + Assertions.assertTrue( + releasePoll.await( + JOIN_TIMEOUT_MS, TimeUnit.MILLISECONDS)); + } + pollReturned.countDown(); + return null; + }) + .when(reader) + .pollNext(Mockito.any()); + SourceAction<SeaTunnelRow, SourceSplit, ?> action = + new SourceAction<>( + 1L, "source", source, Collections.emptySet(), Collections.emptySet()); + flow = + new SourceFlowLifeCycle<>( + action, + 0, + location, + task, + location, + new CompletableFuture<>(), + metrics, + 1000L); + flow.setCollector(collector); + flow.init(); + flow.hook(); + } + } +}
