[JENKINS-70103] `WeakHashMap` broken for looking up Jetty `Session` in WebSocket response (#7412)

(cherry picked from commit aedacbab)

core/src/main/java/jenkins/websocket/WebSockets.java

@@ -68,12 +68,20 @@ public class WebSockets {
         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();

websocket/jetty10/src/main/java/jenkins/websocket/Jetty10Provider.java

@@ -26,9 +26,6 @@ package jenkins.websocket;
 
 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<Listener, Session> sessions = Collections.synchronizedMap(new WeakHashMap<>());
-
     public Jetty10Provider() {
         JettyWebSocketServerContainer.class.hashCode();
     }
@@ -101,7 +95,7 @@ public class Jetty10Provider implements Provider {
             }
 
             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 class Jetty10Provider implements Provider {
 
             @Override
             public void onWebSocketConnect(Session session) {
-                sessions.put(listener, session);
-                listener.onWebSocketConnect();
+                listener.onWebSocketConnect(session);
             }
 
             @Override

websocket/jetty9/src/main/java/jenkins/websocket/Jetty9Provider.java

@@ -26,9 +26,6 @@ package jenkins.websocket;
 
 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<Listener, Session> sessions = Collections.synchronizedMap(new WeakHashMap<>());
-
     private WebSocketServletFactory factory;
 
     public Jetty9Provider() {
@@ -104,7 +98,7 @@ public class Jetty9Provider implements Provider {
             }
 
             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 class Jetty9Provider implements Provider {
 
             @Override
             public void onWebSocketConnect(Session session) {
-                sessions.put(listener, session);
-                listener.onWebSocketConnect();
+                listener.onWebSocketConnect(session);
             }
 
             @Override

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);