From 20d5f628bc08a3077367a1766a81fe98ae9b7a50 Mon Sep 17 00:00:00 2001
From: Artem Bilan <abilan@vmware.com>
Date: Thu, 13 Jul 2023 10:47:00 -0400
Subject: [PATCH] GH-8644: WebSocketHandlerReg: Remove ThreadLocal (#8667)

Fixes https://github.com/spring-projects/spring-integration/issues/8644

* Introduce an inner `DynamicHandlerRegistrationProxy` class
to create an `IntegrationDynamicWebSocketHandlerRegistration`
and map it into a provided `WebSocketHandler` when a `dynamicHandlerMapping`
is in action.

Co-authored-by: pziobron <64628007+pziobron@users.noreply.github.com>
---
 ...rationServletWebSocketHandlerRegistry.java | 49 ++++++++++---------
 1 file changed, 27 insertions(+), 22 deletions(-)

diff --git a/spring-integration-websocket/src/main/java/org/springframework/integration/websocket/config/IntegrationServletWebSocketHandlerRegistry.java b/spring-integration-websocket/src/main/java/org/springframework/integration/websocket/config/IntegrationServletWebSocketHandlerRegistry.java
index 8c486f5abdd..ad26f0f1a78 100644
--- a/spring-integration-websocket/src/main/java/org/springframework/integration/websocket/config/IntegrationServletWebSocketHandlerRegistry.java
+++ b/spring-integration-websocket/src/main/java/org/springframework/integration/websocket/config/IntegrationServletWebSocketHandlerRegistry.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2021 the original author or authors.
+ * Copyright 2021-2023 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.
@@ -34,6 +34,7 @@
 import org.springframework.web.socket.config.annotation.ServletWebSocketHandlerRegistration;
 import org.springframework.web.socket.config.annotation.ServletWebSocketHandlerRegistry;
 import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistration;
+import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;
 
 /**
  * The {@link ServletWebSocketHandlerRegistry} extension for Spring Integration purpose, especially
@@ -46,8 +47,6 @@
 class IntegrationServletWebSocketHandlerRegistry extends ServletWebSocketHandlerRegistry
 		implements ApplicationContextAware, DestructionAwareBeanPostProcessor {
 
-	private final ThreadLocal<IntegrationDynamicWebSocketHandlerRegistration> currentRegistration = new ThreadLocal<>();
-
 	private final Map<WebSocketHandler, List<String>> dynamicRegistrations = new HashMap<>();
 
 	private ApplicationContext applicationContext;
@@ -78,30 +77,17 @@ public AbstractHandlerMapping getHandlerMapping() {
 		return originHandlerMapping;
 	}
 
-	@Override
-	public WebSocketHandlerRegistration addHandler(WebSocketHandler handler, String... paths) {
-		if (this.dynamicHandlerMapping != null) {
-			IntegrationDynamicWebSocketHandlerRegistration registration =
-					new IntegrationDynamicWebSocketHandlerRegistration();
-			registration.addHandler(handler, paths);
-			this.currentRegistration.set(registration);
-			return registration;
-		}
-		else {
-			return super.addHandler(handler, paths);
-		}
-	}
-
 	@Override
 	public Object postProcessAfterInitialization(Object bean, String beanName) throws BeansException {
-		if (this.dynamicHandlerMapping != null && bean instanceof ServerWebSocketContainer) {
-			ServerWebSocketContainer serverWebSocketContainer = (ServerWebSocketContainer) bean;
+		if (this.dynamicHandlerMapping != null && bean instanceof ServerWebSocketContainer serverWebSocketContainer) {
 			if (serverWebSocketContainer.getSockJsTaskScheduler() == null) {
 				serverWebSocketContainer.setSockJsTaskScheduler(this.sockJsTaskScheduler);
 			}
-			serverWebSocketContainer.registerWebSocketHandlers(this);
-			IntegrationDynamicWebSocketHandlerRegistration registration = this.currentRegistration.get();
-			this.currentRegistration.remove();
+
+			DynamicHandlerRegistrationProxy dynamicHandlerRegistrationProxy = new DynamicHandlerRegistrationProxy();
+
+			serverWebSocketContainer.registerWebSocketHandlers(dynamicHandlerRegistrationProxy);
+			IntegrationDynamicWebSocketHandlerRegistration registration = dynamicHandlerRegistrationProxy.registration;
 			MultiValueMap<HttpRequestHandler, String> mappings = registration.getMapping();
 			for (Map.Entry<HttpRequestHandler, List<String>> entry : mappings.entrySet()) {
 				HttpRequestHandler httpHandler = entry.getKey();
@@ -136,11 +122,30 @@ void removeRegistration(ServerWebSocketContainer serverWebSocketContainer) {
 		}
 	}
 
+	private static final class DynamicHandlerRegistrationProxy implements WebSocketHandlerRegistry {
+
+		private IntegrationDynamicWebSocketHandlerRegistration registration;
+
+		DynamicHandlerRegistrationProxy() {
+		}
+
+		@Override
+		public WebSocketHandlerRegistration addHandler(WebSocketHandler webSocketHandler, String... paths) {
+			this.registration = new IntegrationDynamicWebSocketHandlerRegistration();
+			this.registration.addHandler(webSocketHandler, paths);
+			return this.registration;
+		}
+
+	}
+
 	private static final class IntegrationDynamicWebSocketHandlerRegistration
 			extends ServletWebSocketHandlerRegistration {
 
 		private WebSocketHandler handler;
 
+		IntegrationDynamicWebSocketHandlerRegistration() {
+		}
+
 		@Override
 		public WebSocketHandlerRegistration addHandler(WebSocketHandler handler, String... paths) {
 			// The IntegrationWebSocketContainer comes only with a single WebSocketHandler