如何为Flink Task Manager实现自定义健康检查以在连接异常时自动重启？

一、核心思路 #

通过自定义REST端点定期校验S3连接状态，结合Kubernetes的Liveness探针触发Task Manager自动重启，替代手动重启的临时方案。

二、分步实现 #

1. 编写S3连接校验工具类 #

实现轻量的S3连接校验逻辑，通过执行简单操作（如列出测试桶前缀）判断连接是否正常：

import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider;
import software.amazon.awssdk.regions.Region;
import software.amazon.awssdk.services.s3.S3Client;
import software.amazon.awssdk.services.s3.model.ListObjectsV2Request;

public class S3HealthChecker {
    // 替换为实际测试桶和区域
    private static final String TEST_BUCKET = "your-s3-test-bucket";
    private static final String TEST_PREFIX = "health-check-marker/";
    private static final String AWS_REGION = "your-aws-region";

    public static boolean isConnectionHealthy() {
        try (S3Client s3Client = S3Client.builder()
                .credentialsProvider(DefaultCredentialsProvider.create())
                .region(Region.of(AWS_REGION))
                .build()) {
            // 执行轻量操作校验连接有效性
            ListObjectsV2Request request = ListObjectsV2Request.builder()
                    .bucket(TEST_BUCKET)
                    .prefix(TEST_PREFIX)
                    .maxKeys(1)
                    .build();
            s3Client.listObjectsV2(request);
            return true;
        } catch (Exception e) {
            System.err.println("S3 connection check failed: " + e.getMessage());
            return false;
        }
    }
}

2. 自定义Task Manager REST健康检查端点 #

基于Flink的RestEndpointExtension扩展机制，添加专属健康检查接口：

import org.apache.flink.runtime.rest.RestEndpoint;
import org.apache.flink.runtime.rest.RestEndpointExtension;
import org.apache.flink.runtime.rest.handler.RestHandler;
import org.apache.flink.runtime.rest.handler.RestHandlerSpecification;
import org.apache.flink.runtime.rest.messages.EmptyRequestBody;
import org.apache.flink.runtime.rest.messages.MessageHeaders;
import org.apache.flink.runtime.rest.messages.MessageParameters;
import org.apache.flink.runtime.rest.messages.ResponseBody;
import org.apache.flink.shaded.netty4.io.netty.handler.codec.http.HttpResponseStatus;

import javax.annotation.Nonnull;
import java.util.Collection;
import java.util.Collections;
import java.util.Map;
import java.util.concurrent.CompletableFuture;

// 注册REST端点扩展
public class S3HealthCheckExtension implements RestEndpointExtension {
    @Override
    public Collection<RestHandler<?, ?, ?>> getRestHandlers(RestEndpoint restEndpoint) {
        return Collections.singletonList(new S3HealthCheckHandler());
    }
}

// 健康检查请求处理器
class S3HealthCheckHandler implements RestHandler<EmptyRequestBody, S3HealthResponse, MessageParameters> {
    @Override
    public CompletableFuture<S3HealthResponse> handleRequest(EmptyRequestBody request, Map<String, String> pathParams, Map<String, String> queryParams, MessageParameters messageParameters) {
        boolean isHealthy = S3HealthChecker.isConnectionHealthy();
        HttpResponseStatus status = isHealthy ? HttpResponseStatus.OK : HttpResponseStatus.SERVICE_UNAVAILABLE;
        return CompletableFuture.completedFuture(new S3HealthResponse(isHealthy, status.code()));
    }

    @Override
    public RestHandlerSpecification getHandlerSpecification() {
        return RestHandlerSpecification.builder()
                .path("/taskmanager/s3/health")
                .method(org.apache.flink.shaded.netty4.io.netty.handler.codec.http.HttpMethod.GET)
                .build();
    }

    @Override
    public MessageHeaders<EmptyRequestBody, S3HealthResponse, MessageParameters> getMessageHeaders() {
        return new MessageHeaders<EmptyRequestBody, S3HealthResponse, MessageParameters>() {
            @Override
            public Class<EmptyRequestBody> getRequestClass() { return EmptyRequestBody.class; }
            @Override
            public Class<S3HealthResponse> getResponseClass() { return S3HealthResponse.class; }
            @Override
            public HttpResponseStatus getResponseStatusCode() { return HttpResponseStatus.OK; }
            @Nonnull
            @Override
            public MessageParameters getUnresolvedMessageParameters() { return MessageParameters.empty(); }
        };
    }
}

// 健康检查响应体
class S3HealthResponse implements ResponseBody {
    private final boolean healthy;
    private final int statusCode;

    public S3HealthResponse(boolean healthy, int statusCode) {
        this.healthy = healthy;
        this.statusCode = statusCode;
    }

    public boolean isHealthy() { return healthy; }
    public int getStatusCode() { return statusCode; }
}

然后在项目的META-INF/services目录下创建文件org.apache.flink.runtime.rest.RestEndpointExtension，写入扩展类全路径：

com.your.package.S3HealthCheckExtension

3. 配置Kubernetes Liveness探针 #

修改FlinkDeployment CR配置，为Task Manager添加指向自定义端点的Liveness探针：

apiVersion: flink.apache.org/v1beta1
kind: FlinkDeployment
metadata:
  name: your-flink-cluster
spec:
  taskManager:
    replicas: 3
    resources:
      memory: "4096m"
      cpu: 2
    # 添加S3健康检查探针
    livenessProbe:
      httpGet:
        path: /taskmanager/s3/health
        port: rest
      initialDelaySeconds: 60  # 等待Task Manager完全启动后开始检查
      periodSeconds: 15        # 每15秒检查一次
      failureThreshold: 2      # 连续2次失败触发重启

4. 打包部署 #

• 将代码打包成JAR，放入Flink集群的plugins目录，或通过FlinkDeployment的taskManager.spec.jar字段指定加载。
• 确保Task Manager Pod拥有访问S3测试桶的权限（通过IAM角色绑定或环境变量配置AWS凭证）。

三、备选方案：基于自定义指标的监控触发 #

如果已有Prometheus监控体系，可采用以下方式：

• 在Task Manager中定期执行S3连接校验，将结果作为自定义Gauge指标（如s3_connection_healthy，1=健康，0=异常）暴露。
• 配置Prometheus Alertmanager，当指标连续3次为0时触发Kubernetes API调用重启对应Task Manager Pod。

内容的提问来源于stack exchange，提问作者Techie97