From 6cdb50be9a6da0ff8b2a73460701585e3887959b Mon Sep 17 00:00:00 2001 From: Sangkyoon Nam Date: Wed, 28 Oct 2026 17:35:43 +0900 Subject: [PATCH] [runtime] Add cross-release state migration test for ActionExecutionOperator Generated-by: Claude Code 2.1.312 (Claude Opus 5.5) --- runtime/pom.xml | 45 ++++ .../runtime/migration/FlinkAgentsVersion.java | 69 +++++ .../MigrationTestsSnapshotGenerator.java | 102 ++++++++ .../runtime/migration/SnapshotsGenerator.java | 34 +++ .../ActionExecutionOperatorMigrationTest.java | 239 ++++++++++++++++++ ...migration-at-rest-flink-agents0.4-snapshot | Bin 0 -> 5496 bytes 6 files changed, 489 insertions(+) create mode 100644 runtime/src/test/java/org/apache/flink/agents/runtime/migration/FlinkAgentsVersion.java create mode 100644 runtime/src/test/java/org/apache/flink/agents/runtime/migration/MigrationTestsSnapshotGenerator.java create mode 100644 runtime/src/test/java/org/apache/flink/agents/runtime/migration/SnapshotsGenerator.java create mode 100644 runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorMigrationTest.java create mode 100644 runtime/src/test/resources/action-execution-operator-migration-at-rest-flink-agents0.4-snapshot diff --git a/runtime/pom.xml b/runtime/pom.xml index 96ba564d0..0e1a2cb2e 100644 --- a/runtime/pom.xml +++ b/runtime/pom.xml @@ -246,6 +246,51 @@ under the License. + + + generate-migration-test-data + + org.apache.flink.agents.runtime.operator.ActionExecutionOperatorMigrationTest + + + + + org.apache.maven.plugins + maven-antrun-plugin + + + generate-migration-test-data + package + + run + + + + + + + + + + + + + + + + + + + + + + + java-21 diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/migration/FlinkAgentsVersion.java b/runtime/src/test/java/org/apache/flink/agents/runtime/migration/FlinkAgentsVersion.java new file mode 100644 index 000000000..cc22e549f --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/migration/FlinkAgentsVersion.java @@ -0,0 +1,69 @@ +/* + * 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.agents.runtime.migration; + +import java.util.Arrays; +import java.util.List; +import java.util.regex.Matcher; +import java.util.regex.Pattern; + +/** + * Flink Agents releases that have reference snapshots for migration tests. + * + *

Flink's {@code FlinkVersion} can't name Flink Agents releases, so migration tests are + * parameterized over this enum instead. After a release, add its constant here and generate its + * snapshots on the release branch with the {@code generate-migration-test-data} profile. + */ +public enum FlinkAgentsVersion { + v0_4("0.4"); + + private static final Pattern VERSION_PATTERN = Pattern.compile("v?(\\d+)[._](\\d+)"); + + private final String versionStr; + + FlinkAgentsVersion(String versionStr) { + this.versionStr = versionStr; + } + + @Override + public String toString() { + return versionStr; + } + + /** Returns every version with reference snapshots, oldest first. */ + public static List publishedVersions() { + return Arrays.asList(values()); + } + + /** Parses names like {@code 0.4}, {@code 0_4} or {@code v0_4}. */ + public static FlinkAgentsVersion parse(String name) { + Matcher matcher = VERSION_PATTERN.matcher(name); + if (matcher.matches()) { + String versionStr = matcher.group(1) + "." + matcher.group(2); + for (FlinkAgentsVersion version : values()) { + if (version.versionStr.equals(versionStr)) { + return version; + } + } + } + throw new IllegalArgumentException( + "Unknown Flink Agents version " + + name + + ". Add it to FlinkAgentsVersion before generating its snapshots."); + } +} diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/migration/MigrationTestsSnapshotGenerator.java b/runtime/src/test/java/org/apache/flink/agents/runtime/migration/MigrationTestsSnapshotGenerator.java new file mode 100644 index 000000000..2e419bea4 --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/migration/MigrationTestsSnapshotGenerator.java @@ -0,0 +1,102 @@ +/* + * 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.agents.runtime.migration; + +import java.lang.reflect.Constructor; +import java.lang.reflect.InvocationTargetException; +import java.lang.reflect.Method; +import java.util.ArrayList; +import java.util.List; + +/** + * Writes the reference snapshots of migration tests for one Flink Agents version. + * + *

Run it through the {@code generate-migration-test-data} profile from the repository root, on + * the release branch of the version being generated. The profile runs the generator in the {@code + * runtime} module directory, so the snapshots land in {@code runtime/src/test/resources}: + * + * 

+ * mvn clean package -Pgenerate-migration-test-data -Dgenerate.version=0.4 -DskipTests -pl runtime -am
+ *
+ * + *

Arguments: {@code --version --classes [,...]}. + */ +public class MigrationTestsSnapshotGenerator { + + private static final String USAGE = + "Usage: MigrationTestsSnapshotGenerator --version --classes [,...]"; + + public static void main(String[] args) throws Throwable { + String versionName = null; + String classes = null; + if (args.length % 2 != 0) { + throw new IllegalArgumentException(USAGE); + } + for (int i = 0; i < args.length; i += 2) { + switch (args[i]) { + case "--version": + versionName = args[i + 1]; + break; + case "--classes": + classes = args[i + 1]; + break; + default: + throw new IllegalArgumentException("Unknown argument " + args[i]); + } + } + List classNames = new ArrayList<>(); + if (classes != null) { + for (String className : classes.split(",")) { + if (!className.trim().isEmpty()) { + classNames.add(className.trim()); + } + } + } + if (versionName == null || classNames.isEmpty()) { + throw new IllegalArgumentException(USAGE); + } + + FlinkAgentsVersion version = FlinkAgentsVersion.parse(versionName); + for (String className : classNames) { + generate(Class.forName(className), version); + } + } + + private static void generate(Class testClass, FlinkAgentsVersion version) throws Throwable { + Constructor constructor = testClass.getDeclaredConstructor(); + constructor.setAccessible(true); + Object test = constructor.newInstance(); + int generated = 0; + for (Method method : testClass.getDeclaredMethods()) { + if (!method.isAnnotationPresent(SnapshotsGenerator.class)) { + continue; + } + method.setAccessible(true); + try { + method.invoke(test, version); + } catch (InvocationTargetException e) { + throw e.getCause(); + } + generated++; + } + if (generated == 0) { + throw new IllegalArgumentException( + testClass.getName() + " has no @SnapshotsGenerator method"); + } + } +} diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/migration/SnapshotsGenerator.java b/runtime/src/test/java/org/apache/flink/agents/runtime/migration/SnapshotsGenerator.java new file mode 100644 index 000000000..07432216a --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/migration/SnapshotsGenerator.java @@ -0,0 +1,34 @@ +/* + * 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.agents.runtime.migration; + +import java.lang.annotation.ElementType; +import java.lang.annotation.Retention; +import java.lang.annotation.RetentionPolicy; +import java.lang.annotation.Target; + +/** + * Marks a method that writes the reference snapshots of a migration test. + * + *

The method must take a single {@link FlinkAgentsVersion} argument, the version the snapshots + * are written for. {@link MigrationTestsSnapshotGenerator} calls it on a new instance of the test + * class. + */ +@Target(ElementType.METHOD) +@Retention(RetentionPolicy.RUNTIME) +public @interface SnapshotsGenerator {} diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorMigrationTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorMigrationTest.java new file mode 100644 index 000000000..ea9077483 --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorMigrationTest.java @@ -0,0 +1,239 @@ +/* + * 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.agents.runtime.operator; + +import org.apache.flink.agents.api.Event; +import org.apache.flink.agents.api.InputEvent; +import org.apache.flink.agents.api.OutputEvent; +import org.apache.flink.agents.api.context.MemoryObject; +import org.apache.flink.agents.api.context.RunnerContext; +import org.apache.flink.agents.plan.AgentPlan; +import org.apache.flink.agents.plan.JavaFunction; +import org.apache.flink.agents.plan.actions.Action; +import org.apache.flink.agents.runtime.migration.FlinkAgentsVersion; +import org.apache.flink.agents.runtime.migration.SnapshotsGenerator; +import org.apache.flink.api.common.JobID; +import org.apache.flink.api.common.state.ListState; +import org.apache.flink.api.common.state.ListStateDescriptor; +import org.apache.flink.api.common.state.State; +import org.apache.flink.api.common.state.StateDescriptor; +import org.apache.flink.api.common.state.ValueState; +import org.apache.flink.api.common.state.ValueStateDescriptor; +import org.apache.flink.api.common.typeinfo.TypeInformation; +import org.apache.flink.api.java.functions.KeySelector; +import org.apache.flink.runtime.state.KeyedStateBackend; +import org.apache.flink.runtime.state.VoidNamespace; +import org.apache.flink.runtime.state.VoidNamespaceSerializer; +import org.apache.flink.streaming.api.operators.AbstractStreamOperator; +import org.apache.flink.streaming.runtime.streamrecord.StreamRecord; +import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness; +import org.apache.flink.streaming.util.OperatorSnapshotUtil; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.MethodSource; + +import java.lang.reflect.Field; +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Restores {@link ActionExecutionOperator} snapshots written by earlier Flink Agents releases. + * + *

The snapshots live in {@code src/test/resources} and are written by the {@link + * SnapshotsGenerator} methods on the release branch of each version. The test agent's actions and + * the generated input must stay the same across releases, so a restore failure points at a state or + * serializer change rather than a change in this test. + */ +class ActionExecutionOperatorMigrationTest { + + static Collection versions() { + return FlinkAgentsVersion.publishedVersions(); + } + + /** + * Snapshot taken after every input finished, so no action or input is pending. Inputs are keyed + * by parity: key 1 has seen 1 and 3, key 0 has seen 2. + */ + @SnapshotsGenerator + void writeAtRestSnapshot(FlinkAgentsVersion version) throws Exception { + try (KeyedOneInputStreamOperatorTestHarness harness = harness()) { + harness.open(); + processAndWait(harness, 1L); + processAndWait(harness, 3L); + processAndWait(harness, 2L); + assertThat(outputValues(harness)).containsExactly(1L, 4L, 2L); + + OperatorSnapshotUtil.writeStateHandle( + harness.snapshot(0L, 0L), "src/test/resources/" + atRestSnapshot(version)); + } + } + + @ParameterizedTest(name = "Migration snapshot: {0}") + @MethodSource("versions") + void restoresStateAtRest(FlinkAgentsVersion version) throws Exception { + try (KeyedOneInputStreamOperatorTestHarness harness = harness()) { + harness.initializeState( + OperatorSnapshotUtil.getResourceFilename(atRestSnapshot(version))); + harness.open(); + + AbstractStreamOperator operator = (AbstractStreamOperator) harness.getOperator(); + assertThat(sequenceNumber(operator, 1L)).isEqualTo(1L); + assertThat(sequenceNumber(operator, 0L)).isEqualTo(0L); + for (long key : new long[] {0L, 1L}) { + assertThat(pending(operator, key, OperatorStateManager.ACTION_TASK_STATE_NAME)) + .isEmpty(); + assertThat( + pending( + operator, + key, + OperatorStateManager.PENDING_INPUT_EVENT_STATE_NAME)) + .isEmpty(); + } + assertThat(operatorState(operator, "currentProcessingKeys", Object.class)).isEmpty(); + + // The operator keeps the job identifier from the snapshot instead of taking this + // harness's job ID. The Java runner context doesn't expose it, so read the field. + List storedJobIdentifier = + operatorState(operator, "identifier_state", String.class); + assertThat(storedJobIdentifier).hasSize(1); + assertThat(jobIdentifier(operator)).isEqualTo(storedJobIdentifier.get(0)); + assertThat(JobID.fromHexString(storedJobIdentifier.get(0))) + .isNotEqualTo(harness.getEnvironment().getJobID()); + + // The short-term memory keeps each key's running total, so processing continues from + // it. + processAndWait(harness, 1L); + processAndWait(harness, 2L); + processAndWait(harness, 5L); + assertThat(outputValues(harness)).containsExactly(5L, 4L, 10L); + assertThat(sequenceNumber(operator, 1L)).isEqualTo(3L); + assertThat(sequenceNumber(operator, 0L)).isEqualTo(1L); + } + } + + private static String atRestSnapshot(FlinkAgentsVersion version) { + return "action-execution-operator-migration-at-rest-flink-agents" + version + "-snapshot"; + } + + private static KeyedOneInputStreamOperatorTestHarness harness() + throws Exception { + return new KeyedOneInputStreamOperatorTestHarness<>( + new ActionExecutionOperatorFactory(MigrationTestAgent.agentPlan(), true), + (KeySelector) value -> value % 2, + TypeInformation.of(Long.class)); + } + + private static void processAndWait( + KeyedOneInputStreamOperatorTestHarness harness, long input) + throws Exception { + harness.processElement(new StreamRecord<>(input)); + ((ActionExecutionOperator) harness.getOperator()) + .waitInFlightEventsFinished(); + } + + private static List outputValues( + KeyedOneInputStreamOperatorTestHarness harness) { + List values = new ArrayList<>(); + for (Object record : harness.extractOutputStreamRecords()) { + values.add(((StreamRecord) record).getValue()); + } + harness.getOutput().clear(); + return values; + } + + private static Long sequenceNumber(AbstractStreamOperator operator, long key) + throws Exception { + ValueState state = + keyedState( + operator, + key, + new ValueStateDescriptor<>( + OperatorStateManager.MESSAGE_SEQUENCE_NUMBER_STATE_NAME, + Long.class)); + return state.value(); + } + + private static List pending( + AbstractStreamOperator operator, long key, String stateName) throws Exception { + ListState state = + keyedState( + operator, + key, + new ListStateDescriptor<>(stateName, TypeInformation.of(Object.class))); + List values = new ArrayList<>(); + Iterable stored = state.get(); + if (stored != null) { + stored.forEach(values::add); + } + return values; + } + + private static S keyedState( + AbstractStreamOperator operator, long key, StateDescriptor descriptor) + throws Exception { + KeyedStateBackend backend = operator.getKeyedStateBackend(); + backend.setCurrentKey(key); + return backend.getPartitionedState( + VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, descriptor); + } + + private static String jobIdentifier(AbstractStreamOperator operator) throws Exception { + Field field = ActionExecutionOperator.class.getDeclaredField("jobIdentifier"); + field.setAccessible(true); + return (String) field.get(operator); + } + + private static List operatorState( + AbstractStreamOperator operator, String stateName, Class type) throws Exception { + List values = new ArrayList<>(); + operator.getOperatorStateBackend() + .getUnionListState(new ListStateDescriptor<>(stateName, type)) + .get() + .forEach(values::add); + return values; + } + + /** Agent with one action that keeps a running total per key in short-term memory. */ + public static class MigrationTestAgent { + + private static final String TOTAL = "total"; + + static AgentPlan agentPlan() throws Exception { + Action accumulate = + new Action( + "accumulate", + new JavaFunction( + MigrationTestAgent.class, + "accumulate", + new Class[] {Event.class, RunnerContext.class}), + List.of(InputEvent.EVENT_TYPE)); + return new AgentPlan(Map.of(accumulate.getName(), accumulate)); + } + + public static void accumulate(Event event, RunnerContext context) throws Exception { + long input = (Long) InputEvent.fromEvent(event).getInput(); + MemoryObject memory = context.getShortTermMemory(); + long total = memory.isExist(TOTAL) ? (Long) memory.get(TOTAL).getValue() : 0L; + memory.set(TOTAL, total + input); + context.sendEvent(new OutputEvent(total + input)); + } + } +} diff --git a/runtime/src/test/resources/action-execution-operator-migration-at-rest-flink-agents0.4-snapshot b/runtime/src/test/resources/action-execution-operator-migration-at-rest-flink-agents0.4-snapshot new file mode 100644 index 0000000000000000000000000000000000000000..dc67cdfe9fcb80e56288af7a4c037d40905f4de9 GIT binary patch literal 5496 zcmeHL-EP}96sDY{*svALigwp6W?&cw3I37f_+r4+SZIy-54MsbyD=z=a@6Q=q#VF| zls(8EVXv~AT@AY(wnI6NrOM8lKt+bl1SC<1hligI`JF=%A+(3Ci?k1w3Oe*08YZ6W zQ8wiXPN)Q8aY3p4g`SZ)W1ui%5uu!W;p~Jic-k(Zs@9^qiW`pHZZsOQdZejxTi04L zv1=_CH`Q8=)M3cyVC&Jw!JuI|-t*ze92?efJe^pEWllTYi8b!NJhi&R!L&c@nCRei z0QI!P;hE7pHK!AE+%z5SftEo2)-oB ztL3wU6I?YI!wX^Mn>*_Nj7~H5*Hd|$JJzn?)y_N!76!l?)0xK;%BXYfW6lAJgwD$` zop=AW`5Q08NZI;}rmePYwW%M~?Yi4kb+t)t8y~4hF4n1L6VMk$`gHkR@$WDOj$uMn z-Ko2>+EAOas<&MkOts`%i`1ywZeZMKAoO_%U4dhQ2&0XkxS0sOC`W_*Yi3pLp#$`o z(~w7O(We32Z=fEYnCE8anzz>Z$V7X6W0WGHiV!P` zgo;67zXvHD+8{q|Oy5QU;`a32gRs6vnL|?y3lb`oyv-ucJb9y-GMq$AF>d^lW9xkk z2zqMhzCGLb>M17jj+y)ap6MTkfKOr?Iu}T zM3>LtH<%?BWx;OWzdBU7?N2nn-ZOiXQugV&n6}&Z?+zG5n0lDNAVwBBQ-r~F5$B6H z@5|2kYGsBCL>V!Jh6bBMNr@Hd4p B