Support graceful shutdown for gRPC server

starter/sisyphus-grpc-server-starter/src/main/kotlin/com/bybutter/sisyphus/starter/grpc/ServerLifecycle.kt

@@ -0,0 +1,54 @@
+package com.bybutter.sisyphus.starter.grpc
+
+import io.grpc.Server
+import kotlin.concurrent.thread
+import org.slf4j.LoggerFactory
+import org.springframework.boot.web.server.Shutdown
+import org.springframework.context.SmartLifecycle
+
+class ServerLifecycle(private val server: Server, private val shutdown: Shutdown) : SmartLifecycle {
+    private var running = false
+
+    override fun getPhase(): Int {
+        return super.getPhase() - 10
+    }
+
+    override fun isRunning(): Boolean {
+        if (server.isTerminated) return false
+        if (server.isShutdown) return false
+        return running
+    }
+
+    override fun start() {
+        server.start()
+        running = true
+        logger.info("Running gRPC server via netty on port: ${server.port}")
+    }
+
+    override fun stop() {
+        server.shutdownNow()
+    }
+
+    override fun stop(callback: Runnable) {
+        when (shutdown) {
+            Shutdown.GRACEFUL -> {
+                logger.info("Commencing graceful shutdown for gRPC server. Waiting for active requests to complete")
+
+                server.shutdown()
+                thread(name = "grpc-shutdown") {
+                    server.awaitTermination()
+                    logger.info("Graceful shutdown complete for gRPC server.")
+                    callback.run()
+                }
+            }
+            else -> {
+                stop()
+                callback.run()
+            }
+        }
+    }
+
+    companion object {
+        private val logger = LoggerFactory.getLogger(ServerLifecycle::class.java)
+    }
+}

starter/sisyphus-grpc-server-starter/src/main/kotlin/com/bybutter/sisyphus/starter/grpc/ServiceRegistrar.kt

@@ -16,21 +16,14 @@ import org.springframework.beans.factory.config.ConfigurableListableBeanFactory
 import org.springframework.beans.factory.support.BeanDefinitionBuilder
 import org.springframework.beans.factory.support.BeanDefinitionRegistry
 import org.springframework.beans.factory.support.BeanDefinitionRegistryPostProcessor
+import org.springframework.boot.web.server.Shutdown
 import org.springframework.context.EnvironmentAware
+import org.springframework.context.Lifecycle
 import org.springframework.core.env.Environment
 import org.springframework.stereotype.Component
 
 @Component
 class ServiceRegistrar : BeanDefinitionRegistryPostProcessor, EnvironmentAware {
-    companion object {
-        private val logger = LoggerFactory.getLogger(ServiceRegistrar::class.java)
-
-        const val GRPC_PORT_PROPERTY = "grpc.port"
-        const val DEFAULT_GRPC_PORT = "9090"
-
-        const val QUALIFIER_AUTO_CONFIGURED_GRPC_SERVER = "sisyphus:grpc:server"
-    }
-
     private lateinit var environment: Environment
 
     override fun setEnvironment(environment: Environment) {
@@ -47,7 +40,7 @@ class ServiceRegistrar : BeanDefinitionRegistryPostProcessor, EnvironmentAware {
             }
 
             val services = beanFactory.getBeansWithAnnotation(RpcServiceImpl::class.java)
-            logger.info("${services.size} grpc services registered: ${services.keys.joinToString(", ")}")
+            logger.info("${services.size} gRPC services registered: ${services.keys.joinToString(", ")}")
             for ((_, service) in services) {
                 builder = when (service) {
                     is BindableService -> {
@@ -78,10 +71,28 @@ class ServiceRegistrar : BeanDefinitionRegistryPostProcessor, EnvironmentAware {
 
             builder.build()
         }
-
         (beanFactory as BeanDefinitionRegistry).registerBeanDefinition(QUALIFIER_AUTO_CONFIGURED_GRPC_SERVER, definitionBuilder.beanDefinition)
+
+        val lifecycleBuilder = BeanDefinitionBuilder.genericBeanDefinition(Lifecycle::class.java) {
+            val server = beanFactory.getBean(QUALIFIER_AUTO_CONFIGURED_GRPC_SERVER) as Server
+            val shutdown = environment.getProperty("server.shutdown", Shutdown::class.java)
+            ServerLifecycle(server, shutdown ?: Shutdown.IMMEDIATE)
+        }
+        (beanFactory as BeanDefinitionRegistry).registerBeanDefinition(QUALIFIER_AUTO_CONFIGURED_GRPC_SERVER_LIFECYCLE, lifecycleBuilder.beanDefinition)
     }
 
     override fun postProcessBeanDefinitionRegistry(registry: BeanDefinitionRegistry) {
     }
+
+    companion object {
+        private val logger = LoggerFactory.getLogger(ServiceRegistrar::class.java)
+
+        const val GRPC_PORT_PROPERTY = "server.grpc.port"
+
+        const val DEFAULT_GRPC_PORT = "9090"
+
+        const val QUALIFIER_AUTO_CONFIGURED_GRPC_SERVER = "sisyphus:grpc:server"
+
+        const val QUALIFIER_AUTO_CONFIGURED_GRPC_SERVER_LIFECYCLE = "sisyphus:grpc:server-lifecycle"
+    }
 }

starter/sisyphus-webflux-starter/src/main/kotlin/com/bybutter/sisyphus/starter/webflux/AutoGracefulShutdownConfigurator.kt

@@ -0,0 +1,18 @@
+package com.bybutter.sisyphus.starter.webflux
+
+import java.util.Properties
+import org.springframework.boot.SpringApplication
+import org.springframework.boot.env.EnvironmentPostProcessor
+import org.springframework.boot.web.server.Shutdown
+import org.springframework.core.env.ConfigurableEnvironment
+import org.springframework.core.env.PropertiesPropertySource
+
+class AutoGracefulShutdownConfigurator : EnvironmentPostProcessor {
+    override fun postProcessEnvironment(environment: ConfigurableEnvironment, application: SpringApplication) {
+        if (environment.containsProperty("server.shutdown")) return
+
+        environment.propertySources.addLast(PropertiesPropertySource("AutoGracefulShutdownConfiguration", Properties().apply {
+            this["server.shutdown"] = Shutdown.GRACEFUL.name
+        }))
+    }
+}

starter/sisyphus-webflux-starter/src/main/resources/META-INF/spring.factories

@@ -1,2 +1,4 @@
 org.springframework.boot.autoconfigure.EnableAutoConfiguration=\
-com.bybutter.sisyphus.starter.webflux.SisyphusWebfluxAutoConfiguration
+com.bybutter.sisyphus.starter.webflux.SisyphusWebfluxAutoConfiguration
+org.springframework.boot.env.EnvironmentPostProcessor=\
+com.bybutter.sisyphus.starter.webflux.AutoGracefulShutdownConfigurator