11package org.gitanimals.core
22
33import jakarta.annotation.PreDestroy
4- import kotlinx.coroutines.*
4+ import kotlinx.coroutines.CoroutineDispatcher
5+ import kotlinx.coroutines.CoroutineScope
6+ import kotlinx.coroutines.asCoroutineDispatcher
7+ import kotlinx.coroutines.launch
58import kotlinx.coroutines.slf4j.MDCContext
6- import org.gitanimals.core.GracefulShutdownDispatcher.executorService
9+ import org.gitanimals.core.GracefulShutdownDispatcher.graceFulShutdownExecutorServices
710import org.slf4j.LoggerFactory
811import org.springframework.stereotype.Component
12+ import java.util.concurrent.ExecutorService
913import java.util.concurrent.Executors
1014import java.util.concurrent.TimeUnit
1115
1216object GracefulShutdownDispatcher {
1317
14- val executorService = Executors .newFixedThreadPool(10 ) { runnable ->
18+ val graceFulShutdownExecutorServices: MutableList <ExecutorService > = mutableListOf ()
19+
20+ private val executorService = Executors .newFixedThreadPool(10 ) { runnable ->
1521 Thread (runnable, " gitanimals-gracefulshutdown" ).apply { isDaemon = false }
16- }
22+ }.withGracefulShutdown()
1723
18- val dispatcher : CoroutineDispatcher = executorService.asCoroutineDispatcher()
24+ private val defaultDispatcher : CoroutineDispatcher = executorService.asCoroutineDispatcher()
1925
20- fun gracefulLaunch (block : suspend CoroutineScope .() -> Unit ) {
26+ fun ExecutorService.withGracefulShutdown (): ExecutorService {
27+ graceFulShutdownExecutorServices.add(this )
28+ return this
29+ }
30+
31+ fun gracefulLaunch (
32+ dispatcher : CoroutineDispatcher = defaultDispatcher,
33+ block : suspend CoroutineScope .() -> Unit
34+ ) {
2135 CoroutineScope (dispatcher + MDCContext ()).launch(block = block)
2236 }
2337}
@@ -30,18 +44,28 @@ class GracefulShutdownHook {
3044 @PreDestroy
3145 fun tryGracefulShutdown () {
3246 logger.info(" Shutting down dispatcher..." )
33- executorService.shutdown()
47+ graceFulShutdownExecutorServices.forEach {
48+ it.shutdown()
49+ }
3450 runCatching {
35- if (! executorService.awaitTermination(60 , TimeUnit .SECONDS )) {
51+ if (
52+ graceFulShutdownExecutorServices.any {
53+ it.awaitTermination(60 , TimeUnit .SECONDS ).not ()
54+ }
55+ ) {
3656 logger.warn(" Forcing shutdown..." )
37- executorService.shutdownNow()
57+ graceFulShutdownExecutorServices.forEach {
58+ it.shutdown()
59+ }
3860 } else {
3961 logger.info(" Shutdown completed gracefully." )
4062 }
4163 }.onFailure {
4264 if (it is InterruptedException ) {
4365 logger.warn(" Shutdown interrupted. Forcing shutdown..." )
44- executorService.shutdownNow()
66+ graceFulShutdownExecutorServices.forEach {
67+ it.shutdownNow()
68+ }
4569 Thread .currentThread().interrupt()
4670 }
4771 }
0 commit comments