Skip to content

[Refactor] Decouple HA container and bucket assignment store #4367

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
zhoujinsong wants to merge 5 commits into
base: master
Choose a base branch
from
Open

[Refactor] Decouple HA container and bucket assignment store #4367

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
6a3cb26
Decouple HA container and bucket assign stor
zhoujinsong Aug 19, 2026
f3e0f99
Polish some logging output for AmoroServiceContainer
zhoujinsong Sep 8, 2026
59e7612
Fix issues in unit test classed
zhoujinsong Sep 8, 2026
4655682
Fix some unit test errors
zhoujinsong Sep 8, 2026
7017010
Fix checkstyle issues
zhoujinsong Sep 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
59 changes: 40 additions & 19 deletions amoro-ams/src/main/java/org/apache/amoro/server/AmoroServiceContainer.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -127,11 +128,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) {
Expand Down Expand Up @@ -220,7 +221,14 @@ 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) {
AmsServerInfo amsServerInfo = haContainer.getOptimizingServiceServerInfo();
bucketAssignStore.registerNode(amsServerInfo);
LOG.info("Registered node {} to bucket assignment store", amsServerInfo);
}
}
}

Expand All @@ -241,11 +249,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<ActionCoordinator> actionCoordinators = defaultRuntimeFactory.supportedCoordinators();

tableService =
Expand Down Expand Up @@ -293,16 +296,14 @@ 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);
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();
Expand All @@ -322,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;
}
Expand All @@ -332,6 +333,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 bucket assignment store");
} catch (Exception e) {
LOG.warn("Failed to unregister node from bucket assignment store", e);
}
try {
bucketAssignStore.close();
} catch (Exception e) {
LOG.warn("Failed to close bucket assignment store", e);
}
bucketAssignStore = null;
}
disposeOptimizingService();
}
}

private void addHandlerChain(RuntimeHandlerChain chain) {
if (chain != null) {
tableService.addHandlerChain(chain);
Expand Down Expand Up @@ -388,12 +410,11 @@ public void disposeRestService() {

public void dispose() {
stopLeaderServices();
disposeOptimizingService();
disposeRestService();
stopBaseServices();
}

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);
}
Expand Down Expand Up @@ -606,9 +627,9 @@ public void init() throws Exception {
}

private void initServiceConfig(Map<String, Object> 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));
Expand All @@ -631,7 +652,7 @@ private void initServiceConfig(Map<String, Object> envConfig) throws Exception {
}

private Map<String, Object> initEnvConfig() {
LOG.info("initializing system env configuration...");
LOG.info("Initializing system env configuration...");
Map<String, String> envs = System.getenv();
envs.forEach((k, v) -> LOG.info("export {}={}", k, v));
String prefix = AmoroManagementConf.SYSTEM_CONFIG.toUpperCase();
Expand Down Expand Up @@ -663,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<ContainerMetadata> containerList = new ArrayList<>();
if (containers != null && containers.isArray()) {
Expand Down
67 changes: 39 additions & 28 deletions amoro-ams/src/main/java/org/apache/amoro/server/AmsAssignService.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand All @@ -66,41 +63,25 @@ 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 =
serviceConfig.get(AmoroManagementConf.HA_NODE_OFFLINE_TIMEOUT).toMillis();
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;
Expand Down Expand Up @@ -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<AmsServerInfo> aliveNodes = haContainer.getAliveNodes();
List<AmsServerInfo> aliveNodes = assignStore.getAliveNodes();
if (aliveNodes.isEmpty()) {
LOG.debug("No alive nodes found, skip bucket assignment");
return;
Expand All @@ -161,6 +137,8 @@ public void doAssign() {
Map<AmsServerInfo, List<String>> 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);
Expand All @@ -175,6 +153,15 @@ public void doAssign() {

List<String> 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<String> allBuckets = generateBucketIds();
Map<AmsServerInfo, List<String>> newAssignments =
buildNewAssignments(aliveNodes, change.offlineNodes, normalized.assignments);
Expand Down Expand Up @@ -676,6 +663,30 @@ private void persistAssignments(Map<AmsServerInfo, List<String>> 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<AmsServerInfo> aliveNodes, Map<AmsServerInfo, List<String>> currentAssignments) {
Set<String> 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.
Expand Down
19 changes: 19 additions & 0 deletions amoro-ams/src/main/java/org/apache/amoro/server/BucketAssignStore.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,22 @@ void saveAssignments(AmsServerInfo nodeInfo, List<String> bucketIds)
*/
List<AmsServerInfo> 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.
*
Expand All @@ -90,4 +106,7 @@ void saveAssignments(AmsServerInfo nodeInfo, List<String> 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();
}
39 changes: 22 additions & 17 deletions amoro-ams/src/main/java/org/apache/amoro/server/BucketAssignStoreFactory.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,48 +19,53 @@
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.
*
* <p>Supports different storage backends (ZK, database) according to HA type.
* <p>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);

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();
Expand Down
Loading
Loading