租户上下文维护优化

pigx-common/pigx-common-data/src/main/java/com/pig4cloud/pigx/common/data/tenant/TenantBroker.java

@@ -0,0 +1,139 @@
+package com.pig4cloud.pigx.common.data.tenant;
+
+import lombok.experimental.UtilityClass;
+import lombok.extern.slf4j.Slf4j;
+
+import java.util.function.Supplier;
+
+/**
+ * 租户运行时代理<br/>
+ * 这是一个工具类，用于切换租户运行时，保护租户ID上下文<br/>
+ * 下面这段代码演示问题所在
+ * <pre>
+ *     void methodA(){
+ *         // 因为某些特殊原因，需要手动指定租户
+ *         TenantContextHolder.setTenantId(1);
+ *         // do something ...
+ *     }
+ *     void methodB(){
+ *         // 因为某些特殊原因，需要手动指定租户
+ *         TenantContextHolder.setTenantId(2);
+ *         methodA();
+ *         // 此时租户ID已经变成 1
+ *         // do something ...
+ *     }
+ * </pre>
+ * 嵌套设置租户ID会导致租户上下文难以维护,并且很难察觉，容易导致数据错乱。
+ * 推荐的写法：
+ * <pre>
+ *     void methodA(){
+ *         TenantBroker.RunAs(1,() -> {
+ *             // do something ...
+ *         });
+ *     }
+ *     void methodB(){
+ *         TenantBroker.RunAs(2,() -> {
+ *              methodA();
+ *             // do something ...
+ *         });
+ *     }
+ * </pre>
+ *
+ * @author CJ (jclazz@outlook.com)
+ * @date 2020/6/12
+ * @since 3.9
+ */
+@Slf4j
+@UtilityClass
+public class TenantBroker {
+
+	@FunctionalInterface
+	public interface RunAs<T> {
+		/**
+		 * 执行业务逻辑
+		 * @param tenantId
+		 * @throws Exception
+		 */
+		void run(T tenantId) throws Exception;
+	}
+
+	@FunctionalInterface
+	public interface ApplyAs<T,R> {
+		/**
+		 * 执行业务逻辑,返回一个值
+		 * @param tenantId
+		 * @return
+		 * @throws Exception
+		 */
+		R apply(T tenantId) throws Exception;
+	}
+
+	/**
+	 * 以某个租户的身份运行
+	 * @param tenant 租户ID
+	 * @param func
+	 */
+	public  void runAs(Integer tenant,RunAs<Integer> func){
+		final Integer pre = TenantContextHolder.getTenantId();
+		try {
+			log.trace("TenantBroker 切换租户{} -> {}",pre,tenant);
+			TenantContextHolder.setTenantId(tenant);
+			func.run(tenant);
+		} catch (Exception e){
+			throw new TenantBrokerExceptionWrapper(e.getMessage(),e);
+		} finally {
+			log.trace("TenantBroker 还原租户{} <- {}",pre,tenant);
+			TenantContextHolder.setTenantId(pre);
+		}
+	}
+
+	/**
+	 * 以某个租户的身份运行
+	 * @param tenant 租户ID
+	 * @param func
+	 * @param <T> 返回数据类型
+	 * @return
+	 */
+	public <T> T applyAs(Integer tenant, ApplyAs<Integer,T> func){
+		final Integer pre = TenantContextHolder.getTenantId();
+		try {
+			log.trace("TenantBroker 切换租户{} -> {}",pre,tenant);
+			TenantContextHolder.setTenantId(tenant);
+			return func.apply(tenant);
+		} catch (Exception e){
+			throw new TenantBrokerExceptionWrapper(e.getMessage(),e);
+		} finally {
+			log.trace("TenantBroker 还原租户{} <- {}",pre,tenant);
+			TenantContextHolder.setTenantId(pre);
+		}
+	}
+
+	/**
+	 * 以某个租户的身份运行
+	 * @param supplier
+	 * @param func
+	 */
+	public  void runAs(Supplier<Integer> supplier, RunAs<Integer> func){
+		runAs(supplier.get(),func);
+	}
+
+	/**
+	 * 以某个租户的身份运行
+	 * @param supplier
+	 * @param func
+	 * @param <T> 返回数据类型
+	 * @return
+	 */
+	public <T> T applyAs(Supplier<Integer> supplier, ApplyAs<Integer,T> func){
+		return applyAs(supplier.get(),func);
+	}
+
+	public static class TenantBrokerExceptionWrapper extends RuntimeException {
+		public TenantBrokerExceptionWrapper(String message, Throwable cause) {
+			super(message, cause);
+		}
+		public TenantBrokerExceptionWrapper(Throwable cause) {
+			super(cause);
+		}
+	}
+}

pigx-common/pigx-common-data/src/main/java/com/pig4cloud/pigx/common/data/tenant/TenantContextHolder.java

@@ -32,9 +32,10 @@ public class TenantContextHolder {
 
 
 	/**
-	 * TTL 设置租户ID
-	 *
+	 * TTL 设置租户ID<br/>
+	 * <b>谨慎使用此方法,避免嵌套调用。尽量使用 {@code TenantBroker} </b>
 	 * @param tenantId
+	 * @see TenantBroker
 	 */
 	public void setTenantId(Integer tenantId) {
 		THREAD_LOCAL_TENANT.set(tenantId);

pigx-upms/pigx-upms-biz/src/main/java/com/pig4cloud/pigx/admin/service/impl/SysTenantServiceImpl.java

@@ -27,7 +27,7 @@ import com.pig4cloud.pigx.admin.service.*;
 import com.pig4cloud.pigx.common.core.constant.CacheConstants;
 import com.pig4cloud.pigx.common.core.constant.CommonConstants;
 import com.pig4cloud.pigx.common.data.resolver.ParamResolver;
-import com.pig4cloud.pigx.common.data.tenant.TenantContextHolder;
+import com.pig4cloud.pigx.common.data.tenant.TenantBroker;
 import lombok.AllArgsConstructor;
 import org.springframework.beans.BeanUtils;
 import org.springframework.cache.annotation.CacheEvict;
@@ -110,73 +110,84 @@ public class SysTenantServiceImpl extends ServiceImpl<SysTenantMapper, SysTenant
 		String defaultRoleCode = ParamResolver.getStr("TENANT_DEFAULT_ROLECODE", "ROLE_ADMIN");
 		String defaultRoleName = ParamResolver.getStr("TENANT_DEFAULT_ROLENAME", "租户默认角色");
 
-		// 查询系统内置字典
-		TenantContextHolder.setTenantId(defaultId);
-		List<SysDict> dictList = new ArrayList<>(dictService.list());
-		// 查询系统内置字典项目
-		List<Integer> dictIdList = dictList.stream().map(SysDict::getId)
-				.collect(Collectors.toList());
-		List<SysDictItem> dictItemList = new ArrayList<>(dictItemService
-				.list(Wrappers.<SysDictItem>lambdaQuery()
-						.in(SysDictItem::getDictId, dictIdList)));
-		// 查询当前租户菜单
-		List<SysMenu> menuList = menuService.list();
-		// 查询客户端配置
-		List<SysOauthClientDetails> clientDetailsList = clientServices.list();
-		// 查询系统参数配置
-		List<SysPublicParam> publicParamList = paramService.list();
+		List<SysDict> dictList = new ArrayList<>(32);
+		List<Integer> dictIdList = new ArrayList<>(32);
+		List<SysDictItem> dictItemList = new ArrayList<>(64);
+		List<SysMenu> menuList = new ArrayList<>(128);
+		List<SysOauthClientDetails> clientDetailsList = new ArrayList<>(16);
+		List<SysPublicParam> publicParamList = new ArrayList<>(64);
+
+		TenantBroker.runAs(defaultId,(id) -> {
+			// 查询系统内置字典
+			dictList.addAll(dictService.list());
+			// 查询系统内置字典项目
+			dictIdList.addAll(dictList.stream().map(SysDict::getId)
+					.collect(Collectors.toList()));
+			dictItemList.addAll(dictItemService
+					.list(Wrappers.<SysDictItem>lambdaQuery()
+							.in(SysDictItem::getDictId, dictIdList)));
+			// 查询当前租户菜单
+			menuList.addAll(menuService.list());
+			// 查询客户端配置
+			clientDetailsList.addAll(clientServices.list());
+			// 查询系统参数配置
+			publicParamList.addAll(paramService.list());
+		});
+
 		// 保证插入租户为新的租户
-		TenantContextHolder.setTenantId(sysTenant.getId());
+		return TenantBroker.applyAs(sysTenant.getId(),(id -> {
 
-		// 插入部门
-		SysDept dept = new SysDept();
-		dept.setName(defaultDeptName);
-		dept.setParentId(0);
-		deptService.save(dept);
-		//维护部门关系
-		deptRelationService.insertDeptRelation(dept);
-		// 构造初始化用户
-		SysUser user = new SysUser();
-		user.setUsername(defaultUsername);
-		user.setPassword(ENCODER.encode(defaultPassword));
-		user.setDeptId(dept.getDeptId());
-		userService.save(user);
-		// 构造新角色
-		SysRole role = new SysRole();
-		role.setRoleCode(defaultRoleCode);
-		role.setRoleName(defaultRoleName);
-		roleService.save(role);
-		// 用户角色关系
-		SysUserRole userRole = new SysUserRole();
-		userRole.setUserId(user.getUserId());
-		userRole.setRoleId(role.getRoleId());
-		userRoleService.save(userRole);
-		// 插入新的菜单
-		saveTenantMenu(TreeUtil.buildTree(menuList, CommonConstants.MENU_TREE_ROOT_ID), CommonConstants.MENU_TREE_ROOT_ID);
-		List<SysMenu> newMenuList = menuService.list();
+			// 插入部门
+			SysDept dept = new SysDept();
+			dept.setName(defaultDeptName);
+			dept.setParentId(0);
+			deptService.save(dept);
+			//维护部门关系
+			deptRelationService.insertDeptRelation(dept);
+			// 构造初始化用户
+			SysUser user = new SysUser();
+			user.setUsername(defaultUsername);
+			user.setPassword(ENCODER.encode(defaultPassword));
+			user.setDeptId(dept.getDeptId());
+			userService.save(user);
+			// 构造新角色
+			SysRole role = new SysRole();
+			role.setRoleCode(defaultRoleCode);
+			role.setRoleName(defaultRoleName);
+			roleService.save(role);
+			// 用户角色关系
+			SysUserRole userRole = new SysUserRole();
+			userRole.setUserId(user.getUserId());
+			userRole.setRoleId(role.getRoleId());
+			userRoleService.save(userRole);
+			// 插入新的菜单
+			saveTenantMenu(TreeUtil.buildTree(menuList, CommonConstants.MENU_TREE_ROOT_ID), CommonConstants.MENU_TREE_ROOT_ID);
+			List<SysMenu> newMenuList = menuService.list();
 
-		// 查询全部菜单,构造角色菜单关系
-		List<SysRoleMenu> collect = newMenuList.stream().map(menu -> {
-			SysRoleMenu roleMenu = new SysRoleMenu();
-			roleMenu.setRoleId(role.getRoleId());
-			roleMenu.setMenuId(menu.getMenuId());
-			return roleMenu;
-		}).collect(Collectors.toList());
-		roleMenuService.saveBatch(collect);
-		// 插入系统字典
-		dictService.saveBatch(dictList);
-		// 处理字典项最新关联的字典ID
-		List<SysDictItem> itemList = dictList.stream()
-				.flatMap(dict -> dictItemList.stream()
-						.filter(item -> item.getType().equals(dict.getType()))
-						.peek(item -> item.setDictId(dict.getId())))
-				.collect(Collectors.toList());
+			// 查询全部菜单,构造角色菜单关系
+			List<SysRoleMenu> collect = newMenuList.stream().map(menu -> {
+				SysRoleMenu roleMenu = new SysRoleMenu();
+				roleMenu.setRoleId(role.getRoleId());
+				roleMenu.setMenuId(menu.getMenuId());
+				return roleMenu;
+			}).collect(Collectors.toList());
+			roleMenuService.saveBatch(collect);
+			// 插入系统字典
+			dictService.saveBatch(dictList);
+			// 处理字典项最新关联的字典ID
+			List<SysDictItem> itemList = dictList.stream()
+					.flatMap(dict -> dictItemList.stream()
+							.filter(item -> item.getType().equals(dict.getType()))
+							.peek(item -> item.setDictId(dict.getId())))
+					.collect(Collectors.toList());
+
+			//插入客户端
+			clientServices.saveBatch(clientDetailsList);
+			// 插入系统配置
+			paramService.saveBatch(publicParamList);
+			return dictItemService.saveBatch(itemList);
+		}));
 
-		//插入客户端
-		clientServices.saveBatch(clientDetailsList);
-		// 插入系统配置
-		paramService.saveBatch(publicParamList);
-		return dictItemService.saveBatch(itemList);
 	}
 
 	/**

pigx-visual/pigx-mp-platform/src/main/java/com/pig4cloud/pigx/mp/config/WxMpInitConfigRunner.java

@@ -4,7 +4,7 @@ import com.google.common.collect.Maps;
 import com.pig4cloud.pigx.admin.api.feign.RemoteTenantService;
 import com.pig4cloud.pigx.common.core.constant.CacheConstants;
 import com.pig4cloud.pigx.common.core.constant.SecurityConstants;
-import com.pig4cloud.pigx.common.data.tenant.TenantContextHolder;
+import com.pig4cloud.pigx.common.data.tenant.TenantBroker;
 import com.pig4cloud.pigx.mp.entity.WxAccount;
 import com.pig4cloud.pigx.mp.handler.*;
 import com.pig4cloud.pigx.mp.service.WxAccountService;
@@ -73,10 +73,7 @@ public class WxMpInitConfigRunner {
 		// 获取全部租户 遍历所有租户对应的公众号列表
 		List<WxAccount> accountList = new ArrayList<>();
 		tenantService.list(SecurityConstants.FROM_IN).getData()
-				.forEach(tenant -> {
-					TenantContextHolder.setTenantId(tenant.getId());
-					accountList.addAll(accountService.list());
-				});
+				.forEach( tenant -> TenantBroker.runAs(tenant.getId(), (id) -> accountList.addAll(accountService.list())));
 
 		mpServices = accountList.stream().map(a -> {
 			WxMpInRedisConfigStorage configStorage = new WxMpInRedisConfigStorage(redisTemplate);

pigx-visual/pigx-mp-platform/src/main/java/com/pig4cloud/pigx/mp/controller/WxPortalController.java

@@ -2,7 +2,7 @@ package com.pig4cloud.pigx.mp.controller;
 
 import cn.hutool.core.util.StrUtil;
 import com.pig4cloud.pigx.common.core.constant.SecurityConstants;
-import com.pig4cloud.pigx.common.data.tenant.TenantContextHolder;
+import com.pig4cloud.pigx.common.data.tenant.TenantBroker;
 import com.pig4cloud.pigx.common.security.annotation.Inner;
 import com.pig4cloud.pigx.mp.config.WxMpContextHolder;
 import com.pig4cloud.pigx.mp.config.WxMpInitConfigRunner;
@@ -86,9 +86,33 @@ public class WxPortalController {
 					   @RequestParam("openid") String openid,
 					   @RequestParam(name = "encrypt_type", required = false) String encType,
 					   @RequestParam(name = "msg_signature", required = false) String msgSignature) {
-		// 查询当前 appId 对应的租户 保存到TTL 中
-		Integer tenantId = WxMpInitConfigRunner.getTenants().get(appId);
-		TenantContextHolder.setTenantId(tenantId);
+		return TenantBroker.applyAs(
+				() -> WxMpInitConfigRunner.getTenants().get(appId),
+				(id) -> handleMessage(appId,requestBody,signature,timestamp,nonce,openid,encType,msgSignature)
+		);
+	}
+
+	/**
+	 * 微信消息处理
+	 *
+	 * @param appId        多公众号标志位
+	 * @param requestBody  请求报文体
+	 * @param signature    微信签名
+	 * @param encType      加签方式
+	 * @param msgSignature 微信签名
+	 * @param timestamp    时间戳
+	 * @param nonce        随机数
+	 * @return
+	 */
+	public String handleMessage(String appId,
+								String requestBody,
+								String signature,
+								String timestamp,
+								String nonce,
+								String openid,
+								String encType,
+								String msgSignature) {
+
 		WxMpContextHolder.setAppId(appId);
 
 		final WxMpService wxService = WxMpInitConfigRunner.getMpServices().get(appId);

pigx-visual/pigx-mp-platform/src/main/java/com/pig4cloud/pigx/mp/service/impl/WxAccountFansServiceImpl.java

@@ -21,7 +21,7 @@ import cn.hutool.core.util.BooleanUtil;
 import cn.hutool.core.util.StrUtil;
 import com.baomidou.mybatisplus.core.toolkit.Wrappers;
 import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
-import com.pig4cloud.pigx.common.data.tenant.TenantContextHolder;
+import com.pig4cloud.pigx.common.data.tenant.TenantBroker;
 import com.pig4cloud.pigx.mp.config.WxMpInitConfigRunner;
 import com.pig4cloud.pigx.mp.entity.WxAccount;
 import com.pig4cloud.pigx.mp.entity.WxAccountFans;
@@ -77,11 +77,7 @@ public class WxAccountFansServiceImpl extends ServiceImpl<WxAccountFansMapper, W
 		WxMpService wxMpService = WxMpInitConfigRunner.getMpServices().get(appId);
 		WxMpUserService wxMpUserService = wxMpService.getUserService();
 
-		// 获取对应公众号租户ID
-		Integer tenantId = WxMpInitConfigRunner.getTenants().get(appId);
-		TenantContextHolder.setTenantId(tenantId);
-
-		fetchUser(null, wxAccount, wxMpUserService);
+		TenantBroker.runAs(()->WxMpInitConfigRunner.getTenants().get(appId),(id)->fetchUser(null, wxAccount, wxMpUserService));
 		return Boolean.TRUE;
 	}
 

pigx-visual/pigx-pay-platform/src/main/java/com/pig4cloud/pigx/pay/config/PayConfigParmaInitRunner.java

@@ -14,7 +14,7 @@ import com.pig4cloud.pigx.admin.api.feign.RemoteTenantService;
 import com.pig4cloud.pigx.common.core.constant.CacheConstants;
 import com.pig4cloud.pigx.common.core.constant.CommonConstants;
 import com.pig4cloud.pigx.common.core.constant.SecurityConstants;
-import com.pig4cloud.pigx.common.data.tenant.TenantContextHolder;
+import com.pig4cloud.pigx.common.data.tenant.TenantBroker;
 import com.pig4cloud.pigx.pay.entity.PayChannel;
 import com.pig4cloud.pigx.pay.service.PayChannelService;
 import com.pig4cloud.pigx.pay.utils.PayChannelNameEnum;
@@ -61,13 +61,12 @@ public class PayConfigParmaInitRunner {
 
 		List<PayChannel> channelList = new ArrayList<>();
 		tenantService.list(SecurityConstants.FROM_IN).getData()
-				.forEach(tenant -> {
-					TenantContextHolder.setTenantId(tenant.getId());
+				.forEach(tenant -> TenantBroker.runAs(tenant.getId(), (id) -> {
 					List<PayChannel> payChannelList = channelService
 							.list(Wrappers.<PayChannel>lambdaQuery()
 									.eq(PayChannel::getState, CommonConstants.STATUS_NORMAL));
 					channelList.addAll(payChannelList);
-				});
+				}));
 
 		channelList.forEach(channel -> {
 			JSONObject params = JSONUtil.parseObj(channel.getParam());