[FLINK-40297][runtime] Route TTL-aware value migration through the migrate hook

flink-core/src/main/java/org/apache/flink/api/common/typeutils/TypeSerializerSnapshot.java

@@ -134,6 +134,41 @@ void readSnapshot(int readVersion, DataInputView in, ClassLoader userCodeClassLo
     TypeSerializerSchemaCompatibility<T> resolveSchemaCompatibility(
             TypeSerializerSnapshot<T> oldSerializerSnapshot);
 
+    /**
+     * Migrates a single state value from the schema described by {@code oldSerializerSnapshot} to
+     * the schema described by this (new) snapshot. Like {@link
+     * #resolveSchemaCompatibility(TypeSerializerSnapshot)}, this is invoked on the new snapshot and
+     * receives the old snapshot as its argument.
+     *
+     * <p>The default implementation returns the value unchanged: a value already deserialized with
+     * the prior serializer is structurally compatible with the current serializer, so the caller
+     * can re-serialize it as-is. A serializer whose in-memory representation is coupled to its
+     * schema should override this to transform the value into the new layout -- for example by
+     * inserting nulls for added fields or reordering fields by name. An implementation may return
+     * the given value or a new instance.
+     *
+     * <p>The migration is not applied recursively to nested serializers. The snapshot of a
+     * composite type returns its value unchanged unless it overrides this method to decompose the
+     * value and migrate each part, so a caller that needs a nested value migrated must reach the
+     * nested snapshot itself. An implementation that does so should not assume that the old and the
+     * new snapshot expose nested snapshots of the same type: restoring may have replaced those of
+     * the old snapshot with decorators, so nested snapshots are best matched by position or name
+     * rather than by class.
+     *
+     * @param oldSerializerSnapshot snapshot of the serializer that wrote the value. A caller that
+     *     holds the snapshot persisted with the state should pass that one in preference to a
+     *     snapshot re-derived from a serializer restored from it, because that round trip does not
+     *     always reproduce the schema that was written.
+     * @param value the value, already deserialized with the prior serializer. It may be {@code
+     *     null} wherever the prior serializer can produce {@code null}. An implementation that
+     *     decomposes a composite value may likewise pass {@code null} to a nested snapshot for an
+     *     absent part, even where that part's serializer would reject {@code null} at top level.
+     * @return the value adapted to the schema of the current serializer.
+     */
+    default T migrate(TypeSerializerSnapshot<T> oldSerializerSnapshot, T value) {
+        return value;
+    }
+
     // ------------------------------------------------------------------------
     //  read / write utilities
     // ------------------------------------------------------------------------

flink-core/src/test/java/org/apache/flink/api/common/typeutils/TypeSerializerSnapshotTest.java

@@ -54,6 +54,15 @@ public TypeSerializerSchemaCompatibility<Integer> resolveSchemaCompatibility(
                 .isTrue();
     }
 
+    @Test
+    void testMigrateReturnsValueUnchangedByDefault() {
+        TypeSerializerSnapshot<Integer> oldSnapshot = new NotCompletedTypeSerializerSnapshot();
+        TypeSerializerSnapshot<Integer> newSnapshot = new NotCompletedTypeSerializerSnapshot();
+        Integer value = 1000;
+
+        assertThat(newSnapshot.migrate(oldSnapshot, value)).isSameAs(value);
+    }
+
     private static class NotCompletedTypeSerializer extends TypeSerializer<Integer> {
 
         @Override

flink-runtime/src/main/java/org/apache/flink/runtime/state/ttl/TtlAwareSerializer.java

@@ -26,6 +26,8 @@
 import org.apache.flink.core.memory.DataOutputView;
 import org.apache.flink.util.function.SupplierWithException;
 
+import javax.annotation.Nullable;
+
 import java.io.IOException;
 import java.util.List;
 import java.util.Map;
@@ -47,6 +49,13 @@ public class TtlAwareSerializer<T, S extends TypeSerializer<T>> extends TypeSeri
 
     private final S typeSerializer;
 
+    /**
+     * Snapshot of {@link #bareValueSerializer()}, computed on first use. {@link
+     * #migrateValueFromPriorSerializer} runs once per migrated state value while the serializer
+     * stays the same, and taking a snapshot allocates one object per nested serializer.
+     */
+    private transient TypeSerializerSnapshot<?> bareValueSerializerSnapshot;
+
     public TtlAwareSerializer(S typeSerializer) {
         checkArgument(
                 !(typeSerializer instanceof TtlAwareSerializer),
@@ -128,31 +137,129 @@ public int hashCode() {
         return Objects.hash(isTtlEnabled, typeSerializer);
     }
 
-    @SuppressWarnings("unchecked")
+    /**
+     * Reads one state value written by {@code priorTtlAwareSerializer}, adapts it to this
+     * serializer's TTL setting and value schema, and writes it to {@code target}.
+     *
+     * <p>The value is unwrapped to its bare form, passed through {@link
+     * TypeSerializerSnapshot#migrate}, and re-wrapped. The hook returns the value unchanged unless
+     * the value serializer overrides it, so a value whose schema did not change is written back
+     * byte for byte.
+     *
+     * @param priorSerializerSnapshot the snapshot persisted with the state for {@code
+     *     priorTtlAwareSerializer}, or {@code null} for a state that carries none.
+     */
+    @SuppressWarnings({"unchecked", "rawtypes"})
     public void migrateValueFromPriorSerializer(
             TtlAwareSerializer<T, ?> priorTtlAwareSerializer,
+            @Nullable TypeSerializerSnapshot<T> priorSerializerSnapshot,
             SupplierWithException<T, IOException> inputSupplier,
             DataOutputView target,
             TtlTimeProvider ttlTimeProvider)
             throws IOException {
+        T priorValue = inputSupplier.get();
+        Object bareValue =
+                priorTtlAwareSerializer.wrapsTtlValue()
+                        ? ((TtlValue<?>) priorValue).getUserValue()
+                        : priorValue;
+
+        TypeSerializerSnapshot newSnapshot = bareValueSerializerSnapshot();
+        Object migratedValue =
+                newSnapshot.migrate(
+                        priorBareValueSerializerSnapshot(
+                                priorTtlAwareSerializer, priorSerializerSnapshot),
+                        bareValue);
+
         T outputRecord;
-        if (this.isTtlEnabled()) {
-            outputRecord =
-                    priorTtlAwareSerializer.isTtlEnabled
-                            ? inputSupplier.get()
-                            : (T)
-                                    new TtlValue<>(
-                                            inputSupplier.get(),
-                                            ttlTimeProvider.currentTimestamp());
+        if (this.wrapsTtlValue()) {
+            // Carrying the prior timestamp over keeps the value's expiry where it was; migration
+            // is not a state access.
+            long lastAccessTimestamp =
+                    priorTtlAwareSerializer.wrapsTtlValue()
+                            ? ((TtlValue<?>) priorValue).getLastAccessTimestamp()
+                            : ttlTimeProvider.currentTimestamp();
+            outputRecord = (T) new TtlValue<>(migratedValue, lastAccessTimestamp);
         } else {
-            outputRecord =
-                    priorTtlAwareSerializer.isTtlEnabled
-                            ? ((TtlValue<T>) inputSupplier.get()).getUserValue()
-                            : inputSupplier.get();
+            outputRecord = (T) migratedValue;
         }
         this.serialize(outputRecord, target);
     }
 
+    /**
+     * The snapshot describing the schema the prior bare value was written with.
+     *
+     * <p>The snapshot persisted with the state is preferred over one re-derived from the prior
+     * serializer, because the prior serializer is itself restored from that snapshot and the round
+     * trip back to a snapshot is not always lossless: a POJO field that no longer exists on the
+     * class returns under a generated placeholder name, which would present a schema that was never
+     * written. Only the absence of a persisted snapshot falls back to the re-derived one: a
+     * persisted snapshot that does not match the prior serializer is an error, not a second reason
+     * to fall back, because re-deriving there would silently reintroduce that lossy round trip.
+     */
+    private static TypeSerializerSnapshot<?> priorBareValueSerializerSnapshot(
+            TtlAwareSerializer<?, ?> priorSerializer,
+            @Nullable TypeSerializerSnapshot<?> priorSerializerSnapshot) {
+        if (priorSerializerSnapshot == null) {
+            return priorSerializer.bareValueSerializerSnapshot();
+        }
+        // TtlAwareSerializerSnapshot is the snapshot counterpart of this class, so the persisted
+        // snapshot carries that layer wherever the serializer carries the wrapper: for a list or
+        // map state it is the element or value snapshot, for a value state the whole snapshot.
+        TypeSerializerSnapshot<?> priorSnapshot =
+                priorSerializerSnapshot instanceof TtlAwareSerializerSnapshot
+                        ? ((TtlAwareSerializerSnapshot<?>) priorSerializerSnapshot)
+                                .getOrinalTypeSerializerSnapshot()
+                        : priorSerializerSnapshot;
+
+        // Thrown rather than checked through Preconditions: this runs once per migrated state
+        // value, so the message must not be built while the check is passing.
+        boolean isTtlSnapshot = priorSnapshot instanceof TtlStateFactory.TtlSerializerSnapshot;
+        if (!priorSerializer.wrapsTtlValue()) {
+            if (isTtlSnapshot) {
+                throw new IllegalArgumentException(
+                        "The prior serializer does not wrap values in TtlValue, but its persisted snapshot is a TtlSerializerSnapshot.");
+            }
+            return priorSnapshot;
+        }
+        if (!isTtlSnapshot) {
+            throw new IllegalArgumentException(
+                    "The prior serializer wraps values in TtlValue, so its persisted snapshot should be a TtlSerializerSnapshot, but was "
+                            + priorSnapshot.getClass().getName()
+                            + ".");
+        }
+        // The persisted snapshot describes the TtlValue envelope, so descend to the user value
+        // the same way bareValueSerializer() descends the serializer.
+        return ((TtlStateFactory.TtlSerializerSnapshot<?>) priorSnapshot)
+                .getValueSerializerSnapshot();
+    }
+
+    private TypeSerializerSnapshot<?> bareValueSerializerSnapshot() {
+        if (bareValueSerializerSnapshot == null) {
+            bareValueSerializerSnapshot = bareValueSerializer().snapshotConfiguration();
+        }
+        return bareValueSerializerSnapshot;
+    }
+
+    /**
+     * The serializer of the bare (non-TTL) value: the user value serializer of a {@link
+     * TtlStateFactory.TtlSerializer}, otherwise the wrapped serializer itself.
+     */
+    private TypeSerializer<?> bareValueSerializer() {
+        return wrapsTtlValue()
+                ? ((TtlStateFactory.TtlSerializer<?>) typeSerializer).getValueSerializer()
+                : typeSerializer;
+    }
+
+    /**
+     * Whether the values this serializer reads and writes are {@link TtlValue} envelopes. Narrower
+     * than {@link #isTtlEnabled()}, which is also true for a list or map serializer whose element
+     * or value serializer is a {@link TtlStateFactory.TtlSerializer}: such a serializer wraps the
+     * collection, not a single {@code TtlValue}.
+     */
+    private boolean wrapsTtlValue() {
+        return typeSerializer instanceof TtlStateFactory.TtlSerializer;
+    }
+
     @Override
     public void copy(DataInputView source, DataOutputView target) throws IOException {
         typeSerializer.copy(source, target);

flink-runtime/src/test/java/org/apache/flink/runtime/state/StateSerializerProviderTest.java

@@ -21,13 +21,18 @@
 import org.apache.flink.api.common.typeutils.TypeSerializer;
 import org.apache.flink.api.common.typeutils.TypeSerializerSchemaCompatibility;
 import org.apache.flink.api.common.typeutils.TypeSerializerSnapshot;
+import org.apache.flink.api.common.typeutils.base.ListSerializer;
+import org.apache.flink.api.common.typeutils.base.ListSerializerSnapshot;
+import org.apache.flink.api.common.typeutils.base.StringSerializer;
 import org.apache.flink.core.memory.DataInputView;
 import org.apache.flink.core.memory.DataOutputView;
+import org.apache.flink.runtime.state.ttl.TtlAwareSerializerSnapshot;
 import org.apache.flink.runtime.testutils.statemigration.TestType;
 
 import org.junit.jupiter.api.Test;
 
 import java.io.IOException;
+import java.util.List;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
@@ -311,6 +316,42 @@ void testEagerlyRegisterIncompatibleSerializer() {
                 .isInstanceOf(IllegalStateException.class);
     }
 
+    // --------------------------------------------------------------------------------
+    //  Tests for the ttl-aware wrapping of the previous serializer snapshot
+    // --------------------------------------------------------------------------------
+
+    /**
+     * Registering a new serializer replaces the nested snapshot of the previous snapshot in place,
+     * so the snapshot a caller still holds no longer reports what the checkpoint wrote: its nested
+     * snapshot becomes a {@link TtlAwareSerializerSnapshot} around the original. Anything that
+     * descends a restored composite snapshot has to expect that layer.
+     */
+    @Test
+    void testRegisterNewSerializerWrapsNestedSnapshotOfPreviousSnapshotInPlace() {
+        ListSerializerSnapshot<String> previousSnapshot =
+                (ListSerializerSnapshot<String>)
+                        new ListSerializer<>(StringSerializer.INSTANCE).snapshotConfiguration();
+        TypeSerializerSnapshot<String> elementSnapshotAsWritten =
+                previousSnapshot.getElementSerializerSnapshot();
+
+        StateSerializerProvider<List<String>> testProvider =
+                StateSerializerProvider.fromPreviousSerializerSnapshot(previousSnapshot);
+        testProvider.registerNewSerializerForRestoredState(
+                new ListSerializer<>(StringSerializer.INSTANCE));
+
+        // The same snapshot instance now reports a different element snapshot than it did above.
+        TypeSerializerSnapshot<String> elementSnapshotAfterRestore =
+                previousSnapshot.getElementSerializerSnapshot();
+        assertThat(elementSnapshotAfterRestore).isNotSameAs(elementSnapshotAsWritten);
+        assertThat(elementSnapshotAfterRestore).isInstanceOf(TtlAwareSerializerSnapshot.class);
+        // The original is carried inside the wrapper, not re-derived: a re-derived snapshot would
+        // be an equal instance of the same class but a different object.
+        assertThat(
+                        ((TtlAwareSerializerSnapshot<String>) elementSnapshotAfterRestore)
+                                .getOrinalTypeSerializerSnapshot())
+                .isSameAs(elementSnapshotAsWritten);
+    }
+
     // --------------------------------------------------------------------------------
     //  Utilities
     // --------------------------------------------------------------------------------