From 3b531655743867cfaef8fc9145669b2d1d2692e2 Mon Sep 17 00:00:00 2001
From: Brian Clozel <brian.clozel@broadcom.com>
Date: Mon, 27 May 2024 16:06:32 +0200
Subject: [PATCH] Expose WebSocketMessageBrokerStats contracts

Prior to this commit, the `WebSocketMessageBrokerStats` would be in
charge of periodically logging WebSocket stats. This class would also
publicly expose each stats type with dedicated methods, as `String`.
This would not allow observation libraries to easily extract information
and turn them into metrics.

This commit introduces new methods exposing the `Stats` types directly
and deprecates the former `String` variants. This will allow
observability libraries like Micrometer to expose this as metrics:

```
MeterRegistry meterRegistry = ...;

Gauge.builder("spring.stomp.frames", stats.getStompSubProtocolStats(),
      StompSubProtocolHandler.Stats::getTotalConnect)
  .tag("type", "CONNECT")
  .description("number of CONNECT frames processed")
  .register(meterRegistry);
```

Closes gh-31604
---
 .../config/WebSocketMessageBrokerStats.java   | 46 +++++++++++++++++--
 1 file changed, 43 insertions(+), 3 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 8f57cbea4e..5cc7c3911a 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
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2023 the original author or authors.
+ * Copyright 2002-2024 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.
@@ -52,6 +52,7 @@ import org.springframework.web.socket.messaging.SubProtocolWebSocketHandler;
  *
  * @author Rossen Stoyanchev
  * @author Sam Brannen
+ * @author Brian Clozel
  * @since 4.1
  */
 public class WebSocketMessageBrokerStats implements SmartInitializingSingleton {
@@ -160,25 +161,63 @@ public class WebSocketMessageBrokerStats implements SmartInitializingSingleton {
 
 	/**
 	 * Get stats about WebSocket sessions.
+	 * @deprecated as of 6.2 in favor of {@link #getWebSocketSessionStats()}.
 	 */
+	@Deprecated(since = "6.2", forRemoval = true)
 	public String getWebSocketSessionStatsInfo() {
 		return (this.webSocketHandler != null ? this.webSocketHandler.getStatsInfo() : "null");
 	}
 
 	/**
-	 * Get stats about STOMP-related WebSocket message processing.
+	 * Get stats about WebSocket sessions.
+	 * Can return {@code null} if no {@link #setSubProtocolWebSocketHandler(SubProtocolWebSocketHandler) WebSocket handler}
+	 * is configured.
+	 * @since 6.2
 	 */
+	@Nullable
+	public SubProtocolWebSocketHandler.Stats getWebSocketSessionStats() {
+		return (this.webSocketHandler != null ? this.webSocketHandler.getStats() : null);
+	}
+
+	/**
+	 * Get stats about STOMP-related WebSocket message processing.
+	 * @deprecated as of 6.2 in favor of {@link #getStompSubProtocolStats()}.
+	 */
+	@Deprecated(since = "6.2", forRemoval = true)
 	public String getStompSubProtocolStatsInfo() {
 		return (this.stompSubProtocolHandler != null ? this.stompSubProtocolHandler.getStatsInfo() : "null");
 	}
 
 	/**
-	 * Get stats about STOMP broker relay (when using a full-featured STOMP broker).
+	 * Get stats about STOMP-related WebSocket message processing.
+	 * Can return {@code null} if no {@link SubProtocolHandler} was found.
+	 * @since 6.2
 	 */
+	@Nullable
+	public StompSubProtocolHandler.Stats getStompSubProtocolStats() {
+		return (this.stompSubProtocolHandler != null ? this.stompSubProtocolHandler.getStats() : null);
+	}
+
+	/**
+	 * Get stats about STOMP broker relay (when using a full-featured STOMP broker).
+	 * @deprecated as of 6.2 in favor of {@link #getStompBrokerRelayStats()}.
+	 */
+	@Deprecated(since = "6.2", forRemoval = true)
 	public String getStompBrokerRelayStatsInfo() {
 		return (this.stompBrokerRelay != null ? this.stompBrokerRelay.getStatsInfo() : "null");
 	}
 
+	/**
+	 * Get stats about STOMP broker relay (when using a full-featured STOMP broker).
+	 * Can return {@code null} if no {@link #setStompBrokerRelay(StompBrokerRelayMessageHandler) STOMP broker relay}
+	 * is configured.
+	 * @since 6.2
+	 */
+	@Nullable
+	public StompBrokerRelayMessageHandler.Stats getStompBrokerRelayStats() {
+		return (this.stompBrokerRelay != null ? this.stompBrokerRelay.getStats() : null);
+	}
+
 	/**
 	 * Get stats about the executor processing incoming messages from WebSocket clients.
 	 */
@@ -231,6 +270,7 @@ public class WebSocketMessageBrokerStats implements SmartInitializingSingleton {
 	}
 
 	@Override
+	@SuppressWarnings("removal")
 	public String toString() {
 		return "WebSocketSession[" + getWebSocketSessionStatsInfo() + "]" +
 				", stompSubProtocol[" + getStompSubProtocolStatsInfo() + "]" +