[ZEPPELIN-6260] Fix memory leak in WebSocket watcher connections

###  What is this PR for?

  This PR fixes a memory leak issue where WebSocket connections that are switched to watcher mode are never removed from the watcherSockets queue when the connection is closed. This causes the queue to grow indefinitely, leading to increased memory
  usage over time and potential OutOfMemoryError in long-running Zeppelin servers.

###  What type of PR is it?

  Bug Fix

###  What is the Jira issue?

  https://issues.apache.org/jira/browse/ZEPPELIN-6260

###  How should this be tested?

  Steps to reproduce the issue:
  1. Open a WebSocket connection to Zeppelin
  2. Send a WATCHER message to switch the connection to watcher mode
  3. Close the connection
  4. The connection remains in watcherSockets queue indefinitely

  Verification after fix:
  - Unit tests have been added to verify the fix:
    - removeWatcherConnectionCleansQueue: Basic functionality test
    - removeWatcherConnectionWithMultipleWatchers: Tests selective removal
    - removeWatcherConnectionConcurrentTest: Tests thread safety
    - switchConnectionToWatcherAndRemove: Tests complete lifecycle

###  Screenshots (if appropriate)

  N/A

###  Questions:

  - Does the license files need update? No
  - Is there breaking changes for older versions? No
  - Does this needs documentation? No

###  Description of changes:

  1. Added removeWatcherConnection() method to ConnectionManager to safely remove connections from the watcherSockets queue
  2. Modified NotebookServer.removeConnection() to call removeWatcherConnection() when any connection is closed
  3. Added debug logging to track watcher connection removal
  4. Added comprehensive unit tests including concurrent access scenarios

  The fix is minimal and safe:
  - The ConcurrentLinkedQueue.remove() operation is safe even if the element doesn't exist
  - The synchronization block is very short, minimizing performance impact
  - The approach ensures no watcher connections are leaked, regardless of how they were closed

###  Related observations:

  During the investigation, I noticed that broadcastToWatchers() doesn't remove watchers when IOException occurs. This could cause performance degradation as closed connections would repeatedly fail. However, this is a separate issue and should be
  addressed in a different PR to keep this fix focused.

Closes #5001 from renechoi/ZEPPELIN-6260-fix-watcher-memory-leak.

Signed-off-by: Philipp Dallig <philipp.dallig@gmail.com>

Author: renechoi

zeppelin-server/src/main/java/org/apache/zeppelin/socket/ConnectionManager.java

@@ -175,6 +175,14 @@ public void removeUserConnection(String user, NotebookSocket conn) {
     }
   }
 
+  public void removeWatcherConnection(NotebookSocket conn) {
+    synchronized (watcherSockets) {
+      if (watcherSockets.remove(conn)) {
+        LOGGER.debug("Removed watcher connection: {}", conn);
+      }
+    }
+  }
+
   public String getAssociatedNoteId(NotebookSocket socket) {
     String associatedNoteId = null;
     synchronized (noteSocketMap) {

zeppelin-server/src/main/java/org/apache/zeppelin/socket/NotebookServer.java

@@ -503,6 +503,7 @@ public void onClose(Session session, CloseReason closeReason) {
 
   private void removeConnection(NotebookSocket notebookSocket) {
     connectionManager.removeConnection(notebookSocket);
+    connectionManager.removeWatcherConnection(notebookSocket);
     connectionManager.removeConnectionFromAllNote(notebookSocket);
     connectionManager.removeUserConnection(notebookSocket.getUser(), notebookSocket);
   }

zeppelin-server/src/test/java/org/apache/zeppelin/socket/ConnectionManagerTest.java

@@ -17,10 +17,21 @@
 package org.apache.zeppelin.socket;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
 
 import org.apache.zeppelin.conf.ZeppelinConfiguration;
 import org.apache.zeppelin.notebook.AuthorizationService;
+import org.apache.zeppelin.util.WatcherSecurityKey;
 import org.junit.jupiter.api.Test;
 
 class ConnectionManagerTest {
@@ -57,4 +68,104 @@ void checkMapGrowRemoveAll() {
     manager.removeConnectionFromAllNote(socket);
     assertEquals(0, manager.noteSocketMap.size());
   }
+
+  @Test
+  void removeWatcherConnectionCleansQueue() {
+    AuthorizationService authService = mock(AuthorizationService.class);
+
+    ConnectionManager manager = new ConnectionManager(authService, ZeppelinConfiguration.load());
+    NotebookSocket socket = mock(NotebookSocket.class);
+
+    manager.watcherSockets.add(socket);
+    assertEquals(1, manager.watcherSockets.size());
+
+    manager.removeWatcherConnection(socket);
+    assertEquals(0, manager.watcherSockets.size());
+  }
+
+  @Test
+  void removeWatcherConnectionWithMultipleWatchers() {
+    AuthorizationService authService = mock(AuthorizationService.class);
+
+    ConnectionManager manager = new ConnectionManager(authService, ZeppelinConfiguration.load());
+    NotebookSocket socket1 = mock(NotebookSocket.class);
+    NotebookSocket socket2 = mock(NotebookSocket.class);
+    NotebookSocket socket3 = mock(NotebookSocket.class);
+
+    // Add multiple watchers
+    manager.watcherSockets.add(socket1);
+    manager.watcherSockets.add(socket2);
+    manager.watcherSockets.add(socket3);
+    assertEquals(3, manager.watcherSockets.size());
+
+    // Remove only socket2
+    manager.removeWatcherConnection(socket2);
+    assertEquals(2, manager.watcherSockets.size());
+    assertTrue(manager.watcherSockets.contains(socket1));
+    assertFalse(manager.watcherSockets.contains(socket2));
+    assertTrue(manager.watcherSockets.contains(socket3));
+  }
+
+  @Test
+  void removeWatcherConnectionConcurrentTest() throws InterruptedException {
+    AuthorizationService authService = mock(AuthorizationService.class);
+    ConnectionManager manager = new ConnectionManager(authService, ZeppelinConfiguration.load());
+    
+    int threadCount = 10;
+    List<NotebookSocket> sockets = new ArrayList<>();
+    
+    // Create and add multiple watcher sockets
+    for (int i = 0; i < threadCount; i++) {
+      NotebookSocket socket = mock(NotebookSocket.class);
+      sockets.add(socket);
+      manager.watcherSockets.add(socket);
+    }
+    
+    assertEquals(threadCount, manager.watcherSockets.size());
+    
+    // Remove sockets concurrently
+    ExecutorService executor = Executors.newFixedThreadPool(threadCount);
+    CountDownLatch latch = new CountDownLatch(threadCount);
+    
+    for (NotebookSocket socket : sockets) {
+      executor.submit(() -> {
+        manager.removeWatcherConnection(socket);
+        latch.countDown();
+      });
+    }
+    
+    // Wait for all threads to complete
+    assertTrue(latch.await(5, TimeUnit.SECONDS));
+    executor.shutdown();
+    
+    // Verify all sockets were removed
+    assertEquals(0, manager.watcherSockets.size());
+  }
+
+  @Test
+  void switchConnectionToWatcherAndRemove() {
+    AuthorizationService authService = mock(AuthorizationService.class);
+    ConnectionManager manager = new ConnectionManager(authService, ZeppelinConfiguration.load());
+    
+    NotebookSocket socket = mock(NotebookSocket.class);
+    when(socket.getUser()).thenReturn("testUser");
+    when(socket.getHeader(WatcherSecurityKey.HTTP_HEADER)).thenReturn(WatcherSecurityKey.getKey());
+    
+    // Add socket as regular connection first
+    manager.addConnection(socket);
+    manager.addUserConnection("testUser", socket);
+    
+    // Switch to watcher
+    manager.switchConnectionToWatcher(socket);
+    
+    // Verify it's in watcher queue
+    assertTrue(manager.watcherSockets.contains(socket));
+    assertFalse(manager.connectedSockets.contains(socket));
+    
+    // Remove watcher connection
+    manager.removeWatcherConnection(socket);
+    
+    // Verify it's completely removed
+    assertFalse(manager.watcherSockets.contains(socket));
+  }
 }