Spring boot Tomcat GracefulShutdown

TomcatGracefulShutdown.java

package com.laobai.boot.tomcat;

import org.apache.catalina.connector.Connector;
import org.springframework.boot.context.embedded.tomcat.TomcatConnectorCustomizer;
import org.springframework.context.ApplicationListener;
import org.springframework.context.event.ContextClosedEvent;

import java.util.concurrent.Executor;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

import lombok.extern.slf4j.Slf4j;

/**
 * @author cqy
 * @since 2017/11/6.
 */
@Slf4j
public class TomcatGracefulShutdown implements TomcatConnectorCustomizer, ApplicationListener<ContextClosedEvent> {

    final long shutdownTimeout;
    final TimeUnit unit;

    private volatile Connector connector;

    public TomcatGracefulShutdown(long shutdownTimeout, TimeUnit unit) {
        this.shutdownTimeout = shutdownTimeout;
        this.unit = unit;
    }

    @Override
    public void customize(Connector connector) {
        this.connector = connector;
    }

    @Override
    public void onApplicationEvent(ContextClosedEvent event) {
        if (this.connector == null) {
            return;
        }
        awaitTermination(this.connector);
    }

    void awaitTermination(Connector connector) {
        connector.pause();
        Executor executor = connector.getProtocolHandler().getExecutor();
        if (executor instanceof ThreadPoolExecutor) {
            log.warn("Context closed. Going to await termination for $shutdownTimeout $unit.");
            try {
                ThreadPoolExecutor ex = (ThreadPoolExecutor) executor;
                ex.shutdown();
                if (!ex.awaitTermination(shutdownTimeout, unit)) {
                    log.warn("Tomcat thread pool did not shut down gracefully within $shutdownTimeout $unit. " +
                            "Proceeding with forceful shutdown");
                }
            } catch (InterruptedException e){
                Thread.currentThread().interrupt();
            }
        }
    }
}

TomcatGracefulShutdownAutoConfiguration.java

package com.laobai.boot.tomcat;

import org.apache.catalina.startup.Tomcat;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
import org.springframework.boot.context.embedded.EmbeddedServletContainerCustomizer;
import org.springframework.boot.context.embedded.tomcat.TomcatEmbeddedServletContainerFactory;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.concurrent.TimeUnit;

import javax.servlet.Servlet;

/**
 * @author cqy
 * @since 2017/11/6.
 */
@Configuration
@ConditionalOnClass({Servlet.class, Tomcat.class})
public class TomcatGracefulShutdownAutoConfiguration {

    @Value("${catalina.threadpool.execution.timeout.seconds:30}")
    long shutdownTimeoutSeconds;

    @Bean
    TomcatGracefulShutdown gracefulShutdown() {
        return new TomcatGracefulShutdown(shutdownTimeoutSeconds, TimeUnit.SECONDS);
    }

    @Bean
    EmbeddedServletContainerCustomizer tomcatCustomizer() {
        return container -> {
            if (container instanceof TomcatEmbeddedServletContainerFactory) {
                ((TomcatEmbeddedServletContainerFactory) container).addConnectorCustomizers(gracefulShutdown());
            }
        };
    }
}