[ConsoleProxy] Performance fixes on concurrency and closing console sessions sockets

services/console-proxy/server/src/main/java/com/cloud/consoleproxy/ConsoleProxy.java

@@ -26,7 +26,6 @@
 import java.net.InetSocketAddress;
 import java.net.URISyntaxException;
 import java.net.URL;
-import java.util.HashSet;
 import java.util.Hashtable;
 import java.util.Map;
 import java.util.Properties;
@@ -83,8 +82,10 @@ public class ConsoleProxy {
     static String encryptorPassword = "Dummy";
     static final String[] skipProperties = new String[]{"certificate", "cacertificate", "keystore_password", "privatekey"};
 
-    static Set<String> allowedSessions = new HashSet<>();
+    static Set<String> allowedSessions = ConcurrentHashMap.newKeySet();
+    private static final Object allowedSessionsLock = new Object();
 
+    // Invoked through reflection
     public static void addAllowedSession(String sessionUuid) {
         allowedSessions.add(sessionUuid);
     }
@@ -209,13 +210,15 @@ public static ConsoleProxyAuthenticationResult authenticateConsoleAccess(Console
         }
 
         String sessionUuid = param.getSessionUuid();
-        if (allowedSessions.contains(sessionUuid)) {
-            LOGGER.debug("Acquiring the session " + sessionUuid + " not available for future use");
-            allowedSessions.remove(sessionUuid);
-        } else {
-            LOGGER.info("Session " + sessionUuid + " has already been used, cannot connect");
-            authResult.setSuccess(false);
-            return authResult;
+        synchronized (allowedSessionsLock) {
+            if (allowedSessions.contains(sessionUuid)) {
+                LOGGER.debug("Acquiring the session " + sessionUuid + " not available for future use");
+                allowedSessions.remove(sessionUuid);
+            } else {
+                LOGGER.info("Session " + sessionUuid + " has already been used, cannot connect");
+                authResult.setSuccess(false);
+                return authResult;
+            }
         }
 
         String websocketUrl = param.getWebsocketUrl();
@@ -625,7 +628,7 @@ public static ConsoleProxyNoVncClient getNoVncViewer(ConsoleProxyClientParam par
                 } catch (IOException e) {
                     LOGGER.error("Exception while disconnect session of novnc viewer object: " + viewer, e);
                 }
-                removeViewer(viewer);
+                viewer.closeClient();
                 viewer = new ConsoleProxyNoVncClient(session);
                 viewer.initClient(param);
                 connectionMap.put(clientKey, viewer);

services/console-proxy/server/src/main/java/com/cloud/consoleproxy/ConsoleProxyNoVNCHandler.java

@@ -175,12 +175,20 @@ private boolean checkSessionSourceIp(final Session session, final String sourceI
 
     @OnWebSocketClose
     public void onClose(Session session, int statusCode, String reason) throws IOException, InterruptedException {
-        String sessionSourceIp = session.getRemoteAddress().getAddress().getHostAddress();
-        logger.debug("Closing WebSocket session [source IP: {}, status code: {}].", sessionSourceIp, statusCode);
         if (viewer != null) {
-            ConsoleProxy.removeViewer(viewer);
+            viewer.closeClient();
+        }
+        String sessionSourceIp = getRemoteAddressSafely(session);
+        logger.debug("WebSocket session [source IP: {}, status code: {}, reason: {}] closed successfully.", sessionSourceIp, statusCode, reason);
+    }
+
+    private String getRemoteAddressSafely(Session session) {
+        try {
+            return session.getRemoteAddress().getAddress().getHostAddress();
+        } catch (Exception e) {
+            logger.debug("Failed to get remote address from WebSocket session", e);
+            return "unknown";
         }
-        logger.debug("WebSocket session [source IP: {}, status code: {}] closed successfully.", sessionSourceIp, statusCode);
     }
 
     @OnWebSocketFrame

services/console-proxy/server/src/main/java/com/cloud/consoleproxy/ConsoleProxyNoVncClient.java

@@ -374,6 +374,9 @@ public void closeClient() {
         this.connectionAlive = false;
         // Clear buffer reference to allow GC when client disconnects
         this.readBuffer = null;
+        if (client != null) {
+            client.close();
+        }
         ConsoleProxy.removeViewer(this);
     }
 

services/console-proxy/server/src/main/java/com/cloud/consoleproxy/vnc/NoVncClient.java

@@ -129,6 +129,28 @@ public void proxyMsgOverWebSocketConnection(ByteBuffer msg) {
         }
     }
 
+    public void close() {
+        if (nioSocketConnection != null) {
+            nioSocketConnection.close();
+        }
+        if (webSocketReverseProxy != null) {
+            webSocketReverseProxy.close();
+        }
+        if (socket != null) {
+            try {
+                if (is != null) {
+                    is.close();
+                }
+                if (os != null) {
+                    os.close();
+                }
+                socket.close();
+            } catch (IOException e) {
+                logger.debug("Error closing socket: " + e.getMessage(), e);
+            }
+        }
+    }
+
     private void setTunnelSocketStreams() throws IOException {
         this.is = new DataInputStream(this.socket.getInputStream());
         this.os = new DataOutputStream(this.socket.getOutputStream());

services/console-proxy/server/src/main/java/com/cloud/consoleproxy/vnc/network/NioSocket.java

@@ -120,4 +120,28 @@ protected int writeToSocketChannel(ByteBuffer buf, int len) {
             return 0;
         }
     }
+
+    public void close() {
+        try {
+            if (socketChannel != null) {
+                socketChannel.close();
+            }
+        } catch (IOException e) {
+            logger.debug("Error closing socket channel: " + e.getMessage(), e);
+        }
+        try {
+            if (readSelector != null) {
+                readSelector.close();
+            }
+        } catch (IOException e) {
+            logger.debug("Error closing read selector: " + e.getMessage(), e);
+        }
+        try {
+            if (writeSelector != null) {
+                writeSelector.close();
+            }
+        } catch (IOException e) {
+            logger.debug("Error closing write selector: " + e.getMessage(), e);
+        }
+    }
 }

services/console-proxy/server/src/main/java/com/cloud/consoleproxy/vnc/network/NioSocketHandler.java

@@ -41,4 +41,5 @@ public interface NioSocketHandler {
     void flushWriteBuffer();
     void startTLSConnection(NioSocketSSLEngineManager sslEngineManager);
     boolean isTLSConnection();
+    void close();
 }

services/console-proxy/server/src/main/java/com/cloud/consoleproxy/vnc/network/NioSocketHandlerImpl.java

@@ -28,10 +28,12 @@ public class NioSocketHandlerImpl implements NioSocketHandler {
     private NioSocketInputStream inputStream;
     private NioSocketOutputStream outputStream;
     private boolean isTLS = false;
+    private final NioSocket socket;
 
     protected Logger logger = LogManager.getLogger(getClass());
 
     public NioSocketHandlerImpl(NioSocket socket) {
+        this.socket = socket;
         this.inputStream = new NioSocketInputStream(ConsoleProxy.defaultBufferSize, socket);
         this.outputStream = new NioSocketOutputStream(ConsoleProxy.defaultBufferSize, socket);
     }
@@ -109,4 +111,9 @@ public NioSocketInputStream getInputStream() {
     public NioSocketOutputStream getOutputStream() {
         return outputStream;
     }
+
+    @Override
+    public void close() {
+        socket.close();
+    }
 }