Merge branch 'master' into upgrade-iceberg-version

xxubai · web-flow · commit f23a886bc5e2 · 2026-03-06T17:48:36.000+08:00
diff --git a/amoro-ams/src/main/java/org/apache/amoro/server/AmoroServiceContainer.java b/amoro-ams/src/main/java/org/apache/amoro/server/AmoroServiceContainer.java
@@ -171,7 +171,7 @@ public static void main(String[] args) {
   }
 
   public void registAndElect() throws Exception {
-    haContainer.registAndElect();
+    haContainer.registerAndElect();
   }
 
   public enum HAState {
diff --git a/amoro-ams/src/main/java/org/apache/amoro/server/ha/DataBaseHighAvailabilityContainer.java b/amoro-ams/src/main/java/org/apache/amoro/server/ha/DataBaseHighAvailabilityContainer.java
@@ -28,6 +28,8 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.util.ArrayList;
+import java.util.List;
 import java.util.UUID;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.Executors;
@@ -135,6 +137,44 @@ public void waitFollowerShip() throws InterruptedException {
     LOG.info("Became the follower of AMS (Database lease)");
   }
 
+  @Override
+  public void registerAndElect() throws Exception {
+    boolean isMasterSlaveMode = serviceConfig.getBoolean(AmoroManagementConf.USE_MASTER_SLAVE_MODE);
+    if (!isMasterSlaveMode) {
+      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
+    long now = System.currentTimeMillis();
+    String optimizingInfoJson = 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);
+    } catch (Exception e) {
+      LOG.error("Failed to register node to database", e);
+      throw e;
+    }
+  }
+
+  @Override
+  public boolean hasLeadership() {
+    return isLeader.get();
+  }
+
   /** Closes the heartbeat executor safely. */
   @Override
   public void close() {
@@ -147,9 +187,6 @@ public void close() {
     }
   }
 
-  @Override
-  public void registAndElect() throws Exception {}
-
   private class HeartbeatRunnable implements Runnable {
     @Override
     public void run() {
@@ -304,6 +341,40 @@ private void onLeaderLost() {
     }
   }
 
+  @Override
+  public List<AmsServerInfo> getAliveNodes() {
+    List<AmsServerInfo> aliveNodes = new ArrayList<>();
+    if (!isLeader.get()) {
+      LOG.warn("Only leader node can get alive nodes list");
+      return aliveNodes;
+    }
+    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);
+            }
+          }
+        }
+      }
+    } catch (Exception e) {
+      LOG.error("Failed to get alive nodes from database", e);
+      throw e;
+    }
+    return aliveNodes;
+  }
+
   private AmsServerInfo buildServerInfo(String host, int thriftBindPort, int restBindPort) {
     AmsServerInfo amsServerInfo = new AmsServerInfo();
     amsServerInfo.setHost(host);
diff --git a/amoro-ams/src/main/java/org/apache/amoro/server/ha/HighAvailabilityContainer.java b/amoro-ams/src/main/java/org/apache/amoro/server/ha/HighAvailabilityContainer.java
@@ -18,6 +18,10 @@
 
 package org.apache.amoro.server.ha;
 
+import org.apache.amoro.client.AmsServerInfo;
+
+import java.util.List;
+
 /**
  * Common interface for high availability (HA) containers.
  *
@@ -49,5 +53,19 @@ public interface HighAvailabilityContainer {
    *
    * @throws Exception If registration fails or participation in the primary election fails.
    */
-  void registAndElect() throws Exception;
+  void registerAndElect() throws Exception;
+
+  /**
+   * Used in master-slave mode to obtain information about all currently registered AMS nodes.
+   *
+   * @return List<AmsServerInfo>
+   */
+  List<AmsServerInfo> getAliveNodes();
+
+  /**
+   * Used to determine whether the current AMS node is the primary node.
+   *
+   * @return true if the current AMS node is the primary node, false otherwise
+   */
+  boolean hasLeadership();
 }
diff --git a/amoro-ams/src/main/java/org/apache/amoro/server/ha/NoopHighAvailabilityContainer.java b/amoro-ams/src/main/java/org/apache/amoro/server/ha/NoopHighAvailabilityContainer.java
@@ -18,9 +18,11 @@
 
 package org.apache.amoro.server.ha;
 
+import org.apache.amoro.client.AmsServerInfo;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.util.List;
 import java.util.concurrent.CountDownLatch;
 
 /** No-op HA container that never blocks and performs no leader election. */
@@ -48,5 +50,15 @@ public void close() {
   }
 
   @Override
-  public void registAndElect() throws Exception {}
+  public void registerAndElect() throws Exception {}
+
+  @Override
+  public List<AmsServerInfo> getAliveNodes() {
+    return List.of();
+  }
+
+  @Override
+  public boolean hasLeadership() {
+    return false;
+  }
 }
diff --git a/amoro-ams/src/main/java/org/apache/amoro/server/ha/ZkHighAvailabilityContainer.java b/amoro-ams/src/main/java/org/apache/amoro/server/ha/ZkHighAvailabilityContainer.java
@@ -46,6 +46,8 @@
 import java.io.File;
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CountDownLatch;
 
@@ -55,13 +57,27 @@ public class ZkHighAvailabilityContainer implements HighAvailabilityContainer, L
 
   private final LeaderLatch leaderLatch;
   private final CuratorFramework zkClient;
+
+  // Package-private accessors for testing
+  CuratorFramework getZkClient() {
+    return zkClient;
+  }
+
+  LeaderLatch getLeaderLatch() {
+    return leaderLatch;
+  }
+
   private final String tableServiceMasterPath;
   private final String optimizingServiceMasterPath;
+  private final String nodesPath;
   private final AmsServerInfo tableServiceServerInfo;
   private final AmsServerInfo optimizingServiceServerInfo;
+  private final boolean isMasterSlaveMode;
   private volatile CountDownLatch followerLatch;
+  private String registeredNodePath;
 
   public ZkHighAvailabilityContainer(Configurations serviceConfig) throws Exception {
+    this.isMasterSlaveMode = serviceConfig.getBoolean(AmoroManagementConf.USE_MASTER_SLAVE_MODE);
     if (serviceConfig.getBoolean(AmoroManagementConf.HA_ENABLE)) {
       String zkServerAddress = serviceConfig.getString(AmoroManagementConf.HA_ZOOKEEPER_ADDRESS);
       int zkSessionTimeout =
@@ -71,6 +87,7 @@ public ZkHighAvailabilityContainer(Configurations serviceConfig) throws Exceptio
       String haClusterName = serviceConfig.getString(AmoroManagementConf.HA_CLUSTER_NAME);
       tableServiceMasterPath = AmsHAProperties.getTableServiceMasterPath(haClusterName);
       optimizingServiceMasterPath = AmsHAProperties.getOptimizingServiceMasterPath(haClusterName);
+      nodesPath = AmsHAProperties.getNodesPath(haClusterName);
       ExponentialBackoffRetry retryPolicy = new ExponentialBackoffRetry(1000, 3, 5000);
       setupZookeeperAuth(serviceConfig);
       this.zkClient =
@@ -83,6 +100,7 @@ public ZkHighAvailabilityContainer(Configurations serviceConfig) throws Exceptio
       zkClient.start();
       createPathIfNeeded(tableServiceMasterPath);
       createPathIfNeeded(optimizingServiceMasterPath);
+      createPathIfNeeded(nodesPath);
       String leaderPath = AmsHAProperties.getLeaderPath(haClusterName);
       createPathIfNeeded(leaderPath);
       leaderLatch = new LeaderLatch(zkClient, leaderPath);
@@ -103,8 +121,10 @@ public ZkHighAvailabilityContainer(Configurations serviceConfig) throws Exceptio
       zkClient = null;
       tableServiceMasterPath = null;
       optimizingServiceMasterPath = null;
+      nodesPath = null;
       tableServiceServerInfo = null;
       optimizingServiceServerInfo = null;
+      registeredNodePath = null;
       // block follower latch forever when ha is disabled
       followerLatch = new CountDownLatch(1);
     }
@@ -141,8 +161,25 @@ public void waitLeaderShip() throws Exception {
   }
 
   @Override
-  public void registAndElect() throws Exception {
-    // TODO Here you can register for AMS and participate in the election.
+  public void registerAndElect() throws Exception {
+    if (!isMasterSlaveMode) {
+      LOG.debug("Master-slave mode is not enabled, skip node registration");
+      return;
+    }
+    if (zkClient == null || nodesPath == null) {
+      LOG.warn("HA is not enabled, skip node registration");
+      return;
+    }
+    // Register node to ZK using ephemeral node
+    // The node will be automatically deleted when the session expires
+    String nodeInfo = JacksonUtil.toJSONString(optimizingServiceServerInfo);
+    registeredNodePath =
+        zkClient
+            .create()
+            .creatingParentsIfNeeded()
+            .withMode(CreateMode.EPHEMERAL_SEQUENTIAL)
+            .forPath(nodesPath + "/node-", nodeInfo.getBytes(StandardCharsets.UTF_8));
+    LOG.info("Registered AMS node to ZK: {}", registeredNodePath);
   }
 
   @Override
@@ -158,6 +195,18 @@ public void waitFollowerShip() throws Exception {
   public void close() {
     if (leaderLatch != null) {
       try {
+        // Unregister node from ZK
+        if (registeredNodePath != null) {
+          try {
+            zkClient.delete().forPath(registeredNodePath);
+            LOG.info("Unregistered AMS node from ZK: {}", registeredNodePath);
+          } catch (KeeperException.NoNodeException e) {
+            // Node already deleted, ignore
+            LOG.debug("Node {} already deleted", registeredNodePath);
+          } catch (Exception e) {
+            LOG.warn("Failed to unregister node from ZK: {}", registeredNodePath, e);
+          }
+        }
         this.leaderLatch.close();
         this.zkClient.close();
       } catch (IOException e) {
@@ -192,6 +241,60 @@ private AmsServerInfo buildServerInfo(String host, int thriftBindPort, int restB
     return amsServerInfo;
   }
 
+  /**
+   * Get list of alive nodes. Only the leader node can call this method.
+   *
+   * @return List of alive node information
+   */
+  public List<AmsServerInfo> getAliveNodes() {
+    List<AmsServerInfo> aliveNodes = new ArrayList<>();
+    if (!isMasterSlaveMode) {
+      LOG.debug("Master-slave mode is not enabled, return empty node list");
+      return aliveNodes;
+    }
+    if (zkClient == null || nodesPath == null) {
+      LOG.warn("HA is not enabled, return empty node list");
+      return aliveNodes;
+    }
+    if (!leaderLatch.hasLeadership()) {
+      LOG.warn("Only leader node can get alive nodes list");
+      return aliveNodes;
+    }
+    try {
+      List<String> nodePaths = zkClient.getChildren().forPath(nodesPath);
+      for (String nodePath : nodePaths) {
+        try {
+          String fullPath = nodesPath + "/" + nodePath;
+          byte[] data = zkClient.getData().forPath(fullPath);
+          if (data != null && data.length > 0) {
+            String nodeInfoJson = new String(data, StandardCharsets.UTF_8);
+            AmsServerInfo nodeInfo = JacksonUtil.parseObject(nodeInfoJson, AmsServerInfo.class);
+            aliveNodes.add(nodeInfo);
+          }
+        } catch (Exception e) {
+          LOG.warn("Failed to get node info for path: {}", nodePath, e);
+        }
+      }
+    } catch (KeeperException.NoNodeException e) {
+      LOG.debug("Nodes path {} does not exist", nodesPath);
+    } catch (Exception e) {
+      throw new RuntimeException(e);
+    }
+    return aliveNodes;
+  }
+
+  /**
+   * Check if current node is the leader.
+   *
+   * @return true if current node is the leader, false otherwise
+   */
+  public boolean hasLeadership() {
+    if (leaderLatch == null) {
+      return false;
+    }
+    return leaderLatch.hasLeadership();
+  }
+
   private void createPathIfNeeded(String path) throws Exception {
     try {
       zkClient.create().creatingParentsIfNeeded().withMode(CreateMode.PERSISTENT).forPath(path);
diff --git a/amoro-ams/src/main/java/org/apache/amoro/server/persistence/mapper/HaLeaseMapper.java b/amoro-ams/src/main/java/org/apache/amoro/server/persistence/mapper/HaLeaseMapper.java
@@ -107,6 +107,20 @@ int renewLease(
   HaLeaseMeta selectLease(
       @Param("clusterName") String clusterName, @Param("serviceName") String serviceName);
 
+  /**
+   * Select all leases for cluster and service.
+   *
+   * @param clusterName cluster name
+   * @param serviceName service name
+   * @return list of lease rows
+   */
+  @Select(
+      "SELECT cluster_name, service_name, node_id, node_ip, server_info_json, lease_expire_ts, version, updated_at "
+          + "FROM ha_lease WHERE cluster_name = #{clusterName} AND service_name = #{serviceName}")
+  @ResultMap("HaLeaseMetaMap")
+  List<HaLeaseMeta> selectLeasesByService(
+      @Param("clusterName") String clusterName, @Param("serviceName") String serviceName);
+
   /**
    * Select current lease for cluster and service.
    *
diff --git a/amoro-ams/src/test/java/org/apache/amoro/server/ha/TestZkHighAvailabilityContainer.java b/amoro-ams/src/test/java/org/apache/amoro/server/ha/TestZkHighAvailabilityContainer.java
diff --git a/amoro-common/src/main/java/org/apache/amoro/properties/AmsHAProperties.java b/amoro-common/src/main/java/org/apache/amoro/properties/AmsHAProperties.java

Original file line number	Diff line number	Diff line change
`@@ -171,7 +171,7 @@ public static void main(String[] args) {`
`171`	`171`	`}`
`172`	`172`
`173`	`173`	`public void registAndElect() throws Exception {`
`174`		`- haContainer.registAndElect();`
	`174`	`+ haContainer.registerAndElect();`
`175`	`175`	`}`
`176`	`176`
`177`	`177`	`public enum HAState {`