From b25637e8adcde375c129bf383b8fa0b497fab658 Mon Sep 17 00:00:00 2001
From: Juergen Hoeller <hoellerj@vmware.com>
Date: Mon, 9 Nov 2020 13:32:34 +0100
Subject: [PATCH] Make setTaskScheduler and setTcpClient chaining-friendly

Closes gh-26049
---
 .../config/StompBrokerRelayRegistration.java  | 47 +++++++++----------
 1 file changed, 23 insertions(+), 24 deletions(-)

diff --git a/spring-messaging/src/main/java/org/springframework/messaging/simp/config/StompBrokerRelayRegistration.java b/spring-messaging/src/main/java/org/springframework/messaging/simp/config/StompBrokerRelayRegistration.java
index 1e58609626b..d24b63e2dd0 100644
--- a/spring-messaging/src/main/java/org/springframework/messaging/simp/config/StompBrokerRelayRegistration.java
+++ b/spring-messaging/src/main/java/org/springframework/messaging/simp/config/StompBrokerRelayRegistration.java
@@ -56,6 +56,9 @@ public class StompBrokerRelayRegistration extends AbstractBrokerRegistration {
 	@Nullable
 	private TcpOperations<byte[]> tcpClient;
 
+	@Nullable
+	private TaskScheduler taskScheduler;
+
 	private boolean autoStartup = true;
 
 	@Nullable
@@ -64,8 +67,6 @@ public class StompBrokerRelayRegistration extends AbstractBrokerRegistration {
 	@Nullable
 	private String userRegistryBroadcast;
 
-	@Nullable
-	private TaskScheduler taskScheduler;
 
 	public StompBrokerRelayRegistration(SubscribableChannel clientInboundChannel,
 			MessageChannel clientOutboundChannel, String[] destinationPrefixes) {
@@ -181,8 +182,27 @@ public class StompBrokerRelayRegistration extends AbstractBrokerRegistration {
 	 * specified are effectively ignored.
 	 * @since 4.3.15
 	 */
-	public void setTcpClient(TcpOperations<byte[]> tcpClient) {
+	public StompBrokerRelayRegistration setTcpClient(TcpOperations<byte[]> tcpClient) {
 		this.tcpClient = tcpClient;
+		return this;
+	}
+
+	/**
+	 * Some STOMP clients (e.g. stomp-js) always send heartbeats at a fixed rate
+	 * but others (Spring STOMP client) do so only when no other messages are
+	 * sent. However messages with a non-broker {@link #getDestinationPrefixes()
+	 * destination prefix} aren't forwarded and as a result the broker may deem
+	 * the connection inactive.
+	 * <p>When this {@link TaskScheduler} is set, it is used to reset a count of
+	 * the number of messages sent from client to broker since the beginning of
+	 * the current heartbeat period. This is then used to decide whether to send
+	 * a heartbeat to the broker when ignoring a message with a non-broker
+	 * destination prefix.
+	 * @since 5.3
+	 */
+	public StompBrokerRelayRegistration setTaskScheduler(@Nullable TaskScheduler taskScheduler) {
+		this.taskScheduler = taskScheduler;
+		return this;
 	}
 
 	/**
@@ -228,26 +248,6 @@ public class StompBrokerRelayRegistration extends AbstractBrokerRegistration {
 		return this;
 	}
 
-	/**
-	 * Some STOMP clients (e.g. stomp-js) always send heartbeats at a fixed rate
-	 * but others (Spring STOMP client) do so only when no other messages are
-	 * sent. However messages with a non-broker {@link #getDestinationPrefixes()
-	 * destination prefix} aren't forwarded and as a result the broker may deem
-	 * the connection inactive.
-	 *
-	 * <p>When this {@link TaskScheduler} is set, it is used to reset a count of
-	 * the number of messages sent from client to broker since the beginning of
-	 * the current heartbeat period. This is then used to decide whether to send
-	 * a heartbeat to the broker when ignoring a message with a non-broker
-	 * destination prefix.
-	 *
-	 * @param taskScheduler the scheduler to use
-	 * @since 5.3
-	 */
-	public void setTaskScheduler(@Nullable TaskScheduler taskScheduler) {
-		this.taskScheduler = taskScheduler;
-	}
-
 	@Nullable
 	protected String getUserRegistryBroadcast() {
 		return this.userRegistryBroadcast;
@@ -256,7 +256,6 @@ public class StompBrokerRelayRegistration extends AbstractBrokerRegistration {
 
 	@Override
 	protected StompBrokerRelayMessageHandler getMessageHandler(SubscribableChannel brokerChannel) {
-
 		StompBrokerRelayMessageHandler handler = new StompBrokerRelayMessageHandler(
 				getClientInboundChannel(), getClientOutboundChannel(),
 				brokerChannel, getDestinationPrefixes());