This is an automated email from the ASF dual-hosted git repository.
wenjin272 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/flink-agents.git
The following commit(s) were added to refs/heads/main by this push:
new 84c76fbc [runtime] Binary serde for MemoryUpdate values via versioned
Kryo envelope (#874)
84c76fbc is described below
commit 84c76fbc3e0accb1e4c12c522f28c01b1758793f
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)
---
.../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");
