微服务架构 | 优雅停机 - [完整方案]

/**
 * spring 容器 ApplicationEvent 监听器
 * 此事件在 spring 容器启动 refreshContext() -> finishRefresh() 时由 publishEvent 发布
 * 此事件发布时，已经完成刷新 sprign 容器上下文，但 callRunners() 还没调用
 *  即 > refreshContext()
 *    < ApplicationRunner
 *    < CommandLineRunner
 */
@Slf4j
@Component
public class SpringRefreshedEnentListrner implements ApplicationListener {

    // ApplicationShutdownHooks 全类名，此类非 public，只能 forname 获取
    private static final String HOOKS_CLASS_NAME = "java.lang.ApplicationShutdownHooks";

    @Resource
    private ConfigurableShutDownHooksSorter sorter;

    @Value("${graceful.graceful:30}")
    private long gracefulLimit;

    @Resource
    ApplicationContext applicationContext;

    @Override
    public void onApplicationEvent(ApplicationEvent event) {
        if (!(event instanceof ContextRefreshedEvent)) return;
        if (!Objects.equals(
                ((ApplicationContextEvent)event).getApplicationContext(), this.applicationContext))
            return;

        log.info("[SpringRefreshedEnentListrner] on event start...");

        log.info("[SpringRefreshedEnentListrner] rearrange application shutdown hooks");
        rearrangeApplicationShutdownHooks();

        log.info("[SpringRefreshedEnentListrner] register application shutdown hook trigger");
        renewApplicationShutdownTrigger();

        log.info("[SpringRefreshedEnentListrner] on event done...");
    }

    /* *******************************
     * 更新应用 shutdownhook，只留一个自定义的，作为触发用
     * 原有内容全部移入新容器
     ******************************* */
    private void renewApplicationShutdownTrigger() {
        TriggerSpringShutdownHook.register(
                gracefulLimit, (ConfigurableApplicationContext) applicationContext);
    }

    /* *******************************
     * 重排应用 shutdownhook 至新容器
     *  获取原 spring 的 shutdownhook 容器
     *  按配置顺序排序后加入新容器
     * 清空原容器
     *  原容器在类加载阶段就会注册自己的将自己的 runHooks() 包装成 Runnable 注册至系统(java.lang.Shutdown)
     *  清空原容器可以保证这个行为失效
     * 新容器中注册的 shutdownhook 会有序串行执行
     ******************************* */
    private void rearrangeApplicationShutdownHooks() {
        try {
            // 找到容器原有的 shutdownhooks
            Class hooksClz = Class.forName(HOOKS_CLASS_NAME);
            Field hooksField = hooksClz.getDeclaredField("hooks");
            hooksField.setAccessible(true);
            IdentityHashMap<Thread, Thread> origHooks =
                    (IdentityHashMap<Thread, Thread>) hooksField.get(hooksClz);

            if(MapUtils.isEmpty(origHooks)){
                log.info("[rearrangeApplicationShutdownHooks] empty application hooks, skiped");
            }

            // shutdownhooks 排序配置
            sorter.init();
            log.info("[rearrangeApplicationShutdownHooks] sorting shutdownhook {} by {}"
                    ,origHooks.size() ,sorter.toString());

            synchronized (hooksClz) {
                for (Thread hook : origHooks.keySet()) {
                    log.info("[rearrangeApplicationShutdownHooks] shutdownHook: [{}]-[{}]"
                            ,hook.getClass().getName(), hook.getName());
                    SortedSerialRunningShutdownHookHolder.register(hook, sorter.sorting(hook.getName()));
                }
                origHooks.clear();
            }
            SortedSerialRunningShutdownHookHolder.report();
            log.info("[rearrangeApplicationShutdownHooks] done");
        } catch (Exception e) {
            log.error("error on rearrange application shutdown hooks : ",e);
            throw new RuntimeException(e);
        }
    }
}

/**
 * 排序的，串行执行的 shutdownhook holder
 * 使用后优先级队列作为 hooks 的容器，元素类型为 元组
 * 此类主抄 ：
 * @see java.lang.ApplicationShutdownHooks
 */
@Slf4j
public class SortedSerialRunningShutdownHookHolder {
    /**
     * 容器，使用优先级队列
     * 默认小顶堆，越小越先 poll，所以需要先 shutdown 的应该分数小
     */
    private static PriorityQueue<Pair<Thread,Integer>> HOOKS =
            new PriorityQueue<>(Comparator.comparingInt(Pair::getValue));

    /**
     * 注册 hook，必须提供对应的分数用于排序
     * @param hook
     * @param score
     */
    public static synchronized void register(Thread hook,int score){
        if (hook.isAlive())
            throw new IllegalArgumentException("Hook already running");

        if (hasBeenRegisted(hook))
            throw new IllegalArgumentException("Hook previously registered");

        HOOKS.add(Pair.of(hook,score));
    }

    /**
     * 注销 hook
     * @param hook
     */
    public static synchronized void unregister(Thread hook){
        if (hook == null)
            throw new IllegalArgumentException();

        Iterator<Pair<Thread,Integer>> hsi = HOOKS.iterator();

        Pair<Thread,Integer> node;
        while(hsi.hasNext()){
            node = hsi.next();
            if(node.getKey() == hook) hsi.remove();
            break;
        }
    }

    /**
     * 执行 hooks
     * 按优先级队列中顺序，
     */
    public static synchronized void runHooks(){
        PriorityQueue<Pair<Thread,Integer>> toBeRuns;

        //执行时，使原引用无效化，从 ApplicationShutdownHooks 抄的
        synchronized(SortedSerialRunningShutdownHookHolder.class) {
            toBeRuns = HOOKS;
            HOOKS = null;
        }

        // 串行执行
        Pair<Thread,Integer> curr = null;
        Thread t = null;
        while( (curr=toBeRuns.poll()) !=null){
            t = curr.getKey();
            try {
                log.info("[shutdown hook holder] run: {} ",t.getName());
                t.start();
                t.join();
                log.info("[shutdown hook holder] finished: {} ",t.getName());
            } catch (InterruptedException e) {
                log.warn("SortedSerialRunningShutdownHooks has been interrupted");
            }
        }
    }

    /**
     * 打印容器信息，形如：
     * [(thread_name,score),(thread_name,score)]
     */
    public static void report() {
        if(CollectionUtils.isEmpty(HOOKS)) {
            log.info("[SortedSerialRunningShutdownHookHolder] report: {} ","");
            return ;
        }

        StringBuilder report = new StringBuilder("[");
        for(Pair<Thread,Integer> hook: HOOKS){
            report.append("(").append(hook.getKey().getName())
                    .append(",").append(hook.getValue()).append("),");
        }
        report.deleteCharAt(report.length()-1);
        report.append("]");

        log.info("[SortedSerialRunningShutdownHookHolder] report: {} ",report);
    }

    /* *******************************
     * 以下是工具
     ******************************* */

    private static boolean hasBeenRegisted(Thread target) {
        for(Pair<Thread,Integer> hook: HOOKS){
            if(hook.getKey()==target) return true;
        }
        return false;
    }
}

@Component
@Data
public class ConfigurableShutDownHooksSorter {

//    @Value("#{'${shutdownhook.sorted}'.split(',')}")
    @Value("${graceful.shutdownhook.sorted}")
    private List<String> sorted;

    @Value("${graceful.shutdownhook.outOfSorting:65535}")
    private int outOfSorting;

    /* *******************************
     * 用于处理有冲突的场景，避免想在 a、b两个hook之间插一个，但这俩挨着的场景
     * 示例：new ConfigurableShutDownHooksSorter(hooks,x->( 2 >> (x+2) ));
     ******************************* */
    private Function<Integer,Integer> processer;
    /* *******************************
     * 以下是构造
     ******************************* */

    public ConfigurableShutDownHooksSorter() {    }

    /**
     * 额外指定处理器，processer
     * 用于处理有冲突的场景，避免想在 a、b两个hook之间插一个，但这俩挨着的场景
     * 示例：new ConfigurableShutDownHooksSorter(hooks,x->( 2 >> (x+2) ));
     * @param sorted
     * @param processer
     */
    public ConfigurableShutDownHooksSorter(List<String> sorted, Function<Integer, Integer> processer) {
        this.sorted = sorted;
        this.processer = processer;
    }
    /* *******************************
     * 以下是工具
     ******************************* */

    /**
     * 安全的初始化方法，仅在首次检查且未初始化时初始化一次
     * 虽然是安全的，但也不要随便玩
     */
    public void init(){
        if(CollectionUtils.isEmpty(sorted)){
            synchronized (ConfigurableShutDownHooksSorter.class){
                if(CollectionUtils.isEmpty(sorted)){
                    sorted = new ArrayList<>();
                    sorted.add("as-shutdown-hooker");
                    sorted.add("DubboShutdownHook-NettyClient");
                    sorted.add("DubboShutdownHook");
                }
            }
        }
    }

    /**
     * 返回 hook 的排序，排序完全由配置指定（有个兜底的默认排序）
     * 没有找到时，返回-1 ，返回 -1 时，意味着不在意排序
     * 这些hook认为需要在最先、最后处理都可以，推荐最后
     * @param hookName
     * @return
     */
    public int sorting(String hookName){
        int score = sorted.indexOf(hookName);
        if(-1 == score) score = outOfSorting;

        return processer==null?score:processer.apply(score);
    }

    @Override
    public String toString() {
        final StringBuilder sb = new StringBuilder("{");
        sb.append("\"sorted\":").append(sorted);
        sb.append(",\"outOfSorting\":").append(outOfSorting);
        sb.append('}');
        return sb.toString();
    }
}

@Slf4j
public class TriggerSpringShutdownHook extends Thread {
    public static TriggerSpringShutdownHook INSTANCE = null;
    public static final String DEFAULT_THREAD_NAME = "triggerSpringShutdownHook";
    public static final String DEFAULT_FIELD_NAME = "shutdownHook";

    private long gracefulLimit;
    private ConfigurableApplicationContext context;

    /**
     * hook 的工作流程分 3 步
     * - 截断所有流量
     *      - dubbo 没有提供直接的截断方式，只能通过从注册中心注销+等待注册中心同步所有消费者实现
     *      - mq 没有提供直接的截断方式，只能通过挂起所有 container 处理
     * - 统一睡眠
     *      - 睡眠时间建议值为所有流量所需等待时间的最大值
     * - 执行 shutdown
     *      - 运行非 spring shutdownhook，被 SortedSerialRunningShutdownHookHolder 纳管
     *      - 运行 spring application context 的 close()，
     */
    @Override
    public void run() {
        //切断流量 dubbo+mq
        log.info("Application shutdown hooking for [flow truncate]");
        AbstractRegistryFactory.destroyAll();//truncate dubbo


        //等待时间，因为不好监控实际流量情况，所以直接等待这个时长
        log.info("Application shutdown hooking for [graceful time]");
        try {
            TimeUnit.SECONDS.sleep(gracefulLimit);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();//仅设置标记，保持礼貌，但下面的hook应该需要继续执行
            log.error("interrupted :", e);
        }

        //通过新容器串行调用 hooks
        log.info("Application shutdown hooking for [other hooks holder]");
        SortedSerialRunningShutdownHookHolder.runHooks();
        //spring 容器本身的关闭
        log.info("Application shutdown hooking for [applicatoin hook]");
        context.close();
    }

    /* *******************************
     * 以下是静态工具
     ******************************* */

    /**
     * 注册 spring 容器自己的 shutdownhook
     * 原生的 registerShutdownHook() 包含 3 个行为，此方法中完整复刻，以防莫名其妙的问题
     *  - 创建 shutdownhook 的线程实例
     *  - 用此实例给 AbstractApplicationContext.shutdownHook 赋值
     *  - 将此实例注册到 ApplicationShutdownHooks.hooks 容器
     * @param gracefulLimit
     * @param context
     */
    public static void register(long gracefulLimit,ConfigurableApplicationContext context){
        try {
            TriggerSpringShutdownHook.get(DEFAULT_THREAD_NAME ,gracefulLimit, context);
            log.info("[TriggerSpringShutdownHook] register: instanced");

            Field hook = AbstractApplicationContext.class.getDeclaredField(DEFAULT_FIELD_NAME);
            hook.setAccessible(true);
            hook.set(context,TriggerSpringShutdownHook.get());
            log.info("[TriggerSpringShutdownHook] register: spring context done");

            Runtime.getRuntime().addShutdownHook(TriggerSpringShutdownHook.get());
            log.info("[TriggerSpringShutdownHook] register: application hooks done");
        } catch (Exception e) {
            log.error("[TriggerSpringShutdownHook] exception: ",e);
        }
    }

    /**
     * 获取唯一实例
     * 但如果没有经过初始化会报错
     * @return
     */
    public static TriggerSpringShutdownHook get(){
        if(null == INSTANCE)
            throw new IllegalStateException("TriggerSpringShutdownHook has not bean initialized");

        return INSTANCE;
    }

    /**
     * 初始化、存储、获取唯一实例
     * @param name
     * @param gracefulLimit
     * @param context
     * @return
     */
    public static TriggerSpringShutdownHook get(String name,long gracefulLimit,ConfigurableApplicationContext context){
        if(null == INSTANCE){
            synchronized (TriggerSpringShutdownHook.class){
                if(null == INSTANCE){
                    INSTANCE = new TriggerSpringShutdownHook(name);
                    INSTANCE.gracefulLimit = gracefulLimit;
                    INSTANCE.context = context;
                }else{
                    log.warn("TriggerSpringShutdownHook is not updated");
                }
            }
        }
        return INSTANCE;
    }


    /* *******************************
     * 以下是构造
     ******************************* */
    private TriggerSpringShutdownHook(String name) {
        super(name);
    }
}