diff --git a/.github/workflows/flink.yml b/.github/workflows/flink.yml index e32cc1b7fcb..d066e6859b7 100644 --- a/.github/workflows/flink.yml +++ b/.github/workflows/flink.yml @@ -87,8 +87,8 @@ jobs: export VELOX_DEPENDENCY_SOURCE=BUNDLED export fmt_SOURCE=BUNDLED export folly_SOURCE=BUNDLED - git clone -b gluten-0530 https://github.com/bigo-sg/velox4j.git - cd velox4j && git reset --hard 3ccc4e61dec4d569dc3dd9408b78ee1d4facf3ba + git clone -b feature/idle-source-handling https://github.com/bigo-sg/velox4j.git + cd velox4j && git reset --hard 9d2ae596a6b31c2f7d81f05d3dcba1ff953c6fab git apply $GITHUB_WORKSPACE/gluten-flink/patches/fix-velox4j.patch $GITHUB_WORKSPACE/build/mvn clean install -DskipTests -Dgpg.skip -Dspotless.skip=true cd .. diff --git a/gluten-flink/runtime/src/main/java/org/apache/gluten/client/OffloadedJobGraphGenerator.java b/gluten-flink/runtime/src/main/java/org/apache/gluten/client/OffloadedJobGraphGenerator.java index a9fa29557dd..e24fb4c0692 100644 --- a/gluten-flink/runtime/src/main/java/org/apache/gluten/client/OffloadedJobGraphGenerator.java +++ b/gluten-flink/runtime/src/main/java/org/apache/gluten/client/OffloadedJobGraphGenerator.java @@ -187,14 +187,16 @@ private OperatorChainSlice createOffloadedOperatorChainSlice( boolean supportsVectorOutput = supportsVectorOutput(sourceChainSlice, chainSliceGraph, jobVertex); Class outClass = supportsVectorOutput ? StatefulRecord.class : RowData.class; - GlutenStreamSource newSourceOp = - new GlutenStreamSource( - new GlutenSourceFunction<>( - planNode, - sourceOperator.getOutputTypes(), - sourceOperator.getId(), - ((GlutenStreamSource) sourceOperator).getConnectorSplit(), - outClass)); + GlutenSourceFunction newFn = + new GlutenSourceFunction<>( + planNode, + sourceOperator.getOutputTypes(), + sourceOperator.getId(), + ((GlutenStreamSource) sourceOperator).getConnectorSplit(), + outClass); + newFn.setShouldCallNoMoreSplits( + ((GlutenStreamSource) sourceOperator).isShouldCallNoMoreSplits()); + GlutenStreamSource newSourceOp = new GlutenStreamSource(newFn); offloadedOpConfig.setStreamOperator(newSourceOp); if (supportsVectorOutput) { setOffloadedOutputSerializer(offloadedOpConfig, sourceOperator); diff --git a/gluten-flink/runtime/src/main/java/org/apache/gluten/streaming/api/operators/GlutenStreamSource.java b/gluten-flink/runtime/src/main/java/org/apache/gluten/streaming/api/operators/GlutenStreamSource.java index acc6b2e60a5..dae9f84410c 100644 --- a/gluten-flink/runtime/src/main/java/org/apache/gluten/streaming/api/operators/GlutenStreamSource.java +++ b/gluten-flink/runtime/src/main/java/org/apache/gluten/streaming/api/operators/GlutenStreamSource.java @@ -79,6 +79,10 @@ public ConnectorSplit getConnectorSplit() { return sourceFunction.getConnectorSplit(); } + public boolean isShouldCallNoMoreSplits() { + return sourceFunction.isShouldCallNoMoreSplits(); + } + @SuppressWarnings("rawtypes") private SourceFunction.SourceContext getSourceContext() { return (SourceFunction.SourceContext) diff --git a/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenSourceFunction.java b/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenSourceFunction.java index a47b6c7c17b..c8c4b2fda86 100644 --- a/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenSourceFunction.java +++ b/gluten-flink/runtime/src/main/java/org/apache/gluten/table/runtime/operators/GlutenSourceFunction.java @@ -31,6 +31,7 @@ import io.github.zhztheplayer.velox4j.stateful.StatefulElement; import io.github.zhztheplayer.velox4j.stateful.StatefulRecord; import io.github.zhztheplayer.velox4j.stateful.StatefulWatermark; +import io.github.zhztheplayer.velox4j.stateful.StatefulWatermarkStatus; import io.github.zhztheplayer.velox4j.type.RowType; import org.apache.flink.api.common.state.ListState; @@ -64,6 +65,13 @@ public class GlutenSourceFunction extends RichParallelSourceFunction private final ConnectorSplit split; private volatile boolean isRunning = true; + /** + * If true, {@code noMoreSplits()} is called after adding the initial split in {@link + * #initSession()}. Set to false for unbounded streaming test scenarios where the task should stay + * alive to detect idleness. + */ + private boolean shouldCallNoMoreSplits = true; + private GlutenSessionResource sessionResource; private Query query; private SerialTask task; @@ -85,6 +93,16 @@ public GlutenSourceFunction( this.outClass = outClass; } + /** Sets whether noMoreSplits() should be called after adding the initial split. */ + public void setShouldCallNoMoreSplits(boolean value) { + this.shouldCallNoMoreSplits = value; + } + + /** Returns whether noMoreSplits() should be called after adding the initial split. */ + public boolean isShouldCallNoMoreSplits() { + return shouldCallNoMoreSplits; + } + public StatefulPlanNode getPlanNode() { return planNode; } @@ -132,8 +150,10 @@ private void processAvailableElement(SourceContext sourceContext) { processRecord(sourceContext, element.asRecord()); } else if (element.isWatermark()) { processWatermark(sourceContext, element.asWatermark()); + } else if (element.isWatermarkStatus()) { + processWatermarkStatus(sourceContext, element.asWatermarkStatus()); } else { - LOG.debug("Ignoring element that is neither record nor watermark"); + LOG.debug("Ignoring element that is neither record, watermark, nor watermark status"); } } finally { element.close(); @@ -169,6 +189,16 @@ private void processWatermark(SourceContext sourceContext, StatefulWatermar sourceContext.emitWatermark(new Watermark(watermark.getTimestamp())); } + /** Processes a watermark status and notifies the source context about idleness. */ + private void processWatermarkStatus( + SourceContext sourceContext, StatefulWatermarkStatus status) { + if (status.isIdle()) { + sourceContext.markAsTemporarilyIdle(); + } + // ACTIVE: no explicit action needed; the source context will resume + // activity tracking when the next record or watermark is emitted. + } + /** Collects a StatefulRecord as RowData by converting the RowVector. */ private void collectAsRowData(SourceContext sourceContext, StatefulRecord record) { List rows = @@ -270,7 +300,9 @@ private void initSession() { VeloxConnectorConfig.getConfig(getRuntimeContext())); task = session.queryOps().execute(query); task.addSplit(id, activeSplit); - task.noMoreSplits(id); + if (shouldCallNoMoreSplits) { + task.noMoreSplits(id); + } metrics = new SourceOperatorMetrics(getRuntimeContext().getMetricGroup()); } } diff --git a/gluten-flink/ut/pom.xml b/gluten-flink/ut/pom.xml index c97c3838e12..0f80c362927 100644 --- a/gluten-flink/ut/pom.xml +++ b/gluten-flink/ut/pom.xml @@ -288,7 +288,7 @@ org.apache.maven.plugins maven-surefire-plugin - 3.0.0-M5 + 3.3.0 false diff --git a/gluten-flink/ut/src/test/java/org/apache/gluten/streaming/api/operators/GlutenOneInputWatermarkAssignerIdleTest.java b/gluten-flink/ut/src/test/java/org/apache/gluten/streaming/api/operators/GlutenOneInputWatermarkAssignerIdleTest.java new file mode 100644 index 00000000000..00f28d52c8a --- /dev/null +++ b/gluten-flink/ut/src/test/java/org/apache/gluten/streaming/api/operators/GlutenOneInputWatermarkAssignerIdleTest.java @@ -0,0 +1,241 @@ +/* + * 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.gluten.streaming.api.operators; + +import org.apache.gluten.rexnode.RexConversionContext; +import org.apache.gluten.rexnode.RexNodeConverter; +import org.apache.gluten.rexnode.Utils; +import org.apache.gluten.table.runtime.operators.GlutenOneInputOperator; +import org.apache.gluten.table.runtime.stream.common.Velox4jEnvironment; +import org.apache.gluten.util.LogicalTypeConverter; +import org.apache.gluten.util.PlanNodeIdGenerator; + +import io.github.zhztheplayer.velox4j.expression.TypedExpr; +import io.github.zhztheplayer.velox4j.plan.EmptyNode; +import io.github.zhztheplayer.velox4j.plan.ProjectNode; +import io.github.zhztheplayer.velox4j.plan.StatefulPlanNode; +import io.github.zhztheplayer.velox4j.plan.WatermarkAssignerNode; + +import org.apache.flink.api.common.serialization.SerializerConfigImpl; +import org.apache.flink.api.common.typeinfo.TypeInformation; +import org.apache.flink.streaming.api.watermark.Watermark; +import org.apache.flink.streaming.runtime.streamrecord.StreamRecord; +import org.apache.flink.streaming.runtime.watermarkstatus.WatermarkStatus; +import org.apache.flink.streaming.util.OneInputStreamOperatorTestHarness; +import org.apache.flink.table.data.GenericRowData; +import org.apache.flink.table.data.RowData; +import org.apache.flink.table.planner.calcite.FlinkRexBuilder; +import org.apache.flink.table.planner.calcite.FlinkTypeFactory; +import org.apache.flink.table.planner.calcite.FlinkTypeSystem; +import org.apache.flink.table.runtime.typeutils.InternalTypeInfo; +import org.apache.flink.table.types.logical.BigIntType; +import org.apache.flink.table.types.logical.IntType; +import org.apache.flink.table.types.logical.LogicalType; +import org.apache.flink.table.types.logical.RowType; + +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.sql.type.SqlTypeName; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; + +import java.util.List; +import java.util.Map; +import java.util.Queue; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Integration test for WatermarkAssigner idle detection. + * + *

Native {@code checkWatermarkStatus} is driven by the {@code next()} → {@code advance()} loop + * inside the {@code GlutenOneInputOperator}'s drain pipeline. Since {@code addInput()} resets the + * idle baseline on every record, the idle check must happen on an {@code advance()} call that + * processes no new input. We achieve this by calling {@code processWatermark()} (which triggers a + * drain cycle without new data) after waiting past the idle timeout. + */ +public class GlutenOneInputWatermarkAssignerIdleTest { + + private static FlinkTypeFactory typeFactory; + private static FlinkRexBuilder rexBuilder; + private static RowType inputFlinkRowType; + private static io.github.zhztheplayer.velox4j.type.RowType inputVeloxType; + private static TypeInformation typeInfo; + + @BeforeAll + static void setUpClass() { + Velox4jEnvironment.initializeOnce(); + + typeFactory = + new FlinkTypeFactory( + Thread.currentThread().getContextClassLoader(), FlinkTypeSystem.INSTANCE); + rexBuilder = new FlinkRexBuilder(typeFactory); + + inputFlinkRowType = + RowType.of(new LogicalType[] {new IntType(), new BigIntType()}, new String[] {"id", "ts"}); + inputVeloxType = + (io.github.zhztheplayer.velox4j.type.RowType) + LogicalTypeConverter.toVLType(inputFlinkRowType); + typeInfo = InternalTypeInfo.of(inputFlinkRowType); + } + + @Test + void testIdleDetectionWithRealTimePassage() throws Exception { + long idleTimeout = 100L; // 100 ms + long watermarkInterval = 50L; + + // ── Watermark expression: reference the ts field (index 1) ── + List fieldNames = Utils.getNamesFromRowType(inputFlinkRowType); + RexNode tsRef = rexBuilder.makeInputRef(typeFactory.createSqlType(SqlTypeName.BIGINT), 1); + TypedExpr watermarkExpr = + RexNodeConverter.toTypedExpr(tsRef, new RexConversionContext(fieldNames)); + + ProjectNode watermarkProject = + new ProjectNode( + PlanNodeIdGenerator.newId(), + List.of(new EmptyNode(inputVeloxType)), + List.of("TIMESTAMP"), + List.of(watermarkExpr)); + + // ── WatermarkAssignerNode ── + WatermarkAssignerNode assignerNode = + new WatermarkAssignerNode( + PlanNodeIdGenerator.newId(), + null, + watermarkProject, + idleTimeout, + 1, // rowtimeFieldIndex (ts at index 1) + watermarkInterval); + + // ── GlutenOneInputOperator ── + GlutenOneInputOperator operator = + new GlutenOneInputOperator<>( + new StatefulPlanNode(assignerNode.getId(), assignerNode), + PlanNodeIdGenerator.newId(), + inputVeloxType, + Map.of(assignerNode.getId(), inputVeloxType), + RowData.class, + RowData.class, + "IdleDetectionTest"); + + // ── Test harness ── + OneInputStreamOperatorTestHarness harness = + new OneInputStreamOperatorTestHarness<>( + operator, typeInfo.createSerializer(new SerializerConfigImpl())); + harness.setup(typeInfo.createSerializer(new SerializerConfigImpl())); + harness.open(); + + try { + // ── Phase 1: feed one record ── + GenericRowData record1 = GenericRowData.of(1, 1000L); + harness.processElement(new StreamRecord<>(record1, 1000L)); + // output: StreamRecord(record1), Watermark(1000) + // After drain, checkWatermarkStatus(now1) scheduled timer at now1+100ms. + + // ── Phase 2: wait past idleTimeout, then trigger a drain WITHOUT new input ── + Thread.sleep(idleTimeout * 2); // 200 ms > 100 ms + harness.processWatermark(new Watermark(0)); + // Inside drain: advance → next → advanceWithFuture → blocked + // → checkWatermarkStatus(now2) → idle detected (now2 - lastRecordTime > 100ms) + // → push WatermarkStatus.IDLE to pendings_ + + // ── Phase 3: feed a second record — the drain first pops pending IDLE ── + GenericRowData record2 = GenericRowData.of(2, 2000L); + harness.processElement(new StreamRecord<>(record2, 2000L)); + // drain: advance → next → pendings_ non-empty (IDLE) → pop IDLE → emit IDLE + // → advance → next → process record2 → addInput → onRecord → idle was true + // → emit ACTIVE → push ACTIVE → advance → push record2 + watermark + // → pop ACTIVE → emit ACTIVE → pop record2 → collect → pop watermark → emit + + // ── Assertions ── + Queue output = harness.getOutput(); + assertThat(output) + .as("Output must contain WatermarkStatus.IDLE after idle timeout") + .anyMatch(e -> e instanceof WatermarkStatus && ((WatermarkStatus) e).isIdle()); + assertThat(output) + .as("Output must contain WatermarkStatus.ACTIVE after idle→active transition") + .anyMatch(e -> e instanceof WatermarkStatus && !((WatermarkStatus) e).isIdle()); + assertThat(output) + .as("All input records must be preserved") + .anyMatch( + e -> + e instanceof StreamRecord + && ((StreamRecord) e).getValue().getInt(0) == 1) + .anyMatch( + e -> + e instanceof StreamRecord + && ((StreamRecord) e).getValue().getInt(0) == 2); + } finally { + harness.close(); + } + } + + @Test + void testNoIdleWithContinuousRecords() throws Exception { + long idleTimeout = 100L; + long watermarkInterval = 50L; + + List fieldNames = Utils.getNamesFromRowType(inputFlinkRowType); + RexNode tsRef = rexBuilder.makeInputRef(typeFactory.createSqlType(SqlTypeName.BIGINT), 1); + TypedExpr watermarkExpr = + RexNodeConverter.toTypedExpr(tsRef, new RexConversionContext(fieldNames)); + + ProjectNode watermarkProject = + new ProjectNode( + PlanNodeIdGenerator.newId(), + List.of(new EmptyNode(inputVeloxType)), + List.of("TIMESTAMP"), + List.of(watermarkExpr)); + + WatermarkAssignerNode assignerNode = + new WatermarkAssignerNode( + PlanNodeIdGenerator.newId(), null, watermarkProject, idleTimeout, 1, watermarkInterval); + + GlutenOneInputOperator operator = + new GlutenOneInputOperator<>( + new StatefulPlanNode(assignerNode.getId(), assignerNode), + PlanNodeIdGenerator.newId(), + inputVeloxType, + Map.of(assignerNode.getId(), inputVeloxType), + RowData.class, + RowData.class, + "NoIdleTest"); + + OneInputStreamOperatorTestHarness harness = + new OneInputStreamOperatorTestHarness<>( + operator, typeInfo.createSerializer(new SerializerConfigImpl())); + harness.setup(typeInfo.createSerializer(new SerializerConfigImpl())); + harness.open(); + + try { + GenericRowData record1 = GenericRowData.of(1, 1000L); + GenericRowData record2 = GenericRowData.of(2, 2000L); + + harness.processElement(new StreamRecord<>(record1, 1000L)); + // Feed second record immediately (no idle gap) → addInput resets baseline + harness.processElement(new StreamRecord<>(record2, 2000L)); + // Then drain without new data — idle timeout has NOT elapsed since record2 + harness.processWatermark(new Watermark(0)); + + Queue output = harness.getOutput(); + assertThat(output) + .as("No WatermarkStatus.IDLE should appear when records arrive continuously") + .noneMatch(e -> e instanceof WatermarkStatus && ((WatermarkStatus) e).isIdle()); + } finally { + harness.close(); + } + } +} diff --git a/gluten-flink/ut/src/test/java/org/apache/gluten/streaming/api/operators/GlutenStreamTwoInputWatermarkStatusTest.java b/gluten-flink/ut/src/test/java/org/apache/gluten/streaming/api/operators/GlutenStreamTwoInputWatermarkStatusTest.java index 8306c9d20d9..a35b21716d5 100644 --- a/gluten-flink/ut/src/test/java/org/apache/gluten/streaming/api/operators/GlutenStreamTwoInputWatermarkStatusTest.java +++ b/gluten-flink/ut/src/test/java/org/apache/gluten/streaming/api/operators/GlutenStreamTwoInputWatermarkStatusTest.java @@ -79,6 +79,32 @@ public void testWatermarkReactivatesIdleNativeJoinOperator() throws Exception { } } + @Test + public void testIdleInputExcludedFromMinWatermark() throws Exception { + // When one input is idle, its watermark is excluded from the combined min-watermark + // calculation. The other active input's watermark can advance freely. + GlutenTwoInputOperator operator = createGlutenJoinOperator(FlinkJoinType.INNER); + + try (TwoInputStreamOperatorTestHarness harness = + new TwoInputStreamOperatorTestHarness<>(operator)) { + harness.setup(); + harness.open(); + + harness.processWatermark1(new Watermark(100L)); + harness.processWatermark2(new Watermark(90L)); + assertThat(harness.getOutput()).containsExactly(new Watermark(90L)); + + harness.processWatermarkStatus1(WatermarkStatus.IDLE); + // Input 1 (watermark=100) is idle and excluded. Combined = input 2 (90). No change. + assertThat(harness.getOutput()).containsExactly(new Watermark(90L)); + + // Input 2 advances to 120. Since input 1 is idle and excluded, combined = 120. + // If input 1 were still active, combined would be min(100, 120) = 100. + harness.processWatermark2(new Watermark(120L)); + assertThat(harness.getOutput()).containsExactly(new Watermark(90L), new Watermark(120L)); + } + } + @Test public void testWatermarksUseNativeTwoInputMinimum() throws Exception { // While both inputs are active, native execution should combine indexed input watermarks by diff --git a/gluten-flink/ut/src/test/java/org/apache/gluten/table/runtime/operators/GlutenSourceFunctionWatermarkStatusE2ETest.java b/gluten-flink/ut/src/test/java/org/apache/gluten/table/runtime/operators/GlutenSourceFunctionWatermarkStatusE2ETest.java new file mode 100644 index 00000000000..d5303eda299 --- /dev/null +++ b/gluten-flink/ut/src/test/java/org/apache/gluten/table/runtime/operators/GlutenSourceFunctionWatermarkStatusE2ETest.java @@ -0,0 +1,236 @@ +/* + * 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.gluten.table.runtime.operators; + +import org.apache.gluten.streaming.api.operators.GlutenStreamSource; +import org.apache.gluten.table.runtime.stream.common.Velox4jEnvironment; + +import io.github.zhztheplayer.velox4j.connector.KafkaConnectorSplit; +import io.github.zhztheplayer.velox4j.connector.KafkaTableHandle; +import io.github.zhztheplayer.velox4j.expression.FieldAccessTypedExpr; +import io.github.zhztheplayer.velox4j.plan.EmptyNode; +import io.github.zhztheplayer.velox4j.plan.ProjectNode; +import io.github.zhztheplayer.velox4j.plan.StatefulPlanNode; +import io.github.zhztheplayer.velox4j.plan.TableScanWithWatermarkNode; +import io.github.zhztheplayer.velox4j.plan.WatermarkPushDownSpec; +import io.github.zhztheplayer.velox4j.stateful.StatefulRecord; +import io.github.zhztheplayer.velox4j.type.BigIntType; +import io.github.zhztheplayer.velox4j.type.RowType; +import io.github.zhztheplayer.velox4j.type.VarCharType; + +import org.apache.flink.api.common.typeinfo.TypeInformation; +import org.apache.flink.streaming.api.datastream.DataStreamSource; +import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment; +import org.apache.flink.streaming.api.operators.AbstractStreamOperator; +import org.apache.flink.streaming.api.operators.OneInputStreamOperator; +import org.apache.flink.streaming.api.watermark.Watermark; +import org.apache.flink.streaming.runtime.streamrecord.StreamRecord; +import org.apache.flink.streaming.runtime.watermarkstatus.WatermarkStatus; + +import com.salesforce.kafka.test.junit5.SharedKafkaTestResource; +import com.salesforce.kafka.test.listeners.PlainListener; +import org.apache.kafka.clients.producer.ProducerRecord; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; + +import java.nio.charset.StandardCharsets; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.concurrent.CopyOnWriteArrayList; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * End-to-end test that verifies WatermarkStatus.IDLE is emitted through GlutenSourceFunction when + * the native Kafka source detects idleness. + * + *

The test produces a few records to an embedded Kafka broker, starts a Flink MiniCluster job + * with the Gluten native source pipeline, waits for the idle timeout to expire, and checks that + * WatermarkStatus.IDLE is captured by a downstream operator. + */ +class GlutenSourceFunctionWatermarkStatusE2ETest { + + private static final int KAFKA_PORT = 19093; + private static final long IDLE_TIMEOUT_MS = 5000; + private static final long WATERMARK_INTERVAL_MS = 500; + + @RegisterExtension + static final SharedKafkaTestResource KAFKA = + new SharedKafkaTestResource() + .withBrokerProperty("host.name", "127.0.0.1") + .withBrokers(1) + .registerListener(new PlainListener().onPorts(KAFKA_PORT)); + + private static final CopyOnWriteArrayList capturedStatuses = + new CopyOnWriteArrayList<>(); + + @BeforeAll + static void setupGluten() { + Velox4jEnvironment.initializeOnce(); + } + + @BeforeEach + void clearCaptured() { + capturedStatuses.clear(); + } + + @AfterEach + void ensureJobCancelled() { + if (jobThread != null && jobThread.isAlive()) { + jobThread.interrupt(); + } + } + + private Thread jobThread; + + @Test + void testIdleDetectionAfterStopWritingToKafka() throws Exception { + String topic = "idle_e2e_" + UUID.randomUUID().toString().replace("-", ""); + KAFKA.getKafkaTestUtils().createTopic(topic, 1, (short) 1); + KAFKA + .getKafkaTestUtils() + .produceRecords( + List.of( + jsonRecord(topic, "{\"id\":1000,\"name\":\"r0\"}"), + jsonRecord(topic, "{\"id\":2000,\"name\":\"r1\"}"), + jsonRecord(topic, "{\"id\":3000,\"name\":\"r2\"}"))); + + // Build pipeline: GlutenStreamSource -> WatermarkStatusCaptureOperator + StreamExecutionEnvironment env = StreamExecutionEnvironment.createLocalEnvironment(1); + env.getConfig().disableClosureCleaner(); + env.setParallelism(1); + + DataStreamSource source = addSourceToEnv(env, topic); + source + .transform("capture", TypeInformation.of(StatefulRecord.class), new StatusCaptureOp()) + .setParallelism(1); + + // Run job in background thread + jobThread = + new Thread( + () -> { + try { + env.execute("IdleDetectionE2ETest"); + } catch (Exception e) { + // ignore cancellation + } + }, + "job-runner"); + jobThread.start(); + + // Wait for initial records to be consumed + Thread.sleep(WATERMARK_INTERVAL_MS * 4); + + // Wait for idle timeout + Thread.sleep(IDLE_TIMEOUT_MS + 3000); + + // Cancel the job + jobThread.interrupt(); + + // Verify idle was detected + assertThat(capturedStatuses) + .as("Should have received WatermarkStatus.IDLE after idle timeout") + .contains(WatermarkStatus.IDLE); + } + + // -- helpers -- + + private static DataStreamSource addSourceToEnv( + StreamExecutionEnvironment env, String topic) { + RowType veloxRowType = + new RowType(List.of("id", "name"), List.of(new BigIntType(), new VarCharType())); + + ProjectNode watermarkProject = + new ProjectNode( + "watermark_project", + List.of(new EmptyNode(veloxRowType)), + List.of("watermark"), + List.of(FieldAccessTypedExpr.create(new BigIntType(), "id"))); + WatermarkPushDownSpec watermarkSpec = + new WatermarkPushDownSpec(watermarkProject, IDLE_TIMEOUT_MS, WATERMARK_INTERVAL_MS, 0); + + Map tableParams = new HashMap<>(); + tableParams.put("bootstrap.servers", "127.0.0.1:" + KAFKA_PORT); + tableParams.put("client.id", "test-client-e2e-" + UUID.randomUUID()); + tableParams.put("group.id", "test-group-e2e"); + tableParams.put("topic", topic); + tableParams.put("format", "json"); + tableParams.put("scan.startup.mode", "earliest-offsets"); + tableParams.put("enable.auto.commit", "false"); + + KafkaTableHandle tableHandle = + new KafkaTableHandle("connector-kafka", topic, veloxRowType, tableParams); + String planId = "plan_" + UUID.randomUUID().toString().replace("-", ""); + TableScanWithWatermarkNode scanNode = + new TableScanWithWatermarkNode(planId, veloxRowType, tableHandle, List.of(), watermarkSpec); + KafkaConnectorSplit connectorSplit = + new KafkaConnectorSplit( + "connector-kafka", + 0, + false, + "127.0.0.1:" + KAFKA_PORT, + "test-group-e2e", + "json", + false, + "earliest-offset", + List.of(new KafkaConnectorSplit.TopicPartitionOffset(topic, 0, -1L))); + + GlutenSourceFunction sourceFn = + new GlutenSourceFunction<>( + new StatefulPlanNode(scanNode.getId(), scanNode), + Map.of(scanNode.getId(), veloxRowType), + scanNode.getId(), + connectorSplit, + StatefulRecord.class); + sourceFn.setShouldCallNoMoreSplits(false); + + GlutenStreamSource sourceOp = new GlutenStreamSource(sourceFn, "KafkaSource"); + return new DataStreamSource( + env, TypeInformation.of(StatefulRecord.class), sourceOp, false, "KafkaSource"); + } + + private static ProducerRecord jsonRecord(String topic, String value) { + return new ProducerRecord<>(topic, value.getBytes(StandardCharsets.UTF_8)); + } + + // -- Capture operator -- + + private static class StatusCaptureOp extends AbstractStreamOperator + implements OneInputStreamOperator { + + @Override + public void processElement(StreamRecord element) throws Exception { + output.collect(element); + } + + @Override + public void processWatermark(Watermark mark) throws Exception { + output.emitWatermark(mark); + } + + @Override + public void processWatermarkStatus(WatermarkStatus status) throws Exception { + capturedStatuses.add(status); + super.processWatermarkStatus(status); + } + } +}