com.firefly.kotlin.ext.log.CoroutineMappedDiagnosticContext.kt Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of firefly-kotlin-ext Show documentation

There is a newer version: 5.0.0-dev6

package com.firefly.kotlin.ext.log

import com.firefly.kotlin.ext.common.CoroutineLocal
import com.firefly.kotlin.ext.http.getAttr
import com.firefly.server.http2.router.RoutingContext
import com.firefly.utils.Assert
import com.firefly.utils.log.MappedDiagnosticContext
import java.util.concurrent.atomic.AtomicBoolean

/**
 * @author Pengtao Qiu
 */
class CoroutineMappedDiagnosticContext : MappedDiagnosticContext {

    val tracingIdKey = "_coroutineMDCKey"

    private var requestCtx: CoroutineLocal? = null
    private var initialized = AtomicBoolean(false)

    init {
        println("new CoroutineMappedDiagnosticContext")
    }

    fun setRequestCtx(reqCtx: CoroutineLocal) {
        Assert.notNull(reqCtx, "the request ctx must be not null")
        if (initialized.compareAndSet(false, true)) {
            requestCtx = reqCtx
        }
    }

    fun getContextMap(): MutableMap? {
        return requestCtx?.get()?.getAttr(tracingIdKey)
    }

    override fun clear() {
        getContextMap()?.clear()
    }

    override fun getCopyOfContextMap(): MutableMap {
        val map = getContextMap()
        return map?.toMutableMap() ?: mutableMapOf()
    }

    override fun put(key: String, value: String) {
        var map = getContextMap()
        if (map == null) {
            map = mutableMapOf()
            setContextMap(map)
        }
        map[key] = value
    }

    override fun setContextMap(contextMap: MutableMap) {
        requestCtx?.get()?.setAttribute(tracingIdKey, contextMap)
    }

    override fun remove(key: String) {
        getContextMap()?.remove(key)
    }

    override fun get(key: String): String? {
        return getContextMap()?.get(key)
    }

    override fun getKeys(): MutableSet? {
        return getContextMap()?.keys
    }

}