[ZEPPELIN-6694] Reap dead websocket connections via pong tracking

conf/zeppelin-site.xml.template

@@ -571,6 +571,12 @@
   <description>Interval in milliseconds at which the server sends a websocket ping frame to each session to keep it alive. Defaults to 60000 (1 minute). Set to 0 or a negative value to disable server-initiated heartbeats.</description>
 </property>
 
+<property>
+  <name>zeppelin.websocket.heartbeat.max.missed.pongs</name>
+  <value>3</value>
+  <description>Number of consecutive heartbeat pings left unanswered (no pong) before the server closes the websocket session as dead. Defaults to 3. Set to 0 or a negative value to disable reaping.</description>
+</property>
+
 <property>
   <name>zeppelin.server.default.dir.allowed</name>
   <value>false</value>

docs/setup/operation/configuration.md

@@ -418,6 +418,12 @@ Sources descending by priority:
     <td>60000</td>
     <td>Interval(in milliseconds) at which the server sends a websocket ping frame to each session to keep it alive. Set to 0 or a negative value to disable server-initiated heartbeats.</td>
   </tr>
+  <tr>
+    <td><h6 class="properties">ZEPPELIN_WEBSOCKET_HEARTBEAT_MAX_MISSED_PONGS</h6></td>
+    <td><h6 class="properties">zeppelin.websocket.heartbeat.max.missed.pongs</h6></td>
+    <td>3</td>
+    <td>Number of consecutive heartbeat pings left unanswered (no pong) before the server closes the websocket session as dead. Set to 0 or a negative value to disable reaping. Has no effect when heartbeats are disabled.</td>
+  </tr>
   <tr>
     <td><h6 class="properties">ZEPPELIN_SERVER_DEFAULT_DIR_ALLOWED</h6></td>
     <td><h6 class="properties">zeppelin.server.default.dir.allowed</h6></td>

zeppelin-server/src/main/java/org/apache/zeppelin/conf/ZeppelinConfiguration.java

@@ -743,6 +743,10 @@ public long getWebsocketHeartbeatInterval() {
     return getLong(ConfVars.ZEPPELIN_WEBSOCKET_HEARTBEAT_INTERVAL);
   }
 
+  public int getWebsocketHeartbeatMaxMissedPongs() {
+    return getInt(ConfVars.ZEPPELIN_WEBSOCKET_HEARTBEAT_MAX_MISSED_PONGS);
+  }
+
   public String getJettyName() {
     return getString(ConfVars.ZEPPELIN_SERVER_JETTY_NAME);
   }
@@ -1105,6 +1109,8 @@ public enum ConfVars {
     // per-connection traffic low. 60s gives 5 pings within the 300s default idle window.
     // <= 0 disables server-initiated heartbeats.
     ZEPPELIN_WEBSOCKET_HEARTBEAT_INTERVAL("zeppelin.websocket.heartbeat.interval", 60000L),
+    // Consecutive unanswered pings before a session is closed as dead. <= 0 disables reaping.
+    ZEPPELIN_WEBSOCKET_HEARTBEAT_MAX_MISSED_PONGS("zeppelin.websocket.heartbeat.max.missed.pongs", 3),
     ZEPPELIN_WEBSOCKET_PARAGRAPH_STATUS_PROGRESS("zeppelin.websocket.paragraph_status_progress.enable", true),
     ZEPPELIN_SERVER_DEFAULT_DIR_ALLOWED("zeppelin.server.default.dir.allowed", false),
     ZEPPELIN_SERVER_XFRAME_OPTIONS("zeppelin.server.xframe.options", "SAMEORIGIN"),

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

@@ -50,6 +50,7 @@
 import jakarta.websocket.OnError;
 import jakarta.websocket.OnMessage;
 import jakarta.websocket.OnOpen;
+import jakarta.websocket.PongMessage;
 import jakarta.websocket.Session;
 import jakarta.websocket.server.ServerEndpoint;
 import org.apache.commons.lang3.StringUtils;
@@ -250,6 +251,7 @@ public void onOpen(Session session, EndpointConfig endpointConfig) throws IOExce
     if (checkOrigin(origin)) {
       NotebookSocket notebookSocket = sessionIdNotebookSocketMap
           .computeIfAbsent(session.getId(), unused -> new NotebookSocket(session, headers));
+      session.addMessageHandler(PongMessage.class, pong -> notebookSocket.onPong());
       onOpen(notebookSocket);
     } else {
       LOGGER.error("Websocket request is not allowed by {} settings. Origin: {}", ZEPPELIN_ALLOWED_ORIGINS,
@@ -307,23 +309,60 @@ public synchronized void stopHeartbeatScheduler() {
    * point of this heartbeat: it keeps connections alive even when the client-side application
    * keep-alive timer is throttled or stopped (e.g. a backgrounded browser tab). A single
    * session failing to receive a ping must not stop the remaining sessions from being pinged.
-   * Pong responses are not tracked; once the heartbeat is enabled, Jetty's idle timeout no
-   * longer determines connection liveness (see ZEPPELIN-6694).
+   *
+   * <p>Because these writes keep resetting Jetty's idle timer, the idle timeout can no longer
+   * detect dead clients. Liveness is therefore tracked explicitly (ZEPPELIN-6694): a session
+   * that has left {@code zeppelin.websocket.heartbeat.max.missed.pongs} consecutive pings
+   * unanswered is closed instead of being pinged again. A session that is still handling a
+   * message is never closed here, because Jetty reads its pongs only after onMessage returns.
    */
   void sendHeartbeat() {
+    int maxMissedPongs = zConf.getWebsocketHeartbeatMaxMissedPongs();
     for (NotebookSocket conn : connectionManager.connectedSockets) {
       try {
+        if (maxMissedPongs > 0 && !conn.isHandlingMessage()
+            && conn.getPingsSinceLastPong() >= maxMissedPongs) {
+          reapDeadConnection(conn, maxMissedPongs);
+          continue;
+        }
         conn.sendPing();
       } catch (RuntimeException e) {
         LOGGER.warn("Failed to send heartbeat ping to {}", conn, e);
       }
     }
   }
 
+  /**
+   * Removes a connection that stopped answering pings and closes its session. The connection is
+   * dropped from all bookkeeping before closing, because a dead peer never completes the close
+   * handshake and {@link #onClose} may therefore arrive late or not at all. Removing it from
+   * {@code sessionIdNotebookSocketMap} first also makes a later {@code onClose} a no-op.
+   */
+  private void reapDeadConnection(NotebookSocket conn, int maxMissedPongs) {
+    LOGGER.warn("Closing websocket to {}: {} consecutive heartbeat pings unanswered, last pong at {}",
+        conn, maxMissedPongs, new Date(conn.getLastPongTimestamp()));
+    String sessionId = conn.getSessionId();
+    if (sessionId != null) {
+      sessionIdNotebookSocketMap.remove(sessionId);
+    }
+    removeConnection(conn);
+    conn.close(new CloseReason(CloseReason.CloseCodes.GOING_AWAY,
+        "No pong received for " + maxMissedPongs + " consecutive pings"));
+  }
+
   @OnMessage
   public void onMessage(Session session, String msg) {
     NotebookSocket conn = sessionIdNotebookSocketMap.get(session.getId());
-    onMessage(conn, msg);
+    if (conn == null) {
+      onMessage(conn, msg);
+      return;
+    }
+    conn.setHandlingMessage(true);
+    try {
+      onMessage(conn, msg);
+    } finally {
+      conn.setHandlingMessage(false);
+    }
   }
 
   public void onMessage(NotebookSocket conn, String msg) {

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

@@ -24,7 +24,9 @@
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
 
+import jakarta.websocket.CloseReason;
 import jakarta.websocket.Session;
 
 /**
@@ -42,13 +44,26 @@ public class NotebookSocket {
   private Map<String, Object> headers;
   private String user;
 
+  // Liveness tracking (ZEPPELIN-6694). Written from the heartbeat thread (sendPing) and from
+  // the websocket container thread (onPong), hence atomic/volatile.
+  private final AtomicInteger unansweredPings = new AtomicInteger();
+  private volatile long lastPongTimestamp;
+  // True while onMessage is running for this session. Jetty reads the next frame (pongs
+  // included) only after onMessage returns, so pongs pile up unread during a long operation.
+  private volatile boolean handlingMessage;
+
   public NotebookSocket(Session session, Map<String, Object> headers) {
     this.session = session;
     this.headers = headers;
     this.user = StringUtils.EMPTY;
+    this.lastPongTimestamp = System.currentTimeMillis();
     LOGGER.debug("NotebookSocket created for session: {}", session.getId());
   }
 
+  public String getSessionId() {
+    return session.getId();
+  }
+
   public String getHeader(String key) {
     return String.valueOf(headers.get(key));
   }
@@ -67,15 +82,55 @@ public void send(String serializeMessage) throws IOException {
    * session resets Jetty's idle timeout as well as any intermediate proxy's idle timer, so no
    * application-level handling is required on the client. Exceptions are swallowed and logged
    * so a single dead session cannot break the caller's heartbeat loop over all sessions.
+   * Every call counts as one outstanding ping until {@link #onPong()} is called, including
+   * calls whose write failed, since a failed write is itself a sign the peer is gone.
    */
   public void sendPing() {
+    unansweredPings.incrementAndGet();
     try {
       session.getBasicRemote().sendPing(PING_PAYLOAD);
     } catch (IOException | IllegalArgumentException | IllegalStateException e) {
       LOGGER.warn("Failed to send heartbeat ping to session {}: {}", session.getId(), e.toString());
     }
   }
 
+  /**
+   * Records a pong frame from the peer. Any pong proves the connection is alive, so the
+   * outstanding-ping counter is reset rather than decremented.
+   */
+  public void onPong() {
+    lastPongTimestamp = System.currentTimeMillis();
+    unansweredPings.set(0);
+  }
+
+  public int getPingsSinceLastPong() {
+    return unansweredPings.get();
+  }
+
+  public long getLastPongTimestamp() {
+    return lastPongTimestamp;
+  }
+
+  public boolean isHandlingMessage() {
+    return handlingMessage;
+  }
+
+  public void setHandlingMessage(boolean handlingMessage) {
+    this.handlingMessage = handlingMessage;
+  }
+
+  /**
+   * Closes the underlying session. Exceptions are swallowed and logged because this is used to
+   * reap connections that are already presumed dead.
+   */
+  public void close(CloseReason closeReason) {
+    try {
+      session.close(closeReason);
+    } catch (IOException | IllegalStateException e) {
+      LOGGER.debug("Failed to close session {}: {}", session.getId(), e.toString());
+    }
+  }
+
   public String getUser() {
     return user;
   }

zeppelin-server/src/test/java/org/apache/zeppelin/conf/ZeppelinConfigurationTest.java

@@ -185,4 +185,17 @@ void getWebsocketHeartbeatIntervalDisabledTest() {
     zConf.setProperty(ConfVars.ZEPPELIN_WEBSOCKET_HEARTBEAT_INTERVAL.getVarName(), "0");
     assertEquals(0L, zConf.getWebsocketHeartbeatInterval());
   }
+
+  @Test
+  void getWebsocketHeartbeatMaxMissedPongsDefaultTest() {
+    ZeppelinConfiguration zConf = ZeppelinConfiguration.load("zeppelin-test-site.xml");
+    assertEquals(3, zConf.getWebsocketHeartbeatMaxMissedPongs());
+  }
+
+  @Test
+  void getWebsocketHeartbeatMaxMissedPongsOverrideTest() {
+    ZeppelinConfiguration zConf = ZeppelinConfiguration.load("zeppelin-test-site.xml");
+    zConf.setProperty(ConfVars.ZEPPELIN_WEBSOCKET_HEARTBEAT_MAX_MISSED_PONGS.getVarName(), "5");
+    assertEquals(5, zConf.getWebsocketHeartbeatMaxMissedPongs());
+  }
 }