malei05 1 mēnesi atpakaļ
vecāks
revīzija
2e8b3a2c13

+ 42 - 11
fn-rabbitmq/src/main/java/net/yaoyi/pipeline/rabbitmq/MultiResultProducer.java

@@ -2,28 +2,34 @@ package net.yaoyi.pipeline.rabbitmq;
 
 import com.aliyun.fc.runtime.Context;
 import com.aliyun.fc.runtime.PojoRequestHandler;
+import com.rabbitmq.client.AMQP;
 import com.rabbitmq.client.Channel;
 import com.rabbitmq.client.Connection;
 import lombok.extern.slf4j.Slf4j;
 import net.yaoyi.pipeline.common.JacksonUtils;
-import net.yaoyi.pipeline.model.AlgoBaseInfo;
 import net.yaoyi.pipeline.model.FnResult;
+import net.yaoyi.pipeline.model.mq.ImageMultiMsg;
 
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 import java.util.List;
+import java.util.UUID;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
 
 @Slf4j
-public class MultiResultProducer implements PojoRequestHandler<List<AlgoBaseInfo>, FnResult<?>> {
+public class MultiResultProducer implements PojoRequestHandler<List<ImageMultiMsg>, FnResult<?>> {
 
     private static final String QUEUE_NAME = "vector.search.task.img.dul.check.result.queue";
     private static final String EXCHANGE_NAME = "default.direct.exchange";
 
+    private static final ConcurrentMap<Long, String> unConfirmMsg = new ConcurrentHashMap<>();
+
     private Connection connection;
     private Channel channel;
 
     @Override
-    public FnResult<?> handleRequest(List<AlgoBaseInfo> algoBaseInfos, Context context) {
+    public FnResult<?> handleRequest(List<ImageMultiMsg> algoBaseInfos, Context context) {
         if (algoBaseInfos == null || algoBaseInfos.isEmpty()) {
             return FnResult.ok();
         }
@@ -32,24 +38,41 @@ public class MultiResultProducer implements PojoRequestHandler<List<AlgoBaseInfo
         initRabbitMQ();
 
         //2.执行消息发送
-        for (AlgoBaseInfo algoBaseInfo : algoBaseInfos) {
-            rabbitSend(JacksonUtils.toJson(algoBaseInfo));
-        }
+        rabbitSend(algoBaseInfos);
 
         //3.关闭生产端
         closeRabbitMQ();
-        return FnResult.ok();
+        return FnResult.ok(unConfirmMsg);
     }
 
     private void initRabbitMQ() {
         connection = RabbitMQHelper.createConnection();
-        channel = RabbitMQHelper.createChannel(connection, QUEUE_NAME);
+        channel = RabbitMQHelper.createPublishChannel(connection, QUEUE_NAME);
+
+        channel.addConfirmListener((seqNo, multiple) -> {
+            if (multiple) {
+                unConfirmMsg.keySet().removeIf(k -> k <= seqNo);
+            } else {
+                unConfirmMsg.remove(seqNo);
+            }
+        }, (seqNo, multiple) -> log.error("消息发送失败,seqNo: {}, multiple: {}", seqNo, multiple));
+    }
+
+    private void rabbitSend(List<ImageMultiMsg> algoBaseInfos) {
+        long seq = channel.getNextPublishSeqNo();
+        for (ImageMultiMsg algoBaseInfo : algoBaseInfos) {
+            String body = JacksonUtils.toJson(algoBaseInfo);
+            rabbitSend(body);
+            unConfirmMsg.put(seq, body);
+        }
     }
 
     private void rabbitSend(String json) {
+        String messageId = UUID.randomUUID().toString().replaceAll("-", "");
         try {
-            channel.basicPublish(EXCHANGE_NAME, QUEUE_NAME, null, json.getBytes(StandardCharsets.UTF_8));
-            log.info("消息发送成功: {}", json);
+            AMQP.BasicProperties props = new AMQP.BasicProperties.Builder().messageId(messageId).build();
+            channel.basicPublish(EXCHANGE_NAME, QUEUE_NAME, props, json.getBytes(StandardCharsets.UTF_8));
+            log.info("消息发送成功,messageId:{}. body:{}", messageId, json);
         } catch (IOException e) {
             log.error("消息发送失败: {}", json, e);
             throw new RuntimeException("消息发送失败: " + e.getMessage(), e);
@@ -57,7 +80,15 @@ public class MultiResultProducer implements PojoRequestHandler<List<AlgoBaseInfo
     }
 
     private void closeRabbitMQ() {
-        RabbitMQHelper.close(channel, connection);
+        try {
+            if (channel.waitForConfirms()) {
+                RabbitMQHelper.close(channel, connection);
+            }
+        } catch (Exception e) {
+            log.error("channel 关闭失败");
+        } finally {
+            RabbitMQHelper.close(channel, connection);
+        }
     }
 
 }

+ 6 - 0
yyc-pipeline-model/src/main/java/net/yaoyi/pipeline/model/FnResult.java

@@ -16,4 +16,10 @@ public class FnResult<T> {
         return new FnResult<>();
     }
 
+    public static <T> FnResult<?> ok(T data) {
+        FnResult<T> result = new FnResult<>();
+        result.setData(data);
+        return result;
+    }
+
 }