[AMORO-3928] Modify the optimizer to support obtaining tasks from each AMS node for processing. (#3950)

* [Subtask]: Use a new configuration item to control whether master & slave mode is enabled. #3845

* [Subtask]: add AmsAssignService to implement balanced bucket allocation in master-slave mode. #3921

* [Subtask]: add AmsAssignService to implement balanced bucket allocation in master-slave mode. #3921

* [Subtask]: Modify DefaultTableService to be compatible with master-slave mode #3923

* [Subtask]: In master-slave mode, each AMS should automatically senses the optimizer. #3929

* [Subtask]: Modify the optimizer to support obtaining tasks from each AMS node for processing. #3928

* [Subtask]: Modify the optimizer to support obtaining tasks from each AMS node for processing. #3928

* [Subtask]: Optimize the logic for retrieving the AMS list from ZooKeeper in master-slave mode. #3928

* This addresses the conflict issues with the latest main branch and introduces a new solution for storing allocation information based on a database.

* [Subtask]: Modify the optimizer to support obtaining tasks from each AMS node for processing. #3928

* [Subtask]: Modify the optimizer to support obtaining tasks from each AMS node for processing. #3928

* [Subtask]: Modify the optimizer to support obtaining tasks from each AMS node for processing. #3928

* [Subtask]: Fixed a legacy bug that could cause unit tests to fail during compilation. #3928

* [Subtask]: Optimized based on CR feedback. #3928

* [Subtask]: Fixed a legacy bug that could cause unit tests to fail during compilation. #3928

* Revert "[Subtask]: Fixed a legacy bug that could cause unit tests to fail during compilation. #3928"

This reverts commit d54c126.

* [Subtask]: Optimized based on CR feedback. #3928

* [Subtask]: Optimized based on CR feedback. #3928

* [Subtask]: Optimized based on CR feedback. #3928

---------

Co-authored-by: wardli <wardli@tencent.com>

Author: wardlican

amoro-ams/src/main/java/org/apache/amoro/server/AmoroServiceContainer.java

@@ -277,7 +277,8 @@ public void startOptimizingService() throws Exception {
             serviceConfig, catalogManager, defaultRuntimeFactory, haContainer, bucketAssignStore);
     processService = new ProcessService(tableService, actionCoordinators, executeEngineManager);
     optimizingService =
-        new DefaultOptimizingService(serviceConfig, catalogManager, optimizerManager, tableService);
+        new DefaultOptimizingService(
+            serviceConfig, catalogManager, optimizerManager, tableService, bucketAssignStore);
 
     LOG.info("Setting up AMS table executors...");
     InlineTableExecutors.getInstance().setup(tableService, serviceConfig);

amoro-ams/src/main/java/org/apache/amoro/server/BucketAssignStore.java

@@ -65,6 +65,15 @@ void saveAssignments(AmsServerInfo nodeInfo, List<String> bucketIds)
    */
   Map<AmsServerInfo, List<String>> getAllAssignments() throws BucketAssignStoreException;
 
+  /**
+   * Get all alive AMS nodes that have bucket assignments. Used by optimizers in master-slave mode
+   * to discover all AMS optimizing endpoints.
+   *
+   * @return List of AmsServerInfo for all nodes with bucket assignments, empty list if none
+   * @throws BucketAssignStoreException If retrieval operation fails
+   */
+  List<AmsServerInfo> getAliveNodes() throws BucketAssignStoreException;
+
   /**
    * Get the last update time for a node's assignments.
    *

amoro-ams/src/main/java/org/apache/amoro/server/BucketAssignStoreFactory.java

@@ -63,8 +63,9 @@ public static BucketAssignStore create(
             "Cannot create ZkBucketAssignStore: ZK client not available or invalid container type");
 
       case AmoroManagementConf.HA_TYPE_DATABASE:
+        long nodeHeartbeatTtlMs = conf.get(AmoroManagementConf.HA_LEASE_TTL).toMillis();
         LOG.info("Creating DBBucketAssignStore for cluster: {}", clusterName);
-        return new DBBucketAssignStore(clusterName);
+        return new DBBucketAssignStore(clusterName, nodeHeartbeatTtlMs);
 
       default:
         throw new IllegalArgumentException(

amoro-ams/src/main/java/org/apache/amoro/server/DBBucketAssignStore.java

@@ -45,9 +45,11 @@ public class DBBucketAssignStore extends PersistentBase implements BucketAssignS
       new TypeReference<List<String>>() {};
 
   private final String clusterName;
+  private final long nodeHeartbeatTtlMs;
 
-  public DBBucketAssignStore(String clusterName) {
+  public DBBucketAssignStore(String clusterName, long nodeHeartbeatTtlMs) {
     this.clusterName = clusterName;
+    this.nodeHeartbeatTtlMs = nodeHeartbeatTtlMs;
   }
 
   @Override
@@ -180,6 +182,32 @@ public void updateLastUpdateTime(AmsServerInfo nodeInfo) throws BucketAssignStor
     }
   }
 
+  @Override
+  public List<AmsServerInfo> getAliveNodes() throws BucketAssignStoreException {
+    try {
+      long cutoff = System.currentTimeMillis() - nodeHeartbeatTtlMs;
+      List<BucketAssignmentMeta> rows =
+          getAs(BucketAssignMapper.class, mapper -> mapper.selectAllByCluster(clusterName));
+      List<AmsServerInfo> nodes = new ArrayList<>();
+      for (BucketAssignmentMeta meta : rows) {
+        Long heartbeatTs = meta.getNodeHeartbeatTs();
+        if (heartbeatTs == null || heartbeatTs < cutoff) {
+          LOG.debug(
+              "Skipping stale node key={}, node_heartbeat_ts={}", meta.getNodeKey(), heartbeatTs);
+          continue;
+        }
+        AmsServerInfo nodeInfo = parseNodeInfo(meta);
+        if (nodeInfo.getThriftBindPort() != null && nodeInfo.getThriftBindPort() > 0) {
+          nodes.add(nodeInfo);
+        }
+      }
+      return nodes;
+    } catch (Exception e) {
+      LOG.error("Failed to get alive nodes", e);
+      throw new BucketAssignStoreException("Failed to get alive nodes", e);
+    }
+  }
+
   private static String getNodeKey(AmsServerInfo nodeInfo) {
     return nodeInfo.getHost() + ":" + nodeInfo.getThriftBindPort();
   }

amoro-ams/src/main/java/org/apache/amoro/server/DefaultOptimizingService.java

@@ -26,6 +26,7 @@
 import org.apache.amoro.api.OptimizingTask;
 import org.apache.amoro.api.OptimizingTaskId;
 import org.apache.amoro.api.OptimizingTaskResult;
+import org.apache.amoro.client.AmsServerInfo;
 import org.apache.amoro.config.Configurations;
 import org.apache.amoro.config.TableConfiguration;
 import org.apache.amoro.exception.ForbiddenException;
@@ -66,6 +67,7 @@
 
 import java.time.Duration;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -116,12 +118,22 @@ public class DefaultOptimizingService extends StatedPersistentBase
   private final TableService tableService;
   private final RuntimeHandlerChain tableHandlerChain;
   private final ExecutorService planExecutor;
+  private final BucketAssignStore bucketAssignStore;
 
   public DefaultOptimizingService(
       Configurations serviceConfig,
       CatalogManager catalogManager,
       OptimizerManager optimizerManager,
       TableService tableService) {
+    this(serviceConfig, catalogManager, optimizerManager, tableService, null);
+  }
+
+  public DefaultOptimizingService(
+      Configurations serviceConfig,
+      CatalogManager catalogManager,
+      OptimizerManager optimizerManager,
+      TableService tableService,
+      BucketAssignStore bucketAssignStore) {
     this.optimizerTouchTimeout =
         serviceConfig.getDurationInMillis(AmoroManagementConf.OPTIMIZER_HB_TIMEOUT);
     this.taskAckTimeout =
@@ -144,6 +156,7 @@ public DefaultOptimizingService(
     this.tableService = tableService;
     this.catalogManager = catalogManager;
     this.optimizerManager = optimizerManager;
+    this.bucketAssignStore = bucketAssignStore;
     this.tableHandlerChain = new TableRuntimeHandlerImpl();
     this.planExecutor =
         Executors.newCachedThreadPool(
@@ -322,6 +335,28 @@ public boolean cancelProcess(long processId) {
     return true;
   }
 
+  @Override
+  public List<String> getOptimizingNodeUrls() {
+    if (bucketAssignStore == null) {
+      return Collections.emptyList();
+    }
+    try {
+      List<AmsServerInfo> nodes = bucketAssignStore.getAliveNodes();
+      List<String> urls = new ArrayList<>(nodes.size());
+      for (AmsServerInfo node : nodes) {
+        if (node.getHost() != null
+            && node.getThriftBindPort() != null
+            && node.getThriftBindPort() > 0) {
+          urls.add(String.format("thrift://%s:%d", node.getHost(), node.getThriftBindPort()));
+        }
+      }
+      return urls;
+    } catch (Exception e) {
+      LOG.warn("Failed to get optimizing node URLs from bucket assign store", e);
+      return Collections.emptyList();
+    }
+  }
+
   /**
    * Get optimizing queue.
    *

amoro-ams/src/main/java/org/apache/amoro/server/ZkBucketAssignStore.java

@@ -170,6 +170,35 @@ public Map<AmsServerInfo, List<String>> getAllAssignments() throws BucketAssignS
     return allAssignments;
   }
 
+  @Override
+  public List<AmsServerInfo> getAliveNodes() throws BucketAssignStoreException {
+    List<AmsServerInfo> nodes = new ArrayList<>();
+    try {
+      if (zkClient.checkExists().forPath(assignmentsBasePath) == null) {
+        return nodes;
+      }
+      List<String> nodeKeys = zkClient.getChildren().forPath(assignmentsBasePath);
+      for (String nodeKey : nodeKeys) {
+        try {
+          AmsServerInfo nodeInfo = parseNodeKey(nodeKey);
+          if (nodeInfo.getThriftBindPort() != null && nodeInfo.getThriftBindPort() > 0) {
+            nodes.add(nodeInfo);
+          }
+        } catch (Exception e) {
+          LOG.warn("Failed to parse node key: {}", nodeKey, e);
+        }
+      }
+    } catch (KeeperException.NoNodeException e) {
+      // path doesn't exist
+    } catch (BucketAssignStoreException e) {
+      throw e;
+    } catch (Exception e) {
+      LOG.error("Failed to get alive nodes", e);
+      throw new BucketAssignStoreException("Failed to get alive nodes", e);
+    }
+    return nodes;
+  }
+
   @Override
   public long getLastUpdateTime(AmsServerInfo nodeInfo) throws BucketAssignStoreException {
     String nodeKey = getNodeKey(nodeInfo);

amoro-ams/src/main/java/org/apache/amoro/server/ha/DataBaseHighAvailabilityContainer.java

@@ -21,8 +21,10 @@
 import org.apache.amoro.client.AmsServerInfo;
 import org.apache.amoro.config.Configurations;
 import org.apache.amoro.server.AmoroManagementConf;
+import org.apache.amoro.server.persistence.BucketAssignmentMeta;
 import org.apache.amoro.server.persistence.HaLeaseMeta;
 import org.apache.amoro.server.persistence.PersistentBase;
+import org.apache.amoro.server.persistence.mapper.BucketAssignMapper;
 import org.apache.amoro.server.persistence.mapper.HaLeaseMapper;
 import org.apache.amoro.utils.JacksonUtil;
 import org.slf4j.Logger;
@@ -144,29 +146,50 @@ public void registerAndElect() throws Exception {
       LOG.debug("Master-slave mode is not enabled, skip node registration");
       return;
     }
-    // In master-slave mode, register node to database by writing OPTIMIZING_SERVICE info
-    // This is similar to ZK mode registering ephemeral nodes
+    // Register this node in bucket_assignments so that all nodes can be discovered via
+    // getAliveNodes(). ha_lease has PK (cluster_name, service_name) and cannot store multiple
+    // nodes for the same service, so we use the per-node bucket_assignments table instead.
+    upsertNodeHeartbeat();
+    LOG.info(
+        "Registered AMS node to bucket_assignments: nodeKey={}, optimizingService={}",
+        getNodeKey(),
+        optimizingServiceServerInfo);
+  }
+
+  /** Returns nodeKey used as the bucket_assignments row identifier: host:optimizingPort. */
+  private String getNodeKey() {
+    return optimizingServiceServerInfo.getHost()
+        + ":"
+        + optimizingServiceServerInfo.getThriftBindPort();
+  }
+
+  /**
+   * Upsert this node's heartbeat row into bucket_assignments. Only updates node_heartbeat_ts (and
+   * server_info_json on first insert); never touches assignments_json so the leader's assignments
+   * are preserved.
+   */
+  private void upsertNodeHeartbeat() {
     long now = System.currentTimeMillis();
-    String optimizingInfoJson = JacksonUtil.toJSONString(optimizingServiceServerInfo);
+    String nodeKey = getNodeKey();
+    String serverInfoJson = JacksonUtil.toJSONString(optimizingServiceServerInfo);
     try {
-      doAsIgnoreError(
-          HaLeaseMapper.class,
-          mapper -> {
-            int updated =
-                mapper.updateServerInfo(
-                    clusterName, OPTIMIZING_SERVICE, nodeId, nodeIp, optimizingInfoJson, now);
-            if (updated == 0) {
-              mapper.insertServerInfoIfAbsent(
-                  clusterName, OPTIMIZING_SERVICE, nodeId, nodeIp, optimizingInfoJson, now);
-            }
-          });
-      LOG.info(
-          "Registered AMS node to database: nodeId={}, optimizingService={}",
-          nodeId,
-          optimizingServiceServerInfo);
+      int updated =
+          updateAs(
+                  BucketAssignMapper.class,
+                  mapper -> mapper.updateNodeHeartbeat(clusterName, nodeKey, now))
+              .intValue();
+      if (updated == 0) {
+        // First registration: insert a new row with empty assignments
+        doAsIgnoreError(
+            BucketAssignMapper.class,
+            mapper ->
+                mapper.insert(
+                    new BucketAssignmentMeta(
+                        clusterName, nodeKey, serverInfoJson, null, now, now)));
+      }
     } catch (Exception e) {
-      LOG.error("Failed to register node to database", e);
-      throw e;
+      LOG.error("Failed to upsert node heartbeat for nodeKey={}", nodeKey, e);
+      throw new RuntimeException("Failed to register node in bucket_assignments", e);
     }
   }
 
@@ -180,7 +203,12 @@ public AmsServerInfo getTableServiceServerInfo() {
     return tableServiceServerInfo;
   }
 
-  /** Closes the heartbeat executor safely. */
+  @Override
+  public AmsServerInfo getOptimizingServiceServerInfo() {
+    return optimizingServiceServerInfo;
+  }
+
+  /** Closes the heartbeat executor safely and removes this node's registration row. */
   @Override
   public void close() {
     try {
@@ -190,6 +218,18 @@ public void close() {
     } catch (Exception e) {
       LOG.error("Close Database HighAvailabilityContainer failed", e);
     }
+    // Remove this node from bucket_assignments so the leader immediately stops seeing it
+    // as alive without waiting for the heartbeat TTL to expire.
+    boolean isMasterSlaveMode = serviceConfig.getBoolean(AmoroManagementConf.USE_MASTER_SLAVE_MODE);
+    if (isMasterSlaveMode) {
+      try {
+        String nodeKey = getNodeKey();
+        doAs(BucketAssignMapper.class, mapper -> mapper.deleteByNode(clusterName, nodeKey));
+        LOG.info("Unregistered AMS node from bucket_assignments: nodeKey={}", nodeKey);
+      } catch (Exception e) {
+        LOG.warn("Failed to unregister node from bucket_assignments on close", e);
+      }
+    }
   }
 
   private class HeartbeatRunnable implements Runnable {
@@ -204,6 +244,14 @@ public void run() {
           return;
         }
 
+        // Each node independently refreshes its own heartbeat in bucket_assignments so that
+        // getAliveNodes() can correctly detect liveness without relying on ha_lease.
+        try {
+          upsertNodeHeartbeat();
+        } catch (Exception e) {
+          LOG.warn("Failed to refresh node heartbeat in bucket_assignments", e);
+        }
+
         if (!isLeader.get()) {
           // First attempt to acquire the lease (similar to candidate/await)
           boolean success = tryAcquireLease(newExpireTs, now);
@@ -353,29 +401,35 @@ public List<AmsServerInfo> getAliveNodes() {
       LOG.warn("Only leader node can get alive nodes list");
       return aliveNodes;
     }
+    // Read alive nodes from bucket_assignments keyed by node_heartbeat_ts. ha_lease has
+    // PK (cluster_name, service_name) which only allows one row per service and cannot
+    // represent multiple AMS nodes. bucket_assignments has PK (cluster_name, node_key) and
+    // stores one row per node; node_heartbeat_ts is updated exclusively by the owning node
+    // so the leader's refreshLastUpdateTime calls cannot mask a dead node's staleness.
     try {
-      long currentTime = System.currentTimeMillis();
-      List<HaLeaseMeta> leases =
-          getAs(
-              HaLeaseMapper.class,
-              mapper -> mapper.selectLeasesByService(clusterName, OPTIMIZING_SERVICE));
-      for (HaLeaseMeta lease : leases) {
-        // Only include nodes with valid (non-expired) leases
-        if (lease.getLeaseExpireTs() != null && lease.getLeaseExpireTs() > currentTime) {
-          if (lease.getServerInfoJson() != null && !lease.getServerInfoJson().isEmpty()) {
-            try {
-              AmsServerInfo nodeInfo =
-                  JacksonUtil.parseObject(lease.getServerInfoJson(), AmsServerInfo.class);
-              aliveNodes.add(nodeInfo);
-            } catch (Exception e) {
-              LOG.warn("Failed to parse server info for node {}", lease.getNodeId(), e);
-            }
+      long cutoff = System.currentTimeMillis() - TimeUnit.SECONDS.toMillis(ttlSeconds);
+      List<BucketAssignmentMeta> rows =
+          getAs(BucketAssignMapper.class, mapper -> mapper.selectAllByCluster(clusterName));
+      for (BucketAssignmentMeta meta : rows) {
+        Long heartbeatTs = meta.getNodeHeartbeatTs();
+        if (heartbeatTs == null || heartbeatTs < cutoff) {
+          LOG.debug(
+              "Skipping stale node key={}, node_heartbeat_ts={}", meta.getNodeKey(), heartbeatTs);
+          continue;
+        }
+        if (meta.getServerInfoJson() != null && !meta.getServerInfoJson().isEmpty()) {
+          try {
+            AmsServerInfo nodeInfo =
+                JacksonUtil.parseObject(meta.getServerInfoJson(), AmsServerInfo.class);
+            aliveNodes.add(nodeInfo);
+          } catch (Exception e) {
+            LOG.warn("Failed to parse server_info_json for node {}", meta.getNodeKey(), e);
           }
         }
       }
     } catch (Exception e) {
-      LOG.error("Failed to get alive nodes from database", e);
-      throw e;
+      LOG.error("Failed to get alive nodes from bucket_assignments", e);
+      throw new RuntimeException("Failed to get alive nodes", e);
     }
     return aliveNodes;
   }

amoro-ams/src/main/java/org/apache/amoro/server/ha/HighAvailabilityContainer.java

@@ -70,9 +70,18 @@ public interface HighAvailabilityContainer {
   boolean hasLeadership();
 
   /**
-   * Get current AMS node information.
+   * Get current AMS node's table service server info (host + table-service Thrift port).
    *
    * @return {@link AmsServerInfo}
    */
   AmsServerInfo getTableServiceServerInfo();
+
+  /**
+   * Get current AMS node's optimizing service server info (host + optimizing Thrift port). This
+   * must be consistent with the {@link AmsServerInfo} written into {@link
+   * org.apache.amoro.server.BucketAssignStore} so that bucket lookups can match the stored nodeKey.
+   *
+   * @return {@link AmsServerInfo}
+   */
+  AmsServerInfo getOptimizingServiceServerInfo();
 }