[runtime] Binary serde for MemoryUpdate values via versioned Kryo envelope

runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateSerde.java

@@ -19,14 +19,29 @@
 
 import com.fasterxml.jackson.annotation.JsonTypeInfo;
 import com.fasterxml.jackson.core.JsonGenerator;
+import com.fasterxml.jackson.core.JsonParser;
+import com.fasterxml.jackson.databind.DeserializationContext;
+import com.fasterxml.jackson.databind.JsonDeserializer;
+import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.JsonSerializer;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.databind.SerializerProvider;
+import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
+import com.fasterxml.jackson.databind.annotation.JsonSerialize;
 import com.fasterxml.jackson.databind.module.SimpleModule;
 import org.apache.flink.agents.api.Event;
+import org.apache.flink.agents.api.context.MemoryUpdate;
 import org.apache.flink.agents.runtime.operator.ActionTask;
+import org.apache.flink.api.common.serialization.SerializerConfigImpl;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.common.typeutils.TypeSerializer;
+import org.apache.flink.core.memory.DataInputViewStreamWrapper;
+import org.apache.flink.core.memory.DataOutputViewStreamWrapper;
 
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
 import java.io.IOException;
+import java.util.Base64;
 
 /**
  * Backend-agnostic serializer/deserializer for {@link ActionState}.
@@ -37,6 +52,37 @@
  */
 public final class ActionStateSerde {
 
+    /**
+     * Recovery-compat contract for the {@link MemoryUpdate#getValue() MemoryUpdate.value} envelope.
+     *
+     * <p>A non-null value is written on disk as {@code
+     * {"serde":"kryo","version":1,"payload":<b64>}} where {@code payload} is the base64 of the
+     * Flink-native (Kryo) binary of the untyped value. That binary is only guaranteed readable by
+     * the same code and same Flink version that wrote it — which is all that ever happens, because
+     * the durable-execution journal is checkpoint-scoped and never survives a code or Flink
+     * upgrade. Bump {@code version} if the payload encoding ever changes; Kryo is pinned within a
+     * Flink major release (supplied by the cluster).
+     */
+    private static final String MEMORY_UPDATE_SERDE = "kryo";
+
+    private static final int MEMORY_UPDATE_VERSION = 1;
+    private static final String ENVELOPE_SERDE_FIELD = "serde";
+    private static final String ENVELOPE_VERSION_FIELD = "version";
+    private static final String ENVELOPE_PAYLOAD_FIELD = "payload";
+
+    // The Flink-native serializer for an untyped Object resolves to Kryo
+    // (GenericTypeInfo<Object> -> KryoSerializer).
+    private static final TypeSerializer<Object> VALUE_SERIALIZER_TEMPLATE =
+            TypeInformation.of(Object.class).createSerializer(new SerializerConfigImpl());
+
+    // KryoSerializer is not thread-safe; only duplicate() (which is safe) is ever called on the
+    // shared template, giving each thread an independent copy for serialize/deserialize. Because
+    // this ThreadLocal is static, a thread's copy can retain (pin) the classes of the values it has
+    // serialized for that thread's lifetime — an accepted tradeoff over rebuilding a heavy Kryo
+    // instance on every call.
+    private static final ThreadLocal<TypeSerializer<Object>> VALUE_SERIALIZER =
+            ThreadLocal.withInitial(VALUE_SERIALIZER_TEMPLATE::duplicate);
+
     private static final ObjectMapper OBJECT_MAPPER = createObjectMapper();
 
     private ActionStateSerde() {}
@@ -70,9 +116,25 @@ private static ObjectMapper createObjectMapper() {
         module.addSerializer(ActionTask.class, new ActionTaskSerializer());
         mapper.registerModule(module);
 
+        // Preserve the concrete runtime type of the untyped MemoryUpdate.value across durable
+        // recovery by wrapping it in a versioned Kryo envelope.
+        mapper.addMixIn(MemoryUpdate.class, MemoryUpdateValueMixin.class);
+
         return mapper;
     }
 
+    private static byte[] kryoSerialize(Object value) throws IOException {
+        ByteArrayOutputStream baos = new ByteArrayOutputStream();
+        VALUE_SERIALIZER.get().serialize(value, new DataOutputViewStreamWrapper(baos));
+        return baos.toByteArray();
+    }
+
+    private static Object kryoDeserialize(byte[] bytes) throws IOException {
+        return VALUE_SERIALIZER
+                .get()
+                .deserialize(new DataInputViewStreamWrapper(new ByteArrayInputStream(bytes)));
+    }
+
     /** Mixin to add type information for Event hierarchy. */
     @JsonTypeInfo(
             use = JsonTypeInfo.Id.CLASS,
@@ -88,4 +150,61 @@ public void serialize(ActionTask value, JsonGenerator gen, SerializerProvider se
             gen.writeNull();
         }
     }
+
+    /**
+     * Binds the Kryo envelope (de)serializer to the {@code value} property of {@link MemoryUpdate}.
+     */
+    abstract static class MemoryUpdateValueMixin {
+        @JsonSerialize(using = MemoryUpdateValueSerializer.class)
+        @JsonDeserialize(using = MemoryUpdateValueDeserializer.class)
+        Object value;
+    }
+
+    /**
+     * Writes a non-null {@link MemoryUpdate} value as the {@code {serde,version,payload}} envelope.
+     */
+    static class MemoryUpdateValueSerializer extends JsonSerializer<Object> {
+        @Override
+        public void serialize(Object value, JsonGenerator gen, SerializerProvider provider)
+                throws IOException {
+            gen.writeStartObject();
+            gen.writeStringField(ENVELOPE_SERDE_FIELD, MEMORY_UPDATE_SERDE);
+            gen.writeNumberField(ENVELOPE_VERSION_FIELD, MEMORY_UPDATE_VERSION);
+            gen.writeStringField(
+                    ENVELOPE_PAYLOAD_FIELD,
+                    Base64.getEncoder().encodeToString(kryoSerialize(value)));
+            gen.writeEndObject();
+        }
+    }
+
+    /** Reads the Kryo envelope back to the concrete value and version-checks it. */
+    static class MemoryUpdateValueDeserializer extends JsonDeserializer<Object> {
+        @Override
+        public Object deserialize(JsonParser p, DeserializationContext ctxt) throws IOException {
+            JsonNode node = ctxt.readTree(p);
+            // Every non-null value this serializer writes is an envelope, and the durable journal
+            // never outlives the code that wrote it (see the recovery-compat contract above), so a
+            // non-envelope node is an unsupported cross-version restore rather than legacy data to
+            // tolerate. Reject it instead of silently binding it to a wrong-typed stock value.
+            if (!isEnvelope(node)) {
+                throw new IOException(
+                        "Expected a Kryo-envelope MemoryUpdate value but found: "
+                                + node.getNodeType());
+            }
+            int version = node.get(ENVELOPE_VERSION_FIELD).asInt();
+            if (version != MEMORY_UPDATE_VERSION) {
+                throw new IOException(
+                        "Unsupported MemoryUpdate value envelope version: " + version);
+            }
+            return kryoDeserialize(
+                    Base64.getDecoder().decode(node.get(ENVELOPE_PAYLOAD_FIELD).asText()));
+        }
+
+        private static boolean isEnvelope(JsonNode n) {
+            return n.isObject()
+                    && MEMORY_UPDATE_SERDE.equals(n.path(ENVELOPE_SERDE_FIELD).asText())
+                    && n.has(ENVELOPE_VERSION_FIELD)
+                    && n.path(ENVELOPE_PAYLOAD_FIELD).isTextual();
+        }
+    }
 }

runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateSerdeTest.java

@@ -320,6 +320,143 @@ public void testDeserializeLegacyCallResultWithoutStatus() throws Exception {
                 "exception".getBytes(StandardCharsets.UTF_8), legacyFailure.getExceptionPayload());
     }
 
+    @Test
+    public void testByteArrayMemoryValuePreserved() throws Exception {
+        byte[] value = new byte[] {1, 2, 3, 4, 5};
+        ActionState originalState = new ActionState(new InputEvent("in"));
+        originalState.addShortTermMemoryUpdate(new MemoryUpdate("stm.bytes", value));
+
+        ActionState deserializedState =
+                ActionStateSerde.deserialize(ActionStateSerde.serialize(originalState));
+
+        Object recovered = deserializedState.getShortTermMemoryUpdates().get(0).getValue();
+        assertInstanceOf(byte[].class, recovered);
+        assertArrayEquals(value, (byte[]) recovered);
+    }
+
+    @Test
+    public void testLongMemoryValuePreserved() throws Exception {
+        Long value = 42L;
+        ActionState originalState = new ActionState(new InputEvent("in"));
+        originalState.addShortTermMemoryUpdate(new MemoryUpdate("stm.long", value));
+
+        ActionState deserializedState =
+                ActionStateSerde.deserialize(ActionStateSerde.serialize(originalState));
+
+        Object recovered = deserializedState.getShortTermMemoryUpdates().get(0).getValue();
+        assertInstanceOf(Long.class, recovered);
+        assertEquals(42L, recovered);
+    }
+
+    @Test
+    public void testNestedCollectionMemoryValuePreserved() throws Exception {
+        byte[] nestedBytes = new byte[] {9, 8, 7};
+        List<Object> innerList = new ArrayList<>();
+        innerList.add(nestedBytes);
+        innerList.add(7L);
+        Map<String, Object> value = new HashMap<>();
+        value.put("items", innerList);
+
+        ActionState originalState = new ActionState(new InputEvent("in"));
+        originalState.addShortTermMemoryUpdate(new MemoryUpdate("stm.nested", value));
+
+        ActionState deserializedState =
+                ActionStateSerde.deserialize(ActionStateSerde.serialize(originalState));
+
+        Object recovered = deserializedState.getShortTermMemoryUpdates().get(0).getValue();
+        assertInstanceOf(Map.class, recovered);
+        @SuppressWarnings("unchecked")
+        List<Object> recoveredList = (List<Object>) ((Map<String, Object>) recovered).get("items");
+        assertInstanceOf(byte[].class, recoveredList.get(0));
+        assertArrayEquals(nestedBytes, (byte[]) recoveredList.get(0));
+        assertInstanceOf(Long.class, recoveredList.get(1));
+        assertEquals(7L, recoveredList.get(1));
+    }
+
+    @Test
+    public void testPojoMemoryValuePreserved() throws Exception {
+        MemoryValuePojo value = new MemoryValuePojo("hello", 99);
+        ActionState originalState = new ActionState(new InputEvent("in"));
+        originalState.addShortTermMemoryUpdate(new MemoryUpdate("stm.pojo", value));
+
+        ActionState deserializedState =
+                ActionStateSerde.deserialize(ActionStateSerde.serialize(originalState));
+
+        Object recovered = deserializedState.getShortTermMemoryUpdates().get(0).getValue();
+        assertInstanceOf(MemoryValuePojo.class, recovered);
+        assertEquals(value, recovered);
+    }
+
+    @Test
+    public void testNullMemoryValuePreserved() throws Exception {
+        ActionState originalState = new ActionState(new InputEvent("in"));
+        originalState.addShortTermMemoryUpdate(new MemoryUpdate("stm.null", null));
+
+        ActionState deserializedState =
+                ActionStateSerde.deserialize(ActionStateSerde.serialize(originalState));
+
+        assertEquals(1, deserializedState.getShortTermMemoryUpdates().size());
+        assertNull(deserializedState.getShortTermMemoryUpdates().get(0).getValue());
+    }
+
+    @Test
+    public void testNonEnvelopeMemoryValueRejected() throws Exception {
+        ActionState originalState = new ActionState(new InputEvent("in"));
+        originalState.addShortTermMemoryUpdate(new MemoryUpdate("stm.raw", "some value"));
+
+        byte[] serialized = ActionStateSerde.serialize(originalState);
+        String json = new String(serialized, StandardCharsets.UTF_8);
+        // Replace the whole envelope object with a bare JSON value, as a legacy pre-envelope
+        // journal would have stored it.
+        int start = json.indexOf("{\"serde\":");
+        int end = json.indexOf('}', start) + 1;
+        byte[] patched =
+                (json.substring(0, start) + "\"some value\"" + json.substring(end))
+                        .getBytes(StandardCharsets.UTF_8);
+
+        assertThrows(RuntimeException.class, () -> ActionStateSerde.deserialize(patched));
+    }
+
+    @Test
+    public void testUnknownEnvelopeVersionRejected() throws Exception {
+        ActionState originalState = new ActionState(new InputEvent("in"));
+        originalState.addShortTermMemoryUpdate(new MemoryUpdate("stm.version", "some value"));
+
+        byte[] serialized = ActionStateSerde.serialize(originalState);
+        String json = new String(serialized, StandardCharsets.UTF_8);
+        assertTrue(json.contains("\"version\":1"));
+        byte[] patched =
+                json.replace("\"version\":1", "\"version\":2").getBytes(StandardCharsets.UTF_8);
+
+        assertThrows(RuntimeException.class, () -> ActionStateSerde.deserialize(patched));
+    }
+
+    /** Serializable POJO used to verify Kryo preserves user types across durable recovery. */
+    public static class MemoryValuePojo implements java.io.Serializable {
+        private String name;
+        private int count;
+
+        public MemoryValuePojo() {}
+
+        public MemoryValuePojo(String name, int count) {
+            this.name = name;
+            this.count = count;
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) return true;
+            if (!(o instanceof MemoryValuePojo)) return false;
+            MemoryValuePojo that = (MemoryValuePojo) o;
+            return count == that.count && java.util.Objects.equals(name, that.name);
+        }
+
+        @Override
+        public int hashCode() {
+            return java.util.Objects.hash(name, count);
+        }
+    }
+
     @Test
     public void testKafkaSederDelegatesToActionStateSerde() throws Exception {
         InputEvent inputEvent = new InputEvent("test delegation");