Pass headers to STOMP receipt callbacks

See gh-28715

Author: schnapster

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2020 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.
@@ -27,6 +27,7 @@
 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;
 
@@ -441,7 +442,7 @@ else if (logger.isDebugEnabled()) {
 					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()));
@@ -546,16 +547,19 @@ private class ReceiptHandler implements Receiptable {
 		@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);
+		private final List<Consumer<StompHeaders>> receiptLostCallbacks = new ArrayList<>(2);
 
 		@Nullable
 		private ScheduledFuture<?> future;
 
 		@Nullable
 		private Boolean result;
 
+		@Nullable
+		private StompHeaders receiptHeaders;
+
 		public ReceiptHandler(@Nullable String receiptId) {
 			this.receiptId = receiptId;
 			if (receiptId != null) {
@@ -578,15 +582,20 @@ public String getReceiptId() {
 
 		@Override
 		public void addReceiptTask(Runnable task) {
+			addTask(h -> task.run(), true);
+		}
+
+		@Override
+		public void addReceiptTask(Consumer<StompHeaders> task) {
 			addTask(task, true);
 		}
 
 		@Override
 		public void addReceiptLostTask(Runnable task) {
-			addTask(task, false);
+			addTask(h -> task.run(), false);
 		}
 
-		private void addTask(Runnable task, boolean successTask) {
+		private void addTask(Consumer<StompHeaders> task, boolean successTask) {
 			Assert.notNull(this.receiptId,
 					"To track receipts, set autoReceiptEnabled=true or add 'receiptId' header");
 			synchronized (this) {
@@ -604,31 +613,32 @@ private void addTask(Runnable task, boolean successTask) {
 			}
 		}
 
-		private void invoke(List<Runnable> callbacks) {
-			for (Runnable runnable : callbacks) {
+		private void invoke(List<Consumer<StompHeaders>> callbacks) {
+			for (Consumer<StompHeaders> consumer : callbacks) {
 				try {
-					runnable.run();
+					consumer.accept(this.receiptHeaders);
 				}
 				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;
+				this.receiptHeaders = receiptHeaders;
 				invoke(result ? this.receiptCallbacks : this.receiptLostCallbacks);
 				DefaultStompSession.this.receiptHandlers.remove(this.receiptId);
 				if (this.future != null) {

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;
 
 /**
@@ -143,6 +145,13 @@ interface Receiptable {
 		 */
 		void addReceiptTask(Runnable runnable);
 
+		/**
+		 * Consumer to invoke when a receipt is received. Accepts the headers of the received RECEIPT frame.
+		 * @throws java.lang.IllegalArgumentException if the receiptId is {@code null}
+		 * @since TBD
+		 */
+		void addReceiptTask(Consumer<StompHeaders> task);
+
 		/**
 		 * Task to invoke when a receipt is not received in the configured time.
 		 * @throws java.lang.IllegalArgumentException if the receiptId is {@code null}

spring-messaging/src/test/java/org/springframework/messaging/simp/stomp/DefaultStompSessionTests.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.
@@ -576,22 +576,30 @@ public void receiptReceived() {
 		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
@@ -600,6 +608,7 @@ public void receiptReceivedBeforeTaskAdded() {
 		this.session.setTaskScheduler(mock(TaskScheduler.class));
 
 		AtomicReference<Boolean> received = new AtomicReference<>();
+		AtomicReference<StompHeaders> receivedHeaders = new AtomicReference<>();
 
 		StompHeaders headers = new StompHeaders();
 		headers.setDestination("/topic/foo");
@@ -608,13 +617,20 @@ public void receiptReceivedBeforeTaskAdded() {
 
 		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

src/docs/asciidoc/web/websocket.adoc

@@ -1347,7 +1347,7 @@ receipt if the server supports it (simple broker does not). For example, with th
 	headers.setDestination("/topic/...");
 	headers.setReceipt("r1");
 	FrameHandler handler = ...;
-	stompSession.subscribe(headers, handler).addReceiptTask(() -> {
+	stompSession.subscribe(headers, handler).addReceiptTask(receiptHeaders -> {
 		// Subscription ready...
 	});
 ----