如何在 Spring Boot 优雅关闭加入一些自定义机制

<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-web</artifactId>
    <exclusions>
        <!--不使用默认的 tomcat 容器-->
        <exclusion>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-tomcat</artifactId>
        </exclusion>
    </exclusions>
</dependency>
<!--使用 undertow 容器-->
<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-undertow</artifactId>
</dependency>

server:
  # 设置关闭方式为优雅关闭
  shutdown: graceful
  
management:
  endpoint:
    health:
      show-details: always
    # actuator 暴露 /actuator/shutdown 接口用于关闭（由于这里开启了优雅关闭所以其实是优雅关闭）
    shutdown:
      enabled: true
  endpoints:
    jmx:
      exposure:
        exclude: '*'
    web:
      exposure:
        include: '*'

static List<HttpHandlerFactory> createHttpHandlerFactories(Compression compression, boolean useForwardHeaders,
			String serverHeader, Shutdown shutdown, HttpHandlerFactory... initialHttpHandlerFactories) {
	List<HttpHandlerFactory> factories = new ArrayList<>(Arrays.asList(initialHttpHandlerFactories));
	if (compression != null && compression.getEnabled()) {
		factories.add(new CompressionHttpHandlerFactory(compression));
	}
	if (useForwardHeaders) {
		factories.add(Handlers::proxyPeerAddress);
	}
	if (StringUtils.hasText(serverHeader)) {
		factories.add((next) -> Handlers.header(next, "Server", serverHeader));
	}
	//如果指定了优雅关闭，则添加 gracefulShutdown
	if (shutdown == Shutdown.GRACEFUL) {
		factories.add(Handlers::gracefulShutdown);
	}
	return factories;
}

public interface HttpHandler {
    void handleRequest(HttpServerExchange exchange) throws Exception;
}

@Override
public void handleRequest(HttpServerExchange exchange) throws Exception {
    //原子更新，请求计数器加一，返回的 snapshot 是包含是否关闭状态位的数字
    long snapshot = stateUpdater.updateAndGet(this, incrementActive);
    //通过状态位判断是否正在关闭
    if (isShutdown(snapshot)) {
        //如果正在关闭，直接请求数原子减一
        decrementRequests();
        //设置响应码为 503
        exchange.setStatusCode(StatusCodes.SERVICE_UNAVAILABLE);
        //标记请求完成
        exchange.endExchange();
        //直接返回，不继续走其他的 HttpHandler
        return;
    }
    //添加请求完成时候的 listener，这个在请求完成返回响应时会被调用，将计数器原子减一
    exchange.addExchangeCompleteListener(listener);
    //继续走下一个 HttpHandler
    next.handleRequest(exchange);
}

private void createWebServer() {
	WebServer webServer = this.webServer;
	ServletContext servletContext = getServletContext();
	if (webServer == null && servletContext == null) {
		StartupStep createWebServer = this.getApplicationStartup().start("spring.boot.webserver.create");
		ServletWebServerFactory factory = getWebServerFactory();
		createWebServer.tag("factory", factory.getClass().toString());
		this.webServer = factory.getWebServer(getSelfInitializer());
		createWebServer.end();
		//就是这里，创建一个 WebServerGracefulShutdownLifecycle 并注册到当前 ApplicationContext 的 BeanFactory 中
		getBeanFactory().registerSingleton("webServerGracefulShutdown",
				new WebServerGracefulShutdownLifecycle(this.webServer));
		getBeanFactory().registerSingleton("webServerStartStop",
				new WebServerStartStopLifecycle(this, this.webServer));
	}
	else if (servletContext != null) {
		try {
			getSelfInitializer().onStartup(servletContext);
		}
		catch (ServletException ex) {
			throw new ApplicationContextException("Cannot initialize servlet context", ex);
		}
	}
	initPropertySources();
}

@Override
public void stop(Runnable callback) {
	this.running = false;
	this.webServer.shutDownGracefully((result) -> callback.run());
}

//这里的这个 GracefulShutdownHandler 就是前面说的在启动时加的 GracefulShutdownHandler
private volatile GracefulShutdownHandler gracefulShutdown;

@Override
public void shutDownGracefully(GracefulShutdownCallback callback) {
    // 如果 GracefulShutdownHandler 不为 null，证明开启了优雅关闭（server.shutdown=graceful）
	if (this.gracefulShutdown == null) {
	    //为 null，就证明没开启优雅关闭，什么都不等
		callback.shutdownComplete(GracefulShutdownResult.IMMEDIATE);
		return;
	}
	//开启优雅关闭，需要等待请求处理完
	logger.info("Commencing graceful shutdown. Waiting for active requests to complete");
	this.gracefulShutdownCallback.set(callback);
	//调用 GracefulShutdownHandler 的 shutdown 进行优雅关闭
	this.gracefulShutdown.shutdown();
	//调用 GracefulShutdownHandler 的 addShutdownListener 添加关闭后调用的操作，这里是调用 notifyGracefulCallback
	//其实就是调用方法参数的 callback（就是外部的回调）
	this.gracefulShutdown.addShutdownListener((success) -> notifyGracefulCallback(success));
}

private void notifyGracefulCallback(boolean success) {
	GracefulShutdownCallback callback = this.gracefulShutdownCallback.getAndSet(null);
	if (callback != null) {
		if (success) {
			logger.info("Graceful shutdown complete");
			callback.shutdownComplete(GracefulShutdownResult.IDLE);
		}
		else {
			logger.info("Graceful shutdown aborted with one or more requests still active");
			callback.shutdownComplete(GracefulShutdownResult.REQUESTS_ACTIVE);
		}
	}
}

public void shutdown() {
    //设置关闭状态位，并原子 + 1
    stateUpdater.updateAndGet(this, incrementActiveAndShutdown);
    //直接请求数原子减一
    decrementRequests();
}

private void decrementRequests() {
    long snapshot = stateUpdater.updateAndGet(this, decrementActive);
    // Shutdown has completed when the activeCount portion is zero, and shutdown is set.
    //如果与 关闭状态位 MASK 完全相等，证明其他位都是 0，证明剩余处理中的请求数量为 0
    if (snapshot == SHUTDOWN_MASK) {
        //调用 shutdownComplete
        shutdownComplete();
    }
}

private void shutdownComplete() {
    synchronized (lock) {
        lock.notifyAll();
        //调用每个 ShutdownListener 的 shutdown 方法
        for (ShutdownListener listener : shutdownListeners) {
            listener.shutdown(true);
        }
        shutdownListeners.clear();
    }
}

/**
 * 这个方法并不只是字面意思，首先如果不是关闭中不能添加 ShutdownListener
 * 然后如果没有请求了，就直接调用传入的 shutdownListener 的 shutdown 方法
 * 如果还有请求，则添加入 shutdownListeners，等其他调用 shutdownComplete 的时候遍历 shutdownListeners 调用 shutdown
 * lock 主要为了 addShutdownListener 与 shutdownComplete 对 shutdownListeners 的访问安全
 * lock 的 wait notify 主要为了实现 awaitShutdown 机制，我们这里没有提
 */
public void addShutdownListener(final ShutdownListener shutdownListener) {
        synchronized (lock) {
            if (!isShutdown(stateUpdater.get(this))) {
                throw UndertowMessages.MESSAGES.handlerNotShutdown();
            }
            long count = activeCount(stateUpdater.get(this));
            if (count == 0) {
                shutdownListener.shutdown(true);
            } else {
                shutdownListeners.add(shutdownListener);
            }
        }
    }

private void stopBeans() {
    //读取所有的 Lifecycle bean，返回的是一个 LinkedHashMap，遍历它的顺序和放入的顺序一样
    //放入的顺序就是从 BeanFactory 读取所有 Lifecycle 的 Bean 的返回顺序，这个和 Bean 加载顺序有关，不太可控，可能这个版本加载顺序升级一个版本就变了
	Map<String, Lifecycle> lifecycleBeans = getLifecycleBeans();
	//按照每个 Lifecycle 的 Phase 值进行分组
	//如果实现了 Phased 接口就通过其 phase 方法返回得出 phase 值
	//如果没有实现 Phased 接口则认为 Phase 是 0
	Map<Integer, LifecycleGroup> phases = new HashMap<>();
	lifecycleBeans.forEach((beanName, bean) -> {
		int shutdownPhase = getPhase(bean);
		LifecycleGroup group = phases.get(shutdownPhase);
		if (group == null) {
			group = new LifecycleGroup(shutdownPhase, this.timeoutPerShutdownPhase, lifecycleBeans, false);
			phases.put(shutdownPhase, group);
		}
		group.add(beanName, bean);
	});
	//如果不为空，证明有需要关闭的 Lifecycle，开始关闭
	if (!phases.isEmpty()) {
	    //按照 Phase 值倒序
		List<Integer> keys = new ArrayList<>(phases.keySet());
		keys.sort(Collections.reverseOrder());
		//挨个关闭
		for (Integer key : keys) {
			phases.get(key).stop();
		}
	}
}

class WebServerGracefulShutdownLifecycle implements SmartLifecycle {
    ....
}

public interface SmartLifecycle extends Lifecycle, Phased {
    int DEFAULT_PHASE = Integer.MAX_VALUE;
    @Override
	default int getPhase() {
		return DEFAULT_PHASE;
	}
}

@Log4j2
@Component
public class BizThreadPoolShutdownLifecycle implements SmartLifecycle {
    private volatile boolean running = false;
    
    @Override
    public int getPhase() {
        //在 WebServerGracefulShutdownLifecycle 那一组之后
        return SmartLifecycle.DEFAULT_PHASE - 1;
    }

    @Override
    public void start() {
        this.running = true;
    }

    @Override
    public void stop() {
        //在这里编写的优雅关闭逻辑
        this.running = false;
    }

    @Override
    public boolean isRunning() {
        return running;
    }
}

@Log4j2
@Componenet
//仅在包含 Undertow 这个类的时候加载
@ConditionalOnClass(name = "io.undertow.Undertow")
public class ThreadPoolFactoryGracefulShutDownHandler implements ApplicationListener<ApplicationEvent> {
    
    //获取操作 UndertowWebServer 的 gracefulShutdown 字段的句柄
    private static VarHandle undertowGracefulShutdown;
    //获取操作 GracefulShutdownHandler 的 shutdownListeners 字段的句柄
    private static VarHandle undertowShutdownListeners;

    static {
        try {
            undertowGracefulShutdown = MethodHandles
                    .privateLookupIn(UndertowWebServer.class, MethodHandles.lookup())
                    .findVarHandle(UndertowWebServer.class, "gracefulShutdown",
                            GracefulShutdownHandler.class);
            undertowShutdownListeners = MethodHandles
                    .privateLookupIn(GracefulShutdownHandler.class, MethodHandles.lookup())
                    .findVarHandle(GracefulShutdownHandler.class, "shutdownListeners",
                            List.class);
        } catch (Exception e) {
            log.warn("ThreadPoolFactoryGracefulShutDownHandler undertow not found, ignore fetch var handles");
        }
    }

    @Override
    public void onApplicationEvent(ApplicationEvent event) {
        //仅处理 WebServerInitializedEvent 事件，这个是在 WebServer 创建并初始化完成后发出的事件
        if (event instanceof WebServerInitializedEvent) {
            WebServer webServer = ((WebServerInitializedEvent) event).getWebServer();
            //检查当前的 web 容器是否是 UnderTow 的
            if (webServer instanceof UndertowWebServer) {
                GracefulShutdownHandler gracefulShutdownHandler = (GracefulShutdownHandler) undertowGracefulShutdown.getVolatile(webServer);
                //如果启用了优雅关闭，则 gracefulShutdownHandler 不为 null
                if (gracefulShutdownHandler != null) {
                    var shutdownListeners = (List<GracefulShutdownHandler.ShutdownListener>) undertowShutdownListeners.getVolatile(gracefulShutdownHandler);
                    shutdownListeners.add(shutdownSuccessful -> {
                        if (shutdownSuccessful) {
                            //添加你的优雅关闭逻辑
                        } else {
                            log.info("ThreadPoolFactoryGracefulShutDownHandler-onApplicationEvent shutdown failed");
                        }
                    });
                }
            }
        }
    }
}

public static boolean isCompleted(ExecutorService executorService) {
    if (executorService instanceof ThreadPoolExecutor) {
        ThreadPoolExecutor threadPoolExecutor = (ThreadPoolExecutor) executorService;
        //对于 ThreadPoolExecutor，就是判断没有任何 active 的线程了
        return threadPoolExecutor.getActiveCount() == 0;
    } else if (executorService instanceof ForkJoinPool) {
        //对于 ForkJoinPool，复杂一些，就是判断既没有活跃线程，也没有运行的线程，队列里面也没有任何任务并且并没有任何等待提交的任务
        ForkJoinPool forkJoinPool = (ForkJoinPool) executorService;
        return forkJoinPool.getActiveThreadCount() == 0
                && forkJoinPool.getRunningThreadCount() == 0
                && forkJoinPool.getQueuedTaskCount() == 0
                && forkJoinPool.getQueuedSubmissionCount() == 0;
    }
    return true;
}

List<ExecutorService> executorServices = 获取所有线程池
for (int i = 0; i < 3; ) {
    //连续三次，以随机乱序检查所有的线程池都完成了，才认为是真正完成
    Collections.shuffle(executorServices);
    if (executorServices.stream().allMatch(ThreadPoolFactory::isCompleted)) {
        i++;
        log.info("all threads pools are completed, i: {}", i);
    } else {
        //连续三次
        i = 0;
        log.info("not all threads pools are completed, wait for 1s");
        try {
            TimeUnit.SECONDS.sleep(1);
        } catch (InterruptedException ignored) {
        }
    }
}

@Override
public int getPhase() {
    // Returning Integer.MAX_VALUE only suggests that
    // we will be the first bean to shutdown and last bean to start
    return Integer.MAX_VALUE;
}
@Override
public void stop(Runnable callback) {
    stop();
    callback.run();
}
@Override
public void stop() {
    if (this.isRunning()) {
        if (Objects.nonNull(consumer)) {
            //关闭消费者
            consumer.shutdown();
        }
        setRunning(false);
    }
}