Merge branch '5.3.x'

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

@@ -19,7 +19,6 @@ package org.springframework.messaging.simp.stomp;
 import java.lang.reflect.Type;
 import java.time.Instant;
 import java.util.ArrayList;
-import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
@@ -28,6 +27,7 @@ import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Consumer;
 
 import org.apache.commons.logging.Log;
 
@@ -439,7 +439,7 @@ public class DefaultStompSession implements ConnectionHandlingStompSession {
 					String receiptId = headers.getReceiptId();
 					ReceiptHandler handler = this.receiptHandlers.get(receiptId);
 					if (handler != null) {
-						handler.handleReceiptReceived();
+						handler.handleReceiptReceived(headers);
 					}
 					else if (logger.isDebugEnabled()) {
 						logger.debug("No matching receipt: " + accessor.getDetailedLogMessage(message.getPayload()));
@@ -544,7 +544,7 @@ public class DefaultStompSession implements ConnectionHandlingStompSession {
 		@Nullable
 		private final String receiptId;
 
-		private final List<Runnable> receiptCallbacks = new ArrayList<>(2);
+		private final List<Consumer<StompHeaders>> receiptCallbacks = new ArrayList<>(2);
 
 		private final List<Runnable> receiptLostCallbacks = new ArrayList<>(2);
 
@@ -554,6 +554,9 @@ public class DefaultStompSession implements ConnectionHandlingStompSession {
 		@Nullable
 		private Boolean result;
 
+		@Nullable
+		private StompHeaders receiptHeaders;
+
 		public ReceiptHandler(@Nullable String receiptId) {
 			this.receiptId = receiptId;
 			if (receiptId != null) {
@@ -576,64 +579,80 @@ public class DefaultStompSession implements ConnectionHandlingStompSession {
 
 		@Override
 		public void addReceiptTask(Runnable task) {
-			addTask(task, true);
+			addReceiptTask(headers -> task.run());
+		}
+
+		@Override
+		public void addReceiptTask(Consumer<StompHeaders> task) {
+			Assert.notNull(this.receiptId, "Set autoReceiptEnabled to track receipts or add a 'receiptId' header");
+			synchronized (this) {
+				if (this.result != null) {
+					if (this.result) {
+						task.accept(this.receiptHeaders);
+					}
+				}
+				else {
+					this.receiptCallbacks.add(task);
+				}
+			}
 		}
 
 		@Override
 		public void addReceiptLostTask(Runnable task) {
-			addTask(task, false);
-		}
-
-		private void addTask(Runnable task, boolean successTask) {
-			Assert.notNull(this.receiptId,
-					"To track receipts, set autoReceiptEnabled=true or add 'receiptId' header");
 			synchronized (this) {
-				if (this.result != null && this.result == successTask) {
-					invoke(Collections.singletonList(task));
+				if (this.result != null) {
+					if (!this.result) {
+						task.run();
+					}
 				}
 				else {
-					if (successTask) {
-						this.receiptCallbacks.add(task);
-					}
-					else {
-						this.receiptLostCallbacks.add(task);
-					}
+					this.receiptLostCallbacks.add(task);
 				}
 			}
 		}
 
-		private void invoke(List<Runnable> callbacks) {
-			for (Runnable runnable : callbacks) {
-				try {
-					runnable.run();
-				}
-				catch (Throwable ex) {
-					// ignore
-				}
-			}
-		}
-
-		public void handleReceiptReceived() {
-			handleInternal(true);
+		public void handleReceiptReceived(StompHeaders receiptHeaders) {
+			handleInternal(true, receiptHeaders);
 		}
 
 		public void handleReceiptNotReceived() {
-			handleInternal(false);
+			handleInternal(false, null);
 		}
 
-		private void handleInternal(boolean result) {
+		private void handleInternal(boolean result, @Nullable StompHeaders receiptHeaders) {
 			synchronized (this) {
 				if (this.result != null) {
 					return;
 				}
 				this.result = result;
-				invoke(result ? this.receiptCallbacks : this.receiptLostCallbacks);
+				this.receiptHeaders = receiptHeaders;
+				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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2019 the original author or authors.
+ * Copyright 2002-2022 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,6 +16,8 @@
 
 package org.springframework.messaging.simp.stomp;
 
+import java.util.function.Consumer;
+
 import org.springframework.lang.Nullable;
 
 /**
@@ -139,16 +141,27 @@ public interface StompSession {
 
 		/**
 		 * 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);
+
+		/**
+		 * 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 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);
 	}
 
 

spring-messaging/src/test/java/org/springframework/messaging/simp/stomp/DefaultStompSessionTests.java

@@ -575,22 +575,30 @@ public class DefaultStompSessionTests {
 		this.session.setTaskScheduler(mock(TaskScheduler.class));
 
 		AtomicReference<Boolean> received = new AtomicReference<>();
+		AtomicReference<StompHeaders> receivedHeaders = new AtomicReference<>();
 
 		StompHeaders headers = new StompHeaders();
 		headers.setDestination("/topic/foo");
 		headers.setReceipt("my-receipt");
 		Subscription subscription = this.session.subscribe(headers, mock(StompFrameHandler.class));
-		subscription.addReceiptTask(() -> received.set(true));
+		subscription.addReceiptTask(receiptHeaders -> {
+			received.set(true);
+			receivedHeaders.set(receiptHeaders);
+		});
 
 		assertThat((Object) received.get()).isNull();
 
 		StompHeaderAccessor accessor = StompHeaderAccessor.create(StompCommand.RECEIPT);
 		accessor.setReceiptId("my-receipt");
+		accessor.setNativeHeader("foo", "bar");
 		accessor.setLeaveMutable(true);
 		this.session.handleMessage(MessageBuilder.createMessage(new byte[0], accessor.getMessageHeaders()));
 
 		assertThat(received.get()).isNotNull();
 		assertThat(received.get()).isTrue();
+		assertThat(receivedHeaders.get()).isNotNull();
+		assertThat(receivedHeaders.get().get("foo").size()).isEqualTo(1);
+		assertThat(receivedHeaders.get().get("foo").get(0)).isEqualTo("bar");
 	}
 
 	@Test
@@ -599,6 +607,7 @@ public class DefaultStompSessionTests {
 		this.session.setTaskScheduler(mock(TaskScheduler.class));
 
 		AtomicReference<Boolean> received = new AtomicReference<>();
+		AtomicReference<StompHeaders> receivedHeaders = new AtomicReference<>();
 
 		StompHeaders headers = new StompHeaders();
 		headers.setDestination("/topic/foo");
@@ -607,13 +616,20 @@ public class DefaultStompSessionTests {
 
 		StompHeaderAccessor accessor = StompHeaderAccessor.create(StompCommand.RECEIPT);
 		accessor.setReceiptId("my-receipt");
+		accessor.setNativeHeader("foo", "bar");
 		accessor.setLeaveMutable(true);
 		this.session.handleMessage(MessageBuilder.createMessage(new byte[0], accessor.getMessageHeaders()));
 
-		subscription.addReceiptTask(() -> received.set(true));
+		subscription.addReceiptTask(receiptHeaders -> {
+			received.set(true);
+			receivedHeaders.set(receiptHeaders);
+		});
 
 		assertThat(received.get()).isNotNull();
 		assertThat(received.get()).isTrue();
+		assertThat(receivedHeaders.get()).isNotNull();
+		assertThat(receivedHeaders.get().get("foo").size()).isEqualTo(1);
+		assertThat(receivedHeaders.get().get("foo").get(0)).isEqualTo("bar");
 	}
 
 	@Test