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>
This commit is contained in:
@@ -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.WebSocketHandler;
|
||||
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 @@ import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistra
|
||||
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 @@ class IntegrationServletWebSocketHandlerRegistry extends ServletWebSocketHandler
|
||||
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 @@ class IntegrationServletWebSocketHandlerRegistry extends ServletWebSocketHandler
|
||||
}
|
||||
}
|
||||
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user