ACK代码优化

mpush-core/src/main/java/com/mpush/core/ack/AckContext.java

@@ -19,50 +19,58 @@
 
 package com.mpush.core.ack;
 
-import com.mpush.common.message.BaseMessage;
-
-import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.Future;
 
 /**
  * Created by ohun on 16/9/5.
  *
  * @author ohun@live.cn (夜色)
  */
 public final class AckContext implements Runnable {
-    private final AtomicBoolean done = new AtomicBoolean(false);
-
     private AckCallback callback;
-    /*package*/ int pushMessageId;
+
+    private int sessionId;
+    private Future<?> future;
 
     public AckContext() {
     }
 
-    public boolean tryDone() {
-        return done.compareAndSet(false, true);
+    public void setSessionId(int sessionId) {
+        this.sessionId = sessionId;
+    }
+
+    public void setFuture(Future<?> future) {
+        this.future = future;
     }
 
     public AckContext setCallback(AckCallback callback) {
         this.callback = callback;
         return this;
     }
 
+    private boolean tryDone() {
+        return future.cancel(true);
+    }
+
     public void success() {
         if (tryDone()) {
             callback.onSuccess(this);
+            callback = null;
         }
     }
 
     public void timeout() {
-        AckContext context = AckMessageQueue.I.getAndRemove(pushMessageId);
+        AckContext context = AckMessageQueue.I.getAndRemove(sessionId);
         if (context != null && tryDone()) {
             callback.onTimeout(this);
+            callback = null;
         }
     }
 
     @Override
     public String toString() {
         return "AckContext{" +
-                ", pushMessageId=" + pushMessageId +
+                ", sessionId=" + sessionId +
                 '}';
     }
 

mpush-core/src/main/java/com/mpush/core/ack/AckMessageQueue.java

@@ -34,22 +34,27 @@
  */
 public final class AckMessageQueue {
     private final Logger logger = LoggerFactory.getLogger(AckMessageQueue.class);
+
     private static final int DEFAULT_TIMEOUT = 3000;
     public static final AckMessageQueue I = new AckMessageQueue();
 
     private final ConcurrentMap<Integer, AckContext> queue = new ConcurrentHashMap<>();
     private final ScheduledExecutorService scheduledExecutor;
 
     private AckMessageQueue() {
-        scheduledExecutor = new ScheduledThreadPoolExecutor(1, new NamedPoolThreadFactory(T_ARK_REQ_TIMER), (r, e) -> {
-            logger.error("one ack context was rejected, context=" + r);
-        });
+        scheduledExecutor = new ScheduledThreadPoolExecutor(1,
+                new NamedPoolThreadFactory(T_ARK_REQ_TIMER),
+                (r, e) -> logger.error("one ack context was rejected, context=" + r)
+        );
     }
 
-    public void put(int sessionId, AckContext context, int timeout) {
+    public void add(int sessionId, AckContext context, int timeout) {
         queue.put(sessionId, context);
-        context.pushMessageId = sessionId;
-        scheduledExecutor.schedule(context, timeout > 0 ? timeout : DEFAULT_TIMEOUT, TimeUnit.MILLISECONDS);
+        context.setSessionId(sessionId);
+        context.setFuture(scheduledExecutor.schedule(context,
+                timeout > 0 ? timeout : DEFAULT_TIMEOUT,
+                TimeUnit.MILLISECONDS
+        ));
     }
 
     public AckContext getAndRemove(int sessionId) {

mpush-core/src/main/java/com/mpush/core/push/SingleUserPushTask.java

@@ -105,7 +105,7 @@ private boolean checkLocal(final GatewayPushMessage message) {
                 if (future.isSuccess()) {//推送成功
 
                     if (message.needAck()) {//需要客户端ACK, 消息进队列等待客户端响应ACK
-                        AckMessageQueue.I.put(pushMessage.getSessionId(), buildAckContext(message), message.timeout);
+                        AckMessageQueue.I.add(pushMessage.getSessionId(), buildAckContext(message), message.timeout);
                     } else {
                         OkMessage.from(message).setData(userId + ',' + clientType).sendRaw();
                     }