[FLINK-40113][checkpoint] Reject restore when operators on a keyed vertex disagree on max parallelism

[rionmonster]

What is the purpose of the change

When operators are chained into a single vertex, each operator's recorded maximum parallelism is reconciled onto that shared vertex individually during state restore. Since FLINK-31996 relaxed chaining to allow operators with differing maximum parallelism to share a vertex, a keyed operator can be silently restored under a key-group count that is not its own if the chaining topology has changed since the checkpoint — its keyed state is then remapped through an incompatible hash % maxParallelism layout.

Flink already rejects a single operator's incompatible maximum-parallelism change on restore, but never validates that operators sharing a vertex agree with each other, so the per-vertex reconciliation silently adopts whichever value is applied last. This change closes that consistency gap by rejecting such a restore with a clear error instead of restoring keyed state improperly.

Brief change log

• Add checkMaxParallelismAgreement in StateAssignmentOperation, invoked before the per-operator reconciliation in checkParallelismPreconditions.
• On a vertex that carries keyed state, fail with a clear IllegalStateException when its operators recorded differing maximum parallelism; the comparison spans all chained operators, since a non-keyed operator's recorded value can win the reconciliation and misroute a keyed operator's state.
• Add ChainingMaxParallelismStateLossITCase covering both directions (explicit value below and above the chain head's) on the HashMap and RocksDB backends.

Verifying this change

This change added tests and can be verified as follows:

• ChainingMaxParallelismStateLossITCase takes a canonical savepoint with operator chaining disabled (the keyed operator is its own vertex at an explicit maximum parallelism), then restores with chaining enabled (the keyed operator chains under the source head, whose derived maximum parallelism differs). It asserts the restore is rejected with the new error for both directions and both state backends.
• Existing StateAssignmentOperationTest and CheckpointCoordinatorRestoringTest continue to pass, confirming the new check is additive and runs only on restore.

Does this pull request potentially affect one of the following parts:

• Dependencies (does it add or upgrade a dependency): no
• The public API, i.e., is any changed class annotated with @Public(Evolving): no
• The serializers: no
• The runtime per-record code paths (performance sensitive): no
• Anything that affects deployment or recovery: JobManager (and its components), Checkpointing, Kubernetes/Yarn, ZooKeeper: yes — adds a restore-time precondition in checkpoint state assignment, executed once per vertex during recovery.
• The S3 file system connector: no

Documentation

• Does this pull request introduce a new feature? no
• If yes, how is the feature documented? not applicable

---

Was generative AI tooling used to co-author this PR?

• Yes (please specify the tool below)

Generated-by: Claude Opus 4.8 (1M context)

[flinkbot]

CI report:

• 5f94cad Azure: FAILURE

Bot commands

The @flinkbot bot supports the following commands:

• @flinkbot run azure re-run the last Azure build

[pnowojski]

Can this issue be reproduced without reinterpretAsKeyedStream? 🤔 Just normal keyBy followed by a chained/not chained operator?

reinterpretAsKeyedStream is a bit hacky and doesn't represents the normal use cases. It's intention is if partitions in your sources are already keyed by something, so you don't need to keyBy that data stream again. Here it only works because you constrain the parallelism to 1 which is also a bit artificial constraint.

[rionmonster]

Yes, I agree and don't typically don't prefer using it.

I went down that road and tried iterating several times with just a plain keyBy to trigger the issue with very no luck. Per Claude:

keyBy inserts a hash (KeyGroupStreamPartitioner) edge, which breaks the chain, so the keyed operator always ends up as a chain head with its own maxParallelism. There's no second operator on the vertex to disagree with, and on restore it just adopts its own recorded value. To hit this the keyed operator has to be a chained non-head sharing a vertex with something that recorded a different maxParallelism, and reinterpretAsKeyedStream (forward edge) is the only pure-streaming way to get that topology.

I'll try exploring a few other ways to reproduce this issue without it and report back. 👍