Monitor your WebSockets and ThreadPools for Performance Tuning via JMX

Last week, after releasing Simple RealTime and having the traffic spike from Hacker News, I needed to setup some basic monitoring for performance tuning.

Even if there’s no silver bullet in performance tuning, there are some clear points you may want to look at like: concurrent connections, buffers and thread pool executors for inbound and outbound message channels (see the configuration and performance section for further information).

The following class (based on this) exposes to JMX the number of concurrent connections and the information of the ThreadPoolTaskExecutor used for outbound messages (modify as needed, including the inbound executor for instance):

@ManagedResource
public class WebSocketMessageBrokerStatsMonitor {

	private final Map<String, WebSocketSession> webSocketSessions;
	private final ThreadPoolExecutor outboundExecutor;

	@SuppressWarnings("unchecked")
	public WebSocketMessageBrokerStatsMonitor(SubProtocolWebSocketHandler webSocketHandler, ThreadPoolTaskExecutor outboundTaskExecutor) {
		this.webSocketSessions = (Map<String, WebSocketSession>) new DirectFieldAccessor(webSocketHandler).getPropertyValue("sessions");
		this.outboundExecutor  = outboundTaskExecutor.getThreadPoolExecutor();
	}

	@ManagedAttribute
	public int getCurrentSessions() {
		return webSocketSessions.size();
	}

	@ManagedAttribute
	public String getSendBufferSize() {
		int sendBufferSize = 0;
		for (WebSocketSession session : this.webSocketSessions.values()) {
			ConcurrentWebSocketSessionDecorator concurrentSession = (ConcurrentWebSocketSessionDecorator) session;
			sendBufferSize += concurrentSession.getBufferSize();
		}
		return formatByteCount(sendBufferSize);
	}

	@ManagedAttribute
	public int getOutboundPoolSize() {
		return outboundExecutor.getPoolSize();
	}

	@ManagedAttribute
	public int getOutboundLargestPoolSize() {
		return outboundExecutor.getLargestPoolSize();
	}

	@ManagedAttribute
	public int getOutboundActiveThreads() {
		return outboundExecutor.getActiveCount();
	}

	@ManagedAttribute
	public int getOutboundQueuedTaskCount() {
		return outboundExecutor.getQueue().size();
	}

	@ManagedAttribute
	public long getOutboundCompletedTaskCount() {
		return outboundExecutor.getCompletedTaskCount();
	}

	private static String formatByteCount(long bytes) {
		int unit = 1024;
		if (bytes < unit) return bytes + " B";
		int exp = (int) (Math.log(bytes) / Math.log(unit));
		return String.format("%.1f %sB", bytes / Math.pow(unit, exp), "KMGTPE".charAt(exp - 1));
	}

}

Define the WebSocketMessageBrokerStatsMonitor as a Spring managed bean:

@Bean
public WebSocketMessageBrokerStatsMonitor statsMonitor() {
	return new WebSocketMessageBrokerStatsMonitor((SubProtocolWebSocketHandler) subProtocolWebSocketHandler(), clientOutboundChannelExecutor());
}

Enable JMX annotations to expose our Spring managed beans by registering the AnnotationMBeanExporter:

@Bean
public AnnotationMBeanExporter annotationMBeanExporter() {
    return new AnnotationMBeanExporter();
}

And you are ready to go: open jConsole / VisualVM to see your live stats (remember to add -Dcom.sun.management.jmxremote as a VM argument):

jConsole WebSockets