flink+es实战优化

package com.juege.business;

import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.shaded.jackson2.com.fasterxml.jackson.databind.JsonNode;
import org.apache.flink.shaded.jackson2.com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.functions.source.SourceFunction;
import org.apache.http.HttpHost;
import org.apache.http.util.EntityUtils;
import org.elasticsearch.client.Request;
import org.elasticsearch.client.Response;
import org.elasticsearch.client.RestClient;

import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;

public class ElasticsearchScrollSourceFunction<T> implements SourceFunction<T>, CheckpointedFunction {

    private final String host;
    private final int port;
    private final String index;
    private final String type;
    private final String query;

    private volatile boolean isRunning = true;
    private volatile String scrollId;

    // 缓存列表和计数器
    private transient List<T> cacheList;
    private transient int count = 0;

    // 保存当前 Scroll ID 的状态后端
    private transient ListState<String> checkpointedState;

    public ElasticsearchScrollSourceFunction(String host, int port, String index, String type, String query) {
        this.host = host;
        this.port = port;
        this.index = index;
        this.type = type;
        this.query = query;
    }

    @Override
    public void run(SourceContext<T> ctx) throws Exception {
        // 创建 Elasticsearch 客户端
        RestClient client = RestClient.builder(new HttpHost(host, port)).build();

        Map<String, String> params = new HashMap<>();
        params.put("scroll", "1m");

        // 第一次查询，获取 Scroll ID
        Request request = new Request("POST", "/" + index + "/" + type + "/_search");
        request.setJsonEntity(query);
        Response response = client.performRequest(request);
        String responseBody = EntityUtils.toString(response.getEntity());
        JsonNode jsonNode = new ObjectMapper().readTree(responseBody);
        scrollId = jsonNode.get("_scroll_id").asText();

        while (isRunning) {
            // 使用 Scroll API 进行分页查询
            request = new Request("POST", "/_search/scroll");
            request.addParameter("scroll", "1m");
            request.setJsonEntity(scrollId);
            response = client.performRequest(request);
            responseBody = EntityUtils.toString(response.getEntity());

            // 解析查询结果并将数据添加到缓存列表中
            List<T> dataList = parseData(responseBody);
            if (dataList != null && !dataList.isEmpty()) {
                cacheList.addAll(dataList);
                count += dataList.size();
            }

            // 如果达到缓存阈值或者到达 Checkpoint，则将缓存中的数据发送出去
            if (count >= 10000) {
                for (T data : cacheList) {
                    ctx.collect(data);
                }
                cacheList.clear();
                count = 0;

                // 将当前的 Scroll ID 保存到状态后端中
                if (checkpointedState != null) {
                    checkpointedState.clear();
                    checkpointedState.add(scrollId);
                }
            }

            // 更新 Scroll ID，准备下一次查询
            jsonNode = new ObjectMapper().readTree(responseBody);
            scrollId = jsonNode.get("_scroll_id").asText();
        }
    }

    @Override
    public void cancel() {
        isRunning = false;
    }

    @Override
    public void snapshotState(FunctionSnapshotContext context) throws Exception {
        // Snapshot 时不需要进行任何操作，因为 Scroll ID 已经保存在 ListState 中

    }

    @Override
    public void initializeState(FunctionInitializationContext context) throws Exception {
        // 从状态后端中获取已保存的 Scroll ID
        checkpointedState = context.getOperatorStateStore()
                .getListState(new ListStateDescriptor<>("scroll-id-list", Types.STRING));
        if (context.isRestored()) {
            for (String scrollId : checkpointedState.get()) {
                this.scrollId = scrollId;
            }
        }

        // 初始化缓存列表
        cacheList = new LinkedList<>();
    }

    private List<T> parseData(String responseBody) {
        // 解析查询结果，并将数据转换成对应的对象列表
        // ...
        return null;
    }

}

所需依赖

implementation "org.apache.flink:flink-core_2.12:1.13.2"
implementation "org.apache.flink:flink-connector-elasticsearch7_2.12:1.13.2"
implementation "org.elasticsearch.client:elasticsearch-rest-high-level-client:7.14.0"
implementation "com.fasterxml.jackson.core:jackson-databind:2.12.5"
implementation "org.apache.flink:flink-json_2.12:1.13.2"