This is an automated email from the ASF dual-hosted git repository.
swuferhong pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/fluss.git
The following commit(s) were added to refs/heads/main by this push:
new 91bea7cd4 [server] Add per-task timeout for RebalanceManager to
prevent indefinite blocking (#3429)
91bea7cd4 is described below
commit 91bea7cd48c5f34ee87adfee897af40b67dc21b0
Author: yunhong <[email protected]>
AuthorDate: Sun Jun 7 12:36:52 2026 +0800
[server] Add per-task timeout for RebalanceManager to prevent indefinite
blocking (#3429)
* [server] Add per-task timeout for RebalanceManager to prevent indefinite
blocking
* address hongshun's comments
* address gujian's comments
---
.../fluss/cluster/rebalance/RebalanceStatus.java | 5 +-
.../coordinator/CoordinatorEventProcessor.java | 14 +-
.../event/RebalanceTaskTimeoutEvent.java | 39 +++++
.../coordinator/rebalance/RebalanceManager.java | 115 ++++++++++++-
.../rebalance/RebalanceManagerTest.java | 189 ++++++++++++++++++++-
website/docs/maintenance/operations/rebalance.md | 1 +
6 files changed, 358 insertions(+), 5 deletions(-)
diff --git
a/fluss-common/src/main/java/org/apache/fluss/cluster/rebalance/RebalanceStatus.java
b/fluss-common/src/main/java/org/apache/fluss/cluster/rebalance/RebalanceStatus.java
index a96b40ddd..5edbafd09 100644
---
a/fluss-common/src/main/java/org/apache/fluss/cluster/rebalance/RebalanceStatus.java
+++
b/fluss-common/src/main/java/org/apache/fluss/cluster/rebalance/RebalanceStatus.java
@@ -34,10 +34,11 @@ public enum RebalanceStatus {
REBALANCING(1),
FAILED(2),
COMPLETED(3),
- CANCELED(4);
+ CANCELED(4),
+ TIMEOUT(5);
public static final Set<RebalanceStatus> FINAL_STATUSES =
- new HashSet<>(Arrays.asList(COMPLETED, CANCELED, FAILED));
+ new HashSet<>(Arrays.asList(COMPLETED, CANCELED, FAILED, TIMEOUT));
private final int code;
diff --git
a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessor.java
b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessor.java
index 00ca807f0..456e7c587 100644
---
a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessor.java
+++
b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessor.java
@@ -85,6 +85,7 @@ import
org.apache.fluss.server.coordinator.event.NotifyKvSnapshotOffsetEvent;
import org.apache.fluss.server.coordinator.event.NotifyLakeTableOffsetEvent;
import
org.apache.fluss.server.coordinator.event.NotifyLeaderAndIsrResponseReceivedEvent;
import org.apache.fluss.server.coordinator.event.RebalanceEvent;
+import org.apache.fluss.server.coordinator.event.RebalanceTaskTimeoutEvent;
import org.apache.fluss.server.coordinator.event.RemoveServerTagEvent;
import org.apache.fluss.server.coordinator.event.SchemaChangeEvent;
import org.apache.fluss.server.coordinator.event.TableRegistrationChangeEvent;
@@ -122,6 +123,7 @@ import org.apache.fluss.server.zk.data.ZkData.TableIdsZNode;
import org.apache.fluss.server.zk.data.lake.LakeTableHelper;
import org.apache.fluss.server.zk.data.lake.LakeTableSnapshot;
import org.apache.fluss.utils.AutoPartitionStrategy;
+import org.apache.fluss.utils.clock.SystemClock;
import org.apache.fluss.utils.types.Tuple2;
import org.slf4j.Logger;
@@ -249,7 +251,9 @@ public class CoordinatorEventProcessor implements
EventProcessor {
this.lakeTableTieringManager = lakeTableTieringManager;
this.coordinatorMetricGroup = coordinatorMetricGroup;
this.internalListenerName =
conf.getString(ConfigOptions.INTERNAL_LISTENER_NAME);
- this.rebalanceManager = new RebalanceManager(this, zooKeeperClient);
+ this.rebalanceManager =
+ new RebalanceManager(
+ this, zooKeeperClient, coordinatorEventManager,
SystemClock.getInstance());
this.ioExecutor = ioExecutor;
this.lakeTableHelper =
new LakeTableHelper(zooKeeperClient,
conf.getString(ConfigOptions.REMOTE_DATA_DIR));
@@ -302,6 +306,7 @@ public class CoordinatorEventProcessor implements
EventProcessor {
// start rebalance manager.
rebalanceManager.startup();
+ rebalanceManager.start();
}
public void shutdown() {
@@ -645,6 +650,13 @@ public class CoordinatorEventProcessor implements
EventProcessor {
completeFromCallable(
cancelRebalanceEvent.getRespCallback(),
() -> processCancelRebalance(cancelRebalanceEvent));
+ } else if (event instanceof RebalanceTaskTimeoutEvent) {
+ RebalanceTaskTimeoutEvent timeoutEvent =
(RebalanceTaskTimeoutEvent) event;
+ LOG.warn(
+ "Rebalance task for {} timed out. Treating as timeout.",
+ timeoutEvent.getTableBucket());
+ rebalanceManager.finishRebalanceTask(
+ timeoutEvent.getTableBucket(), RebalanceStatus.TIMEOUT);
} else if (event instanceof ListRebalanceProgressEvent) {
ListRebalanceProgressEvent listRebalanceProgressEvent =
(ListRebalanceProgressEvent) event;
diff --git
a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/event/RebalanceTaskTimeoutEvent.java
b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/event/RebalanceTaskTimeoutEvent.java
new file mode 100644
index 000000000..c5d962d02
--- /dev/null
+++
b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/event/RebalanceTaskTimeoutEvent.java
@@ -0,0 +1,39 @@
+/*
+ * 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.fluss.server.coordinator.event;
+
+import org.apache.fluss.metadata.TableBucket;
+
+/** An event fired when a rebalance task exceeds the timeout without
completing. */
+public class RebalanceTaskTimeoutEvent implements CoordinatorEvent {
+
+ private final TableBucket tableBucket;
+
+ public RebalanceTaskTimeoutEvent(TableBucket tableBucket) {
+ this.tableBucket = tableBucket;
+ }
+
+ public TableBucket getTableBucket() {
+ return tableBucket;
+ }
+
+ @Override
+ public String toString() {
+ return "RebalanceTaskTimeoutEvent{tableBucket=" + tableBucket + "}";
+ }
+}
diff --git
a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/RebalanceManager.java
b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/RebalanceManager.java
index 1a722ff4d..9f2f350bd 100644
---
a/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/RebalanceManager.java
+++
b/fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/RebalanceManager.java
@@ -27,6 +27,8 @@ import
org.apache.fluss.exception.NoRebalanceInProgressException;
import org.apache.fluss.metadata.TableBucket;
import org.apache.fluss.server.coordinator.CoordinatorContext;
import org.apache.fluss.server.coordinator.CoordinatorEventProcessor;
+import org.apache.fluss.server.coordinator.event.EventManager;
+import org.apache.fluss.server.coordinator.event.RebalanceTaskTimeoutEvent;
import org.apache.fluss.server.coordinator.rebalance.goal.Goal;
import org.apache.fluss.server.coordinator.rebalance.goal.GoalOptimizer;
import org.apache.fluss.server.coordinator.rebalance.model.ClusterModel;
@@ -36,6 +38,9 @@ import org.apache.fluss.server.metadata.ServerInfo;
import org.apache.fluss.server.zk.ZooKeeperClient;
import org.apache.fluss.server.zk.data.LeaderAndIsr;
import org.apache.fluss.server.zk.data.RebalanceTask;
+import org.apache.fluss.utils.clock.Clock;
+import org.apache.fluss.utils.clock.SystemClock;
+import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -53,6 +58,9 @@ import java.util.SortedSet;
import java.util.TreeSet;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
import static org.apache.fluss.cluster.rebalance.RebalanceStatus.CANCELED;
import static org.apache.fluss.cluster.rebalance.RebalanceStatus.COMPLETED;
@@ -70,8 +78,17 @@ import static
org.apache.fluss.utils.Preconditions.checkNotNull;
public class RebalanceManager {
private static final Logger LOG =
LoggerFactory.getLogger(RebalanceManager.class);
+ /** Hardcoded timeout for an in-flight rebalance task: 2 minutes. */
+ private static final long REBALANCE_TASK_TIMEOUT_MS = 2 * 60 * 1000L;
+
+ /** Hardcoded interval for the periodic timeout check: 30 seconds. */
+ private static final long TIMEOUT_CHECK_INTERVAL_MS = 30 * 1000L;
+
private final ZooKeeperClient zkClient;
private final CoordinatorEventProcessor eventProcessor;
+ private final EventManager eventManager;
+ private final Clock clock;
+ private final ScheduledExecutorService timeoutChecker;
/** A queue of in progress table bucket to rebalance. */
private final Queue<TableBucket> inProgressRebalanceTasksQueue = new
ArrayDeque<>();
@@ -90,9 +107,46 @@ public class RebalanceManager {
private volatile @Nullable String currentRebalanceId;
private volatile boolean isClosed = false;
- public RebalanceManager(CoordinatorEventProcessor eventProcessor,
ZooKeeperClient zkClient) {
+ /**
+ * Timestamp when the current in-flight task was started, or -1 if idle.
+ *
+ * <p>Write ordering contract (volatile publication idiom): always write
{@code
+ * inflightTaskStartMs} BEFORE {@code inflightTaskBucket} when setting,
and clear {@code
+ * inflightTaskBucket} BEFORE {@code inflightTaskStartMs} when resetting.
The timeout checker
+ * reads in reverse order (bucket first, then startMs), ensuring it never
observes a stale
+ * startMs paired with a new bucket.
+ */
+ private volatile long inflightTaskStartMs = -1;
+
+ /** The bucket of the current in-flight task, or null if idle. Acts as the
"gate" variable. */
+ private volatile @Nullable TableBucket inflightTaskBucket;
+
+ public RebalanceManager(
+ CoordinatorEventProcessor eventProcessor,
+ ZooKeeperClient zkClient,
+ EventManager eventManager,
+ Clock clock) {
+ this(
+ eventProcessor,
+ zkClient,
+ eventManager,
+ clock,
+ Executors.newScheduledThreadPool(
+ 1, new ExecutorThreadFactory("rebalance-timeout")));
+ }
+
+ @VisibleForTesting
+ RebalanceManager(
+ CoordinatorEventProcessor eventProcessor,
+ ZooKeeperClient zkClient,
+ EventManager eventManager,
+ Clock clock,
+ ScheduledExecutorService timeoutChecker) {
this.eventProcessor = eventProcessor;
this.zkClient = zkClient;
+ this.eventManager = eventManager;
+ this.clock = clock == null ? SystemClock.getInstance() : clock;
+ this.timeoutChecker = timeoutChecker;
this.goalOptimizer = new GoalOptimizer();
}
@@ -101,6 +155,19 @@ public class RebalanceManager {
initialize();
}
+ /** Starts the periodic timeout checker. Call after {@link #startup()}. */
+ public void start() {
+ timeoutChecker.scheduleWithFixedDelay(
+ this::checkTimeoutSafely,
+ TIMEOUT_CHECK_INTERVAL_MS,
+ TIMEOUT_CHECK_INTERVAL_MS,
+ TimeUnit.MILLISECONDS);
+ LOG.info(
+ "RebalanceManager timeout checker started: timeoutMs={},
checkIntervalMs={}",
+ REBALANCE_TASK_TIMEOUT_MS,
+ TIMEOUT_CHECK_INTERVAL_MS);
+ }
+
public @Nullable String getRebalanceId() {
return currentRebalanceId;
}
@@ -132,6 +199,9 @@ public class RebalanceManager {
inProgressRebalanceTasks.clear();
inProgressRebalanceTasksQueue.clear();
finishedRebalanceTasks.clear();
+ // Clear gate (bucket) first, then data (startMs).
+ inflightTaskBucket = null;
+ inflightTaskStartMs = -1;
currentRebalanceId = rebalanceId;
if (rebalancePlan.isEmpty()) {
@@ -170,6 +240,9 @@ public class RebalanceManager {
finishedRebalanceTasks.put(
tableBucket,
RebalanceResultForBucket.of(resultForBucket.plan(),
statusForBucket));
+ // Clear gate (bucket) first, then data (startMs).
+ inflightTaskBucket = null;
+ inflightTaskStartMs = -1;
LOG.info(
"Rebalance task {} in progress: {} tasks pending, {}
completed.",
currentRebalanceId,
@@ -250,6 +323,9 @@ public class RebalanceManager {
rebalanceStatus = CANCELED;
inProgressRebalanceTasksQueue.clear();
inProgressRebalanceTasks.clear();
+ // Clear gate (bucket) first, then data (startMs).
+ inflightTaskBucket = null;
+ inflightTaskStartMs = -1;
// Here, it will not clear finishedRebalanceTasks, because it will be
used by
// listRebalanceProgress. It will be cleared when next register.
@@ -302,6 +378,9 @@ public class RebalanceManager {
private void processNewRebalanceTask() {
TableBucket tableBucket = inProgressRebalanceTasksQueue.peek();
if (tableBucket != null &&
inProgressRebalanceTasks.containsKey(tableBucket)) {
+ // Write data (startMs) first, then publish gate (bucket).
+ inflightTaskStartMs = clock.milliseconds();
+ inflightTaskBucket = tableBucket;
RebalanceResultForBucket resultForBucket =
inProgressRebalanceTasks.get(tableBucket);
RebalanceResultForBucket rebalanceResultForBucket =
RebalanceResultForBucket.of(resultForBucket.plan(),
REBALANCING);
@@ -400,12 +479,46 @@ public class RebalanceManager {
return new ClusterModel(servers);
}
+ private void checkTimeoutSafely() {
+ try {
+ checkTimeout();
+ } catch (Throwable t) {
+ LOG.error("Unexpected error in RebalanceManager timeout check.",
t);
+ }
+ }
+
+ @VisibleForTesting
+ void checkTimeout() {
+ // Read gate (bucket) first, then data (startMs).
+ // If bucket is non-null, happens-before guarantees startMs is at
least as
+ // fresh as the value written before bucket was published.
+ TableBucket bucket = inflightTaskBucket;
+ long startMs = inflightTaskStartMs;
+ if (bucket == null || startMs < 0) {
+ return;
+ }
+ long elapsed = clock.milliseconds() - startMs;
+ if (elapsed > REBALANCE_TASK_TIMEOUT_MS) {
+ LOG.warn(
+ "In-flight rebalance task for {} timed out after {}ms. "
+ + "Treating it as timed out and advancing to the
next task.",
+ bucket,
+ elapsed);
+ // Clear gate (bucket) first, then data (startMs), matching the
+ // publication idiom so the next checkTimeout sees bucket==null.
+ inflightTaskBucket = null;
+ inflightTaskStartMs = -1;
+ eventManager.put(new RebalanceTaskTimeoutEvent(bucket));
+ }
+ }
+
private void checkNotClosed() {
checkArgument(!isClosed, "RebalanceManager is already closed.");
}
public void close() {
isClosed = true;
+ timeoutChecker.shutdownNow();
}
@VisibleForTesting
diff --git
a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/rebalance/RebalanceManagerTest.java
b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/rebalance/RebalanceManagerTest.java
index 725430ede..d550aea1c 100644
---
a/fluss-server/src/test/java/org/apache/fluss/server/coordinator/rebalance/RebalanceManagerTest.java
+++
b/fluss-server/src/test/java/org/apache/fluss/server/coordinator/rebalance/RebalanceManagerTest.java
@@ -17,9 +17,12 @@
package org.apache.fluss.server.coordinator.rebalance;
+import org.apache.fluss.cluster.rebalance.RebalancePlanForBucket;
+import org.apache.fluss.cluster.rebalance.RebalanceResultForBucket;
import org.apache.fluss.cluster.rebalance.RebalanceStatus;
import org.apache.fluss.config.ConfigOptions;
import org.apache.fluss.config.Configuration;
+import org.apache.fluss.metadata.TableBucket;
import org.apache.fluss.server.coordinator.AutoPartitionManager;
import org.apache.fluss.server.coordinator.CoordinatorContext;
import org.apache.fluss.server.coordinator.CoordinatorEventProcessor;
@@ -27,6 +30,9 @@ import
org.apache.fluss.server.coordinator.LakeCatalogDynamicLoader;
import org.apache.fluss.server.coordinator.LakeTableTieringManager;
import org.apache.fluss.server.coordinator.MetadataManager;
import org.apache.fluss.server.coordinator.TestCoordinatorChannelManager;
+import org.apache.fluss.server.coordinator.event.CoordinatorEvent;
+import org.apache.fluss.server.coordinator.event.EventManager;
+import org.apache.fluss.server.coordinator.event.RebalanceTaskTimeoutEvent;
import org.apache.fluss.server.coordinator.lease.KvSnapshotLeaseManager;
import org.apache.fluss.server.coordinator.remote.RemoteDirDynamicLoader;
import org.apache.fluss.server.metadata.CoordinatorMetadataCache;
@@ -37,6 +43,7 @@ import org.apache.fluss.server.zk.ZooKeeperClient;
import org.apache.fluss.server.zk.ZooKeeperExtension;
import org.apache.fluss.server.zk.data.RebalanceTask;
import org.apache.fluss.testutils.common.AllCallbackWrapper;
+import org.apache.fluss.utils.clock.ManualClock;
import org.apache.fluss.utils.clock.SystemClock;
import org.apache.fluss.utils.concurrent.ExecutorThreadFactory;
@@ -47,11 +54,17 @@ import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.RegisterExtension;
import java.time.Duration;
+import java.util.ArrayList;
+import java.util.Arrays;
import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledThreadPoolExecutor;
import static org.apache.fluss.cluster.rebalance.RebalanceStatus.COMPLETED;
import static org.apache.fluss.cluster.rebalance.RebalanceStatus.NOT_STARTED;
+import static org.apache.fluss.cluster.rebalance.RebalanceStatus.TIMEOUT;
import static org.assertj.core.api.Assertions.assertThat;
/** Test for {@link RebalanceManager}. */
@@ -107,12 +120,19 @@ public class RebalanceManagerTest {
lakeTableTieringManager =
new
LakeTableTieringManager(TestingMetricGroups.LAKE_TIERING_METRICS);
CoordinatorEventProcessor eventProcessor =
buildCoordinatorEventProcessor(conf);
- rebalanceManager = new RebalanceManager(eventProcessor,
zookeeperClient);
+ RecordingEventManager recordingEventManager = new
RecordingEventManager();
+ rebalanceManager =
+ new RebalanceManager(
+ eventProcessor,
+ zookeeperClient,
+ recordingEventManager,
+ SystemClock.getInstance());
rebalanceManager.startup();
}
@AfterEach
void afterEach() throws Exception {
+ rebalanceManager.close();
zookeeperClient.deleteRebalanceTask();
metadataManager =
new MetadataManager(
@@ -142,6 +162,143 @@ public class RebalanceManagerTest {
.hasValue(new RebalanceTask(rebalanceId, COMPLETED, new
HashMap<>()));
}
+ @Test
+ void testTimeoutEnqueuesEvent() throws Exception {
+ ManualClock clock = new ManualClock(0L);
+ RecordingEventManager eventManager = new RecordingEventManager();
+ NoOpScheduledExecutor executor = new NoOpScheduledExecutor();
+ CoordinatorEventProcessor eventProcessor =
+ buildCoordinatorEventProcessor(new Configuration());
+
+ RebalanceManager manager =
+ new RebalanceManager(
+ eventProcessor, zookeeperClient, eventManager, clock,
executor);
+ manager.startup();
+
+ TableBucket tb1 = new TableBucket(1L, 0);
+ TableBucket tb2 = new TableBucket(1L, 1);
+ Map<TableBucket, RebalancePlanForBucket> plan = new HashMap<>();
+ plan.put(
+ tb1,
+ new RebalancePlanForBucket(
+ tb1, 0, 0, Arrays.asList(0, 1, 2), Arrays.asList(0, 1,
2)));
+ plan.put(
+ tb2,
+ new RebalancePlanForBucket(
+ tb2, 0, 0, Arrays.asList(0, 1, 2), Arrays.asList(0, 1,
2)));
+
+ zookeeperClient.registerRebalanceTask(new
RebalanceTask("timeout-test", NOT_STARTED, plan));
+ manager.registerRebalance("timeout-test", plan, NOT_STARTED);
+
+ // Not yet timed out.
+ clock.advanceTime(Duration.ofMillis(100_000));
+ manager.checkTimeout();
+ assertThat(eventManager.events).isEmpty();
+
+ // Cross the 2-minute boundary.
+ clock.advanceTime(Duration.ofMillis(30_000));
+ manager.checkTimeout();
+
+ assertThat(eventManager.events).hasSize(1);
+
assertThat(eventManager.events.get(0)).isInstanceOf(RebalanceTaskTimeoutEvent.class);
+ RebalanceTaskTimeoutEvent timeoutEvent =
+ (RebalanceTaskTimeoutEvent) eventManager.events.get(0);
+ assertThat(timeoutEvent.getTableBucket()).isEqualTo(tb1);
+
+ // A second checkTimeout() should NOT enqueue another event because the
+ // inflight state was cleared after the first timeout.
+ clock.advanceTime(Duration.ofMillis(30_000));
+ manager.checkTimeout();
+ assertThat(eventManager.events).hasSize(1);
+
+ manager.close();
+ }
+
+ @Test
+ void testTimeoutAfterCompletionIsNoOp() throws Exception {
+ ManualClock clock = new ManualClock(0L);
+ RecordingEventManager eventManager = new RecordingEventManager();
+ NoOpScheduledExecutor executor = new NoOpScheduledExecutor();
+ CoordinatorEventProcessor eventProcessor =
+ buildCoordinatorEventProcessor(new Configuration());
+
+ RebalanceManager manager =
+ new RebalanceManager(
+ eventProcessor, zookeeperClient, eventManager, clock,
executor);
+ manager.startup();
+
+ TableBucket tb1 = new TableBucket(1L, 0);
+ Map<TableBucket, RebalancePlanForBucket> plan = new HashMap<>();
+ plan.put(
+ tb1,
+ new RebalancePlanForBucket(
+ tb1, 0, 0, Arrays.asList(0, 1, 2), Arrays.asList(0, 1,
2)));
+
+ zookeeperClient.registerRebalanceTask(
+ new RebalanceTask("completion-test", NOT_STARTED, plan));
+ manager.registerRebalance("completion-test", plan, NOT_STARTED);
+
+ // The task completes normally before timeout.
+ manager.finishRebalanceTask(tb1, COMPLETED);
+
+ // Now the timeout fires, but the task is already done.
+ clock.advanceTime(Duration.ofMillis(130_000));
+ manager.checkTimeout();
+
+ // No timeout event should be enqueued because inflightTaskStartMs was
cleared.
+ assertThat(eventManager.events).isEmpty();
+
+ manager.close();
+ }
+
+ @Test
+ void testTimeoutTreatsTaskAsCompleted() throws Exception {
+ ManualClock clock = new ManualClock(0L);
+ RecordingEventManager eventManager = new RecordingEventManager();
+ NoOpScheduledExecutor executor = new NoOpScheduledExecutor();
+ CoordinatorEventProcessor eventProcessor =
+ buildCoordinatorEventProcessor(new Configuration());
+
+ RebalanceManager manager =
+ new RebalanceManager(
+ eventProcessor, zookeeperClient, eventManager, clock,
executor);
+ manager.startup();
+
+ TableBucket tb1 = new TableBucket(1L, 0);
+ TableBucket tb2 = new TableBucket(1L, 1);
+ Map<TableBucket, RebalancePlanForBucket> plan = new HashMap<>();
+ plan.put(
+ tb1,
+ new RebalancePlanForBucket(
+ tb1, 0, 0, Arrays.asList(0, 1, 2), Arrays.asList(0, 1,
2)));
+ plan.put(
+ tb2,
+ new RebalancePlanForBucket(
+ tb2, 0, 0, Arrays.asList(0, 1, 2), Arrays.asList(0, 1,
2)));
+
+ zookeeperClient.registerRebalanceTask(
+ new RebalanceTask("completed-test", NOT_STARTED, plan));
+ manager.registerRebalance("completed-test", plan, NOT_STARTED);
+
+ // Timeout fires.
+ clock.advanceTime(Duration.ofMillis(130_000));
+ manager.checkTimeout();
+
+ // Simulate the coordinator event thread processing the timeout event.
+ assertThat(eventManager.events).hasSize(1);
+ RebalanceTaskTimeoutEvent timeoutEvent =
+ (RebalanceTaskTimeoutEvent) eventManager.events.get(0);
+ manager.finishRebalanceTask(timeoutEvent.getTableBucket(), TIMEOUT);
+
+ // The timed-out task should be in finishedRebalanceTasks as TIMEOUT.
+ assertThat(manager.hasInProgressRebalance()).isTrue();
+ RebalanceResultForBucket result =
+
manager.listRebalanceProgress(null).progressForBucketMap().get(tb1);
+ assertThat(result.status()).isEqualTo(TIMEOUT);
+
+ manager.close();
+ }
+
private CoordinatorEventProcessor
buildCoordinatorEventProcessor(Configuration conf) {
return new CoordinatorEventProcessor(
zookeeperClient,
@@ -156,4 +313,34 @@ public class RebalanceManagerTest {
metadataManager,
kvSnapshotLeaseManager);
}
+
+ /** Records events put into the coordinator event queue. */
+ private static final class RecordingEventManager implements EventManager {
+ final List<CoordinatorEvent> events = new ArrayList<>();
+
+ @Override
+ public void put(CoordinatorEvent event) {
+ events.add(event);
+ }
+ }
+
+ /**
+ * A scheduled executor that never actually runs scheduled tasks, so tests
retain full control
+ * over when {@link RebalanceManager#checkTimeout()} is invoked.
+ */
+ private static final class NoOpScheduledExecutor extends
ScheduledThreadPoolExecutor {
+
+ NoOpScheduledExecutor() {
+ super(0);
+ }
+
+ @Override
+ public java.util.concurrent.ScheduledFuture<?> scheduleWithFixedDelay(
+ Runnable command,
+ long initialDelay,
+ long delay,
+ java.util.concurrent.TimeUnit unit) {
+ return null;
+ }
+ }
}
diff --git a/website/docs/maintenance/operations/rebalance.md
b/website/docs/maintenance/operations/rebalance.md
index 48d89ac40..418b7c585 100644
--- a/website/docs/maintenance/operations/rebalance.md
+++ b/website/docs/maintenance/operations/rebalance.md
@@ -110,6 +110,7 @@ Rebalance statuses:
- **COMPLETED**: The rebalance has successfully completed
- **FAILED**: The rebalance has failed
- **CANCELED**: The rebalance has been canceled
+- **TIMEOUT**: The rebalance task timed out (e.g., ISR could not converge
within the timeout period)
### 4. Cancel Rebalance (If Needed)
