You are viewing a plain text version of this content. The canonical link for it is here.
Posted to issues@flink.apache.org by "letaoj (via GitHub)" <gi...@apache.org> on 2025/09/17 03:26:50 UTC

[PR] [api] Add TaskActionState API to persist task action state [flink-agents]

letaoj opened a new pull request, #181:
URL: https://github.com/apache/flink-agents/pull/181

   <!--
   * Thank you very much for contributing to Flink Agents.
   * Please add the relevant components in the PR title. E.g., [api], [runtime], [java], [python], [hotfix], etc.
   -->
   
   <!-- Please link the PR to the relevant issue(s). Hotfix doesn't need this. -->
   Linked issue: #153 
   
   ### Purpose of change
   
   <!-- What is the purpose of this change? -->
   
   ### Tests
   
   <!-- How is this change verified? -->
   
   ### API
   
   <!-- Does this change touches any public APIs? -->
   
   ### Documentation
   
   <!-- Should this change be covered by the user documentation?-->
   


-- 
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: issues-unsubscribe@flink.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [Feature] Implemenet Kafka-based state store [flink-agents]

Posted by "xintongsong (via GitHub)" <gi...@apache.org>.
xintongsong commented on code in PR #181:
URL: https://github.com/apache/flink-agents/pull/181#discussion_r2371573191


##########
runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java:
##########
@@ -190,9 +190,17 @@ public void open() throws Exception {
                 && KAFKA.getType()
                         .equalsIgnoreCase(agentPlan.getConfig().get(ACTION_STATE_STORE_BACKEND))) {
             LOG.info("Using Kafka as backend of action state store.");
-            actionStateStore = new KafkaActionStateStore();
+            actionStateStore = new KafkaActionStateStore(agentPlan.getConfig());
         }
 
+        // init recovery marker state for recovery marker persistence
+        recoveryMarkerOpState =
+                getOperatorStateBackend()
+                        .getUnionListState(
+                                new ListStateDescriptor<>(
+                                        RECOVERY_MARKER_STATE_NAME,
+                                        TypeInformation.of(Object.class)));

Review Comment:
   Minor: We might not need this if `actionStateStore == null`



##########
runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java:
##########
@@ -350,13 +358,12 @@ private void processActionTaskForKey(Object key) throws Exception {
         long sequenceNumber = sequenceNumberKState.value();
         boolean isFinished;
         List<Event> outputEvents;
-        Optional<ActionTask> generatedActionTaskOpt;
+        Optional<ActionTask> generatedActionTaskOpt = Optional.empty();
         ActionState actionState =
                 maybeGetActionState(key, sequenceNumber, actionTask.action, actionTask.event);
-        if (actionState != null && actionState.getGeneratedActionTask().isEmpty()) {
+        if (actionState != null) {

Review Comment:
   I think we need to either check that action task is finished before calling `maybePersistTaskResult()`, or skip the persisting for unfinished actions inside `maybePersistTaskResult()`.



##########
runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java:
##########
@@ -350,13 +358,12 @@ private void processActionTaskForKey(Object key) throws Exception {
         long sequenceNumber = sequenceNumberKState.value();
         boolean isFinished;
         List<Event> outputEvents;
-        Optional<ActionTask> generatedActionTaskOpt;
+        Optional<ActionTask> generatedActionTaskOpt = Optional.empty();
         ActionState actionState =
                 maybeGetActionState(key, sequenceNumber, actionTask.action, actionTask.event);
-        if (actionState != null && actionState.getGeneratedActionTask().isEmpty()) {
+        if (actionState != null) {

Review Comment:
   I assume this means we expect an incompleted action should not have a action state in the store? How is this guaranteed? Looks like we always call `maybePersistTaskResult` after `actionTask.invoke()`. 



-- 
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: issues-unsubscribe@flink.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [Feature] Implemenet Kafka-based state store [flink-agents]

Posted by "Sxnan (via GitHub)" <gi...@apache.org>.
Sxnan commented on code in PR #181:
URL: https://github.com/apache/flink-agents/pull/181#discussion_r2367358405


##########
runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/KafkaActionStateStore.java:
##########
@@ -33,38 +57,303 @@
  */
 public class KafkaActionStateStore implements ActionStateStore {
 
-    // In memory action state for quick state retrival, this map is only used during recovery
-    private final Map<String, Map<String, ActionState>> keyedActionStates;
+    private static final Duration CONSUMER_POLL_TIMEOUT = Duration.ofMillis(1000);
+    private static final Logger LOG = LoggerFactory.getLogger(KafkaActionStateStore.class);
+    private static final Long DEFAULT_FUTURE_GET_TIMEOUT_MS = 100L;
+
+    private final AgentConfiguration agentConfiguration;
+
+    // In memory action state for quick state retrieval
+    private final Map<String, ActionState> actionStates;
+
+    // Record the lastest sequence number for each key that should be considered as valid
+    private final Map<String, Long> latestKeySeqNum;
+
+    // Kafka producer
+    private final Producer<String, ActionState> producer;
+    // Kafka consumer
+    private final Consumer<String, ActionState> consumer;
+
+    // Kafka topic that stores action states
+    private final String topic;
 
     @VisibleForTesting
-    KafkaActionStateStore(Map<String, Map<String, ActionState>> keyedActionStates) {
-        this.keyedActionStates = keyedActionStates;
+    KafkaActionStateStore(
+            Map<String, ActionState> actionStates,
+            AgentConfiguration agentConfiguration,
+            Producer<String, ActionState> producer,
+            Consumer<String, ActionState> consumer,
+            String topic) {
+        this.actionStates = actionStates;
+        this.producer = producer;
+        this.consumer = consumer;
+        this.topic = topic;
+        this.latestKeySeqNum = new HashMap<>();
+        this.agentConfiguration = agentConfiguration;
     }
 
-    /** Constructs a new KafkaActionStateStore with an empty in-memory action state map. */
-    public KafkaActionStateStore() {
-        this(new HashMap<>());
+    /** Constructs a new KafkaActionStateStore with custom Kafka configuration. */
+    public KafkaActionStateStore(AgentConfiguration agentConfiguration) {
+        this.actionStates = new HashMap<>();
+        this.latestKeySeqNum = new HashMap<>();
+        this.agentConfiguration = agentConfiguration;
+        this.topic = agentConfiguration.get(KAFKA_ACTION_STATE_TOPIC);
+        // create the topic if not exists
+        maybeCreateTopic();

Review Comment:
   Should we support creating the topic if it is not exist, or we require that the topic is created beforehand.



##########
runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateKafkaDeserializer.java:
##########
@@ -0,0 +1,110 @@
+/*
+ * 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.actionstate;
+
+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.runtime.operator.ActionTask;
+import org.apache.flink.shaded.jackson2.com.fasterxml.jackson.annotation.JsonTypeInfo;
+import org.apache.flink.shaded.jackson2.com.fasterxml.jackson.core.JsonParser;
+import org.apache.flink.shaded.jackson2.com.fasterxml.jackson.databind.DeserializationContext;
+import org.apache.flink.shaded.jackson2.com.fasterxml.jackson.databind.JsonDeserializer;
+import org.apache.flink.shaded.jackson2.com.fasterxml.jackson.databind.ObjectMapper;
+import org.apache.flink.shaded.jackson2.com.fasterxml.jackson.databind.module.SimpleModule;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.util.Map;
+
+/**
+ * Kafka deserializer for {@link ActionState}.
+ *
+ * <p>This deserializer handles the deserialization of byte arrays from Kafka back to ActionState
+ * instances. It uses Jackson ObjectMapper with custom deserializers to handle polymorphic Event
+ * types and ensures ActionTask is deserialized as null.
+ */
+public class ActionStateKafkaDeserializer implements Deserializer<ActionState> {
+
+    private static final Logger LOG = LoggerFactory.getLogger(ActionStateKafkaDeserializer.class);
+    private static final ObjectMapper OBJECT_MAPPER = createObjectMapper();
+
+    @Override
+    public void configure(Map<String, ?> configs, boolean isKey) {
+        // No configuration needed
+    }
+
+    @Override
+    public ActionState deserialize(String topic, byte[] data) {
+        if (data == null) {
+            return null;
+        }
+
+        try {
+            return OBJECT_MAPPER.readValue(data, ActionState.class);
+        } catch (Exception e) {
+            LOG.error("Failed to deserialize ActionState for topic: {}", topic, e);
+            throw new RuntimeException("Failed to deserialize ActionState", e);
+        }
+    }
+
+    @Override
+    public void close() {
+        // No resources to close
+    }
+
+    /** Creates and configures the ObjectMapper for ActionState deserialization. */
+    private static ObjectMapper createObjectMapper() {
+        ObjectMapper mapper = new ObjectMapper();
+
+        // Add type information for polymorphic Event deserialization
+        mapper.addMixIn(Event.class, EventTypeInfoMixin.class);
+        mapper.addMixIn(InputEvent.class, EventTypeInfoMixin.class);
+        mapper.addMixIn(OutputEvent.class, EventTypeInfoMixin.class);
+
+        // Create a module for custom deserializers
+        SimpleModule module = new SimpleModule();
+
+        // Custom deserializer for ActionTask - always deserialize as null
+        module.addDeserializer(ActionTask.class, new ActionTaskDeserializer());
+
+        mapper.registerModule(module);
+
+        return mapper;
+    }
+
+    /** Mixin to add type information for Event hierarchy. */
+    @JsonTypeInfo(
+            use = JsonTypeInfo.Id.CLASS,
+            include = JsonTypeInfo.As.PROPERTY,
+            property = "@class")
+    public abstract static class EventTypeInfoMixin {}

Review Comment:
   Same as the one in `ActionStateKafkaSerializer`, can we reuse?



##########
runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateStore.java:
##########
@@ -24,7 +24,7 @@
 import java.util.List;
 
 /** Interface for storing and retrieving the state of actions performed by agents. */
-public interface ActionStateStore {
+public interface ActionStateStore extends AutoCloseable {

Review Comment:
   `close` method is not called. Should it be call in the `ActionExecutionOperator`?



##########
runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/KafkaActionStateStore.java:
##########
@@ -33,38 +57,303 @@
  */
 public class KafkaActionStateStore implements ActionStateStore {
 
-    // In memory action state for quick state retrival, this map is only used during recovery
-    private final Map<String, Map<String, ActionState>> keyedActionStates;
+    private static final Duration CONSUMER_POLL_TIMEOUT = Duration.ofMillis(1000);
+    private static final Logger LOG = LoggerFactory.getLogger(KafkaActionStateStore.class);
+    private static final Long DEFAULT_FUTURE_GET_TIMEOUT_MS = 100L;
+
+    private final AgentConfiguration agentConfiguration;
+
+    // In memory action state for quick state retrieval
+    private final Map<String, ActionState> actionStates;
+
+    // Record the lastest sequence number for each key that should be considered as valid
+    private final Map<String, Long> latestKeySeqNum;
+
+    // Kafka producer
+    private final Producer<String, ActionState> producer;
+    // Kafka consumer
+    private final Consumer<String, ActionState> consumer;
+
+    // Kafka topic that stores action states
+    private final String topic;
 
     @VisibleForTesting
-    KafkaActionStateStore(Map<String, Map<String, ActionState>> keyedActionStates) {
-        this.keyedActionStates = keyedActionStates;
+    KafkaActionStateStore(
+            Map<String, ActionState> actionStates,
+            AgentConfiguration agentConfiguration,
+            Producer<String, ActionState> producer,
+            Consumer<String, ActionState> consumer,
+            String topic) {
+        this.actionStates = actionStates;
+        this.producer = producer;
+        this.consumer = consumer;
+        this.topic = topic;
+        this.latestKeySeqNum = new HashMap<>();
+        this.agentConfiguration = agentConfiguration;
     }
 
-    /** Constructs a new KafkaActionStateStore with an empty in-memory action state map. */
-    public KafkaActionStateStore() {
-        this(new HashMap<>());
+    /** Constructs a new KafkaActionStateStore with custom Kafka configuration. */
+    public KafkaActionStateStore(AgentConfiguration agentConfiguration) {
+        this.actionStates = new HashMap<>();
+        this.latestKeySeqNum = new HashMap<>();
+        this.agentConfiguration = agentConfiguration;
+        this.topic = agentConfiguration.get(KAFKA_ACTION_STATE_TOPIC);
+        // create the topic if not exists
+        maybeCreateTopic();
+        Properties producerProp = createProducerProp();
+        this.producer = new KafkaProducer<>(producerProp);
+        Properties consumerProp = createConsumerProp();
+        this.consumer = new KafkaConsumer<>(consumerProp);
+        consumer.subscribe(Collections.singletonList(topic));
+        LOG.info("Initialized KafkaActionStateStore with topic: {}", topic);
     }
 
     @Override
     public void put(Object key, long seqNum, Action action, Event event, ActionState state)
-            throws IOException {
-        // TODO: Implement me
+            throws Exception {
+        if (producer == null) {
+            LOG.error("Producer is null, cannot put action state to Kafka");
+            return;
+        }
+
+        String stateKey = generateKey(key, seqNum, action, event);
+        try {
+            ProducerRecord<String, ActionState> kafkaRecord =
+                    new ProducerRecord<>(topic, stateKey, state);
+            Future<RecordMetadata> recordFuture = producer.send(kafkaRecord);
+            RecordMetadata recordMetadata =
+                    recordFuture.get(DEFAULT_FUTURE_GET_TIMEOUT_MS, TimeUnit.MILLISECONDS);
+            LOG.debug("Sent action state to Kafka for key: {}", stateKey);
+        } catch (Exception e) {
+            throw new RuntimeException("Failed to send action state to Kafka", e);
+        }
     }
 
     @Override
-    public ActionState get(Object key, long seqNum, Action action, Event event) throws IOException {
-        // TODO: Implement me
-        return null;
+    public ActionState get(Object key, long seqNum, Action action, Event event) throws Exception {
+        String stateKey = generateKey(key, seqNum, action, event);
+
+        boolean hasDivergence = checkDivergence(key.toString(), seqNum);

Review Comment:
   Can you elaborate on what case we have divergence? Why do we need to check?



##########
runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/KafkaActionStateStore.java:
##########
@@ -33,38 +57,303 @@
  */
 public class KafkaActionStateStore implements ActionStateStore {
 
-    // In memory action state for quick state retrival, this map is only used during recovery
-    private final Map<String, Map<String, ActionState>> keyedActionStates;
+    private static final Duration CONSUMER_POLL_TIMEOUT = Duration.ofMillis(1000);
+    private static final Logger LOG = LoggerFactory.getLogger(KafkaActionStateStore.class);
+    private static final Long DEFAULT_FUTURE_GET_TIMEOUT_MS = 100L;
+
+    private final AgentConfiguration agentConfiguration;
+
+    // In memory action state for quick state retrieval
+    private final Map<String, ActionState> actionStates;
+
+    // Record the lastest sequence number for each key that should be considered as valid
+    private final Map<String, Long> latestKeySeqNum;
+
+    // Kafka producer
+    private final Producer<String, ActionState> producer;
+    // Kafka consumer
+    private final Consumer<String, ActionState> consumer;
+
+    // Kafka topic that stores action states
+    private final String topic;
 
     @VisibleForTesting
-    KafkaActionStateStore(Map<String, Map<String, ActionState>> keyedActionStates) {
-        this.keyedActionStates = keyedActionStates;
+    KafkaActionStateStore(
+            Map<String, ActionState> actionStates,
+            AgentConfiguration agentConfiguration,
+            Producer<String, ActionState> producer,
+            Consumer<String, ActionState> consumer,
+            String topic) {
+        this.actionStates = actionStates;
+        this.producer = producer;
+        this.consumer = consumer;
+        this.topic = topic;
+        this.latestKeySeqNum = new HashMap<>();
+        this.agentConfiguration = agentConfiguration;
     }
 
-    /** Constructs a new KafkaActionStateStore with an empty in-memory action state map. */
-    public KafkaActionStateStore() {
-        this(new HashMap<>());
+    /** Constructs a new KafkaActionStateStore with custom Kafka configuration. */
+    public KafkaActionStateStore(AgentConfiguration agentConfiguration) {
+        this.actionStates = new HashMap<>();
+        this.latestKeySeqNum = new HashMap<>();
+        this.agentConfiguration = agentConfiguration;
+        this.topic = agentConfiguration.get(KAFKA_ACTION_STATE_TOPIC);
+        // create the topic if not exists
+        maybeCreateTopic();
+        Properties producerProp = createProducerProp();
+        this.producer = new KafkaProducer<>(producerProp);
+        Properties consumerProp = createConsumerProp();
+        this.consumer = new KafkaConsumer<>(consumerProp);
+        consumer.subscribe(Collections.singletonList(topic));
+        LOG.info("Initialized KafkaActionStateStore with topic: {}", topic);
     }
 
     @Override
     public void put(Object key, long seqNum, Action action, Event event, ActionState state)
-            throws IOException {
-        // TODO: Implement me
+            throws Exception {
+        if (producer == null) {
+            LOG.error("Producer is null, cannot put action state to Kafka");
+            return;
+        }
+
+        String stateKey = generateKey(key, seqNum, action, event);
+        try {
+            ProducerRecord<String, ActionState> kafkaRecord =
+                    new ProducerRecord<>(topic, stateKey, state);
+            Future<RecordMetadata> recordFuture = producer.send(kafkaRecord);
+            RecordMetadata recordMetadata =

Review Comment:
   The `recordMetadata` is unused.



##########
tools/lint.sh:
##########
@@ -67,7 +67,7 @@ install_python_deps() {
   if command -v uv >/dev/null 2>&1; then
     echo "Using uv for dependency management"
     pushd python
-    uv sync --extra lint
+    uv sync --extra lint --index-strategy unsafe-best-match

Review Comment:
   Why do we need this change?



##########
api/src/main/java/org/apache/flink/agents/api/configuration/AgentConfigOptions.java:
##########
@@ -27,4 +27,20 @@ public class AgentConfigOptions {
     /** The config parameter specifies the backend for action state store. */
     public static final ConfigOption<String> ACTION_STATE_STORE_BACKEND =
             new ConfigOption<>("actionStateStoreBackend", String.class, null);
+
+    /** The config parameter specifies the Kafka bootstrap server. */
+    public static final ConfigOption<String> KAFKA_BOOTSTRAP_SERVERS =
+            new ConfigOption<>("kafkaBootstrapServers", String.class, "localhost:9092");
+
+    /** The config parameter specifies the Kafka topic for action state. */
+    public static final ConfigOption<String> KAFKA_ACTION_STATE_TOPIC =
+            new ConfigOption<>("kafkaActionStateTopic", String.class, "flink-agents-action-state");
+
+    /** The config parameter sepcifies the number of partitions for the Kafka action state topic. */
+    public static final ConfigOption<Integer> KAFKA_ACTION_STATE_TOPIC_NUM_PARTITIONS =

Review Comment:
   Why do we need to specify the number of partitions? Should we allow creating Kafka topic during job run?



##########
api/src/main/java/org/apache/flink/agents/api/configuration/AgentConfigOptions.java:
##########
@@ -27,4 +27,20 @@ public class AgentConfigOptions {
     /** The config parameter specifies the backend for action state store. */
     public static final ConfigOption<String> ACTION_STATE_STORE_BACKEND =
             new ConfigOption<>("actionStateStoreBackend", String.class, null);
+
+    /** The config parameter specifies the Kafka bootstrap server. */
+    public static final ConfigOption<String> KAFKA_BOOTSTRAP_SERVERS =
+            new ConfigOption<>("kafkaBootstrapServers", String.class, "localhost:9092");
+
+    /** The config parameter specifies the Kafka topic for action state. */
+    public static final ConfigOption<String> KAFKA_ACTION_STATE_TOPIC =
+            new ConfigOption<>("kafkaActionStateTopic", String.class, "flink-agents-action-state");

Review Comment:
   I think it is better to use `null` as the default value and throw an exception if it is null. Otherwise, it is problematic when multiple jobs use the same default topic.



-- 
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: issues-unsubscribe@flink.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [Feature] Implemenet Kafka-based state store [flink-agents]

Posted by "xintongsong (via GitHub)" <gi...@apache.org>.
xintongsong merged PR #181:
URL: https://github.com/apache/flink-agents/pull/181


-- 
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: issues-unsubscribe@flink.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [Feature] Implemenet Kafka-based state store [flink-agents]

Posted by "letaoj (via GitHub)" <gi...@apache.org>.
letaoj commented on code in PR #181:
URL: https://github.com/apache/flink-agents/pull/181#discussion_r2369836768


##########
api/src/main/java/org/apache/flink/agents/api/configuration/AgentConfigOptions.java:
##########
@@ -27,4 +27,20 @@ public class AgentConfigOptions {
     /** The config parameter specifies the backend for action state store. */
     public static final ConfigOption<String> ACTION_STATE_STORE_BACKEND =
             new ConfigOption<>("actionStateStoreBackend", String.class, null);
+
+    /** The config parameter specifies the Kafka bootstrap server. */
+    public static final ConfigOption<String> KAFKA_BOOTSTRAP_SERVERS =
+            new ConfigOption<>("kafkaBootstrapServers", String.class, "localhost:9092");
+
+    /** The config parameter specifies the Kafka topic for action state. */
+    public static final ConfigOption<String> KAFKA_ACTION_STATE_TOPIC =
+            new ConfigOption<>("kafkaActionStateTopic", String.class, "flink-agents-action-state");
+
+    /** The config parameter sepcifies the number of partitions for the Kafka action state topic. */
+    public static final ConfigOption<Integer> KAFKA_ACTION_STATE_TOPIC_NUM_PARTITIONS =

Review Comment:
   Because we are auto-creating the state table, it will be better to control the number of partitions by the framework. The partition will be use perform the custom partition for kafka (i can probably retrieve it if the topic is already created)



-- 
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: issues-unsubscribe@flink.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [Feature] Implemenet Kafka-based state store [flink-agents]

Posted by "letaoj (via GitHub)" <gi...@apache.org>.
letaoj commented on code in PR #181:
URL: https://github.com/apache/flink-agents/pull/181#discussion_r2370229726


##########
runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateKafkaSeder.java:
##########
@@ -0,0 +1,110 @@
+/*
+ * 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.actionstate;
+
+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.runtime.operator.ActionTask;
+import org.apache.flink.shaded.jackson2.com.fasterxml.jackson.annotation.JsonTypeInfo;
+import org.apache.flink.shaded.jackson2.com.fasterxml.jackson.core.JsonParser;
+import org.apache.flink.shaded.jackson2.com.fasterxml.jackson.databind.DeserializationContext;
+import org.apache.flink.shaded.jackson2.com.fasterxml.jackson.databind.JsonDeserializer;
+import org.apache.flink.shaded.jackson2.com.fasterxml.jackson.databind.ObjectMapper;
+import org.apache.flink.shaded.jackson2.com.fasterxml.jackson.databind.module.SimpleModule;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.util.Map;
+
+/**
+ * Kafka deserializer for {@link ActionState}.
+ *
+ * <p>This deserializer handles the deserialization of byte arrays from Kafka back to ActionState
+ * instances. It uses Jackson ObjectMapper with custom deserializers to handle polymorphic Event
+ * types and ensures ActionTask is deserialized as null.
+ */
+public class ActionStateKafkaDeserializer implements Deserializer<ActionState> {
+
+    private static final Logger LOG = LoggerFactory.getLogger(ActionStateKafkaDeserializer.class);
+    private static final ObjectMapper OBJECT_MAPPER = createObjectMapper();
+
+    @Override
+    public void configure(Map<String, ?> configs, boolean isKey) {
+        // No configuration needed
+    }
+
+    @Override
+    public ActionState deserialize(String topic, byte[] data) {
+        if (data == null) {
+            return null;
+        }
+
+        try {
+            return OBJECT_MAPPER.readValue(data, ActionState.class);
+        } catch (Exception e) {
+            LOG.error("Failed to deserialize ActionState for topic: {}", topic, e);
+            throw new RuntimeException("Failed to deserialize ActionState", e);
+        }
+    }
+
+    @Override
+    public void close() {
+        // No resources to close
+    }
+
+    /** Creates and configures the ObjectMapper for ActionState deserialization. */
+    private static ObjectMapper createObjectMapper() {
+        ObjectMapper mapper = new ObjectMapper();
+
+        // Add type information for polymorphic Event deserialization
+        mapper.addMixIn(Event.class, EventTypeInfoMixin.class);
+        mapper.addMixIn(InputEvent.class, EventTypeInfoMixin.class);
+        mapper.addMixIn(OutputEvent.class, EventTypeInfoMixin.class);
+
+        // Create a module for custom deserializers
+        SimpleModule module = new SimpleModule();
+
+        // Custom deserializer for ActionTask - always deserialize as null
+        module.addDeserializer(ActionTask.class, new ActionTaskDeserializer());
+
+        mapper.registerModule(module);
+
+        return mapper;
+    }
+
+    /** Mixin to add type information for Event hierarchy. */
+    @JsonTypeInfo(
+            use = JsonTypeInfo.Id.CLASS,
+            include = JsonTypeInfo.As.PROPERTY,
+            property = "@class")
+    public abstract static class EventTypeInfoMixin {}

Review Comment:
   Sure, I combined the two class together. This class is a abstract class, i do not want to cross reference the class in two different classes...



##########
api/src/main/java/org/apache/flink/agents/api/configuration/AgentConfigOptions.java:
##########
@@ -27,4 +27,20 @@ public class AgentConfigOptions {
     /** The config parameter specifies the backend for action state store. */
     public static final ConfigOption<String> ACTION_STATE_STORE_BACKEND =
             new ConfigOption<>("actionStateStoreBackend", String.class, null);
+
+    /** The config parameter specifies the Kafka bootstrap server. */
+    public static final ConfigOption<String> KAFKA_BOOTSTRAP_SERVERS =
+            new ConfigOption<>("kafkaBootstrapServers", String.class, "localhost:9092");
+
+    /** The config parameter specifies the Kafka topic for action state. */
+    public static final ConfigOption<String> KAFKA_ACTION_STATE_TOPIC =
+            new ConfigOption<>("kafkaActionStateTopic", String.class, "flink-agents-action-state");
+
+    /** The config parameter sepcifies the number of partitions for the Kafka action state topic. */
+    public static final ConfigOption<Integer> KAFKA_ACTION_STATE_TOPIC_NUM_PARTITIONS =

Review Comment:
   We are auto-creating the state topic if it is not created thus I think it would be beneficial if we can let the user to specify the number of partition. 



##########
runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/KafkaActionStateStore.java:
##########
@@ -33,38 +57,303 @@
  */
 public class KafkaActionStateStore implements ActionStateStore {
 
-    // In memory action state for quick state retrival, this map is only used during recovery
-    private final Map<String, Map<String, ActionState>> keyedActionStates;
+    private static final Duration CONSUMER_POLL_TIMEOUT = Duration.ofMillis(1000);
+    private static final Logger LOG = LoggerFactory.getLogger(KafkaActionStateStore.class);
+    private static final Long DEFAULT_FUTURE_GET_TIMEOUT_MS = 100L;
+
+    private final AgentConfiguration agentConfiguration;
+
+    // In memory action state for quick state retrieval
+    private final Map<String, ActionState> actionStates;
+
+    // Record the lastest sequence number for each key that should be considered as valid
+    private final Map<String, Long> latestKeySeqNum;
+
+    // Kafka producer
+    private final Producer<String, ActionState> producer;
+    // Kafka consumer
+    private final Consumer<String, ActionState> consumer;
+
+    // Kafka topic that stores action states
+    private final String topic;
 
     @VisibleForTesting
-    KafkaActionStateStore(Map<String, Map<String, ActionState>> keyedActionStates) {
-        this.keyedActionStates = keyedActionStates;
+    KafkaActionStateStore(
+            Map<String, ActionState> actionStates,
+            AgentConfiguration agentConfiguration,
+            Producer<String, ActionState> producer,
+            Consumer<String, ActionState> consumer,
+            String topic) {
+        this.actionStates = actionStates;
+        this.producer = producer;
+        this.consumer = consumer;
+        this.topic = topic;
+        this.latestKeySeqNum = new HashMap<>();
+        this.agentConfiguration = agentConfiguration;
     }
 
-    /** Constructs a new KafkaActionStateStore with an empty in-memory action state map. */
-    public KafkaActionStateStore() {
-        this(new HashMap<>());
+    /** Constructs a new KafkaActionStateStore with custom Kafka configuration. */
+    public KafkaActionStateStore(AgentConfiguration agentConfiguration) {
+        this.actionStates = new HashMap<>();
+        this.latestKeySeqNum = new HashMap<>();
+        this.agentConfiguration = agentConfiguration;
+        this.topic = agentConfiguration.get(KAFKA_ACTION_STATE_TOPIC);
+        // create the topic if not exists
+        maybeCreateTopic();
+        Properties producerProp = createProducerProp();
+        this.producer = new KafkaProducer<>(producerProp);
+        Properties consumerProp = createConsumerProp();
+        this.consumer = new KafkaConsumer<>(consumerProp);
+        consumer.subscribe(Collections.singletonList(topic));
+        LOG.info("Initialized KafkaActionStateStore with topic: {}", topic);
     }
 
     @Override
     public void put(Object key, long seqNum, Action action, Event event, ActionState state)
-            throws IOException {
-        // TODO: Implement me
+            throws Exception {
+        if (producer == null) {
+            LOG.error("Producer is null, cannot put action state to Kafka");
+            return;
+        }
+
+        String stateKey = generateKey(key, seqNum, action, event);
+        try {
+            ProducerRecord<String, ActionState> kafkaRecord =
+                    new ProducerRecord<>(topic, stateKey, state);
+            Future<RecordMetadata> recordFuture = producer.send(kafkaRecord);
+            RecordMetadata recordMetadata =

Review Comment:
   Yep. Removed. I was using it to get the partition -> offset map but we are now getting the map in the `getRecoveryMarker` method



##########
runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/KafkaActionStateStore.java:
##########
@@ -33,38 +57,303 @@
  */
 public class KafkaActionStateStore implements ActionStateStore {
 
-    // In memory action state for quick state retrival, this map is only used during recovery
-    private final Map<String, Map<String, ActionState>> keyedActionStates;
+    private static final Duration CONSUMER_POLL_TIMEOUT = Duration.ofMillis(1000);
+    private static final Logger LOG = LoggerFactory.getLogger(KafkaActionStateStore.class);
+    private static final Long DEFAULT_FUTURE_GET_TIMEOUT_MS = 100L;
+
+    private final AgentConfiguration agentConfiguration;
+
+    // In memory action state for quick state retrieval
+    private final Map<String, ActionState> actionStates;
+
+    // Record the lastest sequence number for each key that should be considered as valid
+    private final Map<String, Long> latestKeySeqNum;
+
+    // Kafka producer
+    private final Producer<String, ActionState> producer;
+    // Kafka consumer
+    private final Consumer<String, ActionState> consumer;
+
+    // Kafka topic that stores action states
+    private final String topic;
 
     @VisibleForTesting
-    KafkaActionStateStore(Map<String, Map<String, ActionState>> keyedActionStates) {
-        this.keyedActionStates = keyedActionStates;
+    KafkaActionStateStore(
+            Map<String, ActionState> actionStates,
+            AgentConfiguration agentConfiguration,
+            Producer<String, ActionState> producer,
+            Consumer<String, ActionState> consumer,
+            String topic) {
+        this.actionStates = actionStates;
+        this.producer = producer;
+        this.consumer = consumer;
+        this.topic = topic;
+        this.latestKeySeqNum = new HashMap<>();
+        this.agentConfiguration = agentConfiguration;
     }
 
-    /** Constructs a new KafkaActionStateStore with an empty in-memory action state map. */
-    public KafkaActionStateStore() {
-        this(new HashMap<>());
+    /** Constructs a new KafkaActionStateStore with custom Kafka configuration. */
+    public KafkaActionStateStore(AgentConfiguration agentConfiguration) {
+        this.actionStates = new HashMap<>();
+        this.latestKeySeqNum = new HashMap<>();
+        this.agentConfiguration = agentConfiguration;
+        this.topic = agentConfiguration.get(KAFKA_ACTION_STATE_TOPIC);
+        // create the topic if not exists
+        maybeCreateTopic();
+        Properties producerProp = createProducerProp();
+        this.producer = new KafkaProducer<>(producerProp);
+        Properties consumerProp = createConsumerProp();
+        this.consumer = new KafkaConsumer<>(consumerProp);
+        consumer.subscribe(Collections.singletonList(topic));
+        LOG.info("Initialized KafkaActionStateStore with topic: {}", topic);
     }
 
     @Override
     public void put(Object key, long seqNum, Action action, Event event, ActionState state)
-            throws IOException {
-        // TODO: Implement me
+            throws Exception {
+        if (producer == null) {
+            LOG.error("Producer is null, cannot put action state to Kafka");
+            return;
+        }
+
+        String stateKey = generateKey(key, seqNum, action, event);
+        try {
+            ProducerRecord<String, ActionState> kafkaRecord =
+                    new ProducerRecord<>(topic, stateKey, state);
+            Future<RecordMetadata> recordFuture = producer.send(kafkaRecord);
+            RecordMetadata recordMetadata =
+                    recordFuture.get(DEFAULT_FUTURE_GET_TIMEOUT_MS, TimeUnit.MILLISECONDS);
+            LOG.debug("Sent action state to Kafka for key: {}", stateKey);
+        } catch (Exception e) {
+            throw new RuntimeException("Failed to send action state to Kafka", e);
+        }
     }
 
     @Override
-    public ActionState get(Object key, long seqNum, Action action, Event event) throws IOException {
-        // TODO: Implement me
-        return null;
+    public ActionState get(Object key, long seqNum, Action action, Event event) throws Exception {
+        String stateKey = generateKey(key, seqNum, action, event);
+
+        boolean hasDivergence = checkDivergence(key.toString(), seqNum);

Review Comment:
   This is what @xintongsong and I was discussing, think about a case like this
   
   1 -> 2 (crash)
      -> 2' -> 3 -> 4 (crash)
   
   When the second crash happened and we are trying to recover from the state and the input event comes with an order of `1, 2, 3, 4` (action 2 also generate action 3,4 in this case), we should not reuse the state of 3,4 because state of action 3, 4 was based on 2'. So to make 100% correct, we will ignore states after the diverge happens



##########
api/src/main/java/org/apache/flink/agents/api/configuration/AgentConfigOptions.java:
##########
@@ -27,4 +27,20 @@ public class AgentConfigOptions {
     /** The config parameter specifies the backend for action state store. */
     public static final ConfigOption<String> ACTION_STATE_STORE_BACKEND =
             new ConfigOption<>("actionStateStoreBackend", String.class, null);
+
+    /** The config parameter specifies the Kafka bootstrap server. */
+    public static final ConfigOption<String> KAFKA_BOOTSTRAP_SERVERS =
+            new ConfigOption<>("kafkaBootstrapServers", String.class, "localhost:9092");
+
+    /** The config parameter specifies the Kafka topic for action state. */
+    public static final ConfigOption<String> KAFKA_ACTION_STATE_TOPIC =
+            new ConfigOption<>("kafkaActionStateTopic", String.class, "flink-agents-action-state");

Review Comment:
   Make sense. Updated.



##########
runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/KafkaActionStateStore.java:
##########
@@ -33,38 +57,303 @@
  */
 public class KafkaActionStateStore implements ActionStateStore {
 
-    // In memory action state for quick state retrival, this map is only used during recovery
-    private final Map<String, Map<String, ActionState>> keyedActionStates;
+    private static final Duration CONSUMER_POLL_TIMEOUT = Duration.ofMillis(1000);
+    private static final Logger LOG = LoggerFactory.getLogger(KafkaActionStateStore.class);
+    private static final Long DEFAULT_FUTURE_GET_TIMEOUT_MS = 100L;
+
+    private final AgentConfiguration agentConfiguration;
+
+    // In memory action state for quick state retrieval
+    private final Map<String, ActionState> actionStates;
+
+    // Record the lastest sequence number for each key that should be considered as valid
+    private final Map<String, Long> latestKeySeqNum;
+
+    // Kafka producer
+    private final Producer<String, ActionState> producer;
+    // Kafka consumer
+    private final Consumer<String, ActionState> consumer;
+
+    // Kafka topic that stores action states
+    private final String topic;
 
     @VisibleForTesting
-    KafkaActionStateStore(Map<String, Map<String, ActionState>> keyedActionStates) {
-        this.keyedActionStates = keyedActionStates;
+    KafkaActionStateStore(
+            Map<String, ActionState> actionStates,
+            AgentConfiguration agentConfiguration,
+            Producer<String, ActionState> producer,
+            Consumer<String, ActionState> consumer,
+            String topic) {
+        this.actionStates = actionStates;
+        this.producer = producer;
+        this.consumer = consumer;
+        this.topic = topic;
+        this.latestKeySeqNum = new HashMap<>();
+        this.agentConfiguration = agentConfiguration;
     }
 
-    /** Constructs a new KafkaActionStateStore with an empty in-memory action state map. */
-    public KafkaActionStateStore() {
-        this(new HashMap<>());
+    /** Constructs a new KafkaActionStateStore with custom Kafka configuration. */
+    public KafkaActionStateStore(AgentConfiguration agentConfiguration) {
+        this.actionStates = new HashMap<>();
+        this.latestKeySeqNum = new HashMap<>();
+        this.agentConfiguration = agentConfiguration;
+        this.topic = agentConfiguration.get(KAFKA_ACTION_STATE_TOPIC);
+        // create the topic if not exists
+        maybeCreateTopic();

Review Comment:
   maybeCreateTopic method will create the topic if the topic does not exist with default config. 



##########
runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/ActionStateStore.java:
##########
@@ -24,7 +24,7 @@
 import java.util.List;
 
 /** Interface for storing and retrieving the state of actions performed by agents. */
-public interface ActionStateStore {
+public interface ActionStateStore extends AutoCloseable {

Review Comment:
   Yes, thank you for catching this!



-- 
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: issues-unsubscribe@flink.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [Feature] Implemenet Kafka-based state store [flink-agents]

Posted by "letaoj (via GitHub)" <gi...@apache.org>.
letaoj commented on code in PR #181:
URL: https://github.com/apache/flink-agents/pull/181#discussion_r2369836768


##########
api/src/main/java/org/apache/flink/agents/api/configuration/AgentConfigOptions.java:
##########
@@ -27,4 +27,20 @@ public class AgentConfigOptions {
     /** The config parameter specifies the backend for action state store. */
     public static final ConfigOption<String> ACTION_STATE_STORE_BACKEND =
             new ConfigOption<>("actionStateStoreBackend", String.class, null);
+
+    /** The config parameter specifies the Kafka bootstrap server. */
+    public static final ConfigOption<String> KAFKA_BOOTSTRAP_SERVERS =
+            new ConfigOption<>("kafkaBootstrapServers", String.class, "localhost:9092");
+
+    /** The config parameter specifies the Kafka topic for action state. */
+    public static final ConfigOption<String> KAFKA_ACTION_STATE_TOPIC =
+            new ConfigOption<>("kafkaActionStateTopic", String.class, "flink-agents-action-state");
+
+    /** The config parameter sepcifies the number of partitions for the Kafka action state topic. */
+    public static final ConfigOption<Integer> KAFKA_ACTION_STATE_TOPIC_NUM_PARTITIONS =

Review Comment:
   Because we are auto-creating the state table, it will be better to control the number of partitions by the framework. The partition will be use perform the custom partition for kafka (i can probably retrieve it if the topic is already created)



-- 
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: issues-unsubscribe@flink.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [api] Add TaskActionState API to persist task action state [flink-agents]

Posted by "letaoj (via GitHub)" <gi...@apache.org>.
letaoj commented on PR #181:
URL: https://github.com/apache/flink-agents/pull/181#issuecomment-3301136420

   I will need to rebase after the API change is merged..


-- 
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: issues-unsubscribe@flink.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [Feature] Implemenet Kafka-based state store [flink-agents]

Posted by "letaoj (via GitHub)" <gi...@apache.org>.
letaoj commented on code in PR #181:
URL: https://github.com/apache/flink-agents/pull/181#discussion_r2374060175


##########
runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java:
##########
@@ -350,13 +358,12 @@ private void processActionTaskForKey(Object key) throws Exception {
         long sequenceNumber = sequenceNumberKState.value();
         boolean isFinished;
         List<Event> outputEvents;
-        Optional<ActionTask> generatedActionTaskOpt;
+        Optional<ActionTask> generatedActionTaskOpt = Optional.empty();
         ActionState actionState =
                 maybeGetActionState(key, sequenceNumber, actionTask.action, actionTask.event);
-        if (actionState != null && actionState.getGeneratedActionTask().isEmpty()) {
+        if (actionState != null) {

Review Comment:
   > I assume this means we expect an incompleted action should not have a action state in the store?
   
   Yes, this is because I realize it will not be easy to serializer the ActionTask when the async action is only completed halfway. I will add a check before calling `maybePersistTaskResult()` so that unfinished job will not be persisted.



##########
runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java:
##########
@@ -350,13 +358,12 @@ private void processActionTaskForKey(Object key) throws Exception {
         long sequenceNumber = sequenceNumberKState.value();
         boolean isFinished;
         List<Event> outputEvents;
-        Optional<ActionTask> generatedActionTaskOpt;
+        Optional<ActionTask> generatedActionTaskOpt = Optional.empty();
         ActionState actionState =
                 maybeGetActionState(key, sequenceNumber, actionTask.action, actionTask.event);
-        if (actionState != null && actionState.getGeneratedActionTask().isEmpty()) {
+        if (actionState != null) {

Review Comment:
   > I assume this means we expect an incompleted action should not have a action state in the store?
   
   Yes, this is because I realize it will not be easy to serializer the ActionTask when the async action is only completed halfway. I will add a check in `maybePersistTaskResult()` so that unfinished job will not be persisted.



-- 
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: issues-unsubscribe@flink.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [Feature] Implemenet Kafka-based state store [flink-agents]

Posted by "letaoj (via GitHub)" <gi...@apache.org>.
letaoj commented on code in PR #181:
URL: https://github.com/apache/flink-agents/pull/181#discussion_r2374045119


##########
runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java:
##########
@@ -190,9 +190,17 @@ public void open() throws Exception {
                 && KAFKA.getType()
                         .equalsIgnoreCase(agentPlan.getConfig().get(ACTION_STATE_STORE_BACKEND))) {
             LOG.info("Using Kafka as backend of action state store.");
-            actionStateStore = new KafkaActionStateStore();
+            actionStateStore = new KafkaActionStateStore(agentPlan.getConfig());
         }
 
+        // init recovery marker state for recovery marker persistence
+        recoveryMarkerOpState =
+                getOperatorStateBackend()
+                        .getUnionListState(
+                                new ListStateDescriptor<>(
+                                        RECOVERY_MARKER_STATE_NAME,
+                                        TypeInformation.of(Object.class)));

Review Comment:
   Updated



-- 
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: issues-unsubscribe@flink.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [Feature] Implemenet Kafka-based state store [flink-agents]

Posted by "Sxnan (via GitHub)" <gi...@apache.org>.
Sxnan commented on code in PR #181:
URL: https://github.com/apache/flink-agents/pull/181#discussion_r2371306579


##########
runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/KafkaActionStateStore.java:
##########
@@ -33,38 +57,303 @@
  */
 public class KafkaActionStateStore implements ActionStateStore {
 
-    // In memory action state for quick state retrival, this map is only used during recovery
-    private final Map<String, Map<String, ActionState>> keyedActionStates;
+    private static final Duration CONSUMER_POLL_TIMEOUT = Duration.ofMillis(1000);
+    private static final Logger LOG = LoggerFactory.getLogger(KafkaActionStateStore.class);
+    private static final Long DEFAULT_FUTURE_GET_TIMEOUT_MS = 100L;
+
+    private final AgentConfiguration agentConfiguration;
+
+    // In memory action state for quick state retrieval
+    private final Map<String, ActionState> actionStates;
+
+    // Record the lastest sequence number for each key that should be considered as valid
+    private final Map<String, Long> latestKeySeqNum;
+
+    // Kafka producer
+    private final Producer<String, ActionState> producer;
+    // Kafka consumer
+    private final Consumer<String, ActionState> consumer;
+
+    // Kafka topic that stores action states
+    private final String topic;
 
     @VisibleForTesting
-    KafkaActionStateStore(Map<String, Map<String, ActionState>> keyedActionStates) {
-        this.keyedActionStates = keyedActionStates;
+    KafkaActionStateStore(
+            Map<String, ActionState> actionStates,
+            AgentConfiguration agentConfiguration,
+            Producer<String, ActionState> producer,
+            Consumer<String, ActionState> consumer,
+            String topic) {
+        this.actionStates = actionStates;
+        this.producer = producer;
+        this.consumer = consumer;
+        this.topic = topic;
+        this.latestKeySeqNum = new HashMap<>();
+        this.agentConfiguration = agentConfiguration;
     }
 
-    /** Constructs a new KafkaActionStateStore with an empty in-memory action state map. */
-    public KafkaActionStateStore() {
-        this(new HashMap<>());
+    /** Constructs a new KafkaActionStateStore with custom Kafka configuration. */
+    public KafkaActionStateStore(AgentConfiguration agentConfiguration) {
+        this.actionStates = new HashMap<>();
+        this.latestKeySeqNum = new HashMap<>();
+        this.agentConfiguration = agentConfiguration;
+        this.topic = agentConfiguration.get(KAFKA_ACTION_STATE_TOPIC);
+        // create the topic if not exists
+        maybeCreateTopic();
+        Properties producerProp = createProducerProp();
+        this.producer = new KafkaProducer<>(producerProp);
+        Properties consumerProp = createConsumerProp();
+        this.consumer = new KafkaConsumer<>(consumerProp);
+        consumer.subscribe(Collections.singletonList(topic));
+        LOG.info("Initialized KafkaActionStateStore with topic: {}", topic);
     }
 
     @Override
     public void put(Object key, long seqNum, Action action, Event event, ActionState state)
-            throws IOException {
-        // TODO: Implement me
+            throws Exception {
+        if (producer == null) {
+            LOG.error("Producer is null, cannot put action state to Kafka");
+            return;
+        }
+
+        String stateKey = generateKey(key, seqNum, action, event);
+        try {
+            ProducerRecord<String, ActionState> kafkaRecord =
+                    new ProducerRecord<>(topic, stateKey, state);
+            Future<RecordMetadata> recordFuture = producer.send(kafkaRecord);
+            RecordMetadata recordMetadata =
+                    recordFuture.get(DEFAULT_FUTURE_GET_TIMEOUT_MS, TimeUnit.MILLISECONDS);
+            LOG.debug("Sent action state to Kafka for key: {}", stateKey);
+        } catch (Exception e) {
+            throw new RuntimeException("Failed to send action state to Kafka", e);
+        }
     }
 
     @Override
-    public ActionState get(Object key, long seqNum, Action action, Event event) throws IOException {
-        // TODO: Implement me
-        return null;
+    public ActionState get(Object key, long seqNum, Action action, Event event) throws Exception {
+        String stateKey = generateKey(key, seqNum, action, event);
+
+        boolean hasDivergence = checkDivergence(key.toString(), seqNum);

Review Comment:
   Thanks for the explanation! LGTM



-- 
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: issues-unsubscribe@flink.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [Feature] Implemenet Kafka-based state store [flink-agents]

Posted by "letaoj (via GitHub)" <gi...@apache.org>.
letaoj commented on code in PR #181:
URL: https://github.com/apache/flink-agents/pull/181#discussion_r2369694666


##########
tools/lint.sh:
##########
@@ -67,7 +67,7 @@ install_python_deps() {
   if command -v uv >/dev/null 2>&1; then
     echo "Using uv for dependency management"
     pushd python
-    uv sync --extra lint
+    uv sync --extra lint --index-strategy unsafe-best-match

Review Comment:
   :( This is to mitigate the issue where our internal pip repo does not have the necessary dependency thus i have to add this to the lint for the lint to work. I will revert it to keep the main branch clean. 



-- 
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: issues-unsubscribe@flink.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org