Skip to content

[runtime] Preserve byte[] MemoryUpdate values across durable-execution serde (#865) #871

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
weiqingy wants to merge 1 commit into from
Closed

[runtime] Preserve byte[] MemoryUpdate values across durable-execution serde (#865) #871

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
75 changes: 75 additions & 0 deletions runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateSerde.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,16 +19,24 @@

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.InputEvent;
import org.apache.flink.agents.api.OutputEvent;
import org.apache.flink.agents.api.context.MemoryUpdate;
import org.apache.flink.agents.runtime.operator.ActionTask;

import java.io.IOException;
import java.util.Base64;

/**
* Backend-agnostic serializer/deserializer for {@link ActionState}.
Expand All @@ -39,6 +47,13 @@
*/
public final class ActionStateSerde {

/**
* Reserved envelope key marking a base64-encoded {@code byte[]} {@link MemoryUpdate} value. A
* namespaced key keeps the residual collision with a genuine single-entry user {@code Map}
* negligible and clearly framework-reserved.
*/
private static final String MEMORY_UPDATE_BYTES_KEY = "__flink_agents_bytes__";

private static final ObjectMapper OBJECT_MAPPER = createObjectMapper();

private ActionStateSerde() {}
Expand Down Expand Up @@ -74,6 +89,9 @@ private static ObjectMapper createObjectMapper() {
module.addSerializer(ActionTask.class, new ActionTaskSerializer());
mapper.registerModule(module);

// Preserve byte[] MemoryUpdate values across the untyped Object round-trip
mapper.addMixIn(MemoryUpdate.class, MemoryUpdateValueMixin.class);

return mapper;
}

Expand All @@ -92,4 +110,61 @@ public void serialize(ActionTask value, JsonGenerator gen, SerializerProvider se
gen.writeNull();
}
}

/** Binds the byte[]-preserving (de)serializer to {@link MemoryUpdate#getValue()} only. */
abstract static class MemoryUpdateValueMixin {
@JsonSerialize(using = MemoryUpdateValueSerializer.class)
@JsonDeserialize(using = MemoryUpdateValueDeserializer.class)
Object value;
}

/**
* Serializes a {@code byte[]} {@link MemoryUpdate} value as a one-key base64 envelope so it can
* be recovered as a {@code byte[]} rather than the base64 {@code String} that untyped {@code
* Object} serialization would yield; all other types delegate to default serialization and stay
* byte-identical on disk.
*
* <p>Only a top-level {@code byte[]} value is preserved; a {@code byte[]} nested inside a
* {@code List} or {@code Map} value goes through the default serializers and is not preserved.
*/
static class MemoryUpdateValueSerializer extends JsonSerializer<Object> {
@Override
public void serialize(Object value, JsonGenerator gen, SerializerProvider provider)
throws IOException {
if (value instanceof byte[]) {
gen.writeStartObject();
gen.writeStringField(
MEMORY_UPDATE_BYTES_KEY,
Base64.getEncoder().encodeToString((byte[]) value));
gen.writeEndObject();
} else {
provider.defaultSerializeValue(value, gen);
}
}
}

/**
* Reverses the base64 envelope written by {@link MemoryUpdateValueSerializer} back to a {@code
* byte[]}; every other shape reproduces stock untyped {@code Object} binding.
*
* <p>Only a top-level {@code byte[]} value is recovered; a {@code byte[]} nested inside a
* {@code List} or {@code Map} value is not.
*/
static class MemoryUpdateValueDeserializer extends JsonDeserializer<Object> {
@Override
public Object deserialize(JsonParser p, DeserializationContext ctxt) throws IOException {
JsonNode node = ctxt.readTree(p);
if (node.isObject()
&& node.size() == 1
&& node.has(MEMORY_UPDATE_BYTES_KEY)
&& node.get(MEMORY_UPDATE_BYTES_KEY).isTextual()) {
try {
return Base64.getDecoder().decode(node.get(MEMORY_UPDATE_BYTES_KEY).asText());
} catch (IllegalArgumentException notBase64) {
// Not a real envelope - fall through to generic binding.
}
}
return ctxt.readTreeAsValue(node, Object.class);
}
}
}
110 changes: 110 additions & 0 deletions runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/ActionStateSerdeTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,8 +25,10 @@

import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Base64;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

Expand Down Expand Up @@ -136,6 +138,114 @@ public void testActionStateWithComplexAttributes() throws Exception {
assertEquals(42, deserializedComplexAttr.get("number"));
}

@Test
public void testByteArrayShortTermMemoryValueRoundTrip() throws Exception {
// Verbatim issue repro: a byte[] short-term memory value must survive the durable serde
// round-trip as a byte[], not silently decay to a base64 String.
byte[] original = "hello".getBytes(StandardCharsets.UTF_8);
ActionState state = new ActionState(null, null, null, null, null, false);
state.addShortTermMemoryUpdate(new MemoryUpdate("stm.path", original));

ActionState recovered = ActionStateSerde.deserialize(ActionStateSerde.serialize(state));

Object value = recovered.getShortTermMemoryUpdates().get(0).getValue();
assertInstanceOf(byte[].class, value);
assertArrayEquals(original, (byte[]) value);
}

@Test
public void testByteArraySensoryMemoryValueRoundTrip() throws Exception {
byte[] original = "hello".getBytes(StandardCharsets.UTF_8);
ActionState state = new ActionState(null, null, null, null, null, false);
state.addSensoryMemoryUpdate(new MemoryUpdate("sm.path", original));

ActionState recovered = ActionStateSerde.deserialize(ActionStateSerde.serialize(state));

Object value = recovered.getSensoryMemoryUpdates().get(0).getValue();
assertInstanceOf(byte[].class, value);
assertArrayEquals(original, (byte[]) value);
}

@Test
public void testEmptyByteArrayValueRoundTrip() throws Exception {
byte[] original = new byte[0];
ActionState state = new ActionState(null, null, null, null, null, false);
state.addShortTermMemoryUpdate(new MemoryUpdate("stm.path", original));

ActionState recovered = ActionStateSerde.deserialize(ActionStateSerde.serialize(state));

Object value = recovered.getShortTermMemoryUpdates().get(0).getValue();
assertInstanceOf(byte[].class, value);
assertEquals(0, ((byte[]) value).length);
}

@Test
public void testStringMemoryValueUnchanged() throws Exception {
// Guard: the envelope must not disturb the common String case.
ActionState state = new ActionState(null, null, null, null, null, false);
state.addShortTermMemoryUpdate(new MemoryUpdate("stm.path", "plain string value"));

ActionState recovered = ActionStateSerde.deserialize(ActionStateSerde.serialize(state));

Object value = recovered.getShortTermMemoryUpdates().get(0).getValue();
assertInstanceOf(String.class, value);
assertEquals("plain string value", value);
}

@Test
public void testMapMemoryValueUnchanged() throws Exception {
// Guard: a Map value round-trips as a Map, not misread as a byte[] envelope.
Map<String, Object> original = new LinkedHashMap<>();
original.put("a", 1);
original.put("b", Arrays.asList(1, 2, 3));
ActionState state = new ActionState(null, null, null, null, null, false);
state.addShortTermMemoryUpdate(new MemoryUpdate("stm.path", original));

ActionState recovered = ActionStateSerde.deserialize(ActionStateSerde.serialize(state));

Object value = recovered.getShortTermMemoryUpdates().get(0).getValue();
assertInstanceOf(Map.class, value);
@SuppressWarnings("unchecked")
Map<String, Object> recoveredMap = (Map<String, Object>) value;
assertEquals(1, recoveredMap.get("a"));
assertEquals(Arrays.asList(1, 2, 3), recoveredMap.get("b"));
}

@Test
public void testReservedKeyMapNotByteArrayIsDocumentedResidual() throws Exception {
// Documented residual, not a bug: because the reserved key is the collision guard, a user
// Map that is EXACTLY {"__flink_agents_bytes__":"<valid-base64>"} is indistinguishable from
// a byte[] envelope and recovers as byte[]. A namespaced key makes this negligible. A map
// with any second key falls outside the single-key predicate and stays a Map.
String reservedKey = "__flink_agents_bytes__";

Map<String, Object> collidingMap = new LinkedHashMap<>();
collidingMap.put(reservedKey, "aGVsbG8="); // base64 of "hello"
ActionState collidingState = new ActionState(null, null, null, null, null, false);
collidingState.addShortTermMemoryUpdate(new MemoryUpdate("stm.path", collidingMap));

ActionState recoveredColliding =
ActionStateSerde.deserialize(ActionStateSerde.serialize(collidingState));
Object collidingValue = recoveredColliding.getShortTermMemoryUpdates().get(0).getValue();
assertInstanceOf(byte[].class, collidingValue);
assertArrayEquals("hello".getBytes(StandardCharsets.UTF_8), (byte[]) collidingValue);

Map<String, Object> twoKeyMap = new LinkedHashMap<>();
twoKeyMap.put(reservedKey, "aGVsbG8=");
twoKeyMap.put("x", 1);
ActionState twoKeyState = new ActionState(null, null, null, null, null, false);
twoKeyState.addShortTermMemoryUpdate(new MemoryUpdate("stm.path", twoKeyMap));

ActionState recoveredTwoKey =
ActionStateSerde.deserialize(ActionStateSerde.serialize(twoKeyState));
Object twoKeyValue = recoveredTwoKey.getShortTermMemoryUpdates().get(0).getValue();
assertInstanceOf(Map.class, twoKeyValue);
@SuppressWarnings("unchecked")
Map<String, Object> recoveredTwoKeyMap = (Map<String, Object>) twoKeyValue;
assertEquals("aGVsbG8=", recoveredTwoKeyMap.get(reservedKey));
assertEquals(1, recoveredTwoKeyMap.get("x"));
}

@Test
public void testActionStateWithCallResults() throws Exception {
// Create ActionState with call results
Expand Down
Loading