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")));
+ }
}