Merge pull request #16428 from stuartwdouglas/16427

Fix websockets class loading issue
quarkusio · Apr 12, 2021 · f32102d · f32102d
2 parents 851ef0a + b541a8c
commit f32102d
Showing 1 changed file with 16 additions and 2 deletions.
diff --git a/...ckets/runtime/src/main/java/io/quarkus/undertow/websockets/runtime/WebsocketRecorder.java b/...ckets/runtime/src/main/java/io/quarkus/undertow/websockets/runtime/WebsocketRecorder.java
@@ -19,6 +19,7 @@
 import io.quarkus.arc.runtime.BeanContainer;
 import io.quarkus.runtime.ExecutorRecorder;
 import io.quarkus.runtime.annotations.Recorder;
+import io.undertow.websockets.UndertowContainerProvider;
 import io.undertow.websockets.WebSocketDeploymentInfo;
 import io.undertow.websockets.util.ContextSetupHandler;
 import io.undertow.websockets.util.ObjectFactory;
@@ -107,6 +108,7 @@ public WebSocketDeploymentInfo createDeploymentInfo(Set<String> annotatedEndpoin
 
     public Handler<RoutingContext> createHandler(BeanContainer beanContainer, Supplier<EventLoopGroup> eventLoopGroupSupplier,
             WebSocketDeploymentInfo info) throws DeploymentException {
+        ClassLoader cl = Thread.currentThread().getContextClassLoader();
         ManagedContext requestContext = Arc.container().requestContext();
         VertxServerWebSocketContainer container = new VertxServerWebSocketContainer(new ObjectIntrospecter() {
             @Override
@@ -137,11 +139,22 @@ public <T, C> Action<T, C> create(Action<T, C> action) {
                         return new Action<T, C>() {
                             @Override
                             public T call(C context) throws Exception {
-                                requestContext.activate();
+                                ClassLoader old = Thread.currentThread().getContextClassLoader();
+                                Thread.currentThread().setContextClassLoader(cl);
+                                boolean required = !requestContext.isActive();
+                                if (required) {
+                                    requestContext.activate();
+                                }
                                 try {
                                     return action.call(context);
                                 } finally {
-                                    requestContext.terminate();
+                                    try {
+                                        if (required) {
+                                            requestContext.terminate();
+                                        }
+                                    } finally {
+                                        Thread.currentThread().setContextClassLoader(old);
+                                    }
                                 }
                             }
                         };
@@ -159,6 +172,7 @@ public Executor get() {
         for (ServerEndpointConfig i : info.getProgramaticEndpoints()) {
             container.addEndpoint(i);
         }
+        UndertowContainerProvider.setDefaultContainer(container);
         return new VertxWebSocketHandler(container, info);
     }