spring boot @Async非同步註解上下文透傳

_否極泰來發表於2021-07-01

上一篇文章說到,之前使用了@Async註解,子執行緒無法獲取到上下文資訊,導致流量無法打到灰度,然後改成 執行緒池的方式,每次呼叫非同步呼叫的時候都手動透傳 上下文(硬編碼)解決了問題。

後面查閱了資料,找到了方案不用每次硬編碼,來上下文透傳資料了。

方案一:
繼承執行緒池,重寫相應的方法,透傳上下文。
方案二:(推薦)
執行緒池ThreadPoolTaskExecutor,有一個TaskDecorator裝飾器,實現這個介面,透傳上下文。

方案一:繼承執行緒池,重寫相應的方法,透傳上下文。

1、ThreadPoolTaskExecutor spring封裝的執行緒池

ThreadPoolTaskExecutor 執行緒池程式碼如下:

    @Bean(ExecutorConstant.simpleExecutor_3)
    public Executor asyncExecutor3() {
        MyThreadPoolTaskExecutor taskExecutor = new MyThreadPoolTaskExecutor();
        taskExecutor.setCorePoolSize(corePoolSize);
        taskExecutor.setMaxPoolSize(maxPoolSize);
        taskExecutor.setQueueCapacity(queueCapacity);
        taskExecutor.setThreadNamePrefix(threadNamePrefix_3);
        taskExecutor.initialize();
        return taskExecutor;
    }

    //------- 繼承父類 重寫對應的方法 start
    class MyCallable<T> implements Callable<T> {
        private Callable<T> task;
        private RequestAttributes context;

        public MyCallable(Callable<T> task, RequestAttributes context) {
            this.task = task;
            this.context = context;
        }

        @Override
        public T call() throws Exception {
            if (context != null) {
                RequestContextHolder.setRequestAttributes(context);
            }

            try {
                return task.call();
            } finally {
                RequestContextHolder.resetRequestAttributes();
            }
        }
    }
    class MyThreadPoolTaskExecutor extends ThreadPoolTaskExecutor{

        @Override
        public <T> Future<T> submit(Callable<T> task) {
            return super.submit(new MyCallable(task, RequestContextHolder.currentRequestAttributes()));
        }

        @Override
        public <T> ListenableFuture<T> submitListenable(Callable<T> task) {
            return super.submitListenable(new MyCallable(task, RequestContextHolder.currentRequestAttributes()));
        }
    }
    //------- 繼承父類 重寫對應的方法 end

1、MyCallable是繼承Callable,建立MyCallable物件的時候已經把Attributes物件賦值給屬性context了(建立MyCallable物件的時候因為實在當前主執行緒建立的,所以是能獲取到請求的Attributes),在執行call方法前,先執行了RequestContextHolder.setRequestAttributes(context); 【把這個MyCallable物件的屬性context 設定到setRequestAttributes中】 所以在執行具體業務時,當前執行緒(子執行緒)就能取得主執行緒的Attributes

2、MyThreadPoolTaskExecutor類是繼承了ThreadPoolTaskExecutor 重寫了submit和submitListenable方法

為什麼是重寫submit和submitListenable這兩個方法了?

@Async AOP原始碼的方法位置是在:AsyncExecutionInterceptor.invoke

doSubmit方法能看出來

無返回值呼叫的是執行緒池方法:submit()
有返回值,根據不同的返回型別也知道:

  1. 返回值型別是:Future.class 呼叫的是方法:submit()
  2. 返回值型別是:ListenableFuture.class 呼叫的方法是:submitListenable(task)
  3. 返回值型別是:CompletableFuture.class呼叫的是CompletableFuture.supplyAsync這個在非同步註解中暫時用不上的,就不考慮重寫了。
public Object invoke(final MethodInvocation invocation) throws Throwable {
		Class<?> targetClass = (invocation.getThis() != null ? AopUtils.getTargetClass(invocation.getThis()) : null);
		Method specificMethod = ClassUtils.getMostSpecificMethod(invocation.getMethod(), targetClass);
		final Method userDeclaredMethod = BridgeMethodResolver.findBridgedMethod(specificMethod);

		AsyncTaskExecutor executor = determineAsyncExecutor(userDeclaredMethod);
		if (executor == null) {
			throw new IllegalStateException(
					"No executor specified and no default executor set on AsyncExecutionInterceptor either");
		}

		Callable<Object> task = () -> {
			try {
				Object result = invocation.proceed();
				if (result instanceof Future) {
					return ((Future<?>) result).get();
				}
			}
			catch (ExecutionException ex) {
				handleError(ex.getCause(), userDeclaredMethod, invocation.getArguments());
			}
			catch (Throwable ex) {
				handleError(ex, userDeclaredMethod, invocation.getArguments());
			}
			return null;
		};

		return doSubmit(task, executor, invocation.getMethod().getReturnType());
	}

	@Nullable
	protected Object doSubmit(Callable<Object> task, AsyncTaskExecutor executor, Class<?> returnType) {
		if (CompletableFuture.class.isAssignableFrom(returnType)) {
			return CompletableFuture.supplyAsync(() -> {
				try {
					return task.call();
				}
				catch (Throwable ex) {
					throw new CompletionException(ex);
				}
			}, executor);
		}
		else if (ListenableFuture.class.isAssignableFrom(returnType)) {
			return ((AsyncListenableTaskExecutor) executor).submitListenable(task);
		}
		else if (Future.class.isAssignableFrom(returnType)) {
			return executor.submit(task);
		}
		else {
			executor.submit(task);
			return null;
		}
	}

2、ThreadPoolExecutor 原生執行緒池

ThreadPoolExecutor執行緒池程式碼如下:

//------- ThreadPoolExecutor 繼承父類 重寫對應的方法 start
    class MyRunnable implements Runnable {
        private Runnable runnable;
        private RequestAttributes context;

        public MyRunnable(Runnable runnable, RequestAttributes context) {
            this.runnable = runnable;
            this.context = context;
        }

        @Override
        public void run() {
            if (context != null) {
                RequestContextHolder.setRequestAttributes(context);
            }
            try {
                runnable.run();
            } finally {
                RequestContextHolder.resetRequestAttributes();
            }
        }
    }

    class MyThreadPoolExecutor extends ThreadPoolExecutor{
        @Override
        public void execute(Runnable command) {
            if(!(command instanceof MyRunnable)){
                command = new MyRunnable(command,RequestContextHolder.currentRequestAttributes())
            }
            super.execute(command);
        }

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

        public MyThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory) {
            super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory);
        }

        public MyThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, RejectedExecutionHandler handler) {
            super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, handler);
        }

        public MyThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory, RejectedExecutionHandler handler) {
            super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory, handler);
        }
    }
    //------- ThreadPoolExecutor 繼承父類 重寫對應的方法 end

像ThreadPoolExecutor主要重寫execute方法,在啟動新執行緒的時候先把Attributes取到放到MyRunnable物件的一個屬性中,MyRunnable在具體執行run方法的時候,把屬性Attributes賦值到子執行緒中,當run方法執行完了在把Attributes清空掉。

為什麼只要重寫了execute方法就可以了?

ThreadPoolExecutor大家都知道主要是由submit和execute方法來執行的。
看ThreadPoolExecutor類的submit具體執行方法是由父類AbstractExecutorService#submit來實現。
具體程式碼在下面貼出來了,可以看到submit實際上最後呼叫的還是execute方法,所以我們重寫execute方法就好了。

submit方法路徑及原始碼:
java.util.concurrent.AbstractExecutorService#submit(java.lang.Runnable)

    public Future<?> submit(Runnable task) {
        if (task == null) throw new NullPointerException();
        RunnableFuture<Void> ftask = newTaskFor(task, null);
        execute(ftask);
        return ftask;
    }

    /**
     * @throws RejectedExecutionException {@inheritDoc}
     * @throws NullPointerException       {@inheritDoc}
     */
    public <T> Future<T> submit(Runnable task, T result) {
        if (task == null) throw new NullPointerException();
        RunnableFuture<T> ftask = newTaskFor(task, result);
        execute(ftask);
        return ftask;
    }

    /**
     * @throws RejectedExecutionException {@inheritDoc}
     * @throws NullPointerException       {@inheritDoc}
     */
    public <T> Future<T> submit(Callable<T> task) {
        if (task == null) throw new NullPointerException();
        RunnableFuture<T> ftask = newTaskFor(task);
        execute(ftask);
        return ftask;
    }

方案二:(推薦)

ThreadPoolTaskExecutor執行緒池

實現TaskDecorator介面,把實現類設定到taskExecutor.setTaskDecorator(new MyTaskDecorator());

    //------- 實現TaskDecorator 介面 start

    @Bean(ExecutorConstant.simpleExecutor_4)
    public Executor asyncExecutor4() {
        MyThreadPoolTaskExecutor taskExecutor = new MyThreadPoolTaskExecutor();
        taskExecutor.setCorePoolSize(corePoolSize);
        taskExecutor.setMaxPoolSize(maxPoolSize);
        taskExecutor.setQueueCapacity(queueCapacity);
        taskExecutor.setThreadNamePrefix(threadNamePrefix_4);
        taskExecutor.setTaskDecorator(new MyTaskDecorator());
        taskExecutor.initialize();
        return taskExecutor;
    }

    class MyTaskDecorator implements TaskDecorator{

        @Override
        public Runnable decorate(Runnable runnable) {
            try {
                RequestAttributes attributes = RequestContextHolder.getRequestAttributes();
                return () -> {
                    try {
                        RequestContextHolder.setRequestAttributes(attributes);
                        runnable.run();
                    } finally {
                        RequestContextHolder.resetRequestAttributes();
                    }
                };
            } catch (IllegalStateException e) {
                return runnable;
            }
        }
    }
    //------- 實現TaskDecorator 介面 end

為什麼設定了setTaskDecorator就能實現透傳資料了?

主要還是看taskExecutor.initialize()方法,主要是重寫了ThreadPoolExecutor的execute方法,用裝飾器模式 增強了Runnable介面,原始碼如下:

	@Nullable
	private ThreadPoolExecutor threadPoolExecutor;

	//初始化方法
	public void initialize() {
		if (logger.isDebugEnabled()) {
			logger.debug("Initializing ExecutorService" + (this.beanName != null ? " '" + this.beanName + "'" : ""));
		}
		if (!this.threadNamePrefixSet && this.beanName != null) {
			setThreadNamePrefix(this.beanName + "-");
		}
		this.executor = initializeExecutor(this.threadFactory, this.rejectedExecutionHandler);
	}

	@Override
	protected ExecutorService initializeExecutor(
			ThreadFactory threadFactory, RejectedExecutionHandler rejectedExecutionHandler) {

		BlockingQueue<Runnable> queue = createQueue(this.queueCapacity);

		ThreadPoolExecutor executor;

		//判斷是否設定了,taskDecorator裝飾器
		if (this.taskDecorator != null) {
			executor = new ThreadPoolExecutor(
					this.corePoolSize, this.maxPoolSize, this.keepAliveSeconds, TimeUnit.SECONDS,
					queue, threadFactory, rejectedExecutionHandler) {
				@Override
				public void execute(Runnable command) {
					//執行裝飾器方法包裝Runnable介面
					Runnable decorated = taskDecorator.decorate(command);
					if (decorated != command) {
						decoratedTaskMap.put(decorated, command);
					}
					super.execute(decorated);
				}
			};
		}
		else {
			executor = new ThreadPoolExecutor(
					this.corePoolSize, this.maxPoolSize, this.keepAliveSeconds, TimeUnit.SECONDS,
					queue, threadFactory, rejectedExecutionHandler);

		}

		if (this.allowCoreThreadTimeOut) {
			executor.allowCoreThreadTimeOut(true);
		}
		//把初始化好的ThreadPoolExecutor執行緒池賦值給 當前類屬性threadPoolExecutor
		this.threadPoolExecutor = executor;
		return executor;
	}

總結

無論是方案1還是方案2,原理都是先在當前執行緒獲取到Attributes,然後把Attributes賦值到Runnable的一個屬性中,在起一個子執行緒後,具體執行run方法的時候,把Attributes設定給當子執行緒,當run方法執行完了,在清空Attributes。

方案2實現比較優雅,所以推薦使用它。

工作沒多久的時候覺得spring的使用很麻煩,但是工作久了慢慢發現spring一些小細節、設計模式運用的非常巧妙,很容易解決遇到的問題,只能說spring ??。

我已經把上述程式碼例子放到gitee了,大家感興趣可以clone 傳送門~

相關文章