[kubernetes] Graceful TaskManager termination + fix wrong-signal TODO in ResourceManager

docs/content/docs/deployment/resource-providers/native_kubernetes.md

@@ -468,6 +468,32 @@ Configure the value of <a href="{{< ref "docs/deployment/config" >}}#kubernetes-
 It will help to achieve faster recovery.
 Notice that high availability should be enabled when starting standby JobManagers.
 
+### Graceful TaskManager Termination
+
+By default, Flink's Kubernetes deployment reacts to a lost TaskManager the same way whether the
+pod was evicted on purpose (a node drain, a spot/preemptible reclaim, a rolling upgrade) or crashed
+outright: the ResourceManager only notices once the TaskManager's heartbeat times out (50 seconds
+by default), and the TaskManager pod itself is torn down as soon as Kubernetes sends it `SIGTERM`,
+with no chance for currently-running tasks to keep making progress in the meantime.
+
+Setting <a href="{{< ref "docs/deployment/config" >}}#kubernetes-taskmanager-termination-grace-period">kubernetes.taskmanager.termination-grace-period</a>
+changes this for voluntary evictions. It sets the pod's `terminationGracePeriodSeconds` and adds a
+`preStop` lifecycle hook that:
+
+1. Signals the TaskExecutor to disconnect from the ResourceManager immediately, instead of waiting
+   to be timed out, so the ResourceManager can free its slots and schedule a replacement right away.
+2. Keeps the pod alive for most of the configured grace period afterwards (a few seconds are
+   reserved at the end for the JVM's own shutdown sequence to run once Kubernetes sends
+   `SIGTERM`), so tasks that are already running continue to make progress - and get a chance to
+   complete an in-flight checkpoint - instead of being killed the instant the pod is scheduled for
+   deletion.
+
+This only takes effect on a graceful `kubectl delete pod` / node drain, where Kubernetes actually
+runs `preStop` and waits out the grace period; it does not change how quickly Flink notices a pod
+that has actually crashed or become unreachable, since there is no signal to send to a process
+that is no longer running. It has no effect if the pod template already sets
+`terminationGracePeriodSeconds` or a `preStop` hook.
+
 ### Manual Resource Cleanup
 
 Flink uses [Kubernetes OwnerReference's](https://kubernetes.io/docs/concepts/workloads/controllers/garbage-collection/) to clean up all cluster components.

docs/layouts/shortcodes/generated/kubernetes_config_configuration.html

@@ -326,6 +326,12 @@
             <td>String</td>
             <td>Service account that is used by taskmanager within kubernetes cluster. The task manager uses this service account when watching config maps on the API server to retrieve leader address of jobmanager and resourcemanager. If not explicitly configured, config option 'kubernetes.service-account' will be used.</td>
         </tr>
+        <tr>
+            <td><h5>kubernetes.taskmanager.termination-grace-period</h5></td>
+            <td style="word-wrap: break-word;">(none)</td>
+            <td>Duration</td>
+            <td>If set, the TaskManager pod's <code class="highlighter-rouge">terminationGracePeriodSeconds</code> is set to this value, and a <code class="highlighter-rouge">preStop</code> lifecycle hook is added that signals the TaskExecutor to proactively disconnect from the ResourceManager - instead of waiting to be timed out - so its slots are freed immediately rather than after the heartbeat timeout elapses. Left unset, no lifecycle hook is added and Kubernetes' own default grace period behavior applies. Has no effect if the pod template already configures <code class="highlighter-rouge">terminationGracePeriodSeconds</code> or a <code class="highlighter-rouge">preStop</code> hook.</td>
+        </tr>
         <tr>
             <td><h5>kubernetes.taskmanager.tolerations</h5></td>
             <td style="word-wrap: break-word;">(none)</td>

flink-kubernetes/src/main/java/org/apache/flink/kubernetes/configuration/KubernetesConfigOptions.java

@@ -656,6 +656,27 @@ public class KubernetesConfigOptions {
      * #KUBERNETES_PERSISTENT_VOLUME_CLAIMS}. If you need different access modes for different PVCs,
      * consider using pod templates instead.
      */
+    public static final ConfigOption<Duration> TASK_MANAGER_TERMINATION_GRACE_PERIOD =
+            key("kubernetes.taskmanager.termination-grace-period")
+                    .durationType()
+                    .noDefaultValue()
+                    .withDescription(
+                            Description.builder()
+                                    .text(
+                                            "If set, the TaskManager pod's %s is set to this value, and a %s "
+                                                    + "lifecycle hook is added that signals the TaskExecutor to "
+                                                    + "proactively disconnect from the ResourceManager - instead of "
+                                                    + "waiting to be timed out - so its slots are freed immediately "
+                                                    + "rather than after the heartbeat timeout elapses. Left unset, "
+                                                    + "no lifecycle hook is added and Kubernetes' own default grace "
+                                                    + "period behavior applies. Has no effect if the pod template "
+                                                    + "already configures %s or a %s hook.",
+                                            code("terminationGracePeriodSeconds"),
+                                            code("preStop"),
+                                            code("terminationGracePeriodSeconds"),
+                                            code("preStop"))
+                                    .build());
+
     public static final ConfigOption<Boolean> KUBERNETES_PERSISTENT_VOLUME_CLAIM_READ_ONLY =
             key("kubernetes.persistent-volume-claim-read-only")
                     .booleanType()

flink-kubernetes/src/main/java/org/apache/flink/kubernetes/kubeclient/decorators/TerminationGracePeriodDecorator.java

@@ -0,0 +1,143 @@
+/*
+ * 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.flink.kubernetes.kubeclient.decorators;
+
+import org.apache.flink.configuration.ConfigOption;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.kubernetes.kubeclient.FlinkPod;
+import org.apache.flink.kubernetes.kubeclient.parameters.AbstractKubernetesParameters;
+
+import io.fabric8.kubernetes.api.model.Container;
+import io.fabric8.kubernetes.api.model.ContainerBuilder;
+import io.fabric8.kubernetes.api.model.ExecActionBuilder;
+import io.fabric8.kubernetes.api.model.PodBuilder;
+
+import java.time.Duration;
+
+import static org.apache.flink.util.Preconditions.checkNotNull;
+
+/**
+ * Sets {@code terminationGracePeriodSeconds} and a {@code preStop} lifecycle hook on a pod, gated
+ * by a termination-grace-period {@link ConfigOption} (see {@code
+ * KubernetesConfigOptions#TASK_MANAGER_TERMINATION_GRACE_PERIOD}). Only wired in for the
+ * TaskManager today; the {@link ConfigOption} is taken as a constructor argument so the same
+ * decorator can be reused for the JobManager once it has an equivalent {@code SIGUSR2} handler.
+ *
+ * <p>The {@code preStop} hook sends {@code SIGUSR2} to the main container process and then sleeps
+ * for most of the remaining grace period before returning - Kubernetes only sends {@code SIGTERM}
+ * once {@code preStop} returns, so without that sleep the grace period configured here would never
+ * actually be available to already-running tasks; {@code SIGTERM} would follow the signal almost
+ * immediately. {@code TaskManagerRunner} interprets {@code SIGUSR2} as a request to prepare for
+ * graceful termination: it disconnects from the ResourceManager immediately instead of waiting to
+ * be timed out, so the ResourceManager frees this TaskExecutor's slots right away rather than after
+ * the heartbeat timeout elapses, while currently-running tasks keep running for the rest of the
+ * sleep. A {@code SIGUSR2} handler must exist on the receiving side before this decorator is wired
+ * in for a given pod type - Kubernetes' and the JVM's default disposition for an unhandled {@code
+ * SIGUSR2} is to terminate the process, which would make termination less graceful, not more.
+ *
+ * <p>Does nothing if the corresponding option is unset, and never overwrites a {@code
+ * terminationGracePeriodSeconds} or {@code preStop} hook already present in the user's pod
+ * template.
+ */
+public class TerminationGracePeriodDecorator extends AbstractKubernetesStepDecorator {
+
+    /**
+     * How much of the grace period to reserve for the JVM's own shutdown sequence to run after
+     * {@code SIGTERM}, once the {@code preStop} hook returns and Kubernetes sends it.
+     */
+    private static final long SAFETY_MARGIN_SECONDS = 5;
+
+    private final AbstractKubernetesParameters kubernetesParameters;
+    private final ConfigOption<Duration> terminationGracePeriodOption;
+
+    public TerminationGracePeriodDecorator(
+            AbstractKubernetesParameters kubernetesParameters,
+            ConfigOption<Duration> terminationGracePeriodOption) {
+        this.kubernetesParameters = checkNotNull(kubernetesParameters);
+        this.terminationGracePeriodOption = checkNotNull(terminationGracePeriodOption);
+    }
+
+    @Override
+    public FlinkPod decorateFlinkPod(FlinkPod flinkPod) {
+        final Configuration flinkConfig = kubernetesParameters.getFlinkConfiguration();
+        if (!flinkConfig.contains(terminationGracePeriodOption)) {
+            return flinkPod;
+        }
+        final Duration terminationGracePeriod = flinkConfig.get(terminationGracePeriodOption);
+
+        final PodBuilder podBuilder = new PodBuilder(flinkPod.getPodWithoutMainContainer());
+        if (flinkPod.getPodWithoutMainContainer().getSpec().getTerminationGracePeriodSeconds()
+                == null) {
+            podBuilder
+                    .editOrNewSpec()
+                    .withTerminationGracePeriodSeconds(terminationGracePeriod.getSeconds())
+                    .endSpec();
+        } else {
+            logger.info(
+                    "Not overwriting terminationGracePeriodSeconds already set by the pod "
+                            + "template; '{}' has no effect on the grace period.",
+                    terminationGracePeriodOption.key());
+        }
+
+        final Container mainContainer = flinkPod.getMainContainer();
+        final Container decoratedMainContainer;
+        if (mainContainer.getLifecycle() != null
+                && mainContainer.getLifecycle().getPreStop() != null) {
+            logger.info(
+                    "Not overwriting the preStop hook already set by the pod template; '{}' has "
+                            + "no effect on the container lifecycle.",
+                    terminationGracePeriodOption.key());
+            decoratedMainContainer = mainContainer;
+        } else {
+            decoratedMainContainer =
+                    new ContainerBuilder(mainContainer)
+                            .editOrNewLifecycle()
+                            .withNewPreStop()
+                            .withExec(
+                                    new ExecActionBuilder()
+                                            .withCommand(
+                                                    "sh",
+                                                    "-c",
+                                                    buildGracefulTerminationCommand(
+                                                            terminationGracePeriod))
+                                            .build())
+                            .endPreStop()
+                            .endLifecycle()
+                            .build();
+        }
+
+        return new FlinkPod.Builder(flinkPod)
+                .withPod(podBuilder.build())
+                .withMainContainer(decoratedMainContainer)
+                .build();
+    }
+
+    /**
+     * {@code 1} is the main container process because Flink's docker-entrypoint.sh execs it
+     * directly rather than running it as a child of the shell. The sleep keeps {@code preStop} from
+     * returning - and therefore Kubernetes from sending {@code SIGTERM} - until most of the grace
+     * period has elapsed, leaving {@link #SAFETY_MARGIN_SECONDS} for the JVM's own shutdown
+     * sequence to run afterwards.
+     */
+    private static String buildGracefulTerminationCommand(Duration terminationGracePeriod) {
+        final long sleepSeconds =
+                Math.max(0, terminationGracePeriod.getSeconds() - SAFETY_MARGIN_SECONDS);
+        return String.format("kill -USR2 1 2>/dev/null || true; sleep %d", sleepSeconds);
+    }
+}

flink-kubernetes/src/main/java/org/apache/flink/kubernetes/kubeclient/factory/KubernetesTaskManagerFactory.java

@@ -29,6 +29,7 @@
 import org.apache.flink.kubernetes.kubeclient.decorators.KubernetesStepDecorator;
 import org.apache.flink.kubernetes.kubeclient.decorators.MountSecretsDecorator;
 import org.apache.flink.kubernetes.kubeclient.decorators.PersistentVolumeClaimMountDecorator;
+import org.apache.flink.kubernetes.kubeclient.decorators.TerminationGracePeriodDecorator;
 import org.apache.flink.kubernetes.kubeclient.parameters.KubernetesTaskManagerParameters;
 import org.apache.flink.kubernetes.kubeclient.resources.KubernetesPod;
 import org.apache.flink.util.Preconditions;
@@ -42,6 +43,7 @@
 
 import static org.apache.flink.kubernetes.configuration.KubernetesConfigOptions.KUBERNETES_HADOOP_CONF_MOUNT_DECORATOR_ENABLED;
 import static org.apache.flink.kubernetes.configuration.KubernetesConfigOptions.KUBERNETES_KERBEROS_MOUNT_DECORATOR_ENABLED;
+import static org.apache.flink.kubernetes.configuration.KubernetesConfigOptions.TASK_MANAGER_TERMINATION_GRACE_PERIOD;
 
 /** Utility class for constructing the TaskManager Pod on the JobManager. */
 public class KubernetesTaskManagerFactory {
@@ -58,7 +60,10 @@ public static KubernetesPod buildTaskManagerKubernetesPod(
                                 new MountSecretsDecorator(kubernetesTaskManagerParameters),
                                 new PersistentVolumeClaimMountDecorator(
                                         kubernetesTaskManagerParameters),
-                                new CmdTaskManagerDecorator(kubernetesTaskManagerParameters)));
+                                new CmdTaskManagerDecorator(kubernetesTaskManagerParameters),
+                                new TerminationGracePeriodDecorator(
+                                        kubernetesTaskManagerParameters,
+                                        TASK_MANAGER_TERMINATION_GRACE_PERIOD)));
 
         Configuration configuration = kubernetesTaskManagerParameters.getFlinkConfiguration();
         if (configuration.get(KUBERNETES_HADOOP_CONF_MOUNT_DECORATOR_ENABLED)) {