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