Skip to content

Commit 287f458

Browse files
j1wonparkzhoujinsong
j1wonpark
and
zhoujinsong
authored
[AMORO-4295] AIP-5 Phase 3+4: idle-driven scale-down with graceful drain for dynamic allocation groups (#4296)
* [AMORO-4295] Expose per-token in-flight counts in the dynamic allocation load snapshot Scale-down needs to know which optimizer instances are busy, not only how many threads are. Aggregate SCHEDULED/ACKED task counts per optimizer token in the same scan collectDynamicAllocationLoad() already performs, so one snapshot supplies both scale directions. Recovered tasks keep their token across an AMS restart, which makes the very first snapshot after a restart accurate without any rebuild step. Signed-off-by: Jiwon Park <jpark92@outlook.kr> * [AMORO-4295] Add snapshot-derived idle observation and the scale-down decision Track per-token idle time by observation instead of event counters: each keeper round feeds the task snapshot into observe(), which refreshes busy timestamps, seeds first-seen tokens as idle (an instance that never receives a task must still become removable), and prunes unregistered ones. Event-maintained counters were rejected because three existing paths reclaim assigned tasks without any service-level decrement hook (queue-internal stale-ACKED resets, ack/exec-timeout reclaims of live optimizers, and process close), each silently inflating a counter until the instance can never look idle again. computeScaleDown picks at most one candidate per round: the longest-idle instance past executor-idle-timeout, rate-limited by scale-down-cooldown, whose removal keeps registered-minus-draining threads at or above min-parallelism. Draining threads count as already gone so consecutive removals cannot pass the floor check together and land below the floor once they all complete. validate() now also requires sustained-backlog-timeout to be at most half of executor-idle-timeout: the keeper cadence is the observation resolution, and sampling slower than that misjudges an instance that worked between samples as continuously idle. Signed-off-by: Jiwon Park <jpark92@outlook.kr> * [AMORO-4295] Add graceful drain plumbing: poll blocking and optimizer removal A draining optimizer's token enters a pending-removal set consulted twice in pollTask: once on entry, and once after the queue poll returns, because the entry check cannot stop a thread already parked in the queue's long-poll — it may fetch a task after the drain began, and that task is handed back (PLANNED again) instead of being assigned. touch/ack/complete stay open so in-flight tasks finish normally. executeRemoval releases the container resource, deletes the persisted row, and unregisters. A missing resource row is handled by releasing through the optimizer instance itself (it extends Resource and carries the container identity): a pod that self-registered after a persist failure has no row and never will, so treating that as a retryable error would drain-block it forever while its heartbeat keeps it registered. A failed container release keeps the drain state and retries on a later round — Kubernetes deletion is idempotent. The mock container used by the keeper tests now records released resources and can simulate release failures. Signed-off-by: Jiwon Park <jpark92@outlook.kr> * [AMORO-4295] Integrate scale-down and drain progress into the scale keeper round Each round now runs in a fixed order. Drain progress first and unconditionally — a drain that completed or hit its deadline converts to a removal even in rounds that scale up, or a busy drain would linger to its full timeout while backlog persists. Draining instances are then accounted as already gone on both sides of the demand math: leaving their threads in the capacity undercounts demand by up to their thread count, and leaving their tasks in the load keeps the future-demand signal (busy >= effective) from ever firing mid-drain. Idle observation runs every round including scale-up ones, so an instance busy through a burst does not come out of it looking idle since before the burst began. Scale-down runs only in rounds with no demand signal at all — including demand still held back by the backlog gate, which computeScaleUp now exposes: removing a warm instance right before the gate opens would free exactly the capacity the next round re-requests. A selected victim begins a graceful drain and is removed in the same round only if a snapshot taken after the token entered the pending-removal set shows it idle; the pre-insert snapshot may miss a task fetched by a long-poll racing the drain start. The keeper tests drive rounds with injected times because the validated minimum executor-idle-timeout (30s) puts real idle waits beyond sane test durations. Signed-off-by: Jiwon Park <jpark92@outlook.kr> * [AMORO-4295] Clear drain state on unregistration and on dynamic allocation disable Two lifecycle paths could strand a token in the pending-removal set. An optimizer that dies mid-drain is unregistered by heartbeat expiry, and its token can never be matched again (the replacement pod registers under a fresh one), so unregisterOptimizer now clears the drain state. And a group whose dynamic allocation is disabled mid-drain returns to the legacy floor keeper, where a leftover drain block would starve the still-running pod forever — unwatching the group re-admits its tokens to task assignment. Signed-off-by: Jiwon Park <jpark92@outlook.kr> * [AMORO-4295] Fix concurrent optimizer unregistration Signed-off-by: Jiwon Park <jpark92@outlook.kr> * [AMORO-4295] Document scale-down and drain behavior in the optimizer group property table Signed-off-by: Jiwon Park <jpark92@outlook.kr> --------- Signed-off-by: Jiwon Park <jpark92@outlook.kr> Co-authored-by: ZhouJinsong <zhoujinsong0505@163.com>
1 parent c26fd6e commit 287f458

12 files changed

Lines changed: 1064 additions & 16 deletions

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

Lines changed: 229 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,16 @@ public class DefaultOptimizingService extends StatedPersistentBase
116116
private final Map<String, OptimizingQueue> optimizingQueueByGroup = new ConcurrentHashMap<>();
117117
private final Map<String, OptimizingQueue> optimizingQueueByToken = new ConcurrentHashMap<>();
118118
private final Map<String, OptimizerInstance> authOptimizers = new ConcurrentHashMap<>();
119+
120+
/**
121+
* Tokens of draining optimizers (AIP-5 scale-down): {@link #pollTask} returns {@code null} for
122+
* them, blocking new assignments while in-flight tasks complete normally.
123+
*/
124+
private final Set<String> pendingRemovalTokens = ConcurrentHashMap.newKeySet();
125+
126+
/** Force-removal deadline per draining token, the {@code drain-timeout} safety net. */
127+
private final Map<String, Long> drainDeadlines = new ConcurrentHashMap<>();
128+
119129
private final OptimizerKeeper optimizerKeeper = new OptimizerKeeper("optimizer-keeper-thread");
120130
private final OptimizerGroupKeeper optimizerGroupKeeper =
121131
new OptimizerGroupKeeper("optimizer-group-keeper-thread");
@@ -242,9 +252,18 @@ private void unregisterOptimizer(String token) {
242252
doAs(OptimizerMapper.class, mapper -> mapper.deleteOptimizer(token));
243253
OptimizingQueue optimizingQueue = optimizingQueueByToken.remove(token);
244254
OptimizerInstance optimizer = authOptimizers.remove(token);
245-
if (optimizingQueue != null) {
246-
optimizingQueue.removeOptimizer(optimizer);
255+
if (optimizer != null) {
256+
if (optimizingQueue == null) {
257+
optimizingQueue = optimizingQueueByGroup.get(optimizer.getGroupName());
258+
}
259+
if (optimizingQueue != null) {
260+
optimizingQueue.removeOptimizer(optimizer);
261+
}
247262
}
263+
// An optimizer that dies mid-drain is unregistered here by heartbeat expiry; its token can
264+
// never be matched again, so leftover drain state would sit in the pending-removal set
265+
// forever (its replacement pod registers under a fresh token).
266+
cancelDrain(token);
248267
}
249268

250269
@Override
@@ -269,17 +288,118 @@ private OptimizerInstance getAuthenticatedOptimizer(String authToken) {
269288

270289
@Override
271290
public OptimizingTask pollTask(String authToken, int threadId) {
291+
if (pendingRemovalTokens.contains(authToken)) {
292+
return null;
293+
}
272294
LOG.debug("Optimizer {} (threadId {}) try polling task", authToken, threadId);
273295
OptimizerThread optimizerThread = getAuthenticatedOptimizer(authToken).getThread(threadId);
274296
OptimizingQueue queue = getQueueByToken(authToken);
275-
TaskRuntime<?> task = queue.pollTask(optimizerThread, pollingTimeout, breakQuotaLimit);
297+
TaskRuntime<?> task =
298+
guardDrainedPoll(
299+
authToken, queue.pollTask(optimizerThread, pollingTimeout, breakQuotaLimit));
276300
if (task != null) {
277301
LOG.info("OptimizerThread {} polled task {}", optimizerThread, task.getTaskId());
278302
return task.extractProtocolTask();
279303
}
280304
return null;
281305
}
282306

307+
/**
308+
* Close the long-poll race on drain start: the entry check above cannot stop a thread already
309+
* parked inside the queue's poll, which may fetch a task after its token entered the
310+
* pending-removal set. Hand such a task back instead of assigning it to a draining optimizer.
311+
*/
312+
@VisibleForTesting
313+
TaskRuntime<?> guardDrainedPoll(String authToken, TaskRuntime<?> task) {
314+
if (task == null || !pendingRemovalTokens.contains(authToken)) {
315+
return task;
316+
}
317+
OptimizingQueue queue = optimizingQueueByToken.get(authToken);
318+
if (queue != null) {
319+
try {
320+
queue.retryTask(task);
321+
} catch (Exception e) {
322+
// The existing suspending-task safety net will still reclaim it after the removal.
323+
LOG.warn(
324+
"Failed to hand back task {} from draining optimizer {}",
325+
task.getTaskId(),
326+
authToken,
327+
e);
328+
}
329+
}
330+
return null;
331+
}
332+
333+
/** Block new task assignments to the token; in-flight tasks keep completing normally. */
334+
void beginGracefulDrain(String token, long deadlineMs) {
335+
drainDeadlines.put(token, deadlineMs);
336+
pendingRemovalTokens.add(token);
337+
LOG.info("Optimizer {} begins graceful drain", token);
338+
}
339+
340+
/** Re-admit the token to task assignment, e.g. when dynamic allocation is disabled mid-drain. */
341+
void cancelDrain(String token) {
342+
pendingRemovalTokens.remove(token);
343+
drainDeadlines.remove(token);
344+
}
345+
346+
@VisibleForTesting
347+
boolean isDraining(String token) {
348+
return pendingRemovalTokens.contains(token);
349+
}
350+
351+
/**
352+
* Run one dynamic-allocation round for the group at an injected time. The production cadence is
353+
* driven by the scale keeper's delay queue with the wall clock; tests inject times because the
354+
* validated minimum {@code executor-idle-timeout} (30s) puts real idle waits beyond sane test
355+
* durations.
356+
*/
357+
@VisibleForTesting
358+
void evaluateDynamicAllocation(String groupName, long nowMs) {
359+
ResourceGroup resourceGroup = optimizerManager.getResourceGroup(groupName);
360+
OptimizingQueue queue = optimizingQueueByGroup.get(groupName);
361+
optimizerScaleKeeper.scaleIfNeeded(
362+
resourceGroup, queue, DynamicAllocationConfig.parse(resourceGroup), nowMs);
363+
}
364+
365+
/**
366+
* Remove a drained optimizer: release the container resource, delete the persisted resource row,
367+
* and unregister. A missing resource row (the pod self-registered after a persist failure, or a
368+
* manual release raced the row away) is not an error — the instance itself carries the
369+
* container-side identity, so release through it and only skip the row delete; treating this as a
370+
* retryable failure would loop forever on a pod whose row can never reappear. A container release
371+
* failure keeps the drain state so a later round retries the idempotent deletion.
372+
*/
373+
void executeRemoval(String token) {
374+
OptimizerInstance optimizer = authOptimizers.get(token);
375+
if (optimizer == null || optimizer.getResourceId() == null) {
376+
// Already unregistered, or externally launched: nothing for AMS to release.
377+
cancelDrain(token);
378+
return;
379+
}
380+
try {
381+
Resource resource = optimizerManager.getResource(optimizer.getResourceId());
382+
if (resource != null) {
383+
resource.getProperties().putAll(optimizer.getProperties());
384+
((AbstractOptimizerContainer) Containers.get(resource.getContainerName()))
385+
.releaseResource(resource);
386+
optimizerManager.deleteResource(optimizer.getResourceId());
387+
} else {
388+
((AbstractOptimizerContainer) Containers.get(optimizer.getContainerName()))
389+
.releaseResource(optimizer);
390+
}
391+
} catch (Throwable t) {
392+
LOG.warn(
393+
"Failed to release optimizer {} (resource {}), will retry",
394+
token,
395+
optimizer.getResourceId(),
396+
t);
397+
return;
398+
}
399+
unregisterOptimizer(token);
400+
LOG.info("Optimizer {} (resource {}) removed by scale-down", token, optimizer.getResourceId());
401+
}
402+
283403
@Override
284404
public void ackTask(String authToken, int threadId, OptimizingTaskId taskId) {
285405
LOG.info("Ack task {} by optimizer {} (threadId {})", taskId, authToken, threadId);
@@ -1046,6 +1166,12 @@ private void unwatch(String groupName) {
10461166
watchedGroups.remove(groupName);
10471167
scaleStates.remove(groupName);
10481168
planningBoundStreaks.remove(groupName);
1169+
// A drain block left behind would starve the group's pods forever once the legacy floor
1170+
// keeper resumes duty for the disabled group: re-admit them to task assignment.
1171+
authOptimizers.values().stream()
1172+
.filter(optimizer -> groupName.equals(optimizer.getGroupName()))
1173+
.map(OptimizerInstance::getToken)
1174+
.forEach(DefaultOptimizingService.this::cancelDrain);
10491175
// pendingRegistrations is deliberately kept: a pod requested before a disable survives its
10501176
// boot window, so re-enabling within it does not re-request the same capacity. Entries
10511177
// self-prune past their deadline.
@@ -1096,7 +1222,7 @@ protected void processTask(DraScaleTask task) {
10961222
}
10971223
DynamicAllocationConfig config = DynamicAllocationConfig.parse(resourceGroup);
10981224
try {
1099-
scaleIfNeeded(resourceGroup, queue, config);
1225+
scaleIfNeeded(resourceGroup, queue, config, System.currentTimeMillis());
11001226
} catch (Throwable t) {
11011227
LOG.error("Dynamic allocation scale evaluation failed for group {}", task.groupName, t);
11021228
} finally {
@@ -1111,6 +1237,81 @@ private int pendingThreads(String groupName) {
11111237
return pending == null ? 0 : pending.pendingThreads(System.currentTimeMillis());
11121238
}
11131239

1240+
/**
1241+
* Advance this group's drains: an entry whose in-flight count reached zero, or whose {@code
1242+
* drain-timeout} deadline passed, executes its removal now (a force-removed instance's orphaned
1243+
* tasks are reclaimed by the existing suspending-task safety net). Returns the thread and
1244+
* busy-task counts of instances still draining afterwards — a failed release keeps its instance
1245+
* in both, since it remains registered.
1246+
*/
1247+
private int[] processDrainProgress(
1248+
String groupName, DynamicAllocationState.GroupLoad load, long now) {
1249+
int drainingThreads = 0;
1250+
int drainingBusy = 0;
1251+
for (String token : pendingRemovalTokens) {
1252+
OptimizerInstance optimizer = authOptimizers.get(token);
1253+
if (optimizer == null) {
1254+
// Unregistered mid-drain (e.g. its heartbeat expired): nothing left to remove.
1255+
cancelDrain(token);
1256+
continue;
1257+
}
1258+
if (!groupName.equals(optimizer.getGroupName())) {
1259+
continue;
1260+
}
1261+
int inFlight = load.getInFlightByToken().getOrDefault(token, 0);
1262+
Long deadline = drainDeadlines.get(token);
1263+
if (inFlight == 0 || (deadline != null && now >= deadline)) {
1264+
executeRemoval(token);
1265+
if (!authOptimizers.containsKey(token)) {
1266+
continue;
1267+
}
1268+
}
1269+
drainingThreads += optimizer.getThreadCount();
1270+
drainingBusy += inFlight;
1271+
}
1272+
return new int[] {drainingThreads, drainingBusy};
1273+
}
1274+
1275+
private Set<String> registeredTokens(String groupName) {
1276+
return authOptimizers.values().stream()
1277+
.filter(optimizer -> groupName.equals(optimizer.getGroupName()))
1278+
.map(OptimizerInstance::getToken)
1279+
.collect(Collectors.toSet());
1280+
}
1281+
1282+
private void evaluateScaleDown(
1283+
String groupName,
1284+
OptimizingQueue queue,
1285+
DynamicAllocationState state,
1286+
DynamicAllocationConfig config,
1287+
int registeredThreads,
1288+
int drainingThreads,
1289+
long now) {
1290+
List<DynamicAllocationState.RemovalCandidate> candidates =
1291+
authOptimizers.values().stream()
1292+
// Externally-registered optimizers (no resourceId) are not AMS's to remove.
1293+
.filter(optimizer -> groupName.equals(optimizer.getGroupName()))
1294+
.filter(optimizer -> optimizer.getResourceId() != null)
1295+
.filter(optimizer -> !pendingRemovalTokens.contains(optimizer.getToken()))
1296+
.map(
1297+
optimizer ->
1298+
new DynamicAllocationState.RemovalCandidate(
1299+
optimizer.getToken(), optimizer.getThreadCount()))
1300+
.collect(Collectors.toList());
1301+
String victim =
1302+
state.computeScaleDown(candidates, registeredThreads, drainingThreads, config, now);
1303+
if (victim == null) {
1304+
return;
1305+
}
1306+
beginGracefulDrain(victim, now + config.getDrainTimeout().toMillis());
1307+
// Only a snapshot taken after the token entered the pending-removal set can prove idleness:
1308+
// the pre-insert one may miss a task fetched by a long-poll racing the drain start.
1309+
DynamicAllocationState.GroupLoad fresh = queue.collectDynamicAllocationLoad();
1310+
if (fresh.getInFlightByToken().getOrDefault(victim, 0) == 0) {
1311+
executeRemoval(victim);
1312+
}
1313+
}
1314+
11141315
private void recheckAfterUnwatch(String groupName) {
11151316
try {
11161317
ResourceGroup fresh = optimizerManager.getResourceGroup(groupName);
@@ -1157,27 +1358,46 @@ private void warnOnPlanningBoundTransition(
11571358
}
11581359

11591360
private void scaleIfNeeded(
1160-
ResourceGroup resourceGroup, OptimizingQueue queue, DynamicAllocationConfig config) {
1361+
ResourceGroup resourceGroup,
1362+
OptimizingQueue queue,
1363+
DynamicAllocationConfig config,
1364+
long now) {
11611365
String groupName = resourceGroup.getName();
1162-
long now = System.currentTimeMillis();
11631366
PendingRegistrations pending =
11641367
pendingRegistrations.computeIfAbsent(
11651368
groupName, name -> new PendingRegistrations(BOOT_TIMEOUT_MS));
11661369
DynamicAllocationState state =
11671370
scaleStates.computeIfAbsent(groupName, name -> new DynamicAllocationState());
1168-
int registeredThreads = getTotalQuota(groupName);
1169-
int effectiveThreads = registeredThreads + pending.pendingThreads(now);
11701371
DynamicAllocationState.GroupLoad load = queue.collectDynamicAllocationLoad();
1372+
// Drain progress runs before anything else and unconditionally: a completed or expired
1373+
// drain must convert to a removal even in rounds that scale up, or a busy drain would
1374+
// linger to its full timeout while backlog persists.
1375+
int[] draining = processDrainProgress(groupName, load, now);
1376+
int drainingThreads = draining[0];
1377+
int drainingBusy = draining[1];
1378+
int registeredThreads = getTotalQuota(groupName);
1379+
// A draining instance takes no new work, so it is accounted as already gone on both sides:
1380+
// leaving its threads in the capacity undercounts demand by up to their count, and leaving
1381+
// its tasks in the load keeps future demand (busy >= effective) from ever firing mid-drain.
1382+
int effectiveThreads = registeredThreads - drainingThreads + pending.pendingThreads(now);
1383+
int busyThreads = load.getBusyThreads() - drainingBusy;
1384+
// Observed every round, including scale-up ones: an instance busy through a burst must not
1385+
// come out of it looking idle since before the burst began.
1386+
state.observe(registeredTokens(groupName), load.getInFlightByToken(), now);
11711387
warnOnPlanningBoundTransition(groupName, registeredThreads, load);
11721388
int addInstances =
11731389
state.computeScaleUp(
11741390
effectiveThreads,
1175-
load.getBusyThreads(),
1391+
busyThreads,
11761392
load.getServiceablePlanned(),
11771393
load.getPendingTables(),
11781394
config,
11791395
now);
11801396
if (addInstances <= 0) {
1397+
if (!state.wasDemandActive()) {
1398+
evaluateScaleDown(
1399+
groupName, queue, state, config, registeredThreads, drainingThreads, now);
1400+
}
11811401
return;
11821402
}
11831403
int threadsPerInstance = config.getExecutorParallelism();

‎amoro-ams/src/main/java/org/apache/amoro/server/optimizing/OptimizingQueue.java‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -428,11 +428,13 @@ public List<TaskRuntime<?>> collectTasks(Predicate<TaskRuntime<?>> predicate) {
428428
public DynamicAllocationState.GroupLoad collectDynamicAllocationLoad() {
429429
Map<Long, Integer> plannedByTable = Maps.newHashMap();
430430
Map<Long, Integer> occupiedByTable = Maps.newHashMap();
431+
Map<String, Integer> inFlightByToken = Maps.newHashMap();
431432
int busyThreads = 0;
432433
for (TaskRuntime<?> task : collectTasks()) {
433434
if (DynamicAllocationState.occupiesThread(task.getStatus())) {
434435
busyThreads++;
435436
occupiedByTable.merge(task.getTableId(), 1, Integer::sum);
437+
inFlightByToken.merge(task.getToken(), 1, Integer::sum);
436438
} else if (task.getStatus() == Status.PLANNED) {
437439
plannedByTable.merge(task.getTableId(), 1, Integer::sum);
438440
}
@@ -457,7 +459,10 @@ public DynamicAllocationState.GroupLoad collectDynamicAllocationLoad() {
457459
targetQuotaByTable.getOrDefault(tableId, 1.0),
458460
occupiedByTable.getOrDefault(tableId, 0))));
459461
return new DynamicAllocationState.GroupLoad(
460-
busyThreads, DynamicAllocationState.serviceablePlannedCount(demands), pendingTables);
462+
busyThreads,
463+
DynamicAllocationState.serviceablePlannedCount(demands),
464+
pendingTables,
465+
inFlightByToken);
461466
}
462467

463468
public void retryTask(TaskRuntime<?> taskRuntime) {

‎amoro-ams/src/main/java/org/apache/amoro/server/optimizing/dra/DynamicAllocationConfig.java‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -338,6 +338,19 @@ public void validate() {
338338
OptimizerProperties.DYNAMIC_ALLOCATION_SUSTAINED_BACKLOG_TIMEOUT, sustainedBacklogTimeout);
339339
requirePositive(OptimizerProperties.DYNAMIC_ALLOCATION_SCALE_DOWN_COOLDOWN, scaleDownCooldown);
340340
requirePositive(OptimizerProperties.DYNAMIC_ALLOCATION_DRAIN_TIMEOUT, drainTimeout);
341+
// The scale keeper evaluates each group at sustained-backlog-timeout cadence, which is also
342+
// the idle observation resolution: sampling slower than half the idle timeout lets an
343+
// instance that worked between samples be misjudged as continuously idle and drained.
344+
if (sustainedBacklogTimeout.toMillis() * 2 > executorIdleTimeout.toMillis()) {
345+
throw new IllegalArgumentException(
346+
String.format(
347+
"Resource group:%s '%s'(%s) must be <= half of '%s'(%s).",
348+
groupName,
349+
OptimizerProperties.DYNAMIC_ALLOCATION_SUSTAINED_BACKLOG_TIMEOUT,
350+
sustainedBacklogTimeout,
351+
OptimizerProperties.DYNAMIC_ALLOCATION_EXECUTOR_IDLE_TIMEOUT,
352+
executorIdleTimeout));
353+
}
341354
}
342355

343356
private void requirePositive(String property, Duration value) {

0 commit comments

Comments
 (0)