Polishing contribution

Closes gh-28715

Author: rstoyanchev

spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/DefaultStompSession.java

@@ -18,7 +18,6 @@
 
 import java.lang.reflect.Type;
 import java.util.ArrayList;
-import java.util.Collections;
 import java.util.Date;
 import java.util.List;
 import java.util.Map;
@@ -549,7 +548,7 @@ private class ReceiptHandler implements Receiptable {
 
 		private final List<Consumer<StompHeaders>> receiptCallbacks = new ArrayList<>(2);
 
-		private final List<Consumer<StompHeaders>> receiptLostCallbacks = new ArrayList<>(2);
+		private final List<Runnable> receiptLostCallbacks = new ArrayList<>(2);
 
 		@Nullable
 		private ScheduledFuture<?> future;
@@ -582,44 +581,34 @@ public String getReceiptId() {
 
 		@Override
 		public void addReceiptTask(Runnable task) {
-			addTask(h -> task.run(), true);
+			addReceiptTask(headers -> task.run());
 		}
 
 		@Override
 		public void addReceiptTask(Consumer<StompHeaders> task) {
-			addTask(task, true);
-		}
-
-		@Override
-		public void addReceiptLostTask(Runnable task) {
-			addTask(h -> task.run(), false);
-		}
-
-		private void addTask(Consumer<StompHeaders> task, boolean successTask) {
-			Assert.notNull(this.receiptId,
-					"To track receipts, set autoReceiptEnabled=true or add 'receiptId' header");
+			Assert.notNull(this.receiptId, "Set autoReceiptEnabled to track receipts or add a 'receiptId' header");
 			synchronized (this) {
-				if (this.result != null && this.result == successTask) {
-					invoke(Collections.singletonList(task));
+				if (this.result != null) {
+					if (this.result) {
+						task.accept(this.receiptHeaders);
+					}
 				}
 				else {
-					if (successTask) {
-						this.receiptCallbacks.add(task);
-					}
-					else {
-						this.receiptLostCallbacks.add(task);
-					}
+					this.receiptCallbacks.add(task);
 				}
 			}
 		}
 
-		private void invoke(List<Consumer<StompHeaders>> callbacks) {
-			for (Consumer<StompHeaders> consumer : callbacks) {
-				try {
-					consumer.accept(this.receiptHeaders);
+		@Override
+		public void addReceiptLostTask(Runnable task) {
+			synchronized (this) {
+				if (this.result != null) {
+					if (!this.result) {
+						task.run();
+					}
 				}
-				catch (Throwable ex) {
-					// ignore
+				else {
+					this.receiptLostCallbacks.add(task);
 				}
 			}
 		}
@@ -639,13 +628,33 @@ private void handleInternal(boolean result, @Nullable StompHeaders receiptHeader
 				}
 				this.result = result;
 				this.receiptHeaders = receiptHeaders;
-				invoke(result ? this.receiptCallbacks : this.receiptLostCallbacks);
+				if (result) {
+					this.receiptCallbacks.forEach(consumer -> {
+						try {
+							consumer.accept(this.receiptHeaders);
+						}
+						catch (Throwable ex) {
+							// ignore
+						}
+					});
+				}
+				else {
+					this.receiptLostCallbacks.forEach(task -> {
+						try {
+							task.run();
+						}
+						catch (Throwable ex) {
+							// ignore
+						}
+					});
+				}
 				DefaultStompSession.this.receiptHandlers.remove(this.receiptId);
 				if (this.future != null) {
 					this.future.cancel(true);
 				}
 			}
 		}
+
 	}
 
 

spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompSession.java

@@ -141,23 +141,27 @@ interface Receiptable {
 
 		/**
 		 * Task to invoke when a receipt is received.
+		 * @param task the task to invoke
 		 * @throws java.lang.IllegalArgumentException if the receiptId is {@code null}
 		 */
-		void addReceiptTask(Runnable runnable);
+		void addReceiptTask(Runnable task);
 
 		/**
-		 * Consumer to invoke when a receipt is received. Accepts the headers of the received RECEIPT frame.
+		 * Variant of {@link #addReceiptTask(Runnable)} with a {@link Consumer}
+		 * of the headers from the {@code RECEIPT} frame.
+		 * @param task the consumer to invoke
 		 * @throws java.lang.IllegalArgumentException if the receiptId is {@code null}
-		 * @since TBD
+		 * @since 5.3.23
 		 */
 		void addReceiptTask(Consumer<StompHeaders> task);
 
 		/**
 		 * Task to invoke when a receipt is not received in the configured time.
+		 * @param task the task to invoke
 		 * @throws java.lang.IllegalArgumentException if the receiptId is {@code null}
 		 * @see org.springframework.messaging.simp.stomp.StompClientSupport#setReceiptTimeLimit(long)
 		 */
-		void addReceiptLostTask(Runnable runnable);
+		void addReceiptLostTask(Runnable task);
 	}