[FLINK-37553] Make ForStKeyedStateBackend update NamespaceSerializer during restoring state

flink-runtime/src/main/java/org/apache/flink/runtime/state/v2/RegisteredKeyValueStateBackendMetaInfo.java

@@ -138,6 +138,13 @@ public TypeSerializerSchemaCompatibility<S> updateStateSerializer(
         return stateSerializerProvider.registerNewSerializerForRestoredState(newStateSerializer);
     }
 
+    @Nonnull
+    public TypeSerializerSchemaCompatibility<N> updateNamespaceSerializer(
+            TypeSerializer<N> newNamespaceSerializer) {
+        return namespaceSerializerProvider.registerNewSerializerForRestoredState(
+                newNamespaceSerializer);
+    }
+
     @Override
     public boolean equals(Object o) {
         if (this == o) {

flink-state-backends/flink-statebackend-forst/src/main/java/org/apache/flink/state/forst/ForStKeyedStateBackend.java

@@ -445,6 +445,22 @@ private <N, UK, SV> RegisteredKeyValueStateBackendMetaInfo<N, SV> updateRestored
 
         RegisteredKeyValueStateBackendMetaInfo<N, SV> restoredKvStateMetaInfo = oldStateInfo.f1;
 
+        // fetch current namespace serializer now because if it is incompatible, we can't access
+        // it anymore to improve the error message
+        TypeSerializer<N> previousNamespaceSerializer =
+                restoredKvStateMetaInfo.getNamespaceSerializer();
+
+        TypeSerializerSchemaCompatibility<N> s =
+                restoredKvStateMetaInfo.updateNamespaceSerializer(namespaceSerializer);
+        if (s.isCompatibleAfterMigration() || s.isIncompatible()) {
+            throw new StateMigrationException(
+                    "The new namespace serializer ("
+                            + namespaceSerializer
+                            + ") must be compatible with the old namespace serializer ("
+                            + previousNamespaceSerializer
+                            + ").");
+        }
+
         restoredKvStateMetaInfo.checkStateMetaInfo(stateDesc);
 
         // fetch current serializer now because if it is incompatible, we can't access it anymore to

flink-state-backends/flink-statebackend-forst/src/test/java/org/apache/flink/state/forst/ForStStateMigrationTest.java

@@ -265,4 +265,56 @@ void testKryoRestoreResilienceWithDifferentRegistrationOrder(@TempDir File newTm
         }
         snapshot1.get().discardState();
     }
+
+    @Test
+    void testStateNamespaceSerializerChanged() throws Exception {
+        MapStateDescriptor<Integer, String> descriptor =
+                new MapStateDescriptor<>(
+                        "testState", IntSerializer.INSTANCE, StringSerializer.INSTANCE);
+
+        // set the old namespace serializer to IntSerializer.INSTANCE
+        MapState<Integer, String> mapState =
+                keyedBackend.createState(1, IntSerializer.INSTANCE, descriptor);
+
+        RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot =
+                keyedBackend.snapshot(
+                        1L,
+                        System.currentTimeMillis(),
+                        env.getCheckpointStorageAccess()
+                                .resolveCheckpointStorageLocation(
+                                        1L, CheckpointStorageLocationReference.getDefault()),
+                        CheckpointOptions.forCheckpointWithDefaultLocation());
+
+        if (!snapshot.isDone()) {
+            snapshot.run();
+        }
+        SnapshotResult<KeyedStateHandle> snapshotResult = snapshot.get();
+        KeyedStateHandle stateHandle = snapshotResult.getJobManagerOwnedSnapshot();
+        IOUtils.closeQuietly(keyedBackend);
+        keyedBackend.dispose();
+
+        FileSystem.initialize(new Configuration(), null);
+        Configuration configuration = new Configuration();
+        ForStStateBackend forStStateBackend =
+                new ForStStateBackend().configure(configuration, null);
+        keyedBackend =
+                createKeyedStateBackend(
+                        forStStateBackend,
+                        env,
+                        StringSerializer.INSTANCE,
+                        Collections.singletonList(stateHandle));
+        keyedBackend.setup(aec);
+        try {
+
+            // change new NSSerializer to StringSerializer
+            keyedBackend.createState("String", StringSerializer.INSTANCE, descriptor);
+            fail("Expected a state migration exception.");
+        } catch (Exception e) {
+            if (CommonTestUtils.containsCause(e, StateMigrationException.class)) {
+                // StateMigrationException expected
+            } else {
+                throw e;
+            }
+        }
+    }
 }