循环(分页)调度接口工具

说明

这个工具是几个月前帮公司重构数据采集工程时写的，主要作用就是把分页查询以及重试机制封装在一起进行调度。写这工具的目的实在是因为祖传代码看不下去了，动不动就是几百行的冗余代码，眼睛看得疼啊。

该工具由三部分组成

1. Pagination 分页对象，已经定义好了
2. SearchUtil 循环查询工具
3. 调用代码，需要调用方自己去写。本文中暂时就不列出了。

目前将代码维护在：https://gitee.com/sulbert/ketools
有时间会逐步完善

Pagination类定义

package com.ssy.search.entity;

import lombok.AllArgsConstructor;
import lombok.Builder;
import lombok.Data;
import lombok.NoArgsConstructor;

import java.io.Serial;
import java.io.Serializable;

/**
 * @author SiYang
 * @description 分页对象
 */
@Data
@Builder
@NoArgsConstructor
@AllArgsConstructor
public class Pagination implements Serializable {

    @Serial private static final long serialVersionUID = -3263304549357348931L;
    private Integer pageNo;
    private Integer pageSize;
    private Integer total;

    /** 是否翻页，有些接口并不通过上面几个属性来控制翻页，而是其他方式。 该字段用来体现：通过其他方式计算得到的最终结果——是否需要翻页 */
    private Boolean turnPage;

    public Pagination(Integer pageNo, Integer pageSize) {
        this.pageNo = pageNo;
        this.pageSize = pageSize;
    }

    public static Pagination getDefaultPagination() {
        return new Pagination(1, 100);
    }

    public void increasePage() {
        setPageNo(getPageNo() + 1);
    }
}

循环查询的工具类

package com.ssy.search.service;

import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;

import org.apache.commons.lang3.tuple.Pair;
import org.springframework.util.CollectionUtils;
import org.springframework.util.ObjectUtils;

import java.time.Duration;
import java.util.*;
import java.util.concurrent.*;
import java.util.function.Consumer;
import java.util.function.Function;

/**
 * @author SiYang
 * @description 查询工具
 */
@Slf4j
public class SearchUtil {

    /** 总共需要查询的次数 */
    private static final ThreadLocal<Integer> LOOP_TIMES = new ThreadLocal<>();

    /** 当前这次请求的开始时间 */
    private static final ThreadLocal<Long> CURRENT_SEARCH_START_TIME = new ThreadLocal<>();

    /** 并行查询时采用的线程池 */
    private static final ThreadLocal<ExecutorService> QUERY_EXECUTOR = new ThreadLocal<>();

    /**
     * 支持查询策略
     *
     * @param searchFunc 调度实际查询的函数
     * @param pagination 分页查询参数对象
     * @param retryTimes 失败时的重试次数
     * @param retryBeforeHandler 在执行重试之前执行的函数
     * @param errorHandler 最终查询错误的处理器
     * @param searchStrategy 查询策略，{@link SearchStrategy} 暂时只有普通查询（单线程）和并行查询
     */
    public static <T> List<T> search(
            Function<Pagination, Pair<Pagination, List<T>>> searchFunc,
            Pagination pagination,
            Integer retryTimes,
            Consumer<Pair<Exception, Duration>> retryBeforeHandler,
            Consumer<Exception> errorHandler,
            SearchStrategy searchStrategy) {
        List<T> resultData = new ArrayList<>();
        return search(resultData, searchFunc, pagination, retryTimes, retryBeforeHandler, errorHandler, searchStrategy);
    }

    public static <T> List<T> search(Function<Pagination, Pair<Pagination, List<T>>> searchFunc, Integer retryTimes) {
        List<T> resultData = new ArrayList<>();
        return search(resultData, searchFunc, null, retryTimes, null, null, SearchStrategy.SIMPLE);
    }

    public static <T> List<T> search(
            List<T> resultData,
            Function<Pagination, Pair<Pagination, List<T>>> searchFunc,
            Pagination pagination,
            Integer retryTimes,
            Consumer<Pair<Exception, Duration>> retryBeforeHandler,
            Consumer<Exception> errorHandler,
            SearchStrategy searchStrategy) {
        try {
            doSearch(resultData, searchFunc, pagination, searchStrategy);
            return resultData;
        } catch (Exception e) {
            retryBeforeHandler = Optional.ofNullable(retryBeforeHandler).orElse(defaultRetryBeforeHandler());
            errorHandler = Optional.ofNullable(errorHandler).orElse(defaultErrorHandler());
            if (retryTimes > 0) {
                log.debug("start to retry: cause {}, msg: {}", e.getStackTrace()[0], e.getMessage());
                Duration duration = Duration.ofMillis(System.currentTimeMillis() - CURRENT_SEARCH_START_TIME.get());
                retryBeforeHandler.accept(Pair.of(e, duration));
                CURRENT_SEARCH_START_TIME.remove();
                return search(resultData, searchFunc, pagination, retryTimes - 1, retryBeforeHandler, errorHandler, searchStrategy);
            } else {
                errorHandler.accept(e);
                return Collections.emptyList();
            }
        }
    }

    private static Consumer<Pair<Exception, Duration>> defaultRetryBeforeHandler() {
        return r -> {
            Exception exc = r.getKey();
            log.warn(
                    "search error: {}, errMsg: {}, start to retry",
                    exc.getStackTrace()[0],
                    exc.getMessage());
        };
    }

    private static Consumer<Exception> defaultErrorHandler() {
        return exc -> log.error("search error: ", exc);
    }

    private static <T> void doSearch(
            List<T> resultData,
            Function<Pagination, Pair<Pagination, List<T>>> searchFunc,
            Pagination pagination,
            SearchStrategy searchStrategy) {
        CURRENT_SEARCH_START_TIME.set(System.currentTimeMillis());
        pagination = Optional.ofNullable(pagination).orElseGet(Pagination::getDefaultPagination);
        searchStrategy.doSearch(resultData, searchFunc, pagination);
    }

    private static Integer getLoopTimes(Integer total, Integer pageSize) {
        if (total == null || total == 0) {
            return 1;
        }
        return (total + pageSize - 1) / pageSize;
    }

    public enum SearchStrategy {
        SIMPLE {
            /**
             * 结束查询的三个充分条件：
             * 1. 本次查询的是最后一页
             * 2. 本次查询返回的数据列表为空
             * 3. turnPage为false
             */
            @Override
            <T> void doSearch(List<T> resultData, Function<Pagination, Pair<Pagination, List<T>>> searchFunc, Pagination pagination) {
                boolean doSearch = true;
                do {
                    Pair<Pagination, List<T>> searchResult = searchFunc.apply(pagination);
                    List<T> resultList = searchResult.getValue();
                    if (!ObjectUtils.isEmpty(resultList) && !CollectionUtils.isEmpty(resultList)) {
                        resultData.addAll(resultList);
                    }

                    if (Objects.isNull(LOOP_TIMES.get())) {
                        LOOP_TIMES.set(
                                getLoopTimes(
                                        searchResult.getKey().getTotal(),
                                        pagination.getPageSize()));
                    }

                    if (Objects.equals(pagination.getPageNo(), LOOP_TIMES.get()) || resultList.isEmpty()) {
                        doSearch = false;
                    }

                    // 优先级最高，覆盖上面的计算结果
                    if (Objects.nonNull(pagination.getTurnPage())) {
                        doSearch = pagination.getTurnPage();
                    }

                    pagination.increasePage();
                } while (doSearch);
                LOOP_TIMES.remove();
                CURRENT_SEARCH_START_TIME.remove();
            }
        },

        PARALLEL {
            @Override
            @SneakyThrows
            <T> void doSearch(List<T> resultData, Function<Pagination, Pair<Pagination, List<T>>> searchFunc, Pagination pagination) {
                // 设置线程池
                if (Objects.isNull(QUERY_EXECUTOR.get())) {
                    QUERY_EXECUTOR.set(Executors.newFixedThreadPool(5));
                }

                // 第一次查询，为了获得totalCount
                Pair<Pagination, List<T>> searchResult = searchFunc.apply(pagination);
                List<T> resultList = searchResult.getValue();

                // resultList是ArrayList，并发场景下用Vector转一下，新定义一个data
                List<T> data = new Vector<>(resultList);

                // 针对page进行并发查询
                int pageCount = getLoopTimes(searchResult.getKey().getTotal(), pagination.getPageSize());

                CountDownLatch latch = new CountDownLatch(pageCount - 1);
                for (int i = 2; i <= pageCount; i++) {
                    // 有些接口同一时间点并发分页查询，返回结果会缺胳膊少腿，稍稍错开一点即可
                    TimeUnit.MILLISECONDS.sleep(50L);

                    // 每个线程都需要有自己的参数对象
                    Pagination p = new Pagination(i, pagination.getPageSize());

                    QUERY_EXECUTOR
                            .get()
                            .execute(
                                    () -> {
                                        try {
                                            executeSearchAndRetry(data, searchFunc, p, 3);
                                        } catch (Exception e) {
                                            log.warn(
                                                    "PARALLEL Search (retried 3 times) Finally Error: {}, loss data on page {}",
                                                    e.getMessage(),
                                                    p.getPageNo());
                                        } finally {
                                            latch.countDown();
                                        }
                                    });
                }

                latch.await();
                QUERY_EXECUTOR.get().shutdown();

                CURRENT_SEARCH_START_TIME.remove();
                QUERY_EXECUTOR.remove();

                resultData.addAll(data);
            }

            @SneakyThrows
            private <T> void executeSearchAndRetry(List<T> resultData, Function<Pagination, Pair<Pagination, List<T>>> searchFunc, Pagination pagination, int retryTimes) {
                try {
                    resultData.addAll(searchFunc.apply(pagination).getValue());
                } catch (Exception e) {
                    log.warn(
                            "PARALLEL Search Error: {}, start to retry for page {}",
                            e.getMessage(),
                            pagination.getPageNo());
                    TimeUnit.SECONDS.sleep(1);
                    if (retryTimes > 0) {
                        executeSearchAndRetry(resultData, searchFunc, pagination, retryTimes - 1);
                    }
                }
            }
        };

        abstract <T> void doSearch(List<T> resultData, Function<Pagination, Pair<Pagination, List<T>>> searchFunc, Pagination pagination);
    }
}