weiqingy commented on code in PR #950:
URL: https://github.com/apache/flink-agents/pull/950#discussion_r3890875875
##########
runtime/src/test/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperatorTest.java:
##########
@@ -144,6 +145,52 @@ void testExecuteAgent() throws Exception {
}
}
+ @Test
+ void testJavaEventAttachmentsAreOffloadedAndResolvedBetweenActions()
throws Exception {
+ InMemoryActionStateStore actionStateStore = new
InMemoryActionStateStore(false);
+ AgentPlan agentPlan = TestAgent.getEventAttachmentAgentPlan();
+ try (KeyedOneInputStreamOperatorTestHarness<Long, Long, Object>
testHarness =
+ new KeyedOneInputStreamOperatorTestHarness<>(
+ new ActionExecutionOperatorFactory<>(agentPlan, true,
actionStateStore),
+ (KeySelector<Long, Long>) value -> value,
+ TypeInformation.of(Long.class))) {
+ testHarness.open();
+ ActionExecutionOperator<Long, Object> operator =
+ (ActionExecutionOperator<Long, Object>)
testHarness.getOperator();
+
+ List<Event> eventsAtSendBoundary = new ArrayList<>();
+ operator.getEventRouter()
+ .addEventListener(
+ (context, event) -> {
+ if
(TestAgent.ATTACHMENT_EVENT_TYPE.equals(event.getType())) {
+ eventsAtSendBoundary.add(event);
+ }
+ });
+
+ testHarness.processElement(new StreamRecord<>(1L));
+ operator.waitInFlightEventsFinished();
+
+ assertThat(eventsAtSendBoundary).hasSize(1);
+ Event runtimeEvent = eventsAtSendBoundary.get(0);
+ Object reference =
runtimeEvent.getAttachment(TestAgent.ATTACHMENT_KEY);
+ assertThat(reference).isInstanceOf(MemoryRef.class);
+ List<StreamRecord<Object>> recordOutput =
+ (List<StreamRecord<Object>>) testHarness.getRecordOutput();
+ assertThat(recordOutput).hasSize(1);
+
assertThat(recordOutput.get(0).getValue()).isEqualTo(Map.of("value", 1L));
+
+ ActionState actionState =
+ actionStateStore.get(
+ 1L,
+ 0L,
+
agentPlan.getActions().get("receiveEventAttachment"),
+ runtimeEvent);
+ assertThat(actionState).isNotNull();
+
assertThat(actionState.getTaskEvent().getAttachment(TestAgent.ATTACHMENT_KEY))
+ .isSameAs(reference);
Review Comment:
I think this one can never fail. `actionState.getTaskEvent()` is the same
object as `runtimeEvent`.
`processEvent` passes one `event` instance both to
`eventRouter.notifyEventProcessed` (`ActionExecutionOperator.java:298`) and to
`createActionTask` (`:328`). `ActionTask.event` is `final`
(`ActionTask.java:54`), and the resolved copy goes to a different field,
`JavaActionTask.invocationEvent` (`JavaActionTask.java:45`, assigned at `:75`),
so the task event is never replaced. `new ActionState(event)` keeps that same
reference (`DurableExecutionManager.java:231`), and `InMemoryActionStateStore`
is a plain `HashMap` with no serialization. So this compares `reference` with
itself.
Line 176 is the check that does the work. It already fails if resolution
leaks back into the runtime event.
There is also a shape question. The identity only holds because this store
does not serialize. A Kafka or Fluss store would round-trip the event and break
`isSameAs` even though the contract still holds. `ActionStateSerdeTest:96-97`
uses `assertInstanceOf` plus `assertEquals` for the same idea. Would that fit
better here?
##########
runtime/src/main/java/org/apache/flink/agents/runtime/memory/EventAttachmentUtils.java:
##########
@@ -0,0 +1,146 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.flink.agents.runtime.memory;
+
+import org.apache.flink.agents.api.Event;
+import org.apache.flink.agents.api.OutputEvent;
+import org.apache.flink.agents.api.context.MemoryObject;
+import org.apache.flink.agents.api.context.MemoryRef;
+import org.apache.flink.agents.api.context.RunnerContext;
+import org.apache.flink.api.common.typeutils.TypeSerializer;
+
+import java.nio.charset.StandardCharsets;
+import java.security.MessageDigest;
+import java.security.NoSuchAlgorithmException;
+import java.util.Map;
+import java.util.Objects;
+import java.util.UUID;
+import java.util.stream.Collectors;
+
+/** Stores event attachments in sensory memory while events cross action
boundaries. */
+public final class EventAttachmentUtils {
+
+ private static final String ATTACHMENT_ROOT = "__event_attachments__";
+
+ private EventAttachmentUtils() {}
+
+ /** Stores concrete attachment values and replaces them with
sensory-memory references. */
+ public static void storeEventAttachments(Event event, RunnerContext
context) throws Exception {
+ if (event.getAttachments().isEmpty()) {
+ return;
+ }
+
+ if (OutputEvent.EVENT_TYPE.equals(event.getType())) {
+ String keys =
+ event.getAttachments().keySet().stream()
+ .sorted()
+ .collect(Collectors.joining(", "));
+ throw new IllegalArgumentException(
+ "Output events cannot carry attachments: event_id="
+ + event.getId()
+ + ", event_type="
+ + event.getType()
+ + ", key="
+ + keys);
+ }
+
+ for (Map.Entry<String, Object> entry :
event.getAttachments().entrySet()) {
+ String key = entry.getKey();
+ Object value = entry.getValue();
+ if (value instanceof MemoryRef) {
+ MemoryRef reference = (MemoryRef) value;
+ if
(!MemoryObject.MemoryType.SENSORY.equals(reference.getType())) {
+ throw new IllegalArgumentException(
+ "Event attachments must use sensory memory
references: event_id="
+ + event.getId()
+ + ", event_type="
+ + event.getType()
+ + ", key="
+ + key
+ + ", memory_type="
+ + reference.getType());
+ }
+ continue;
+ }
+
+ MemoryRef reference =
+
context.getSensoryMemory().set(buildAttachmentPath(event.getId(), key), value);
+
+ event.getAttachments().put(key, reference);
+ }
+ }
+
+ /** Returns an Event with sensory-memory references resolved on a copy
when necessary. */
+ public static Event loadEventAttachments(
+ Event event, RunnerContext context, TypeSerializer<Event>
eventSerializer)
+ throws Exception {
+ boolean requiresResolution =
+
event.getAttachments().values().stream().anyMatch(MemoryRef.class::isInstance);
+ if (!requiresResolution) {
+ return event;
+ }
+
+ Event actionEvent =
+ Objects.requireNonNull(
+ eventSerializer,
+ "Event serializer is required to resolve
attachment references.")
+ .copy(event);
+ for (Map.Entry<String, Object> entry :
actionEvent.getAttachments().entrySet()) {
+ Object value = entry.getValue();
+ if (!(value instanceof MemoryRef)) {
+ continue;
+ }
+ MemoryRef reference = (MemoryRef) value;
+
+ MemoryObject attachment =
context.getSensoryMemory().get(reference);
+ if (attachment == null) {
+ throw new IllegalStateException(
Review Comment:
nit: Python covers this branch with
`test_load_rejects_missing_event_attachment`
(`test_event_attachment_utils.py:133`), but I could not find a Java twin.
`EventAttachmentUtilsTest` has seven tests and none reach the null path. Worth
a short one next to `loadsEventAttachments`?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
