aidlab/aida_back
7
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

1、完善超分功能

Browse Source 
2、完善积分系统
3、新增任务列表
This commit is contained in:
xupei
2024-03-26 14:58:43 +08:00
parent 305324fe1a
commit d411b428f8
26 changed files with 807 additions and 328 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
src/main/java/com/ai/da/common/RabbitMQ/MQConfig.java
View File

@@ -8,12 +8,16 @@ import org.springframework.context.annotation.Configuration;
public class MQConfig {
public static final String GENERATE_EXCHANGE_FANOUT = "generate-exchange";
// public static final String GENERATE_QUEUE = "generate-queue-prod";
// public static final String GENERATE_QUEUE = "generate-queue-prod";
// public static final String GENERATE_QUEUE = "generate-queue-test";
// public static final String GENERATE_QUEUE = "generate-queue-dev";
public static final String GENERATE_QUEUE = "generate-queue";
public static final String GENERATE_QUEUE = "generate-queue-dev";
// public static final String GENERATE_QUEUE = "generate-queue";
public static final String SR_QUEUE = "SR-queue-dev";
public static final String SR_QUEUE = "SR-queue-dev";
// public static final String SR_QUEUE = "SR-queue-local";
// public static final String SR_RESULT_QUEUE = "SuperResolution-local";
public static final String SR_RESULT_QUEUE = "SuperResolution-dev";
public MQConfig() {
}
@@ -36,6 +40,11 @@ public class MQConfig {
return new Queue(SR_QUEUE);
}
@Bean
public Queue SRResultQueue() {
return new Queue(SR_RESULT_QUEUE);
}
/**
* 将队列绑定到交换机上【队列订阅交换机】
*/

164
src/main/java/com/ai/da/common/RabbitMQ/SRConsumer.java
View File

@@ -2,13 +2,15 @@ package com.ai.da.common.RabbitMQ;
import com.ai.da.common.config.exception.BusinessException;
import com.ai.da.common.utils.RedisUtil;
import com.ai.da.model.dto.GenerateThroughImageTextDTO;
import com.ai.da.model.dto.SuperResolutionDTO;
import com.ai.da.model.vo.GenerateCollectionVO;
import com.ai.da.model.dto.TaskDTO;
import com.ai.da.service.SuperResolutionService;
import com.ai.da.service.TaskListService;
import com.alibaba.fastjson.JSONException;
import com.alibaba.fastjson.JSONObject;
import com.google.gson.Gson;
import com.google.gson.reflect.TypeToken;
import com.rabbitmq.client.Channel;
import io.netty.util.internal.StringUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.amqp.core.Message;
import org.springframework.amqp.rabbit.annotation.RabbitHandler;
@@ -18,6 +20,7 @@ import org.springframework.stereotype.Component;
import javax.annotation.Resource;
import java.io.IOException;
import java.lang.reflect.Type;
import java.util.HashMap;
@Slf4j
@@ -27,6 +30,9 @@ public class SRConsumer {
@Resource
private RedisUtil redisUtil;
@Resource
private TaskListService taskListService;
@Value("${redis.key.orderForSR}")
private String consumptionOrderKey;
@@ -36,21 +42,33 @@ public class SRConsumer {
@Value("${redis.key.SRExceptionMap}")
private String exceptionMapKey;
@Value("${redis.key.resultMap}")
private String resultMapKey;
@Value("${redis.key.taskList}")
private String taskListKey;
@Resource
private SuperResolutionService superResolutionService;
public void superResolution(Message msg, Channel channel, String consumerName){
log.info("============start listening==========");
/**
* 请求超分处理
*/
public void superResolution(Message msg, Channel channel, String consumerName) {
log.info("============SR start listening==========");
long start = System.currentTimeMillis();
SuperResolutionDTO superResolutionDTO = JSONObject.parseObject(msg.getBody(), SuperResolutionDTO.class);
String uniqueId = superResolutionDTO.getUniqueId();
log.info("From " + consumerName + " : " + uniqueId);
SuperResolutionDTO superResolutionDTO;
String uniqueId = null;
try {
superResolutionDTO = JSONObject.parseObject(msg.getBody(), SuperResolutionDTO.class);
uniqueId = superResolutionDTO.getUniqueId();
log.info("From " + consumerName + " : " + uniqueId);
superResolutionService.updateSROutput(uniqueId, "Executing", null);
taskListService.updateTaskStatusOrOutputRedis(uniqueId, "Executing", null);
/*try {
Thread.sleep(2 * 60 * 1000);
} catch (InterruptedException e) {
throw new RuntimeException(e);
}*/
// 2、判断当前消息是否在取消列表中
Boolean isMember = redisUtil.isElementExistsInSet(cancelSetKey, uniqueId);
if (isMember) {
@@ -61,45 +79,112 @@ public class SRConsumer {
log.error("手动确认，不返回队列重新消费");
}
} else {
/*try {
Thread.sleep(15000);
} catch (InterruptedException e) {
throw new RuntimeException(e);
}*/
String srOutput = superResolutionService.SR(superResolutionDTO);
// 将消息从redis排队队列中删除,需保证被消费的消息存储到db之后再从redis删除
redisUtil.removeFromZSet(consumptionOrderKey, uniqueId);
if (!StringUtil.isNullOrEmpty(srOutput)) {
HashMap<String, String> generateResult = new HashMap<>();
generateResult.put(uniqueId, srOutput);
// 将结果存在redis中 ,为空时不要存
redisUtil.addToMap(resultMapKey, generateResult);
}
// 请求python端进行超分
superResolutionService.SR(superResolutionDTO);
}
} catch (BusinessException e) {
log.error(e.getMsg());
superResolutionDTO = JSONObject.parseObject(msg.getBody(), SuperResolutionDTO.class);
// channel.basicNack() 为不确认deliveryTag对应的消息，第二个参数是否应用于多消息，第三个参数是否requeue
setErrorMessage(msg, channel, e.getMsg(), superResolutionDTO);
} catch (JSONException e) {
log.error(e.getMessage());
setErrorMessage(msg, channel, e.getMessage(), null);
} catch (Exception e) {
log.error(e.getMessage());
superResolutionDTO = JSONObject.parseObject(msg.getBody(), SuperResolutionDTO.class);
setErrorMessage(msg, channel, e.getMessage(), superResolutionDTO);
}
long end = System.currentTimeMillis();
log.info(" task_id： " + uniqueId + "----------" + consumerName + " 执行时长：" + (end - start) + "毫秒");
log.info("=============SR end listening===========");
}
/**
* 获取超分结果
*/
public void getSRResult(Message msg, Channel channel, String consumerName) {
log.info("============SRResult start listening==========");
long start = System.currentTimeMillis();
JSONObject result = null;
String taskId = null;
try {
result = JSONObject.parseObject(msg.getBody(), JSONObject.class);
taskId = result.get("tasks_id").toString();
} catch (JSONException e) {
log.error("SRResult 返回数据格式不合规范");
log.error(e.getMessage());
setErrorMessage(msg, channel, e.getMessage(), null);
}
try {
// channel.basicAck(msg.getMessageProperties().getDeliveryTag(), false);
// 2、判断状态是否成功
if ("SUCCESS".equals(result.get("status").toString())) {
String output = result.get("data").toString();
superResolutionService.setSRResult(taskId, output, "success");
taskListService.updateTaskStatusOrOutputRedis(taskId, "success", output);
} else {
superResolutionService.setSRResult(taskId, null, "fail");
taskListService.updateTaskStatusOrOutputRedis(taskId, "fail", null);
HashMap<String, String> exceptionInfo = new HashMap<>();
// 获取输入信息
String task = redisUtil.getFromString(taskListKey + taskId + taskId.substring(taskId.lastIndexOf("-") + 1));
Gson gson = new Gson();
Type type = new TypeToken<TaskDTO<SuperResolutionDTO>>() {
}.getType();
TaskDTO<SuperResolutionDTO> taskDTO = gson.fromJson(task, type);
// 将输入信息和报错信息均存入redis todo 加判空
exceptionInfo.put(taskId, "Input ==> " + taskDTO.getInputParam() + "Fail Message ==> " + result.get("message").toString());
// 将报错信息存入redis
redisUtil.addToMap(exceptionMapKey, exceptionInfo);
}
} catch (Exception e) {
log.error(e.getMessage());
// channel.basicNack() 为不确认deliveryTag对应的消息，第二个参数是否应用于多消息，第三个参数是否requeue
try {
// 第二个参数，是否批量确认消息，当传false时，只确认当前 deliveryTag对应的消息;当传true时，会确认当前及之前所有未确认的消息。
channel.basicAck(msg.getMessageProperties().getDeliveryTag(), false);
// 将消息从redis排队队列中删除,需保证被消费的消息存储到db之后再从redis删除
redisUtil.removeFromZSet(consumptionOrderKey, uniqueId);
} catch (IOException exception) {
log.error("手动确认，取消返回队列，不再重新消费");
}
// 将入参和错误信息存入数据库
String exceptionMessage = JSONObject.toJSONString(superResolutionDTO) +
" Exception message ： " + e.getMsg();
HashMap<String, String> exceptionInfo = new HashMap<>();
exceptionInfo.put(String.valueOf(uniqueId), exceptionMessage);
// 存redis
redisUtil.addToMap(exceptionMapKey, exceptionInfo);
}
long end = System.currentTimeMillis();
log.info(" task_id： " + uniqueId + "----------" + consumerName + " 执行时长：" + (end - start) + "毫秒");
log.info("=============end listening===========");
log.info(" task_id： " + taskId + "----------" + consumerName + " 执行时长：" + (end - start) + "毫秒");
log.info("=============SRResult end listening===========");
}
private void setErrorMessage(Message msg, Channel channel, String message, SuperResolutionDTO superResolutionDTO) {
String uniqueId;
try {
// 第二个参数，是否批量确认消息，当传false时，只确认当前 deliveryTag对应的消息;当传true时，会确认当前及之前所有未确认的消息。
channel.basicAck(msg.getMessageProperties().getDeliveryTag(), false);
uniqueId = superResolutionDTO.getUniqueId();
// 将消息从redis排队队列中删除,需保证被消费的消息存储到db之后再从redis删除
redisUtil.removeFromZSet(consumptionOrderKey, uniqueId);
} catch (Exception exception) {
log.error("手动确认，取消返回队列，不再重新消费");
throw new BusinessException("发生错误，手动确认消息");
}
// 将入参和错误信息存入redis
String exceptionMessage = JSONObject.toJSONString(superResolutionDTO) +
" Exception message ： " + message;
// " Exception message ： " + e.getMessage();
HashMap<String, String> exceptionInfo = new HashMap<>();
uniqueId = superResolutionDTO.getUniqueId();
exceptionInfo.put(String.valueOf(uniqueId), exceptionMessage);
// 存redis
redisUtil.addToMap(exceptionMapKey, exceptionInfo);
taskListService.updateTaskStatusOrOutputRedis(uniqueId, "fail", null);
}
@RabbitListener(queues = MQConfig.SR_QUEUE)
@@ -108,4 +193,11 @@ public class SRConsumer {
superResolution(msg, channel, "consumer 1");
}
@RabbitListener(queues = MQConfig.SR_RESULT_QUEUE)
@RabbitHandler
public void SRResultConsumer1(Message msg, Channel channel) {
getSRResult(msg, channel, "consumer 1");
}
}

9
src/main/java/com/ai/da/common/constant/PayPalCheckoutConstant.java
View File

@@ -2,7 +2,8 @@ package com.ai.da.common.constant;
public class PayPalCheckoutConstant {
public static String MODE = "sandbox";
// public static String MODE = "sandbox";
public static String MODE = "live";
public static final String CAPTURE = "CAPTURE";
/**
@@ -167,7 +168,11 @@ public class PayPalCheckoutConstant {
public final static String CMD_NOTIFY_VALIDATE = "_notify-validate";
public final static String WEBHOOK_ID = "31797347YC028794L";
// public final static String WEBHOOK_ID = "31797347YC028794L";
// kim-sandbox
// public final static String WEBHOOK_ID = "1WH327112B602422N";
// kim-live
public final static String WEBHOOK_ID = "41L14847MC833625B";
public final static String PAYPAL_TOKEN_KEY = "PayPalAccessToken";

9
src/main/java/com/ai/da/common/enums/CreditsEventsEnum.java
View File

@@ -7,17 +7,18 @@ import lombok.Getter;
@Getter
public enum CreditsEventsEnum {
PRICE("price","2"),
PRICE("price","1"),
// 6USD -> 1000 points ==> 10HKD -> 215 points ==> 2HKD -> 43points
BUY_CREDITS("Buy Credits","43"),
INIT("init", "1000"),
DAILY_CHECKIN("Daily Check-In", "50"),
SOCIAL_MEDIA_SHARING("Social Media Sharing","50"),
// 6USD -> 1000 points ==> 10HKD -> 215 points ==> 2HKD -> 43points
BUY_CREDITS("Buy Credits","43"),
SUPER_RESOLUTION("Super Resolution","300"),
// SUPER_RESOLUTION("Super Resolution","300"),
SUPER_RESOLUTION("Super Resolution","30"),
OTHER("Other","10");

49
src/main/java/com/ai/da/common/utils/AsyncCallerUtil.java
View File

@@ -2,7 +2,6 @@ package com.ai.da.common.utils;
import com.ai.da.common.config.exception.BusinessException;
import com.ai.da.model.dto.GenerateToPythonDTO;
import com.ai.da.model.dto.SuperResolutionDTO;
import com.ai.da.python.PythonService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
@@ -72,52 +71,4 @@ public class AsyncCallerUtil {
}
}
public CompletableFuture<String> callSRAsync(SuperResolutionDTO superResolutionDTO) {
return CompletableFuture.supplyAsync(() -> pythonService.superResolution(superResolutionDTO));
}
public String SR(SuperResolutionDTO superResolutionDTO) {
ScheduledExecutorService scheduledExecutorService = Executors.newScheduledThreadPool(5);
String taskId = superResolutionDTO.getUniqueId();
ScheduledFuture<?> timeoutTask = null;
if (!waitingStatus.containsKey(taskId)) waitingStatus.put(taskId, true);
try {
CompletableFuture<String> generateResult = callSRAsync(superResolutionDTO);
// 5秒后第一次确认，之后每隔10秒确认一次用户选择结果
timeoutTask = scheduledExecutorService.scheduleAtFixedRate(() -> {
// 调用另一个接口获取用户的选择
if (!waitingStatus.get(taskId)) {
// 如果用户选择取消，则取消对generate的调用
generateResult.cancel(true);
waitingStatus.remove(taskId);
} else log.info("===============持续等待===============");
}, 5, 10, TimeUnit.SECONDS);
log.info("阻塞等待结果...");
// 阻塞，等待结果
String result = generateResult.get();
// 取消定时任务
timeoutTask.cancel(true);
waitingStatus.remove(taskId);
return result;
} catch (CancellationException e) {
// generateResult.cancel(true);通过抛出异常取消该任务
log.info("==========成功取消generate任务==========");
return null;
} catch (InterruptedException | ExecutionException | BusinessException e) {
// 处理异常
log.error("发生错误 ： " + e, e);
// 取消定时任务
assert timeoutTask != null;
timeoutTask.cancel(true);
throw new BusinessException(e.getMessage());
} finally {
// 关闭线程池
// executorService.shutdown();
// scheduledExecutorService.shutdown();
}
}
}

20
src/main/java/com/ai/da/common/utils/RedisUtil.java
View File

@@ -7,6 +7,7 @@ import org.springframework.stereotype.Component;
import org.springframework.util.CollectionUtils;
import javax.annotation.Resource;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.TimeUnit;
@@ -69,10 +70,15 @@ public class RedisUtil {
/**
* 获取当前ZSet中数据量的总和
*/
public Long getZSetTotal(String key) {
public Long getZSetTotalCount(String key) {
return redisTemplate.opsForZSet().zCard(key);
}
public Set<String> getZSetTotalData(String key){
return redisTemplate.opsForZSet().range(key, 0, -1);
}
//- - - - - - - - - - - - - - - - - - - - - set类型 - - - - - - - - - - - - - - - - - - - -
/**
@@ -138,4 +144,16 @@ public class RedisUtil {
return redisTemplate.opsForValue().get(key);
}
public Set<String> getKeysFromString(String key){
return redisTemplate.keys(key);
}
public List<String> getMultiValue(Set<String> keys){
return redisTemplate.opsForValue().multiGet(keys);
}
public Long getExpire(String key){
return redisTemplate.getExpire(key);
}
}