From 61ac912f6715d4c556127d06e752c4c9ce588994 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   | 71 ++++++++++---------
 1 file changed, 39 insertions(+), 32 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 16e1b2f51393..7bb3cfc828fc 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 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 @@
  * @author Sam Brannen
  * @since 4.1
  */
-public class WebSocketMessageBrokerStats {
+public class WebSocketMessageBrokerStats implements SmartInitializingSingleton {
 
 	private static final Log logger = LogFactory.getLog(WebSocketMessageBrokerStats.class);
 
@@ -84,24 +85,6 @@ public class WebSocketMessageBrokerStats {
 
 	public void setSubProtocolWebSocketHandler(SubProtocolWebSocketHandler webSocketHandler) {
 		this.webSocketHandler = webSocketHandler;
-		this.stompSubProtocolHandler = initStompSubProtocolHandler();
-	}
-
-	@Nullable
-	private StompSubProtocolHandler initStompSubProtocolHandler() {
-		if (this.webSocketHandler == null) {
-			return null;
-		}
-		for (SubProtocolHandler handler : this.webSocketHandler.getProtocolHandlers()) {
-			if (handler instanceof StompSubProtocolHandler) {
-				return (StompSubProtocolHandler) handler;
-			}
-		}
-		SubProtocolHandler defaultHandler = this.webSocketHandler.getDefaultProtocolHandler();
-		if (defaultHandler instanceof StompSubProtocolHandler) {
-			return (StompSubProtocolHandler) defaultHandler;
-		}
-		return null;
 	}
 
 	public void setStompBrokerRelay(StompBrokerRelayMessageHandler stompBrokerRelay) {
@@ -118,17 +101,6 @@ public void setOutboundChannelExecutor(TaskExecutor 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()) {
-			return this.sockJsTaskScheduler.scheduleWithFixedDelay(
-					() -> logger.info(WebSocketMessageBrokerStats.this.toString()),
-					Instant.now().plusMillis(initialDelay), Duration.ofMillis(this.loggingPeriod));
-		}
-		return null;
 	}
 
 	/**
@@ -137,11 +109,11 @@ private ScheduledFuture<?> initLoggingTask(long initialDelay) {
 	 * <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);
 		}
-		this.loggingPeriod = period;
-		this.loggingTask = initLoggingTask(0);
 	}
 
 	/**
@@ -151,6 +123,41 @@ public long getLoggingPeriod() {
 		return this.loggingPeriod;
 	}
 
+
+	@Override
+	public void afterSingletonsInstantiated() {
+		this.stompSubProtocolHandler = initStompSubProtocolHandler();
+		this.loggingTask = initLoggingTask(TimeUnit.MINUTES.toMillis(1));
+	}
+
+	@Nullable
+	private StompSubProtocolHandler initStompSubProtocolHandler() {
+		if (this.webSocketHandler == null) {
+			return null;
+		}
+		for (SubProtocolHandler handler : this.webSocketHandler.getProtocolHandlers()) {
+			if (handler instanceof StompSubProtocolHandler) {
+				return (StompSubProtocolHandler) handler;
+			}
+		}
+		SubProtocolHandler defaultHandler = this.webSocketHandler.getDefaultProtocolHandler();
+		if (defaultHandler instanceof StompSubProtocolHandler) {
+			return (StompSubProtocolHandler) defaultHandler;
+		}
+		return null;
+	}
+
+	@Nullable
+	private ScheduledFuture<?> initLoggingTask(long initialDelay) {
+		if (this.sockJsTaskScheduler != null && this.loggingPeriod > 0 && logger.isInfoEnabled()) {
+			return this.sockJsTaskScheduler.scheduleWithFixedDelay(
+					() -> logger.info(WebSocketMessageBrokerStats.this.toString()),
+					Instant.now().plusMillis(initialDelay), Duration.ofMillis(this.loggingPeriod));
+		}
+		return null;
+	}
+
+
 	/**
 	 * Get stats about WebSocket sessions.
 	 */