利用ThreadLocal實(shí)現(xiàn)一個(gè)上下文管理組件

更新時(shí)間：2022年10月13日 10:34:47 作者：localhost80

本文基于ThreadLocal原理，實(shí)現(xiàn)了一個(gè)上下文狀態(tài)管理組件Scope，通過開啟一個(gè)自定義的Scope，在Scope范圍內(nèi)，可以通過Scope各個(gè)方法讀寫數(shù)據(jù)，感興趣的可以了解一下

1 ThreadLocal原理

ThreadLocal主要作用就是實(shí)現(xiàn)線程間變量隔離，對(duì)于一個(gè)變量，每個(gè)線程維護(hù)一個(gè)自己的實(shí)例，防止多線程環(huán)境下的資源競爭，那ThreadLocal是如何實(shí)現(xiàn)這一特性的呢？

圖1

從上圖可知：

每個(gè)Thread對(duì)象中都包含一個(gè)ThreadLocal.ThreadLocalMap類型的threadlocals成員變量；
該map對(duì)應(yīng)的每個(gè)元素Entry對(duì)象中：key是ThreadLocal對(duì)象的弱引用，value是該threadlocal變量在當(dāng)前線程中的對(duì)應(yīng)的變量實(shí)體；
當(dāng)某一線程執(zhí)行獲取該ThreadLocal對(duì)象對(duì)應(yīng)的變量時(shí)，首先從當(dāng)前線程對(duì)象中獲取對(duì)應(yīng)的threadlocals哈希表，再以該ThreadLocal對(duì)象為key查詢哈希表中對(duì)應(yīng)的value；
由于每個(gè)線程獨(dú)占一個(gè)threadlocals哈希表，因此線程間ThreadLocal對(duì)象對(duì)應(yīng)的變量實(shí)體也是獨(dú)占的，不存在競爭問題，也就避免了多線程問題。

有人可能會(huì)問：ThreadLocalMap是Thread成員變量(非public，只有包訪問權(quán)限，Thread和Threadlocal都在java.lang 包下，Thread可以訪問ThreadLocal.ThreadLocalMap)，定義卻在ThreadLocal中，為什么要這么設(shè)計(jì)？

源碼的注釋給出了解釋：ThreadLocalMap就是維護(hù)線程本地變量設(shè)計(jì)的，就是讓使用者知道ThreadLocalMap就只做保存線程局部變量這一件事。

set() 方法

public void set(T value) {
    Thread t = Thread.currentThread();	//獲取當(dāng)前線程
    ThreadLocalMap map = getMap(t);	//從當(dāng)前線程對(duì)象中獲取threadlocals，該map保存了所用的變量實(shí)例
    if (map != null) {
        map.set(this, value);
    } else {
        createMap(t, value);	//初始threadlocals，并設(shè)置當(dāng)前變量
    }
}

ThreadLocalMap getMap(Thread t) {
    return t.threadLocals;
}

void createMap(Thread t, T firstValue) {
    t.threadLocals = new ThreadLocalMap(this, firstValue);
}

get() 方法

public T get() {
    Thread t = Thread.currentThread();
    ThreadLocalMap map = getMap(t); //從當(dāng)前線程對(duì)象中獲取threadlocals，該map保存了所用的變量實(shí)體
    if (map != null) {
        // 獲取當(dāng)前threadlocal對(duì)象對(duì)應(yīng)的變量實(shí)體
        ThreadLocalMap.Entry e = map.getEntry(this);
        if (e != null) {
            @SuppressWarnings("unchecked")
            T result = (T)e.value;
            return result;
        }
    }
    // 如果map沒有初始化，那么在這里初始化一下
    return setInitialValue();
}

withInitial()方法

由于通過 ThreadLocal 的 set() 設(shè)置的值，只會(huì)設(shè)置當(dāng)前線程對(duì)應(yīng)變量實(shí)體，無法實(shí)現(xiàn)統(tǒng)一初始化所有線程的ThreadLocal的值。ThreadLocal提供了一個(gè) withInitial() 方法實(shí)現(xiàn)這一功能：

ThreadLocal<String> initValue = ThreadLocal.withInitial(() -> "initValue");

public static <S> ThreadLocal<S> withInitial(Supplier<? extends S> supplier) {
    // 返回SuppliedThreadLocal類型對(duì)象
    return new SuppliedThreadLocal<>(supplier);
}

static final class SuppliedThreadLocal<T> extends ThreadLocal<T> {

    private final Supplier<? extends T> supplier;

    SuppliedThreadLocal(Supplier<? extends T> supplier) {
        this.supplier = Objects.requireNonNull(supplier);
    }

    @Override
    protected T initialValue() {
        // 獲取初始化值
        return supplier.get();
    }
}

ThreadLocal中的內(nèi)存泄漏問題

由圖1可知，ThreadLocal.ThreadLocalMap 對(duì)應(yīng)的Entry中，key為ThreadLocal對(duì)象的弱引用，方法執(zhí)行對(duì)應(yīng)棧幀中的ThreadLocal引用為強(qiáng)引用。當(dāng)方法執(zhí)行過程中，由于棧幀銷毀或者主動(dòng)釋放等原因，釋放了ThreadLocal對(duì)象的強(qiáng)引用，即表示該ThreadLocal對(duì)象可以被回收了。又因?yàn)?code>Entry中key為ThreadLocal對(duì)象的弱引用，所以當(dāng)jvm執(zhí)行GC操作時(shí)是能夠回收該ThreadLocal對(duì)象的。

而Entry中value對(duì)應(yīng)的是變量實(shí)體對(duì)象的強(qiáng)引用，因此釋放一個(gè)ThreadLocal對(duì)象，是無法釋放ThreadLocal.ThreadLocalMap中對(duì)應(yīng)的value對(duì)象的，也就造成了內(nèi)存泄漏。除非釋放當(dāng)前線程對(duì)象，這樣整個(gè)threadlocals都被回收了。但是日常開發(fā)中會(huì)經(jīng)常使用線程池等線程池化技術(shù)，釋放線程對(duì)象的條件往往無法達(dá)到。

JDK處理的方法是，在ThreadLocalMap進(jìn)行set()、get()、remove()的時(shí)候，都會(huì)進(jìn)行清理：

private Entry getEntry(ThreadLocal<?> key) {
    int i = key.threadLocalHashCode & (table.length - 1);
    Entry e = table[i];
    if (e != null && e.get() == key)
        return e;
    else
        return getEntryAfterMiss(key, i, e);
}

private Entry getEntryAfterMiss(ThreadLocal<?> key, int i, Entry e) {
    Entry[] tab = table;
    int len = tab.length;

    while (e != null) {
        ThreadLocal<?> k = e.get();
        if (k == key)
            return e;
        if (k == null)
            //如果key為null，對(duì)應(yīng)的threadlocal對(duì)象已經(jīng)被回收，清理該Entry
            expungeStaleEntry(i);
        else
            i = nextIndex(i, len);
        e = tab[i];
    }
    return null;
}

2 自定義上下文Scope

在工作中，我們經(jīng)常需要維護(hù)一些上下文，這樣可以避免在方法調(diào)用過程中傳入過多的參數(shù)，需要查詢/修改一些數(shù)據(jù)的時(shí)候，直接在當(dāng)前上下文中操作就行了。舉個(gè)具體點(diǎn)的例子：當(dāng)web服務(wù)器收到一個(gè)請求時(shí)，需要解析當(dāng)前登錄態(tài)的用戶，在后續(xù)的業(yè)務(wù)執(zhí)行流程中都需要這個(gè)用戶名。

如果只需要維護(hù)一個(gè)上下文狀態(tài)數(shù)據(jù)還比較好處理，可以通過方法傳參的形式，執(zhí)行每個(gè)業(yè)務(wù)方法的時(shí)候都通過添加一個(gè)表示用戶名方法參數(shù)傳遞進(jìn)去，但是如果需要維護(hù)上下文狀態(tài)數(shù)據(jù)比較多的話，這個(gè)方式就不太優(yōu)雅了。

一個(gè)可行的方案是通過Threadlocal實(shí)現(xiàn)請求線程的上下文，只要是同一線程的執(zhí)行過程，不同方法間不傳遞上下文狀態(tài)變量，直接操作ThreadLocal對(duì)象實(shí)現(xiàn)狀態(tài)數(shù)據(jù)的讀寫。當(dāng)存在多個(gè)上下文狀態(tài)的話，則需要維護(hù)多個(gè)ThreadLocal，似乎也可以勉強(qiáng)接受。但是當(dāng)遇到業(yè)務(wù)流程中使用線程池的情況下，從Tomcat傳遞這些ThreadLocal到線程池中的線程里就變的比較麻煩了。

基于以上考慮，下面介紹一種基于Threadlocal實(shí)現(xiàn)的上下文管理組件Scope：

Scope.java

public class Scope {

    // 靜態(tài)變量，維護(hù)不同線程的上下文Scope
    private static final ThreadLocal<Scope> SCOPE_THREAD_LOCAL = new ThreadLocal<>();

    // 實(shí)例變量，維護(hù)每個(gè)上下文中所有的狀態(tài)數(shù)據(jù)，為了區(qū)分不同的狀態(tài)數(shù)據(jù)，使用ScopeKey類型的實(shí)例作為key
    private final ConcurrentMap<ScopeKey<?>, Object> values = new ConcurrentHashMap<>();

    // 獲取當(dāng)前上下文
    public static Scope getCurrentScope() {
        return SCOPE_THREAD_LOCAL.get();
    }

    // 在當(dāng)前上下文設(shè)置一個(gè)狀態(tài)數(shù)據(jù)
    public <T> void set(ScopeKey<T> key, T value) {
        if (value != null) {
            values.put(key, value);
        } else {
            values.remove(key);
        }
    }

    // 在當(dāng)前上下文讀取一個(gè)狀態(tài)數(shù)據(jù)
    public <T> T get(ScopeKey<T> key) {
        T value = (T) values.get(key);
        if (value == null && key.initializer() != null) {
            value = key.initializer().get();
        }
        return value;
    }

    // 開啟一個(gè)上下文
    public static Scope beginScope() {
        Scope scope = SCOPE_THREAD_LOCAL.get();
        if (scope != null) {
            throw new IllegalStateException("start a scope in an exist scope.");
        }
        scope = new Scope();
        SCOPE_THREAD_LOCAL.set(scope);
        return scope;
    }

    // 關(guān)閉當(dāng)前上下文
    public static void endScope() {
        SCOPE_THREAD_LOCAL.remove();
    }
}

ScopeKey.java

public final class ScopeKey<T> {

    // 初始化器，參考 ThreadLocal 的 withInitial()
    private final Supplier<T> initializer;

    public ScopeKey() {
        this(null);
    }

    public ScopeKey(Supplier<T> initializer) {
        this.initializer = initializer;
    }

    // 統(tǒng)一初始化所有線程的 ScopeKey 對(duì)應(yīng)的值，參考 ThreadLocal 的 withInitial()
    public static <T> ScopeKey<T> withInitial(Supplier<T> initializer) {
        return new ScopeKey<>(initializer);
    }

    public Supplier<T> initializer() {
        return this.initializer;
    }

    // 獲取當(dāng)前上下文中 ScopeKey 對(duì)應(yīng)的變量
    public T get() {
        Scope currentScope = getCurrentScope();
        return currentScope.get(this);
    }

    // 設(shè)置當(dāng)前上下文中 ScopeKey 對(duì)應(yīng)的變量
    public boolean set(T value) {
        Scope currentScope = getCurrentScope();
        if (currentScope != null) {
            currentScope.set(this, value);
            return true;
        } else {
            return false;
        }
    }
}

使用方式

@Test
public void testScopeKey() {
    ScopeKey<String> localThreadName = new ScopeKey<>();

    // 不同線程中執(zhí)行時(shí)，開啟獨(dú)占的 Scope
    Runnable r = () -> {
        // 開啟 Scope
        Scope.beginScope();
        try {
            String currentThreadName = Thread.currentThread().getName();
            localThreadName.set(currentThreadName);
            log.info("currentThread: {}", localThreadName.get());
        } finally {
            // 關(guān)閉 Scope
            Scope.endScope();
        }
    };

    new Thread(r, "thread-1").start();
    new Thread(r, "thread-2").start();

    /** 執(zhí)行結(jié)果
     * [thread-1] INFO com.example.demo.testscope.TestScope - currentThread: thread-1
     * [thread-2] INFO com.example.demo.testscope.TestScope - currentThread: thread-2
     */
}

@Test
public void testWithInitial() {
    ScopeKey<String> initValue = ScopeKey.withInitial(() -> "initVal");

    Runnable r = () -> {
        Scope.beginScope();
        try {
            log.info("initValue: {}", initValue.get());
        } finally {
            Scope.endScope();
        }
    };

    new Thread(r, "thread-1").start();
    new Thread(r, "thread-2").start();

    /** 執(zhí)行結(jié)果
     * [thread-1] INFO com.example.demo.testscope.TestScope - initValue: initVal
     * [thread-2] INFO com.example.demo.testscope.TestScope - initValue: initVal
     */
}

上面的測試用例中在代碼中手動(dòng)開啟和關(guān)閉Scope不太優(yōu)雅，可以在Scope中添加兩個(gè)個(gè)靜態(tài)方法包裝下Runnable和Supplier接口：

public static <X extends Throwable> void runWithNewScope(@Nonnull ThrowableRunnable<X> runnable)
        throws X {
    supplyWithNewScope(() -> {
        runnable.run();
        return null;
    });
}

public static <T, X extends Throwable> T
        supplyWithNewScope(@Nonnull ThrowableSupplier<T, X> supplier) throws X {
    beginScope();
    try {
        return supplier.get();
    } finally {
        endScope();
    }
}

@FunctionalInterface
public interface ThrowableRunnable<X extends Throwable> {
    void run() throws X;
}

public interface ThrowableSupplier<T, X extends Throwable> {
    T get() throws X;
}

以新的Scope執(zhí)行，可以這樣寫：

@Test
public void testRunWithNewScope() {
    ScopeKey<String> localThreadName = new ScopeKey<>();

    ThrowableRunnable r = () -> {
        String currentThreadName = Thread.currentThread().getName();
        localThreadName.set(currentThreadName);
        log.info("currentThread: {}", localThreadName.get());
    };

    // 不同線程中執(zhí)行時(shí)，開啟獨(dú)占的 Scope
    new Thread(() -> Scope.runWithNewScope(r), "thread-1").start();
    new Thread(() -> Scope.runWithNewScope(r), "thread-2").start();

    /** 執(zhí)行結(jié)果
     * [thread-2] INFO com.example.demo.TestScope.testscope - currentThread: thread-2
     * [thread-1] INFO com.example.demo.TestScope.testscope - currentThread: thread-1
     */
}

3 在線程池中傳遞Scope

在上一節(jié)中實(shí)現(xiàn)的Scope，通過ThreadLocal實(shí)現(xiàn)了了一個(gè)自定義的上下文組件，在同一個(gè)線程中通過ScopeKey.set() / ScopeKey.get()讀寫同一個(gè)上下文中的狀態(tài)數(shù)據(jù)。

現(xiàn)在需要實(shí)現(xiàn)這樣一個(gè)功能，在一個(gè)線程執(zhí)行過程中開啟了一個(gè)Scope，隨后使用線程池執(zhí)行任務(wù)，要求在線程池中也能獲取當(dāng)前Scope中的狀態(tài)數(shù)據(jù)。典型的使用場景是：服務(wù)收到一個(gè)用戶請求，通過Scope將登陸態(tài)數(shù)據(jù)存到當(dāng)前線程的上下文中，隨后使用線程池執(zhí)行一些耗時(shí)的操作，希望線程池中的線程也能拿到Scope中的登陸態(tài)數(shù)據(jù)。

由于線程池中的線程和請求線程不是一個(gè)線程，按照目前的實(shí)現(xiàn)，線程池中的線程是無法拿到請求線程上下文中的數(shù)據(jù)的。

解決方法是，在提交runnable時(shí)，將當(dāng)前的Scope引用存到runnable對(duì)象中，當(dāng)獲得線程執(zhí)行時(shí)，將Scope替換到執(zhí)行線程中，執(zhí)行完成后，再恢復(fù)現(xiàn)場。在Scope中新增如下靜態(tài)方法：

// 以給定的上下文執(zhí)行 Runnable
public static <X extends Throwable> void runWithExistScope(Scope scope, ThrowableRunnable<X> runnable) throws X {
    supplyWithExistScope(scope, () -> {
        runnable.run();
        return null;
    });
}

// 以給定的上下文執(zhí)行 Supplier
public static <T, X extends Throwable> T supplyWithExistScope(Scope scope, ThrowableSupplier<T, X> supplier) throws X {
    // 保留現(xiàn)場
    Scope oldScope = SCOPE_THREAD_LOCAL.get();
    // 替換成外部傳入的 Scope
    SCOPE_THREAD_LOCAL.set(scope);
    try {
        return supplier.get();
    } finally {
        if (oldScope != null) {
            // 恢復(fù)線程
            SCOPE_THREAD_LOCAL.set(oldScope);
        } else {
            SCOPE_THREAD_LOCAL.remove();
        }
    }
}

實(shí)現(xiàn)支持Scope切換的自定義線程池ScopeThreadPoolExecutor：

public class ScopeThreadPoolExecutor extends ThreadPoolExecutor {

    ScopeThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime,
                            TimeUnit unit, BlockingQueue<Runnable> workQueue) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue);
    }

    public static ScopeThreadPoolExecutor newFixedThreadPool(int nThreads) {
        return new ScopeThreadPoolExecutor(nThreads, nThreads, 0L, TimeUnit.MILLISECONDS,
                new LinkedBlockingQueue<Runnable>());
    }

    /**
     * 只要override這一個(gè)方法就可以
     * 所有submit, invokeAll等方法都會(huì)代理到這里來
     */
    @Override
    public void execute(Runnable command) {
        Scope scope = getCurrentScope();
        // 提交任務(wù)時(shí)，把執(zhí)行 execute 方法的線程中的 Scope 傳進(jìn)去
        super.execute(() -> runWithExistScope(scope, command::run));
    }
}

測試下ScopeThreadPoolExecutor是否生效：

@Test
public void testScopeThreadPoolExecutor() {
    ScopeKey<String> localVariable = new ScopeKey<>();
    Scope.beginScope();

    try {
        localVariable.set("value out of thread pool");
        Runnable r = () -> log.info("localVariable in thread pool: {}", localVariable.get());

        // 使用線程池執(zhí)行，能獲取到外部Scope中的數(shù)據(jù)
        ExecutorService executor = ScopeThreadPoolExecutor.newFixedThreadPool(10);
        executor.execute(r);
        executor.submit(r);

    } finally {
        Scope.endScope();
    }

    /** 執(zhí)行結(jié)果
     * [pool-1-thread-1] INFO com.example.demo.testscope.TestScope - localVariable in thread pool: value out of thread pool
     * [pool-1-thread-2] INFO com.example.demo.testscope.TestScope - localVariable in thread pool: value out of thread pool
     */
}

@Test
public void testScopeThreadPoolExecutor2() {
    ScopeKey<String> localVariable = new ScopeKey<>();
    Scope.runWithNewScope(() -> {
        localVariable.set("value out of thread pool");
        Runnable r = () -> log.info("localVariable in thread pool: {}", localVariable.get());

        // 使用線程池執(zhí)行，能獲取到外部Scope中的數(shù)據(jù)
        ExecutorService executor = ScopeThreadPoolExecutor.newFixedThreadPool(10);
        executor.execute(r);
        executor.submit(r);
    });

    /** 執(zhí)行結(jié)果
     * [pool-1-thread-2] INFO com.example.demo.testscope.TestScope - localVariable in thread pool: value out of thread pool
     * [pool-1-thread-1] INFO com.example.demo.testscope.TestScope - localVariable in thread pool: value out of thread pool
     */
}

以上兩個(gè)測試用例，分別通過手動(dòng)開啟Scope、借助runWithNewScope工具方法自動(dòng)開啟Scope兩種方式驗(yàn)證了自定義線程池ScopeThreadPoolExecutor的Scope可傳遞性。

4 通過Filter、Scope實(shí)現(xiàn)Request上下文

接下來介紹如何通過Filter和Scope實(shí)現(xiàn)Request粒度的Scope上下文。思路是：通過注入一個(gè)攔截器，在進(jìn)入攔截器后開啟Scope作為一個(gè)請求的上下文，解析Request對(duì)象獲取獲取相關(guān)狀態(tài)信息(如登陸用戶)，并在Scope中設(shè)置，在離開攔截器時(shí)關(guān)閉Scope。

AuthScope.java

// 獲取登錄態(tài)的工具類
public class AuthScope {
    private static final ScopeKey<String> LOGIN_USER = new ScopeKey<>();

    public static String getLoginUser() {
        return LOGIN_USER.get();
    }

    public static void setLoginUser(String loginUser) {
        if (loginUser == null) {
            loginUser = "unknownUser";
        }
        LOGIN_USER.set(loginUser);
    }
}

ScopeFilter.java

@Lazy
@Order(0)
@Service("scopeFilter")
public class ScopeFilter extends OncePerRequestFilter {

    @Override
    protected String getAlreadyFilteredAttributeName() {
        return this.getClass().getName();
    }

    @Override
    protected void doFilterInternal(HttpServletRequest request, HttpServletResponse response,
                                    FilterChain filterChain) throws ServletException, IOException {
        // 開啟Scope
        beginScope();
        try {
            Cookie[] cookies = request.getCookies();
            String loginUser = "unknownUser";
            if (cookies != null) {
                for (Cookie cookie : cookies) {
                    if (cookie.getName().equals("login_user")) {
                        loginUser = cookie.getValue();
                        break;
                    }
                }
            }

            // 設(shè)置該 Request 上下文對(duì)用的登陸用戶
            AuthScope.setLoginUser(loginUser);

            filterChain.doFilter(request, response);
        } finally {
            // 關(guān)閉Scope
            endScope();
        }
    }
}

注入Filter

@Slf4j
@Configuration
public class FilterConfig {

    @Bean
    public FilterRegistrationBean<ScopeFilter> scopeFilterRegistration() {
        FilterRegistrationBean<ScopeFilter> registration = new FilterRegistrationBean<>();
        registration.setFilter(new ScopeFilter());
        registration.addUrlPatterns("/rest/*");
        registration.setOrder(0);
        log.info("scope filter registered");
        return registration;
    }
}

UserController.java

@Slf4j
@RestController
@RequestMapping("/rest")
public class UserController {

    // curl --location --request GET 'localhost:8080/rest/getLoginUser' --header 'Cookie: login_user=zhangsan'
    @GetMapping("/getLoginUser")
    public String getLoginUser() {
        return AuthScope.getLoginUser();
    }

    // curl --location --request GET 'localhost:8080/rest/getLoginUserInThreadPool' --header 'Cookie: login_user=zhangsan'
    @GetMapping("/getLoginUserInThreadPool")
    public String getLoginUserInThreadPool() {
        ScopeThreadPoolExecutor executor = ScopeThreadPoolExecutor.newFixedThreadPool(4);
        executor.execute(() -> {
            log.info("get login user in thread pool: {}", AuthScope.getLoginUser());
        });

        return AuthScope.getLoginUser();
    }
}

通過以下請求驗(yàn)證，請求線程和線程池線程是否能獲取登錄態(tài)，其中登錄態(tài)通過Cookie模擬：

curl --location --request GET 'localhost:8080/rest/getLoginUser' --header 'Cookie: login_user=zhangsan'
curl --location --request GET 'localhost:8080/rest/getLoginUserInThreadPool' --header 'Cookie: login_user=zhangsan'