joeyutong commented on code in PR #628: URL: https://github.com/apache/flink-agents/pull/628#discussion_r3144759850
########## 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) { + LOG.info("Rebuilding action state from Fluss log table"); + actionStates.clear(); + + int numBuckets = table.getTableInfo().getNumBuckets(); + try (LogScanner scanner = table.newScan().createLogScanner()) { + for (int b = 0; b < numBuckets; b++) { + scanner.subscribeFromBeginning(b); + } + + int emptyPollCount = 0; + while (emptyPollCount < EMPTY_POLL_THRESHOLD) { + ScanRecords records = scanner.poll(POLL_TIMEOUT); + if (records.isEmpty()) { + emptyPollCount++; + } else { + emptyPollCount = 0; + for (TableBucket bucket : records.buckets()) { + for (ScanRecord record : records.records(bucket)) { + InternalRow row = record.getRow(); + String stateKey = row.getString(COL_STATE_KEY).toString(); + byte[] payload = row.getBytes(COL_STATE_PAYLOAD); + try { + ActionState state = + OBJECT_MAPPER.readValue(payload, ActionState.class); + actionStates.put(stateKey, state); + } catch (Exception e) { + LOG.warn( + "Failed to deserialize action state for key: {}, skipping", + stateKey, + e); + } + } + } + } + } + } catch (Exception e) { + throw new RuntimeException("Failed to rebuild state from Fluss log table", e); + } + + LOG.info("Completed rebuilding state, recovered {} states", actionStates.size()); + } + + /** + * Returns null. Recovery always rebuilds from the Fluss log beginning, so no positional markers + * are needed in Flink operator state. + */ + @Override + public Object getRecoveryMarker() { + return null; + } + + /** + * Evicts pruned states from the in-memory cache. The Fluss log is append-only; physical cleanup + * relies on Fluss log retention configuration. + */ + @Override + public void pruneState(Object key, long seqNum) { + LOG.debug("Pruning in-memory state for key: {} up to seqNum: {}", key, seqNum); + + actionStates + .entrySet() + .removeIf( + entry -> { + String stateKey = entry.getKey(); + if (stateKey.startsWith(key.toString() + "_")) { + try { + List<String> parts = ActionStateUtil.parseKey(stateKey); + if (parts.size() >= 2) { + long stateSeqNum = Long.parseLong(parts.get(1)); + return stateSeqNum <= seqNum; + } + } catch (NumberFormatException e) { + LOG.warn("Failed to parse seqNum from state key: {}", stateKey); + } + } + return false; + }); + + LOG.debug("Pruned in-memory state for key: {} up to seqNum: {}", key, seqNum); + } + + @Override + public void close() throws Exception { + Exception firstException = null; + + try { + if (writer != null) { + writer.flush(); + } + } catch (Exception e) { + firstException = e; + } + + try { + if (table != null) { + table.close(); + } + } catch (Exception e) { + if (firstException == null) { + firstException = e; + } + } + + try { + if (connection != null) { + connection.close(); + } + } catch (Exception e) { + if (firstException == null) { + firstException = e; + } + } + + if (firstException != null) { + throw firstException; + } + } + + private void maybeCreateDatabaseAndTable() { + try (Admin admin = connection.getAdmin()) { + if (!admin.databaseExists(databaseName).get()) { + admin.createDatabase(databaseName, DatabaseDescriptor.EMPTY, true).get(); + LOG.info("Created Fluss database: {}", databaseName); + } + + if (!admin.tableExists(tablePath).get()) { + // No primaryKey() call — this creates an append-only log table in Fluss. + Schema schema = Review Comment: Do we need all these metadata columns in the log table? In the current implementation, rebuild only reads `state_key` and `state_payload`; `seq_num`, `completed`, and `updated_at` seem either derivable from `state_key` / payload or unused by the store logic. If `agent_key` is only needed for bucket placement, maybe we can keep `state_key`, `state_payload`, and `agent_key`, and use `agent_key` as the distribution key. -- 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]
