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

TASK:异步调用generate及取消generate

Browse Source
This commit is contained in:
xupei
2024-01-24 11:43:56 +08:00
parent a9ce35200c
commit 96858c2cc3
10 changed files with 159 additions and 112 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
src/main/java/com/ai/da/service/GenerateService.java
View File

@@ -25,12 +25,12 @@ public interface GenerateService extends IService<Generate> {
List<GenerateDetail> selectBatchByLibraryId(List<Long> libraryId);
GenerateCollectionVO getGenerateResult(Long uniqueId);
GenerateCollectionVO getGenerateResult(String uniqueId);
Long prepareForGenerate(GenerateThroughImageTextDTO generateThroughImageTextDTO);
String prepareForGenerate(GenerateThroughImageTextDTO generateThroughImageTextDTO);
Long getRankPosition(Long uniqueId);
Long getRankPosition(String uniqueId);
void cancelGenerate(Long uniqueId);
void cancelGenerate(String uniqueId);
}

141
src/main/java/com/ai/da/service/impl/GenerateServiceImpl.java
View File

@@ -24,6 +24,7 @@ import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import io.minio.errors.MinioException;
import io.netty.util.internal.StringUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
@@ -35,6 +36,7 @@ import java.util.*;
import static com.ai.da.common.enums.CollectionLevel1TypeEnum.*;
@Slf4j
@Service
public class GenerateServiceImpl extends ServiceImpl<GenerateMapper, Generate> implements GenerateService {
@@ -77,7 +79,7 @@ public class GenerateServiceImpl extends ServiceImpl<GenerateMapper, Generate> i
@Override
public GenerateCaptionVO generateCaption(Long sketchElementId) {
CollectionElement collectionElement = collectionElementMapper.selectById(sketchElementId);
if (Objects.isNull(collectionElement)){
if (Objects.isNull(collectionElement)) {
throw new BusinessException("the.image.does.not.exist.please.reselect");
}
String url = collectionElement.getUrl();
@@ -101,14 +103,14 @@ public class GenerateServiceImpl extends ServiceImpl<GenerateMapper, Generate> i
generate.setLevel1Type(generateThroughImageTextDTO.getLevel1Type());
// 当level1type是sketchboard时，存数据库需要加上当前性别
generate.setGenerateType(generate.getLevel1Type().equals(SKETCH_BOARD.getRealName()) ?
generateType + " (" +generateThroughImageTextDTO.getGender() + ")":
generateType + " (" + generateThroughImageTextDTO.getGender() + ")" :
generateType);
generate.setModelName(StringUtil.isNullOrEmpty(generateThroughImageTextDTO.getVersion()) ? ModelNameEnum.MODEL_0.getCode() : generateThroughImageTextDTO.getVersion());
generate.setCreateDate(DateUtil.getByTimeZone(generateThroughImageTextDTO.getTimeZone()));
String text = generateThroughImageTextDTO.getText();
Long elementId = generateThroughImageTextDTO.getCollectionElementId();
validateGeneraType(generate, text, elementId,generateType);
validateGeneraType(generate, text, elementId, generateType);
// 2.1 sketch或print在t_collection_element表中的信息是否需要更新 如 level2Type
CollectionElement collectionElement = collectionElementService.editLevel2Type(elementId, generateThroughImageTextDTO.getLevel2Type());
@@ -120,11 +122,11 @@ public class GenerateServiceImpl extends ServiceImpl<GenerateMapper, Generate> i
String category = generateThroughImageTextDTO.getLevel1Type().equals(SKETCH_BOARD.getRealName()) ? "sketch" :
generateThroughImageTextDTO.getLevel1Type().equals(PRINT_BOARD.getRealName()) ? "print" : "moodboard";
AsyncCallerUtil asyncCallerUtil = new AsyncCallerUtil();
List<String> generatedSketchUrl = asyncCallerUtil.generate(new GenerateToPythonDTO(accountId, Objects.isNull(collectionElement) ? null : collectionElement.getUrl(),
category, text, mode, "1", generateThroughImageTextDTO.getGender()),generateThroughImageTextDTO.getUniqueId());
List<String> generatedSketchUrl = asyncCallerUtil.generate(new GenerateToPythonDTO(accountId, Objects.isNull(collectionElement) ? "" : collectionElement.getUrl(),
category, text, mode, "1", generateThroughImageTextDTO.getGender() ,generateThroughImageTextDTO.getUniqueId()));
// List<String> generatedSketchUrl = pythonService.generateSketchOrPrint(new GenerateToPythonDTO(accountId, Objects.isNull(elementId) ? null : collectionElement.getUrl(),
// category, text, mode, "1", generateThroughImageTextDTO.getGender()));
if (CollectionUtils.isEmpty(generatedSketchUrl)){
if (CollectionUtils.isEmpty(generatedSketchUrl)) {
return null;
}
@@ -139,8 +141,8 @@ public class GenerateServiceImpl extends ServiceImpl<GenerateMapper, Generate> i
GenerateCollectionItemVO generateCollectionItemVO = new GenerateCollectionItemVO();
String md5 = MD5Utils.encryptFile(minioUtil.getPresignedUrl(item, 24 * 60), Boolean.FALSE);
// 通过MD5值和level1Type,判断不同level1Type下相同的图片是否被like过
List<Map<String,Long>> libraryIdList = generateDetailMapper.getLibraryIdThroughMD5(md5, generateThroughImageTextDTO.getLevel1Type());
if (!libraryIdList.isEmpty()){
List<Map<String, Long>> libraryIdList = generateDetailMapper.getLibraryIdThroughMD5(md5, generateThroughImageTextDTO.getLevel1Type());
if (!libraryIdList.isEmpty()) {
generateDetail.setIsLike((byte) 1);
generateDetail.setLibraryId(libraryIdList.get(0).get("library_id"));
generateCollectionItemVO.setIsLiked(Boolean.TRUE);
@@ -161,22 +163,22 @@ public class GenerateServiceImpl extends ServiceImpl<GenerateMapper, Generate> i
return new GenerateCollectionVO(generate.getId(), collectionId, generatedCollectionItems);
}
private void validateGeneraType(Generate generate, String text, Long elementId,String generateType) {
private void validateGeneraType(Generate generate, String text, Long elementId, String generateType) {
switch (generateType) {
case "text":
if (StringUtil.isNullOrEmpty(text)){
if (StringUtil.isNullOrEmpty(text)) {
throw new BusinessException("please.input.the.caption");
}
generate.setText(text);
break;
case "image":
if (Objects.isNull(elementId)){
if (Objects.isNull(elementId)) {
throw new BusinessException("please.choose.an.image");
}
generate.setCollectionElementId(elementId);
break;
case "text-image":
if (StringUtil.isNullOrEmpty(text) || Objects.isNull(elementId)){
if (StringUtil.isNullOrEmpty(text) || Objects.isNull(elementId)) {
throw new BusinessException("please.input.the.caption.and.choose.an.image");
}
generate.setText(text);
@@ -191,21 +193,21 @@ public class GenerateServiceImpl extends ServiceImpl<GenerateMapper, Generate> i
// 1、判断参数是否正确
// 1.1 必须参数是否非空
if (SKETCH_BOARD.getRealName().equals(generateLikeDTO.getLevel1Type())) {
if (StringUtil.isNullOrEmpty(generateLikeDTO.getLevel2Type())){
if (StringUtil.isNullOrEmpty(generateLikeDTO.getLevel2Type())) {
throw new BusinessException("level2Type.cannot.be.empty");
}
if (StringUtil.isNullOrEmpty(generateLikeDTO.getGender())){
if (StringUtil.isNullOrEmpty(generateLikeDTO.getGender())) {
throw new BusinessException("gender.cannot.be.empty");
}
}
// 1.2 判断参数是否真实有效
Long generateDetailId = generateLikeDTO.getGenerateDetailId();
GenerateDetail generateDetail = generateDetailMapper.selectById(generateDetailId);
if (Objects.isNull(generateDetail)){
if (Objects.isNull(generateDetail)) {
throw new BusinessException("generateItem.does.not.exist");
}
Generate generate = getById(generateDetail.getGenerateId());
if (!generateLikeDTO.getLevel1Type().equals(generate.getLevel1Type())){
if (!generateLikeDTO.getLevel1Type().equals(generate.getLevel1Type())) {
throw new BusinessException("level1Type.does.not.match");
}
@@ -213,8 +215,8 @@ public class GenerateServiceImpl extends ServiceImpl<GenerateMapper, Generate> i
// 2.1、不能重复喜欢
// 2.1.1 判断该图片是否被喜欢过
Library libraryDetail = libraryService.getById(generateDetail.getLibraryId());
if ( (Objects.nonNull(generateDetail.getLibraryId()) && !generateDetail.getLibraryId().equals(0L))
|| Objects.nonNull(libraryDetail)){
if ((Objects.nonNull(generateDetail.getLibraryId()) && !generateDetail.getLibraryId().equals(0L))
|| Objects.nonNull(libraryDetail)) {
throw new BusinessException("duplicate.likes.are.not.allowed");
}
@@ -237,7 +239,7 @@ public class GenerateServiceImpl extends ServiceImpl<GenerateMapper, Generate> i
public Boolean generateDislike(Long generateDetailId, String timeZone) {
// 1、确定generateDetail中是否有这条记录
GenerateDetail generateDetail = generateDetailMapper.selectById(generateDetailId);
if (Objects.isNull(generateDetail)){
if (Objects.isNull(generateDetail)) {
throw new BusinessException("generateItem.does.not.exist");
}
@@ -287,7 +289,7 @@ public class GenerateServiceImpl extends ServiceImpl<GenerateMapper, Generate> i
generateDetailMapper.update(generateDetail, queryWrapper);
}
public void updateLikeStatusBatch(List<Long> generateDetailIdList, Byte hasLike, Long libraryId, String timeZone){
public void updateLikeStatusBatch(List<Long> generateDetailIdList, Byte hasLike, Long libraryId, String timeZone) {
QueryWrapper<GenerateDetail> queryWrapper = new QueryWrapper<>();
queryWrapper.in("id", generateDetailIdList);
@@ -299,93 +301,101 @@ public class GenerateServiceImpl extends ServiceImpl<GenerateMapper, Generate> i
generateDetailMapper.update(generateDetail, queryWrapper);
}
public List<GenerateDetail> selectBatchByLibraryId(List<Long> libraryId){
public List<GenerateDetail> selectBatchByLibraryId(List<Long> libraryId) {
QueryWrapper<GenerateDetail> qw = new QueryWrapper<>();
qw.in("library_id",libraryId);
qw.in("library_id", libraryId);
return generateDetailMapper.selectList(qw);
}
@Override
public Long prepareForGenerate(GenerateThroughImageTextDTO generateThroughImageTextDTO) {
public String prepareForGenerate(GenerateThroughImageTextDTO generateThroughImageTextDTO) {
// 1、参数检查，判断必须参数是否为空
if (Objects.isNull(generateThroughImageTextDTO.getUserId())){
if (Objects.isNull(generateThroughImageTextDTO.getUserId())) {
throw new BusinessException("userId cannot be empty");
}
String generateType = generateThroughImageTextDTO.getGenerateType();
if (!GenerateModeEnum.getGenerateModeList().contains(generateType)){
if (!GenerateModeEnum.getGenerateModeList().contains(generateType)) {
throw new BusinessException("unknown.generate.type");
}
String text = generateThroughImageTextDTO.getText();
Long elementId = generateThroughImageTextDTO.getCollectionElementId();
validateGeneraType(new Generate(), text, elementId,generateType);
validateGeneraType(new Generate(), text, elementId, generateType);
// 2、确定当前排队人数总数，超过15个，暂停使用当前功能
Long zSetTotal = redisUtil.getZSetTotal(consumptionOrderKey);
if (zSetTotal.equals(15L)){
return null;
// 2、生成唯一id 使用uuid
String uuid = UUID.randomUUID().toString();
// SnowflakeUtil idWorker = new SnowflakeUtil(0, 0);
// long snowflakeId = idWorker.nextId();
int num = 1;
// 判断与已经正常生成结果的uuid中有相同的id
while (redisUtil.isElementExistsInMap(resultMapKey, uuid) && num < 10) {
uuid = UUID.randomUUID().toString();
num++;
}
// 3、生成唯一id
SnowflakeUtil idWorker = new SnowflakeUtil(0, 0);
long snowflakeId = idWorker.nextId();
if (AsyncCallerUtil.waitingStatus.containsKey(snowflakeId)){
snowflakeId = idWorker.nextId();
// 无依据确定的数字
if (num > 10){
try {
Thread.sleep(1000);
} catch (InterruptedException e) {
throw new RuntimeException(e);
}
uuid = UUID.randomUUID().toString();
}
generateThroughImageTextDTO.setUniqueId(snowflakeId);
generateThroughImageTextDTO.setUniqueId(uuid);
String jsonString = JSON.toJSONString(generateThroughImageTextDTO);
// 4、加入redis排队，便于获取实时排队信息
// 3、加入redis排队，便于获取实时排队信息
Double maxScore = redisUtil.getMaxScore(consumptionOrderKey);
redisUtil.addToZSet(consumptionOrderKey, String.valueOf(snowflakeId),maxScore);
redisUtil.addToZSet(consumptionOrderKey, uuid, maxScore);
// 5、将消息发布到MQ消息队列
// 4、将消息发布到MQ消息队列
rabbitMQService.publishMessage(jsonString);
// 6、返回唯一id
return snowflakeId;
// 5、返回唯一id
return uuid;
}
@Override
public Long getRankPosition(Long uniqueId) {
return redisUtil.getRank(consumptionOrderKey, String.valueOf(uniqueId));
public Long getRankPosition(String uniqueId) {
return redisUtil.getRank(consumptionOrderKey, uniqueId);
}
@Override
public GenerateCollectionVO getGenerateResult(Long uniqueId) {
public GenerateCollectionVO getGenerateResult(String uniqueId) {
// 1、判断该请求是否已经异常
Boolean isMember = redisUtil.isElementExistsInMap(exceptionMapKey, String.valueOf(uniqueId));
if (isMember){
Boolean isMember = redisUtil.isElementExistsInMap(exceptionMapKey, uniqueId);
if (isMember) {
throw new BusinessException("generate.interface.error");
}
// 2、判断该请求是否还在排队
Boolean existsInZSet = redisUtil.isElementExistsInZSet(consumptionOrderKey, String.valueOf(uniqueId));
if (existsInZSet){
Boolean existsInZSet = redisUtil.isElementExistsInZSet(consumptionOrderKey, uniqueId);
if (existsInZSet) {
// 排队中，给出当前排序位置
return new GenerateCollectionVO(getRankPosition(uniqueId) + 1L);
}
// 3、判断redis中有没有
boolean hasHashKey = redisUtil.isElementExistsInMap(resultMapKey, String.valueOf(uniqueId));
if (hasHashKey){
boolean hasHashKey = redisUtil.isElementExistsInMap(resultMapKey, uniqueId);
if (hasHashKey) {
// 3.1 有直接从redis中拿
String resultString = redisUtil.getMapValue(resultMapKey, String.valueOf(uniqueId));
return JSONObject.parseObject(resultString,GenerateCollectionVO.class);
String resultString = redisUtil.getMapValue(resultMapKey, uniqueId);
return JSONObject.parseObject(resultString, GenerateCollectionVO.class);
}
// 3.2 判断数据库中有没有
Generate generate = selectByUniqueId(uniqueId);
if (Objects.isNull(generate)){
if (Objects.isNull(generate)) {
// 3.3 还没执行完，给出当前位置
return new GenerateCollectionVO(0L);
}
Long generateId = generate.getId();
QueryWrapper<GenerateDetail> qw = new QueryWrapper<>();
qw.eq("generate_id",generateId);
qw.eq("generate_id", generateId);
List<GenerateDetail> generateDetails = generateDetailMapper.selectList(qw);
if (CollectionUtils.isEmpty(generateDetails)){
if (CollectionUtils.isEmpty(generateDetails)) {
// 会有这种情况吗？存到generate中，但是还没存到generateDetail中
return new GenerateCollectionVO(0L);
}
@@ -401,7 +411,7 @@ public class GenerateServiceImpl extends ServiceImpl<GenerateMapper, Generate> i
return new GenerateCollectionVO(generateId, null, generatedCollectionItems);
}
public Generate selectByUniqueId(Long uniqueId){
public Generate selectByUniqueId(String uniqueId){
QueryWrapper<Generate> qw = new QueryWrapper<>();
qw.eq("unique_id",uniqueId);
@@ -409,23 +419,24 @@ public class GenerateServiceImpl extends ServiceImpl<GenerateMapper, Generate> i
}
@Override
public void cancelGenerate(Long uniqueId) {
public void cancelGenerate(String uniqueId) {
// 1、确认当前消息是否还在排队中
Boolean exists = redisUtil.isElementExistsInZSet(consumptionOrderKey, String.valueOf(uniqueId));
if (exists){
Boolean exists = redisUtil.isElementExistsInZSet(consumptionOrderKey, uniqueId);
if (exists) {
// 1.1、将需要取消的唯一id加入redis，以便及时取消生成
redisUtil.addToSet(cancelSetKey, String.valueOf(uniqueId));
redisUtil.addToSet(cancelSetKey, uniqueId);
// 1.2 将需要取消的id从redis的ConsumptionOrder中删除
redisUtil.removeFromZSet(consumptionOrderKey, String.valueOf(uniqueId));
redisUtil.removeFromZSet(consumptionOrderKey, uniqueId);
}else {
// 2、判断该消息是否异常
boolean hasKey = redisUtil.isElementExistsInMap(exceptionMapKey, String.valueOf(uniqueId));
boolean hasKey = redisUtil.isElementExistsInMap(exceptionMapKey, uniqueId);
// 3、判断该消息是否已经消费结束
Boolean existsInResult = redisUtil.isElementExistsInMap(resultMapKey, String.valueOf(uniqueId));
Boolean existsInResult = redisUtil.isElementExistsInMap(resultMapKey, uniqueId);
if (!hasKey && !existsInResult){
// 设置取等待状态为false
AsyncCallerUtil.waitingStatus.put(uniqueId,false);
// 3、直接发送取消请求到python端
pythonService.cancelGenerateTask(uniqueId);
}
}
}