From 8bfdf6005d0123d9332fe68214e707ccbf295bac Mon Sep 17 00:00:00 2001 From: HwangRock Date: Sat, 22 Aug 2026 01:44:49 +0900 Subject: [PATCH] [ZEPPELIN-6092] Send server-initiated websocket ping frames to keep connections alive Keep-alive is client-driven only: both UIs send an application-level {"op":"PING"} every 10 seconds and the server never writes first. When the client timer stops -- a backgrounded tab under Chrome's intensive throttling, a discarded tab, a sleeping laptop -- nothing resets the idle timer and the connection dies. Send a WebSocket protocol ping frame from the server on a schedule. The peer answers automatically per RFC 6455 section 5.5.2, so no client change is needed, and writing to the session resets Jetty's idle timeout along with any intermediate proxy's idle timer. Also make the idle timeout configurable. setupNotebookServer() never called setDefaultMaxSessionIdleTimeout(), so the effective value was whatever Jetty defaulted to -- 300000ms under Jetty 9, 30 seconds under Jetty 11. Measured with a client that sends nothing after the handshake, against Jetty 11's 30s default: before: closed at 30.0s, code=1001 'Connection Idle Timeout' after: still connected at 120.0s --- conf/zeppelin-env.sh.template | 2 + conf/zeppelin-site.xml.template | 12 ++ docs/setup/operation/configuration.md | 12 ++ .../zeppelin/conf/ZeppelinConfiguration.java | 15 +++ .../zeppelin/server/ZeppelinServer.java | 1 + .../zeppelin/socket/NotebookServer.java | 74 ++++++++++++ .../zeppelin/socket/NotebookSocket.java | 21 ++++ .../conf/ZeppelinConfigurationTest.java | 33 ++++++ .../socket/NotebookServerHeartbeatTest.java | 109 ++++++++++++++++++ .../zeppelin/socket/NotebookSocketTest.java | 62 ++++++++++ 10 files changed, 341 insertions(+) create mode 100644 zeppelin-server/src/test/java/org/apache/zeppelin/socket/NotebookServerHeartbeatTest.java create mode 100644 zeppelin-server/src/test/java/org/apache/zeppelin/socket/NotebookSocketTest.java diff --git a/conf/zeppelin-env.sh.template b/conf/zeppelin-env.sh.template index e8160b563c9..38b60cfa4d7 100644 --- a/conf/zeppelin-env.sh.template +++ b/conf/zeppelin-env.sh.template @@ -100,6 +100,8 @@ # export ZEPPELIN_SPARK_IMPORTIMPLICIT # Import implicits, UDF collection, and sql if set true. true by default. # export ZEPPELIN_SPARK_MAXRESULT # Max number of Spark SQL result to display. 1000 by default. # export ZEPPELIN_WEBSOCKET_MAX_TEXT_MESSAGE_SIZE # Size in characters of the maximum text message to be received by websocket. Defaults to 1024000 +# export ZEPPELIN_WEBSOCKET_IDLE_TIMEOUT # Time in milliseconds before an idle websocket session is closed. Defaults to 300000 +# export ZEPPELIN_WEBSOCKET_HEARTBEAT_INTERVAL # Interval in milliseconds at which the server sends a websocket ping frame to keep each session alive. Defaults to 60000. Set to 0 or a negative value to disable. #### HBase interpreter configuration #### diff --git a/conf/zeppelin-site.xml.template b/conf/zeppelin-site.xml.template index d04aee833e6..4107a6799cf 100755 --- a/conf/zeppelin-site.xml.template +++ b/conf/zeppelin-site.xml.template @@ -559,6 +559,18 @@ Size in characters of the maximum text message to be received by websocket. Defaults to 10240000 + + zeppelin.websocket.idle.timeout + 300000 + Time in milliseconds before an idle websocket session is closed. Defaults to 300000 (5 minutes) + + + + zeppelin.websocket.heartbeat.interval + 60000 + 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.server.default.dir.allowed false diff --git a/docs/setup/operation/configuration.md b/docs/setup/operation/configuration.md index 1e994e0263e..4215222c40d 100644 --- a/docs/setup/operation/configuration.md +++ b/docs/setup/operation/configuration.md @@ -406,6 +406,18 @@ Sources descending by priority: 1024000 Size(in characters) of the maximum text message that can be received by websocket. + +
ZEPPELIN_WEBSOCKET_IDLE_TIMEOUT
+
zeppelin.websocket.idle.timeout
+ 300000 + Time(in milliseconds) before an idle websocket session is closed. + + +
ZEPPELIN_WEBSOCKET_HEARTBEAT_INTERVAL
+
zeppelin.websocket.heartbeat.interval
+ 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_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 b2e15160b52..01ad388ebac 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 @@ -735,6 +735,14 @@ public String getWebsocketMaxTextMessageSize() { return getString(ConfVars.ZEPPELIN_WEBSOCKET_MAX_TEXT_MESSAGE_SIZE); } + public long getWebsocketIdleTimeout() { + return getLong(ConfVars.ZEPPELIN_WEBSOCKET_IDLE_TIMEOUT); + } + + public long getWebsocketHeartbeatInterval() { + return getLong(ConfVars.ZEPPELIN_WEBSOCKET_HEARTBEAT_INTERVAL); + } + public String getJettyName() { return getString(ConfVars.ZEPPELIN_SERVER_JETTY_NAME); } @@ -1090,6 +1098,13 @@ public enum ConfVars { ZEPPELIN_CREDENTIALS_PERSIST("zeppelin.credentials.persist", true), ZEPPELIN_CREDENTIALS_ENCRYPT_KEY("zeppelin.credentials.encryptKey", null), ZEPPELIN_WEBSOCKET_MAX_TEXT_MESSAGE_SIZE("zeppelin.websocket.max.text.message.size", "10240000"), + ZEPPELIN_WEBSOCKET_IDLE_TIMEOUT("zeppelin.websocket.idle.timeout", 300000L), + // Server-initiated websocket protocol ping interval, in milliseconds. Writing a ping frame + // resets the Jetty idle timer (see ZEPPELIN_WEBSOCKET_IDLE_TIMEOUT above) and any intermediate + // proxy's idle timer, so the default must stay well below that timeout while still keeping + // 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), 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/server/ZeppelinServer.java b/zeppelin-server/src/main/java/org/apache/zeppelin/server/ZeppelinServer.java index b3f78816aec..6dece8a13d3 100644 --- a/zeppelin-server/src/main/java/org/apache/zeppelin/server/ZeppelinServer.java +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/server/ZeppelinServer.java @@ -475,6 +475,7 @@ private void setupNotebookServer(WebAppContext webapp) { JakartaWebSocketServletContainerInitializer .configure(webapp, (servletContext, wsContainer) -> { wsContainer.setDefaultMaxTextMessageBufferSize(Integer.parseInt(maxTextMessageSize)); + wsContainer.setDefaultMaxSessionIdleTimeout(zConf.getWebsocketIdleTimeout()); wsContainer.addEndpoint(ServerEndpointConfig.Builder.create(NotebookServer.class, "/ws") .configurator(new SessionConfigurator(sharedServiceLocator)).build()); }); 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 85a552e7f45..7f0a02039e8 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 @@ -39,6 +39,8 @@ import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicReference; import jakarta.inject.Inject; import jakarta.inject.Provider; @@ -146,6 +148,11 @@ String getKey() { private final ExecutorService executorService = Executors.newFixedThreadPool(10); + // Package-private (not private) so NotebookServerHeartbeatTest can observe scheduler + // lifecycle without exposing it as part of the public API. + ScheduledExecutorService heartbeatScheduler; + private Thread heartbeatShutdownHook; + // TODO(jl): This will be removed by handling session directly private final Map sessionIdNotebookSocketMap = Metrics.gaugeMapSize("zeppelin_session_id_notebook_sockets", Tags.empty(), new ConcurrentHashMap<>()); private ConnectionManager connectionManager; @@ -265,6 +272,73 @@ public void onOpen(Session session, EndpointConfig endpointConfig) throws IOExce public void onOpen(NotebookSocket conn) { connectionManager.addConnection(conn); + startHeartbeatScheduler(); + } + + /** + * Starts the websocket heartbeat scheduler on first use. Idempotent: a second call while the + * scheduler is already running is a no-op. zConf and connectionManager are both required and + * are set via setter injection before any real connection can open, so starting lazily here + * (rather than from the injected setters, whose call order is not guaranteed) is safe. + */ + synchronized void startHeartbeatScheduler() { + if (heartbeatScheduler != null) { + return; + } + long intervalMs = zConf.getWebsocketHeartbeatInterval(); + if (intervalMs <= 0) { + LOGGER.info("Websocket heartbeat is disabled (zeppelin.websocket.heartbeat.interval={})", intervalMs); + return; + } + heartbeatScheduler = Executors.newSingleThreadScheduledExecutor(r -> { + Thread thread = new Thread(r, "NotebookServer-Heartbeat"); + thread.setDaemon(true); + return thread; + }); + heartbeatScheduler.scheduleAtFixedRate( + this::sendHeartbeat, intervalMs, intervalMs, TimeUnit.MILLISECONDS); + heartbeatShutdownHook = new Thread(this::stopHeartbeatScheduler); + Runtime.getRuntime().addShutdownHook(heartbeatShutdownHook); + LOGGER.info("Started websocket heartbeat scheduler with interval {} ms", intervalMs); + } + + /** + * Stops the websocket heartbeat scheduler, if running, and deregisters its shutdown hook so + * repeated start/stop cycles do not accumulate hooks. Safe to call multiple times and safe + * to call when the scheduler was never started. + */ + synchronized void stopHeartbeatScheduler() { + if (heartbeatScheduler != null) { + heartbeatScheduler.shutdownNow(); + heartbeatScheduler = null; + } + if (heartbeatShutdownHook != null && Thread.currentThread() != heartbeatShutdownHook) { + try { + Runtime.getRuntime().removeShutdownHook(heartbeatShutdownHook); + } catch (IllegalStateException e) { + // JVM is already shutting down; the hook will simply run (as a harmless no-op). + } + heartbeatShutdownHook = null; + } + } + + /** + * Sends a WebSocket protocol ping frame to every connected session. Writing to a session + * resets Jetty's idle timeout (and any intermediate proxy's idle timer), which is the whole + * 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. + */ + void sendHeartbeat() { + synchronized (connectionManager.connectedSockets) { + for (NotebookSocket conn : connectionManager.connectedSockets) { + try { + conn.sendPing(); + } catch (RuntimeException e) { + LOGGER.warn("Failed to send heartbeat ping to {}", conn, e); + } + } + } } @OnMessage 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 1805ce456f1..57edf1d79b1 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 @@ -22,6 +22,7 @@ import org.slf4j.LoggerFactory; import java.io.IOException; +import java.nio.ByteBuffer; import java.util.Map; import jakarta.websocket.Session; @@ -32,6 +33,11 @@ public class NotebookSocket { private static final Logger LOGGER = LoggerFactory.getLogger(NotebookSocket.class); + // WebSocket protocol ping frames (RFC 6455 5.5.2) carry no meaningful payload here, so a + // single empty, effectively immutable (zero remaining bytes) buffer can be reused for every + // send instead of allocating one per heartbeat tick. + private static final ByteBuffer PING_PAYLOAD = ByteBuffer.allocate(0); + private Session session; private Map headers; private String user; @@ -55,6 +61,21 @@ public void send(String serializeMessage) throws IOException { }); } + /** + * Sends a WebSocket protocol ping frame to keep this connection alive. The peer's WebSocket + * implementation answers automatically with a pong (RFC 6455 5.5.2), and writing to the + * 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. + */ + public void sendPing() { + try { + session.getBasicRemote().sendPing(PING_PAYLOAD); + } catch (IOException | IllegalArgumentException | IllegalStateException e) { + LOGGER.warn("Failed to send heartbeat ping to 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 a5cb0037fd0..f1e4d0d4daa 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 @@ -152,4 +152,37 @@ void checkParseException() { // then assertEquals(12345, zConf.getServerPort()); } + + @Test + void getWebsocketIdleTimeoutDefaultTest() { + ZeppelinConfiguration zConf = ZeppelinConfiguration.load("zeppelin-test-site.xml"); + assertEquals(300000L, zConf.getWebsocketIdleTimeout()); + } + + @Test + void getWebsocketIdleTimeoutOverrideTest() { + ZeppelinConfiguration zConf = ZeppelinConfiguration.load("zeppelin-test-site.xml"); + zConf.setProperty(ConfVars.ZEPPELIN_WEBSOCKET_IDLE_TIMEOUT.getVarName(), "600000"); + assertEquals(600000L, zConf.getWebsocketIdleTimeout()); + } + + @Test + void getWebsocketHeartbeatIntervalDefaultTest() { + ZeppelinConfiguration zConf = ZeppelinConfiguration.load("zeppelin-test-site.xml"); + assertEquals(60000L, zConf.getWebsocketHeartbeatInterval()); + } + + @Test + void getWebsocketHeartbeatIntervalOverrideTest() { + ZeppelinConfiguration zConf = ZeppelinConfiguration.load("zeppelin-test-site.xml"); + zConf.setProperty(ConfVars.ZEPPELIN_WEBSOCKET_HEARTBEAT_INTERVAL.getVarName(), "30000"); + assertEquals(30000L, zConf.getWebsocketHeartbeatInterval()); + } + + @Test + void getWebsocketHeartbeatIntervalDisabledTest() { + ZeppelinConfiguration zConf = ZeppelinConfiguration.load("zeppelin-test-site.xml"); + zConf.setProperty(ConfVars.ZEPPELIN_WEBSOCKET_HEARTBEAT_INTERVAL.getVarName(), "0"); + assertEquals(0L, zConf.getWebsocketHeartbeatInterval()); + } } 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 new file mode 100644 index 00000000000..6f53668261b --- /dev/null +++ b/zeppelin-server/src/test/java/org/apache/zeppelin/socket/NotebookServerHeartbeatTest.java @@ -0,0 +1,109 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.zeppelin.socket; + +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import org.apache.zeppelin.conf.ZeppelinConfiguration; +import org.apache.zeppelin.notebook.AuthorizationService; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; + +class NotebookServerHeartbeatTest { + + private NotebookServer notebookServer; + + @AfterEach + void tearDown() { + if (notebookServer != null) { + notebookServer.stopHeartbeatScheduler(); + } + } + + private NotebookServer buildNotebookServer(long heartbeatIntervalMs) { + ZeppelinConfiguration zConf = mock(ZeppelinConfiguration.class); + when(zConf.getWebsocketHeartbeatInterval()).thenReturn(heartbeatIntervalMs); + AuthorizationService authorizationService = mock(AuthorizationService.class); + ConnectionManager connectionManager = new ConnectionManager(authorizationService, zConf); + + notebookServer = new NotebookServer(); + notebookServer.setZeppelinConfiguration(zConf); + notebookServer.setConnectionManager(connectionManager); + return notebookServer; + } + + @Test + void sendHeartbeatSendsPingToEveryConnectedSocket() { + NotebookServer server = buildNotebookServer(60000L); + NotebookSocket first = mock(NotebookSocket.class); + NotebookSocket second = mock(NotebookSocket.class); + server.getConnectionManager().addConnection(first); + server.getConnectionManager().addConnection(second); + + server.sendHeartbeat(); + + verify(first).sendPing(); + verify(second).sendPing(); + } + + @Test + void sendHeartbeatContinuesWhenOneSocketThrows() { + NotebookServer server = buildNotebookServer(60000L); + NotebookSocket failing = mock(NotebookSocket.class); + NotebookSocket healthy = mock(NotebookSocket.class); + doThrow(new RuntimeException("connection reset")).when(failing).sendPing(); + server.getConnectionManager().addConnection(failing); + server.getConnectionManager().addConnection(healthy); + + assertDoesNotThrow(server::sendHeartbeat); + + verify(healthy).sendPing(); + } + + @Test + void startHeartbeatSchedulerStartsWhenIntervalPositive() { + NotebookServer server = buildNotebookServer(50L); + + server.startHeartbeatScheduler(); + + assertNotNull(server.heartbeatScheduler); + } + + @Test + void startHeartbeatSchedulerDoesNotStartWhenIntervalIsZero() { + NotebookServer server = buildNotebookServer(0L); + + server.startHeartbeatScheduler(); + + assertNull(server.heartbeatScheduler); + } + + @Test + void startHeartbeatSchedulerDoesNotStartWhenIntervalIsNegative() { + NotebookServer server = buildNotebookServer(-1L); + + server.startHeartbeatScheduler(); + + assertNull(server.heartbeatScheduler); + } +} 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 new file mode 100644 index 00000000000..4382e0dafe3 --- /dev/null +++ b/zeppelin-server/src/test/java/org/apache/zeppelin/socket/NotebookSocketTest.java @@ -0,0 +1,62 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.zeppelin.socket; + +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.util.Collections; + +import jakarta.websocket.RemoteEndpoint; +import jakarta.websocket.Session; + +import org.junit.jupiter.api.Test; + +class NotebookSocketTest { + + @Test + void sendPingWritesEmptyPingFrameToBasicRemote() throws IOException { + Session session = mock(Session.class); + RemoteEndpoint.Basic basicRemote = mock(RemoteEndpoint.Basic.class); + when(session.getId()).thenReturn("session-1"); + when(session.getBasicRemote()).thenReturn(basicRemote); + NotebookSocket notebookSocket = new NotebookSocket(session, Collections.emptyMap()); + + notebookSocket.sendPing(); + + verify(basicRemote).sendPing(any(ByteBuffer.class)); + } + + @Test + void sendPingSwallowsIOExceptionFromDeadSession() throws IOException { + Session session = mock(Session.class); + RemoteEndpoint.Basic basicRemote = mock(RemoteEndpoint.Basic.class); + when(session.getId()).thenReturn("session-2"); + when(session.getBasicRemote()).thenReturn(basicRemote); + doThrow(new IOException("session already closed")) + .when(basicRemote).sendPing(any(ByteBuffer.class)); + NotebookSocket notebookSocket = new NotebookSocket(session, Collections.emptyMap()); + + assertDoesNotThrow(notebookSocket::sendPing); + } +}