Merge pull request #111 from volkanto/chore-fix-jms-timestamp

eddy-aws · web-flow · commit 2d09d8568571 · 2022-05-26T16:23:41.000-07:00
JMSTimestamp could be set to SentTimestamp
diff --git a/src/main/java/com/amazon/sqs/javamessaging/SQSMessageConsumerPrefetch.java b/src/main/java/com/amazon/sqs/javamessaging/SQSMessageConsumerPrefetch.java
@@ -19,6 +19,7 @@
 import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import java.util.UUID;
 
@@ -94,14 +95,14 @@ public class SQSMessageConsumerPrefetch implements Runnable, PrefetchManager {
      * Counter on how many messages are prefetched into internal messageQueue.
      */
     protected int messagesPrefetched = 0;
-    
+
     /**
      * Counter on how many messages have been explicitly requested.
      * TODO: Consider renaming this class and several other variables now that
      * this logic factors in message requests as well as prefetching.
      */
     protected int messagesRequested = 0;
-    
+
     /**
      * States of the prefetch thread
      */
@@ -121,7 +122,7 @@ public class SQSMessageConsumerPrefetch implements Runnable, PrefetchManager {
      * 25ms delayInterval.
      */
     protected ExponentialBackoffStrategy backoffStrategy = new ExponentialBackoffStrategy(25,25,2000);
-    
+
     SQSMessageConsumerPrefetch(SQSSessionCallbackScheduler sqsSessionRunnable, Acknowledger acknowledger,
                                NegativeAcknowledger negativeAcknowledger, SQSQueueDestination sqsDestination,
                                AmazonSQSMessagingClientWrapper amazonSQSClient, int numberOfMessagesToPrefetch) {
@@ -139,16 +140,16 @@ public class SQSMessageConsumerPrefetch implements Runnable, PrefetchManager {
     MessageListener getMessageListener() {
         return messageListener;
     }
-    
+
     void setMessageConsumer(SQSMessageConsumer messageConsumer) {
         this.messageConsumer = messageConsumer;
     }
-    
+
     @Override
     public SQSMessageConsumer getMessageConsumer() {
         return messageConsumer;
     }
-     
+
     /**
      * Sets the message listener.
      * <P>
@@ -168,7 +169,7 @@ protected void setMessageListener(MessageListener messageListener) {
             if (!running || isClosed()) {
                 return;
             }
-            
+
             List<MessageManager> allPrefetchedMessages = new ArrayList<MessageManager>(messageQueue);
             sqsSessionRunnable.scheduleCallBacks(messageListener, allPrefetchedMessages);
             messageQueue.clear();
@@ -179,7 +180,7 @@ protected void setMessageListener(MessageListener messageListener) {
             messageListenerReady();
         }
     }
-    
+
     /**
      * Determine the number of messages we should attempt to fetch from SQS.
      * Returns the difference between the number of messages needed (either for
@@ -189,7 +190,7 @@ private int numberOfMessagesToFetch() {
         int numberOfMessagesNeeded = Math.max(numberOfMessagesToPrefetch, messagesRequested);
         return Math.max(numberOfMessagesNeeded - messagesPrefetched, 0);
     }
-    
+
     /**
      * Runs until the message consumer is closed and in-progress SQS
      * <code>receiveMessage</code> call returns.
@@ -210,7 +211,7 @@ public void run() {
                 if (isClosed()) {
                     break;
                 }
-                
+
                 synchronized (stateLock) {
                     waitForStart();
                     waitForPrefetch();
@@ -238,7 +239,7 @@ public void run() {
             }
         }
     }
-    
+
     /**
      * Call <code>receiveMessage</code> with the given wait time.
      */
@@ -259,7 +260,7 @@ protected List<Message> getMessages(int batchSize, int waitTimeSeconds) throws J
         ReceiveMessageResult receivedMessageResult = amazonSQSClient.receiveMessage(receiveMessageRequest);
         return receivedMessageResult.getMessages();
     }
-    
+
     /**
      * Converts the received message to JMS message, and pushes to messages to
      * either callback scheduler for asynchronous message delivery or to
@@ -278,14 +279,14 @@ protected void processReceivedMessages(List<Message> messages) {
                 nackMessages.add(message.getReceiptHandle());
             }
         }
-        
+
         synchronized (stateLock) {
             if (messageListener != null) {
                 sqsSessionRunnable.scheduleCallBacks(messageListener, messageManagers);
             } else {
                 messageQueue.addAll(messageManagers);
             }
-            
+
             messagesPrefetched += messageManagers.size();
             notifyStateChange();
         }
@@ -359,9 +360,9 @@ protected javax.jms.Message convertToJMSMessage(Message message) throws JMSExcep
                 throw new JMSException("Not a supported JMS message type");
             }
         }
-        
+
         jmsMessage.setJMSDestination(sqsDestination);
-        
+
         MessageAttributeValue replyToQueueNameAttribute = message.getMessageAttributes().get(
                 SQSMessage.JMS_SQS_REPLY_TO_QUEUE_NAME);
         MessageAttributeValue replyToQueueUrlAttribute = message.getMessageAttributes().get(
@@ -372,16 +373,27 @@ protected javax.jms.Message convertToJMSMessage(Message message) throws JMSExcep
             Destination replyToQueue = new SQSQueueDestination(replyToQueueName, replyToQueueUrl);
             jmsMessage.setJMSReplyTo(replyToQueue);
         }
-        
+
         MessageAttributeValue correlationIdAttribute = message.getMessageAttributes().get(
                 SQSMessage.JMS_SQS_CORRELATION_ID);
         if (correlationIdAttribute != null) {
                 jmsMessage.setJMSCorrelationID(correlationIdAttribute.getStringValue());
         }
-        
+
+        jmsMessage.setJMSTimestamp(getJMSTimestamp(message));
         return jmsMessage;
     }
 
+    private long getJMSTimestamp(Message message) {
+        Map<String, String> systemAttributes = message.getAttributes();
+        String timestamp = systemAttributes.get(SQSMessagingClientConstants.SENT_TIMESTAMP);
+        if (timestamp != null) {
+            return Long.parseLong(timestamp);
+        } else {
+            return 0L;
+        }
+    }
+
     protected void nackQueueMessages() {
         // Also nack messages already in the messageQueue
         synchronized (stateLock) {
@@ -407,7 +419,7 @@ protected void waitForStart() throws InterruptedException {
             }
         }
     }
-        
+
     @Override
     public void messageDispatched() {
         synchronized (stateLock) {
@@ -429,7 +441,7 @@ public void messageListenerReady() {
             }
         }
     }
-    
+
     void requestMessage() {
         synchronized (stateLock) {
             messagesRequested++;
@@ -443,7 +455,7 @@ private void unrequestMessage() {
             notifyStateChange();
         }
     }
-    
+
     public static class MessageManager {
 
         private final PrefetchManager prefetchManager;
@@ -467,7 +479,7 @@ public javax.jms.Message getMessage() {
     javax.jms.Message receive() throws JMSException {
         return receive(0);
     }
-    
+
     javax.jms.Message receive(long timeout) throws JMSException {
         if (cannotDeliver()) {
             return null;
@@ -476,7 +488,7 @@ javax.jms.Message receive(long timeout) throws JMSException {
         if (timeout < 0) {
             timeout = 0;
         }
-        
+
         MessageManager messageManager = null;
         synchronized (stateLock) {
             requestMessage();
@@ -486,7 +498,7 @@ javax.jms.Message receive(long timeout) throws JMSException {
                     messageManager = messageQueue.pollFirst();
                 } else {
             	    long startTime = System.currentTimeMillis();
-    
+
                     long waitTime = 0;
                     while (messageQueue.isEmpty() && !isClosed() &&
                             (timeout == 0 || (waitTime = getWaitTime(timeout, startTime)) > 0)) {
@@ -525,7 +537,7 @@ javax.jms.Message receiveNoWait() throws JMSException {
         if (cannotDeliver()) {
             return null;
         }
-        
+
         MessageManager messageManager;
         synchronized (stateLock) {
             if (messageQueue.isEmpty() && numberOfMessagesToPrefetch == 0) {
@@ -572,22 +584,22 @@ void close() {
             messageListener = null;
         }
     }
-    
+
     /**
      * Helper that notifies PrefetchThread that message is dispatched and AutoAcknowledge
      */
     private javax.jms.Message messageHandler(MessageManager messageManager) throws JMSException {
         if (messageManager == null) {
             return null;
-        }        
+        }
         javax.jms.Message message = messageManager.getMessage();
-        
+
         // Notify PrefetchThread that message is dispatched
         this.messageDispatched();
         acknowledger.notifyMessageReceived((SQSMessage) message);
         return message;
     }
-    
+
     private boolean cannotDeliver() throws JMSException {
         if (!running) {
             return true;
@@ -602,7 +614,7 @@ private boolean cannotDeliver() throws JMSException {
         }
         return false;
     }
-    
+
     /**
      * Sleeps for the configured time.
      */
@@ -613,7 +625,7 @@ protected void sleep(long sleepTimeMillis) throws InterruptedException {
             throw e;
         }
     }
-    
+
     protected boolean isClosed() {
         return closed;
     }
@@ -641,7 +653,7 @@ List<SQSMessageIdentifier> purgePrefetchedMessagesWithGroups(Set<String> affecte
 
             notifyStateChange();
         }
-        
+
         return purgedMessages;
     }
 }
diff --git a/src/main/java/com/amazon/sqs/javamessaging/SQSMessagingClientConstants.java b/src/main/java/com/amazon/sqs/javamessaging/SQSMessagingClientConstants.java
@@ -70,6 +70,8 @@ public class SQSMessagingClientConstants {
 
     public static final String APPROXIMATE_RECEIVE_COUNT = "ApproximateReceiveCount";
 
+    public static final String SENT_TIMESTAMP = "SentTimestamp";
+
     public static final String MESSAGE_DEDUPLICATION_ID = "MessageDeduplicationId";
 
     public static final String MESSAGE_GROUP_ID = "MessageGroupId";
diff --git a/src/test/java/com/amazon/sqs/javamessaging/SQSMessageConsumerPrefetchTest.java b/src/test/java/com/amazon/sqs/javamessaging/SQSMessageConsumerPrefetchTest.java
@@ -56,6 +56,7 @@
 import javax.jms.MessageListener;
 import javax.jms.ObjectMessage;
 
+import org.joda.time.DateTime;
 import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
@@ -1066,6 +1067,9 @@ public void testConvertToJMSMessageTextTypeAttribute() throws JMSException, IOEx
 
         Map<String, String> mapAttributes = new HashMap<String, String>();
         mapAttributes.put(SQSMessagingClientConstants.APPROXIMATE_RECEIVE_COUNT, "1");
+        Long now = DateTime.now().getMillis();
+        mapAttributes.put(SQSMessagingClientConstants.SENT_TIMESTAMP, now.toString());
+
 
         com.amazonaws.services.sqs.model.Message message = mock(com.amazonaws.services.sqs.model.Message.class);
         // Return message attributes with message type 'TEXT'
@@ -1083,6 +1087,7 @@ public void testConvertToJMSMessageTextTypeAttribute() throws JMSException, IOEx
          */
         assertTrue(jsmMessage instanceof SQSTextMessage);
         assertEquals(message.getBody(), "MessageBody");
+        assertEquals(jsmMessage.getJMSTimestamp(), now.longValue());
     }
 
     /**
