joeyutong commented on code in PR #628: URL: https://github.com/apache/flink-agents/pull/628#discussion_r3144647757
########## runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStoreIT.java: ########## @@ -0,0 +1,411 @@ +/* + * 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.context.RunnerContext; +import org.apache.flink.agents.plan.AgentConfiguration; +import org.apache.flink.agents.plan.JavaFunction; +import org.apache.flink.agents.plan.actions.Action; +import org.apache.fluss.client.admin.Admin; +import org.apache.fluss.metadata.TableInfo; +import org.apache.fluss.metadata.TablePath; +import org.apache.fluss.server.testutils.FlussClusterExtension; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; + +import java.nio.charset.StandardCharsets; +import java.util.Collections; +import java.util.List; + +import static org.apache.flink.agents.api.configuration.AgentConfigOptions.FLUSS_ACTION_STATE_DATABASE; +import static org.apache.flink.agents.api.configuration.AgentConfigOptions.FLUSS_ACTION_STATE_TABLE; +import static org.apache.flink.agents.api.configuration.AgentConfigOptions.FLUSS_ACTION_STATE_TABLE_BUCKETS; +import static org.apache.flink.agents.api.configuration.AgentConfigOptions.FLUSS_BOOTSTRAP_SERVERS; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatNoException; + +/** Integration tests for {@link FlussActionStateStore} against an embedded Fluss cluster. */ +public class FlussActionStateStoreIT { Review Comment: [P1] This class mixes real Fluss IT coverage with cases that are closer to unit tests, such as cache lookup, serialization, and lifecycle behavior. Could we split it into a focused UT and a smaller Fluss IT? ########## runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStoreIT.java: ########## @@ -0,0 +1,411 @@ +/* + * 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.context.RunnerContext; +import org.apache.flink.agents.plan.AgentConfiguration; +import org.apache.flink.agents.plan.JavaFunction; +import org.apache.flink.agents.plan.actions.Action; +import org.apache.fluss.client.admin.Admin; +import org.apache.fluss.metadata.TableInfo; +import org.apache.fluss.metadata.TablePath; +import org.apache.fluss.server.testutils.FlussClusterExtension; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; + +import java.nio.charset.StandardCharsets; +import java.util.Collections; +import java.util.List; + +import static org.apache.flink.agents.api.configuration.AgentConfigOptions.FLUSS_ACTION_STATE_DATABASE; +import static org.apache.flink.agents.api.configuration.AgentConfigOptions.FLUSS_ACTION_STATE_TABLE; +import static org.apache.flink.agents.api.configuration.AgentConfigOptions.FLUSS_ACTION_STATE_TABLE_BUCKETS; +import static org.apache.flink.agents.api.configuration.AgentConfigOptions.FLUSS_BOOTSTRAP_SERVERS; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatNoException; + +/** Integration tests for {@link FlussActionStateStore} against an embedded Fluss cluster. */ +public class FlussActionStateStoreIT { + + private static final String TEST_DATABASE = "test_flink_agents"; + private static final String TEST_TABLE = "action_state_it"; + private static final String TEST_KEY = "test-key"; + + @RegisterExtension Review Comment: [P1] `FlussActionStateStoreIT` needs the required JVM opens to run reliably. When force-running it on JDK 17+, Fluss/Arrow requires `--add-opens=java.base/java.nio=ALL-UNNAMED`, but the current surefire config does not include it. Could we add this to the test configuration used by the runtime module? ########## runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/FlussActionStateStore.java: ########## @@ -0,0 +1,377 @@ +/* + * 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 com.fasterxml.jackson.annotation.JsonTypeInfo; +import com.fasterxml.jackson.databind.ObjectMapper; +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.plan.AgentConfiguration; +import org.apache.flink.agents.plan.actions.Action; +import org.apache.fluss.client.Connection; +import org.apache.fluss.client.ConnectionFactory; +import org.apache.fluss.client.admin.Admin; +import org.apache.fluss.client.table.Table; +import org.apache.fluss.client.table.scanner.ScanRecord; +import org.apache.fluss.client.table.scanner.log.LogScanner; +import org.apache.fluss.client.table.scanner.log.ScanRecords; +import org.apache.fluss.client.table.writer.AppendWriter; +import org.apache.fluss.config.Configuration; +import org.apache.fluss.metadata.DatabaseDescriptor; +import org.apache.fluss.metadata.Schema; +import org.apache.fluss.metadata.TableBucket; +import org.apache.fluss.metadata.TableDescriptor; +import org.apache.fluss.metadata.TablePath; +import org.apache.fluss.row.BinaryString; +import org.apache.fluss.row.GenericRow; +import org.apache.fluss.row.InternalRow; +import org.apache.fluss.types.DataTypes; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.time.Duration; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import static org.apache.flink.agents.api.configuration.AgentConfigOptions.FLUSS_ACTION_STATE_DATABASE; +import static org.apache.flink.agents.api.configuration.AgentConfigOptions.FLUSS_ACTION_STATE_TABLE; +import static org.apache.flink.agents.api.configuration.AgentConfigOptions.FLUSS_ACTION_STATE_TABLE_BUCKETS; +import static org.apache.flink.agents.api.configuration.AgentConfigOptions.FLUSS_BOOTSTRAP_SERVERS; +import static org.apache.flink.agents.api.configuration.AgentConfigOptions.FLUSS_SASL_JAAS_CONFIG; +import static org.apache.flink.agents.api.configuration.AgentConfigOptions.FLUSS_SASL_MECHANISM; +import static org.apache.flink.agents.api.configuration.AgentConfigOptions.FLUSS_SASL_PASSWORD; +import static org.apache.flink.agents.api.configuration.AgentConfigOptions.FLUSS_SASL_USERNAME; +import static org.apache.flink.agents.api.configuration.AgentConfigOptions.FLUSS_SECURITY_PROTOCOL; +import static org.apache.flink.agents.runtime.actionstate.ActionStateUtil.generateKey; + +/** + * An implementation of {@link ActionStateStore} that uses an Apache Fluss log table as the backend. + * All state is maintained in an in-memory HashMap for O(1) lookup, with the Fluss log table + * providing durability and recovery support. + * + * <p>Compared to the previous KV-table implementation: + * + * <ul> + * <li><b>Read</b>: O(1) in-memory HashMap lookup; no network I/O per query. + * <li><b>Write</b>: Append to log + update in-memory HashMap. + * <li><b>Recovery</b>: Scan log from beginning to rebuild in-memory state. + * <li><b>Prune</b>: In-memory eviction only; physical cleanup relies on Fluss log retention. + * </ul> + * + * <p>The Fluss log table schema (no primary key — this creates a log/append-only table): + * + * <pre> + * state_key STRING -- Generated by ActionStateUtil.generateKey() + * state_payload BYTES -- ActionState JSON serialized bytes + * seq_num BIGINT -- Sequence number for pruneState() + * agent_key STRING -- Original key; used as bucket key so all records for the same + * agent are co-located in the same bucket + * completed BOOLEAN -- Whether action completed + * updated_at BIGINT -- Last update timestamp + * </pre> + */ +public class FlussActionStateStore implements ActionStateStore { + + private static final Logger LOG = LoggerFactory.getLogger(FlussActionStateStore.class); + private static final ObjectMapper OBJECT_MAPPER = createObjectMapper(); + + /** + * Number of consecutive empty polls before considering the log fully consumed during recovery. + */ + private static final int EMPTY_POLL_THRESHOLD = 3; + + private static final Duration POLL_TIMEOUT = Duration.ofSeconds(1); + + /** Creates and configures the ObjectMapper for ActionState serialization. */ + private static ObjectMapper createObjectMapper() { + ObjectMapper mapper = new ObjectMapper(); + mapper.addMixIn(Event.class, EventTypeInfoMixin.class); + mapper.addMixIn(InputEvent.class, EventTypeInfoMixin.class); + mapper.addMixIn(OutputEvent.class, EventTypeInfoMixin.class); + return mapper; + } + + /** Mixin to add type information for Event hierarchy. */ + @JsonTypeInfo( + use = JsonTypeInfo.Id.CLASS, + include = JsonTypeInfo.As.PROPERTY, + property = "@class") + abstract static class EventTypeInfoMixin {} + + // Column indices in the Fluss table schema + private static final int COL_STATE_KEY = 0; + private static final int COL_STATE_PAYLOAD = 1; + private static final int COL_SEQ_NUM = 2; + private static final int COL_AGENT_KEY = 3; + private static final int COL_COMPLETED = 4; + private static final int COL_UPDATED_AT = 5; + private static final int NUM_COLUMNS = 6; + + private final AgentConfiguration agentConfiguration; + private final String databaseName; + private final String tableName; + private final TablePath tablePath; + + private final Connection connection; + private final Table table; + private final AppendWriter writer; + + /** In-memory cache for O(1) state lookups; rebuilt from Fluss log on recovery. */ + private final Map<String, ActionState> actionStates; + + public FlussActionStateStore(AgentConfiguration agentConfiguration) { + this.agentConfiguration = agentConfiguration; + this.databaseName = agentConfiguration.get(FLUSS_ACTION_STATE_DATABASE); + this.tableName = agentConfiguration.get(FLUSS_ACTION_STATE_TABLE); + this.tablePath = TablePath.of(databaseName, tableName); + this.actionStates = new HashMap<>(); + + Configuration flussConf = new Configuration(); + flussConf.setString("bootstrap.servers", agentConfiguration.get(FLUSS_BOOTSTRAP_SERVERS)); + + String securityProtocol = agentConfiguration.get(FLUSS_SECURITY_PROTOCOL); + if (securityProtocol != null) { + flussConf.setString("client.security.protocol", securityProtocol); + } + String saslMechanism = agentConfiguration.get(FLUSS_SASL_MECHANISM); + if (saslMechanism != null) { + flussConf.setString("client.security.sasl.mechanism", saslMechanism); + } + String jaasConfig = agentConfiguration.get(FLUSS_SASL_JAAS_CONFIG); + if (jaasConfig != null) { + flussConf.setString("client.security.sasl.jaas.config", jaasConfig); + } + String username = agentConfiguration.get(FLUSS_SASL_USERNAME); + if (username != null) { + flussConf.setString("client.security.sasl.username", username); + } + String password = agentConfiguration.get(FLUSS_SASL_PASSWORD); + if (password != null) { + flussConf.setString("client.security.sasl.password", password); + } + + this.connection = ConnectionFactory.createConnection(flussConf); + maybeCreateDatabaseAndTable(); + this.table = connection.getTable(tablePath); + this.writer = table.newAppend().createWriter(); + + LOG.info( + "Initialized FlussActionStateStore (log table) with table: {}.{}", + databaseName, + tableName); + } + + @Override + public void put(Object key, long seqNum, Action action, Event event, ActionState state) + throws Exception { + String stateKey = generateKey(key, seqNum, action, event); + byte[] payload = OBJECT_MAPPER.writeValueAsBytes(state); + + GenericRow row = new GenericRow(NUM_COLUMNS); + row.setField(COL_STATE_KEY, BinaryString.fromString(stateKey)); + row.setField(COL_STATE_PAYLOAD, payload); + row.setField(COL_SEQ_NUM, seqNum); + row.setField(COL_AGENT_KEY, BinaryString.fromString(key.toString())); + row.setField(COL_COMPLETED, state.isCompleted()); + row.setField(COL_UPDATED_AT, System.currentTimeMillis()); + + // Append to Fluss log for durability, then update in-memory cache. + // Synchronous write ensures the record is durable before returning. + writer.append(row).get(); + actionStates.put(stateKey, state); + + LOG.debug("Stored action state: key={}, isCompleted={}", stateKey, state.isCompleted()); + } + + @Override + public ActionState get(Object key, long seqNum, Action action, Event event) throws Exception { + String stateKey = generateKey(key, seqNum, action, event); + ActionState state = actionStates.get(stateKey); + LOG.debug("Lookup action state: key={}, found={}", stateKey, state != null); + return state; + } + + /** + * Rebuilds in-memory state by scanning the Fluss log table from the beginning. Reads until the + * log is exhausted ({@link #EMPTY_POLL_THRESHOLD} consecutive empty polls). For the same state + * key appearing multiple times in the log, the latest record wins (last-write-wins). + */ + @Override + public void rebuildState(List<Object> recoveryMarkers) { Review Comment: [P1] `rebuildState()` currently ignores `recoveryMarkers` and always scans the Fluss log from the beginning. This misses the checkpoint boundary semantics expected by `ActionStateStore` and makes recovery depend on retained history. Could we store and consume Fluss offsets as recovery markers, similar to the Kafka backend? -- 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]
