diff --git a/amoro-ams/src/main/java/org/apache/amoro/server/manager/AbstractOptimizerContainer.java b/amoro-ams/src/main/java/org/apache/amoro/server/manager/AbstractOptimizerContainer.java index 05745c79b3..118fa25cc7 100644 --- a/amoro-ams/src/main/java/org/apache/amoro/server/manager/AbstractOptimizerContainer.java +++ b/amoro-ams/src/main/java/org/apache/amoro/server/manager/AbstractOptimizerContainer.java @@ -92,6 +92,11 @@ protected String buildOptimizerStartupArgsString(Resource resource) { .append(" -hb ") .append(resource.getProperties().get(OptimizerProperties.OPTIMIZER_HEART_BEAT_INTERVAL)); } + if (resource.getProperties().containsKey(OptimizerProperties.OPTIMIZER_SHUTDOWN_TIMEOUT_MS)) { + stringBuilder + .append(" -st ") + .append(resource.getProperties().get(OptimizerProperties.OPTIMIZER_SHUTDOWN_TIMEOUT_MS)); + } if (PropertyUtil.propertyAsBoolean( resource.getProperties(), OptimizerProperties.OPTIMIZER_EXTEND_DISK_STORAGE, diff --git a/amoro-ams/src/main/java/org/apache/amoro/server/manager/KubernetesOptimizerContainer.java b/amoro-ams/src/main/java/org/apache/amoro/server/manager/KubernetesOptimizerContainer.java index 850fba4eb3..8326688115 100644 --- a/amoro-ams/src/main/java/org/apache/amoro/server/manager/KubernetesOptimizerContainer.java +++ b/amoro-ams/src/main/java/org/apache/amoro/server/manager/KubernetesOptimizerContainer.java @@ -34,10 +34,12 @@ import io.fabric8.kubernetes.client.ConfigBuilder; import io.fabric8.kubernetes.client.KubernetesClient; import io.fabric8.kubernetes.client.KubernetesClientBuilder; +import org.apache.amoro.OptimizerProperties; import org.apache.amoro.resource.Resource; import org.apache.amoro.shade.guava32.com.google.common.base.Preconditions; import org.apache.amoro.shade.guava32.com.google.common.collect.ImmutableMap; import org.apache.amoro.shade.guava32.com.google.common.collect.Maps; +import org.apache.amoro.utils.PropertyUtil; import org.apache.commons.io.IOUtils; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -75,6 +77,13 @@ public class KubernetesOptimizerContainer extends AbstractOptimizerContainer { private static final String EXTRA_PROPERTY_PREFIX = "extra."; + /** + * Buffer added on top of optimizer's shutdown-timeout when deriving K8s + * terminationGracePeriodSeconds, to leave room for thread join, best-effort task result + * reporting, and log flush before SIGKILL. + */ + static final long TERMINATION_GRACE_BUFFER_SECONDS = 30L; + private static final Map EXTRA_PROPERTY_DEFAULTS = new HashMap<>(); static { @@ -86,6 +95,16 @@ private String getExtraProperty(Map properties, String key) { EXTRA_PROPERTY_PREFIX + key, EXTRA_PROPERTY_DEFAULTS.getOrDefault(key, null)); } + static long resolveTerminationGracePeriodSeconds(Map properties) { + long shutdownTimeoutMs = + PropertyUtil.propertyAsLong( + properties, + OptimizerProperties.OPTIMIZER_SHUTDOWN_TIMEOUT_MS, + OptimizerProperties.OPTIMIZER_SHUTDOWN_TIMEOUT_MS_DEFAULT); + long shutdownTimeoutSeconds = (shutdownTimeoutMs + 999L) / 1000L; + return shutdownTimeoutSeconds + TERMINATION_GRACE_BUFFER_SECONDS; + } + private KubernetesClient client; @Override @@ -118,6 +137,7 @@ protected Map doScaleOut(Resource resource) { String groupName = argsList.get("groupName").toString(); String resourceId = argsList.get("resourceId").toString(); String startUpArgs = argsList.get("startUpArgs").toString(); + long terminationGraceSeconds = (long) argsList.get("terminationGraceSeconds"); String kubernetesName = NAME_PREFIX + resourceId; Deployment deployment; @@ -136,7 +156,8 @@ protected Map doScaleOut(Resource resource) { resourceId, startUpArgs, memory, - imagePullSecretsList); + imagePullSecretsList, + terminationGraceSeconds); } else { deployment = initPodTemplateWithoutConfig( @@ -147,7 +168,8 @@ protected Map doScaleOut(Resource resource) { resourceId, startUpArgs, memory, - imagePullSecretsList); + imagePullSecretsList, + terminationGraceSeconds); } client.apps().deployments().inNamespace(namespace).resource(deployment).create(); @@ -206,6 +228,11 @@ public Map generatePodStartArgs( argsList.put("resourceId", resourceId); argsList.put("groupName", groupName); argsList.put("startUpArgs", startUpArgs); + // Derive the grace period from the same resource properties that drive the -st startup arg + // (buildOptimizerStartupArgsString), so the pod's SIGKILL deadline can never fall below the + // shutdown timeout the optimizer JVM actually honors. + argsList.put( + "terminationGraceSeconds", resolveTerminationGracePeriodSeconds(resource.getProperties())); return argsList; } @@ -218,7 +245,8 @@ public Deployment initPodTemplateWithoutConfig( String resourceId, String startUpArgs, long memory, - List imagePullSecretsList) { + List imagePullSecretsList, + long terminationGraceSeconds) { DeploymentBuilder deploymentBuilder = new DeploymentBuilder() @@ -234,11 +262,12 @@ public Deployment initPodTemplateWithoutConfig( .addToLabels("AmoroResourceId", resourceId) .endMetadata() .withNewSpec() + .withTerminationGracePeriodSeconds(terminationGraceSeconds) .addNewContainer() .withName("optimizer") .withImage(image) .withImagePullPolicy(pullPolicy) - .withCommand("sh", "-c", startUpArgs) + .withCommand("sh", "-c", "exec " + startUpArgs) .withResources( new ResourceRequirementsBuilder() .withLimits( @@ -288,7 +317,8 @@ public Deployment initPodTemplateFromFrontEnd( String resourceId, String startUpArgs, long memory, - List imagePullSecretsList) { + List imagePullSecretsList, + long terminationGraceSeconds) { podTemplate .getTemplate() .getMetadata() @@ -305,7 +335,7 @@ public Deployment initPodTemplateFromFrontEnd( container.setName("optimizer"); container.setImage(image); container.setImagePullPolicy(pullPolicy); - container.setCommand(new ArrayList<>(Arrays.asList("sh", "-c", startUpArgs))); + container.setCommand(new ArrayList<>(Arrays.asList("sh", "-c", "exec " + startUpArgs))); ResourceRequirements resourceRequirements = new ResourceRequirements(); resourceRequirements.setLimits( @@ -324,6 +354,10 @@ public Deployment initPodTemplateFromFrontEnd( podTemplate.getTemplate().getSpec().setImagePullSecrets(imagePullSecretsList); } + if (podTemplate.getTemplate().getSpec().getTerminationGracePeriodSeconds() == null) { + podTemplate.getTemplate().getSpec().setTerminationGracePeriodSeconds(terminationGraceSeconds); + } + DeploymentSpec deploymentSpec = new DeploymentSpec(); deploymentSpec.setTemplate(podTemplate.getTemplate()); diff --git a/amoro-ams/src/test/java/org/apache/amoro/server/manager/TestKubernetesOptimizerContainer.java b/amoro-ams/src/test/java/org/apache/amoro/server/manager/TestKubernetesOptimizerContainer.java index 4178870526..7c117b8470 100644 --- a/amoro-ams/src/test/java/org/apache/amoro/server/manager/TestKubernetesOptimizerContainer.java +++ b/amoro-ams/src/test/java/org/apache/amoro/server/manager/TestKubernetesOptimizerContainer.java @@ -130,6 +130,7 @@ public void testBuildPodTemplateWithResourceSetMemoryMb() { String groupName = argsList.get("groupName").toString(); String resourceId = argsList.get("resourceId").toString(); String startUpArgs = argsList.get("startUpArgs").toString(); + long terminationGraceSeconds = (long) argsList.get("terminationGraceSeconds"); Deployment deployment = kubernetesOptimizerContainer.initPodTemplateFromFrontEnd( @@ -141,7 +142,8 @@ public void testBuildPodTemplateWithResourceSetMemoryMb() { resourceId, startUpArgs, memory, - imagePullSecretsList); + imagePullSecretsList, + terminationGraceSeconds); Assert.assertEquals(1, deployment.getSpec().getReplicas().intValue()); Assert.assertNotEquals( @@ -217,6 +219,8 @@ public void testBuildPodTemplateConfig() { String resourceId = resource.getResourceId(); String groupName = resource.getGroupName(); + long terminationGraceSeconds = + KubernetesOptimizerContainer.resolveTerminationGracePeriodSeconds(groupProperties); Assert.assertEquals(1, podTemplate.getTemplate().getSpec().getContainers().size()); @@ -234,7 +238,8 @@ public void testBuildPodTemplateConfig() { resourceId, startUpArgs, memory, - imagePullSecretsList); + imagePullSecretsList, + terminationGraceSeconds); Assert.assertEquals("amoro-optimizer-" + resourceId, deployment.getMetadata().getName()); Assert.assertEquals( @@ -278,6 +283,269 @@ public void testBuildPodTemplateConfig() { .toString()); } + @Test + public void testContainerCommandUsesExecForSignalForwarding() { + ResourceType resourceType = ResourceType.OPTIMIZER; + Map properties = Maps.newHashMap(); + properties.put("memory", "1024"); + Resource resource = + new Resource.Builder("KubernetesContainer", "k8s", resourceType) + .setMemoryMb(1024) + .setThreadCount(1) + .setProperties(properties) + .build(); + groupProperties.putAll(resource.getProperties()); + + Map argsList = + kubernetesOptimizerContainer.generatePodStartArgs(resource, groupProperties); + String image = argsList.get(IMAGE).toString(); + String pullPolicy = argsList.get(PULL_POLICY).toString(); + List imagePullSecretsList = + (List) argsList.get(PULL_SECRETS); + int cpuLimit = (int) argsList.get("cpuLimit"); + long memory = (long) argsList.get(MEMORY_PROPERTY); + String groupName = argsList.get("groupName").toString(); + String resourceId = argsList.get("resourceId").toString(); + String startUpArgs = argsList.get("startUpArgs").toString(); + long terminationGraceSeconds = (long) argsList.get("terminationGraceSeconds"); + + Deployment deployment = + kubernetesOptimizerContainer.initPodTemplateWithoutConfig( + image, + pullPolicy, + cpuLimit, + groupName, + resourceId, + startUpArgs, + memory, + imagePullSecretsList, + terminationGraceSeconds); + + List command = + deployment.getSpec().getTemplate().getSpec().getContainers().get(0).getCommand(); + Assert.assertEquals(Arrays.asList("sh", "-c", "exec " + startUpArgs), command); + } + + @Test + public void testContainerCommandUsesExecWithPodTemplate() { + PodTemplate podTemplate = + kubernetesOptimizerContainer.initPodTemplateFromLocal(groupProperties); + + ResourceType resourceType = ResourceType.OPTIMIZER; + Map properties = Maps.newHashMap(); + properties.put("memory", "1024"); + Resource resource = + new Resource.Builder("KubernetesContainer", "k8s", resourceType) + .setMemoryMb(1024) + .setThreadCount(1) + .setProperties(properties) + .build(); + groupProperties.putAll(resource.getProperties()); + + Map argsList = + kubernetesOptimizerContainer.generatePodStartArgs(resource, groupProperties); + String image = argsList.get(IMAGE).toString(); + String pullPolicy = argsList.get(PULL_POLICY).toString(); + List imagePullSecretsList = + (List) argsList.get(PULL_SECRETS); + int cpuLimit = (int) argsList.get("cpuLimit"); + long memory = (long) argsList.get(MEMORY_PROPERTY); + String groupName = argsList.get("groupName").toString(); + String resourceId = argsList.get("resourceId").toString(); + String startUpArgs = argsList.get("startUpArgs").toString(); + long terminationGraceSeconds = (long) argsList.get("terminationGraceSeconds"); + + Deployment deployment = + kubernetesOptimizerContainer.initPodTemplateFromFrontEnd( + podTemplate, + image, + pullPolicy, + cpuLimit, + groupName, + resourceId, + startUpArgs, + memory, + imagePullSecretsList, + terminationGraceSeconds); + + List command = + deployment.getSpec().getTemplate().getSpec().getContainers().get(0).getCommand(); + Assert.assertEquals(Arrays.asList("sh", "-c", "exec " + startUpArgs), command); + } + + @Test + public void testTerminationGracePeriodFromDefaultShutdownTimeout() { + ResourceType resourceType = ResourceType.OPTIMIZER; + Map properties = Maps.newHashMap(); + properties.put("memory", "1024"); + Resource resource = + new Resource.Builder("KubernetesContainer", "k8s", resourceType) + .setMemoryMb(1024) + .setThreadCount(1) + .setProperties(properties) + .build(); + groupProperties.putAll(resource.getProperties()); + + Map argsList = + kubernetesOptimizerContainer.generatePodStartArgs(resource, groupProperties); + String image = argsList.get(IMAGE).toString(); + String pullPolicy = argsList.get(PULL_POLICY).toString(); + List imagePullSecretsList = + (List) argsList.get(PULL_SECRETS); + int cpuLimit = (int) argsList.get("cpuLimit"); + long memory = (long) argsList.get(MEMORY_PROPERTY); + String groupName = argsList.get("groupName").toString(); + String resourceId = argsList.get("resourceId").toString(); + String startUpArgs = argsList.get("startUpArgs").toString(); + long terminationGraceSeconds = (long) argsList.get("terminationGraceSeconds"); + + Assert.assertFalse( + "startUpArgs should not contain the -st flag when the property is unset, but was: " + + startUpArgs, + startUpArgs.contains(" -st ")); + + Deployment deployment = + kubernetesOptimizerContainer.initPodTemplateWithoutConfig( + image, + pullPolicy, + cpuLimit, + groupName, + resourceId, + startUpArgs, + memory, + imagePullSecretsList, + terminationGraceSeconds); + + Long grace = deployment.getSpec().getTemplate().getSpec().getTerminationGracePeriodSeconds(); + Assert.assertNotNull(grace); + // default shutdown-timeout = 600_000ms → 600s + 30s buffer + Assert.assertEquals(630L, grace.longValue()); + } + + @Test + public void testShutdownTimeoutPropertyWiredToStartupArgsAndGracePeriod() { + ResourceType resourceType = ResourceType.OPTIMIZER; + Map properties = Maps.newHashMap(); + properties.put("memory", "1024"); + properties.put(OptimizerProperties.OPTIMIZER_SHUTDOWN_TIMEOUT_MS, "120000"); + Resource resource = + new Resource.Builder("KubernetesContainer", "k8s", resourceType) + .setMemoryMb(1024) + .setThreadCount(1) + .setProperties(properties) + .build(); + groupProperties.putAll(resource.getProperties()); + + Map argsList = + kubernetesOptimizerContainer.generatePodStartArgs(resource, groupProperties); + String image = argsList.get(IMAGE).toString(); + String pullPolicy = argsList.get(PULL_POLICY).toString(); + List imagePullSecretsList = + (List) argsList.get(PULL_SECRETS); + int cpuLimit = (int) argsList.get("cpuLimit"); + long memory = (long) argsList.get(MEMORY_PROPERTY); + String groupName = argsList.get("groupName").toString(); + String resourceId = argsList.get("resourceId").toString(); + String startUpArgs = argsList.get("startUpArgs").toString(); + long terminationGraceSeconds = (long) argsList.get("terminationGraceSeconds"); + + Assert.assertTrue( + "startUpArgs should contain the -st flag, but was: " + startUpArgs, + startUpArgs.contains(" -st 120000")); + + Deployment deployment = + kubernetesOptimizerContainer.initPodTemplateWithoutConfig( + image, + pullPolicy, + cpuLimit, + groupName, + resourceId, + startUpArgs, + memory, + imagePullSecretsList, + terminationGraceSeconds); + + Long grace = deployment.getSpec().getTemplate().getSpec().getTerminationGracePeriodSeconds(); + // 120_000ms → 120s + 30s buffer + Assert.assertEquals(Long.valueOf(150L), grace); + } + + @Test + public void testContainerLevelShutdownTimeoutDoesNotDesyncGracePeriod() { + ResourceType resourceType = ResourceType.OPTIMIZER; + Map properties = Maps.newHashMap(); + properties.put("memory", "1024"); + Resource resource = + new Resource.Builder("KubernetesContainer", "k8s", resourceType) + .setMemoryMb(1024) + .setThreadCount(1) + .setProperties(properties) + .build(); + groupProperties.putAll(resource.getProperties()); + // Simulate a container-level property: present in the merged group properties but NOT in + // the resource properties that drive the -st startup arg. The grace period must stay in + // sync with the arg (i.e. keep the default), not silently shrink below the JVM's timeout. + groupProperties.put(OptimizerProperties.OPTIMIZER_SHUTDOWN_TIMEOUT_MS, "120000"); + + Map argsList = + kubernetesOptimizerContainer.generatePodStartArgs(resource, groupProperties); + String startUpArgs = argsList.get("startUpArgs").toString(); + long terminationGraceSeconds = (long) argsList.get("terminationGraceSeconds"); + + Assert.assertFalse( + "startUpArgs should not contain -st for a container-level property: " + startUpArgs, + startUpArgs.contains(" -st ")); + // 600_000ms default → 600s + 30s buffer, consistent with the args the pod receives + Assert.assertEquals(630L, terminationGraceSeconds); + } + + @Test + public void testTerminationGracePeriodFromUserPodTemplateRespected() { + PodTemplate podTemplate = + kubernetesOptimizerContainer.initPodTemplateFromLocal(groupProperties); + podTemplate.getTemplate().getSpec().setTerminationGracePeriodSeconds(900L); + + ResourceType resourceType = ResourceType.OPTIMIZER; + Map properties = Maps.newHashMap(); + properties.put("memory", "1024"); + Resource resource = + new Resource.Builder("KubernetesContainer", "k8s", resourceType) + .setMemoryMb(1024) + .setThreadCount(1) + .setProperties(properties) + .build(); + groupProperties.putAll(resource.getProperties()); + + Map argsList = + kubernetesOptimizerContainer.generatePodStartArgs(resource, groupProperties); + String image = argsList.get(IMAGE).toString(); + String pullPolicy = argsList.get(PULL_POLICY).toString(); + List imagePullSecretsList = + (List) argsList.get(PULL_SECRETS); + int cpuLimit = (int) argsList.get("cpuLimit"); + long memory = (long) argsList.get(MEMORY_PROPERTY); + String groupName = argsList.get("groupName").toString(); + String resourceId = argsList.get("resourceId").toString(); + String startUpArgs = argsList.get("startUpArgs").toString(); + long terminationGraceSeconds = (long) argsList.get("terminationGraceSeconds"); + + Deployment deployment = + kubernetesOptimizerContainer.initPodTemplateFromFrontEnd( + podTemplate, + image, + pullPolicy, + cpuLimit, + groupName, + resourceId, + startUpArgs, + memory, + imagePullSecretsList, + terminationGraceSeconds); + + Long grace = deployment.getSpec().getTemplate().getSpec().getTerminationGracePeriodSeconds(); + Assert.assertEquals(Long.valueOf(900L), grace); + } + @Test public void testAMSWithConfigMap() throws Exception { ConfigMap configMap = buildConfigMap(); @@ -325,6 +593,8 @@ public void testAMSWithConfigMap() throws Exception { String resourceId = resource.getResourceId(); String groupName = resource.getGroupName(); + long terminationGraceSeconds = + KubernetesOptimizerContainer.resolveTerminationGracePeriodSeconds(groupProperties); Assert.assertEquals(1, podTemplate.getTemplate().getSpec().getContainers().size()); Assert.assertEquals( @@ -340,7 +610,8 @@ public void testAMSWithConfigMap() throws Exception { resourceId, startUpArgs, memory, - imagePullSecretsList); + imagePullSecretsList, + terminationGraceSeconds); // Assert the Deployment Assert.assertEquals("amoro-optimizer-" + resourceId, deployment.getMetadata().getName()); diff --git a/amoro-common/src/main/java/org/apache/amoro/OptimizerProperties.java b/amoro-common/src/main/java/org/apache/amoro/OptimizerProperties.java index dd4a6f624e..3c5e7f9008 100644 --- a/amoro-common/src/main/java/org/apache/amoro/OptimizerProperties.java +++ b/amoro-common/src/main/java/org/apache/amoro/OptimizerProperties.java @@ -49,4 +49,6 @@ public class OptimizerProperties { public static final String OPTIMIZER_CACHE_TIMEOUT = "cache-timeout"; public static final String OPTIMIZER_CACHE_TIMEOUT_DEFAULT = "10min"; public static final String OPTIMIZER_MASTER_SLAVE_MODE_ENABLED = "master-slave-mode-enabled"; + public static final String OPTIMIZER_SHUTDOWN_TIMEOUT_MS = "shutdown-timeout-ms"; + public static final long OPTIMIZER_SHUTDOWN_TIMEOUT_MS_DEFAULT = 600_000L; // 10 min } diff --git a/amoro-common/src/test/java/org/apache/amoro/MockAmoroManagementServer.java b/amoro-common/src/test/java/org/apache/amoro/MockAmoroManagementServer.java index ae8f462859..058b598545 100644 --- a/amoro-common/src/test/java/org/apache/amoro/MockAmoroManagementServer.java +++ b/amoro-common/src/test/java/org/apache/amoro/MockAmoroManagementServer.java @@ -412,6 +412,15 @@ public static class OptimizerManagerHandler implements OptimizingService.Iface { new ConcurrentHashMap<>(); private final Map> completedTasks = new ConcurrentHashMap<>(); + private final AtomicInteger completeTaskFailures = new AtomicInteger(0); + + /** + * Injects transient failures: the next {@code count} completeTask calls throw a retryable + * error. + */ + public void failNextCompleteTasks(int count) { + completeTaskFailures.set(count); + } public void cleanUp() {} @@ -449,6 +458,12 @@ public void ackTask(String authToken, int threadId, OptimizingTaskId taskId) thr @Override public void completeTask(String authToken, OptimizingTaskResult taskResult) throws TException { checkToken(authToken); + if (completeTaskFailures.getAndUpdate(c -> c > 0 ? c - 1 : 0) > 0) { + throw new AmoroException( + ErrorCodes.PERSISTENCE_ERROR_CODE, + "InjectedTransientError", + "injected transient error for testing"); + } executingTasks.get(authToken).remove(taskResult.getThreadId()); if (!completedTasks.containsKey(authToken)) { completedTasks.putIfAbsent(authToken, new CopyOnWriteArrayList<>()); diff --git a/amoro-optimizer/amoro-optimizer-common/src/main/java/org/apache/amoro/optimizer/common/AbstractOptimizerOperator.java b/amoro-optimizer/amoro-optimizer-common/src/main/java/org/apache/amoro/optimizer/common/AbstractOptimizerOperator.java index af002bc0fd..93acb42b49 100644 --- a/amoro-optimizer/amoro-optimizer-common/src/main/java/org/apache/amoro/optimizer/common/AbstractOptimizerOperator.java +++ b/amoro-optimizer/amoro-optimizer-common/src/main/java/org/apache/amoro/optimizer/common/AbstractOptimizerOperator.java @@ -33,6 +33,7 @@ import java.util.List; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicReference; +import java.util.function.BooleanSupplier; public class AbstractOptimizerOperator implements Serializable { private static final Logger LOG = LoggerFactory.getLogger(AbstractOptimizerOperator.class); @@ -219,11 +220,44 @@ private boolean shouldReturnNull(Throwable t) { */ protected T callAuthenticatedAms(String amsUrl, AmsAuthenticatedCallOperation operation) throws TException { + return callAuthenticatedAms(amsUrl, operation, this::isStarted); + } + + /** + * Variant for the shutdown drain: keeps the retry loop alive for a bounded window even after + * stop(), so a result produced by an in-flight task is not dropped on the first transient error + * or because stop() landed between the caller's check and the retry loop. The window is anchored + * at the moment stop() is first observed, not at call entry — a call that already spent time + * retrying while running still gets the full drain window once shutdown begins. Aborts early when + * the calling thread is force-interrupted by the shutdown deadline. + */ + protected T callAuthenticatedAmsWithDrain( + String amsUrl, AmsAuthenticatedCallOperation operation, long drainTimeoutMs) + throws TException { + long[] drainDeadline = {-1}; + return callAuthenticatedAms( + amsUrl, + operation, + () -> { + if (isStarted()) { + return true; + } + if (drainDeadline[0] < 0) { + drainDeadline[0] = System.currentTimeMillis() + drainTimeoutMs; + } + return System.currentTimeMillis() < drainDeadline[0] + && !Thread.currentThread().isInterrupted(); + }); + } + + private T callAuthenticatedAms( + String amsUrl, AmsAuthenticatedCallOperation operation, BooleanSupplier proceed) + throws TException { // Per-node retry budget: in master-slave mode, limit consecutive shouldRetryLater retries so // that a permanently unreachable node does not block the multi-node polling loop indefinitely. int consecutiveRetries = 0; - while (isStarted()) { + while (proceed.getAsBoolean()) { if (tokenIsReady()) { String token = getToken(); try { @@ -345,7 +379,13 @@ protected void waitAShortTime(long waitTime) { try { TimeUnit.MILLISECONDS.sleep(waitTime); } catch (InterruptedException e) { - // ignore + // Preserve the flag only when stopping, so the shutdown path skips residual sleeps. + // While still running, swallow a stray interrupt as before: no caller loop clears + // the flag, so preserving it would turn every later sleep into an instant return + // and the poll/retry loops into a busy spin. + if (!isStarted()) { + Thread.currentThread().interrupt(); + } } } diff --git a/amoro-optimizer/amoro-optimizer-common/src/main/java/org/apache/amoro/optimizer/common/Optimizer.java b/amoro-optimizer/amoro-optimizer-common/src/main/java/org/apache/amoro/optimizer/common/Optimizer.java index fbff34fc46..cdf1f2a450 100644 --- a/amoro-optimizer/amoro-optimizer-common/src/main/java/org/apache/amoro/optimizer/common/Optimizer.java +++ b/amoro-optimizer/amoro-optimizer-common/src/main/java/org/apache/amoro/optimizer/common/Optimizer.java @@ -33,6 +33,7 @@ public class Optimizer { private final OptimizerConfig config; private final OptimizerToucher toucher; private final OptimizerExecutor[] executors; + private volatile Thread[] executorThreads; public Optimizer(OptimizerConfig config) { this(config, () -> new OptimizerToucher(config), (i) -> new OptimizerExecutor(config, i)); @@ -54,20 +55,89 @@ protected Optimizer( public void startOptimizing() { LOG.info("Starting optimizer with configuration:{}", config); - Arrays.stream(executors) - .forEach( - optimizerExecutor -> { - new Thread( - optimizerExecutor::start, - String.format("Optimizer-executor-%d", optimizerExecutor.getThreadId())) - .start(); - }); + Thread[] threads = new Thread[executors.length]; + for (int i = 0; i < executors.length; i++) { + threads[i] = + new Thread( + executors[i]::start, + String.format("Optimizer-executor-%d", executors[i].getThreadId())); + } + // Publish the fully-built array before starting any thread: a concurrent shutdown hook + // reads executorThreads, and volatile orders only the array reference, not element + // writes made after publication. + executorThreads = threads; + for (Thread thread : threads) { + thread.start(); + } toucher.withTokenChangeListener(new SetTokenToExecutors()).start(); } public void stopOptimizing() { - toucher.stop(); + LOG.info("Stopping optimizer, waiting for in-progress tasks to complete..."); + // Stop executors first so they don't poll new tasks, but keep the toucher alive + // so it keeps sending heartbeats. Otherwise AMS hits its heartbeat-timeout during + // long-running in-flight tasks, unregisters this optimizer, and the subsequent + // best-effort completeTask fails with "Optimizer has not been authenticated". + // Drain mode stops the toucher from re-registering if AMS has already unregistered + // this optimizer (scale-down releases the resource before the pod gets SIGTERM). + toucher.enterDrainMode(); Arrays.stream(executors).forEach(OptimizerExecutor::stop); + + Thread[] threads = executorThreads; + if (threads == null) { + toucher.stop(); + LOG.info("Optimizer stopped (no executor threads to wait for)"); + return; + } + + long shutdownTimeoutMs = config.getShutdownTimeoutMs(); + boolean selfInterrupted = !joinAll(threads, shutdownTimeoutMs); + + // Force-interrupt pass: interrupt ALL survivors first, then wait for them against one + // shared deadline. Interleaving interrupt with per-thread joins would both serialize + // the waits (N extra seconds past the deadline) and, when this thread was itself + // interrupted above, skip interrupting the remaining threads because join() throws + // instantly while the interrupt flag is set. + boolean anyAlive = false; + for (Thread t : threads) { + if (t.isAlive()) { + anyAlive = true; + LOG.warn( + "Executor thread {} did not terminate within {}ms timeout, force-interrupting", + t.getName(), + shutdownTimeoutMs); + t.interrupt(); + } + } + if (anyAlive) { + selfInterrupted |= !joinAll(threads, 1_000); + } + toucher.stop(); + if (selfInterrupted) { + Thread.currentThread().interrupt(); + } + LOG.info("Optimizer stopped"); + } + + /** + * Joins all threads against one shared deadline. Returns false if the calling thread was + * interrupted while waiting; the interrupt flag is left cleared so subsequent joins work. + */ + private static boolean joinAll(Thread[] threads, long timeoutMs) { + long deadline = System.currentTimeMillis() + timeoutMs; + for (Thread t : threads) { + long remaining = deadline - System.currentTimeMillis(); + if (remaining <= 0) { + return true; + } + try { + t.join(remaining); + } catch (InterruptedException e) { + LOG.warn("Interrupted while waiting for executor thread {} to finish", t.getName()); + return false; + } + } + return true; } public OptimizerToucher getToucher() { diff --git a/amoro-optimizer/amoro-optimizer-common/src/main/java/org/apache/amoro/optimizer/common/OptimizerConfig.java b/amoro-optimizer/amoro-optimizer-common/src/main/java/org/apache/amoro/optimizer/common/OptimizerConfig.java index 2c70a2652a..ebe8c63d71 100644 --- a/amoro-optimizer/amoro-optimizer-common/src/main/java/org/apache/amoro/optimizer/common/OptimizerConfig.java +++ b/amoro-optimizer/amoro-optimizer-common/src/main/java/org/apache/amoro/optimizer/common/OptimizerConfig.java @@ -109,6 +109,15 @@ public class OptimizerConfig implements Serializable { usage = "Enable master-slave mode") private boolean masterSlaveMode = false; + @Option( + name = "-st", + aliases = "--" + OptimizerProperties.OPTIMIZER_SHUTDOWN_TIMEOUT_MS, + usage = + "Graceful shutdown timeout(ms) — wait this long for in-progress tasks before " + + "force-interrupting. Default 600000 (10min). Align with K8s " + + "terminationGracePeriodSeconds in containerized deployments.") + private long shutdownTimeoutMs = OptimizerProperties.OPTIMIZER_SHUTDOWN_TIMEOUT_MS_DEFAULT; + public OptimizerConfig() {} public OptimizerConfig(String[] args) throws CmdLineException { @@ -228,6 +237,14 @@ public void setMasterSlaveMode(boolean masterSlaveMode) { this.masterSlaveMode = masterSlaveMode; } + public long getShutdownTimeoutMs() { + return shutdownTimeoutMs; + } + + public void setShutdownTimeoutMs(long shutdownTimeoutMs) { + this.shutdownTimeoutMs = shutdownTimeoutMs; + } + @Override public String toString() { return MoreObjects.toStringHelper(this) @@ -245,6 +262,7 @@ public String toString() { .add("cacheMaxEntrySize", cacheMaxEntrySize) .add("cacheTimeout", cacheTimeout) .add("masterSlaveMode", masterSlaveMode) + .add("shutdownTimeoutMs", shutdownTimeoutMs) .toString(); } } diff --git a/amoro-optimizer/amoro-optimizer-common/src/main/java/org/apache/amoro/optimizer/common/OptimizerExecutor.java b/amoro-optimizer/amoro-optimizer-common/src/main/java/org/apache/amoro/optimizer/common/OptimizerExecutor.java index 26eda902cc..2868cf0839 100644 --- a/amoro-optimizer/amoro-optimizer-common/src/main/java/org/apache/amoro/optimizer/common/OptimizerExecutor.java +++ b/amoro-optimizer/amoro-optimizer-common/src/main/java/org/apache/amoro/optimizer/common/OptimizerExecutor.java @@ -46,6 +46,14 @@ public class OptimizerExecutor extends AbstractOptimizerOperator { private static final Logger LOG = LoggerFactory.getLogger(OptimizerExecutor.class); protected static final int ERROR_MESSAGE_MAX_LENGTH = 4000; + /** + * How long completeTask keeps retrying through transient AMS errors, measured from the moment + * shutdown is first observed by the retry loop. Sized to ride out a short network blip or AMS + * failover; note the force-interrupt issued at the stopOptimizing deadline aborts the window + * early, so retries never extend the overall shutdown beyond the configured timeout. + */ + static final long COMPLETE_TASK_DRAIN_TIMEOUT_MS = 30_000L; + private final int threadId; public OptimizerExecutor(OptimizerConfig config, int threadId) { @@ -250,12 +258,17 @@ protected OptimizingTaskResult executeTask(OptimizingTask task) { protected void completeTask(String amsUrl, OptimizingTaskResult optimizingTaskResult) { try { - callAuthenticatedAms( + // Report the result even when shutdown was requested before or during this call: a + // finished task is exactly what the graceful drain exists to preserve. The bounded + // drain window keeps retrying through transient errors instead of dropping the result + // on the first failure. + callAuthenticatedAmsWithDrain( amsUrl, (client, token) -> { client.completeTask(token, optimizingTaskResult); return null; - }); + }, + COMPLETE_TASK_DRAIN_TIMEOUT_MS); LOG.info( "Optimizer executor[{}] completed task[{}](status: {}) to AMS {}", threadId, diff --git a/amoro-optimizer/amoro-optimizer-common/src/main/java/org/apache/amoro/optimizer/common/OptimizerToucher.java b/amoro-optimizer/amoro-optimizer-common/src/main/java/org/apache/amoro/optimizer/common/OptimizerToucher.java index 9fd24e22ec..b6a75e1117 100644 --- a/amoro-optimizer/amoro-optimizer-common/src/main/java/org/apache/amoro/optimizer/common/OptimizerToucher.java +++ b/amoro-optimizer/amoro-optimizer-common/src/main/java/org/apache/amoro/optimizer/common/OptimizerToucher.java @@ -35,6 +35,9 @@ public class OptimizerToucher extends AbstractOptimizerOperator { private transient TokenChangeListener tokenChangeListener; private final Map registerProperties = Maps.newHashMap(); private final long startTime; + private transient volatile Thread runnerThread; + private volatile boolean draining = false; + private transient boolean drainTokenLossLogged = false; public OptimizerToucher(OptimizerConfig config) { super(config); @@ -54,21 +57,59 @@ public OptimizerToucher withRegisterProperty(String name, String value) { public void start() { LOG.info("Starting optimizer toucher with configuration:{}", getConfig()); - while (isStarted()) { - try { - if (checkToken()) { - touch(); + runnerThread = Thread.currentThread(); + try { + while (isStarted()) { + try { + if (checkToken()) { + touch(); + } + waitAShortTime(getConfig().getHeartBeat()); + } catch (Throwable t) { + LOG.error("Optimizer toucher got an unexpected error", t); } - waitAShortTime(getConfig().getHeartBeat()); - } catch (Throwable t) { - LOG.error("Optimizer toucher got an unexpected error", t); } + } finally { + runnerThread = null; } LOG.info("Optimizer toucher stopped"); } + /** + * Enters drain mode: keep heartbeating with the current token but never re-register. During + * graceful shutdown AMS may have already unregistered this optimizer (a scale-down releases the + * resource before the pod receives SIGTERM); re-registering on the resulting auth error would + * create a ghost optimizer AMS never asked for and rotate the executors' token, so their final + * completeTask would be rejected as coming from the wrong optimizer. + */ + public void enterDrainMode() { + this.draining = true; + } + + @Override + public void stop() { + super.stop(); + // Wake the runner immediately if it is sleeping in waitAShortTime, so the heartbeat + // loop terminates without waiting up to one full heartbeat interval. waitAShortTime + // preserves the interrupt flag so the loop exits cleanly on its next isStarted() check. + Thread t = runnerThread; + if (t != null) { + t.interrupt(); + } + } + private boolean checkToken() { if (!tokenIsReady()) { + if (draining) { + if (!drainTokenLossLogged) { + drainTokenLossLogged = true; + LOG.warn( + "Optimizer token became invalid while draining; skip re-registering to AMS {} " + + "to avoid creating a ghost optimizer", + getConfig().getAmsUrl()); + } + return false; + } LOG.info( "Registering optimizer to AMS {} (group: {}, mode: {}, threads: {}, memory: {}MB) ...", getConfig().getAmsUrl(), diff --git a/amoro-optimizer/amoro-optimizer-common/src/test/java/org/apache/amoro/optimizer/common/OptimizerTestBase.java b/amoro-optimizer/amoro-optimizer-common/src/test/java/org/apache/amoro/optimizer/common/OptimizerTestBase.java index 95b0cbc58c..d8b83db54e 100644 --- a/amoro-optimizer/amoro-optimizer-common/src/test/java/org/apache/amoro/optimizer/common/OptimizerTestBase.java +++ b/amoro-optimizer/amoro-optimizer-common/src/test/java/org/apache/amoro/optimizer/common/OptimizerTestBase.java @@ -43,5 +43,6 @@ public void clearTasks() { TEST_AMS.getOptimizerHandler().getPendingTasks().clear(); TEST_AMS.getOptimizerHandler().getExecutingTasks().clear(); TEST_AMS.getOptimizerHandler().getCompletedTasks().clear(); + TEST_AMS.getOptimizerHandler().failNextCompleteTasks(0); } } diff --git a/amoro-optimizer/amoro-optimizer-common/src/test/java/org/apache/amoro/optimizer/common/TestAbstractOptimizerOperator.java b/amoro-optimizer/amoro-optimizer-common/src/test/java/org/apache/amoro/optimizer/common/TestAbstractOptimizerOperator.java new file mode 100644 index 0000000000..326ac60048 --- /dev/null +++ b/amoro-optimizer/amoro-optimizer-common/src/test/java/org/apache/amoro/optimizer/common/TestAbstractOptimizerOperator.java @@ -0,0 +1,57 @@ +/* + * 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.optimizer.common; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +public class TestAbstractOptimizerOperator { + + @Test + public void testStrayInterruptWhileStartedIsSwallowed() { + AbstractOptimizerOperator operator = new AbstractOptimizerOperator(new OptimizerConfig()); + Thread.currentThread().interrupt(); + try { + operator.waitAShortTime(10); + // A stray interrupt while the operator is still running must not leave the flag set, + // otherwise every subsequent sleep returns instantly and the poll/retry loops busy-spin + // for the remaining lifetime of the process. + Assertions.assertFalse( + Thread.currentThread().isInterrupted(), + "Interrupt flag must be swallowed while the operator is started"); + } finally { + Thread.interrupted(); + } + } + + @Test + public void testInterruptDuringShutdownIsPreserved() { + AbstractOptimizerOperator operator = new AbstractOptimizerOperator(new OptimizerConfig()); + operator.stop(); + Thread.currentThread().interrupt(); + try { + operator.waitAShortTime(10); + Assertions.assertTrue( + Thread.currentThread().isInterrupted(), + "Interrupt flag must be preserved after stop so shutdown skips residual sleeps"); + } finally { + Thread.interrupted(); + } + } +} diff --git a/amoro-optimizer/amoro-optimizer-common/src/test/java/org/apache/amoro/optimizer/common/TestOptimizer.java b/amoro-optimizer/amoro-optimizer-common/src/test/java/org/apache/amoro/optimizer/common/TestOptimizer.java index 1feeef57fb..c7cf05bed8 100644 --- a/amoro-optimizer/amoro-optimizer-common/src/test/java/org/apache/amoro/optimizer/common/TestOptimizer.java +++ b/amoro-optimizer/amoro-optimizer-common/src/test/java/org/apache/amoro/optimizer/common/TestOptimizer.java @@ -23,7 +23,9 @@ import org.junit.jupiter.api.Test; import java.util.List; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; public class TestOptimizer extends OptimizerTestBase { @@ -48,4 +50,141 @@ public void testStartOptimizer() throws InterruptedException { Assertions.assertEquals(2, taskResults.size()); optimizer.stopOptimizing(); } + + @Test + public void testGracefulShutdown() throws InterruptedException { + OptimizerConfig optimizerConfig = + OptimizerTestHelpers.buildOptimizerConfig(TEST_AMS.getServerUrl()); + Optimizer optimizer = new Optimizer(optimizerConfig); + Thread optimizerThread = new Thread(optimizer::startOptimizing); + optimizerThread.start(); + TimeUnit.SECONDS.sleep(1); + + TEST_AMS + .getOptimizerHandler() + .offerTask(TestOptimizerExecutor.TestOptimizingInput.successInput(1).toTask(0, 0)); + TimeUnit.MILLISECONDS.sleep(OptimizerTestHelpers.CALL_AMS_INTERVAL * 5); + + optimizer.stopOptimizing(); + optimizerThread.join(5000); + + Assertions.assertFalse(optimizerThread.isAlive(), "Optimizer thread should have terminated"); + + String token = optimizer.getToucher().getToken(); + List taskResults = + TEST_AMS.getOptimizerHandler().getCompletedTasks().get(token); + Assertions.assertNotNull(taskResults, "Task results should be reported before shutdown"); + Assertions.assertEquals(1, taskResults.size()); + } + + @Test + public void testGracefulShutdownWaitsForInProgressTask() throws InterruptedException { + OptimizerConfig optimizerConfig = + OptimizerTestHelpers.buildOptimizerConfig(TEST_AMS.getServerUrl()); + Optimizer optimizer = new Optimizer(optimizerConfig); + Thread optimizerThread = new Thread(optimizer::startOptimizing); + optimizerThread.start(); + TimeUnit.SECONDS.sleep(1); + + long taskExecutionMs = 3_000; + CountDownLatch executionStarted = new CountDownLatch(1); + TestOptimizerExecutor.slowExecutionStartedLatch = executionStarted; + try { + TEST_AMS + .getOptimizerHandler() + .offerTask( + TestOptimizerExecutor.TestOptimizingInput.slowSuccessInput(1, taskExecutionMs) + .toTask(0, 0)); + + // Wait until the executor has polled, acked and entered execute(), so stop is called + // while the task is deterministically in progress. + Assertions.assertTrue( + executionStarted.await(10, TimeUnit.SECONDS), + "Task should have started executing before stop"); + } finally { + TestOptimizerExecutor.slowExecutionStartedLatch = null; + } + + long startStop = System.currentTimeMillis(); + optimizer.stopOptimizing(); + long elapsed = System.currentTimeMillis() - startStop; + optimizerThread.join(5_000); + + Assertions.assertFalse(optimizerThread.isAlive(), "Optimizer thread should have terminated"); + Assertions.assertTrue( + elapsed >= 1_000, + "stopOptimizing should block until in-progress task completes (elapsed=" + elapsed + "ms)"); + + String token = optimizer.getToucher().getToken(); + List taskResults = + TEST_AMS.getOptimizerHandler().getCompletedTasks().get(token); + Assertions.assertNotNull(taskResults, "Task results should be reported before shutdown"); + Assertions.assertEquals(1, taskResults.size()); + Assertions.assertNull( + taskResults.get(0).getErrorMessage(), + "In-progress task must complete successfully, not be interrupted"); + } + + @Test + public void testForceInterruptReachesAllExecutorsWhenStopperInterrupted() + throws InterruptedException { + OptimizerConfig optimizerConfig = + OptimizerTestHelpers.buildOptimizerConfig(TEST_AMS.getServerUrl()); + optimizerConfig.setExecutionParallel(2); + + CountDownLatch executorsRunning = new CountDownLatch(2); + AtomicBoolean[] executorInterrupted = {new AtomicBoolean(), new AtomicBoolean()}; + Thread[] executorThreads = new Thread[2]; + Optimizer optimizer = + new Optimizer( + optimizerConfig, + () -> new OptimizerToucher(optimizerConfig), + (i) -> + new OptimizerExecutor(optimizerConfig, i) { + @Override + public void start() { + executorThreads[getThreadId()] = Thread.currentThread(); + executorsRunning.countDown(); + try { + TimeUnit.HOURS.sleep(1); + } catch (InterruptedException e) { + executorInterrupted[getThreadId()].set(true); + } + } + }); + Thread optimizerThread = new Thread(optimizer::startOptimizing); + optimizerThread.start(); + try { + Assertions.assertTrue( + executorsRunning.await(5, TimeUnit.SECONDS), "Executor threads should have started"); + + Thread stopper = new Thread(optimizer::stopOptimizing, "stopper"); + stopper.start(); + // Let the stopper enter its join-with-deadline wait, then interrupt it the same way + // the Hadoop ShutdownHookManager watchdog cancels a hook that exceeded its timeout. + TimeUnit.MILLISECONDS.sleep(500); + stopper.interrupt(); + stopper.join(10_000); + Assertions.assertFalse( + stopper.isAlive(), "stopOptimizing should return promptly after being interrupted"); + + long deadline = System.currentTimeMillis() + 5_000; + while (System.currentTimeMillis() < deadline + && !(executorInterrupted[0].get() && executorInterrupted[1].get())) { + TimeUnit.MILLISECONDS.sleep(50); + } + Assertions.assertTrue( + executorInterrupted[0].get(), "First executor thread must be force-interrupted"); + Assertions.assertTrue( + executorInterrupted[1].get(), + "All executor threads must be force-interrupted, not only the first"); + optimizerThread.join(5_000); + } finally { + for (Thread t : executorThreads) { + if (t != null) { + t.interrupt(); + } + } + } + } } diff --git a/amoro-optimizer/amoro-optimizer-common/src/test/java/org/apache/amoro/optimizer/common/TestOptimizerExecutor.java b/amoro-optimizer/amoro-optimizer-common/src/test/java/org/apache/amoro/optimizer/common/TestOptimizerExecutor.java index 08f5a370c9..866ad655ab 100644 --- a/amoro-optimizer/amoro-optimizer-common/src/test/java/org/apache/amoro/optimizer/common/TestOptimizerExecutor.java +++ b/amoro-optimizer/amoro-optimizer-common/src/test/java/org/apache/amoro/optimizer/common/TestOptimizerExecutor.java @@ -36,12 +36,20 @@ import org.junit.jupiter.api.Test; import java.util.Map; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; public class TestOptimizerExecutor extends OptimizerTestBase { private static final String FAILED_TASK_MESSAGE = "Execute Task failed"; + /** + * When set, slow test tasks count this down as soon as execute() begins, letting tests wait + * deterministically for "the task is now in progress" instead of sleeping fixed intervals. + */ + static volatile CountDownLatch slowExecutionStartedLatch; + private OptimizerExecutor optimizerExecutor; @BeforeEach @@ -85,6 +93,71 @@ public void testExecuteTaskSuccess() throws InterruptedException, TException { Assertions.assertEquals(1, output.inputId()); } + @Test + public void testCompleteTaskRetriesTransientErrorDuringShutdownDrain() throws TException { + TEST_AMS.getOptimizerHandler().authenticate(new OptimizerRegisterInfo()); + String token = + TEST_AMS.getOptimizerHandler().getRegisteredOptimizers().keySet().iterator().next(); + optimizerExecutor.setToken(token); + TEST_AMS.getOptimizerHandler().ackTask(token, 0, new OptimizingTaskId(0, 0)); + + // Shutdown already requested when the in-flight task finishes and reports its result. + optimizerExecutor.stop(); + TEST_AMS.getOptimizerHandler().failNextCompleteTasks(2); + + OptimizingTaskResult result = new OptimizingTaskResult(new OptimizingTaskId(0, 0), 0); + optimizerExecutor.completeTask(TEST_AMS.getServerUrl(), result); + + Assertions.assertNotNull( + TEST_AMS.getOptimizerHandler().getCompletedTasks().get(token), + "Task result must be reported despite transient AMS errors during the shutdown drain"); + Assertions.assertEquals( + 1, TEST_AMS.getOptimizerHandler().getCompletedTasks().get(token).size()); + } + + @Test + public void testDrainWindowAnchorsAtStopTime() throws InterruptedException, TException { + TEST_AMS.getOptimizerHandler().authenticate(new OptimizerRegisterInfo()); + String token = + TEST_AMS.getOptimizerHandler().getRegisteredOptimizers().keySet().iterator().next(); + optimizerExecutor.setToken(token); + TEST_AMS.getOptimizerHandler().ackTask(token, 0, new OptimizingTaskId(0, 0)); + + // AMS keeps failing while the optimizer is still running, for longer than the drain budget. + TEST_AMS.getOptimizerHandler().failNextCompleteTasks(Integer.MAX_VALUE); + OptimizingTaskResult result = new OptimizingTaskResult(new OptimizingTaskId(0, 0), 0); + AtomicReference error = new AtomicReference<>(); + Thread caller = + new Thread( + () -> { + try { + optimizerExecutor.callAuthenticatedAmsWithDrain( + TEST_AMS.getServerUrl(), + (client, callToken) -> { + client.completeTask(callToken, result); + return null; + }, + 2_000); + } catch (Exception e) { + error.set(e); + } + }); + caller.start(); + TimeUnit.MILLISECONDS.sleep(2_500); + + // Shutdown begins and AMS recovers: the drain window must start counting from now, + // not from call entry (which is already past the 2s budget). + optimizerExecutor.stop(); + TEST_AMS.getOptimizerHandler().failNextCompleteTasks(0); + caller.join(5_000); + + Assertions.assertFalse(caller.isAlive(), "Drain call should have finished"); + Assertions.assertNull( + error.get(), "Result must be reported within a full drain window anchored at stop time"); + Assertions.assertEquals( + 1, TEST_AMS.getOptimizerHandler().getCompletedTasks().get(token).size()); + } + @Test public void testExecuteTaskFailed() throws InterruptedException, TException { TEST_AMS.getOptimizerHandler().authenticate(new OptimizerRegisterInfo()); @@ -108,10 +181,16 @@ public void testExecuteTaskFailed() throws InterruptedException, TException { public static class TestOptimizingInput extends BaseOptimizingInput { private final int inputId; private final boolean executeSuccess; + private final long executionTimeMs; private TestOptimizingInput(int inputId, boolean executeSuccess) { + this(inputId, executeSuccess, 0L); + } + + private TestOptimizingInput(int inputId, boolean executeSuccess, long executionTimeMs) { this.inputId = inputId; this.executeSuccess = executeSuccess; + this.executionTimeMs = executionTimeMs; } public static TestOptimizingInput successInput(int inputId) { @@ -122,10 +201,18 @@ public static TestOptimizingInput failedInput(int inputId) { return new TestOptimizingInput(inputId, false); } + public static TestOptimizingInput slowSuccessInput(int inputId, long executionTimeMs) { + return new TestOptimizingInput(inputId, true, executionTimeMs); + } + private int inputId() { return inputId; } + private long executionTimeMs() { + return executionTimeMs; + } + public OptimizingTask toTask(long processId, int taskId) { OptimizingTask optimizingTask = new OptimizingTask(new OptimizingTaskId(processId, taskId)); optimizingTask.setTaskInput(SerializationUtil.simpleSerialize(this)); @@ -159,6 +246,18 @@ private TestOptimizingExecutor(TestOptimizingInput input) { @Override public TestOptimizingOutput execute() { + if (input.executionTimeMs() > 0) { + CountDownLatch latch = slowExecutionStartedLatch; + if (latch != null) { + latch.countDown(); + } + try { + Thread.sleep(input.executionTimeMs()); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IllegalStateException("Task was interrupted before completion", e); + } + } if (input.executeSuccess) { return new TestOptimizingOutput(input.inputId()); } else { diff --git a/amoro-optimizer/amoro-optimizer-common/src/test/java/org/apache/amoro/optimizer/common/TestOptimizerToucher.java b/amoro-optimizer/amoro-optimizer-common/src/test/java/org/apache/amoro/optimizer/common/TestOptimizerToucher.java index 749dec4990..76a38714e3 100644 --- a/amoro-optimizer/amoro-optimizer-common/src/test/java/org/apache/amoro/optimizer/common/TestOptimizerToucher.java +++ b/amoro-optimizer/amoro-optimizer-common/src/test/java/org/apache/amoro/optimizer/common/TestOptimizerToucher.java @@ -62,6 +62,38 @@ public void testRegisterOptimizer() throws InterruptedException { optimizerToucher.stop(); } + @Test + public void testNoReRegistrationInDrainMode() throws InterruptedException { + OptimizerConfig optimizerConfig = + OptimizerTestHelpers.buildOptimizerConfig(TEST_AMS.getServerUrl()); + OptimizerToucher optimizerToucher = new OptimizerToucher(optimizerConfig); + TestTokenChangeListener tokenChangeListener = new TestTokenChangeListener(); + optimizerToucher.withTokenChangeListener(tokenChangeListener); + new Thread(optimizerToucher::start).start(); + try { + tokenChangeListener.waitForTokenChange(); + Assertions.assertEquals(1, TEST_AMS.getOptimizerHandler().getRegisteredOptimizers().size()); + + // Drain begins (graceful shutdown keeps the toucher alive for heartbeats), then AMS + // unregisters this optimizer the way a scale-down does before deleting the pod. + optimizerToucher.enterDrainMode(); + TEST_AMS.getOptimizerHandler().getRegisteredOptimizers().clear(); + + // Give the heartbeat loop several cycles to hit the auth error. + TimeUnit.MILLISECONDS.sleep(3_000); + + Assertions.assertTrue( + TEST_AMS.getOptimizerHandler().getRegisteredOptimizers().isEmpty(), + "Toucher must not re-register a ghost optimizer while draining"); + Assertions.assertEquals( + 1, + tokenChangeListener.tokenList().size(), + "Executor tokens must not be rotated during drain"); + } finally { + optimizerToucher.stop(); + } + } + private void validateRegisteredOptimizer( String token, OptimizerConfig registerConfig, Map optimizerProperties) { Map registeredOptimizerMap = diff --git a/amoro-optimizer/amoro-optimizer-flink/src/main/java/org/apache/amoro/optimizer/flink/FlinkExecutor.java b/amoro-optimizer/amoro-optimizer-flink/src/main/java/org/apache/amoro/optimizer/flink/FlinkExecutor.java index ddf3b1842c..90733ab256 100644 --- a/amoro-optimizer/amoro-optimizer-flink/src/main/java/org/apache/amoro/optimizer/flink/FlinkExecutor.java +++ b/amoro-optimizer/amoro-optimizer-flink/src/main/java/org/apache/amoro/optimizer/flink/FlinkExecutor.java @@ -19,8 +19,10 @@ package org.apache.amoro.optimizer.flink; import static org.apache.flink.configuration.HighAvailabilityOptions.HA_CLUSTER_ID; +import static org.apache.flink.configuration.TaskManagerOptions.TASK_CANCELLATION_TIMEOUT; import static org.apache.flink.configuration.TaskManagerOptions.TASK_MANAGER_RESOURCE_ID; +import org.apache.amoro.optimizer.common.OptimizerConfig; import org.apache.amoro.optimizer.common.OptimizerExecutor; import org.apache.flink.streaming.api.operators.AbstractStreamOperator; import org.apache.flink.streaming.api.operators.OneInputStreamOperator; @@ -29,14 +31,26 @@ public class FlinkExecutor extends AbstractStreamOperator implements OneInputStreamOperator { + /** + * Kept below Flink's cancellation watchdog deadline (task.cancellation.timeout) so the + * force-interrupt and thread exit complete before Flink fails the whole TaskManager. + */ + static final long CANCELLATION_SAFETY_MARGIN_MS = 10_000L; + private final OptimizerExecutor[] allExecutors; + private final long shutdownTimeoutMs; + private final long heartbeatIntervalMs; private FlinkOptimizerExecutor executor; private String optimizeGroupName; private Thread optimizerThread; + private long drainTimeoutMs; - public FlinkExecutor(OptimizerExecutor[] allExecutors, String optimizeGroupName) { + public FlinkExecutor( + OptimizerExecutor[] allExecutors, String optimizeGroupName, OptimizerConfig config) { this.allExecutors = allExecutors; this.optimizeGroupName = optimizeGroupName; + this.shutdownTimeoutMs = config.getShutdownTimeoutMs(); + this.heartbeatIntervalMs = config.getHeartBeat(); } @Override @@ -62,6 +76,12 @@ public void open() throws Exception { // add label optimize_group; getMetricGroup().getAllVariables().put("", optimizeGroupName); executor.initOperatorMetric(getMetricGroup()); + long taskCancellationTimeoutMs = + getRuntimeContext() + .getTaskManagerRuntimeInfo() + .getConfiguration() + .get(TASK_CANCELLATION_TIMEOUT); + drainTimeoutMs = effectiveDrainTimeoutMs(shutdownTimeoutMs, taskCancellationTimeoutMs); optimizerThread = new Thread(() -> executor.start(), "flink-optimizer-executor-" + subTaskIndex); optimizerThread.setDaemon(true); @@ -73,14 +93,80 @@ public void close() throws Exception { if (executor != null) { executor.stop(); } - if (optimizerThread != null && optimizerThread.isAlive()) { + // The FlinkToucher source is cancelled before this operator, so heartbeats have already + // stopped when the drain begins. Keep the registration alive by touching AMS with the + // existing token from here — otherwise AMS expires the optimizer after its heartbeat + // timeout (default 1 min) and resets the in-flight task, dropping the drained result. + Runnable drainHeartbeat = executor != null ? executor::bestEffortTouch : null; + drainThenForceStop(optimizerThread, drainTimeoutMs, heartbeatIntervalMs, drainHeartbeat); + } + + /** + * Bounds the graceful drain by Flink's cancellation watchdog: when a cancelled task has not + * terminated within task.cancellation.timeout (default 180s), Flink kills the entire TaskManager, + * so waiting any longer than that would turn a graceful drain into a TM failure. A value of 0 + * disables the watchdog, in which case the full shutdown timeout applies. + */ + static long effectiveDrainTimeoutMs(long shutdownTimeoutMs, long taskCancellationTimeoutMs) { + if (taskCancellationTimeoutMs <= 0) { + return shutdownTimeoutMs; + } + long cap = taskCancellationTimeoutMs - CANCELLATION_SAFETY_MARGIN_MS; + return Math.max(0, Math.min(shutdownTimeoutMs, cap)); + } + + static void drainThenForceStop( + Thread optimizerThread, + long drainTimeoutMs, + long heartbeatIntervalMs, + Runnable drainHeartbeat) { + if (optimizerThread == null || !optimizerThread.isAlive()) { + return; + } + // Flink's cancellation machinery (TaskCanceler/TaskInterrupter) interrupts the task thread + // running close() to unblock I/O; absorb those interrupts and keep waiting until the drain + // deadline — the budget is capped below the cancellation watchdog, so close() always + // returns before Flink escalates to failing the TaskManager. + boolean selfInterrupted = + joinUntil(optimizerThread, drainTimeoutMs, heartbeatIntervalMs, drainHeartbeat); + if (optimizerThread.isAlive()) { + LOG.warn( + "Optimizer executor thread {} did not finish in-progress work within {}ms, " + + "force-interrupting", + optimizerThread.getName(), + drainTimeoutMs); optimizerThread.interrupt(); + selfInterrupted |= joinUntil(optimizerThread, 1_000, 1_000, null); + } + if (selfInterrupted) { + Thread.currentThread().interrupt(); + } + } + + /** + * Joins the thread until the deadline, absorbing interrupts of the calling thread and running the + * heartbeat action between join slices. Returns whether an interrupt was absorbed so the caller + * can restore the flag afterwards. + */ + private static boolean joinUntil( + Thread thread, long timeoutMs, long heartbeatIntervalMs, Runnable heartbeat) { + boolean interrupted = false; + long deadline = System.currentTimeMillis() + timeoutMs; + while (thread.isAlive()) { + long remaining = deadline - System.currentTimeMillis(); + if (remaining <= 0) { + break; + } try { - optimizerThread.join(5000); + thread.join(Math.min(remaining, heartbeatIntervalMs)); } catch (InterruptedException e) { - Thread.currentThread().interrupt(); + interrupted = true; + } + if (heartbeat != null && thread.isAlive()) { + heartbeat.run(); } } + return interrupted; } @Override diff --git a/amoro-optimizer/amoro-optimizer-flink/src/main/java/org/apache/amoro/optimizer/flink/FlinkOptimizer.java b/amoro-optimizer/amoro-optimizer-flink/src/main/java/org/apache/amoro/optimizer/flink/FlinkOptimizer.java index 6b46e7dda8..4b6c2bab8a 100644 --- a/amoro-optimizer/amoro-optimizer-flink/src/main/java/org/apache/amoro/optimizer/flink/FlinkOptimizer.java +++ b/amoro-optimizer/amoro-optimizer-flink/src/main/java/org/apache/amoro/optimizer/flink/FlinkOptimizer.java @@ -57,7 +57,8 @@ public static void main(String[] args) throws CmdLineException { .transform( FlinkExecutor.class.getName(), Types.VOID, - new FlinkExecutor(optimizer.getExecutors(), optimizerConfig.getGroupName())) + new FlinkExecutor( + optimizer.getExecutors(), optimizerConfig.getGroupName(), optimizerConfig)) .setParallelism(optimizerConfig.getExecutionParallel()) .addSink(new DiscardingSink<>()) .name("Optimizer empty sink") diff --git a/amoro-optimizer/amoro-optimizer-flink/src/main/java/org/apache/amoro/optimizer/flink/FlinkOptimizerExecutor.java b/amoro-optimizer/amoro-optimizer-flink/src/main/java/org/apache/amoro/optimizer/flink/FlinkOptimizerExecutor.java index cb2fff27e0..88ecddd795 100644 --- a/amoro-optimizer/amoro-optimizer-flink/src/main/java/org/apache/amoro/optimizer/flink/FlinkOptimizerExecutor.java +++ b/amoro-optimizer/amoro-optimizer-flink/src/main/java/org/apache/amoro/optimizer/flink/FlinkOptimizerExecutor.java @@ -20,6 +20,7 @@ import org.apache.amoro.api.OptimizingTask; import org.apache.amoro.api.OptimizingTaskResult; +import org.apache.amoro.client.OptimizingClientPools; import org.apache.amoro.optimizer.common.OptimizerConfig; import org.apache.amoro.optimizer.common.OptimizerExecutor; import org.apache.amoro.shade.guava32.com.google.common.base.Strings; @@ -53,6 +54,26 @@ public void addRuntimeContext(String key, String value) { runtimeContext.put(key, value); } + /** + * Best-effort heartbeat with the current token, usable after stop(). Flink cancels the + * FlinkToucher source before this operator, so the drain in {@link FlinkExecutor#close()} keeps + * the registration alive by touching AMS directly — otherwise AMS expires the optimizer after its + * heartbeat timeout and resets the in-flight task, dropping the drained result. Failures are + * swallowed; the next drain heartbeat simply retries. This never re-registers, consistent with + * the drain-mode toucher. + */ + void bestEffortTouch() { + String currentToken = getToken(); + if (currentToken == null) { + return; + } + try { + OptimizingClientPools.getClient(getConfig().getAmsUrl()).touch(currentToken); + } catch (Exception e) { + LOG.debug("Best-effort touch to AMS failed, will retry on the next drain heartbeat", e); + } + } + public void initOperatorMetric(MetricGroup metricGroup) { this.operatorMetricGroup = metricGroup; taskCounter = this.operatorMetricGroup.addGroup("amoro").addGroup("optimizer").counter("tasks"); diff --git a/amoro-optimizer/amoro-optimizer-flink/src/test/java/org/apache/amoro/optimizer/flink/TestFlinkExecutor.java b/amoro-optimizer/amoro-optimizer-flink/src/test/java/org/apache/amoro/optimizer/flink/TestFlinkExecutor.java new file mode 100644 index 0000000000..82fb8c64ec --- /dev/null +++ b/amoro-optimizer/amoro-optimizer-flink/src/test/java/org/apache/amoro/optimizer/flink/TestFlinkExecutor.java @@ -0,0 +1,178 @@ +/* + * 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.optimizer.flink; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; + +public class TestFlinkExecutor { + + @Test + public void testDrainWaitsForInProgressTask() throws InterruptedException { + AtomicBoolean taskCompleted = new AtomicBoolean(false); + AtomicBoolean interrupted = new AtomicBoolean(false); + Thread worker = + new Thread( + () -> { + try { + TimeUnit.MILLISECONDS.sleep(500); + taskCompleted.set(true); + } catch (InterruptedException e) { + interrupted.set(true); + } + }); + worker.start(); + + FlinkExecutor.drainThenForceStop(worker, 10_000, 10_000, null); + + Assertions.assertFalse(worker.isAlive(), "Worker thread should have terminated"); + Assertions.assertTrue(taskCompleted.get(), "In-progress task should complete within budget"); + Assertions.assertFalse(interrupted.get(), "Task within budget must not be interrupted"); + } + + @Test + public void testDrainKeepsHeartbeatingWhileWaiting() throws InterruptedException { + AtomicInteger heartbeats = new AtomicInteger(); + Thread worker = + new Thread( + () -> { + try { + TimeUnit.MILLISECONDS.sleep(900); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + }); + worker.start(); + + FlinkExecutor.drainThenForceStop(worker, 10_000, 200, heartbeats::incrementAndGet); + + Assertions.assertFalse(worker.isAlive(), "Worker thread should have terminated"); + Assertions.assertTrue( + heartbeats.get() >= 2, + "Drain must keep sending heartbeats periodically, got " + heartbeats.get()); + } + + @Test + public void testDrainForceInterruptsAfterTimeout() throws InterruptedException { + AtomicBoolean interrupted = new AtomicBoolean(false); + Thread worker = + new Thread( + () -> { + try { + TimeUnit.MINUTES.sleep(10); + } catch (InterruptedException e) { + interrupted.set(true); + } + }); + worker.start(); + + long start = System.currentTimeMillis(); + FlinkExecutor.drainThenForceStop(worker, 300, 10_000, null); + long elapsed = System.currentTimeMillis() - start; + + Assertions.assertFalse(worker.isAlive(), "Worker thread should have been force-stopped"); + Assertions.assertTrue(interrupted.get(), "Task exceeding budget should be interrupted"); + Assertions.assertTrue( + elapsed < 5_000, "Drain should return promptly after budget, elapsed=" + elapsed + "ms"); + } + + @Test + public void testDrainSurvivesCallerInterrupts() throws InterruptedException { + AtomicBoolean taskCompleted = new AtomicBoolean(false); + AtomicBoolean workerInterrupted = new AtomicBoolean(false); + Thread worker = + new Thread( + () -> { + try { + TimeUnit.MILLISECONDS.sleep(800); + taskCompleted.set(true); + } catch (InterruptedException e) { + workerInterrupted.set(true); + } + }); + worker.start(); + + // Simulate Flink's TaskCanceler/TaskInterrupter interrupting the task thread that is + // running close() during job cancellation: the drain must absorb it and keep waiting. + Thread drainThread = Thread.currentThread(); + Thread canceler = + new Thread( + () -> { + try { + TimeUnit.MILLISECONDS.sleep(200); + } catch (InterruptedException e) { + return; + } + drainThread.interrupt(); + }); + canceler.start(); + try { + FlinkExecutor.drainThenForceStop(worker, 10_000, 10_000, null); + + Assertions.assertTrue( + Thread.interrupted(), "Self-interrupt must be restored after the drain completes"); + Assertions.assertFalse(worker.isAlive(), "Worker thread should have terminated"); + Assertions.assertTrue( + taskCompleted.get(), "Drain must keep waiting across caller interrupts"); + Assertions.assertFalse( + workerInterrupted.get(), "Task within budget must not be force-interrupted"); + } finally { + Thread.interrupted(); + canceler.interrupt(); + } + } + + @Test + public void testDrainToleratesDeadOrNullThread() { + FlinkExecutor.drainThenForceStop(null, 1_000, 1_000, null); + + Thread finished = new Thread(() -> {}); + finished.start(); + FlinkExecutor.drainThenForceStop(finished, 1_000, 1_000, null); + } + + @Test + public void testEffectiveDrainTimeoutCappedByCancellationTimeout() { + // Flink kills the TaskManager when a cancelled task does not terminate within + // task.cancellation.timeout (default 180s), so the drain budget must stay below it. + Assertions.assertEquals( + 180_000 - FlinkExecutor.CANCELLATION_SAFETY_MARGIN_MS, + FlinkExecutor.effectiveDrainTimeoutMs(600_000, 180_000)); + } + + @Test + public void testEffectiveDrainTimeoutUsesShutdownTimeoutWhenSmaller() { + Assertions.assertEquals(60_000, FlinkExecutor.effectiveDrainTimeoutMs(60_000, 180_000)); + } + + @Test + public void testEffectiveDrainTimeoutUnlimitedCancellationWatchdog() { + // task.cancellation.timeout = 0 disables Flink's cancellation watchdog entirely. + Assertions.assertEquals(600_000, FlinkExecutor.effectiveDrainTimeoutMs(600_000, 0)); + } + + @Test + public void testEffectiveDrainTimeoutNeverNegative() { + Assertions.assertEquals(0, FlinkExecutor.effectiveDrainTimeoutMs(600_000, 1_000)); + } +} diff --git a/amoro-optimizer/amoro-optimizer-flink/src/test/java/org/apache/amoro/optimizer/flink/TestFlinkOptimizerExecutor.java b/amoro-optimizer/amoro-optimizer-flink/src/test/java/org/apache/amoro/optimizer/flink/TestFlinkOptimizerExecutor.java new file mode 100644 index 0000000000..5f85ac41bb --- /dev/null +++ b/amoro-optimizer/amoro-optimizer-flink/src/test/java/org/apache/amoro/optimizer/flink/TestFlinkOptimizerExecutor.java @@ -0,0 +1,66 @@ +/* + * 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.optimizer.flink; + +import org.apache.amoro.TestAms; +import org.apache.amoro.api.OptimizerRegisterInfo; +import org.apache.amoro.optimizer.common.OptimizerConfig; +import org.apache.amoro.shade.thrift.org.apache.thrift.TException; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import org.kohsuke.args4j.CmdLineException; + +public class TestFlinkOptimizerExecutor { + + private static final TestAms TEST_AMS = new TestAms(); + + @BeforeAll + public static void startTestAms() throws Exception { + TEST_AMS.before(); + } + + @AfterAll + public static void stopTestAms() { + TEST_AMS.after(); + } + + @Test + public void testBestEffortTouchWorksWhileStoppedAndSwallowsErrors() + throws CmdLineException, TException { + OptimizerConfig config = + new OptimizerConfig(new String[] {"-a", TEST_AMS.getServerUrl(), "-p", "1", "-g", "g1"}); + FlinkOptimizerExecutor executor = new FlinkOptimizerExecutor(config, 0); + TEST_AMS.getOptimizerHandler().authenticate(new OptimizerRegisterInfo()); + String token = + TEST_AMS.getOptimizerHandler().getRegisteredOptimizers().keySet().iterator().next(); + executor.setToken(token); + executor.stop(); + + // Must work after stop: the drain in FlinkExecutor.close() uses it to keep the + // registration alive while the in-flight task finishes. + executor.bestEffortTouch(); + + // Must swallow failures: unknown token (AMS already expired us) and missing token. + TEST_AMS.getOptimizerHandler().getRegisteredOptimizers().clear(); + executor.bestEffortTouch(); + executor.setToken(null); + executor.bestEffortTouch(); + } +} diff --git a/amoro-optimizer/amoro-optimizer-spark/src/main/java/org/apache/amoro/optimizer/spark/SparkOptimizer.java b/amoro-optimizer/amoro-optimizer-spark/src/main/java/org/apache/amoro/optimizer/spark/SparkOptimizer.java index 2eb38c0283..0e67830bf3 100644 --- a/amoro-optimizer/amoro-optimizer-spark/src/main/java/org/apache/amoro/optimizer/spark/SparkOptimizer.java +++ b/amoro-optimizer/amoro-optimizer-spark/src/main/java/org/apache/amoro/optimizer/spark/SparkOptimizer.java @@ -22,12 +22,15 @@ import org.apache.amoro.optimizer.common.OptimizerConfig; import org.apache.amoro.optimizer.common.OptimizerToucher; import org.apache.amoro.resource.Resource; +import org.apache.hadoop.util.ShutdownHookManager; import org.apache.spark.api.java.JavaSparkContext; import org.apache.spark.sql.SparkSession; import org.apache.spark.util.Utils; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.util.concurrent.TimeUnit; + /** The {@code SparkOptimizer} acts as an entrypoint of the spark program */ public class SparkOptimizer extends Optimizer { private static final Logger LOG = LoggerFactory.getLogger(SparkOptimizer.class); @@ -65,6 +68,29 @@ public static void main(String[] args) throws Exception { OptimizerToucher toucher = optimizer.getToucher(); toucher.withRegisterProperty(Resource.PROPERTY_JOB_ID, spark.sparkContext().applicationId()); + // Register with Hadoop's ShutdownHookManager so that our hook runs to completion + // before downstream cleanup. ShutdownHookManager runs hooks in priority descending + // order, sequentially. Two cleanups must come AFTER our graceful shutdown: + // - Hadoop FileSystem cache close (priority FS_CACHE = 10) — would otherwise + // close in-flight HDFS writers and cause ClosedChannelException on flush. + // - SparkContext.stop (Spark's SPARK_CONTEXT_SHUTDOWN_PRIORITY = 50) — would + // otherwise tear down executors mid-task, failing in-flight RDD actions. + // Use SPARK_CONTEXT_SHUTDOWN_PRIORITY + 10 to guarantee both ordering constraints + // in a single value (60 > 50 > 10). + // Pass an explicit timeout — the 2-arg overload uses hadoop.service.shutdown.timeout + // (default 30s), which would cancel our hook well before stopOptimizing's + // shutdownTimeoutMs deadline. + int shutdownPriority = + org.apache.spark.util.ShutdownHookManager.SPARK_CONTEXT_SHUTDOWN_PRIORITY() + 10; + ShutdownHookManager.get() + .addShutdownHook( + () -> { + LOG.info("Received shutdown signal, initiating graceful shutdown..."); + optimizer.stopOptimizing(); + }, + shutdownPriority, + config.getShutdownTimeoutMs() + 10_000L, + TimeUnit.MILLISECONDS); LOG.info("Starting the spark optimizer with configuration:{}", config); optimizer.startOptimizing(); } diff --git a/amoro-optimizer/amoro-optimizer-standalone/src/main/java/org/apache/amoro/optimizer/standalone/StandaloneOptimizer.java b/amoro-optimizer/amoro-optimizer-standalone/src/main/java/org/apache/amoro/optimizer/standalone/StandaloneOptimizer.java index b38d0fc42b..cc7723a80a 100644 --- a/amoro-optimizer/amoro-optimizer-standalone/src/main/java/org/apache/amoro/optimizer/standalone/StandaloneOptimizer.java +++ b/amoro-optimizer/amoro-optimizer-standalone/src/main/java/org/apache/amoro/optimizer/standalone/StandaloneOptimizer.java @@ -21,13 +21,20 @@ import org.apache.amoro.optimizer.common.Optimizer; import org.apache.amoro.optimizer.common.OptimizerConfig; import org.apache.amoro.resource.Resource; +import org.apache.hadoop.fs.FileSystem; +import org.apache.hadoop.util.ShutdownHookManager; import org.kohsuke.args4j.CmdLineException; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import java.lang.management.ManagementFactory; import java.lang.management.RuntimeMXBean; +import java.util.concurrent.TimeUnit; public class StandaloneOptimizer { + private static final Logger LOG = LoggerFactory.getLogger(StandaloneOptimizer.class); + public static void main(String[] args) throws CmdLineException { OptimizerConfig optimizerConfig = new OptimizerConfig(args); Optimizer optimizer = new Optimizer(optimizerConfig); @@ -39,6 +46,23 @@ public static void main(String[] args) throws CmdLineException { RuntimeMXBean runtimeMXBean = ManagementFactory.getRuntimeMXBean(); String processId = runtimeMXBean.getName().split("@")[0]; optimizer.getToucher().withRegisterProperty(Resource.PROPERTY_JOB_ID, processId); + + // Register with Hadoop's ShutdownHookManager (priority > FS_CACHE) so that our hook + // runs to completion before Hadoop closes cached FileSystems. JVM Runtime shutdown + // hooks fire concurrently and lose this race, causing in-flight writers to hit + // ClosedChannelException during row-group flush. + // Pass an explicit timeout — the 2-arg overload uses hadoop.service.shutdown.timeout + // (default 30s), which would cancel our hook well before stopOptimizing's + // shutdownTimeoutMs deadline. + ShutdownHookManager.get() + .addShutdownHook( + () -> { + LOG.info("Received shutdown signal, initiating graceful shutdown..."); + optimizer.stopOptimizing(); + }, + FileSystem.SHUTDOWN_HOOK_PRIORITY + 10, + optimizerConfig.getShutdownTimeoutMs() + 10_000L, + TimeUnit.MILLISECONDS); optimizer.startOptimizing(); } } diff --git a/dist/src/main/amoro-bin/bin/optimizer.sh b/dist/src/main/amoro-bin/bin/optimizer.sh index 8e01c619c6..30591bd907 100755 --- a/dist/src/main/amoro-bin/bin/optimizer.sh +++ b/dist/src/main/amoro-bin/bin/optimizer.sh @@ -96,7 +96,7 @@ start() { } start-foreground() { - $CMDS + exec $CMDS } #0:pid bad and proc OK; 1:pid ok and proc bad; 2:pid bad diff --git a/docs/admin-guides/managing-optimizers.md b/docs/admin-guides/managing-optimizers.md index fd708e8c2d..35a0d935fc 100644 --- a/docs/admin-guides/managing-optimizers.md +++ b/docs/admin-guides/managing-optimizers.md @@ -289,6 +289,7 @@ The optimizer group supports the following properties: | cache-max-entry-size | All | No | 64mb | Max entry size in optimizer cache. | | cache-timeout | All | No | 10min | Timeout in optimizer cache. | | min-parallelism | All | No | 0 | The minimum total parallelism (CPU cores) that the optimizer group should maintain. When the total cores of running optimizers fall below this value, `OptimizerGroupKeeper` will automatically scale out new optimizers. Set to `0` to disable auto-scaling. Note: The behavior of the auto-scaling mechanism is controlled by the AMS-level configurations `optimizer-group.min-parallelism-check-interval` and `optimizer-group.max-keeping-attempts`.| +| shutdown-timeout-ms | All | No | 600000(10min) | Graceful shutdown timeout in milliseconds. On shutdown the optimizer waits up to this long for in-progress tasks to complete before force-interrupting them. For Kubernetes optimizers, the pod's `terminationGracePeriodSeconds` is derived from this value plus a 30s buffer. For Flink optimizers, the effective wait is additionally capped below `task.cancellation.timeout`. | | memory | Local | Yes | N/A | The max memory of JVM for local optimizer, in MBs. | | flink-conf.\ | Flink | No | N/A | Any flink config options could be overwritten, priority is optimizing-group > optimizing-container > flink-conf.yaml. | | spark-conf.\ | Spark | No | N/A | Any spark config options could be overwritten, priority is optimizing-group > optimizing-container > spark-defaults.conf. | @@ -298,6 +299,7 @@ To better utilize the resources of Flink Optimizer, it is recommended to add the * Set `flink-conf.taskmanager.memory.managed.size` to `32mb` as Flink optimizer does not have any computation logic, it does not need to occupy managed memory. * Set `flink-conf.taskmanager.memory.network.max` to `32mb` as there is no need for communication between operators in Flink Optimizer. * Set `flink-conf.taskmanager.memory.network.min` to `32mb` as there is no need for communication between operators in Flink Optimizer. +* When using `shutdown-timeout-ms` with a Flink Optimizer, also raise `flink-conf.task.cancellation.timeout` (default 180000) accordingly — the graceful drain on job cancellation is capped below Flink's cancellation watchdog, which otherwise fails the whole TaskManager. {{< /hint >}} ### Edit optimizer group @@ -360,6 +362,7 @@ The description of the relevant parameters is shown in the following table: | -cmts | No | Max total size in optimier cache, default 128MB. | | -cmes | No | Max entry size in optimizer cache, default 64MB. | | -ct | No | Timeout in optimizer cache, default 10Min. | +| -st | No | Graceful shutdown timeout(ms) — on job cancellation, wait this long for in-progress tasks to complete before force-interrupting them, default 600000(ms). The effective wait is capped below task.cancellation.timeout (default 180000). | Or you can submit optimizer in your own Spark task development platform or local Spark environment with the following configuration. The main parameters include: @@ -391,3 +394,4 @@ The description of the relevant parameters is shown in the following table: | -cmts | No | Max total size in optimier cache, default 128MB. | | -cmes | No | Max entry size in optimizer cache, default 64MB. | | -ct | No | Timeout in optimizer cache, default 10Min. | +| -st | No | Graceful shutdown timeout(ms) — on shutdown, wait this long for in-progress tasks to complete before force-interrupting them, default 600000(ms). |