org.slf4j.TrpcMDCAdapter Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of trpc-core Show documentation
The core module of tencent trpc Java project
There is a newer version: 1.3.1
/*
 * Tencent is pleased to support the open source community by making tRPC available.
 *
 * Copyright (C) 2023 THL A29 Limited, a Tencent company.
 * All rights reserved.
 *
 * If you have downloaded a copy of the tRPC source code from Tencent,
 * please note that tRPC source code is licensed under the Apache 2.0 License,
 * A copy of the Apache 2.0 License can be found in the LICENSE file.
 */

package org.slf4j;

import com.alibaba.ttl.TransmittableThreadLocal;
import com.tencent.trpc.core.utils.StringUtils;
import java.util.HashMap;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;
import org.slf4j.spi.MDCAdapter;

/**
 * TrpcMDCAdapter
 *
 * trpc internal cross-thread context tracking.
 *
 * Due to the implementation of the MDC class, the MDC.mdcAdapter does not provide read and write interfaces, so
 * the package is org.slf4j.
 *
 * @link https://github.com/alibaba/transmittable-thread-local
 */
public class TrpcMDCAdapter implements MDCAdapter {

    private static final int WRITE_OPERATION = 1;

    private static final int READ_OPERATION = 2;

    private static final TrpcMDCAdapter mtcMDCAdapter;

    private final ThreadLocal> copyOnInheritThreadLocal = new TransmittableThreadLocal<>();

    /**
     * Keeps track of the last operation performed
     */
    private final ThreadLocal lastOperation = new ThreadLocal<>();

    static {
        mtcMDCAdapter = new TrpcMDCAdapter();
        MDC.mdcAdapter = mtcMDCAdapter;
    }

    public static MDCAdapter init() {
        return mtcMDCAdapter;
    }

    /**
     * Put a context value (the val parameter) as identified with the
     * key parameter into the current thread's context map. Note that
     * contrary to log4j, the val parameter can be null.
     * 
     * 

     * If the current thread does not have a context map it is created as a side
     * effect of this call.
     *
     * @throws NullPointerException in case the "key" parameter is null
     */
    @Override
    public void put(String key, String val) {
        Objects.requireNonNull(key, "key cannot be null");
        Map oldMap = copyOnInheritThreadLocal.get();
        Integer lastOp = getAndSetLastOperation();
        if (wasLastOpReadOrNull(lastOp) || oldMap == null) {
            Map newMap = duplicateAndInsertNewMap(oldMap);
            newMap.put(key, val);
        } else {
            oldMap.put(key, val);
        }
    }

    /**
     * Remove the context identified by the key parameter.
     * 

     */
    @Override
    public void remove(String key) {
        Objects.requireNonNull(key, "key cannot be null");
        Map oldMap = copyOnInheritThreadLocal.get();
        if (oldMap == null) {
            return;
        }
        Integer lastOp = getAndSetLastOperation();
        if (wasLastOpReadOrNull(lastOp)) {
            Map newMap = duplicateAndInsertNewMap(oldMap);
            newMap.remove(key);
        } else {
            oldMap.remove(key);
        }
    }


    /**
     * Clear all entries in the MDC.
     */
    @Override
    public void clear() {
        lastOperation.set(WRITE_OPERATION);
        copyOnInheritThreadLocal.remove();
    }

    /**
     * Get the context identified by the key parameter.
     * 
     */
    @Override
    public String get(String key) {
        if (StringUtils.isEmpty(key)) {
            return null;
        }
        return Optional.ofNullable(getPropertyMap())
                .map(propertyMap -> propertyMap.get(key)).orElse(null);
    }

    /**
     * Get the current thread's MDC as a map. This method is intended to be used
     * internally.
     */
    public Map getPropertyMap() {
        lastOperation.set(READ_OPERATION);
        return copyOnInheritThreadLocal.get();
    }

    /**
     * Return a copy of the current thread's context map. Returned value may be
     * null.
     */
    @Override
    public Map getCopyOfContextMap() {
        lastOperation.set(READ_OPERATION);
        return Optional.ofNullable(copyOnInheritThreadLocal.get())
                .map(HashMap::new).orElse(null);
    }

    @SuppressWarnings("unchecked")
    @Override
    public void setContextMap(Map contextMap) {
        lastOperation.set(WRITE_OPERATION);
        Map newMap = new ConcurrentHashMap(contextMap);
        // the newMap replaces the old one for serialisation's sake
        copyOnInheritThreadLocal.set(newMap);
    }

    private Integer getAndSetLastOperation() {
        Integer lastOp = lastOperation.get();
        lastOperation.set(TrpcMDCAdapter.WRITE_OPERATION);
        return lastOp;
    }

    private static boolean wasLastOpReadOrNull(Integer lastOp) {
        return lastOp == null || lastOp == READ_OPERATION;
    }

    /**
     * Remove duplicate Map parameters and lock the oldMap to avoid being modified by other threads.
     *
     * @param oldMap {@code Map}
     * @return Map {@code Map}
     */
    private Map duplicateAndInsertNewMap(Map oldMap) {
        Map newMap = new ConcurrentHashMap<>();
        if (oldMap != null) {
            // we don't want the parent thread modifying oldMap while we are iterating over it
            synchronized (oldMap) {
                newMap.putAll(oldMap);
            }
        }
        copyOnInheritThreadLocal.set(newMap);
        return newMap;
    }

}