如何使用ThreadLocal上下文解决查询性能问题

2023-07-21 10:10:54 作者：我一直在流浪

这篇文章主要介绍了利用ThreadLocal上下文解决查询性能问题,有两种解决方案，一种是使用ThreadLocal上下文，另一种是使用Redis缓存，需要的朋友可以参考下

问题背景：

安全事件列表中数据右键操作需要控制工单处置记录和SOAR处置记录是否置灰，置灰的逻辑是判断当前用户是否有相应的菜单权限以及当前租户是否有相应的授权权限。这个判断逻辑在每次刷新安全事件列表时都需要调用grpc接口判断租户是否授权，为了提高性能，所以禁止每次都去请求调用grpc接口，有两种解决方案，一种是使用ThreadLocal上下文，另一种是使用Redis缓存。简单介绍下如何使用ThreadLocal上下文解决性能问题。

01. 定义实体 Subscription

@Generated
public final class TenantManagement {
    @Generated
    public static final class Subscription {
        private volatile Object id;
        private volatile Object name;
        private volatile Object productId;
        private long startTime;
        private long expiryTime;
        private long createTime;
        private long renewTime;
        private volatile Object info;
        private volatile Object status;
        private volatile Object subscribeType;
    }
}

02. 定义 SubscribeInfoListContext 上下文

利用ThreadLocal上下文存储需要缓存的数据 List<TenantManagement.Subscription>

public class SubscribeInfoListContext {
    private static final ThreadLocal<List<TenantManagement.Subscription>> SUBSCRIBE_INFO_LIST_THREAD_LOCAL = new ThreadLocal<>();
    public static void set(List<TenantManagement.Subscription> subscribeInfoList) {
        SUBSCRIBE_INFO_LIST_THREAD_LOCAL.set(subscribeInfoList);
    }
    public static List<TenantManagement.Subscription> get() {
        return SUBSCRIBE_INFO_LIST_THREAD_LOCAL.get();
    }
    public static void remove() {
        SUBSCRIBE_INFO_LIST_THREAD_LOCAL.remove();
    }
    private SubscribeInfoListContext() {
    }
}

03. 定义SubscribeInfoService 接口

进行grpc接口调用获取 List<TenantManagement.Subscription> 信息

public interface SubscribeInfoService {
    /**
     * 获取授权信息
     *
     * @return 授权信息
     */
    List<TenantManagement.Subscription> getSubscribeInfoList();
}

1. 定义 SubscribeInfoServiceImpl 实现类

@CustomLog
@Service
public class SubscribeInfoServiceImpl implements SubscribeInfoService {
    @Setter(onMethod_ = @Autowired)
    private TenantManageClient tenantManageClient;
    @NotNull
    private static final List<String> SUBSCRIPTION_NAME_LIST = List.of(
            SubscriptionConstant.ORDER_SUBSCRIPTION,
            SubscriptionConstant.SOAR_SUBSCRIPTION
    );
    @Override
    public List<TenantManagement.Subscription> getSubscribeInfoList() {
        // 先从上下文SubscribeInfoListContext中获取List<TenantManagement.Subscription>
        List<TenantManagement.Subscription> subscriptions = SubscribeInfoListContext.get();
        // 如果上下文SubscribeInfoListContext中获取不到，再去调用grpc接口获取
        if (!CollectionUtil.isEmpty(subscriptions)) {
            return subscriptions;
        }
        String tenantId = Objects.requireNonNull(TenantInfoContext.getTenantInfo()).getTenantId();
        List<TenantManagement.Subscription> subscribeInfos = tenantManageClient.getSubscribeInfoByIds(tenantId, SUBSCRIPTION_NAME_LIST);
        // 获取到List<TenantManagement.Subscription> 后存入上下文SubscribeInfoListContext中
        SubscribeInfoListContext.set(subscribeInfos);
        return subscribeInfos;
    }
}

public interface SubscriptionInfoConstant {
    // 工单
    String ORDER_SUBSCRIPTION = "XDR_ORDER_MANAGEMENT";
    // SOAR
    String SOAR_SUBSCRIPTION = "XDR_SOAR_SECURITY_MANAGEMENT";
}

2. 判断是否授权SOAR和是否有SOAR菜单权限 GenerateSoarDisableFlagConverter

@NoArgsConstructor
@AllArgsConstructor
public class GenerateSoarDisableFlagConverter implements ResultConverter {
    @Getter
    @Setter
    private Supplier<LicenseInfoService> licenseInfoServiceSupplier;
    @Getter
    @Setter
    private Supplier<SessionNoRelatedService> sessionNoRelatedService;
    @Getter
    @Setter
    private Supplier<SubscribeInfoService> subscribeInfoServiceSupplier;
    private static final String SUPER_ADMIN_POLICY = "superAdmin";
    private static final String SOAR_QUERY_POLICY = "soarQuery";
    private static final String SOAR_SUBSCRIPTION = "XDR_SOAR_SECURITY_MANAGEMENT";
    @Nullable
    @Override
    public Map<?, ?> convert(@NotNull SplTranslateContext<?> splTranslateContext, @Nullable Map<?, ?> map, boolean ifUsingSplReturnEntity) {
        LinkedHashMap<Object, Object> result;
        if (map == null) {
            result = new LinkedHashMap<>();
        } else {
            result = new LinkedHashMap<>(map);
        }
        Object value = !isSoarAvailable();
        result.put(
                "soarDisableFlag",
                ifUsingSplReturnEntity ? SplArrayReturnEntity.of(value)
                        : value
        );
        return result;
    }
    private Boolean isSoarAvailable() {
        return hasSoarPolicy() && hasSubscription(SOAR_SUBSCRIPTION);
    }
    public Boolean hasSubscription(String subscriptionName) {
        LicenseInfoService licenseInfoService = licenseInfoServiceSupplier.get();
        if (licenseInfoService.ifLicenseInfoIsSaas()) {
            List<TenantManagement.Subscription> subscribeInfos = subscribeInfoServiceSupplier.get().getSubscribeInfoList();
            if (CollectionUtils.isEmpty(subscribeInfos)) {
                return false;
            }
            List<TenantManagement.Subscription> subscriptionList = subscribeInfos.stream()
                    .filter(subscriptionInfo -> subscriptionInfo.getProductID().equals(subscriptionName) && checkSubscribeStatus(subscriptionInfo.getStatus()))
                    .collect(Collectors.toList());
            if (CollectionUtils.isEmpty(subscriptionList)) {
                return false;
            }
            return true;
        }
        Collection<String> moduleCodes = licenseInfoServiceSupplier.get().getActiveModuleCodes();
        return moduleCodes.contains(subscriptionName);
    }
    private boolean checkSubscribeStatus(String status) {
        return SubscriptionStatusEnum.getActivateSubscriptionStatus().contains(status);
    }
    public Boolean hasSoarPolicy() {
        SessionContextResponseVo sessionContextNoRelated = sessionNoRelatedService.get().getSessionContextNoRelated();
        if (sessionContextNoRelated == null) {
            return false;
        }
        Set<String> policies = sessionContextNoRelated.getPolicies();
        if (policies.isEmpty()) {
            return false;
        }
        return policies.contains(SOAR_QUERY_POLICY) || policies.contains(SUPER_ADMIN_POLICY);
    }
}

3. 判断是否授权工单和是否有工单菜单权限 GenerateSoarDisableFlagConverter

@NoArgsConstructor
@AllArgsConstructor
public class GenerateOrderDisableFlagConverter implements ResultConverter {
    @Getter
    @Setter
    private Supplier<LicenseInfoService> licenseInfoServiceSupplier;
    @Getter
    @Setter
    private Supplier<SessionNoRelatedService> sessionNoRelatedServiceSupplier;
    @Getter
    @Setter
    private Supplier<SubscribeInfoService> subscribeInfoServiceSupplier;
    private static final String SUPER_ADMIN_POLICY = "superAdmin";
    private static final String ORDER_QUERY_POLICY = "safeOperationQuery";
    private static final String ORDER_SUBSCRIPTION = "XDR_ORDER_MANAGEMENT";
    @Nullable
    @Override
    public Map<?, ?> convert(@NotNull SplTranslateContext<?> splTranslateContext, @Nullable Map<?, ?> map, boolean ifUsingSplReturnEntity) {
        LinkedHashMap<Object, Object> result;
        if (map == null) {
            result = new LinkedHashMap<>();
        } else {
            result = new LinkedHashMap<>(map);
        }
        Object value = !isOrderAvailable();
        result.put(
                "orderDisableFlag",
                ifUsingSplReturnEntity ? SplArrayReturnEntity.of(value)
                        : value
        );
        return result;
    }
    private Boolean isOrderAvailable() {
        return hasOrderPolicy() && hasSubscription(ORDER_SUBSCRIPTION);
    }
    public Boolean hasSubscription(String subscriptionName) {
        LicenseInfoService licenseInfoService = licenseInfoServiceSupplier.get();
        if (licenseInfoService.ifLicenseInfoIsSaas()) {
            List<TenantManagement.Subscription> subscribeInfos = subscribeInfoServiceSupplier.get().getSubscribeInfoList();
            if (CollectionUtils.isEmpty(subscribeInfos)) {
                return false;
            }
            List<TenantManagement.Subscription> subscriptionList = subscribeInfos.stream()
                    .filter(subscriptionInfo -> subscriptionInfo.getProductID().equals(subscriptionName) && checkSubscribeStatus(subscriptionInfo.getStatus()))
                    .collect(Collectors.toList());
            if (CollectionUtils.isEmpty(subscriptionList)) {
                return false;
            }
            return true;
        }
        Collection<String> moduleCodes = licenseInfoService.getActiveModuleCodes();
        return moduleCodes.contains(subscriptionName);
    }
    private boolean checkSubscribeStatus(String status) {
        return SubscriptionStatusEnum.getActivateSubscriptionStatus().contains(status);
    }
    public Boolean hasOrderPolicy() {
        SessionContextResponseVo sessionContextNoRelated = sessionNoRelatedServiceSupplier.get().getSessionContextNoRelated();
        if (sessionContextNoRelated == null) {
            return false;
        }
        Set<String> policies = sessionContextNoRelated.getPolicies();
        if (policies.isEmpty()) {
            return false;
        }
        return policies.contains(ORDER_QUERY_POLICY) || policies.contains(SUPER_ADMIN_POLICY);
    }
}

4. 查询安全事件列表接口中添加 ResultConverter

@CustomLog
@Service
public class IncidentSplTableHandlingServiceImpl extends SplTableHandlingServiceImpl implements InitializingBean {
    @Getter
    @Setter(onMethod_ = @Autowired)
    private LicenseInfoService licenseInfoService;
    @Getter
    @Setter(onMethod_ = @Autowired)
    private SessionNoRelatedService sessionNoRelatedService;
    @Getter
    @Setter(onMethod_ = @Autowired)
    private SubscribeInfoService subscribeInfoService;
    @Override
    public void afterPropertiesSet() throws Exception {
        List<ResultConverter> originalResultConverters = this.getResultConverters();
        List<ResultConverter> resultConverters = new ArrayList<>(originalResultConverters);
        resultConverters.add(
                new GenerateOrderDisableFlagConverter(
                        this::getLicenseInfoService,
                        this::getSessionNoRelatedService,
                        this::getSubscribeInfoService
                )
        );
        resultConverters.add(
                new GenerateSoarDisableFlagConverter(
                        this::getLicenseInfoService,
                        this::getSessionNoRelatedService,
                        this::getSubscribeInfoService
                )
        );
        this.setResultConverters(resultConverters);
    }
    // 省略.....
}

04. 添加拦截器 IncidentInterceptorConfig

在请求开始和请求结束时清除上下文SubscribeInfoListContext中的数据，消除不同租户间数据的影响。

@Data
@Configuration
@CustomLog
public class IncidentInterceptorConfig implements WebMvcConfigurer {
    @Override
    public void addInterceptors( @NotNull InterceptorRegistry registry) {
        registry.addInterceptor(
                new HandlerInterceptorAdapter() {
                    // 在请求开始清除上下文SubscribeInfoListContext中的数据
                    @Override
                    public boolean preHandle(
                            @NotNull HttpServletRequest request,
                            @NotNull HttpServletResponse response,
                            @NotNull Object handler
                    ) throws Exception {
                        SubscribeInfoListContext.remove();
                        return true;
                    }
					// 在请求结束时清除上下文SubscribeInfoListContext中的数据
                    @Override
                    public void afterCompletion(
                            @NotNull HttpServletRequest request,
                            @NotNull HttpServletResponse response,
                            @NotNull Object handler,
                            @Nullable Exception ex
                    ) throws Exception {
                        SubscribeInfoListContext.remove();
                    }
                }
        );
    }
}

05. 添加 SaasThreadContextDataHolderSubscription 存储缓存信息

public interface SaasThreadContextDataHolder {
}

@Data
@AllArgsConstructor
public class SaasThreadContextDataHolderSubscription implements SaasThreadContextDataHolder {
    @Nullable
    private final List<TenantManagement.Subscription> subscriptionList;
}

06. 添加 SaasThreadContextHolderSubscriptionInfo 以便后续扩展和使用

public interface SaasThreadContextHolder<T extends SaasThreadContextDataHolder> {
    /**
     * 获取data holder类
     *
     * @return data holder类
     */
    @NotNull
    Class<T> getSaasThreadContextDataHolderClass();
    /**
     * 尝试加载SaasThreadContextDataHolder
     *
     * @param holder SaasThreadContextDataHolder
     * @return 可以加载则true, 否则false
     */
    default boolean tryLoad(@NotNull SaasThreadContextDataHolder holder) {
        if (!getSaasThreadContextDataHolderClass().isInstance(holder)) {
            return false;
        }
        // noinspection unchecked
        this.load((T) holder);
        return true;
    }
    /**
     * 加载SaasThreadContextDataHolder
     *
     * @param holder SaasThreadContextDataHolder
     */
    void load(@NotNull T holder);
    /**
     * 存档SaasThreadContextDataHolder
     *
     * @return SaasThreadContextDataHolder
     */
    @NotNull
    T save();
    /**
     * 清理SaasThreadContextDataHolder
     */
    void remove();
}

@AutoService(SaasThreadContextHolder.class)
public class SaasThreadContextHolderSubscriptionInfo implements SaasThreadContextHolder<SaasThreadContextDataHolderSubscriptionInfo> {
    @Override
    public @NotNull Class<SaasThreadContextDataHolderSubscriptionInfo> getSaasThreadContextDataHolderClass() {
        return SaasThreadContextDataHolderSubscriptionInfo.class;
    }
    @Override
    public void load(@NotNull SaasThreadContextDataHolderSubscriptionInfo holder) {
        List<TenantManagement.Subscription> subscriptionInfoList = holder.getSubscriptionInfoList();
        if(!CollectionUtils.isEmpty(subscriptionInfoList)){
            SubscriptionInfoContext.set(subscriptionInfoList);
        }
    }
    @Override
    public @NotNull SaasThreadContextDataHolderSubscriptionInfo save() {
        return new SaasThreadContextDataHolderSubscriptionInfo(
                SubscriptionInfoContext.get()
        );
    }
    @Override
    public void remove() {
        SubscriptionInfoContext.remove();
    }
}

07. 添加 SaasThreadContextUtil 工具类

public class SaasThreadContextUtil {
    @NotNull
    static List<SaasThreadContextHolder<?>> getSaasThreadContextHolders() {
        // IterableUtils.toList 方法将 ServiceLoader.load 返回的 Iterable 转换成了 List
        return (List) IterableUtils.toList(
            	// 加载所有实现了 SaasThreadContextHolder 接口的类，并将它们转换成 List 返回
            	ServiceLoader.load(SaasThreadContextHolder.class)
        );
    }
    @NotNull
    public static List<SaasThreadContextDataHolder> save() {
        List<SaasThreadContextHolder<?>> saasThreadContextHolders = getSaasThreadContextHolders();
        List<SaasThreadContextDataHolder> saasThreadContextDataHolders = new ArrayList<>(
                saasThreadContextHolders.size()
        );
        for (SaasThreadContextHolder<?> saasThreadContextHolder : saasThreadContextHolders) {
            saasThreadContextDataHolders.add(saasThreadContextHolder.save());
        }
        return saasThreadContextDataHolders;
    }
    public static void load(
            @NotNull List<SaasThreadContextDataHolder> saasThreadContextDataHolders
    ) {
        for (SaasThreadContextHolder<?> saasThreadContextHolder : getSaasThreadContextHolders()) {
            for (SaasThreadContextDataHolder saasThreadContextDataHolder : saasThreadContextDataHolders) {
                if (saasThreadContextHolder.tryLoad(saasThreadContextDataHolder)) {
                    break;
                }
            }
        }
    }
    public static void remove() {
        for (SaasThreadContextHolder<?> saasThreadContextHolder : getSaasThreadContextHolders()) {
            saasThreadContextHolder.remove();
        }
    }
    private SaasThreadContextUtil() {
    }
}

08. 使用 SaasThreadContextUtil

@Override
public FusionAlertVo countFusionAlerts(FusionAlertQo fusionAlertQo) {
    // 调用 SaasThreadContextUtil.save()
    List<SaasThreadContextDataHolder> threadContextDataHolders = SaasThreadContextUtil.save();
    CompletableFuture<Long> fusionAlertCountFuture = CompletableFuture.supplyAsync(() -> {
        try {
            // 调用 SaasThreadContextUtil.load(threadContextDataHolders);
            SaasThreadContextUtil.load(threadContextDataHolders);
            SearchRequest searchRequest = new SearchRequest();
            searchRequest.indices(SaasEsFactory.getTenantIndex(DatabaseConstants.ALERT));
            SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
            searchRequest.source(searchSourceBuilder);
            searchSourceBuilder.trackTotalHits(true);
            searchSourceBuilder.query(createBoolQueryBuilder(timeRange));
            return getSearchResponseTotalHits(fusionAlertQo, searchRequest);
        } finally {
            // 调用 SaasThreadContextUtil.remove();
            SaasThreadContextUtil.remove();
        }
    }, THREAD_POOL_EXECUTOR);
   CompletableFuture<Long> multiSourceAssociateAlertCountFuture = CompletableFuture.supplyAsync(() -> {
        try {
            // 调用 SaasThreadContextUtil.load(threadContextDataHolders);
            SaasThreadContextUtil.load(threadContextDataHolders);
            SearchRequest searchRequest = new SearchRequest();
            searchRequest.indices(SaasEsFactory.getTenantIndex(DatabaseConstants.ALERT));
            SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
            searchRequest.source(searchSourceBuilder);
            searchSourceBuilder.trackTotalHits(true);
            BoolQueryBuilder boolQueryBuilder = createBoolQueryBuilder(timeRange);
            searchSourceBuilder.query(boolQueryBuilder);
            TermQueryBuilder termQueryBuilder = QueryBuilders.termQuery("fusionAlert", true);
            boolQueryBuilder.must(termQueryBuilder);
            return getSearchResponseTotalHits(fusionAlertQo, searchRequest);
        } finally {
            // 调用 SaasThreadContextUtil.remove();
            SaasThreadContextUtil.remove();
        }
    }, THREAD_POOL_EXECUT
    // ...                                                                                            
}

到此这篇关于利用ThreadLocal上下文解决查询性能问题的文章就介绍到这了,更多相关ThreadLocal查询性能内容请搜索脚本之家以前的文章或继续浏览下面的相关文章希望大家以后多多支持脚本之家！