From b978ce97895f30ed02130f9ca5a2f9f93f6f182f Mon Sep 17 00:00:00 2001 From: JangAyeon Date: Fri, 25 Sep 2026 12:30:46 +0900 Subject: [PATCH] [ZEPPELIN-6694] Reap dead websocket connections via pong tracking Since ZEPPELIN-6092, the heartbeat pings keep resetting Jetty's idle timer, so the idle timeout no longer detects dead clients and a dead session stays open until TCP retransmission gives up. Track liveness explicitly: register a per-session PongMessage handler, count pings sent since the last pong, and close a session once it reaches zeppelin.websocket.heartbeat.max.missed.pongs (default 3, <= 0 disables). Sessions that are handling a message are not reaped, because Jetty reads their pongs only after onMessage returns and a long-running op would otherwise look like a dead peer. --- conf/zeppelin-site.xml.template | 6 + docs/setup/operation/configuration.md | 6 + .../zeppelin/conf/ZeppelinConfiguration.java | 6 + .../zeppelin/socket/NotebookServer.java | 45 +++++- .../zeppelin/socket/NotebookSocket.java | 55 ++++++++ .../conf/ZeppelinConfigurationTest.java | 13 ++ .../socket/NotebookServerHeartbeatTest.java | 128 ++++++++++++++++++ .../zeppelin/socket/NotebookSocketTest.java | 42 ++++++ 8 files changed, 298 insertions(+), 3 deletions(-) diff --git a/conf/zeppelin-site.xml.template b/conf/zeppelin-site.xml.template index 4107a6799cf..3f1f7f8430a 100755 --- a/conf/zeppelin-site.xml.template +++ b/conf/zeppelin-site.xml.template @@ -571,6 +571,12 @@ 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. + + zeppelin.websocket.heartbeat.max.missed.pongs + 3 + 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. + + zeppelin.server.default.dir.allowed false diff --git a/docs/setup/operation/configuration.md b/docs/setup/operation/configuration.md index 4215222c40d..a6debaaa12c 100644 --- a/docs/setup/operation/configuration.md +++ b/docs/setup/operation/configuration.md @@ -418,6 +418,12 @@ Sources descending by priority: 60000 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. + +
ZEPPELIN_WEBSOCKET_HEARTBEAT_MAX_MISSED_PONGS
+
zeppelin.websocket.heartbeat.max.missed.pongs
+ 3 + 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. +
ZEPPELIN_SERVER_DEFAULT_DIR_ALLOWED
zeppelin.server.default.dir.allowed
diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/conf/ZeppelinConfiguration.java b/zeppelin-server/src/main/java/org/apache/zeppelin/conf/ZeppelinConfiguration.java index 01ad388ebac..6a7e1b7909f 100644 --- a/zeppelin-server/src/main/java/org/apache/zeppelin/conf/ZeppelinConfiguration.java +++ b/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"), diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/socket/NotebookServer.java b/zeppelin-server/src/main/java/org/apache/zeppelin/socket/NotebookServer.java index d55eb271967..aa1774dff84 100644 --- a/zeppelin-server/src/main/java/org/apache/zeppelin/socket/NotebookServer.java +++ b/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,12 +309,22 @@ 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). + * + *

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); @@ -320,10 +332,37 @@ void sendHeartbeat() { } } + /** + * 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) { diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/socket/NotebookSocket.java b/zeppelin-server/src/main/java/org/apache/zeppelin/socket/NotebookSocket.java index 57edf1d79b1..db70d4fc2c2 100644 --- a/zeppelin-server/src/main/java/org/apache/zeppelin/socket/NotebookSocket.java +++ b/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 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 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,8 +82,11 @@ 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) { @@ -76,6 +94,43 @@ public void sendPing() { } } + /** + * 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; } diff --git a/zeppelin-server/src/test/java/org/apache/zeppelin/conf/ZeppelinConfigurationTest.java b/zeppelin-server/src/test/java/org/apache/zeppelin/conf/ZeppelinConfigurationTest.java index f1e4d0d4daa..e230152b366 100644 --- a/zeppelin-server/src/test/java/org/apache/zeppelin/conf/ZeppelinConfigurationTest.java +++ b/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()); + } } diff --git a/zeppelin-server/src/test/java/org/apache/zeppelin/socket/NotebookServerHeartbeatTest.java b/zeppelin-server/src/test/java/org/apache/zeppelin/socket/NotebookServerHeartbeatTest.java index f932b1fa2dd..fd7bac091da 100644 --- a/zeppelin-server/src/test/java/org/apache/zeppelin/socket/NotebookServerHeartbeatTest.java +++ b/zeppelin-server/src/test/java/org/apache/zeppelin/socket/NotebookServerHeartbeatTest.java @@ -17,20 +17,42 @@ package org.apache.zeppelin.socket; import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.spy; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import java.io.IOException; +import java.util.HashMap; +import java.util.Map; import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.atomic.AtomicBoolean; + +import jakarta.websocket.CloseReason; +import jakarta.websocket.EndpointConfig; +import jakarta.websocket.MessageHandler; +import jakarta.websocket.PongMessage; +import jakarta.websocket.RemoteEndpoint; +import jakarta.websocket.Session; + import org.apache.zeppelin.MiniZeppelinServer; import org.apache.zeppelin.conf.ZeppelinConfiguration; import org.apache.zeppelin.notebook.AuthorizationService; +import org.apache.zeppelin.utils.CorsUtils; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; class NotebookServerHeartbeatTest { @@ -44,8 +66,13 @@ void tearDown() { } private NotebookServer buildNotebookServer(long heartbeatIntervalMs) { + return buildNotebookServer(heartbeatIntervalMs, 0); + } + + private NotebookServer buildNotebookServer(long heartbeatIntervalMs, int maxMissedPongs) { ZeppelinConfiguration zConf = mock(ZeppelinConfiguration.class); when(zConf.getWebsocketHeartbeatInterval()).thenReturn(heartbeatIntervalMs); + when(zConf.getWebsocketHeartbeatMaxMissedPongs()).thenReturn(maxMissedPongs); AuthorizationService authorizationService = mock(AuthorizationService.class); ConnectionManager connectionManager = new ConnectionManager(authorizationService, zConf); @@ -83,6 +110,107 @@ void sendHeartbeatContinuesWhenOneSocketThrows() { verify(healthy).sendPing(); } + @Test + void sendHeartbeatReapsSocketThatMissedTooManyPongs() { + NotebookServer server = buildNotebookServer(60000L, 3); + NotebookSocket dead = mock(NotebookSocket.class); + NotebookSocket alive = mock(NotebookSocket.class); + when(dead.getSessionId()).thenReturn("dead-session"); + when(dead.getPingsSinceLastPong()).thenReturn(3); + when(alive.getPingsSinceLastPong()).thenReturn(2); + server.getConnectionManager().addConnection(dead); + server.getConnectionManager().addConnection(alive); + + server.sendHeartbeat(); + + verify(dead).close(any(CloseReason.class)); + verify(dead, never()).sendPing(); + verify(alive).sendPing(); + verify(alive, never()).close(any(CloseReason.class)); + assertFalse(server.getConnectionManager().connectedSockets.contains(dead)); + assertTrue(server.getConnectionManager().connectedSockets.contains(alive)); + } + + @Test + void sendHeartbeatDoesNotReapWhenReapingDisabled() { + NotebookServer server = buildNotebookServer(60000L, 0); + NotebookSocket silent = mock(NotebookSocket.class); + when(silent.getPingsSinceLastPong()).thenReturn(100); + server.getConnectionManager().addConnection(silent); + + server.sendHeartbeat(); + + verify(silent).sendPing(); + verify(silent, never()).close(any(CloseReason.class)); + } + + @Test + void sendHeartbeatDoesNotReapSocketThatIsHandlingMessage() { + NotebookServer server = buildNotebookServer(60000L, 3); + NotebookSocket busy = mock(NotebookSocket.class); + when(busy.isHandlingMessage()).thenReturn(true); + when(busy.getPingsSinceLastPong()).thenReturn(5); + server.getConnectionManager().addConnection(busy); + + server.sendHeartbeat(); + + verify(busy, never()).close(any(CloseReason.class)); + verify(busy).sendPing(); + assertTrue(server.getConnectionManager().connectedSockets.contains(busy)); + } + + private Session openSession(NotebookServer server, String sessionId) throws IOException { + Session session = mock(Session.class); + when(session.getId()).thenReturn(sessionId); + when(session.getBasicRemote()).thenReturn(mock(RemoteEndpoint.Basic.class)); + Map headers = new HashMap<>(); + headers.put(CorsUtils.HEADER_ORIGIN, "http://localhost:8080"); + EndpointConfig config = mock(EndpointConfig.class); + when(config.getUserProperties()).thenReturn(headers); + server.onOpen(session, config); + return session; + } + + @Test + @SuppressWarnings("unchecked") + void onOpenRegistersPongHandlerThatResetsPingCount() throws IOException { + NotebookServer server = buildNotebookServer(60000L, 3); + Session session = openSession(server, "session-pong"); + NotebookSocket conn = server.getConnectionManager().connectedSockets.peek(); + assertNotNull(conn); + + ArgumentCaptor> captor = + ArgumentCaptor.forClass(MessageHandler.Whole.class); + verify(session).addMessageHandler(eq(PongMessage.class), captor.capture()); + + conn.sendPing(); + conn.sendPing(); + assertEquals(2, conn.getPingsSinceLastPong()); + + captor.getValue().onMessage(mock(PongMessage.class)); + assertEquals(0, conn.getPingsSinceLastPong()); + } + + @Test + void onMessageMarksSocketAsHandlingMessageUntilItReturns() throws IOException { + NotebookServer server = spy(buildNotebookServer(60000L, 3)); + notebookServer = server; + Session session = openSession(server, "session-busy"); + NotebookSocket conn = server.getConnectionManager().connectedSockets.peek(); + assertNotNull(conn); + + AtomicBoolean handlingDuringOp = new AtomicBoolean(); + doAnswer(invocation -> { + handlingDuringOp.set(conn.isHandlingMessage()); + return null; + }).when(server).onMessage(any(NotebookSocket.class), anyString()); + + server.onMessage(session, "{}"); + + assertTrue(handlingDuringOp.get()); + assertFalse(conn.isHandlingMessage()); + } + @Test void startHeartbeatSchedulerStartsWhenIntervalPositive() { NotebookServer server = buildNotebookServer(50L); diff --git a/zeppelin-server/src/test/java/org/apache/zeppelin/socket/NotebookSocketTest.java b/zeppelin-server/src/test/java/org/apache/zeppelin/socket/NotebookSocketTest.java index 4382e0dafe3..53447b5d163 100644 --- a/zeppelin-server/src/test/java/org/apache/zeppelin/socket/NotebookSocketTest.java +++ b/zeppelin-server/src/test/java/org/apache/zeppelin/socket/NotebookSocketTest.java @@ -17,6 +17,7 @@ package org.apache.zeppelin.socket; import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertEquals; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; @@ -27,6 +28,7 @@ import java.nio.ByteBuffer; import java.util.Collections; +import jakarta.websocket.CloseReason; import jakarta.websocket.RemoteEndpoint; import jakarta.websocket.Session; @@ -59,4 +61,44 @@ void sendPingSwallowsIOExceptionFromDeadSession() throws IOException { assertDoesNotThrow(notebookSocket::sendPing); } + + @Test + void sendPingCountsUnansweredPingsAndPongResetsCount() { + Session session = mock(Session.class); + when(session.getId()).thenReturn("session-3"); + when(session.getBasicRemote()).thenReturn(mock(RemoteEndpoint.Basic.class)); + NotebookSocket notebookSocket = new NotebookSocket(session, Collections.emptyMap()); + + notebookSocket.sendPing(); + notebookSocket.sendPing(); + assertEquals(2, notebookSocket.getPingsSinceLastPong()); + + notebookSocket.onPong(); + assertEquals(0, notebookSocket.getPingsSinceLastPong()); + } + + @Test + void failedPingStillCountsAsUnanswered() throws IOException { + Session session = mock(Session.class); + RemoteEndpoint.Basic basicRemote = mock(RemoteEndpoint.Basic.class); + when(session.getId()).thenReturn("session-4"); + when(session.getBasicRemote()).thenReturn(basicRemote); + doThrow(new IOException("broken pipe")).when(basicRemote).sendPing(any(ByteBuffer.class)); + NotebookSocket notebookSocket = new NotebookSocket(session, Collections.emptyMap()); + + notebookSocket.sendPing(); + + assertEquals(1, notebookSocket.getPingsSinceLastPong()); + } + + @Test + void closeSwallowsIOException() throws IOException { + Session session = mock(Session.class); + when(session.getId()).thenReturn("session-5"); + doThrow(new IOException("already closed")).when(session).close(any(CloseReason.class)); + NotebookSocket notebookSocket = new NotebookSocket(session, Collections.emptyMap()); + + assertDoesNotThrow(() -> notebookSocket.close( + new CloseReason(CloseReason.CloseCodes.GOING_AWAY, "test"))); + } }