refactor coroutine query count logic

Author: Hyeon9mak

build.gradle.kts

@@ -8,7 +8,7 @@ plugins {
 }
 
 group = "hyeon9mak"
-version = "2.2.0-spring-boot-2"
+version = "2.3.2-rc1-spring-boot-2"
 
 java {
     sourceCompatibility = JavaVersion.VERSION_1_8

src/main/java/hyeon9mak/multidatasourcequerycounter/CoroutineQueryCountContextElement.kt

@@ -1,28 +1,25 @@
 package hyeon9mak.multidatasourcequerycounter
 
 import kotlinx.coroutines.ThreadContextElement
-import org.springframework.web.context.request.RequestAttributes
-import org.springframework.web.context.request.RequestContextHolder
 import kotlin.coroutines.CoroutineContext
 
 /**
  * https://kotlinlang.org/api/kotlinx.coroutines/kotlinx-coroutines-core/kotlinx.coroutines/-thread-context-element/
  */
 class CoroutineQueryCountContextElement(
-    private val requestAttributes: RequestAttributes = RequestContextHolder.currentRequestAttributes(),
-) : ThreadContextElement<RequestAttributes> {
+    var queryCountPerRequest: QueryCountPerRequest? = null,
+) : ThreadContextElement<QueryCountPerRequest?> {
 
     companion object Key : CoroutineContext.Key<CoroutineQueryCountContextElement>
 
     override val key: CoroutineContext.Key<CoroutineQueryCountContextElement>
         get() = Key
 
-    override fun updateThreadContext(context: CoroutineContext): RequestAttributes {
-        RequestContextHolder.setRequestAttributes(requestAttributes)
-        return requestAttributes
+    override fun updateThreadContext(context: CoroutineContext): QueryCountPerRequest? {
+        return queryCountPerRequest
     }
 
-    override fun restoreThreadContext(context: CoroutineContext, oldState: RequestAttributes) {
-        RequestContextHolder.setRequestAttributes(oldState)
+    override fun restoreThreadContext(context: CoroutineContext, oldState: QueryCountPerRequest?) {
+        queryCountPerRequest = oldState
     }
 }

src/main/java/hyeon9mak/multidatasourcequerycounter/HikariDataSourceQueryCounter.kt

@@ -11,20 +11,19 @@ import org.springframework.web.context.request.ServletRequestAttributes
 @Aspect
 @Component
 class HikariDataSourceQueryCounter(
-    private val queryCountPerRequest: QueryCountPerRequest,
     private val queryCountLogger: QueryCountLogger,
 ) {
     @Around("execution( * com.zaxxer.hikari.HikariDataSource.getConnection())")
     fun aroundConnection(joinPoint: ProceedingJoinPoint): Any {
         val connection = joinPoint.proceed()
+        val queryCountPerRequest = QueryCountPerRequestHolder.get() ?: return connection
         val connectionQueryMonitor = ConnectionQueryMonitor(queryCountPerRequest, connection)
         return connectionQueryMonitor.getProxy()
     }
 
     @Around("@annotation(countQueries)")
     fun aroundCountQueriesMethod(joinPoint: ProceedingJoinPoint, countQueries: CountQueries): Any {
-        val result = joinPoint.proceed()
-
+        val queryCountPerRequest = QueryCountPerRequest()
         val attributes = RequestContextHolder.getRequestAttributes() as ServletRequestAttributes?
 
         if (attributes.isInRequestScope) {
@@ -37,8 +36,10 @@ class HikariDataSourceQueryCounter(
             queryCountPerRequest.apiUrl = countQueries.prefix + className + "." + methodName
         }
 
+        QueryCountPerRequestHolder.set(queryCountPerRequest)
+        val result = joinPoint.proceed()
         queryCountLogger.logQueryCount(queryCountPerRequest)
-
+        QueryCountPerRequestHolder.remove()
         return result
     }
 

src/main/java/hyeon9mak/multidatasourcequerycounter/PreparedStatementQueryMonitor.kt

@@ -2,15 +2,14 @@ package hyeon9mak.multidatasourcequerycounter
 
 import org.aopalliance.intercept.MethodInterceptor
 import org.aopalliance.intercept.MethodInvocation
-import org.springframework.web.context.request.RequestContextHolder
 import java.lang.System.*
 
 class PreparedStatementQueryMonitor(
     private val queryCountPerRequest: QueryCountPerRequest,
 ) : MethodInterceptor {
 
     override fun invoke(invocation: MethodInvocation): Any? {
-        if (QUERY_METHODS.contains(invocation.method.name) && isRequestScope()) {
+        if (QUERY_METHODS.contains(invocation.method.name)) {
             val startTime = currentTimeMillis()
             val result = invocation.proceed()
             val endTime = currentTimeMillis()
@@ -23,8 +22,6 @@ class PreparedStatementQueryMonitor(
         return invocation.proceed()
     }
 
-    private fun isRequestScope(): Boolean = RequestContextHolder.getRequestAttributes() != null
-
     companion object {
         private val QUERY_METHODS = listOf("executeQuery", "execute", "executeUpdate")
     }

src/main/java/hyeon9mak/multidatasourcequerycounter/QueryCountPerRequest.kt

@@ -1,17 +1,21 @@
 package hyeon9mak.multidatasourcequerycounter
 
-import org.springframework.stereotype.Component
-import org.springframework.web.context.annotation.RequestScope
+import java.util.concurrent.atomic.AtomicInteger
+import java.util.concurrent.atomic.AtomicLong
 
-@RequestScope
-@Component
 data class QueryCountPerRequest(
     var apiUrl: String = "",
-    var totalQueryCount: Int = 0,
-    var totalQueryMilliSeconds: Long = 0L,
+    private var _totalQueryCount: AtomicInteger = AtomicInteger(0),
+    private var _totalQueryMilliSeconds: AtomicLong = AtomicLong(0L),
 ) {
+    val totalQueryCount: Int
+        get() = _totalQueryCount.get()
+
+    val totalQueryMilliSeconds: Long
+        get() = _totalQueryMilliSeconds.get()
+
     fun incrementQueryCount(executionMilliSeconds: Long) {
-        totalQueryCount++
-        totalQueryMilliSeconds += executionMilliSeconds
+        _totalQueryCount.incrementAndGet()
+        _totalQueryMilliSeconds.addAndGet(executionMilliSeconds)
     }
 }

src/main/java/hyeon9mak/multidatasourcequerycounter/QueryCountPerRequestHolder.kt

@@ -0,0 +1,20 @@
+package hyeon9mak.multidatasourcequerycounter
+
+object QueryCountPerRequestHolder {
+
+    private val queryCountPerRequestHolder = ThreadLocal<QueryCountPerRequest?>()
+    private val coroutineContextElement = CoroutineQueryCountContextElement()
+
+    fun set(queryCountPerRequest: QueryCountPerRequest) {
+        queryCountPerRequestHolder.set(queryCountPerRequest)
+        coroutineContextElement.queryCountPerRequest = queryCountPerRequest
+    }
+
+    fun get(): QueryCountPerRequest? = queryCountPerRequestHolder.get()
+        ?: coroutineContextElement.queryCountPerRequest
+
+    fun remove() {
+        queryCountPerRequestHolder.remove()
+        coroutineContextElement.queryCountPerRequest = null
+    }
+}