From 3d6e38bb43fe86cc58bc816eff13b968c1e0884f Mon Sep 17 00:00:00 2001
From: Rossen Stoyanchev <rstoyanchev@pivotal.io>
Date: Tue, 10 Jul 2018 08:12:23 -0400
Subject: [PATCH] Conditional registration of task scheduler for SockJS

Issue: SPR-16189
---
 .../ServletWebSocketHandlerRegistry.java      |  8 +-
 .../WebSocketConfigurationSupport.java        | 76 +++++--------------
 2 files changed, 21 insertions(+), 63 deletions(-)

diff --git a/spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/ServletWebSocketHandlerRegistry.java b/spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/ServletWebSocketHandlerRegistry.java
index 0918d8604107..53089e38bd51 100644
--- a/spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/ServletWebSocketHandlerRegistry.java
+++ b/spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/ServletWebSocketHandlerRegistry.java
@@ -100,16 +100,16 @@ protected boolean requiresTaskScheduler() {
 	}
 
 	/**
-	 * Set a TaskScheduler is set on each SockJS registration that hasn't had one
-	 * registered explicitly. This method needs to be invoked prior to calling
-	 * {@link #getHandlerMapping()}.
+	 * Provide the TaskScheduler to use for SockJS endpoints for which a task
+	 * scheduler has not been explicitly registered. This method must be called
+	 * prior to {@link #getHandlerMapping()}.
 	 */
 	protected void setTaskScheduler(TaskScheduler scheduler) {
 		this.registrations.stream()
 				.map(ServletWebSocketHandlerRegistration::getSockJsServiceRegistration)
 				.filter(Objects::nonNull)
 				.filter(r -> r.getTaskScheduler() == null)
-				.forEach(r -> r.setTaskScheduler(scheduler));
+				.forEach(registration -> registration.setTaskScheduler(scheduler));
 	}
 
 	public AbstractHandlerMapping getHandlerMapping() {
diff --git a/spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/WebSocketConfigurationSupport.java b/spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/WebSocketConfigurationSupport.java
index d89357e6bc79..7f07e551966f 100644
--- a/spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/WebSocketConfigurationSupport.java
+++ b/spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/WebSocketConfigurationSupport.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2017 the original author or authors.
+ * Copyright 2002-2018 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.
@@ -16,14 +16,11 @@
 
 package org.springframework.web.socket.config.annotation;
 
-import java.util.Date;
-import java.util.concurrent.ScheduledFuture;
-
 import org.springframework.context.annotation.Bean;
 import org.springframework.lang.Nullable;
 import org.springframework.scheduling.TaskScheduler;
-import org.springframework.scheduling.Trigger;
 import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
+import org.springframework.util.Assert;
 import org.springframework.web.servlet.HandlerMapping;
 
 /**
@@ -45,7 +42,9 @@ public class WebSocketConfigurationSupport {
 	public HandlerMapping webSocketHandlerMapping() {
 		ServletWebSocketHandlerRegistry registry = initHandlerRegistry();
 		if (registry.requiresTaskScheduler()) {
-			registry.setTaskScheduler(initTaskScheduler());
+			TaskScheduler scheduler = defaultSockJsTaskScheduler();
+			Assert.notNull(scheduler, "Expected default TaskScheduler bean");
+			registry.setTaskScheduler(scheduler);
 		}
 		return registry.getHandlerMapping();
 	}
@@ -62,15 +61,17 @@ protected void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
 	}
 
 	/**
-	 * The default TaskScheduler to use if none is configured via
-	 * {@link SockJsServiceRegistration#setTaskScheduler}, i.e.
+	 * The default TaskScheduler to use if none is registered explicitly via
+	 * {@link SockJsServiceRegistration#setTaskScheduler}:
 	 * <pre class="code">
 	 * &#064;Configuration
 	 * &#064;EnableWebSocket
 	 * public class WebSocketConfig implements WebSocketConfigurer {
 	 *
 	 *   public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
-	 *     registry.addHandler(myWsHandler(), "/echo").withSockJS().setTaskScheduler(myScheduler());
+	 *     registry.addHandler(myHandler(), "/echo")
+	 *             .withSockJS()
+	 *             .setTaskScheduler(myScheduler());
 	 *   }
 	 *
 	 *   // ...
@@ -78,59 +79,16 @@ protected void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
 	 * </pre>
 	 */
 	@Bean
+	@Nullable
 	public TaskScheduler defaultSockJsTaskScheduler() {
-		return initTaskScheduler();
-	}
-
-	private TaskScheduler initTaskScheduler() {
-		if (this.scheduler == null) {
-			ServletWebSocketHandlerRegistry registry = initHandlerRegistry();
-			if (registry.requiresTaskScheduler()) {
-				ThreadPoolTaskScheduler threadPoolScheduler = new ThreadPoolTaskScheduler();
-				threadPoolScheduler.setThreadNamePrefix("SockJS-");
-				threadPoolScheduler.setPoolSize(Runtime.getRuntime().availableProcessors());
-				threadPoolScheduler.setRemoveOnCancelPolicy(true);
-				this.scheduler = threadPoolScheduler;
-			}
-			else {
-				this.scheduler = new NoOpScheduler();
-			}
+		if (initHandlerRegistry().requiresTaskScheduler()) {
+			ThreadPoolTaskScheduler threadPoolScheduler = new ThreadPoolTaskScheduler();
+			threadPoolScheduler.setThreadNamePrefix("SockJS-");
+			threadPoolScheduler.setPoolSize(Runtime.getRuntime().availableProcessors());
+			threadPoolScheduler.setRemoveOnCancelPolicy(true);
+			this.scheduler = threadPoolScheduler;
 		}
 		return this.scheduler;
 	}
 
-
-	private static class NoOpScheduler implements TaskScheduler {
-
-		@Override
-		@Nullable
-		public ScheduledFuture<?> schedule(Runnable task, Trigger trigger) {
-			throw new IllegalStateException("Unexpected use of scheduler.");
-		}
-
-		@Override
-		public ScheduledFuture<?> schedule(Runnable task, Date startTime) {
-			throw new IllegalStateException("Unexpected use of scheduler.");
-		}
-
-		@Override
-		public ScheduledFuture<?> scheduleAtFixedRate(Runnable task, Date startTime, long period) {
-			throw new IllegalStateException("Unexpected use of scheduler.");
-		}
-
-		@Override
-		public ScheduledFuture<?> scheduleAtFixedRate(Runnable task, long period) {
-			throw new IllegalStateException("Unexpected use of scheduler.");
-		}
-
-		@Override
-		public ScheduledFuture<?> scheduleWithFixedDelay(Runnable task, Date startTime, long delay) {
-			throw new IllegalStateException("Unexpected use of scheduler.");
-		}
-
-		@Override
-		public ScheduledFuture<?> scheduleWithFixedDelay(Runnable task, long delay) {
-			throw new IllegalStateException("Unexpected use of scheduler.");
-		}
-	}
 }