BJ_ZTEITS-Developers / parking-cloud-yudao

Commit 86137492b86680cdcc9d1b5ac37e749fd6903e1e

Authored by YunaiV
1 parent e585349a

fix(data-permission): 隔离并发任务的权限上下文

Inline Side-by-side
Showing 2 changed files with 168 additions and 1 deletions
yudao-framework/yudao-spring-boot-starter-biz-data-permission/src/main/java/cn/iocoder/yudao/framework/datapermission/core/aop/DataPermissionContextHolder.java
  View file @8613749
... ... @@ -15,9 +15,13 @@ public class DataPermissionContextHolder {
15 15  
16 16 /**
17 17 * 使用 List 的原因，可能存在方法的嵌套调用
  18 + *
  19 + * 1. 直接创建子线程时，不继承，使用空 List。避免线程池中懒创建的线程，永久保留创建时的 DataPermission
  20 + * 2. TtlRunnable、TtlCallable 传递时，复制一份 List。避免父子线程共享同一个 List，导致子线程的入栈、出栈影响父线程
18 21 */
19 22 private static final ThreadLocal<LinkedList<DataPermission>> DATA_PERMISSIONS =
20   - TransmittableThreadLocal.withInitial(LinkedList::new);
  23 + TransmittableThreadLocal.withInitialAndCopier(LinkedList::new,
  24 + parentValue -> new LinkedList<>(), parentValue -> new LinkedList<>(parentValue));
21 25  
22 26 /**
23 27 * 获得当前的 DataPermission 注解
... ...
yudao-framework/yudao-spring-boot-starter-biz-data-permission/src/test/java/cn/iocoder/yudao/framework/datapermission/core/aop/DataPermissionContextHolderTest.java
  View file @8613749
1 1 package cn.iocoder.yudao.framework.datapermission.core.aop;
2 2  
3 3 import cn.iocoder.yudao.framework.datapermission.core.annotation.DataPermission;
  4 +import com.alibaba.ttl.TtlRunnable;
4 5 import org.junit.jupiter.api.BeforeEach;
5 6 import org.junit.jupiter.api.Test;
6 7  
  8 +import java.util.ArrayList;
  9 +import java.util.Arrays;
  10 +import java.util.List;
  11 +import java.util.concurrent.Callable;
  12 +import java.util.concurrent.CountDownLatch;
  13 +import java.util.concurrent.ExecutorService;
  14 +import java.util.concurrent.Executors;
  15 +import java.util.concurrent.Future;
  16 +import java.util.concurrent.FutureTask;
  17 +import java.util.concurrent.TimeUnit;
  18 +import java.util.concurrent.atomic.AtomicReference;
  19 +
  20 +import static java.util.Collections.singletonList;
7 21 import static org.junit.jupiter.api.Assertions.assertEquals;
  22 +import static org.junit.jupiter.api.Assertions.assertNotSame;
  23 +import static org.junit.jupiter.api.Assertions.assertNull;
8 24 import static org.junit.jupiter.api.Assertions.assertSame;
  25 +import static org.junit.jupiter.api.Assertions.assertTrue;
9 26 import static org.mockito.Mockito.mock;
10 27  
11 28 /**
... ... @@ -15,6 +32,11 @@ import static org.mockito.Mockito.mock;
15 32 */
16 33 class DataPermissionContextHolderTest {
17 34  
  35 + private static final long TIMEOUT_SECONDS = 5;
  36 +
  37 + private static final DataPermission ENABLE = EnableClass.class.getAnnotation(DataPermission.class);
  38 + private static final DataPermission DISABLE = DisableClass.class.getAnnotation(DataPermission.class);
  39 +
18 40 @BeforeEach
19 41 public void setUp() {
20 42 DataPermissionContextHolder.clear();
... ... @@ -63,4 +85,145 @@ class DataPermissionContextHolderTest {
63 85 assertEquals(1, DataPermissionContextHolder.getAll().size());
64 86 }
65 87  
  88 + @Test
  89 + public void testTtlRunnable_preStartedExecutor() throws Exception {
  90 + ExecutorService executor = Executors.newSingleThreadExecutor();
  91 + CountDownLatch entered = new CountDownLatch(1);
  92 + CountDownLatch release = new CountDownLatch(1);
  93 + try {
  94 + // 准备参数：预启动线程池的线程，此时父线程还没有 DataPermission，不会通过 InheritableThreadLocal 继承
  95 + executor.submit(() -> {}).get(TIMEOUT_SECONDS, TimeUnit.SECONDS);
  96 + DataPermissionContextHolder.add(ENABLE);
  97 + AtomicReference<List<DataPermission>> childStack = new AtomicReference<>();
  98 +
  99 + // 调用：子任务入栈 enable = false 后暂停
  100 + Future<?> future = executor.submit(TtlRunnable.get(() -> {
  101 + DataPermissionContextHolder.add(DISABLE);
  102 + childStack.set(DataPermissionContextHolder.getAll());
  103 + entered.countDown();
  104 + try {
  105 + await(release);
  106 + } finally {
  107 + DataPermissionContextHolder.remove();
  108 + }
  109 + }));
  110 + await(entered);
  111 + // 断言：子任务继承父线程的 DataPermission，但父线程不受子任务入栈影响
  112 + assertNotSame(DataPermissionContextHolder.getAll(), childStack.get());
  113 + assertEquals(Arrays.asList(ENABLE, DISABLE), childStack.get());
  114 + assertEquals(singletonList(ENABLE), DataPermissionContextHolder.getAll());
  115 + assertTrue(DataPermissionContextHolder.get().enable());
  116 +
  117 + // 调用：子任务出栈并结束
  118 + release.countDown();
  119 + future.get(TIMEOUT_SECONDS, TimeUnit.SECONDS);
  120 + // 断言：父线程不受子任务出栈影响
  121 + assertEquals(singletonList(ENABLE), DataPermissionContextHolder.getAll());
  122 + // 断言：线程池的线程，执行完后恢复为原上下文
  123 + assertNull(executor.submit((Callable<DataPermission>) DataPermissionContextHolder::get)
  124 + .get(TIMEOUT_SECONDS, TimeUnit.SECONDS));
  125 + } finally {
  126 + release.countDown();
  127 + executor.shutdownNow();
  128 + }
  129 + }
  130 +
  131 + @Test
  132 + public void testTtlRunnable_captureSnapshot() throws Exception {
  133 + ExecutorService executor = Executors.newSingleThreadExecutor();
  134 + try {
  135 + // 准备参数：预启动线程池的线程
  136 + executor.submit(() -> {}).get(TIMEOUT_SECONDS, TimeUnit.SECONDS);
  137 + DataPermissionContextHolder.add(ENABLE);
  138 + AtomicReference<List<DataPermission>> childValues = new AtomicReference<>();
  139 + Runnable task = TtlRunnable.get(() -> childValues.set(new ArrayList<>(DataPermissionContextHolder.getAll())));
  140 +
  141 + // 调用：捕获后，父线程嵌套入栈，再执行子任务
  142 + DataPermissionContextHolder.add(DISABLE);
  143 + executor.submit(task).get(TIMEOUT_SECONDS, TimeUnit.SECONDS);
  144 + // 断言：子任务使用捕获时的 DataPermission，不受父线程后续入栈影响
  145 + assertEquals(singletonList(ENABLE), childValues.get());
  146 +
  147 + // 调用 + 断言：父线程嵌套出栈，出栈至空后清理上下文
  148 + assertSame(DISABLE, DataPermissionContextHolder.remove());
  149 + assertSame(ENABLE, DataPermissionContextHolder.remove());
  150 + assertTrue(DataPermissionContextHolder.getAll().isEmpty());
  151 + } finally {
  152 + executor.shutdownNow();
  153 + }
  154 + }
  155 +
  156 + @Test
  157 + public void testChildThread() throws Exception {
  158 + // 准备参数
  159 + DataPermissionContextHolder.add(ENABLE);
  160 + CountDownLatch entered = new CountDownLatch(1);
  161 + CountDownLatch release = new CountDownLatch(1);
  162 + AtomicReference<List<DataPermission>> childStack = new AtomicReference<>();
  163 + FutureTask<Void> task = new FutureTask<>(() -> {
  164 + // 断言：直接创建的子线程，不继承父线程的 DataPermission，即默认开启数据权限
  165 + assertNull(DataPermissionContextHolder.get());
  166 + DataPermissionContextHolder.add(DISABLE);
  167 + childStack.set(DataPermissionContextHolder.getAll());
  168 + entered.countDown();
  169 + await(release);
  170 + // 出栈至空，清理子线程的上下文
  171 + assertSame(DISABLE, DataPermissionContextHolder.remove());
  172 + assertNull(DataPermissionContextHolder.get());
  173 + return null;
  174 + });
  175 +
  176 + // 调用：直接创建子线程，子线程入栈 enable = false 后暂停
  177 + new Thread(task).start();
  178 + try {
  179 + await(entered);
  180 + // 断言：父线程不受子线程入栈影响
  181 + assertNotSame(DataPermissionContextHolder.getAll(), childStack.get());
  182 + assertEquals(singletonList(ENABLE), DataPermissionContextHolder.getAll());
  183 + assertTrue(DataPermissionContextHolder.get().enable());
  184 + } finally {
  185 + release.countDown();
  186 + }
  187 + task.get(TIMEOUT_SECONDS, TimeUnit.SECONDS);
  188 + // 断言：父线程不受子线程出栈、清理影响
  189 + assertEquals(singletonList(ENABLE), DataPermissionContextHolder.getAll());
  190 + assertTrue(DataPermissionContextHolder.get().enable());
  191 + }
  192 +
  193 + @Test
  194 + public void testUnmanagedWorker_notRetainParent() throws Exception {
  195 + ExecutorService executor = Executors.newSingleThreadExecutor();
  196 + try {
  197 + // 准备参数：在 enable = false 的作用域内，懒创建线程池的线程，且任务没有使用 TtlRunnable 包装
  198 + DataPermissionContextHolder.add(DISABLE);
  199 + assertNull(executor.submit((Callable<DataPermission>) DataPermissionContextHolder::get)
  200 + .get(TIMEOUT_SECONDS, TimeUnit.SECONDS));
  201 + // 调用：父线程退出作用域
  202 + assertSame(DISABLE, DataPermissionContextHolder.remove());
  203 +
  204 + // 断言：线程池的线程，没有永久保留创建时的 enable = false，仍为默认开启数据权限
  205 + assertNull(executor.submit((Callable<DataPermission>) DataPermissionContextHolder::get)
  206 + .get(TIMEOUT_SECONDS, TimeUnit.SECONDS));
  207 + } finally {
  208 + executor.shutdownNow();
  209 + }
  210 + }
  211 +
  212 + private static void await(CountDownLatch latch) {
  213 + try {
  214 + assertTrue(latch.await(TIMEOUT_SECONDS, TimeUnit.SECONDS));
  215 + } catch (InterruptedException e) {
  216 + Thread.currentThread().interrupt();
  217 + throw new IllegalStateException(e);
  218 + }
  219 + }
  220 +
  221 + @DataPermission
  222 + private static class EnableClass {
  223 + }
  224 +
  225 + @DataPermission(enable = false)
  226 + private static class DisableClass {
  227 + }
  228 +
66 229 }
... ...