This is an automated email from the ASF dual-hosted git repository. guoweijie pushed a commit to branch master in repository https://gitbox.apache.org/repos/asf/flink.git
commit efbeaaa15761cb3fd5499c58d6b9ab842860ac5c Author: Xu Huang <[email protected]> AuthorDate: Thu Nov 28 15:12:54 2024 +0800 [FLINK-36815][API] Open and close user function in DSv2 ProcessOperators --- .../datastream/impl/operators/ProcessOperator.java | 7 + .../TwoInputBroadcastProcessOperator.java | 7 + .../TwoInputNonBroadcastProcessOperator.java | 7 + .../impl/operators/TwoOutputProcessOperator.java | 7 + .../impl/functions/ProcessFunctionTest.java | 200 +++++++++++++++++++++ 5 files changed, 228 insertions(+) diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/ProcessOperator.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/ProcessOperator.java index 10876fbeaf2..ce703012b25 100644 --- a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/ProcessOperator.java +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/ProcessOperator.java @@ -76,6 +76,7 @@ public class ProcessOperator<IN, OUT> getOperatorStateBackend()); outputCollector = getOutputCollector(); nonPartitionedContext = getNonPartitionedContext(); + userFunction.open(); } @Override @@ -105,4 +106,10 @@ public class ProcessOperator<IN, OUT> return new DefaultNonPartitionedContext<>( context, partitionedContext, outputCollector, false, null); } + + @Override + public void close() throws Exception { + super.close(); + userFunction.close(); + } } diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoInputBroadcastProcessOperator.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoInputBroadcastProcessOperator.java index dea6d2cd2d0..cf707f8dedf 100644 --- a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoInputBroadcastProcessOperator.java +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoInputBroadcastProcessOperator.java @@ -80,6 +80,7 @@ public class TwoInputBroadcastProcessOperator<IN1, IN2, OUT> operatorContext, getOperatorStateBackend()); this.nonPartitionedContext = getNonPartitionedContext(); + this.userFunction.open(); } @Override @@ -122,4 +123,10 @@ public class TwoInputBroadcastProcessOperator<IN1, IN2, OUT> protected ProcessingTimeManager getProcessingTimeManager() { return UnsupportedProcessingTimeManager.INSTANCE; } + + @Override + public void close() throws Exception { + super.close(); + userFunction.close(); + } } diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoInputNonBroadcastProcessOperator.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoInputNonBroadcastProcessOperator.java index c3cfb438792..d952efed124 100644 --- a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoInputNonBroadcastProcessOperator.java +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoInputNonBroadcastProcessOperator.java @@ -83,6 +83,7 @@ public class TwoInputNonBroadcastProcessOperator<IN1, IN2, OUT> operatorContext, operatorStateBackend); this.nonPartitionedContext = getNonPartitionedContext(); + this.userFunction.open(); } @Override @@ -125,4 +126,10 @@ public class TwoInputNonBroadcastProcessOperator<IN1, IN2, OUT> protected ProcessingTimeManager getProcessingTimeManager() { return UnsupportedProcessingTimeManager.INSTANCE; } + + @Override + public void close() throws Exception { + super.close(); + userFunction.close(); + } } diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoOutputProcessOperator.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoOutputProcessOperator.java index 3f4f8e05bec..fb6209ef18f 100644 --- a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoOutputProcessOperator.java +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/TwoOutputProcessOperator.java @@ -92,6 +92,7 @@ public class TwoOutputProcessOperator<IN, OUT_MAIN, OUT_SIDE> operatorContext, operatorStateStore); this.nonPartitionedContext = getNonPartitionedContext(); + this.userFunction.open(); } @Override @@ -150,4 +151,10 @@ public class TwoOutputProcessOperator<IN, OUT_MAIN, OUT_SIDE> output.collect(outputTag, reuse.replace(record)); } } + + @Override + public void close() throws Exception { + super.close(); + userFunction.close(); + } } diff --git a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/functions/ProcessFunctionTest.java b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/functions/ProcessFunctionTest.java new file mode 100644 index 00000000000..2d4330e204e --- /dev/null +++ b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/functions/ProcessFunctionTest.java @@ -0,0 +1,200 @@ +/* + * 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.flink.datastream.impl.functions; + +import org.apache.flink.api.common.typeinfo.Types; +import org.apache.flink.datastream.api.common.Collector; +import org.apache.flink.datastream.api.context.NonPartitionedContext; +import org.apache.flink.datastream.api.context.PartitionedContext; +import org.apache.flink.datastream.api.function.OneInputStreamProcessFunction; +import org.apache.flink.datastream.api.function.ProcessFunction; +import org.apache.flink.datastream.api.function.TwoInputBroadcastStreamProcessFunction; +import org.apache.flink.datastream.api.function.TwoInputNonBroadcastStreamProcessFunction; +import org.apache.flink.datastream.api.function.TwoOutputStreamProcessFunction; +import org.apache.flink.datastream.impl.operators.ProcessOperator; +import org.apache.flink.datastream.impl.operators.TwoInputBroadcastProcessOperator; +import org.apache.flink.datastream.impl.operators.TwoInputNonBroadcastProcessOperator; +import org.apache.flink.datastream.impl.operators.TwoOutputProcessOperator; +import org.apache.flink.streaming.util.OneInputStreamOperatorTestHarness; +import org.apache.flink.streaming.util.TwoInputStreamOperatorTestHarness; +import org.apache.flink.util.OutputTag; + +import org.junit.jupiter.api.Test; + +import java.util.concurrent.atomic.AtomicBoolean; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Test for {@link ProcessFunction}. */ +public class ProcessFunctionTest { + @Test + void testOneInputStreamProcessFunctionOpenAndCloseInvoked() throws Exception { + AtomicBoolean openInvoked = new AtomicBoolean(false); + AtomicBoolean closeInvoked = new AtomicBoolean(false); + + OneInputStreamProcessFunction<Integer, Integer> processFunction = + new OneInputStreamProcessFunction<>() { + @Override + public void open() throws Exception { + openInvoked.set(true); + } + + @Override + public void processRecord( + Integer record, Collector<Integer> output, PartitionedContext ctx) + throws Exception {} + + @Override + public void close() throws Exception { + closeInvoked.set(true); + } + }; + + ProcessOperator<Integer, Integer> processOperator = new ProcessOperator<>(processFunction); + + try (OneInputStreamOperatorTestHarness<Integer, Integer> testHarness = + new OneInputStreamOperatorTestHarness<>(processOperator)) { + testHarness.open(); + assertThat(openInvoked).isTrue(); + testHarness.close(); + assertThat(closeInvoked).isTrue(); + } + } + + @Test + void testTwoInputBroadcastStreamProcessFunctionOpenAndCloseInvoked() throws Exception { + AtomicBoolean openInvoked = new AtomicBoolean(false); + AtomicBoolean closeInvoked = new AtomicBoolean(false); + + TwoInputBroadcastStreamProcessFunction<Integer, Integer, Integer> processFunction = + new TwoInputBroadcastStreamProcessFunction<>() { + + @Override + public void open() throws Exception { + openInvoked.set(true); + } + + @Override + public void processRecordFromNonBroadcastInput( + Integer record, Collector<Integer> output, PartitionedContext ctx) + throws Exception {} + + @Override + public void processRecordFromBroadcastInput( + Integer record, NonPartitionedContext<Integer> ctx) throws Exception {} + + @Override + public void close() throws Exception { + closeInvoked.set(true); + } + }; + + TwoInputBroadcastProcessOperator<Integer, Integer, Integer> processOperator = + new TwoInputBroadcastProcessOperator<>(processFunction); + + try (TwoInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness = + new TwoInputStreamOperatorTestHarness<>(processOperator)) { + testHarness.open(); + assertThat(openInvoked).isTrue(); + testHarness.close(); + assertThat(closeInvoked).isTrue(); + } + } + + @Test + void testTwoInputNonBroadcastStreamProcessFunctionOpenAndCloseInvoked() throws Exception { + AtomicBoolean openInvoked = new AtomicBoolean(false); + AtomicBoolean closeInvoked = new AtomicBoolean(false); + + TwoInputNonBroadcastStreamProcessFunction<Integer, Integer, Integer> processFunction = + new TwoInputNonBroadcastStreamProcessFunction<>() { + + @Override + public void open() throws Exception { + openInvoked.set(true); + } + + @Override + public void processRecordFromFirstInput( + Integer record, Collector<Integer> output, PartitionedContext ctx) + throws Exception {} + + @Override + public void processRecordFromSecondInput( + Integer record, Collector<Integer> output, PartitionedContext ctx) + throws Exception {} + + @Override + public void close() throws Exception { + closeInvoked.set(true); + } + }; + + TwoInputNonBroadcastProcessOperator<Integer, Integer, Integer> processOperator = + new TwoInputNonBroadcastProcessOperator<>(processFunction); + + try (TwoInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness = + new TwoInputStreamOperatorTestHarness<>(processOperator)) { + testHarness.open(); + assertThat(openInvoked).isTrue(); + testHarness.close(); + assertThat(closeInvoked).isTrue(); + } + } + + @Test + void testTwoOutputStreamProcessFunctionOpenAndCloseInvoked() throws Exception { + AtomicBoolean openInvoked = new AtomicBoolean(false); + AtomicBoolean closeInvoked = new AtomicBoolean(false); + + TwoOutputStreamProcessFunction<Integer, Integer, Integer> processFunction = + new TwoOutputStreamProcessFunction<>() { + + @Override + public void open() throws Exception { + openInvoked.set(true); + } + + @Override + public void processRecord( + Integer record, + Collector<Integer> output1, + Collector<Integer> output2, + PartitionedContext ctx) + throws Exception {} + + @Override + public void close() throws Exception { + closeInvoked.set(true); + } + }; + + TwoOutputProcessOperator<Integer, Integer, Integer> processOperator = + new TwoOutputProcessOperator<>( + processFunction, new OutputTag<Integer>("side-output", Types.INT)); + + try (OneInputStreamOperatorTestHarness<Integer, Integer> testHarness = + new OneInputStreamOperatorTestHarness<>(processOperator)) { + testHarness.open(); + assertThat(openInvoked).isTrue(); + testHarness.close(); + assertThat(closeInvoked).isTrue(); + } + } +}
