Adjust SQS message transaction lifetime to handle flux flow (#4205)

* Adjust transaction lifetime to handle flux flow

* fix span end

* update CHANGELOG

* Update apm-agent-plugins/apm-aws-sdk/apm-aws-sdk-common/src/main/java/co/elastic/apm/agent/awssdk/common/AbstractMessageIteratorWrapper.java

Co-authored-by: SylvainJuge <[email protected]>

* review feedback

---------

Co-authored-by: SylvainJuge <[email protected]>

Author: jackshirazi

CHANGELOG.next-release.md

@@ -11,6 +11,7 @@ This file contains all changes which are not released yet.
 <!--FIXES-START-->
 * Fix async httpclient 5.x instrumentation - [#4185](https://github.com/elastic/apm-agent-java/pull/4185)
 * Prevent `FileSystemAlreadyExistsException` on single-jar application startup - [#4204](https://github.com/elastic/apm-agent-java/pull/4204)
+* Improve SQS async transaction handling to be better with flux flow - [#4205](https://github.com/elastic/apm-agent-java/pull/4205)
 
 <!--FIXES-END-->
 # Features and enhancements

apm-agent-plugins/apm-aws-sdk/apm-aws-sdk-common/src/main/java/co/elastic/apm/agent/awssdk/common/AbstractMessageIteratorWrapper.java

@@ -29,6 +29,8 @@
 
 import javax.annotation.Nullable;
 import java.util.Iterator;
+import java.util.NoSuchElementException;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import static co.elastic.apm.agent.awssdk.common.AbstractSQSInstrumentationHelper.MESSAGE_PROCESSING_ACTION;
 import static co.elastic.apm.agent.awssdk.common.AbstractSQSInstrumentationHelper.MESSAGING_TYPE;
@@ -43,6 +45,7 @@ public abstract class AbstractMessageIteratorWrapper<Message> implements Iterato
     private final String queueName;
     private final AbstractSQSInstrumentationHelper<?, ?, Message> sqsInstrumentationHelper;
     private final TextHeaderGetter<Message> textHeaderGetter;
+    private final AtomicBoolean iterationEnded;
 
     public AbstractMessageIteratorWrapper(Iterator<Message> delegate, Tracer tracer,
                                           String queueName,
@@ -53,21 +56,25 @@ public AbstractMessageIteratorWrapper(Iterator<Message> delegate, Tracer tracer,
         this.queueName = queueName;
         this.sqsInstrumentationHelper = sqsInstrumentationHelper;
         this.textHeaderGetter = textHeaderGetter;
+        this.iterationEnded = new AtomicBoolean(false);
     }
 
     @Override
     public boolean hasNext() {
-        endCurrentTransaction();
-        endMessageProcessingSpan();
-        return delegate.hasNext();
+        boolean hasNext = delegate.hasNext();
+        if (!hasNext && iterationEnded.compareAndSet(false, true)) {
+            endCurrentTransaction();
+            endMessageProcessingSpan();
+        }
+        return hasNext;
     }
 
     @Nullable
     public Transaction<?> endCurrentTransaction() {
         Transaction<?> transaction = null;
         try {
             transaction = tracer.currentTransaction();
-            if (transaction != null && MESSAGING_TYPE.equals(transaction.getType())) {
+            if (transaction != null && !transaction.isFinished() && MESSAGING_TYPE.equals(transaction.getType())) {
                 transaction.deactivate().end();
                 return null;
             }
@@ -93,12 +100,21 @@ public void endMessageProcessingSpan() {
 
     @Override
     public Message next() {
-        Transaction<?> currentTransaction = endCurrentTransaction();
-        Message sqsMessage = delegate.next();
-        if (currentTransaction == null) {
-            sqsInstrumentationHelper.startTransactionOnMessage(sqsMessage, queueName, textHeaderGetter);
+        try {
+            Transaction<?> currentTransaction = endCurrentTransaction();
+            Message sqsMessage = delegate.next();
+            if (currentTransaction == null) {
+                sqsInstrumentationHelper.startTransactionOnMessage(sqsMessage, queueName, textHeaderGetter);
+            }
+            return sqsMessage;
+        } catch (NoSuchElementException e) {
+            // end the transaction when the caller loops with a try/catch until an exception is thrown
+            if (iterationEnded.compareAndSet(false, true)) {
+                endCurrentTransaction();
+                endMessageProcessingSpan();
+            }
+            throw e;
         }
-        return sqsMessage;
     }