wenjin272 commented on code in PR #546: URL: https://github.com/apache/flink-agents/pull/546#discussion_r3165421647
########## runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java: ########## @@ -121,54 +72,25 @@ * and the resulting output event is collected for further processing. */ public class ActionExecutionOperator<IN, OUT> extends AbstractStreamOperator<OUT> Review Comment: They are initialized in the Open method of ActionExecutionOperator, and all calls occur after Open. ########## runtime/src/main/java/org/apache/flink/agents/runtime/operator/DurableExecutionManager.java: ########## @@ -0,0 +1,351 @@ +/* + * 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.operator; + +import org.apache.flink.agents.api.Event; +import org.apache.flink.agents.api.context.MemoryUpdate; +import org.apache.flink.agents.plan.AgentConfiguration; +import org.apache.flink.agents.plan.actions.Action; +import org.apache.flink.agents.runtime.actionstate.ActionState; +import org.apache.flink.agents.runtime.actionstate.ActionStateStore; +import org.apache.flink.agents.runtime.actionstate.KafkaActionStateStore; +import org.apache.flink.agents.runtime.context.ActionStatePersister; +import org.apache.flink.agents.runtime.context.RunnerContextImpl; +import org.apache.flink.annotation.VisibleForTesting; +import org.apache.flink.api.common.state.ListState; +import org.apache.flink.api.common.state.ListStateDescriptor; +import org.apache.flink.api.common.typeinfo.TypeInformation; +import org.apache.flink.runtime.state.OperatorStateBackend; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import javax.annotation.Nullable; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import static org.apache.flink.agents.api.configuration.AgentConfigOptions.ACTION_STATE_STORE_BACKEND; +import static org.apache.flink.agents.runtime.actionstate.ActionStateStore.BackendType.KAFKA; + +/** + * Owns the durable-execution side of {@link ActionExecutionOperator}: the optional {@link + * ActionStateStore}, the recovery-marker operator state, the per-checkpoint sequence-number map, + * and the per-{@link ActionTask} {@link RunnerContextImpl.DurableExecutionContext} map. + * + * <p>Durable mode is optional. If no {@link ActionStateStore} is configured (and none is + * pre-injected via the constructor), all {@code maybe*} methods are no-ops and {@link + * #hasDurableStore()} returns {@code false}. This lets the operator stay agnostic of whether + * durable execution is enabled. + * + * <p>Lifecycle: instantiated in the operator constructor. {@link + * #maybeInitActionStateStore(AgentConfiguration)} runs from BOTH the operator's {@code + * initializeState()} and {@code open()} — recovery requires the store to be configured before + * {@link #handleRecovery(OperatorStateBackend)} reads from it, and the {@code open()} call ensures + * the store is also available on the normal (non-recovery) path. The method creates a default + * Kafka-backed store when one was not pre-injected, and is idempotent on the second call. {@link + * #handleRecovery(OperatorStateBackend)} runs from the operator's {@code initializeState()} during + * recovery. {@link #initRecoveryMarkerState(OperatorStateBackend)} runs from the operator's {@code + * open()}. {@link #close()} closes the underlying store. + * + * <p>Design constraint: package-private; no manager-to-manager held references. Cross-cutting data + * flows via method parameters. In particular, {@link + * ActionTaskContextManager#transferContexts(ActionTask, ActionTask, DurableExecutionManager)} + * accepts this manager as a parameter so that the durable-context map can stay encapsulated here. + */ +class DurableExecutionManager implements ActionStatePersister, AutoCloseable { + + private static final Logger LOG = LoggerFactory.getLogger(DurableExecutionManager.class); + + private static final String RECOVERY_MARKER_STATE_NAME = "recoveryMarker"; + + private ActionStateStore actionStateStore; + private transient ListState<Object> recoveryMarkerOpState; + private final Map<Long, Map<Object, Long>> checkpointIdToSeqNums; + + private final Map<ActionTask, RunnerContextImpl.DurableExecutionContext> + actionTaskDurableContexts; + + /** + * @param actionStateStore an optional pre-injected store, primarily for tests. When {@code + * null}, {@link #maybeInitActionStateStore(AgentConfiguration)} may create a default store + * based on configuration; otherwise durable execution is disabled. + */ + DurableExecutionManager(@Nullable ActionStateStore actionStateStore) { + this.actionStateStore = actionStateStore; + this.checkpointIdToSeqNums = new HashMap<>(); + this.actionTaskDurableContexts = new HashMap<>(); + } + + /** + * Lazily creates a default {@link ActionStateStore} from configuration if none was + * pre-injected. + * + * <p>Only creates a store when this manager was constructed without one and the configuration + * selects a recognized backend (currently Kafka). Otherwise this is a no-op, which leaves + * durable execution disabled. + * + * @param config the agent configuration carrying the backend selection. + */ + void maybeInitActionStateStore(AgentConfiguration config) { + if (actionStateStore == null + && KAFKA.getType().equalsIgnoreCase(config.get(ACTION_STATE_STORE_BACKEND))) { + LOG.info("Using Kafka as backend of action state store."); + actionStateStore = new KafkaActionStateStore(config); + } + } + + boolean hasDurableStore() { + return actionStateStore != null; + } + + void initRecoveryMarkerState(OperatorStateBackend operatorStateBackend) throws Exception { + if (actionStateStore != null) { + recoveryMarkerOpState = + operatorStateBackend.getUnionListState( + new ListStateDescriptor<>( + RECOVERY_MARKER_STATE_NAME, TypeInformation.of(Object.class))); + } + } + + /** + * Replays recovery markers from the operator's union-list state to rebuild durable action + * state. + * + * <p>Called from the operator's {@code initializeState()}, which runs before {@code open()}. + * This means {@link #recoveryMarkerOpState} is not yet initialized, so the union-list state + * descriptor is re-created here using the same descriptor name — Flink returns the same + * underlying state. No-op when durable execution is disabled. + * + * @param operatorStateBackend the operator state backend used to obtain the recovery-marker + * union-list state. + */ + void handleRecovery(OperatorStateBackend operatorStateBackend) throws Exception { + if (actionStateStore != null) { + List<Object> markers = new ArrayList<>(); + ListState<Object> markerState = + operatorStateBackend.getUnionListState( + new ListStateDescriptor<>( + RECOVERY_MARKER_STATE_NAME, TypeInformation.of(Object.class))); + Iterable<Object> recoveryMarkers = markerState.get(); + if (recoveryMarkers != null) { + recoveryMarkers.forEach(markers::add); + } + LOG.info("Rebuilding action state from {} recovery markers", markers.size()); + actionStateStore.rebuildState(markers); + } + } + + @Nullable + ActionState maybeGetActionState(Object key, long sequenceNum, Action action, Event event) + throws Exception { + return actionStateStore == null + ? null + : actionStateStore.get(key.toString(), sequenceNum, action, event); + } + + void maybeInitActionState(Object key, long sequenceNum, Action action, Event event) + throws Exception { + if (actionStateStore != null) { + if (actionStateStore.get(key, sequenceNum, action, event) == null) { + actionStateStore.put(key, sequenceNum, action, event, new ActionState(event)); + } + } + } + + /** + * Persists the result of a finished {@link ActionTask} to the durable store. + * + * <p>No-op when no store is configured or when the task did not finish (e.g. it suspended on a + * continuation). On finish, accumulates the task's memory updates and emitted output events + * into the {@link ActionState}, marks it completed, persists it, and clears the in-context + * durable bookkeeping. + * + * @param key the key under which the action ran. + * @param sequenceNum the per-key message sequence number. + * @param action the action being persisted. + * @param event the input event that triggered this action. + * @param context the runner context whose memory updates will be folded into the action state. + * @param actionTaskResult the result of running the action task. + */ + void maybePersistTaskResult( + Object key, + long sequenceNum, + Action action, + Event event, + RunnerContextImpl context, + ActionTask.ActionTaskResult actionTaskResult) + throws Exception { + if (actionStateStore == null) { + return; + } + + if (!actionTaskResult.isFinished()) { + return; + } + + ActionState actionState = actionStateStore.get(key, sequenceNum, action, event); + + for (MemoryUpdate memoryUpdate : context.getSensoryMemoryUpdates()) { + actionState.addSensoryMemoryUpdate(memoryUpdate); + } + + for (MemoryUpdate memoryUpdate : context.getShortTermMemoryUpdates()) { + actionState.addShortTermMemoryUpdate(memoryUpdate); + } + + for (Event outputEvent : actionTaskResult.getOutputEvents()) { + actionState.addEvent(outputEvent); + } + + actionState.markCompleted(); + + actionStateStore.put(key, sequenceNum, action, event, actionState); + + context.clearDurableExecutionContext(); + } + + /** + * Wires a {@link RunnerContextImpl.DurableExecutionContext} onto the given action task's runner + * context. + * + * <p>Returns immediately when no durable store is configured. Otherwise reuses an existing + * {@link RunnerContextImpl.DurableExecutionContext} held in the per-task map (i.e. when + * resuming a continuation), or creates a fresh one bound to this manager so that nested + * persists route back through {@link #persist}. + * + * @param actionTask the action task to attach the context to. + * @param actionState the action state for this (key, sequenceNum, action, event). + * @param seqNum the per-key sequence number. + */ + void setupDurableExecutionContext(ActionTask actionTask, ActionState actionState, long seqNum) { + if (actionStateStore == null) { + return; + } + + RunnerContextImpl.DurableExecutionContext durableContext; + if (actionTaskDurableContexts.containsKey(actionTask)) { + durableContext = actionTaskDurableContexts.get(actionTask); + } else { + durableContext = + new RunnerContextImpl.DurableExecutionContext( + actionTask.getKey(), + seqNum, + actionTask.action, + actionTask.event, + actionState, + this); + } + + actionTask.getRunnerContext().setDurableExecutionContext(durableContext); + } + + @Override + public void persist( + Object key, long sequenceNumber, Action action, Event event, ActionState actionState) { + try { + actionStateStore.put(key, sequenceNumber, action, event, actionState); + } catch (Exception e) { + LOG.error("Failed to persist ActionState", e); + throw new RuntimeException("Failed to persist ActionState", e); + } + } + + void maybePruneState(Object key, long sequenceNum) throws Exception { + if (actionStateStore != null) { + actionStateStore.pruneState(key, sequenceNum); + } + } + + /** + * Prunes durable state for all per-key sequence numbers that were captured at the time of the + * given checkpoint. + * + * <p>The mapping from checkpoint id to per-key sequence numbers must have been recorded earlier + * via {@link #recordCheckpointSequenceNumbers}. After pruning, the entry for that checkpoint is + * removed. No-op when durable execution is disabled. + * + * @param checkpointId the id of the completed checkpoint. + */ + void notifyCheckpointComplete(long checkpointId) { + if (actionStateStore != null) { + Map<Object, Long> keyToSeqNum = + checkpointIdToSeqNums.getOrDefault(checkpointId, new HashMap<>()); + for (Map.Entry<Object, Long> entry : keyToSeqNum.entrySet()) { + actionStateStore.pruneState(entry.getKey(), entry.getValue()); + } + checkpointIdToSeqNums.remove(checkpointId); Review Comment: It appears that `checkpointIdToSeqNums` is only cleaned up when `actionStateStore` is not null, but there is no check for whether `actionStateStore` is null in `recordCheckpointSequenceNumbers`. This seems likely to cause a memory leak. I checked, and this issue existed before the refactoring. We can open a separate PR to fix it. ########## runtime/src/main/java/org/apache/flink/agents/runtime/operator/EventRouter.java: ########## @@ -0,0 +1,260 @@ +/* + * 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.operator; + +import org.apache.flink.agents.api.Event; +import org.apache.flink.agents.api.EventContext; +import org.apache.flink.agents.api.InputEvent; +import org.apache.flink.agents.api.OutputEvent; +import org.apache.flink.agents.api.listener.EventListener; +import org.apache.flink.agents.api.logger.EventLogger; +import org.apache.flink.agents.api.logger.EventLoggerConfig; +import org.apache.flink.agents.api.logger.EventLoggerFactory; +import org.apache.flink.agents.api.logger.EventLoggerOpenParams; +import org.apache.flink.agents.plan.AgentPlan; +import org.apache.flink.agents.plan.actions.Action; +import org.apache.flink.agents.runtime.eventlog.FileEventLogger; +import org.apache.flink.agents.runtime.metrics.BuiltInMetrics; +import org.apache.flink.agents.runtime.operator.queue.SegmentedQueue; +import org.apache.flink.agents.runtime.python.event.PythonEvent; +import org.apache.flink.agents.runtime.python.utils.PythonActionExecutor; +import org.apache.flink.agents.runtime.utils.EventUtil; +import org.apache.flink.annotation.VisibleForTesting; +import org.apache.flink.streaming.api.operators.StreamingRuntimeContext; +import org.apache.flink.streaming.api.watermark.Watermark; +import org.apache.flink.streaming.runtime.streamrecord.StreamRecord; +import org.apache.flink.types.Row; + +import javax.annotation.Nullable; + +import java.util.ArrayList; +import java.util.List; + +import static org.apache.flink.agents.api.configuration.AgentConfigOptions.BASE_LOG_DIR; +import static org.apache.flink.agents.api.configuration.AgentConfigOptions.PRETTY_PRINT; +import static org.apache.flink.util.Preconditions.checkState; + +/** + * Handles event-side concerns for {@link ActionExecutionOperator}: input/output transformation + * between Java/Python representations, action lookup against the {@link AgentPlan}, event-logger + * and event-listener notification, and watermark draining via the per-key segment queue. + * + * <p>Owned state: + * + * <ul> + * <li>The {@link EventLogger} created from the agent plan's logging configuration (may be {@code + * null} when logging is disabled). + * <li>The list of registered {@link EventListener}s. + * <li>A reused {@link StreamRecord} used to emit outputs without per-record allocation. + * <li>The {@link SegmentedQueue} that orders watermarks behind in-flight keys so a watermark is + * only emitted once all keys ahead of it have finished. + * <li>The late-bound {@link BuiltInMetrics} provided in {@link #open(BuiltInMetrics)}. + * </ul> + * + * <p>Lifecycle: instantiated in the operator constructor (which decides {@link #inputIsJava}). + * {@link #open(BuiltInMetrics)} runs from the operator's {@code open()} once metrics are available. + * {@link #initEventLogger} also runs from the operator's {@code open()} once the runtime context is + * available (after metrics have been built). {@link #close()} closes the event logger. + * + * <p>Design constraint: package-private; no manager-to-manager held references. + * + * @param <IN> input record type + * @param <OUT> output record type + */ +class EventRouter<IN, OUT> implements AutoCloseable { + + private final boolean inputIsJava; + private final EventLogger eventLogger; + private final List<EventListener> eventListeners; + private StreamRecord<OUT> reusedStreamRecord; + private SegmentedQueue keySegmentQueue; + private BuiltInMetrics builtInMetrics; + + EventRouter(AgentPlan agentPlan, boolean inputIsJava) { + this.inputIsJava = inputIsJava; + this.eventLogger = createEventLogger(agentPlan); + this.eventListeners = new ArrayList<>(); + } + + /** + * Initializes mutable runtime state that depends on metrics being available. + * + * <p>Allocates the reused stream record and the segmented watermark queue, and stores the + * supplied {@link BuiltInMetrics} for use in {@link #notifyEventProcessed(Event)}. Called from + * the operator's {@code open()} once metric groups are constructed. + * + * @param builtInMetrics the operator's built-in metrics handle. + */ + void open(BuiltInMetrics builtInMetrics) { + this.reusedStreamRecord = new StreamRecord<>(null); + this.keySegmentQueue = new SegmentedQueue(); + this.builtInMetrics = builtInMetrics; + } + + void initEventLogger(StreamingRuntimeContext runtimeContext) throws Exception { + if (eventLogger == null) { + return; + } + eventLogger.open(new EventLoggerOpenParams(runtimeContext)); + } + + /** + * Wraps an incoming record into an {@link Event} suitable for action dispatch. + * + * <p>Java pipelines wrap the raw input directly into a Java {@link InputEvent}. Python + * pipelines expect a two-field {@link Row} where the first field is the key and the second is + * the actual payload; the payload is converted to a Python event via the supplied {@link + * PythonActionExecutor}. + * + * @param input the raw input record. + * @param pythonActionExecutor the Python action executor (used only when input originates from + * Python). + * @return the wrapped input event. + */ + @SuppressWarnings("unchecked") + Event wrapToInputEvent(IN input, PythonActionExecutor pythonActionExecutor) { + if (inputIsJava) { + return new InputEvent(input); + } else { + // the input data must originate from Python and be of type Row with two fields — the + // first representing the key, and the second representing the actual data payload. + checkState(input instanceof Row && ((Row) input).getArity() == 2); + return pythonActionExecutor.wrapToInputEvent(((Row) input).getField(1)); + } + } + + /** + * Extracts the downstream output payload from an {@link OutputEvent}. + * + * <p>For a Java {@link OutputEvent}, returns the payload directly. For a Python {@link + * PythonEvent}, delegates to the supplied {@link PythonActionExecutor} to convert the Python + * output back into the Java output type. + * + * @param event the output event (must satisfy {@link EventUtil#isOutputEvent(Event)}). + * @param pythonActionExecutor the Python action executor (used only for Python events). + * @return the typed output payload. + * @throws IllegalStateException if the event is not a recognized output-event type. + */ + @SuppressWarnings("unchecked") + OUT getOutputFromOutputEvent(Event event, PythonActionExecutor pythonActionExecutor) { + checkState(EventUtil.isOutputEvent(event)); + if (event instanceof OutputEvent) { + return (OUT) ((OutputEvent) event).getOutput(); + } else if (event instanceof PythonEvent) { + Object outputFromOutputEvent = + pythonActionExecutor.getOutputFromOutputEvent(((PythonEvent) event).getEvent()); + return (OUT) outputFromOutputEvent; + } else { + throw new IllegalStateException( + "Unsupported event type: " + event.getClass().getName()); + } + } + + List<Action> getActionsTriggeredBy(Event event, AgentPlan agentPlan) { + if (event instanceof PythonEvent) { + return agentPlan.getActionsTriggeredBy(((PythonEvent) event).getEventType()); + } else { + return agentPlan.getActionsTriggeredBy(event.getClass().getName()); + } + } + + /** + * Notifies the configured event sinks (logger, listeners, metrics) that an event was processed. + * + * <p>If event logging is enabled, appends and immediately flushes the event. Then notifies + * every registered {@link EventListener}. Finally increments the {@code eventProcessed} + * built-in metric. The event logger is flushed per call as a temporary measure pending a + * batched flush mechanism. + * + * @param event the event that was just processed. + */ + void notifyEventProcessed(Event event) throws Exception { + EventContext eventContext = new EventContext(event); + if (eventLogger != null) { + // If event logging is enabled, we log the event along with its context. + eventLogger.append(eventContext, event); + // For now, we flush the event logger after each event to ensure immediate logging. + // This is a temporary solution to ensure that events are logged immediately. + // TODO: In the future, we may want to implement a more efficient batching mechanism. + eventLogger.flush(); + } + if (eventListeners != null) { + // Notify all registered event listeners about the event. + for (EventListener listener : eventListeners) { + listener.onEventProcessed(eventContext, event); + } + } + builtInMetrics.markEventProcessed(); + } + + /** + * Drains all watermarks from the segmented queue that are now eligible to be emitted. + * + * <p>A watermark becomes eligible once every key in the segment ahead of it has finished + * processing. This method pops watermarks in order and forwards each to the supplied {@link + * WatermarkEmitter}. + * + * @param watermarkEmitter callback that emits a watermark downstream. + */ + void processEligibleWatermarks(WatermarkEmitter watermarkEmitter) throws Exception { + Watermark mark = keySegmentQueue.popOldestWatermark(); + while (mark != null) { + watermarkEmitter.emit(mark); + mark = keySegmentQueue.popOldestWatermark(); + } + } + + @Nullable + SegmentedQueue getKeySegmentQueue() { Review Comment: Is the `@Nullable` annotation here necessary? For methods like `maybeGetActionState` and `pollNextPendingInputEvent`, I think the Nullable annotation is appropriate: It may return null when called, and the caller will perform a check. But for methods like `getKeySegmentQueue` and `getSensoryMemState`, they are initialized in the `open` method of `ActionExecutionOperator`, and all accesses occur after `open`. -- 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]
