This is an automated email from the ASF dual-hosted git repository. wenjin272 pushed a commit to branch release-0.3 in repository https://gitbox.apache.org/repos/asf/flink-agents.git
commit 06866c56fd344a7d9ad6a89d6118c265adcb4f8e Author: Weiqing Yang <[email protected]> AuthorDate: Wed Jul 8 02:47:49 2026 -0700 [runtime] Binary serde for MemoryUpdate values via versioned Kryo envelope (#874) (cherry picked from commit 84c76fbc3e0accb1e4c12c522f28c01b1758793f) --- .../runtime/actionstate/ActionStateSerde.java | 119 ++++++++++++++++++ .../runtime/actionstate/ActionStateSerdeTest.java | 137 +++++++++++++++++++++ 2 files changed, 256 insertions(+) diff --git a/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateSerde.java b/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateSerde.java index 90885e31..d38a4495 100644 --- a/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateSerde.java +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateSerde.java @@ -19,14 +19,29 @@ package org.apache.flink.agents.runtime.actionstate; 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 @@ import java.io.IOException; */ 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 @@ public final class ActionStateSerde { 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 final class ActionStateSerde { 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(); + } + } } diff --git a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateSerdeTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateSerdeTest.java index 116aeeb3..f8dbc6fa 100644 --- a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateSerdeTest.java +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateSerdeTest.java @@ -320,6 +320,143 @@ public class ActionStateSerdeTest { "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");
