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.CoroutineLocalContext
import com.firefly.utils.log.MappedDiagnosticContext
import java.util.concurrent.ConcurrentHashMap

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

    val tracingIdKey = "_coroutineMDCKey"

    init {
        println("new CoroutineMappedDiagnosticContext")
    }

    private fun getContextMap(): MutableMap? {
        return CoroutineLocalContext.computeIfAbsent(tracingIdKey) { ConcurrentHashMap() }
    }

    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) {
        if (contextMap is ConcurrentHashMap) {
            CoroutineLocalContext.getAttributes()?.put(tracingIdKey, contextMap)
        } else {
            CoroutineLocalContext.getAttributes()?.put(tracingIdKey, ConcurrentHashMap(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
    }

}