1
<dependency>
2
    <groupId>com.alibaba</groupId>
3
    <artifactId>transmittable-thread-local</artifactId>
4
    <version>2.14.2</version>
5
</dependency>

1
package org.slf4j;
2

3

4
import com.alibaba.ttl.TransmittableThreadLocal;
5
import org.slf4j.spi.MDCAdapter;
6

7
import java.util.Collections;
8
import java.util.HashMap;
9
import java.util.Map;
10
import java.util.Set;
11

12
/**
13
 * 重写logback的LogbackMDCAdapter，用TransmittableThreadLocal替换ThreadLocal，解决异步线程traceId传递
14
 *
15
 * @author liwf106
16
 * @since 2024/1/25
17
 */
18
@SuppressWarnings("all")
19
public class TtlMDCAdapter implements MDCAdapter {
20

21
    private final ThreadLocal<Map<String, String>> copyOnInheritThreadLocal = new TransmittableThreadLocal<>();
22

23
    private static final int WRITE_OPERATION = 1;
24
    private static final int MAP_COPY_OPERATION = 2;
25

26
    static TtlMDCAdapter mtcMDCAdapter;
27

28
    static {
29
        mtcMDCAdapter = new TtlMDCAdapter();
30
        // 替换MDC的MDCAdapter
31
        MDC.mdcAdapter = mtcMDCAdapter;
32
    }
33

34
    public static MDCAdapter getInstance() {
35
        return mtcMDCAdapter;
36
    }
37

38
    final ThreadLocal<Integer> lastOperation = new ThreadLocal<Integer>();
39

40
    private Integer getAndSetLastOperation(int op) {
41
        Integer lastOp = lastOperation.get();
42
        lastOperation.set(op);
43
        return lastOp;
44
    }
45

46
    private boolean wasLastOpReadOrNull(Integer lastOp) {
47
        return lastOp == null || lastOp.intValue() == MAP_COPY_OPERATION;
48
    }
49

50
    private Map<String, String> duplicateAndInsertNewMap(Map<String, String> oldMap) {
51
        Map<String, String> newMap = Collections.synchronizedMap(new HashMap<String, String>());
52
        if (oldMap != null) {
53
            // we don't want the parent thread modifying oldMap while we are
54
            // iterating over it
55
            synchronized (oldMap) {
56
                newMap.putAll(oldMap);
57
            }
58
        }
59

60
        copyOnInheritThreadLocal.set(newMap);
61
        return newMap;
62
    }
63

64
    /**
65
     * Put a context value (the <code>val</code> parameter) as identified with the
66
     * <code>key</code> parameter into the current thread's context map. Note that
67
     * contrary to log4j, the <code>val</code> parameter can be null.
68
     * <p/>
69
     * <p/>
70
     * If the current thread does not have a context map it is created as a side
71
     * effect of this call.
72
     *
73
     * @throws IllegalArgumentException in case of the "key" parameter is null
74
     */
75
    @Override
76
    public void put(String key, String val) throws IllegalArgumentException {
77
        if (key == null) {
78
            throw new IllegalArgumentException("key cannot be null");
79
        }
80

81
        Map<String, String> oldMap = copyOnInheritThreadLocal.get();
82
        Integer lastOp = getAndSetLastOperation(WRITE_OPERATION);
83

84
        if (wasLastOpReadOrNull(lastOp) || oldMap == null) {
85
            Map<String, String> newMap = duplicateAndInsertNewMap(oldMap);
86
            newMap.put(key, val);
87
        } else {
88
            oldMap.put(key, val);
89
        }
90
    }
91

92
    /**
93
     * Remove the the context identified by the <code>key</code> parameter.
94
     * <p/>
95
     */
96
    @Override
97
    public void remove(String key) {
98
        if (key == null) {
99
            return;
100
        }
101
        Map<String, String> oldMap = copyOnInheritThreadLocal.get();
102
        if (oldMap == null)
103
            return;
104

105
        Integer lastOp = getAndSetLastOperation(WRITE_OPERATION);
106

107
        if (wasLastOpReadOrNull(lastOp)) {
108
            Map<String, String> newMap = duplicateAndInsertNewMap(oldMap);
109
            newMap.remove(key);
110
        } else {
111
            oldMap.remove(key);
112
        }
113
    }
114

115
    /**
116
     * Clear all entries in the MDC.
117
     */
118
    @Override
119
    public void clear() {
120
        lastOperation.set(WRITE_OPERATION);
121
        copyOnInheritThreadLocal.remove();
122
    }
123

124
    /**
125
     * Get the context identified by the <code>key</code> parameter.
126
     * <p/>
127
     */
128
    @Override
129
    public String get(String key) {
130
        final Map<String, String> map = copyOnInheritThreadLocal.get();
131
        if ((map != null) && (key != null)) {
132
            return map.get(key);
133
        } else {
134
            return null;
135
        }
136
    }
137

138
    /**
139
     * Get the current thread's MDC as a map. This method is intended to be used
140
     * internally.
141
     */
142
    public Map<String, String> getPropertyMap() {
143
        lastOperation.set(MAP_COPY_OPERATION);
144
        return copyOnInheritThreadLocal.get();
145
    }
146

147
    /**
148
     * Returns the keys in the MDC as a {@link Set}. The returned value can be
149
     * null.
150
     */
151
    public Set<String> getKeys() {
152
        Map<String, String> map = getPropertyMap();
153

154
        if (map != null) {
155
            return map.keySet();
156
        } else {
157
            return null;
158
        }
159
    }
160

161
    /**
162
     * Return a copy of the current thread's context map. Returned value may be
163
     * null.
164
     */
165
    @Override
166
    public Map<String, String> getCopyOfContextMap() {
167
        Map<String, String> hashMap = copyOnInheritThreadLocal.get();
168
        if (hashMap == null) {
169
            return null;
170
        } else {
171
            return new HashMap<>(hashMap);
172
        }
173
    }
174

175
    @Override
176
    public void setContextMap(Map<String, String> contextMap) {
177
        lastOperation.set(WRITE_OPERATION);
178

179
        Map<String, String> newMap = Collections.synchronizedMap(new HashMap<String, String>());
180
        newMap.putAll(contextMap);
181

182
        // the newMap replaces the old one for serialisation's sake
183
        copyOnInheritThreadLocal.set(newMap);
184
    }
185
}

1
package com.midea.escm.core.config;
2

3
import org.slf4j.TtlMDCAdapter;
4
import org.springframework.context.ApplicationContextInitializer;
5
import org.springframework.context.ConfigurableApplicationContext;
6

7
/**
8
 * 初始化自定义MDCAdapter
9
 *
10
 * @author liwf106
11
 * @since 2024/1/25
12
 */
13
@SuppressWarnings("all")
14
public class TtlMDCAdapterInitializer implements ApplicationContextInitializer<ConfigurableApplicationContext> {
15

16
    @Override
17
    public void initialize(ConfigurableApplicationContext configurableApplicationContext) {
18
        // 加载自定义的MDCAdapter
19
        TtlMDCAdapter.getInstance();
20
    }
21
}

1
public static void main(String[] args) {
2
    SpringApplication springApplication = new SpringApplication(VbpApplication.class);
3

4
    springApplication.addInitializers(new TtlMDCAdapterInitializer());
5

6
    springApplication.run(args);
7
}

1
public void addInterceptors(InterceptorRegistry registry) {
2

3
    registry.addWebRequestInterceptor(new WebRequestInterceptor() {
4

5
        @Override
6
        public void preHandle(WebRequest request) throws Exception {
7
            String cookie = request.getHeader("Cookie");
8
            TokenUtil.setToken(cookie);
9
            MDC.put("traceId", TraceContext.getTraceId());
10
        }
11

12
        @Override
13
        public void postHandle(WebRequest request, ModelMap model) throws Exception {
14
        }
15

16
        @Override
17
        public void afterCompletion(WebRequest request, Exception ex) throws Exception {
18
            TokenUtil.removeToken();
19
        }
20
    });
21
}

1
<property name="log.colorPattern"
2
          value="%date{yyyy-MM-dd HH:mm:ss} | %cyan(%-5level) | %yellow(%thread) | %red([TraceId:%X{traceId}] | %magenta(%file:%line) | %green(%logger{96}) | (%msg%n%n)"/>