[grid] shutdown the client related to a websocket

Author: joerg1985

java/src/org/openqa/selenium/grid/node/ProxyNodeWebsockets.java

@@ -230,7 +230,20 @@ private Consumer<Message> createWsEndPoint(
         client.openSocket(
             new HttpRequest(GET, uri.toString()),
             new ForwardingListener(downstream, sessionConsumer, sessionId));
-    return upstream::send;
+
+    return (msg) -> {
+      try {
+        upstream.send(msg);
+      } finally {
+        if (msg instanceof CloseMessage) {
+          try {
+            client.close();
+          } catch (Exception e) {
+            LOG.log(Level.WARNING, "Failed to shutdown the client of " + uri, e);
+          }
+        }
+      }
+    };
   }
 
   private static class ForwardingListener implements WebSocket.Listener {

java/src/org/openqa/selenium/grid/router/ProxyWebsocketsIntoGrid.java

@@ -71,7 +71,20 @@ public Optional<Consumer<Message>> apply(String uri, Consumer<Message> downstrea
       WebSocket upstream =
           client.openSocket(new HttpRequest(GET, uri), new ForwardingListener(downstream));
 
-      return Optional.of(upstream::send);
+      return Optional.of(
+          (msg) -> {
+            try {
+              upstream.send(msg);
+            } finally {
+              if (msg instanceof CloseMessage) {
+                try {
+                  client.close();
+                } catch (Exception e) {
+                  LOG.log(Level.WARNING, "Failed to shutdown the client of " + sessionUri, e);
+                }
+              }
+            }
+          });
 
     } catch (NoSuchSessionException e) {
       LOG.warning("Attempt to connect to non-existent session: " + uri);