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 c431efefdb0341abb118e45480d94caea2f44263 Author: Weijie Guo <res...@163.com> AuthorDate: Tue Apr 9 11:58:51 2024 +0800 [FLINK-34549][API] Implement StateManager and expose it via PartitionedContext --- .../datastream/api/context/PartitionedContext.java | 5 +- .../impl/context/DefaultPartitionedContext.java | 16 ++++- .../impl/context/DefaultStateManager.java | 73 ++++++++++++++++++++++ .../impl/operators/KeyedProcessOperator.java | 5 ++ .../KeyedTwoInputBroadcastProcessOperator.java | 5 ++ .../KeyedTwoInputNonBroadcastProcessOperator.java | 5 ++ .../operators/KeyedTwoOutputProcessOperator.java | 5 ++ .../datastream/impl/operators/ProcessOperator.java | 7 ++- .../TwoInputBroadcastProcessOperator.java | 7 ++- .../TwoInputNonBroadcastProcessOperator.java | 7 ++- .../impl/operators/TwoOutputProcessOperator.java | 7 ++- .../impl/context/DefaultStateManagerTest.java | 61 ++++++++++++++++++ 12 files changed, 197 insertions(+), 6 deletions(-) diff --git a/flink-datastream-api/src/main/java/org/apache/flink/datastream/api/context/PartitionedContext.java b/flink-datastream-api/src/main/java/org/apache/flink/datastream/api/context/PartitionedContext.java index 9538fddd38b..e1b59a4a371 100644 --- a/flink-datastream-api/src/main/java/org/apache/flink/datastream/api/context/PartitionedContext.java +++ b/flink-datastream-api/src/main/java/org/apache/flink/datastream/api/context/PartitionedContext.java @@ -25,4 +25,7 @@ import org.apache.flink.annotation.Experimental; * partition-wise execution information, such as getting state, registering timer, etc. */ @Experimental -public interface PartitionedContext extends RuntimeContext {} +public interface PartitionedContext extends RuntimeContext { + /** Get the {@link StateManager} of this process function. */ + StateManager getStateManager(); +} diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultPartitionedContext.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultPartitionedContext.java index e6c4b2d02a9..4ba359e5e66 100644 --- a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultPartitionedContext.java +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultPartitionedContext.java @@ -23,12 +23,21 @@ import org.apache.flink.datastream.api.context.PartitionedContext; import org.apache.flink.datastream.api.context.RuntimeContext; import org.apache.flink.datastream.api.context.TaskInfo; +import java.util.function.Consumer; +import java.util.function.Supplier; + /** The default implementation of {@link PartitionedContext}. */ public class DefaultPartitionedContext implements PartitionedContext { private final RuntimeContext context; - public DefaultPartitionedContext(RuntimeContext context) { + private final DefaultStateManager stateManager; + + public DefaultPartitionedContext( + RuntimeContext context, + Supplier<Object> currentKeySupplier, + Consumer<Object> currentKeySetter) { this.context = context; + this.stateManager = new DefaultStateManager(currentKeySupplier, currentKeySetter); } @Override @@ -40,4 +49,9 @@ public class DefaultPartitionedContext implements PartitionedContext { public TaskInfo getTaskInfo() { return context.getTaskInfo(); } + + @Override + public DefaultStateManager getStateManager() { + return stateManager; + } } diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultStateManager.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultStateManager.java new file mode 100644 index 00000000000..a6ce75361f4 --- /dev/null +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/context/DefaultStateManager.java @@ -0,0 +1,73 @@ +/* + * 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.context; + +import org.apache.flink.datastream.api.context.StateManager; + +import java.util.function.Consumer; +import java.util.function.Supplier; + +/** + * The default implementation of {@link StateManager}. This class supports eagerly set and reset the + * current key. + */ +public class DefaultStateManager implements StateManager { + + /** + * Retrieve the current key. When {@link #currentKeySetter} receives a key, this must return + * that key until it is reset. + */ + private final Supplier<Object> currentKeySupplier; + + private final Consumer<Object> currentKeySetter; + + public DefaultStateManager( + Supplier<Object> currentKeySupplier, Consumer<Object> currentKeySetter) { + this.currentKeySupplier = currentKeySupplier; + this.currentKeySetter = currentKeySetter; + } + + @Override + @SuppressWarnings("unchecked") + public <K> K getCurrentKey() { + return (K) currentKeySupplier.get(); + } + + /** + * This method should be used to run a block of code with a specific key context. The original + * key must be reset after the block is executed. + */ + public void executeInKeyContext(Runnable runnable, Object key) { + final Object oldKey = currentKeySupplier.get(); + setCurrentKey(key); + try { + runnable.run(); + } finally { + resetCurrentKey(oldKey); + } + } + + private void setCurrentKey(Object key) { + currentKeySetter.accept(key); + } + + private void resetCurrentKey(Object oldKey) { + currentKeySetter.accept(oldKey); + } +} diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/KeyedProcessOperator.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/KeyedProcessOperator.java index 92a07cf8623..fcd29667bcd 100644 --- a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/KeyedProcessOperator.java +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/KeyedProcessOperator.java @@ -50,4 +50,9 @@ public class KeyedProcessOperator<KEY, IN, OUT> extends ProcessOperator<IN, OUT> new OutputCollector<>(output), outKeySelector, () -> (KEY) getCurrentKey()) : new OutputCollector<>(output); } + + @Override + protected Object currentKey() { + return getCurrentKey(); + } } diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/KeyedTwoInputBroadcastProcessOperator.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/KeyedTwoInputBroadcastProcessOperator.java index c3a5664f287..f1852121ab0 100644 --- a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/KeyedTwoInputBroadcastProcessOperator.java +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/KeyedTwoInputBroadcastProcessOperator.java @@ -51,4 +51,9 @@ public class KeyedTwoInputBroadcastProcessOperator<KEY, IN1, IN2, OUT> : new KeyCheckedOutputCollector<>( new OutputCollector<>(output), outKeySelector, () -> (KEY) getCurrentKey()); } + + @Override + protected Object currentKey() { + return getCurrentKey(); + } } diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/KeyedTwoInputNonBroadcastProcessOperator.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/KeyedTwoInputNonBroadcastProcessOperator.java index 5e4945cfe19..2b636461f50 100644 --- a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/KeyedTwoInputNonBroadcastProcessOperator.java +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/KeyedTwoInputNonBroadcastProcessOperator.java @@ -53,4 +53,9 @@ public class KeyedTwoInputNonBroadcastProcessOperator<KEY, IN1, IN2, OUT> : new KeyCheckedOutputCollector<>( new OutputCollector<>(output), outKeySelector, () -> (KEY) getCurrentKey()); } + + @Override + protected Object currentKey() { + return getCurrentKey(); + } } diff --git a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/KeyedTwoOutputProcessOperator.java b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/KeyedTwoOutputProcessOperator.java index 76a2941f5a7..b4c4ec786fe 100644 --- a/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/KeyedTwoOutputProcessOperator.java +++ b/flink-datastream/src/main/java/org/apache/flink/datastream/impl/operators/KeyedTwoOutputProcessOperator.java @@ -75,4 +75,9 @@ public class KeyedTwoOutputProcessOperator<KEY, IN, OUT_MAIN, OUT_SIDE> () -> (KEY) getCurrentKey()) : new SideOutputCollector(output); } + + @Override + protected Object currentKey() { + return getCurrentKey(); + } } 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 a53bf41e3a6..50228806f3e 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 @@ -63,7 +63,8 @@ public class ProcessOperator<IN, OUT> taskInfo.getNumberOfParallelSubtasks(), taskInfo.getMaxNumberOfParallelSubtasks(), taskInfo.getTaskName()); - partitionedContext = new DefaultPartitionedContext(context); + partitionedContext = + new DefaultPartitionedContext(context, this::currentKey, this::setCurrentKey); nonPartitionedContext = new DefaultNonPartitionedContext<>(context); outputCollector = getOutputCollector(); } @@ -82,4 +83,8 @@ public class ProcessOperator<IN, OUT> public void endInput() throws Exception { userFunction.endInput(nonPartitionedContext); } + + protected Object currentKey() { + throw new UnsupportedOperationException("The key is only defined for keyed operator"); + } } 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 1486d84e65a..a31f4c624bf 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 @@ -67,7 +67,8 @@ public class TwoInputBroadcastProcessOperator<IN1, IN2, OUT> taskInfo.getNumberOfParallelSubtasks(), taskInfo.getMaxNumberOfParallelSubtasks(), taskInfo.getTaskName()); - this.partitionedContext = new DefaultPartitionedContext(context); + this.partitionedContext = + new DefaultPartitionedContext(context, this::currentKey, this::setCurrentKey); this.nonPartitionedContext = new DefaultNonPartitionedContext<>(context); } @@ -98,4 +99,8 @@ public class TwoInputBroadcastProcessOperator<IN1, IN2, OUT> userFunction.endBroadcastInput(nonPartitionedContext); } } + + protected Object currentKey() { + throw new UnsupportedOperationException("The key is only defined for keyed operator"); + } } 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 93f667d875f..6d34c9daf72 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 @@ -67,7 +67,8 @@ public class TwoInputNonBroadcastProcessOperator<IN1, IN2, OUT> taskInfo.getNumberOfParallelSubtasks(), taskInfo.getMaxNumberOfParallelSubtasks(), taskInfo.getTaskName()); - this.partitionedContext = new DefaultPartitionedContext(context); + this.partitionedContext = + new DefaultPartitionedContext(context, this::currentKey, this::setCurrentKey); this.nonPartitionedContext = new DefaultNonPartitionedContext<>(context); } @@ -98,4 +99,8 @@ public class TwoInputNonBroadcastProcessOperator<IN1, IN2, OUT> userFunction.endSecondInput(nonPartitionedContext); } } + + protected Object currentKey() { + throw new UnsupportedOperationException("The key is only defined for keyed operator"); + } } 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 fff9086d94b..4356bbd78aa 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 @@ -78,7 +78,8 @@ public class TwoOutputProcessOperator<IN, OUT_MAIN, OUT_SIDE> taskInfo.getNumberOfParallelSubtasks(), taskInfo.getMaxNumberOfParallelSubtasks(), taskInfo.getTaskName()); - this.partitionedContext = new DefaultPartitionedContext(context); + this.partitionedContext = + new DefaultPartitionedContext(context, this::currentKey, this::setCurrentKey); this.nonPartitionedContext = new DefaultTwoOutputNonPartitionedContext<>(context); } @@ -103,6 +104,10 @@ public class TwoOutputProcessOperator<IN, OUT_MAIN, OUT_SIDE> return new SideOutputCollector(output); } + protected Object currentKey() { + throw new UnsupportedOperationException("The key is only defined for keyed operator"); + } + /** * This is a special implementation of {@link TimestampCollector} that using side-output * mechanism to emit data. diff --git a/flink-datastream/src/test/java/org/apache/flink/datastream/impl/context/DefaultStateManagerTest.java b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/context/DefaultStateManagerTest.java new file mode 100644 index 00000000000..7528d8c0b16 --- /dev/null +++ b/flink-datastream/src/test/java/org/apache/flink/datastream/impl/context/DefaultStateManagerTest.java @@ -0,0 +1,61 @@ +/* + * 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.context; + +import org.junit.jupiter.api.Test; + +import java.util.concurrent.atomic.AtomicInteger; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** Tests for {@link DefaultStateManager}. */ +class DefaultStateManagerTest { + @Test + void testGetCurrentKey() { + final String key = "key"; + DefaultStateManager stateManager = new DefaultStateManager(() -> key, ignore -> {}); + assertThat((String) stateManager.getCurrentKey()).isEqualTo(key); + } + + @Test + void testErrorInGetCurrentKey() { + DefaultStateManager stateManager = + new DefaultStateManager( + () -> { + throw new RuntimeException("Expected Error"); + }, + ignore -> {}); + assertThatThrownBy(stateManager::getCurrentKey) + .isInstanceOf(RuntimeException.class) + .hasMessageContaining("Expected Error"); + } + + @Test + void testExecuteInKeyContext() { + final int oldKey = 1; + final int newKey = 2; + // -1 as unset value + AtomicInteger setKey = new AtomicInteger(-1); + DefaultStateManager stateManager = + new DefaultStateManager(() -> oldKey, k -> setKey.set((Integer) k)); + stateManager.executeInKeyContext(() -> assertThat(setKey).hasValue(newKey), newKey); + assertThat(setKey).hasValue(oldKey); + } +}