Skip to content

Commit

Permalink
WebSocketMessageBrokerStats implements SmartInitializingSingleton
Browse files Browse the repository at this point in the history
Closes gh-26536
  • Loading branch information
rstoyanchev committed Nov 21, 2023
1 parent ed613e7 commit 6db00e6
Showing 1 changed file with 39 additions and 32 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);

Expand Down Expand Up @@ -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 stompHandler) {
return stompHandler;
}
}
SubProtocolHandler defaultHandler = this.webSocketHandler.getDefaultProtocolHandler();
if (defaultHandler instanceof StompSubProtocolHandler stompHandler) {
return stompHandler;
}
return null;
}

public void setStompBrokerRelay(StompBrokerRelayMessageHandler stompBrokerRelay) {
Expand All @@ -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;
}

/**
Expand All @@ -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);
}

/**
Expand All @@ -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 stompHandler) {
return stompHandler;
}
}
SubProtocolHandler defaultHandler = this.webSocketHandler.getDefaultProtocolHandler();
if (defaultHandler instanceof StompSubProtocolHandler stompHandler) {
return stompHandler;
}
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.
*/
Expand Down

0 comments on commit 6db00e6

Please sign in to comment.