build: 发送到task-runner的请求都从task-runner-client中发出

This commit is contained in:
CaptainB 2023-10-30 18:20:44 +08:00 committed by Craftsman
parent 69d528bb7b
commit 280af80e18
2 changed files with 96 additions and 28 deletions

View File

@ -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<TestResourceNodeDTO> 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<String> httpEntity = new HttpEntity<>(headers);
ResponseEntity<ResultHolder> 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;
}

View File

@ -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<String> httpEntity = new HttpEntity<>(headers);
ResponseEntity<ResultHolder> 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<Object> httpEntity = new HttpEntity<>(b, headers);
ResponseEntity<ResultHolder> 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);
}
}