From 280af80e18d1f18fbe7302c9743094ac5af6d848 Mon Sep 17 00:00:00 2001 From: CaptainB Date: Mon, 30 Oct 2023 18:20:44 +0800 Subject: [PATCH] =?UTF-8?q?build:=20=E5=8F=91=E9=80=81=E5=88=B0task-runner?= =?UTF-8?q?=E7=9A=84=E8=AF=B7=E6=B1=82=E9=83=BD=E4=BB=8Etask-runner-client?= =?UTF-8?q?=E4=B8=AD=E5=8F=91=E5=87=BA?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../service/NodeResourcePoolService.java | 33 +------ .../system/utils/TaskRunnerClient.java | 91 +++++++++++++++++++ 2 files changed, 96 insertions(+), 28 deletions(-) create mode 100644 backend/services/system-setting/src/main/java/io/metersphere/system/utils/TaskRunnerClient.java diff --git a/backend/services/system-setting/src/main/java/io/metersphere/system/service/NodeResourcePoolService.java b/backend/services/system-setting/src/main/java/io/metersphere/system/service/NodeResourcePoolService.java index 3d3980eec0..a1e645053a 100644 --- a/backend/services/system-setting/src/main/java/io/metersphere/system/service/NodeResourcePoolService.java +++ b/backend/services/system-setting/src/main/java/io/metersphere/system/service/NodeResourcePoolService.java @@ -1,45 +1,27 @@ package io.metersphere.system.service; -import com.bastiaanjansen.otp.TOTPGenerator; -import io.metersphere.sdk.constants.MsHttpHeaders; -import io.metersphere.system.controller.handler.ResultHolder; -import io.metersphere.system.dto.pool.TestResourceDTO; -import io.metersphere.system.dto.pool.TestResourceNodeDTO; import io.metersphere.sdk.exception.MSException; import io.metersphere.sdk.util.LogUtils; import io.metersphere.sdk.util.Translator; -import jakarta.annotation.Resource; +import io.metersphere.system.controller.handler.ResultHolder; +import io.metersphere.system.dto.pool.TestResourceDTO; +import io.metersphere.system.dto.pool.TestResourceNodeDTO; +import io.metersphere.system.utils.TaskRunnerClient; import org.apache.commons.collections.CollectionUtils; import org.apache.commons.lang3.StringUtils; import org.apache.commons.lang3.tuple.ImmutablePair; -import org.springframework.http.HttpEntity; -import org.springframework.http.HttpHeaders; -import org.springframework.http.HttpMethod; -import org.springframework.http.ResponseEntity; -import org.springframework.http.client.HttpComponentsClientHttpRequestFactory; import org.springframework.stereotype.Service; import org.springframework.transaction.annotation.Transactional; -import org.springframework.web.client.RestTemplate; import java.util.List; @Service @Transactional(rollbackFor = Exception.class) public class NodeResourcePoolService { - @Resource - private TOTPGenerator totpGenerator; private final static String nodeControllerUrl = "http://%s:%s/status"; - private static final RestTemplate restTemplateWithTimeOut = new RestTemplate(); - - static { - HttpComponentsClientHttpRequestFactory httpRequestFactory = new HttpComponentsClientHttpRequestFactory(); - httpRequestFactory.setConnectionRequestTimeout(2000); - httpRequestFactory.setConnectTimeout(2000); - restTemplateWithTimeOut.setRequestFactory(httpRequestFactory); - } public boolean validate(TestResourceDTO testResourceDTO, Boolean usedApiType) { List nodesList = testResourceDTO.getNodesList(); @@ -82,12 +64,7 @@ public class NodeResourcePoolService { private boolean validateNode(TestResourceNodeDTO node) { try { - String token = totpGenerator.now(); - HttpHeaders headers = new HttpHeaders(); - headers.add(MsHttpHeaders.OTP_TOKEN, token); - HttpEntity httpEntity = new HttpEntity<>(headers); - ResponseEntity entity = restTemplateWithTimeOut.exchange(String.format(nodeControllerUrl, node.getIp(), node.getPort()), HttpMethod.GET, httpEntity, ResultHolder.class); - ResultHolder body = entity.getBody(); + ResultHolder body = TaskRunnerClient.get(String.format(nodeControllerUrl, node.getIp(), node.getPort())); if (body == null) { return false; } diff --git a/backend/services/system-setting/src/main/java/io/metersphere/system/utils/TaskRunnerClient.java b/backend/services/system-setting/src/main/java/io/metersphere/system/utils/TaskRunnerClient.java new file mode 100644 index 0000000000..b56940e1d1 --- /dev/null +++ b/backend/services/system-setting/src/main/java/io/metersphere/system/utils/TaskRunnerClient.java @@ -0,0 +1,91 @@ +package io.metersphere.system.utils; + +import com.bastiaanjansen.otp.TOTPGenerator; +import io.metersphere.sdk.constants.MsHttpHeaders; +import io.metersphere.system.controller.handler.ResultHolder; +import io.metersphere.system.controller.handler.result.MsHttpResultCode; +import jakarta.annotation.Resource; +import org.springframework.http.HttpEntity; +import org.springframework.http.HttpHeaders; +import org.springframework.http.HttpMethod; +import org.springframework.http.ResponseEntity; +import org.springframework.http.client.HttpComponentsClientHttpRequestFactory; +import org.springframework.stereotype.Component; +import org.springframework.web.client.RestTemplate; + +import java.util.concurrent.TimeUnit; + +@Component +public class TaskRunnerClient { + private static TOTPGenerator totpGenerator; + + private static final RestTemplate restTemplateWithTimeOut = new RestTemplate(); + private static final int retryCount = 3; + + static { + HttpComponentsClientHttpRequestFactory httpRequestFactory = new HttpComponentsClientHttpRequestFactory(); + httpRequestFactory.setConnectionRequestTimeout(2000); + httpRequestFactory.setConnectTimeout(2000); + restTemplateWithTimeOut.setRequestFactory(httpRequestFactory); + } + + + public static ResultHolder get(String url) throws Exception { + // 定义action + Action action = (u, body) -> { + String token = totpGenerator.now(); + HttpHeaders headers = new HttpHeaders(); + headers.add(MsHttpHeaders.OTP_TOKEN, token); + HttpEntity httpEntity = new HttpEntity<>(headers); + ResponseEntity entity = restTemplateWithTimeOut.exchange(u, HttpMethod.GET, httpEntity, ResultHolder.class); + return entity.getBody(); + }; + // 首次调用 + ResultHolder body = action.execute(url, null); + return retry(url, body, action); + } + + public static ResultHolder post(String url, Object requestBody) throws Exception { + // 定义action + Action action = (u, b) -> { + String token = totpGenerator.now(); + HttpHeaders headers = new HttpHeaders(); + headers.add(MsHttpHeaders.OTP_TOKEN, token); + HttpEntity httpEntity = new HttpEntity<>(b, headers); + ResponseEntity entity = restTemplateWithTimeOut.exchange(u, HttpMethod.POST, httpEntity, ResultHolder.class); + restTemplateWithTimeOut.postForEntity(url, httpEntity, ResultHolder.class); + return entity.getBody(); + }; + // 首次调用 + ResultHolder body = action.execute(url, requestBody); + return retry(url, body, action); + } + + private static ResultHolder retry(String url, ResultHolder body, Action action) throws Exception { + if (body != null && body.getCode() == MsHttpResultCode.SUCCESS.getCode()) { + return body; + } + // 增加token失败重试 + for (int i = 0; i < retryCount; i++) { + TimeUnit.MILLISECONDS.sleep(300); + + body = action.execute(url, null); + // 重试后检查是否成功 + if (body != null && body.getCode() == MsHttpResultCode.SUCCESS.getCode()) { + return body; + } + } + return body; + } + + @Resource + public void setTotpGenerator(TOTPGenerator totpGenerator) { + TaskRunnerClient.totpGenerator = totpGenerator; + } + + @FunctionalInterface + private interface Action { + ResultHolder execute(String url, Object body); + } + +}