diff --git a/core/src/main/java/jenkins/websocket/WebSockets.java b/core/src/main/java/jenkins/websocket/WebSockets.java index aafabcd79c9d..f5a62864e39a 100644 --- a/core/src/main/java/jenkins/websocket/WebSockets.java +++ b/core/src/main/java/jenkins/websocket/WebSockets.java @@ -68,12 +68,20 @@ public static HttpResponse upgrade(WebSocketSession session) { return (req, rsp, node) -> { try { session.handler = provider.handle(req, rsp, new Provider.Listener() { + private Object providerSession; + @Override - public void onWebSocketConnect() { + public void onWebSocketConnect(Object providerSession) { + this.providerSession = providerSession; session.startPings(); session.opened(); } + @Override + public Object getProviderSession() { + return providerSession; + } + @Override public void onWebSocketClose(int statusCode, String reason) { session.stopPings(); diff --git a/websocket/jetty10/src/main/java/jenkins/websocket/Jetty10Provider.java b/websocket/jetty10/src/main/java/jenkins/websocket/Jetty10Provider.java index 8504bd5f1822..19533240e323 100644 --- a/websocket/jetty10/src/main/java/jenkins/websocket/Jetty10Provider.java +++ b/websocket/jetty10/src/main/java/jenkins/websocket/Jetty10Provider.java @@ -26,9 +26,6 @@ import java.io.IOException; import java.nio.ByteBuffer; -import java.util.Collections; -import java.util.Map; -import java.util.WeakHashMap; import java.util.concurrent.CompletableFuture; import java.util.concurrent.Future; import javax.servlet.http.HttpServletRequest; @@ -49,9 +46,6 @@ public class Jetty10Provider implements Provider { private static final String ATTR_LISTENER = Jetty10Provider.class.getName() + ".listener"; - // TODO does not seem possible to use HttpServletRequest.get/setAttribute for this - private static final Map sessions = Collections.synchronizedMap(new WeakHashMap<>()); - public Jetty10Provider() { JettyWebSocketServerContainer.class.hashCode(); } @@ -101,7 +95,7 @@ public void close() throws IOException { } private Session session() { - Session session = sessions.get(listener); + Session session = (Session) listener.getProviderSession(); if (session == null) { throw new IllegalStateException("missing session"); } @@ -151,8 +145,7 @@ public void onWebSocketClose(int statusCode, String reason) { @Override public void onWebSocketConnect(Session session) { - sessions.put(listener, session); - listener.onWebSocketConnect(); + listener.onWebSocketConnect(session); } @Override diff --git a/websocket/jetty9/src/main/java/jenkins/websocket/Jetty9Provider.java b/websocket/jetty9/src/main/java/jenkins/websocket/Jetty9Provider.java index 078e0c2480df..ab0771a973cf 100644 --- a/websocket/jetty9/src/main/java/jenkins/websocket/Jetty9Provider.java +++ b/websocket/jetty9/src/main/java/jenkins/websocket/Jetty9Provider.java @@ -26,9 +26,6 @@ import java.io.IOException; import java.nio.ByteBuffer; -import java.util.Collections; -import java.util.Map; -import java.util.WeakHashMap; import java.util.concurrent.Future; import javax.servlet.http.HttpServletRequest; import javax.servlet.http.HttpServletResponse; @@ -48,9 +45,6 @@ public class Jetty9Provider implements Provider { private static final String ATTR_LISTENER = Jetty9Provider.class.getName() + ".listener"; - // TODO does not seem possible to use HttpServletRequest.get/setAttribute for this - private static final Map sessions = Collections.synchronizedMap(new WeakHashMap<>()); - private WebSocketServletFactory factory; public Jetty9Provider() { @@ -104,7 +98,7 @@ public void close() throws IOException { } private Session session() { - Session session = sessions.get(listener); + Session session = (Session) listener.getProviderSession(); if (session == null) { throw new IllegalStateException("missing session"); } @@ -136,8 +130,7 @@ public void onWebSocketClose(int statusCode, String reason) { @Override public void onWebSocketConnect(Session session) { - sessions.put(listener, session); - listener.onWebSocketConnect(); + listener.onWebSocketConnect(session); } @Override diff --git a/websocket/spi/src/main/java/jenkins/websocket/Provider.java b/websocket/spi/src/main/java/jenkins/websocket/Provider.java index caef2559dfb3..0d7c53aeb5c5 100644 --- a/websocket/spi/src/main/java/jenkins/websocket/Provider.java +++ b/websocket/spi/src/main/java/jenkins/websocket/Provider.java @@ -47,7 +47,9 @@ interface Provider { interface Listener { - void onWebSocketConnect(); + void onWebSocketConnect(Object providerSession); + + Object getProviderSession(); void onWebSocketClose(int statusCode, String reason);