From 8ffbecc7c3b60f5a2af2d8e88feef8bb63d1e14a Mon Sep 17 00:00:00 2001
From: rstoyanchev <rstoyanchev@vmware.com>
Date: Tue, 21 Nov 2023 17:58:00 +0000
Subject: [PATCH] WebSocketMessageBrokerStats implements
 SmartInitializingSingleton

Closes gh-26536
---
 .../config/WebSocketMessageBrokerStats.java   | 81 ++++++++++---------
 1 file changed, 44 insertions(+), 37 deletions(-)

diff --git a/spring-websocket/src/main/java/org/springframework/web/socket/config/WebSocketMessageBrokerStats.java b/spring-websocket/src/main/java/org/springframework/web/socket/config/WebSocketMessageBrokerStats.java
index e5aa7675b3a..8f57cbea4e6 100644
--- a/spring-websocket/src/main/java/org/springframework/web/socket/config/WebSocketMessageBrokerStats.java
+++ b/spring-websocket/src/main/java/org/springframework/web/socket/config/WebSocketMessageBrokerStats.java
@@ -26,6 +26,7 @@ import java.util.concurrent.TimeUnit;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 
+import org.springframework.beans.factory.SmartInitializingSingleton;
 import org.springframework.core.task.TaskExecutor;
 import org.springframework.lang.Nullable;
 import org.springframework.messaging.simp.stomp.StompBrokerRelayMessageHandler;
@@ -53,7 +54,7 @@ import org.springframework.web.socket.messaging.SubProtocolWebSocketHandler;
  * @author Sam Brannen
  * @since 4.1
  */
-public class WebSocketMessageBrokerStats {
+public class WebSocketMessageBrokerStats implements SmartInitializingSingleton {
 
 	private static final Log logger = LogFactory.getLog(WebSocketMessageBrokerStats.class);
 
@@ -84,7 +85,49 @@ public class WebSocketMessageBrokerStats {
 
 	public void setSubProtocolWebSocketHandler(SubProtocolWebSocketHandler webSocketHandler) {
 		this.webSocketHandler = webSocketHandler;
+	}
+
+	public void setStompBrokerRelay(StompBrokerRelayMessageHandler stompBrokerRelay) {
+		this.stompBrokerRelay = stompBrokerRelay;
+	}
+
+	public void setInboundChannelExecutor(TaskExecutor inboundChannelExecutor) {
+		this.inboundChannelExecutor = inboundChannelExecutor;
+	}
+
+	public void setOutboundChannelExecutor(TaskExecutor outboundChannelExecutor) {
+		this.outboundChannelExecutor = outboundChannelExecutor;
+	}
+
+	public void setSockJsTaskScheduler(TaskScheduler sockJsTaskScheduler) {
+		this.sockJsTaskScheduler = sockJsTaskScheduler;
+	}
+
+	/**
+	 * Set the frequency for logging information at INFO level in milliseconds.
+	 * If set 0 or less than 0, the logging task is cancelled.
+	 * <p>By default this property is set to 30 minutes (30 * 60 * 1000).
+	 */
+	public void setLoggingPeriod(long period) {
+		this.loggingPeriod = period;
+		if (this.loggingTask != null) {
+			this.loggingTask.cancel(true);
+			this.loggingTask = initLoggingTask(0);
+		}
+	}
+
+	/**
+	 * Return the configured logging period frequency in milliseconds.
+	 */
+	public long getLoggingPeriod() {
+		return this.loggingPeriod;
+	}
+
+
+	@Override
+	public void afterSingletonsInstantiated() {
 		this.stompSubProtocolHandler = initStompSubProtocolHandler();
+		this.loggingTask = initLoggingTask(TimeUnit.MINUTES.toMillis(1));
 	}
 
 	@Nullable
@@ -104,23 +147,6 @@ public class WebSocketMessageBrokerStats {
 		return null;
 	}
 
-	public void setStompBrokerRelay(StompBrokerRelayMessageHandler stompBrokerRelay) {
-		this.stompBrokerRelay = stompBrokerRelay;
-	}
-
-	public void setInboundChannelExecutor(TaskExecutor inboundChannelExecutor) {
-		this.inboundChannelExecutor = inboundChannelExecutor;
-	}
-
-	public void setOutboundChannelExecutor(TaskExecutor outboundChannelExecutor) {
-		this.outboundChannelExecutor = outboundChannelExecutor;
-	}
-
-	public void setSockJsTaskScheduler(TaskScheduler sockJsTaskScheduler) {
-		this.sockJsTaskScheduler = sockJsTaskScheduler;
-		this.loggingTask = initLoggingTask(TimeUnit.MINUTES.toMillis(1));
-	}
-
 	@Nullable
 	private ScheduledFuture<?> initLoggingTask(long initialDelay) {
 		if (this.sockJsTaskScheduler != null && this.loggingPeriod > 0 && logger.isInfoEnabled()) {
@@ -131,25 +157,6 @@ public class WebSocketMessageBrokerStats {
 		return null;
 	}
 
-	/**
-	 * Set the frequency for logging information at INFO level in milliseconds.
-	 * If set 0 or less than 0, the logging task is cancelled.
-	 * <p>By default this property is set to 30 minutes (30 * 60 * 1000).
-	 */
-	public void setLoggingPeriod(long period) {
-		if (this.loggingTask != null) {
-			this.loggingTask.cancel(true);
-		}
-		this.loggingPeriod = period;
-		this.loggingTask = initLoggingTask(0);
-	}
-
-	/**
-	 * Return the configured logging period frequency in milliseconds.
-	 */
-	public long getLoggingPeriod() {
-		return this.loggingPeriod;
-	}
 
 	/**
 	 * Get stats about WebSocket sessions.