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:
Artem Bilan
2023-07-13 10:47:00 -04:00
committed by GitHub
parent fc3c8d2c8d
commit 20d5f628bc

View File

@@ -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