From 6a3cb26d617a2a3d3a24e2abac7182a9ff368fa6 Mon Sep 17 00:00:00 2001 From: zhoujinsong Date: Wed, 19 Aug 2026 20:16:13 +0800 Subject: [PATCH 1/5] Decouple HA container and bucket assign stor --- .../amoro/server/AmoroServiceContainer.java | 41 +- .../apache/amoro/server/AmsAssignService.java | 67 +-- .../amoro/server/BucketAssignStore.java | 19 + .../server/BucketAssignStoreFactory.java | 39 +- .../amoro/server/DBBucketAssignStore.java | 43 ++ .../amoro/server/ZkBucketAssignStore.java | 34 ++ .../ha/DataBaseHighAvailabilityContainer.java | 60 --- .../server/ha/HighAvailabilityContainer.java | 17 - .../ha/NoopHighAvailabilityContainer.java | 9 - .../ha/ZkHighAvailabilityContainer.java | 95 ---- .../amoro/server/TestAmsAssignService.java | 476 +++++++----------- .../server/TestHighAvailabilityContainer.java | 428 +--------------- .../ha/TestZkHighAvailabilityContainer.java | 433 +--------------- 13 files changed, 372 insertions(+), 1389 deletions(-) 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 index bea5498e92..657f9c90be 100644 --- a/amoro-ams/src/main/java/org/apache/amoro/server/AmoroServiceContainer.java +++ b/amoro-ams/src/main/java/org/apache/amoro/server/AmoroServiceContainer.java @@ -127,11 +127,11 @@ public class AmoroServiceContainer { private AmsServiceMetrics amsServiceMetrics; private HAState haState = HAState.INITIALIZING; private AmsAssignService amsAssignService; + private BucketAssignStore bucketAssignStore; public AmoroServiceContainer() throws Exception { initConfig(); haContainer = HighAvailabilityContainerFactory.create(serviceConfig); - haContainer.registerAndElect(); } public static void main(String[] args) { @@ -220,7 +220,13 @@ public void startRestServices() throws Exception { public void startBaseServices() throws Exception { startRestServices(); if (IS_MASTER_SLAVE_MODE) { + bucketAssignStore = BucketAssignStoreFactory.create(serviceConfig); startOptimizingService(); + // Register this node so AmsAssignService (leader) can discover it and assign buckets. + if (haContainer != null) { + bucketAssignStore.registerNode(haContainer.getOptimizingServiceServerInfo()); + LOG.info("Registered this node to BucketAssignStore"); + } } } @@ -241,11 +247,6 @@ private void startOptimizingService() throws Exception { DefaultTableRuntimeFactory defaultRuntimeFactory = new DefaultTableRuntimeFactory(); defaultRuntimeFactory.initialize(processFactories); - BucketAssignStore bucketAssignStore = null; - if (IS_MASTER_SLAVE_MODE && haContainer != null) { - bucketAssignStore = BucketAssignStoreFactory.create(haContainer, serviceConfig); - } - List actionCoordinators = defaultRuntimeFactory.supportedCoordinators(); tableService = @@ -293,9 +294,7 @@ public void startLeaderServices() throws Exception { // call (leader re-election); recreate it if needed. if (amsAssignService == null && haContainer != null) { try { - BucketAssignStore bucketAssignStore = - BucketAssignStoreFactory.create(haContainer, serviceConfig); - amsAssignService = new AmsAssignService(haContainer, serviceConfig, bucketAssignStore); + amsAssignService = new AmsAssignService(serviceConfig, bucketAssignStore); } catch (Exception e) { LOG.error("Failed to recreate AmsAssignService", e); } @@ -332,6 +331,27 @@ public void stopLeaderServices() { haState = HAState.FOLLOWER; } + public void stopBaseServices() { + disposeRestService(); + if (IS_MASTER_SLAVE_MODE) { + if (bucketAssignStore != null && haContainer != null) { + try { + bucketAssignStore.removeNode(haContainer.getOptimizingServiceServerInfo()); + LOG.info("Unregistered this node from BucketAssignStore"); + } catch (Exception e) { + LOG.warn("Failed to unregister node from BucketAssignStore", e); + } + try { + bucketAssignStore.close(); + } catch (Exception e) { + LOG.warn("Failed to close BucketAssignStore", e); + } + bucketAssignStore = null; + } + disposeOptimizingService(); + } + } + private void addHandlerChain(RuntimeHandlerChain chain) { if (chain != null) { tableService.addHandlerChain(chain); @@ -388,8 +408,7 @@ public void disposeRestService() { public void dispose() { stopLeaderServices(); - disposeOptimizingService(); - disposeRestService(); + stopBaseServices(); } private void initConfig() throws Exception { diff --git a/amoro-ams/src/main/java/org/apache/amoro/server/AmsAssignService.java b/amoro-ams/src/main/java/org/apache/amoro/server/AmsAssignService.java index 018be60a77..226e8d7458 100644 --- a/amoro-ams/src/main/java/org/apache/amoro/server/AmsAssignService.java +++ b/amoro-ams/src/main/java/org/apache/amoro/server/AmsAssignService.java @@ -21,7 +21,6 @@ import org.apache.amoro.client.AmsServerInfo; import org.apache.amoro.config.Configurations; import org.apache.amoro.exception.BucketAssignStoreException; -import org.apache.amoro.server.ha.HighAvailabilityContainer; import org.apache.amoro.shade.guava32.com.google.common.annotations.VisibleForTesting; import org.apache.amoro.shade.guava32.com.google.common.util.concurrent.ThreadFactoryBuilder; import org.slf4j.Logger; @@ -54,9 +53,7 @@ public class AmsAssignService { .setDaemon(true) .build()); - private final HighAvailabilityContainer haContainer; private final BucketAssignStore assignStore; - private final Configurations serviceConfig; private final int bucketIdTotalCount; private final long nodeOfflineTimeoutMs; private final long assignIntervalSeconds; @@ -66,21 +63,11 @@ boolean isRunning() { return running; } - public AmsAssignService(HighAvailabilityContainer haContainer, Configurations serviceConfig) { - this(haContainer, serviceConfig, null); - } - /** * @param assignStore if non-null, used as the bucket assignment store; otherwise one is created - * via {@link BucketAssignStoreFactory} (same instance can be shared with {@code - * DefaultTableService}). + * via {@link BucketAssignStoreFactory}. */ - public AmsAssignService( - HighAvailabilityContainer haContainer, - Configurations serviceConfig, - BucketAssignStore assignStore) { - this.haContainer = haContainer; - this.serviceConfig = serviceConfig; + public AmsAssignService(Configurations serviceConfig, BucketAssignStore assignStore) { this.bucketIdTotalCount = serviceConfig.getInteger(AmoroManagementConf.HA_BUCKET_ID_TOTAL_COUNT); this.nodeOfflineTimeoutMs = @@ -88,19 +75,13 @@ public AmsAssignService( this.assignIntervalSeconds = serviceConfig.get(AmoroManagementConf.HA_ASSIGN_INTERVAL).getSeconds(); this.assignStore = - assignStore != null - ? assignStore - : BucketAssignStoreFactory.create(haContainer, serviceConfig); + assignStore != null ? assignStore : BucketAssignStoreFactory.create(serviceConfig); } /** * Start the assignment service. Only works in master-slave mode and when current node is leader. */ public void start() { - if (!serviceConfig.getBoolean(AmoroManagementConf.HA_USE_MASTER_SLAVE_MODE)) { - LOG.info("Master-slave mode is not enabled, skip starting bucket assignment service"); - return; - } if (running) { LOG.warn("Bucket assignment service is already running"); return; @@ -132,12 +113,7 @@ public void stop() { @VisibleForTesting public void doAssign() { try { - if (!haContainer.hasLeadership()) { - LOG.debug("Current node is not leader, skip bucket assignment"); - return; - } - - List aliveNodes = haContainer.getAliveNodes(); + List aliveNodes = assignStore.getAliveNodes(); if (aliveNodes.isEmpty()) { LOG.debug("No alive nodes found, skip bucket assignment"); return; @@ -161,6 +137,8 @@ public void doAssign() { Map> newAssignments = buildNewAssignments(aliveNodes, new HashSet<>(), normalized.assignments); rebalanceExistingAssignments(aliveNodes, allBuckets, newAssignments); + // Remove assignments for nodes that are no longer alive + removeStaleAssignments(aliveNodes, currentAssignments); persistAssignments(newAssignments); } else { refreshLastUpdateTime(aliveNodes); @@ -175,6 +153,15 @@ public void doAssign() { List bucketsToRedistribute = handleOfflineNodes(change.offlineNodes, currentAssignments); + // Remove assignments for offline nodes so they don't linger in the store. + for (AmsServerInfo offlineNode : change.offlineNodes) { + try { + assignStore.removeAssignments(offlineNode); + LOG.info("Removed assignments for offline node {}", offlineNode); + } catch (Exception e) { + LOG.warn("Failed to remove assignments for offline node {}", offlineNode, e); + } + } List allBuckets = generateBucketIds(); Map> newAssignments = buildNewAssignments(aliveNodes, change.offlineNodes, normalized.assignments); @@ -676,6 +663,30 @@ private void persistAssignments(Map> newAssignments) } } + /** + * Remove assignments for nodes that are no longer in the alive list. This prevents stale + * assignments from lingering in the store when a node goes offline but its lastUpdateTime hasn't + * expired yet (so it wasn't detected as offline by detectNodeChanges). + */ + private void removeStaleAssignments( + List aliveNodes, Map> currentAssignments) { + Set aliveNodeKeys = new HashSet<>(); + for (AmsServerInfo node : aliveNodes) { + aliveNodeKeys.add(getNodeKey(node)); + } + for (AmsServerInfo assignedNode : currentAssignments.keySet()) { + String nodeKey = getNodeKey(assignedNode); + if (!aliveNodeKeys.contains(nodeKey)) { + try { + assignStore.removeAssignments(assignedNode); + LOG.info("Removed stale assignments for node {}", assignedNode); + } catch (Exception e) { + LOG.warn("Failed to remove stale assignments for node {}", assignedNode, e); + } + } + } + } + /** * Refreshes last update time for all alive nodes when no reassignment is needed. Per-node * failures are logged and skipped; the next run will retry. diff --git a/amoro-ams/src/main/java/org/apache/amoro/server/BucketAssignStore.java b/amoro-ams/src/main/java/org/apache/amoro/server/BucketAssignStore.java index b85751bf8d..31ad6ad4f5 100644 --- a/amoro-ams/src/main/java/org/apache/amoro/server/BucketAssignStore.java +++ b/amoro-ams/src/main/java/org/apache/amoro/server/BucketAssignStore.java @@ -74,6 +74,22 @@ void saveAssignments(AmsServerInfo nodeInfo, List bucketIds) */ List getAliveNodes() throws BucketAssignStoreException; + /** + * Register this node in the store so it can be discovered by {@link #getAliveNodes()}. + * + * @param serverInfo this node's server info + * @throws BucketAssignStoreException If registration fails + */ + void registerNode(AmsServerInfo serverInfo) throws BucketAssignStoreException; + + /** + * Remove this node's registration from the store. + * + * @param serverInfo this node's server info + * @throws BucketAssignStoreException If removal fails + */ + void removeNode(AmsServerInfo serverInfo) throws BucketAssignStoreException; + /** * Get the last update time for a node's assignments. * @@ -90,4 +106,7 @@ void saveAssignments(AmsServerInfo nodeInfo, List bucketIds) * @throws BucketAssignStoreException If update operation fails */ void updateLastUpdateTime(AmsServerInfo nodeInfo) throws BucketAssignStoreException; + + /** Close the store and release any resources (e.g. ZK connection). */ + void close(); } diff --git a/amoro-ams/src/main/java/org/apache/amoro/server/BucketAssignStoreFactory.java b/amoro-ams/src/main/java/org/apache/amoro/server/BucketAssignStoreFactory.java index 18f6d00e3e..628a89a230 100644 --- a/amoro-ams/src/main/java/org/apache/amoro/server/BucketAssignStoreFactory.java +++ b/amoro-ams/src/main/java/org/apache/amoro/server/BucketAssignStoreFactory.java @@ -19,16 +19,18 @@ package org.apache.amoro.server; import org.apache.amoro.config.Configurations; -import org.apache.amoro.server.ha.HighAvailabilityContainer; -import org.apache.amoro.server.ha.ZkHighAvailabilityContainer; import org.apache.amoro.shade.zookeeper3.org.apache.curator.framework.CuratorFramework; +import org.apache.amoro.shade.zookeeper3.org.apache.curator.framework.CuratorFrameworkFactory; +import org.apache.amoro.shade.zookeeper3.org.apache.curator.retry.ExponentialBackoffRetry; import org.slf4j.Logger; import org.slf4j.LoggerFactory; /** * Factory for creating BucketAssignStore implementations based on HA configuration. * - *

Supports different storage backends (ZK, database) according to HA type. + *

Supports different storage backends (ZK, database) according to HA type. This factory is + * independent of {@link org.apache.amoro.server.ha.HighAvailabilityContainer} — bucket assignment + * storage and leader election are separate concerns. */ public final class BucketAssignStoreFactory { private static final Logger LOG = LoggerFactory.getLogger(BucketAssignStoreFactory.class); @@ -36,31 +38,34 @@ public final class BucketAssignStoreFactory { private BucketAssignStoreFactory() {} /** - * Creates a BucketAssignStore based on the given HA configuration and container. + * Creates a BucketAssignStore based on the given configuration. * - * @param haContainer the HA container * @param conf service configuration * @return a BucketAssignStore implementation according to HA type * @throws IllegalArgumentException if HA type is unsupported * @throws RuntimeException if the ZK store cannot be created */ - public static BucketAssignStore create( - HighAvailabilityContainer haContainer, Configurations conf) { + public static BucketAssignStore create(Configurations conf) { String haType = conf.getString(AmoroManagementConf.HA_TYPE).toLowerCase(); String clusterName = conf.getString(AmoroManagementConf.HA_CLUSTER_NAME); switch (haType) { case AmoroManagementConf.HA_TYPE_ZK: - if (haContainer instanceof ZkHighAvailabilityContainer) { - ZkHighAvailabilityContainer zkHaContainer = (ZkHighAvailabilityContainer) haContainer; - CuratorFramework zkClient = zkHaContainer.getZkClient(); - if (zkClient != null) { - LOG.info("Creating ZkBucketAssignStore for cluster: {}", clusterName); - return new ZkBucketAssignStore(zkClient, clusterName); - } - } - throw new RuntimeException( - "Cannot create ZkBucketAssignStore: ZK client not available or invalid container type"); + String zkAddress = conf.getString(AmoroManagementConf.HA_ZOOKEEPER_ADDRESS); + int sessionTimeoutMs = + (int) conf.get(AmoroManagementConf.HA_ZOOKEEPER_SESSION_TIMEOUT).toMillis(); + int connectionTimeoutMs = + (int) conf.get(AmoroManagementConf.HA_ZOOKEEPER_CONNECTION_TIMEOUT).toMillis(); + CuratorFramework zkClient = + CuratorFrameworkFactory.builder() + .connectString(zkAddress) + .sessionTimeoutMs(sessionTimeoutMs) + .connectionTimeoutMs(connectionTimeoutMs) + .retryPolicy(new ExponentialBackoffRetry(1000, 3)) + .build(); + zkClient.start(); + LOG.info("Creating ZkBucketAssignStore for cluster: {}", clusterName); + return new ZkBucketAssignStore(zkClient, clusterName); case AmoroManagementConf.HA_TYPE_DATABASE: long nodeHeartbeatTtlMs = conf.get(AmoroManagementConf.HA_LEASE_TTL).toMillis(); diff --git a/amoro-ams/src/main/java/org/apache/amoro/server/DBBucketAssignStore.java b/amoro-ams/src/main/java/org/apache/amoro/server/DBBucketAssignStore.java index 0727e850c2..45c59f682b 100644 --- a/amoro-ams/src/main/java/org/apache/amoro/server/DBBucketAssignStore.java +++ b/amoro-ams/src/main/java/org/apache/amoro/server/DBBucketAssignStore.java @@ -118,6 +118,44 @@ public void removeAssignments(AmsServerInfo nodeInfo) throws BucketAssignStoreEx } } + @Override + public void registerNode(AmsServerInfo serverInfo) throws BucketAssignStoreException { + String nodeKey = getNodeKey(serverInfo); + String serverInfoJson = JacksonUtil.toJSONString(serverInfo); + long now = System.currentTimeMillis(); + try { + int updated = + updateAs( + BucketAssignMapper.class, + mapper -> mapper.updateNodeHeartbeat(clusterName, nodeKey, now)) + .intValue(); + if (updated == 0) { + doAs( + BucketAssignMapper.class, + mapper -> + mapper.insert( + new BucketAssignmentMeta( + clusterName, nodeKey, serverInfoJson, null, now, now))); + } + LOG.debug("Registered node {} in bucket_assignments", nodeKey); + } catch (Exception e) { + LOG.error("Failed to register node {}", nodeKey, e); + throw new BucketAssignStoreException("Failed to register node " + nodeKey, e); + } + } + + @Override + public void removeNode(AmsServerInfo serverInfo) throws BucketAssignStoreException { + String nodeKey = getNodeKey(serverInfo); + try { + doAs(BucketAssignMapper.class, mapper -> mapper.deleteByNode(clusterName, nodeKey)); + LOG.debug("Removed node {} from bucket_assignments", nodeKey); + } catch (Exception e) { + LOG.error("Failed to remove node {}", nodeKey, e); + throw new BucketAssignStoreException("Failed to remove node " + nodeKey, e); + } + } + @Override public Map> getAllAssignments() throws BucketAssignStoreException { try { @@ -236,4 +274,9 @@ private static AmsServerInfo parseNodeKey(String nodeKey) { nodeInfo.setThriftBindPort(Integer.parseInt(parts[1])); return nodeInfo; } + + @Override + public void close() { + // No resources to release — DataSource is shared and managed globally. + } } diff --git a/amoro-ams/src/main/java/org/apache/amoro/server/ZkBucketAssignStore.java b/amoro-ams/src/main/java/org/apache/amoro/server/ZkBucketAssignStore.java index ed498c3e39..0466cdabaf 100644 --- a/amoro-ams/src/main/java/org/apache/amoro/server/ZkBucketAssignStore.java +++ b/amoro-ams/src/main/java/org/apache/amoro/server/ZkBucketAssignStore.java @@ -199,6 +199,35 @@ public List getAliveNodes() throws BucketAssignStoreException { return nodes; } + @Override + public void registerNode(AmsServerInfo serverInfo) throws BucketAssignStoreException { + String nodeKey = getNodeKey(serverInfo); + String nodePath = assignmentsBasePath + "/" + nodeKey; + String serverInfoJson = JacksonUtil.toJSONString(serverInfo); + try { + zkClient + .create() + .creatingParentsIfNeeded() + .withMode(CreateMode.EPHEMERAL) + .forPath(nodePath, serverInfoJson.getBytes(StandardCharsets.UTF_8)); + LOG.debug("Registered node {} in ZK", nodeKey); + } catch (KeeperException.NodeExistsException e) { + // Already registered, update data + try { + zkClient.setData().forPath(nodePath, serverInfoJson.getBytes(StandardCharsets.UTF_8)); + } catch (Exception ex) { + throw new BucketAssignStoreException("Failed to update node registration " + nodeKey, ex); + } + } catch (Exception e) { + throw new BucketAssignStoreException("Failed to register node " + nodeKey, e); + } + } + + @Override + public void removeNode(AmsServerInfo serverInfo) throws BucketAssignStoreException { + removeAssignments(serverInfo); + } + @Override public long getLastUpdateTime(AmsServerInfo nodeInfo) throws BucketAssignStoreException { String nodeKey = getNodeKey(nodeInfo); @@ -271,4 +300,9 @@ private void createPathIfNeeded(String path) throws BucketAssignStoreException { throw new BucketAssignStoreException("Failed to create path: " + path, e); } } + + @Override + public void close() { + zkClient.close(); + } } 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 index 1aeaba5835..1f51039253 100644 --- 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 @@ -30,8 +30,6 @@ 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; @@ -139,24 +137,6 @@ 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.HA_USE_MASTER_SLAVE_MODE); - if (!isMasterSlaveMode) { - LOG.debug("Master-slave mode is not enabled, skip node registration"); - return; - } - // 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() @@ -396,46 +376,6 @@ private void onLeaderLost() { } } - @Override - public List getAliveNodes() { - List aliveNodes = new ArrayList<>(); - if (!isLeader.get()) { - 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 cutoff = System.currentTimeMillis() - TimeUnit.SECONDS.toMillis(ttlSeconds); - List 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 bucket_assignments", e); - throw new RuntimeException("Failed to get alive nodes", 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 index 30d01a6063..66c390c502 100644 --- 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 @@ -20,8 +20,6 @@ import org.apache.amoro.client.AmsServerInfo; -import java.util.List; - /** * Common interface for high availability (HA) containers. * @@ -47,21 +45,6 @@ public interface HighAvailabilityContainer { /** Closes the container and releases resources. */ void close(); - /** - * In master-slave mode, this is used for AMS nodes to register and participate in the master - * election process. - * - * @throws Exception If registration fails or participation in the primary election fails. - */ - void registerAndElect() throws Exception; - - /** - * Used in master-slave mode to obtain information about all currently registered AMS nodes. - * - * @return List - */ - List getAliveNodes(); - /** * Used to determine whether the current AMS node is the primary node. * 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 index 48282f2a9e..5660a8980f 100644 --- 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 @@ -22,7 +22,6 @@ 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. */ @@ -49,14 +48,6 @@ public void close() { LOG.info("Noop HA: closed"); } - @Override - public void registerAndElect() throws Exception {} - - @Override - public List 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 index 9d557737c1..bca32342e0 100644 --- 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,8 +46,6 @@ 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; @@ -59,15 +57,11 @@ public class ZkHighAvailabilityContainer implements HighAvailabilityContainer, L private final CuratorFramework zkClient; 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.HA_USE_MASTER_SLAVE_MODE); if (serviceConfig.getBoolean(AmoroManagementConf.HA_ENABLE)) { String zkServerAddress = serviceConfig.getString(AmoroManagementConf.HA_ZOOKEEPER_ADDRESS); int zkSessionTimeout = @@ -77,7 +71,6 @@ 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 = @@ -90,7 +83,6 @@ 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); @@ -111,10 +103,8 @@ 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); } @@ -150,28 +140,6 @@ public void waitLeaderShip() throws Exception { LOG.info("Became the leader of AMS"); } - @Override - 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 public void waitFollowerShip() throws Exception { LOG.info("Waiting to become the follower of AMS"); @@ -185,18 +153,6 @@ 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) { @@ -254,48 +210,6 @@ 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 getAliveNodes() { - List 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 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. * @@ -328,15 +242,6 @@ public AmsServerInfo getOptimizingServiceServerInfo() { return optimizingServiceServerInfo; } - /** - * Get the ZooKeeper client. This is used for creating BucketAssignStore. - * - * @return The ZooKeeper client, null if HA is not enabled - */ - public CuratorFramework getZkClient() { - return zkClient; - } - private void createPathIfNeeded(String path) throws Exception { try { zkClient.create().creatingParentsIfNeeded().withMode(CreateMode.PERSISTENT).forPath(path); diff --git a/amoro-ams/src/test/java/org/apache/amoro/server/TestAmsAssignService.java b/amoro-ams/src/test/java/org/apache/amoro/server/TestAmsAssignService.java index 91fad64363..0e69e091dd 100644 --- a/amoro-ams/src/test/java/org/apache/amoro/server/TestAmsAssignService.java +++ b/amoro-ams/src/test/java/org/apache/amoro/server/TestAmsAssignService.java @@ -101,6 +101,9 @@ public void setUp() throws Exception { node3.setHost("127.0.0.3"); node3.setThriftBindPort(1263); node3.setRestBindPort(1633); + + // Register node1 by default + mockAssignStore.registerNode(node1); } @After @@ -116,314 +119,189 @@ public void tearDown() throws Exception { @Test public void testInitialAssignment() throws Exception { - // Register nodes - haContainer.registerAndElect(); - - // Create second node - Configurations config2 = createNodeConfig("127.0.0.2", 1262, 1632); - HighAvailabilityContainer haContainer2 = createContainerWithMockZk(config2); - haContainer2.registerAndElect(); - - try { - // Wait a bit for registration - Thread.sleep(100); - - // Trigger assignment manually - assignService.doAssign(); - - // Check assignments - Map> assignments = mockAssignStore.getAllAssignments(); - Assert.assertEquals("Should have assignments for 2 nodes", 2, assignments.size()); - - // Verify buckets are distributed - int totalAssigned = 0; - for (List buckets : assignments.values()) { - totalAssigned += buckets.size(); - Assert.assertTrue("Each node should have buckets", !buckets.isEmpty()); - } - Assert.assertEquals("All buckets should be assigned", 100, totalAssigned); + mockAssignStore.registerNode(node2); - // Verify balance (difference should be at most 1) - List bucketCounts = new ArrayList<>(); - for (List buckets : assignments.values()) { - bucketCounts.add(buckets.size()); - } - int max = bucketCounts.stream().mapToInt(Integer::intValue).max().orElse(0); - int min = bucketCounts.stream().mapToInt(Integer::intValue).min().orElse(0); - Assert.assertTrue("Difference should be at most 1", max - min <= 1); - } finally { - haContainer2.close(); + assignService.doAssign(); + + Map> assignments = mockAssignStore.getAllAssignments(); + Assert.assertEquals("Should have assignments for 2 nodes", 2, assignments.size()); + + int totalAssigned = 0; + for (List buckets : assignments.values()) { + totalAssigned += buckets.size(); + Assert.assertTrue("Each node should have buckets", !buckets.isEmpty()); + } + Assert.assertEquals("All buckets should be assigned", 100, totalAssigned); + + List bucketCounts = new ArrayList<>(); + for (List buckets : assignments.values()) { + bucketCounts.add(buckets.size()); } + int max = bucketCounts.stream().mapToInt(Integer::intValue).max().orElse(0); + int min = bucketCounts.stream().mapToInt(Integer::intValue).min().orElse(0); + Assert.assertTrue("Difference should be at most 1", max - min <= 1); } @Test public void testNodeOfflineReassignment() throws Exception { - // Setup: 2 nodes with initial assignment - haContainer.registerAndElect(); - Configurations config2 = createNodeConfig("127.0.0.2", 1262, 1632); - HighAvailabilityContainer haContainer2 = createContainerWithMockZk(config2); - haContainer2.registerAndElect(); - - try { - Thread.sleep(100); - - // Initial assignment - assignService.doAssign(); - Map> initialAssignments = mockAssignStore.getAllAssignments(); - Assert.assertEquals("Should have 2 nodes", 2, initialAssignments.size()); - - // Verify initial assignment is balanced - List initialCounts = new ArrayList<>(); - for (List buckets : initialAssignments.values()) { - initialCounts.add(buckets.size()); - } - int maxInitial = initialCounts.stream().mapToInt(Integer::intValue).max().orElse(0); - int minInitial = initialCounts.stream().mapToInt(Integer::intValue).min().orElse(0); - Assert.assertTrue("Initial assignment should be balanced", maxInitial - minInitial <= 1); - - // Simulate node2 going offline by removing it from mock state - mockZkState.deleteNodeByHost("127.0.0.2"); - Thread.sleep(100); - - // Trigger reassignment - assignService.doAssign(); - - // Check that node2's buckets are redistributed - Map> newAssignments = mockAssignStore.getAllAssignments(); - Assert.assertEquals("Should have 1 node after offline", 1, newAssignments.size()); - - // The only remaining node (node1) should have all buckets. ZK stores - // optimizingServiceServerInfo (thrift port 1261), not table port (1260), so we - // take the single entry instead of matching by node1's thriftBindPort. - List remainingBuckets = newAssignments.values().iterator().next(); - Assert.assertNotNull("Node1 should have assignments", remainingBuckets); - Assert.assertEquals("Node1 should have all buckets", 100, remainingBuckets.size()); - } finally { - try { - haContainer2.close(); - } catch (Exception e) { - // ignore - } - } + mockAssignStore.registerNode(node2); + + // Initial assignment + assignService.doAssign(); + Map> initialAssignments = mockAssignStore.getAllAssignments(); + Assert.assertEquals("Should have 2 nodes", 2, initialAssignments.size()); + + // Simulate node2 going offline (remove from alive list but keep assignments) + mockAssignStore.simulateOffline(node2); + + // Wait for node offline timeout to expire + Thread.sleep(60); + + // Trigger reassignment + assignService.doAssign(); + + Map> newAssignments = mockAssignStore.getAllAssignments(); + Assert.assertEquals("Should have 1 node after offline", 1, newAssignments.size()); + + List remainingBuckets = newAssignments.values().iterator().next(); + Assert.assertNotNull("Node1 should have assignments", remainingBuckets); + Assert.assertEquals("Node1 should have all buckets", 100, remainingBuckets.size()); } @Test public void testNewNodeIncrementalAssignment() throws Exception { - // Setup: 1 node initially - haContainer.registerAndElect(); - Thread.sleep(100); - // Initial assignment - all buckets to node1 assignService.doAssign(); Map> initialAssignments = mockAssignStore.getAllAssignments(); - // ZK stores optimizing port (1261), not table port (1260); match by host only (single node). List node1InitialBuckets = findBucketsByHost(initialAssignments, node1.getHost()); Assert.assertNotNull("Node1 should have assignments", node1InitialBuckets); Assert.assertEquals("Node1 should have all buckets initially", 100, node1InitialBuckets.size()); // Add new node - Configurations config2 = createNodeConfig("127.0.0.2", 1262, 1632); - HighAvailabilityContainer haContainer2 = createContainerWithMockZk(config2); - haContainer2.registerAndElect(); - - try { - Thread.sleep(100); - - // Trigger reassignment - assignService.doAssign(); - - // Check assignments - Map> newAssignments = mockAssignStore.getAllAssignments(); - Assert.assertEquals("Should have 2 nodes", 2, newAssignments.size()); + mockAssignStore.registerNode(node2); - // Verify incremental assignment - node1 should keep most of its buckets. - // ZK stores optimizing port, not table port; match by host. - List node1NewBuckets = findBucketsByHost(newAssignments, node1.getHost()); - Assert.assertNotNull("Node1 should still have assignments", node1NewBuckets); + // Trigger reassignment + assignService.doAssign(); - // Node1 should have kept most buckets (incremental assignment) - Assert.assertTrue("Node1 should keep some buckets", node1NewBuckets.size() > 0); + Map> newAssignments = mockAssignStore.getAllAssignments(); + Assert.assertEquals("Should have 2 nodes", 2, newAssignments.size()); - // Verify balance - List bucketCounts = new ArrayList<>(); - for (List buckets : newAssignments.values()) { - bucketCounts.add(buckets.size()); - } - int max = bucketCounts.stream().mapToInt(Integer::intValue).max().orElse(0); - int min = bucketCounts.stream().mapToInt(Integer::intValue).min().orElse(0); - Assert.assertTrue("Difference should be at most 1", max - min <= 1); + List node1NewBuckets = findBucketsByHost(newAssignments, node1.getHost()); + Assert.assertNotNull("Node1 should still have assignments", node1NewBuckets); + Assert.assertTrue("Node1 should keep some buckets", node1NewBuckets.size() > 0); - // Verify total - int total = bucketCounts.stream().mapToInt(Integer::intValue).sum(); - Assert.assertEquals("Total buckets should be 100", 100, total); - } finally { - haContainer2.close(); + // Verify balance + List bucketCounts = new ArrayList<>(); + for (List buckets : newAssignments.values()) { + bucketCounts.add(buckets.size()); } + int max = bucketCounts.stream().mapToInt(Integer::intValue).max().orElse(0); + int min = bucketCounts.stream().mapToInt(Integer::intValue).min().orElse(0); + Assert.assertTrue("Difference should be at most 1", max - min <= 1); + + int total = bucketCounts.stream().mapToInt(Integer::intValue).sum(); + Assert.assertEquals("Total buckets should be 100", 100, total); } @Test public void testBalanceAfterNodeChanges() throws Exception { - // Setup: 3 nodes - haContainer.registerAndElect(); - Configurations config2 = createNodeConfig("127.0.0.2", 1262, 1632); - HighAvailabilityContainer haContainer2 = createContainerWithMockZk(config2); - haContainer2.registerAndElect(); - Configurations config3 = createNodeConfig("127.0.0.3", 1263, 1633); - HighAvailabilityContainer haContainer3 = createContainerWithMockZk(config3); - haContainer3.registerAndElect(); - - try { - Thread.sleep(200); - - // Initial assignment - assignService.doAssign(); - - // Verify balance - Map> assignments = mockAssignStore.getAllAssignments(); - Assert.assertEquals("Should have 3 nodes", 3, assignments.size()); - - List bucketCounts = new ArrayList<>(); - for (List buckets : assignments.values()) { - bucketCounts.add(buckets.size()); - } - int max = bucketCounts.stream().mapToInt(Integer::intValue).max().orElse(0); - int min = bucketCounts.stream().mapToInt(Integer::intValue).min().orElse(0); - Assert.assertTrue("Difference should be at most 1", max - min <= 1); + mockAssignStore.registerNode(node2); + mockAssignStore.registerNode(node3); - // Verify all buckets are assigned - int total = bucketCounts.stream().mapToInt(Integer::intValue).sum(); - Assert.assertEquals("All buckets should be assigned", 100, total); - } finally { - haContainer2.close(); - haContainer3.close(); + assignService.doAssign(); + + Map> assignments = mockAssignStore.getAllAssignments(); + Assert.assertEquals("Should have 3 nodes", 3, assignments.size()); + + List bucketCounts = new ArrayList<>(); + for (List buckets : assignments.values()) { + bucketCounts.add(buckets.size()); } + int max = bucketCounts.stream().mapToInt(Integer::intValue).max().orElse(0); + int min = bucketCounts.stream().mapToInt(Integer::intValue).min().orElse(0); + Assert.assertTrue("Difference should be at most 1", max - min <= 1); + + int total = bucketCounts.stream().mapToInt(Integer::intValue).sum(); + Assert.assertEquals("All buckets should be assigned", 100, total); } @Test public void testIncrementalAssignmentMinimizesMigration() throws Exception { - // Setup: 2 nodes initially - haContainer.registerAndElect(); - Configurations config2 = createNodeConfig("127.0.0.2", 1262, 1632); - HighAvailabilityContainer haContainer2 = createContainerWithMockZk(config2); - haContainer2.registerAndElect(); - HighAvailabilityContainer haContainer3 = null; - - try { - Thread.sleep(100); - - // Initial assignment - assignService.doAssign(); - Map> initialAssignments = mockAssignStore.getAllAssignments(); - - // Record initial assignments - Set node1InitialBuckets = new HashSet<>(); - Set node2InitialBuckets = new HashSet<>(); - for (Map.Entry> entry : initialAssignments.entrySet()) { - if (entry.getKey().getHost().equals("127.0.0.1")) { - node1InitialBuckets.addAll(entry.getValue()); - } else { - node2InitialBuckets.addAll(entry.getValue()); - } - } + mockAssignStore.registerNode(node2); - // Add new node - Configurations config3 = createNodeConfig("127.0.0.3", 1263, 1633); - haContainer3 = createContainerWithMockZk(config3); - haContainer3.registerAndElect(); - - Thread.sleep(100); - - // Trigger reassignment - assignService.doAssign(); - - // Check new assignments - Map> newAssignments = mockAssignStore.getAllAssignments(); - - // Calculate migration: buckets that moved from node1 or node2 - Set node1NewBuckets = new HashSet<>(); - Set node2NewBuckets = new HashSet<>(); - Set node3Buckets = new HashSet<>(); - for (Map.Entry> entry : newAssignments.entrySet()) { - if (entry.getKey().getHost().equals("127.0.0.1")) { - node1NewBuckets.addAll(entry.getValue()); - } else if (entry.getKey().getHost().equals("127.0.0.2")) { - node2NewBuckets.addAll(entry.getValue()); - } else { - node3Buckets.addAll(entry.getValue()); - } + // Initial assignment + assignService.doAssign(); + Map> initialAssignments = mockAssignStore.getAllAssignments(); + + Set node1InitialBuckets = new HashSet<>(); + Set node2InitialBuckets = new HashSet<>(); + for (Map.Entry> entry : initialAssignments.entrySet()) { + if (entry.getKey().getHost().equals("127.0.0.1")) { + node1InitialBuckets.addAll(entry.getValue()); + } else { + node2InitialBuckets.addAll(entry.getValue()); } + } - // Node1 and Node2 should keep most of their buckets - Set node1Kept = new HashSet<>(node1InitialBuckets); - node1Kept.retainAll(node1NewBuckets); - Set node2Kept = new HashSet<>(node2InitialBuckets); - node2Kept.retainAll(node2NewBuckets); - - // Verify incremental assignment: nodes should keep most buckets - Assert.assertTrue( - "Node1 should keep most buckets (incremental)", - node1Kept.size() > node1InitialBuckets.size() / 2); - Assert.assertTrue( - "Node2 should keep most buckets (incremental)", - node2Kept.size() > node2InitialBuckets.size() / 2); - - // Node3 should get buckets from both - Assert.assertTrue("Node3 should have buckets", node3Buckets.size() > 0); - } finally { - haContainer2.close(); - if (haContainer3 != null) { - try { - haContainer3.close(); - } catch (Exception e) { - // ignore - } + // Add new node + mockAssignStore.registerNode(node3); + + // Trigger reassignment + assignService.doAssign(); + + Map> newAssignments = mockAssignStore.getAllAssignments(); + + Set node1NewBuckets = new HashSet<>(); + Set node2NewBuckets = new HashSet<>(); + Set node3Buckets = new HashSet<>(); + for (Map.Entry> entry : newAssignments.entrySet()) { + if (entry.getKey().getHost().equals("127.0.0.1")) { + node1NewBuckets.addAll(entry.getValue()); + } else if (entry.getKey().getHost().equals("127.0.0.2")) { + node2NewBuckets.addAll(entry.getValue()); + } else { + node3Buckets.addAll(entry.getValue()); } } + + Set node1Kept = new HashSet<>(node1InitialBuckets); + node1Kept.retainAll(node1NewBuckets); + Set node2Kept = new HashSet<>(node2InitialBuckets); + node2Kept.retainAll(node2NewBuckets); + + Assert.assertTrue( + "Node1 should keep most buckets (incremental)", + node1Kept.size() > node1InitialBuckets.size() / 2); + Assert.assertTrue( + "Node2 should keep most buckets (incremental)", + node2Kept.size() > node2InitialBuckets.size() / 2); + Assert.assertTrue("Node3 should have buckets", node3Buckets.size() > 0); } @Test public void testOfflineNodeWithMissingLastUpdateTime() throws Exception { - // Verify that a node absent from the alive list with lastUpdateTime == 0 - // is still reclaimed instead of being stranded forever. - haContainer.registerAndElect(); - Configurations config2 = createNodeConfig("127.0.0.2", 1262, 1632); - HighAvailabilityContainer haContainer2 = createContainerWithMockZk(config2); - haContainer2.registerAndElect(); - - try { - Thread.sleep(100); - - // Initial assignment — both nodes get buckets - assignService.doAssign(); - Map> initialAssignments = mockAssignStore.getAllAssignments(); - Assert.assertEquals("Should have 2 nodes", 2, initialAssignments.size()); - - // Simulate node2 going offline (remove from ZK) - mockZkState.deleteNodeByHost("127.0.0.2"); - - // Clear node2's lastUpdateTime to simulate the edge case where - // saveAssignments wrote the assignment but crashed before - // updateLastUpdateTime completed (two non-atomic ZK writes). - mockAssignStore.clearLastUpdateTime("127.0.0.2"); - - Thread.sleep(100); - - // Trigger reassignment — node2 should be reclaimed even without a timestamp - assignService.doAssign(); - - Map> newAssignments = mockAssignStore.getAllAssignments(); - Assert.assertEquals( - "Node with missing lastUpdateTime should be reclaimed", 1, newAssignments.size()); - - List remainingBuckets = newAssignments.values().iterator().next(); - Assert.assertEquals("All buckets should be redistributed", 100, remainingBuckets.size()); - } finally { - try { - haContainer2.close(); - } catch (Exception e) { - // ignore - } - } + mockAssignStore.registerNode(node2); + + // Initial assignment + assignService.doAssign(); + Map> initialAssignments = mockAssignStore.getAllAssignments(); + Assert.assertEquals("Should have 2 nodes", 2, initialAssignments.size()); + + // Simulate node2 going offline (remove from alive list but keep assignments) + mockAssignStore.simulateOffline(node2); + + // Clear node2's lastUpdateTime to simulate the edge case + mockAssignStore.clearLastUpdateTime("127.0.0.2"); + + // Trigger reassignment + assignService.doAssign(); + + Map> newAssignments = mockAssignStore.getAllAssignments(); + Assert.assertEquals( + "Node with missing lastUpdateTime should be reclaimed", 1, newAssignments.size()); + + List remainingBuckets = newAssignments.values().iterator().next(); + Assert.assertEquals("All buckets should be redistributed", 100, remainingBuckets.size()); } @Test @@ -438,30 +316,13 @@ public void testServiceStartStop() { @Test public void testServiceSkipsWhenNotLeader() throws Exception { - // Create a non-leader container - mockLeaderLatch = createMockLeaderLatch(false); // Not leader - Configurations nonLeaderConfig = createNodeConfig("127.0.0.2", 1262, 1632); - HighAvailabilityContainer nonLeaderContainer = createContainerWithMockZk(nonLeaderConfig); - nonLeaderContainer.registerAndElect(); - - try { - // Wait a bit - Thread.sleep(100); - - AmsAssignService nonLeaderService = createAssignServiceWithMockStore(nonLeaderContainer); - - // Should not throw exception even if not leader - nonLeaderService.doAssign(); - - // Should not have assignments if not leader - Map> assignments = mockAssignStore.getAllAssignments(); - // Verify that non-leader doesn't create assignments - Assert.assertTrue( - "Non-leader should not create assignments", - assignments.isEmpty() || assignments.size() == 0); - } finally { - nonLeaderContainer.close(); - } + // AmsAssignService no longer checks hasLeadership() internally — Leader gating is done by + // AmoroServiceContainer.startLeaderServices/stopLeaderServices. This test verifies that + // doAssign() works correctly when called. + AmsAssignService nonLeaderService = createAssignServiceWithMockStore(); + + // doAssign() should execute without error + nonLeaderService.doAssign(); } private Configurations createNodeConfig(String host, int thriftPort, int httpPort) { @@ -592,13 +453,7 @@ private AmsServerInfo buildServerInfo(String host, Integer thriftPort, Integer h /** Create AmsAssignService with mock BucketAssignStore. */ private AmsAssignService createAssignServiceWithMockStore() throws Exception { - return createAssignServiceWithMockStore(haContainer); - } - - /** Create AmsAssignService with mock BucketAssignStore. */ - private AmsAssignService createAssignServiceWithMockStore(HighAvailabilityContainer container) - throws Exception { - return new AmsAssignService(container, serviceConfig, mockAssignStore); + return new AmsAssignService(serviceConfig, mockAssignStore); } /** Create a mock CuratorFramework that uses MockZkState for storage. */ @@ -920,12 +775,30 @@ public void updateLastUpdateTime(AmsServerInfo nodeInfo) throws BucketAssignStor @Override public List getAliveNodes() throws BucketAssignStoreException { List nodes = new ArrayList<>(); - for (String nodeKey : assignments.keySet()) { - nodes.add(nodeInfoMap.getOrDefault(nodeKey, parseNodeKey(nodeKey))); + for (String nodeKey : nodeInfoMap.keySet()) { + nodes.add(nodeInfoMap.get(nodeKey)); } return nodes; } + @Override + public void registerNode(AmsServerInfo serverInfo) throws BucketAssignStoreException { + String nodeKey = getNodeKey(serverInfo); + nodeInfoMap.put(nodeKey, serverInfo); + lastUpdateTimes.put(nodeKey, System.currentTimeMillis()); + } + + @Override + public void removeNode(AmsServerInfo serverInfo) throws BucketAssignStoreException { + removeAssignments(serverInfo); + } + + /** Simulate a node going offline: remove from alive list but keep its assignments. */ + void simulateOffline(AmsServerInfo serverInfo) { + String nodeKey = getNodeKey(serverInfo); + nodeInfoMap.remove(nodeKey); + } + /** * Clear lastUpdateTime for all nodes matching the given host, simulating a missing timestamp. */ @@ -949,5 +822,10 @@ private AmsServerInfo parseNodeKey(String nodeKey) { nodeInfo.setThriftBindPort(Integer.parseInt(parts[1])); return nodeInfo; } + + @Override + public void close() { + // No resources to release + } } } diff --git a/amoro-ams/src/test/java/org/apache/amoro/server/TestHighAvailabilityContainer.java b/amoro-ams/src/test/java/org/apache/amoro/server/TestHighAvailabilityContainer.java index 7a8411ae8f..53064a8b1e 100644 --- a/amoro-ams/src/test/java/org/apache/amoro/server/TestHighAvailabilityContainer.java +++ b/amoro-ams/src/test/java/org/apache/amoro/server/TestHighAvailabilityContainer.java @@ -19,50 +19,33 @@ package org.apache.amoro.server; import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; -import org.apache.amoro.client.AmsServerInfo; import org.apache.amoro.config.Configurations; -import org.apache.amoro.properties.AmsHAProperties; import org.apache.amoro.server.ha.HighAvailabilityContainer; import org.apache.amoro.server.ha.ZkHighAvailabilityContainer; import org.apache.amoro.shade.zookeeper3.org.apache.curator.framework.CuratorFramework; import org.apache.amoro.shade.zookeeper3.org.apache.curator.framework.recipes.leader.LeaderLatch; -import org.apache.amoro.shade.zookeeper3.org.apache.zookeeper.CreateMode; -import org.apache.amoro.shade.zookeeper3.org.apache.zookeeper.KeeperException; -import org.apache.amoro.shade.zookeeper3.org.apache.zookeeper.data.Stat; -import org.apache.amoro.utils.JacksonUtil; import org.junit.After; import org.junit.Assert; import org.junit.Before; import org.junit.Test; -import java.nio.charset.StandardCharsets; -import java.util.ArrayList; -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import java.util.concurrent.atomic.AtomicInteger; - -/** Test for HighAvailabilityContainer using mocked ZK to avoid connection issues. */ +/** Test for HighAvailabilityContainer leader election using mocked ZK. */ public class TestHighAvailabilityContainer { private Configurations serviceConfig; private HighAvailabilityContainer haContainer; - private MockZkState mockZkState; private CuratorFramework mockZkClient; private LeaderLatch mockLeaderLatch; @Before public void setUp() throws Exception { - mockZkState = new MockZkState(); mockZkClient = createMockZkClient(); mockLeaderLatch = createMockLeaderLatch(); - // Create test configuration serviceConfig = new Configurations(); serviceConfig.setString(AmoroManagementConf.SERVER_EXPOSE_HOST, "127.0.0.1"); serviceConfig.setInteger(AmoroManagementConf.TABLE_SERVICE_THRIFT_BIND_PORT, 1260); @@ -78,200 +61,33 @@ public void tearDown() throws Exception { if (haContainer != null) { haContainer.close(); } - mockZkState.clear(); - } - - @Test - public void testRegistAndElectWithoutMasterSlaveMode() throws Exception { - // Test that node registration is skipped when master-slave mode is disabled - serviceConfig.setBoolean(AmoroManagementConf.HA_USE_MASTER_SLAVE_MODE, false); - haContainer = createContainerWithMockZk(); - - // Should not throw exception and should not register node - haContainer.registerAndElect(); - - // Verify no node was registered - String nodesPath = AmsHAProperties.getNodesPath("test-cluster"); - List children = mockZkState.getChildren(nodesPath); - Assert.assertEquals( - "No nodes should be registered when master-slave mode is disabled", 0, children.size()); - } - - @Test - public void testRegistAndElectWithMasterSlaveMode() throws Exception { - // Test that node registration works when master-slave mode is enabled - serviceConfig.setBoolean(AmoroManagementConf.HA_USE_MASTER_SLAVE_MODE, true); - haContainer = createContainerWithMockZk(); - - // Register node - haContainer.registerAndElect(); - - // Verify node was registered - String nodesPath = AmsHAProperties.getNodesPath("test-cluster"); - List children = mockZkState.getChildren(nodesPath); - Assert.assertEquals("One node should be registered", 1, children.size()); - - // Verify node data - String nodePath = nodesPath + "/" + children.get(0); - byte[] data = mockZkState.getData(nodePath); - Assert.assertNotNull("Node data should not be null", data); - Assert.assertTrue("Node data should not be empty", data.length > 0); - - // Verify node info - String nodeInfoJson = new String(data, StandardCharsets.UTF_8); - AmsServerInfo nodeInfo = JacksonUtil.parseObject(nodeInfoJson, AmsServerInfo.class); - Assert.assertEquals("Host should match", "127.0.0.1", nodeInfo.getHost()); - Assert.assertEquals( - "Thrift port should match", Integer.valueOf(1261), nodeInfo.getThriftBindPort()); - } - - @Test - public void testGetAliveNodesWithoutMasterSlaveMode() throws Exception { - // Test that getAliveNodes returns empty list when master-slave mode is disabled - serviceConfig.setBoolean(AmoroManagementConf.HA_USE_MASTER_SLAVE_MODE, false); - haContainer = createContainerWithMockZk(); - - List aliveNodes = haContainer.getAliveNodes(); - Assert.assertNotNull("Alive nodes list should not be null", aliveNodes); - Assert.assertEquals( - "Alive nodes list should be empty when master-slave mode is disabled", - 0, - aliveNodes.size()); - } - - @Test - public void testGetAliveNodesWhenNotLeader() throws Exception { - // Test that getAliveNodes returns empty list when not leader - serviceConfig.setBoolean(AmoroManagementConf.HA_USE_MASTER_SLAVE_MODE, true); - mockLeaderLatch = createMockLeaderLatch(false); // Not leader - haContainer = createContainerWithMockZk(); - - // Register node - haContainer.registerAndElect(); - - // Since we're not the leader, should return empty list - List aliveNodes = haContainer.getAliveNodes(); - Assert.assertNotNull("Alive nodes list should not be null", aliveNodes); - Assert.assertEquals("Alive nodes list should be empty when not leader", 0, aliveNodes.size()); - } - - @Test - public void testGetAliveNodesAsLeader() throws Exception { - // Test that getAliveNodes returns nodes when leader - serviceConfig.setBoolean(AmoroManagementConf.HA_USE_MASTER_SLAVE_MODE, true); - mockLeaderLatch = createMockLeaderLatch(true); // Is leader - haContainer = createContainerWithMockZk(); - - // Register node - haContainer.registerAndElect(); - - // Verify we are leader - Assert.assertTrue("Should be leader", haContainer.hasLeadership()); - - // Get alive nodes - List aliveNodes = haContainer.getAliveNodes(); - Assert.assertNotNull("Alive nodes list should not be null", aliveNodes); - Assert.assertEquals("Should have one alive node", 1, aliveNodes.size()); - - // Verify node info - AmsServerInfo nodeInfo = aliveNodes.get(0); - Assert.assertEquals("Host should match", "127.0.0.1", nodeInfo.getHost()); - Assert.assertEquals( - "Thrift port should match", Integer.valueOf(1261), nodeInfo.getThriftBindPort()); - Assert.assertEquals( - "HTTP port should match", Integer.valueOf(1630), nodeInfo.getRestBindPort()); - } - - @Test - public void testGetAliveNodesWithMultipleNodes() throws Exception { - // Test that getAliveNodes returns all registered nodes - serviceConfig.setBoolean(AmoroManagementConf.HA_USE_MASTER_SLAVE_MODE, true); - mockLeaderLatch = createMockLeaderLatch(true); // Is leader - haContainer = createContainerWithMockZk(); - - // Register first node - haContainer.registerAndElect(); - - // Verify first node was registered - String nodesPath = AmsHAProperties.getNodesPath("test-cluster"); - List childrenAfterFirst = mockZkState.getChildren(nodesPath); - Assert.assertEquals("First node should be registered", 1, childrenAfterFirst.size()); - - // Register second node manually in mock state - // Use createNode with sequential path to get the correct sequence number - AmsServerInfo nodeInfo2 = new AmsServerInfo(); - nodeInfo2.setHost("127.0.0.2"); - nodeInfo2.setThriftBindPort(1262); - nodeInfo2.setRestBindPort(1631); - String nodeInfo2Json = JacksonUtil.toJSONString(nodeInfo2); - // Use sequential path ending with "-" to let createNode generate the sequence number - // This ensures the second node gets the correct sequence number (0000000001) - mockZkState.createNode(nodesPath + "/node-", nodeInfo2Json.getBytes(StandardCharsets.UTF_8)); - - // Get alive nodes - List aliveNodes = haContainer.getAliveNodes(); - Assert.assertNotNull("Alive nodes list should not be null", aliveNodes); - Assert.assertEquals("Should have two alive nodes", 2, aliveNodes.size()); - } - - @Test - public void testCloseUnregistersNode() throws Exception { - // Test that close() unregisters the node - serviceConfig.setBoolean(AmoroManagementConf.HA_USE_MASTER_SLAVE_MODE, true); - haContainer = createContainerWithMockZk(); - - // Register node - haContainer.registerAndElect(); - - // Verify node was registered - String nodesPath = AmsHAProperties.getNodesPath("test-cluster"); - List children = mockZkState.getChildren(nodesPath); - Assert.assertEquals("One node should be registered", 1, children.size()); - - // Close container - haContainer.close(); - haContainer = null; - - // Verify node was unregistered - List childrenAfterClose = mockZkState.getChildren(nodesPath); - Assert.assertEquals("No nodes should be registered after close", 0, childrenAfterClose.size()); } @Test public void testHasLeadership() throws Exception { - // Test hasLeadership() method serviceConfig.setBoolean(AmoroManagementConf.HA_USE_MASTER_SLAVE_MODE, true); - mockLeaderLatch = createMockLeaderLatch(false); // Not leader initially + mockLeaderLatch = createMockLeaderLatch(false); haContainer = createContainerWithMockZk(); - // Initially should not be leader Assert.assertFalse("Should not be leader initially", haContainer.hasLeadership()); - // Change to leader mockLeaderLatch = createMockLeaderLatch(true); haContainer = createContainerWithMockZk(); - // Should be leader now Assert.assertTrue("Should be leader", haContainer.hasLeadership()); } @Test - public void testRegistAndElectWithoutHAEnabled() throws Exception { - // Test that registAndElect skips when HA is not enabled + public void testCreateWithoutHAEnabled() throws Exception { serviceConfig.setBoolean(AmoroManagementConf.HA_ENABLE, false); serviceConfig.setBoolean(AmoroManagementConf.HA_USE_MASTER_SLAVE_MODE, true); haContainer = new ZkHighAvailabilityContainer(serviceConfig); - - // Should not throw exception - haContainer.registerAndElect(); } /** Create HighAvailabilityContainer with mocked ZK components using reflection. */ private HighAvailabilityContainer createContainerWithMockZk() throws Exception { - // Create container without ZK connection to avoid any connection attempts HighAvailabilityContainer container = createContainerWithoutZk(); - // Inject mock ZK client and leader latch java.lang.reflect.Field zkClientField = ZkHighAvailabilityContainer.class.getDeclaredField("zkClient"); zkClientField.setAccessible(true); @@ -282,282 +98,46 @@ private HighAvailabilityContainer createContainerWithMockZk() throws Exception { leaderLatchField.setAccessible(true); leaderLatchField.set(container, mockLeaderLatch); - // Note: We don't need to create the paths themselves as nodes in ZK - // ZK paths are logical containers, not actual nodes - // The createPathIfNeeded() calls will be handled by the mock when needed - return container; } - /** - * Create a HighAvailabilityContainer without initializing ZK connection. This is used when we - * want to completely avoid ZK connection attempts. - */ private HighAvailabilityContainer createContainerWithoutZk() throws Exception { - // Use reflection to create ZkHighAvailabilityContainer without calling constructor java.lang.reflect.Constructor constructor = ZkHighAvailabilityContainer.class.getDeclaredConstructor(Configurations.class); - // Create a minimal config that disables HA to avoid ZK connection Configurations tempConfig = new Configurations(serviceConfig); tempConfig.setBoolean(AmoroManagementConf.HA_ENABLE, false); HighAvailabilityContainer container = constructor.newInstance(tempConfig); - // Now set all required fields using reflection java.lang.reflect.Field isMasterSlaveModeField = ZkHighAvailabilityContainer.class.getDeclaredField("isMasterSlaveMode"); isMasterSlaveModeField.setAccessible(true); isMasterSlaveModeField.set( container, serviceConfig.getBoolean(AmoroManagementConf.HA_USE_MASTER_SLAVE_MODE)); - if (serviceConfig.getBoolean(AmoroManagementConf.HA_ENABLE)) { - String haClusterName = serviceConfig.getString(AmoroManagementConf.HA_CLUSTER_NAME); - - java.lang.reflect.Field tableServiceMasterPathField = - ZkHighAvailabilityContainer.class.getDeclaredField("tableServiceMasterPath"); - tableServiceMasterPathField.setAccessible(true); - tableServiceMasterPathField.set( - container, AmsHAProperties.getTableServiceMasterPath(haClusterName)); - - java.lang.reflect.Field optimizingServiceMasterPathField = - ZkHighAvailabilityContainer.class.getDeclaredField("optimizingServiceMasterPath"); - optimizingServiceMasterPathField.setAccessible(true); - optimizingServiceMasterPathField.set( - container, AmsHAProperties.getOptimizingServiceMasterPath(haClusterName)); - - java.lang.reflect.Field nodesPathField = - ZkHighAvailabilityContainer.class.getDeclaredField("nodesPath"); - nodesPathField.setAccessible(true); - nodesPathField.set(container, AmsHAProperties.getNodesPath(haClusterName)); - - java.lang.reflect.Field tableServiceServerInfoField = - ZkHighAvailabilityContainer.class.getDeclaredField("tableServiceServerInfo"); - tableServiceServerInfoField.setAccessible(true); - AmsServerInfo tableServiceServerInfo = - buildServerInfo( - serviceConfig.getString(AmoroManagementConf.SERVER_EXPOSE_HOST), - serviceConfig.getInteger(AmoroManagementConf.TABLE_SERVICE_THRIFT_BIND_PORT), - serviceConfig.getInteger(AmoroManagementConf.HTTP_SERVER_PORT)); - tableServiceServerInfoField.set(container, tableServiceServerInfo); - - java.lang.reflect.Field optimizingServiceServerInfoField = - ZkHighAvailabilityContainer.class.getDeclaredField("optimizingServiceServerInfo"); - optimizingServiceServerInfoField.setAccessible(true); - AmsServerInfo optimizingServiceServerInfo = - buildServerInfo( - serviceConfig.getString(AmoroManagementConf.SERVER_EXPOSE_HOST), - serviceConfig.getInteger(AmoroManagementConf.OPTIMIZING_SERVICE_THRIFT_BIND_PORT), - serviceConfig.getInteger(AmoroManagementConf.HTTP_SERVER_PORT)); - optimizingServiceServerInfoField.set(container, optimizingServiceServerInfo); - } - return container; } - /** Helper method to build AmsServerInfo (copied from HighAvailabilityContainer). */ - private AmsServerInfo buildServerInfo(String host, Integer thriftPort, Integer httpPort) { - AmsServerInfo serverInfo = new AmsServerInfo(); - serverInfo.setHost(host); - serverInfo.setThriftBindPort(thriftPort); - serverInfo.setRestBindPort(httpPort); - return serverInfo; - } - - /** Create a mock CuratorFramework that uses MockZkState for storage. */ @SuppressWarnings("unchecked") private CuratorFramework createMockZkClient() throws Exception { CuratorFramework mockClient = mock(CuratorFramework.class); - - // Mock getChildren() - create a chain of mocks - org.apache.amoro.shade.zookeeper3.org.apache.curator.framework.api.GetChildrenBuilder - getChildrenBuilder = - mock( - org.apache.amoro.shade.zookeeper3.org.apache.curator.framework.api - .GetChildrenBuilder.class); - when(mockClient.getChildren()).thenReturn(getChildrenBuilder); - when(getChildrenBuilder.forPath(anyString())) - .thenAnswer( - invocation -> { - String path = invocation.getArgument(0); - return mockZkState.getChildren(path); - }); - - // Mock getData() - org.apache.amoro.shade.zookeeper3.org.apache.curator.framework.api.GetDataBuilder - getDataBuilder = - mock( - org.apache.amoro.shade.zookeeper3.org.apache.curator.framework.api.GetDataBuilder - .class); - when(mockClient.getData()).thenReturn(getDataBuilder); - when(getDataBuilder.forPath(anyString())) - .thenAnswer( - invocation -> { - String path = invocation.getArgument(0); - return mockZkState.getData(path); - }); - - // Mock create() - manually create the entire fluent API chain to ensure consistency - org.apache.amoro.shade.zookeeper3.org.apache.curator.framework.api.CreateBuilder createBuilder = - mock( - org.apache.amoro.shade.zookeeper3.org.apache.curator.framework.api.CreateBuilder.class); - - @SuppressWarnings("unchecked") - org.apache.amoro.shade.zookeeper3.org.apache.curator.framework.api - .ProtectACLCreateModeStatPathAndBytesable< - String> - pathAndBytesable = - mock( - org.apache.amoro.shade.zookeeper3.org.apache.curator.framework.api - .ProtectACLCreateModeStatPathAndBytesable.class); - - when(mockClient.create()).thenReturn(createBuilder); - - // Mock the chain: creatingParentsIfNeeded() -> withMode() -> forPath() - // Use the same mock object for the entire chain - when(createBuilder.creatingParentsIfNeeded()).thenReturn(pathAndBytesable); - when(pathAndBytesable.withMode(any(CreateMode.class))).thenReturn(pathAndBytesable); - - // Mock forPath(path, data) - used by registAndElect() - when(pathAndBytesable.forPath(anyString(), any(byte[].class))) - .thenAnswer( - invocation -> { - String path = invocation.getArgument(0); - byte[] data = invocation.getArgument(1); - return mockZkState.createNode(path, data); - }); - - // Mock forPath(path) - used by createPathIfNeeded() - // Note: createPathIfNeeded() creates paths without data, but we still need to store them - // so that getChildren() can work correctly - when(pathAndBytesable.forPath(anyString())) - .thenAnswer( - invocation -> { - String path = invocation.getArgument(0); - // Create the path as an empty node (this simulates ZK path creation) - // In real ZK, paths are logical containers, but we need to store them - // to make getChildren() work correctly - if (mockZkState.exists(path) == null) { - mockZkState.createNode(path, new byte[0]); - } - return null; - }); - - // Mock delete() - org.apache.amoro.shade.zookeeper3.org.apache.curator.framework.api.DeleteBuilder deleteBuilder = - mock( - org.apache.amoro.shade.zookeeper3.org.apache.curator.framework.api.DeleteBuilder.class); - when(mockClient.delete()).thenReturn(deleteBuilder); - doAnswer( - invocation -> { - String path = invocation.getArgument(0); - mockZkState.deleteNode(path); - return null; - }) - .when(deleteBuilder) - .forPath(anyString()); - - // Mock checkExists() - @SuppressWarnings("unchecked") - org.apache.amoro.shade.zookeeper3.org.apache.curator.framework.api.ExistsBuilder - checkExistsBuilder = - mock( - org.apache.amoro.shade.zookeeper3.org.apache.curator.framework.api.ExistsBuilder - .class); - when(mockClient.checkExists()).thenReturn(checkExistsBuilder); - when(checkExistsBuilder.forPath(anyString())) - .thenAnswer( - invocation -> { - String path = invocation.getArgument(0); - return mockZkState.exists(path); - }); - - // Mock start() and close() doAnswer(invocation -> null).when(mockClient).start(); doAnswer(invocation -> null).when(mockClient).close(); - return mockClient; } - /** Create a mock LeaderLatch. */ private LeaderLatch createMockLeaderLatch() throws Exception { return createMockLeaderLatch(true); } - /** Create a mock LeaderLatch with specified leadership status. */ private LeaderLatch createMockLeaderLatch(boolean hasLeadership) throws Exception { LeaderLatch mockLatch = mock(LeaderLatch.class); when(mockLatch.hasLeadership()).thenReturn(hasLeadership); doAnswer(invocation -> null).when(mockLatch).addListener(any()); doAnswer(invocation -> null).when(mockLatch).start(); doAnswer(invocation -> null).when(mockLatch).close(); - // Mock await() - it throws IOException and InterruptedException - doAnswer( - invocation -> { - // Mock implementation - doesn't actually wait - return null; - }) - .when(mockLatch) - .await(); + doAnswer(invocation -> null).when(mockLatch).await(); return mockLatch; } - - /** In-memory ZK state simulator. */ - private static class MockZkState { - private final Map nodes = new HashMap<>(); - private final AtomicInteger sequenceCounter = new AtomicInteger(0); - - public List getChildren(String path) throws KeeperException { - List children = new ArrayList<>(); - String prefix = path.endsWith("/") ? path : path + "/"; - for (String nodePath : nodes.keySet()) { - // Only include direct children (not the path itself, and not nested paths) - if (nodePath.startsWith(prefix) && !nodePath.equals(path)) { - String relativePath = nodePath.substring(prefix.length()); - // Only add direct children (no additional slashes) - // This means the path should be exactly: prefix + relativePath - if (!relativePath.contains("/")) { - children.add(relativePath); - } - } - } - // Sort to ensure consistent ordering - children.sort(String::compareTo); - return children; - } - - public byte[] getData(String path) throws KeeperException { - byte[] data = nodes.get(path); - if (data == null) { - throw new KeeperException.NoNodeException(path); - } - return data; - } - - public String createNode(String path, byte[] data) { - // Handle sequential nodes - if (path.endsWith("-")) { - int seq = sequenceCounter.incrementAndGet(); - path = path + String.format("%010d", seq); - } - nodes.put(path, data); - return path; - } - - public void deleteNode(String path) throws KeeperException { - if (!nodes.containsKey(path)) { - throw new KeeperException.NoNodeException(path); - } - nodes.remove(path); - } - - public Stat exists(String path) { - return nodes.containsKey(path) ? new Stat() : null; - } - - public void clear() { - nodes.clear(); - sequenceCounter.set(0); - } - } } 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 index 97b0d6d5c4..fbe1b2f63d 100644 --- 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 @@ -19,49 +19,32 @@ package org.apache.amoro.server.ha; import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; -import org.apache.amoro.client.AmsServerInfo; import org.apache.amoro.config.Configurations; -import org.apache.amoro.properties.AmsHAProperties; import org.apache.amoro.server.AmoroManagementConf; import org.apache.amoro.shade.zookeeper3.org.apache.curator.framework.CuratorFramework; import org.apache.amoro.shade.zookeeper3.org.apache.curator.framework.recipes.leader.LeaderLatch; -import org.apache.amoro.shade.zookeeper3.org.apache.zookeeper.CreateMode; -import org.apache.amoro.shade.zookeeper3.org.apache.zookeeper.KeeperException; -import org.apache.amoro.shade.zookeeper3.org.apache.zookeeper.data.Stat; -import org.apache.amoro.utils.JacksonUtil; import org.junit.After; import org.junit.Assert; import org.junit.Before; import org.junit.Test; -import java.nio.charset.StandardCharsets; -import java.util.ArrayList; -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import java.util.concurrent.atomic.AtomicInteger; - -/** Test for HighAvailabilityContainer using mocked ZK to avoid connection issues. */ +/** Test for ZkHighAvailabilityContainer leader election using mocked ZK. */ public class TestZkHighAvailabilityContainer { private Configurations serviceConfig; private HighAvailabilityContainer haContainer; - private MockZkState mockZkState; private CuratorFramework mockZkClient; private LeaderLatch mockLeaderLatch; @Before public void setUp() throws Exception { - mockZkState = new MockZkState(); mockZkClient = createMockZkClient(); mockLeaderLatch = createMockLeaderLatch(); - // Create test configuration serviceConfig = new Configurations(); serviceConfig.setString(AmoroManagementConf.SERVER_EXPOSE_HOST, "127.0.0.1"); serviceConfig.setInteger(AmoroManagementConf.TABLE_SERVICE_THRIFT_BIND_PORT, 1260); @@ -77,200 +60,32 @@ public void tearDown() throws Exception { if (haContainer != null) { haContainer.close(); } - mockZkState.clear(); - } - - @Test - public void testRegisterAndElectWithoutMasterSlaveMode() throws Exception { - // Test that node registration is skipped when master-slave mode is disabled - serviceConfig.setBoolean(AmoroManagementConf.HA_USE_MASTER_SLAVE_MODE, false); - haContainer = createContainerWithMockZk(); - - // Should not throw exception and should not register node - haContainer.registerAndElect(); - - // Verify no node was registered - String nodesPath = AmsHAProperties.getNodesPath("test-cluster"); - List children = mockZkState.getChildren(nodesPath); - Assert.assertEquals( - "No nodes should be registered when master-slave mode is disabled", 0, children.size()); - } - - @Test - public void testRegisterAndElectWithMasterSlaveMode() throws Exception { - // Test that node registration works when master-slave mode is enabled - serviceConfig.setBoolean(AmoroManagementConf.HA_USE_MASTER_SLAVE_MODE, true); - haContainer = createContainerWithMockZk(); - - // Register node - haContainer.registerAndElect(); - - // Verify node was registered - String nodesPath = AmsHAProperties.getNodesPath("test-cluster"); - List children = mockZkState.getChildren(nodesPath); - Assert.assertEquals("One node should be registered", 1, children.size()); - - // Verify node data - String nodePath = nodesPath + "/" + children.get(0); - byte[] data = mockZkState.getData(nodePath); - Assert.assertNotNull("Node data should not be null", data); - Assert.assertTrue("Node data should not be empty", data.length > 0); - - // Verify node info - String nodeInfoJson = new String(data, StandardCharsets.UTF_8); - AmsServerInfo nodeInfo = JacksonUtil.parseObject(nodeInfoJson, AmsServerInfo.class); - Assert.assertEquals("Host should match", "127.0.0.1", nodeInfo.getHost()); - Assert.assertEquals( - "Thrift port should match", Integer.valueOf(1261), nodeInfo.getThriftBindPort()); - } - - @Test - public void testGetAliveNodesWithoutMasterSlaveMode() throws Exception { - // Test that getAliveNodes returns empty list when master-slave mode is disabled - serviceConfig.setBoolean(AmoroManagementConf.HA_USE_MASTER_SLAVE_MODE, false); - haContainer = createContainerWithMockZk(); - - List aliveNodes = haContainer.getAliveNodes(); - Assert.assertNotNull("Alive nodes list should not be null", aliveNodes); - Assert.assertEquals( - "Alive nodes list should be empty when master-slave mode is disabled", - 0, - aliveNodes.size()); - } - - @Test - public void testGetAliveNodesWhenNotLeader() throws Exception { - // Test that getAliveNodes returns empty list when not leader - serviceConfig.setBoolean(AmoroManagementConf.HA_USE_MASTER_SLAVE_MODE, true); - mockLeaderLatch = createMockLeaderLatch(false); // Not leader - haContainer = createContainerWithMockZk(); - - // Register node - haContainer.registerAndElect(); - - // Since we're not the leader, should return empty list - List aliveNodes = haContainer.getAliveNodes(); - Assert.assertNotNull("Alive nodes list should not be null", aliveNodes); - Assert.assertEquals("Alive nodes list should be empty when not leader", 0, aliveNodes.size()); - } - - @Test - public void testGetAliveNodesAsLeader() throws Exception { - // Test that getAliveNodes returns nodes when leader - serviceConfig.setBoolean(AmoroManagementConf.HA_USE_MASTER_SLAVE_MODE, true); - mockLeaderLatch = createMockLeaderLatch(true); // Is leader - haContainer = createContainerWithMockZk(); - - // Register node - haContainer.registerAndElect(); - - // Verify we are leader - Assert.assertTrue("Should be leader", haContainer.hasLeadership()); - - // Get alive nodes - List aliveNodes = haContainer.getAliveNodes(); - Assert.assertNotNull("Alive nodes list should not be null", aliveNodes); - Assert.assertEquals("Should have one alive node", 1, aliveNodes.size()); - - // Verify node info - AmsServerInfo nodeInfo = aliveNodes.get(0); - Assert.assertEquals("Host should match", "127.0.0.1", nodeInfo.getHost()); - Assert.assertEquals( - "Thrift port should match", Integer.valueOf(1261), nodeInfo.getThriftBindPort()); - Assert.assertEquals( - "HTTP port should match", Integer.valueOf(1630), nodeInfo.getRestBindPort()); - } - - @Test - public void testGetAliveNodesWithMultipleNodes() throws Exception { - // Test that getAliveNodes returns all registered nodes - serviceConfig.setBoolean(AmoroManagementConf.HA_USE_MASTER_SLAVE_MODE, true); - mockLeaderLatch = createMockLeaderLatch(true); // Is leader - haContainer = createContainerWithMockZk(); - - // Register first node - haContainer.registerAndElect(); - - // Verify first node was registered - String nodesPath = AmsHAProperties.getNodesPath("test-cluster"); - List childrenAfterFirst = mockZkState.getChildren(nodesPath); - Assert.assertEquals("First node should be registered", 1, childrenAfterFirst.size()); - - // Register second node manually in mock state - // Use createNode with sequential path to get the correct sequence number - AmsServerInfo nodeInfo2 = new AmsServerInfo(); - nodeInfo2.setHost("127.0.0.2"); - nodeInfo2.setThriftBindPort(1262); - nodeInfo2.setRestBindPort(1631); - String nodeInfo2Json = JacksonUtil.toJSONString(nodeInfo2); - // Use sequential path ending with "-" to let createNode generate the sequence number - // This ensures the second node gets the correct sequence number (0000000001) - mockZkState.createNode(nodesPath + "/node-", nodeInfo2Json.getBytes(StandardCharsets.UTF_8)); - - // Get alive nodes - List aliveNodes = haContainer.getAliveNodes(); - Assert.assertNotNull("Alive nodes list should not be null", aliveNodes); - Assert.assertEquals("Should have two alive nodes", 2, aliveNodes.size()); - } - - @Test - public void testCloseUnregistersNode() throws Exception { - // Test that close() unregisters the node - serviceConfig.setBoolean(AmoroManagementConf.HA_USE_MASTER_SLAVE_MODE, true); - haContainer = createContainerWithMockZk(); - - // Register node - haContainer.registerAndElect(); - - // Verify node was registered - String nodesPath = AmsHAProperties.getNodesPath("test-cluster"); - List children = mockZkState.getChildren(nodesPath); - Assert.assertEquals("One node should be registered", 1, children.size()); - - // Close container - haContainer.close(); - haContainer = null; - - // Verify node was unregistered - List childrenAfterClose = mockZkState.getChildren(nodesPath); - Assert.assertEquals("No nodes should be registered after close", 0, childrenAfterClose.size()); } @Test public void testHasLeadership() throws Exception { - // Test hasLeadership() method serviceConfig.setBoolean(AmoroManagementConf.HA_USE_MASTER_SLAVE_MODE, true); - mockLeaderLatch = createMockLeaderLatch(false); // Not leader initially + mockLeaderLatch = createMockLeaderLatch(false); haContainer = createContainerWithMockZk(); - // Initially should not be leader Assert.assertFalse("Should not be leader initially", haContainer.hasLeadership()); - // Change to leader mockLeaderLatch = createMockLeaderLatch(true); haContainer = createContainerWithMockZk(); - // Should be leader now Assert.assertTrue("Should be leader", haContainer.hasLeadership()); } @Test - public void testRegisterAndElectWithoutHAEnabled() throws Exception { - // Test that registAndElect skips when HA is not enabled + public void testCreateWithoutHAEnabled() throws Exception { serviceConfig.setBoolean(AmoroManagementConf.HA_ENABLE, false); serviceConfig.setBoolean(AmoroManagementConf.HA_USE_MASTER_SLAVE_MODE, true); haContainer = new ZkHighAvailabilityContainer(serviceConfig); - - // Should not throw exception - haContainer.registerAndElect(); } - /** Create HighAvailabilityContainer with mocked ZK components using reflection. */ private HighAvailabilityContainer createContainerWithMockZk() throws Exception { - // Create container without ZK connection to avoid any connection attempts HighAvailabilityContainer container = createContainerWithoutZk(); - // Inject mock ZK client and leader latch (fields are on ZkHighAvailabilityContainer) java.lang.reflect.Field zkClientField = ZkHighAvailabilityContainer.class.getDeclaredField("zkClient"); zkClientField.setAccessible(true); @@ -281,286 +96,46 @@ private HighAvailabilityContainer createContainerWithMockZk() throws Exception { leaderLatchField.setAccessible(true); leaderLatchField.set(container, mockLeaderLatch); - // Note: We don't need to create the paths themselves as nodes in ZK - // ZK paths are logical containers, not actual nodes - // The createPathIfNeeded() calls will be handled by the mock when needed - return container; } - /** - * Create a HighAvailabilityContainer without initializing ZK connection. This is used when we - * want to completely avoid ZK connection attempts. - * - *

Uses ZkHighAvailabilityContainer (which has the constructor and fields); - * HighAvailabilityContainer is an interface without constructors or instance fields. - */ private HighAvailabilityContainer createContainerWithoutZk() throws Exception { - // ZkHighAvailabilityContainer has constructor (Configurations); HighAvailabilityContainer is an - // interface java.lang.reflect.Constructor constructor = ZkHighAvailabilityContainer.class.getDeclaredConstructor(Configurations.class); - // Create a minimal config that disables HA to avoid ZK connection Configurations tempConfig = new Configurations(serviceConfig); tempConfig.setBoolean(AmoroManagementConf.HA_ENABLE, false); HighAvailabilityContainer container = constructor.newInstance(tempConfig); - // Now set all required fields using reflection (fields are on ZkHighAvailabilityContainer) java.lang.reflect.Field isMasterSlaveModeField = ZkHighAvailabilityContainer.class.getDeclaredField("isMasterSlaveMode"); isMasterSlaveModeField.setAccessible(true); isMasterSlaveModeField.set( container, serviceConfig.getBoolean(AmoroManagementConf.HA_USE_MASTER_SLAVE_MODE)); - if (serviceConfig.getBoolean(AmoroManagementConf.HA_ENABLE)) { - String haClusterName = serviceConfig.getString(AmoroManagementConf.HA_CLUSTER_NAME); - - java.lang.reflect.Field tableServiceMasterPathField = - ZkHighAvailabilityContainer.class.getDeclaredField("tableServiceMasterPath"); - tableServiceMasterPathField.setAccessible(true); - tableServiceMasterPathField.set( - container, AmsHAProperties.getTableServiceMasterPath(haClusterName)); - - java.lang.reflect.Field optimizingServiceMasterPathField = - ZkHighAvailabilityContainer.class.getDeclaredField("optimizingServiceMasterPath"); - optimizingServiceMasterPathField.setAccessible(true); - optimizingServiceMasterPathField.set( - container, AmsHAProperties.getOptimizingServiceMasterPath(haClusterName)); - - java.lang.reflect.Field nodesPathField = - ZkHighAvailabilityContainer.class.getDeclaredField("nodesPath"); - nodesPathField.setAccessible(true); - nodesPathField.set(container, AmsHAProperties.getNodesPath(haClusterName)); - - java.lang.reflect.Field tableServiceServerInfoField = - ZkHighAvailabilityContainer.class.getDeclaredField("tableServiceServerInfo"); - tableServiceServerInfoField.setAccessible(true); - AmsServerInfo tableServiceServerInfo = - buildServerInfo( - serviceConfig.getString(AmoroManagementConf.SERVER_EXPOSE_HOST), - serviceConfig.getInteger(AmoroManagementConf.TABLE_SERVICE_THRIFT_BIND_PORT), - serviceConfig.getInteger(AmoroManagementConf.HTTP_SERVER_PORT)); - tableServiceServerInfoField.set(container, tableServiceServerInfo); - - java.lang.reflect.Field optimizingServiceServerInfoField = - ZkHighAvailabilityContainer.class.getDeclaredField("optimizingServiceServerInfo"); - optimizingServiceServerInfoField.setAccessible(true); - AmsServerInfo optimizingServiceServerInfo = - buildServerInfo( - serviceConfig.getString(AmoroManagementConf.SERVER_EXPOSE_HOST), - serviceConfig.getInteger(AmoroManagementConf.OPTIMIZING_SERVICE_THRIFT_BIND_PORT), - serviceConfig.getInteger(AmoroManagementConf.HTTP_SERVER_PORT)); - optimizingServiceServerInfoField.set(container, optimizingServiceServerInfo); - } - return container; } - /** Helper method to build AmsServerInfo (copied from HighAvailabilityContainer). */ - private AmsServerInfo buildServerInfo(String host, Integer thriftPort, Integer httpPort) { - AmsServerInfo serverInfo = new AmsServerInfo(); - serverInfo.setHost(host); - serverInfo.setThriftBindPort(thriftPort); - serverInfo.setRestBindPort(httpPort); - return serverInfo; - } - - /** Create a mock CuratorFramework that uses MockZkState for storage. */ @SuppressWarnings("unchecked") private CuratorFramework createMockZkClient() throws Exception { CuratorFramework mockClient = mock(CuratorFramework.class); - - // Mock getChildren() - create a chain of mocks - org.apache.amoro.shade.zookeeper3.org.apache.curator.framework.api.GetChildrenBuilder - getChildrenBuilder = - mock( - org.apache.amoro.shade.zookeeper3.org.apache.curator.framework.api - .GetChildrenBuilder.class); - when(mockClient.getChildren()).thenReturn(getChildrenBuilder); - when(getChildrenBuilder.forPath(anyString())) - .thenAnswer( - invocation -> { - String path = invocation.getArgument(0); - return mockZkState.getChildren(path); - }); - - // Mock getData() - org.apache.amoro.shade.zookeeper3.org.apache.curator.framework.api.GetDataBuilder - getDataBuilder = - mock( - org.apache.amoro.shade.zookeeper3.org.apache.curator.framework.api.GetDataBuilder - .class); - when(mockClient.getData()).thenReturn(getDataBuilder); - when(getDataBuilder.forPath(anyString())) - .thenAnswer( - invocation -> { - String path = invocation.getArgument(0); - return mockZkState.getData(path); - }); - - // Mock create() - manually create the entire fluent API chain to ensure consistency - org.apache.amoro.shade.zookeeper3.org.apache.curator.framework.api.CreateBuilder createBuilder = - mock( - org.apache.amoro.shade.zookeeper3.org.apache.curator.framework.api.CreateBuilder.class); - - @SuppressWarnings("unchecked") - org.apache.amoro.shade.zookeeper3.org.apache.curator.framework.api - .ProtectACLCreateModeStatPathAndBytesable< - String> - pathAndBytesable = - mock( - org.apache.amoro.shade.zookeeper3.org.apache.curator.framework.api - .ProtectACLCreateModeStatPathAndBytesable.class); - - when(mockClient.create()).thenReturn(createBuilder); - - // Mock the chain: creatingParentsIfNeeded() -> withMode() -> forPath() - // Use the same mock object for the entire chain - when(createBuilder.creatingParentsIfNeeded()).thenReturn(pathAndBytesable); - when(pathAndBytesable.withMode(any(CreateMode.class))).thenReturn(pathAndBytesable); - - // Mock forPath(path, data) - used by registAndElect() - when(pathAndBytesable.forPath(anyString(), any(byte[].class))) - .thenAnswer( - invocation -> { - String path = invocation.getArgument(0); - byte[] data = invocation.getArgument(1); - return mockZkState.createNode(path, data); - }); - - // Mock forPath(path) - used by createPathIfNeeded() - // Note: createPathIfNeeded() creates paths without data, but we still need to store them - // so that getChildren() can work correctly - when(pathAndBytesable.forPath(anyString())) - .thenAnswer( - invocation -> { - String path = invocation.getArgument(0); - // Create the path as an empty node (this simulates ZK path creation) - // In real ZK, paths are logical containers, but we need to store them - // to make getChildren() work correctly - if (mockZkState.exists(path) == null) { - mockZkState.createNode(path, new byte[0]); - } - return null; - }); - - // Mock delete() - org.apache.amoro.shade.zookeeper3.org.apache.curator.framework.api.DeleteBuilder deleteBuilder = - mock( - org.apache.amoro.shade.zookeeper3.org.apache.curator.framework.api.DeleteBuilder.class); - when(mockClient.delete()).thenReturn(deleteBuilder); - doAnswer( - invocation -> { - String path = invocation.getArgument(0); - mockZkState.deleteNode(path); - return null; - }) - .when(deleteBuilder) - .forPath(anyString()); - - // Mock checkExists() - @SuppressWarnings("unchecked") - org.apache.amoro.shade.zookeeper3.org.apache.curator.framework.api.ExistsBuilder - checkExistsBuilder = - mock( - org.apache.amoro.shade.zookeeper3.org.apache.curator.framework.api.ExistsBuilder - .class); - when(mockClient.checkExists()).thenReturn(checkExistsBuilder); - when(checkExistsBuilder.forPath(anyString())) - .thenAnswer( - invocation -> { - String path = invocation.getArgument(0); - return mockZkState.exists(path); - }); - - // Mock start() and close() doAnswer(invocation -> null).when(mockClient).start(); doAnswer(invocation -> null).when(mockClient).close(); - return mockClient; } - /** Create a mock LeaderLatch. */ private LeaderLatch createMockLeaderLatch() throws Exception { return createMockLeaderLatch(true); } - /** Create a mock LeaderLatch with specified leadership status. */ private LeaderLatch createMockLeaderLatch(boolean hasLeadership) throws Exception { LeaderLatch mockLatch = mock(LeaderLatch.class); when(mockLatch.hasLeadership()).thenReturn(hasLeadership); doAnswer(invocation -> null).when(mockLatch).addListener(any()); doAnswer(invocation -> null).when(mockLatch).start(); doAnswer(invocation -> null).when(mockLatch).close(); - // Mock await() - it throws IOException and InterruptedException - doAnswer( - invocation -> { - // Mock implementation - doesn't actually wait - return null; - }) - .when(mockLatch) - .await(); + doAnswer(invocation -> null).when(mockLatch).await(); return mockLatch; } - - /** In-memory ZK state simulator. */ - private static class MockZkState { - private final Map nodes = new HashMap<>(); - private final AtomicInteger sequenceCounter = new AtomicInteger(0); - - public List getChildren(String path) throws KeeperException { - List children = new ArrayList<>(); - String prefix = path.endsWith("/") ? path : path + "/"; - for (String nodePath : nodes.keySet()) { - // Only include direct children (not the path itself, and not nested paths) - if (nodePath.startsWith(prefix) && !nodePath.equals(path)) { - String relativePath = nodePath.substring(prefix.length()); - // Only add direct children (no additional slashes) - // This means the path should be exactly: prefix + relativePath - if (!relativePath.contains("/")) { - children.add(relativePath); - } - } - } - // Sort to ensure consistent ordering - children.sort(String::compareTo); - return children; - } - - public byte[] getData(String path) throws KeeperException { - byte[] data = nodes.get(path); - if (data == null) { - throw new KeeperException.NoNodeException(path); - } - return data; - } - - public String createNode(String path, byte[] data) { - // Handle sequential nodes - if (path.endsWith("-")) { - int seq = sequenceCounter.incrementAndGet(); - path = path + String.format("%010d", seq); - } - nodes.put(path, data); - return path; - } - - public void deleteNode(String path) throws KeeperException { - if (!nodes.containsKey(path)) { - throw new KeeperException.NoNodeException(path); - } - nodes.remove(path); - } - - public Stat exists(String path) { - return nodes.containsKey(path) ? new Stat() : null; - } - - public void clear() { - nodes.clear(); - sequenceCounter.set(0); - } - } } From f3e0f99e398e98811a14ee400d00ec7b0cabdff5 Mon Sep 17 00:00:00 2001 From: zhoujinsong Date: Tue, 8 Sep 2026 16:58:00 +0800 Subject: [PATCH 2/5] Polish some logging output for AmoroServiceContainer --- .../amoro/server/AmoroServiceContainer.java | 28 ++++++++++--------- 1 file changed, 15 insertions(+), 13 deletions(-) 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 index 657f9c90be..4f7860130f 100644 --- a/amoro-ams/src/main/java/org/apache/amoro/server/AmoroServiceContainer.java +++ b/amoro-ams/src/main/java/org/apache/amoro/server/AmoroServiceContainer.java @@ -28,6 +28,7 @@ import org.apache.amoro.OptimizerProperties; import org.apache.amoro.api.AmoroTableMetastore; import org.apache.amoro.api.OptimizingService; +import org.apache.amoro.client.AmsServerInfo; import org.apache.amoro.config.ConfigHelpers; import org.apache.amoro.config.ConfigurationException; import org.apache.amoro.config.Configurations; @@ -224,8 +225,9 @@ public void startBaseServices() throws Exception { startOptimizingService(); // Register this node so AmsAssignService (leader) can discover it and assign buckets. if (haContainer != null) { - bucketAssignStore.registerNode(haContainer.getOptimizingServiceServerInfo()); - LOG.info("Registered this node to BucketAssignStore"); + AmsServerInfo amsServerInfo = haContainer.getOptimizingServiceServerInfo(); + bucketAssignStore.registerNode(amsServerInfo); + LOG.info("Registered node {} to bucket assignment store", amsServerInfo); } } } @@ -296,12 +298,12 @@ public void startLeaderServices() throws Exception { try { amsAssignService = new AmsAssignService(serviceConfig, bucketAssignStore); } catch (Exception e) { - LOG.error("Failed to recreate AmsAssignService", e); + LOG.error("Failed to recreate Ams assign service", e); } } if (amsAssignService != null) { amsAssignService.start(); - LOG.info("AmsAssignService started"); + LOG.info("Ams assign service started"); } } else { startOptimizingService(); @@ -321,7 +323,7 @@ public void stopLeaderServices() { } if (IS_MASTER_SLAVE_MODE) { if (amsAssignService != null) { - LOG.info("Stopping AmsAssignService..."); + LOG.info("Stopping Ams assign service..."); amsAssignService.stop(); amsAssignService = null; } @@ -337,14 +339,14 @@ public void stopBaseServices() { if (bucketAssignStore != null && haContainer != null) { try { bucketAssignStore.removeNode(haContainer.getOptimizingServiceServerInfo()); - LOG.info("Unregistered this node from BucketAssignStore"); + LOG.info("Unregistered this node from bucket assignment store"); } catch (Exception e) { - LOG.warn("Failed to unregister node from BucketAssignStore", e); + LOG.warn("Failed to unregister node from bucket assignment store", e); } try { bucketAssignStore.close(); } catch (Exception e) { - LOG.warn("Failed to close BucketAssignStore", e); + LOG.warn("Failed to close bucket assignment store", e); } bucketAssignStore = null; } @@ -412,7 +414,7 @@ public void dispose() { } private void initConfig() throws Exception { - LOG.info("initializing configurations..."); + LOG.info("Initializing configurations..."); new ConfigurationHelper().init(); IS_MASTER_SLAVE_MODE = serviceConfig.getBoolean(HA_USE_MASTER_SLAVE_MODE); } @@ -625,9 +627,9 @@ public void init() throws Exception { } private void initServiceConfig(Map envConfig) throws Exception { - LOG.info("initializing service configuration..."); + LOG.info("Initializing service configuration..."); String configPath = Environments.getConfigPath() + "/" + SERVER_CONFIG_FILENAME; - LOG.info("load config from path: {}", configPath); + LOG.info("Loaded config from path: {}", configPath); yamlConfig = JacksonUtil.fromObjects( new Yaml().loadAs(Files.newInputStream(Paths.get(configPath)), Map.class)); @@ -650,7 +652,7 @@ private void initServiceConfig(Map envConfig) throws Exception { } private Map initEnvConfig() { - LOG.info("initializing system env configuration..."); + LOG.info("Initializing system env configuration..."); Map envs = System.getenv(); envs.forEach((k, v) -> LOG.info("export {}={}", k, v)); String prefix = AmoroManagementConf.SYSTEM_CONFIG.toUpperCase(); @@ -682,7 +684,7 @@ private void initIcebergThreadPools() { } private void initContainerConfig() { - LOG.info("initializing container configuration..."); + LOG.info("Initializing container configuration..."); JsonNode containers = yamlConfig.get(AmoroManagementConf.CONTAINER_LIST); List containerList = new ArrayList<>(); if (containers != null && containers.isArray()) { From 59e7612e0c8db3a7d926bb63a4d867d39522e3dc Mon Sep 17 00:00:00 2001 From: zhoujinsong Date: Tue, 8 Sep 2026 20:19:40 +0800 Subject: [PATCH 3/5] Fix issues in unit test classed --- .../amoro/server/TestAmsAssignService.java | 55 +------------------ .../server/TestHighAvailabilityContainer.java | 13 +---- .../ha/TestZkHighAvailabilityContainer.java | 13 +---- 3 files changed, 3 insertions(+), 78 deletions(-) diff --git a/amoro-ams/src/test/java/org/apache/amoro/server/TestAmsAssignService.java b/amoro-ams/src/test/java/org/apache/amoro/server/TestAmsAssignService.java index 0e69e091dd..baab94b58e 100644 --- a/amoro-ams/src/test/java/org/apache/amoro/server/TestAmsAssignService.java +++ b/amoro-ams/src/test/java/org/apache/amoro/server/TestAmsAssignService.java @@ -382,64 +382,11 @@ private HighAvailabilityContainer createContainerWithMockZk(Configurations confi /** Create a HighAvailabilityContainer without initializing ZK connection. */ private HighAvailabilityContainer createContainerWithoutZk(Configurations config) throws Exception { - java.lang.reflect.Constructor constructor = - ZkHighAvailabilityContainer.class.getDeclaredConstructor(Configurations.class); - // Create a minimal config that disables HA to avoid ZK connection Configurations tempConfig = new Configurations(config); tempConfig.setBoolean(AmoroManagementConf.HA_ENABLE, false); - HighAvailabilityContainer container = constructor.newInstance(tempConfig); - - // Now set all required fields using reflection - java.lang.reflect.Field isMasterSlaveModeField = - ZkHighAvailabilityContainer.class.getDeclaredField("isMasterSlaveMode"); - isMasterSlaveModeField.setAccessible(true); - isMasterSlaveModeField.set( - container, config.getBoolean(AmoroManagementConf.HA_USE_MASTER_SLAVE_MODE)); - - if (config.getBoolean(AmoroManagementConf.HA_ENABLE)) { - String haClusterName = config.getString(AmoroManagementConf.HA_CLUSTER_NAME); - - java.lang.reflect.Field tableServiceMasterPathField = - ZkHighAvailabilityContainer.class.getDeclaredField("tableServiceMasterPath"); - tableServiceMasterPathField.setAccessible(true); - tableServiceMasterPathField.set( - container, AmsHAProperties.getTableServiceMasterPath(haClusterName)); - - java.lang.reflect.Field optimizingServiceMasterPathField = - ZkHighAvailabilityContainer.class.getDeclaredField("optimizingServiceMasterPath"); - optimizingServiceMasterPathField.setAccessible(true); - optimizingServiceMasterPathField.set( - container, AmsHAProperties.getOptimizingServiceMasterPath(haClusterName)); - - java.lang.reflect.Field nodesPathField = - ZkHighAvailabilityContainer.class.getDeclaredField("nodesPath"); - nodesPathField.setAccessible(true); - nodesPathField.set(container, AmsHAProperties.getNodesPath(haClusterName)); - - java.lang.reflect.Field tableServiceServerInfoField = - ZkHighAvailabilityContainer.class.getDeclaredField("tableServiceServerInfo"); - tableServiceServerInfoField.setAccessible(true); - AmsServerInfo tableServiceServerInfo = - buildServerInfo( - config.getString(AmoroManagementConf.SERVER_EXPOSE_HOST), - config.getInteger(AmoroManagementConf.TABLE_SERVICE_THRIFT_BIND_PORT), - config.getInteger(AmoroManagementConf.HTTP_SERVER_PORT)); - tableServiceServerInfoField.set(container, tableServiceServerInfo); - - java.lang.reflect.Field optimizingServiceServerInfoField = - ZkHighAvailabilityContainer.class.getDeclaredField("optimizingServiceServerInfo"); - optimizingServiceServerInfoField.setAccessible(true); - AmsServerInfo optimizingServiceServerInfo = - buildServerInfo( - config.getString(AmoroManagementConf.SERVER_EXPOSE_HOST), - config.getInteger(AmoroManagementConf.OPTIMIZING_SERVICE_THRIFT_BIND_PORT), - config.getInteger(AmoroManagementConf.HTTP_SERVER_PORT)); - optimizingServiceServerInfoField.set(container, optimizingServiceServerInfo); - } - - return container; + return new ZkHighAvailabilityContainer(tempConfig); } /** Helper method to build AmsServerInfo. */ diff --git a/amoro-ams/src/test/java/org/apache/amoro/server/TestHighAvailabilityContainer.java b/amoro-ams/src/test/java/org/apache/amoro/server/TestHighAvailabilityContainer.java index 53064a8b1e..a01335f4ac 100644 --- a/amoro-ams/src/test/java/org/apache/amoro/server/TestHighAvailabilityContainer.java +++ b/amoro-ams/src/test/java/org/apache/amoro/server/TestHighAvailabilityContainer.java @@ -102,21 +102,10 @@ private HighAvailabilityContainer createContainerWithMockZk() throws Exception { } private HighAvailabilityContainer createContainerWithoutZk() throws Exception { - java.lang.reflect.Constructor constructor = - ZkHighAvailabilityContainer.class.getDeclaredConstructor(Configurations.class); - Configurations tempConfig = new Configurations(serviceConfig); tempConfig.setBoolean(AmoroManagementConf.HA_ENABLE, false); - HighAvailabilityContainer container = constructor.newInstance(tempConfig); - - java.lang.reflect.Field isMasterSlaveModeField = - ZkHighAvailabilityContainer.class.getDeclaredField("isMasterSlaveMode"); - isMasterSlaveModeField.setAccessible(true); - isMasterSlaveModeField.set( - container, serviceConfig.getBoolean(AmoroManagementConf.HA_USE_MASTER_SLAVE_MODE)); - - return container; + return new ZkHighAvailabilityContainer(tempConfig); } @SuppressWarnings("unchecked") 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 index fbe1b2f63d..6ba1eb1711 100644 --- 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 @@ -100,21 +100,10 @@ private HighAvailabilityContainer createContainerWithMockZk() throws Exception { } private HighAvailabilityContainer createContainerWithoutZk() throws Exception { - java.lang.reflect.Constructor constructor = - ZkHighAvailabilityContainer.class.getDeclaredConstructor(Configurations.class); - Configurations tempConfig = new Configurations(serviceConfig); tempConfig.setBoolean(AmoroManagementConf.HA_ENABLE, false); - HighAvailabilityContainer container = constructor.newInstance(tempConfig); - - java.lang.reflect.Field isMasterSlaveModeField = - ZkHighAvailabilityContainer.class.getDeclaredField("isMasterSlaveMode"); - isMasterSlaveModeField.setAccessible(true); - isMasterSlaveModeField.set( - container, serviceConfig.getBoolean(AmoroManagementConf.HA_USE_MASTER_SLAVE_MODE)); - - return container; + return new ZkHighAvailabilityContainer(tempConfig); } @SuppressWarnings("unchecked") From 4655682fb8c3eefb6f183ae09c1dafeea59a151d Mon Sep 17 00:00:00 2001 From: zhoujinsong Date: Tue, 8 Sep 2026 21:06:19 +0800 Subject: [PATCH 4/5] Fix some unit test errors --- .../mapper/BucketAssignMapper.java | 6 +- .../amoro/server/TestDBBucketAssignStore.java | 184 ++++++++++++++++++ .../server/TestHighAvailabilityContainer.java | 63 +++++- .../amoro/server/TestZkBucketAssignStore.java | 51 +++++ 4 files changed, 293 insertions(+), 11 deletions(-) create mode 100644 amoro-ams/src/test/java/org/apache/amoro/server/TestDBBucketAssignStore.java diff --git a/amoro-ams/src/main/java/org/apache/amoro/server/persistence/mapper/BucketAssignMapper.java b/amoro-ams/src/main/java/org/apache/amoro/server/persistence/mapper/BucketAssignMapper.java index e53a3a4edc..1c78efd458 100644 --- a/amoro-ams/src/main/java/org/apache/amoro/server/persistence/mapper/BucketAssignMapper.java +++ b/amoro-ams/src/main/java/org/apache/amoro/server/persistence/mapper/BucketAssignMapper.java @@ -35,12 +35,12 @@ public interface BucketAssignMapper { @Insert( "INSERT INTO bucket_assignments (cluster_name, node_key, server_info_json, assignments_json, last_update_time, node_heartbeat_ts) " - + "VALUES (#{meta.clusterName}, #{meta.nodeKey}, #{meta.serverInfoJson}, #{meta.assignmentsJson}, #{meta.lastUpdateTime}, #{meta.nodeHeartbeatTs})") + + "VALUES (#{meta.clusterName, jdbcType=VARCHAR}, #{meta.nodeKey, jdbcType=VARCHAR}, #{meta.serverInfoJson, jdbcType=VARCHAR}, #{meta.assignmentsJson, jdbcType=VARCHAR}, #{meta.lastUpdateTime}, #{meta.nodeHeartbeatTs})") int insert(@Param("meta") BucketAssignmentMeta meta); @Update( - "UPDATE bucket_assignments SET server_info_json = #{serverInfoJson}, assignments_json = #{assignmentsJson}, last_update_time = #{lastUpdateTime} " - + "WHERE cluster_name = #{clusterName} AND node_key = #{nodeKey}") + "UPDATE bucket_assignments SET server_info_json = #{serverInfoJson, jdbcType=VARCHAR}, assignments_json = #{assignmentsJson, jdbcType=VARCHAR}, last_update_time = #{lastUpdateTime} " + + "WHERE cluster_name = #{clusterName, jdbcType=VARCHAR} AND node_key = #{nodeKey, jdbcType=VARCHAR}") int update( @Param("clusterName") String clusterName, @Param("nodeKey") String nodeKey, diff --git a/amoro-ams/src/test/java/org/apache/amoro/server/TestDBBucketAssignStore.java b/amoro-ams/src/test/java/org/apache/amoro/server/TestDBBucketAssignStore.java new file mode 100644 index 0000000000..c9a0cab2c8 --- /dev/null +++ b/amoro-ams/src/test/java/org/apache/amoro/server/TestDBBucketAssignStore.java @@ -0,0 +1,184 @@ +/* + * 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.amoro.server; + +import org.apache.amoro.client.AmsServerInfo; +import org.apache.amoro.server.table.DerbyPersistence; +import org.junit.After; +import org.junit.Assert; +import org.junit.Before; +import org.junit.ClassRule; +import org.junit.Test; + +import java.util.Arrays; +import java.util.List; +import java.util.Map; +import java.util.concurrent.TimeUnit; + +/** Test for {@link DBBucketAssignStore} using an embedded Derby database. */ +public class TestDBBucketAssignStore { + + private static final String CLUSTER_NAME = "test-cluster"; + private static final long HEARTBEAT_TTL_MS = TimeUnit.SECONDS.toMillis(60); + + @ClassRule public static DerbyPersistence DERBY = new DerbyPersistence(); + + private DBBucketAssignStore assignStore; + private AmsServerInfo node1; + private AmsServerInfo node2; + + @Before + public void setUp() throws Exception { + assignStore = new DBBucketAssignStore(CLUSTER_NAME, HEARTBEAT_TTL_MS); + + node1 = new AmsServerInfo(); + node1.setHost("127.0.0.1"); + node1.setThriftBindPort(1260); + node1.setRestBindPort(1630); + + node2 = new AmsServerInfo(); + node2.setHost("127.0.0.2"); + node2.setThriftBindPort(1261); + node2.setRestBindPort(1631); + } + + @After + public void tearDown() throws Exception { + if (assignStore != null) { + try { + assignStore.removeNode(node1); + assignStore.removeNode(node2); + } catch (Exception e) { + // ignore + } + assignStore.close(); + } + } + + @Test + public void testRegisterNode() throws Exception { + assignStore.registerNode(node1); + + List aliveNodes = assignStore.getAliveNodes(); + Assert.assertEquals("Should have 1 alive node", 1, aliveNodes.size()); + Assert.assertEquals(node1.getHost(), aliveNodes.get(0).getHost()); + Assert.assertEquals(node1.getThriftBindPort(), aliveNodes.get(0).getThriftBindPort()); + } + + @Test + public void testRegisterMultipleNodes() throws Exception { + assignStore.registerNode(node1); + assignStore.registerNode(node2); + + List aliveNodes = assignStore.getAliveNodes(); + Assert.assertEquals("Should have 2 alive nodes", 2, aliveNodes.size()); + } + + @Test + public void testRegisterNodeIdempotent() throws Exception { + assignStore.registerNode(node1); + assignStore.registerNode(node1); // should update heartbeat, not duplicate + + List aliveNodes = assignStore.getAliveNodes(); + Assert.assertEquals("Should still have 1 alive node", 1, aliveNodes.size()); + } + + @Test + public void testGetAliveNodesEmpty() throws Exception { + List aliveNodes = assignStore.getAliveNodes(); + Assert.assertNotNull("Should return empty list", aliveNodes); + Assert.assertTrue("Should be empty", aliveNodes.isEmpty()); + } + + @Test + public void testRemoveNode() throws Exception { + assignStore.registerNode(node1); + Assert.assertEquals(1, assignStore.getAliveNodes().size()); + + assignStore.removeNode(node1); + Assert.assertEquals(0, assignStore.getAliveNodes().size()); + } + + @Test + public void testRemoveNodeNotRegistered() throws Exception { + // Should not throw + assignStore.removeNode(node1); + Assert.assertEquals(0, assignStore.getAliveNodes().size()); + } + + @Test + public void testSaveAndGetAssignments() throws Exception { + List bucketIds = Arrays.asList("1", "2", "3"); + assignStore.saveAssignments(node1, bucketIds); + + List retrieved = assignStore.getAssignments(node1); + Assert.assertEquals(bucketIds, retrieved); + } + + @Test + public void testUpdateAssignments() throws Exception { + List initial = Arrays.asList("1", "2"); + List updated = Arrays.asList("3", "4", "5"); + + assignStore.saveAssignments(node1, initial); + Assert.assertEquals(initial, assignStore.getAssignments(node1)); + + assignStore.saveAssignments(node1, updated); + Assert.assertEquals(updated, assignStore.getAssignments(node1)); + } + + @Test + public void testRemoveAssignments() throws Exception { + assignStore.saveAssignments(node1, Arrays.asList("1", "2")); + Assert.assertFalse(assignStore.getAssignments(node1).isEmpty()); + + assignStore.removeAssignments(node1); + Assert.assertTrue(assignStore.getAssignments(node1).isEmpty()); + } + + @Test + public void testGetAllAssignments() throws Exception { + assignStore.saveAssignments(node1, Arrays.asList("1", "2")); + assignStore.saveAssignments(node2, Arrays.asList("3", "4")); + + Map> all = assignStore.getAllAssignments(); + Assert.assertEquals(2, all.size()); + } + + @Test + public void testGetAllAssignmentsEmpty() throws Exception { + Map> all = assignStore.getAllAssignments(); + Assert.assertTrue(all.isEmpty()); + } + + @Test + public void testLastUpdateTime() throws Exception { + long initial = assignStore.getLastUpdateTime(node1); + Assert.assertEquals(0, initial); + + assignStore.saveAssignments(node1, Arrays.asList("1", "2")); + long afterSave = assignStore.getLastUpdateTime(node1); + Assert.assertTrue(afterSave > 0); + + Thread.sleep(10); + assignStore.updateLastUpdateTime(node1); + long afterUpdate = assignStore.getLastUpdateTime(node1); + Assert.assertTrue(afterUpdate > afterSave); + } +} diff --git a/amoro-ams/src/test/java/org/apache/amoro/server/TestHighAvailabilityContainer.java b/amoro-ams/src/test/java/org/apache/amoro/server/TestHighAvailabilityContainer.java index a01335f4ac..30e0cf5a6a 100644 --- a/amoro-ams/src/test/java/org/apache/amoro/server/TestHighAvailabilityContainer.java +++ b/amoro-ams/src/test/java/org/apache/amoro/server/TestHighAvailabilityContainer.java @@ -23,6 +23,7 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; +import org.apache.amoro.client.AmsServerInfo; import org.apache.amoro.config.Configurations; import org.apache.amoro.server.ha.HighAvailabilityContainer; import org.apache.amoro.server.ha.ZkHighAvailabilityContainer; @@ -84,28 +85,74 @@ public void testCreateWithoutHAEnabled() throws Exception { haContainer = new ZkHighAvailabilityContainer(serviceConfig); } + @Test + public void testGetTableServiceServerInfo() throws Exception { + serviceConfig.setBoolean(AmoroManagementConf.HA_USE_MASTER_SLAVE_MODE, true); + mockLeaderLatch = createMockLeaderLatch(true); + haContainer = createContainerWithMockZk(); + + AmsServerInfo tableServiceInfo = haContainer.getTableServiceServerInfo(); + Assert.assertEquals("127.0.0.1", tableServiceInfo.getHost()); + Assert.assertEquals(Integer.valueOf(1260), tableServiceInfo.getThriftBindPort()); + } + + @Test + public void testGetOptimizingServiceServerInfo() throws Exception { + serviceConfig.setBoolean(AmoroManagementConf.HA_USE_MASTER_SLAVE_MODE, true); + mockLeaderLatch = createMockLeaderLatch(true); + haContainer = createContainerWithMockZk(); + + AmsServerInfo optimizingServiceInfo = haContainer.getOptimizingServiceServerInfo(); + Assert.assertEquals("127.0.0.1", optimizingServiceInfo.getHost()); + Assert.assertEquals(Integer.valueOf(1261), optimizingServiceInfo.getThriftBindPort()); + } + /** Create HighAvailabilityContainer with mocked ZK components using reflection. */ private HighAvailabilityContainer createContainerWithMockZk() throws Exception { - HighAvailabilityContainer container = createContainerWithoutZk(); + // Build with HA disabled to avoid real ZK connection, then inject mocks via reflection + Configurations tempConfig = new Configurations(serviceConfig); + tempConfig.setBoolean(AmoroManagementConf.HA_ENABLE, false); + HighAvailabilityContainer container = new ZkHighAvailabilityContainer(tempConfig); + // Inject mock ZK client java.lang.reflect.Field zkClientField = ZkHighAvailabilityContainer.class.getDeclaredField("zkClient"); zkClientField.setAccessible(true); zkClientField.set(container, mockZkClient); + // Inject mock leader latch java.lang.reflect.Field leaderLatchField = ZkHighAvailabilityContainer.class.getDeclaredField("leaderLatch"); leaderLatchField.setAccessible(true); leaderLatchField.set(container, mockLeaderLatch); - return container; - } + // Inject server info (null when HA disabled, but tests need it) + AmsServerInfo tableServiceInfo = new AmsServerInfo(); + tableServiceInfo.setHost(serviceConfig.getString(AmoroManagementConf.SERVER_EXPOSE_HOST)); + tableServiceInfo.setThriftBindPort( + serviceConfig.getInteger(AmoroManagementConf.TABLE_SERVICE_THRIFT_BIND_PORT)); + tableServiceInfo.setRestBindPort( + serviceConfig.getInteger(AmoroManagementConf.HTTP_SERVER_PORT)); + + AmsServerInfo optimizingServiceInfo = new AmsServerInfo(); + optimizingServiceInfo.setHost( + serviceConfig.getString(AmoroManagementConf.SERVER_EXPOSE_HOST)); + optimizingServiceInfo.setThriftBindPort( + serviceConfig.getInteger(AmoroManagementConf.OPTIMIZING_SERVICE_THRIFT_BIND_PORT)); + optimizingServiceInfo.setRestBindPort( + serviceConfig.getInteger(AmoroManagementConf.HTTP_SERVER_PORT)); + + java.lang.reflect.Field tableServiceField = + ZkHighAvailabilityContainer.class.getDeclaredField("tableServiceServerInfo"); + tableServiceField.setAccessible(true); + tableServiceField.set(container, tableServiceInfo); + + java.lang.reflect.Field optimizingServiceField = + ZkHighAvailabilityContainer.class.getDeclaredField("optimizingServiceServerInfo"); + optimizingServiceField.setAccessible(true); + optimizingServiceField.set(container, optimizingServiceInfo); - private HighAvailabilityContainer createContainerWithoutZk() throws Exception { - Configurations tempConfig = new Configurations(serviceConfig); - tempConfig.setBoolean(AmoroManagementConf.HA_ENABLE, false); - - return new ZkHighAvailabilityContainer(tempConfig); + return container; } @SuppressWarnings("unchecked") diff --git a/amoro-ams/src/test/java/org/apache/amoro/server/TestZkBucketAssignStore.java b/amoro-ams/src/test/java/org/apache/amoro/server/TestZkBucketAssignStore.java index a331001edc..a5e0308ded 100644 --- a/amoro-ams/src/test/java/org/apache/amoro/server/TestZkBucketAssignStore.java +++ b/amoro-ams/src/test/java/org/apache/amoro/server/TestZkBucketAssignStore.java @@ -219,6 +219,57 @@ public void testMultipleNodesWithSameHostDifferentPort() throws Exception { Assert.assertEquals("Should have 2 nodes", 2, allAssignments.size()); } + @Test + public void testRegisterNode() throws Exception { + assignStore.registerNode(node1); + + List aliveNodes = assignStore.getAliveNodes(); + Assert.assertEquals("Should have 1 alive node", 1, aliveNodes.size()); + Assert.assertEquals(node1.getHost(), aliveNodes.get(0).getHost()); + Assert.assertEquals(node1.getThriftBindPort(), aliveNodes.get(0).getThriftBindPort()); + } + + @Test + public void testRegisterMultipleNodes() throws Exception { + assignStore.registerNode(node1); + assignStore.registerNode(node2); + + List aliveNodes = assignStore.getAliveNodes(); + Assert.assertEquals("Should have 2 alive nodes", 2, aliveNodes.size()); + } + + @Test + public void testRegisterNodeIdempotent() throws Exception { + assignStore.registerNode(node1); + assignStore.registerNode(node1); // duplicate should not throw + + List aliveNodes = assignStore.getAliveNodes(); + Assert.assertEquals("Should still have 1 alive node", 1, aliveNodes.size()); + } + + @Test + public void testGetAliveNodesEmpty() throws Exception { + List aliveNodes = assignStore.getAliveNodes(); + Assert.assertNotNull("Should return empty list", aliveNodes); + Assert.assertTrue("Should be empty", aliveNodes.isEmpty()); + } + + @Test + public void testRemoveNode() throws Exception { + assignStore.registerNode(node1); + Assert.assertEquals(1, assignStore.getAliveNodes().size()); + + assignStore.removeNode(node1); + Assert.assertEquals(0, assignStore.getAliveNodes().size()); + } + + @Test + public void testRemoveNodeNotRegistered() throws Exception { + // Should not throw + assignStore.removeNode(node1); + Assert.assertEquals(0, assignStore.getAliveNodes().size()); + } + /** Create a mock CuratorFramework that uses MockZkState for storage. */ @SuppressWarnings("unchecked") private CuratorFramework createMockZkClient() throws Exception { From 7017010b88ed704b53f8000896081608a9497994 Mon Sep 17 00:00:00 2001 From: zhoujinsong Date: Wed, 9 Sep 2026 10:19:03 +0800 Subject: [PATCH 5/5] Fix checkstyle issues --- .../java/org/apache/amoro/server/TestAmsAssignService.java | 1 - .../org/apache/amoro/server/TestHighAvailabilityContainer.java | 3 +-- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/amoro-ams/src/test/java/org/apache/amoro/server/TestAmsAssignService.java b/amoro-ams/src/test/java/org/apache/amoro/server/TestAmsAssignService.java index baab94b58e..b1f4b3636f 100644 --- a/amoro-ams/src/test/java/org/apache/amoro/server/TestAmsAssignService.java +++ b/amoro-ams/src/test/java/org/apache/amoro/server/TestAmsAssignService.java @@ -27,7 +27,6 @@ import org.apache.amoro.client.AmsServerInfo; import org.apache.amoro.config.Configurations; import org.apache.amoro.exception.BucketAssignStoreException; -import org.apache.amoro.properties.AmsHAProperties; import org.apache.amoro.server.ha.HighAvailabilityContainer; import org.apache.amoro.server.ha.ZkHighAvailabilityContainer; import org.apache.amoro.shade.zookeeper3.org.apache.curator.framework.CuratorFramework; diff --git a/amoro-ams/src/test/java/org/apache/amoro/server/TestHighAvailabilityContainer.java b/amoro-ams/src/test/java/org/apache/amoro/server/TestHighAvailabilityContainer.java index 30e0cf5a6a..8571754224 100644 --- a/amoro-ams/src/test/java/org/apache/amoro/server/TestHighAvailabilityContainer.java +++ b/amoro-ams/src/test/java/org/apache/amoro/server/TestHighAvailabilityContainer.java @@ -135,8 +135,7 @@ private HighAvailabilityContainer createContainerWithMockZk() throws Exception { serviceConfig.getInteger(AmoroManagementConf.HTTP_SERVER_PORT)); AmsServerInfo optimizingServiceInfo = new AmsServerInfo(); - optimizingServiceInfo.setHost( - serviceConfig.getString(AmoroManagementConf.SERVER_EXPOSE_HOST)); + optimizingServiceInfo.setHost(serviceConfig.getString(AmoroManagementConf.SERVER_EXPOSE_HOST)); optimizingServiceInfo.setThriftBindPort( serviceConfig.getInteger(AmoroManagementConf.OPTIMIZING_SERVICE_THRIFT_BIND_PORT)); optimizingServiceInfo.setRestBindPort(