This is an automated email from the ASF dual-hosted git repository.
tbonelee pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/zeppelin.git
The following commit(s) were added to refs/heads/master by this push:
new d247e442a4 [ZEPPELIN-6092] Send server-initiated websocket ping frames
to keep connections alive
d247e442a4 is described below
commit d247e442a409ae1100df76126540f45a0627bd0f
Author: HwangRock <[email protected]>
AuthorDate: Fri Sep 4 23:29:40 2026 +0900
[ZEPPELIN-6092] Send server-initiated websocket ping frames to keep
connections alive
### What is this PR for?
Zeppelin's websocket keep-alive is client-driven only. Both UIs send an
application-level `{"op":"PING"}` every 10 seconds
(`websocket-event.factory.js`, `message.ts`) and the server answers nothing —
`case PING:` in `NotebookServer.onMessage()` is a bare `break`. The server
never writes first, so a connection survives only as long as the client's timer
keeps firing.
That timer is not reliable. Chrome's intensive throttling drops
background-tab timers to once per minute once the tab has been hidden for a few
seconds, the chain count reaches 5, and no WebRTC is in use — all of which a
10-second `setInterval` satisfies within a minute. An open websocket is not an
exemption; only WebRTC is. A discarded tab or a sleeping laptop stops the timer
outright. When it stops, nothing resets the idle timer and the connection dies.
This PR has the server send a WebSocket protocol ping frame on a schedule.
Per RFC 6455 section 5.5.2 the peer answers with a pong automatically, so no
client-side change is required, and writing to the session resets Jetty's idle
timeout — `SocketChannelEndPoint.flush()` calls `IdleTimeout.notIdle()` — along
with any intermediate proxy's idle timer. This is what the nginx websocket
proxying guide recommends as well: "the proxied server can be configured to
periodically send WebSocket [...]
The PR also makes the idle timeout itself configurable.
`setupNotebookServer()` in `ZeppelinServer` sets the text message buffer size
but never calls `setDefaultMaxSessionIdleTimeout()`, so the effective value has
always been whatever Jetty defaults to. That default is not stable across
versions: `WebSocketPolicy` used 300000ms under Jetty 9, while
`WebSocketConstants.DEFAULT_IDLE_TIMEOUT` is 30 seconds under Jetty 11. The new
key restores an explicit value and gives operators a knob.
Two new keys, documented in `zeppelin-site.xml.template`,
`zeppelin-env.sh.template`, and `docs/setup/operation/configuration.md`:
| Key | Default | Meaning |
|---|---|---|
| `zeppelin.websocket.heartbeat.interval` | 60000 | Interval in ms between
server-initiated ping frames. `0` or negative disables the heartbeat. |
| `zeppelin.websocket.idle.timeout` | 300000 | Idle timeout in ms applied
to websocket sessions. |
`NotebookSocket.sendPing()` swallows and logs its exceptions so one dead
session cannot break the loop over the others. The scheduler runs on a single
daemon thread and starts lazily on the first connection.
### What type of PR is it?
Improvement
### What is the Jira issue?
[ZEPPELIN-6092](https://issues.apache.org/jira/browse/ZEPPELIN-6092)
### How should this be tested?
Unit tests cover the ping send path, the disabled-when-non-positive case,
isolation of a failing session, and the two configuration keys.
End to end, with a client that sends nothing after the handshake, against
Jetty 11's 30-second default idle timeout:
Before
```
[+0000.0s] connected (HTTP/1.1 101 Switching Protocols)
[+0030.0s] closed code=1001 reason='Connection Idle Timeout'
[+0030.0s] RESULT [A-direct]: closed code=1001 reason='Connection Idle
Timeout'
```
```
ERROR NotebookServer.java[onError] - Error in WebSocket Session to
/[0:0:0:0:0:0:0:1]:59302
org.eclipse.jetty.websocket.core.exception.WebSocketTimeoutException:
Connection Idle Timeout
at
org.eclipse.jetty.websocket.core.internal.WebSocketConnection.onIdleExpired(WebSocketConnection.java:242)
at
org.eclipse.jetty.io.IdleTimeout.checkIdleTimeout(IdleTimeout.java:170)
```
After — same client, same 30-second idle timeout, heartbeat pinned to 10
seconds
```
[+0000.0s] connected (HTTP/1.1 101 Switching Protocols)
[+0120.0s] RESULT [A-direct]: still connected after 120s (no disconnect)
```
```
INFO NotebookServer.java[startHeartbeatScheduler] - Started websocket
heartbeat scheduler with interval 10000 ms
```
Zero `WebSocketTimeoutException` in the server log across the 120-second
window. The idle timeout is identical between the two runs, so the difference
comes from the heartbeat writes alone.
To reproduce: set `zeppelin.websocket.idle.timeout` to `30000` and
`zeppelin.websocket.heartbeat.interval` to `10000`, open a websocket to `/ws`,
send nothing after the handshake, and watch for a close past the 30-second mark.
### Questions:
* Does the license files need updating? No
* Are there breaking changes for older versions? No. The heartbeat uses
protocol ping frames, which every websocket client answers automatically. No
message op was added or changed, and the existing client-driven `PING` path is
untouched.
* Does this need documentation? Yes — both keys are documented in the
config template, the env template, and the configuration docs.
Closes #5432 from HwangRock/ZEPPELIN-6092.
Signed-off-by: ChanHo Lee <[email protected]>
---
conf/zeppelin-env.sh.template | 2 +
conf/zeppelin-site.xml.template | 12 +++
docs/setup/operation/configuration.md | 12 +++
.../zeppelin/conf/ZeppelinConfiguration.java | 15 +++
.../org/apache/zeppelin/server/ZeppelinServer.java | 1 +
.../org/apache/zeppelin/socket/NotebookServer.java | 77 +++++++++++++++
.../org/apache/zeppelin/socket/NotebookSocket.java | 21 ++++
.../zeppelin/conf/ZeppelinConfigurationTest.java | 33 +++++++
.../socket/NotebookServerHeartbeatTest.java | 109 +++++++++++++++++++++
.../apache/zeppelin/socket/NotebookSocketTest.java | 62 ++++++++++++
10 files changed, 344 insertions(+)
diff --git a/conf/zeppelin-env.sh.template b/conf/zeppelin-env.sh.template
index e8160b563c..38b60cfa4d 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 d04aee833e..4107a6799c 100755
--- a/conf/zeppelin-site.xml.template
+++ b/conf/zeppelin-site.xml.template
@@ -559,6 +559,18 @@
<description>Size in characters of the maximum text message to be received
by websocket. Defaults to 10240000</description>
</property>
+<property>
+ <name>zeppelin.websocket.idle.timeout</name>
+ <value>300000</value>
+ <description>Time in milliseconds before an idle websocket session is
closed. Defaults to 300000 (5 minutes)</description>
+</property>
+
+<property>
+ <name>zeppelin.websocket.heartbeat.interval</name>
+ <value>60000</value>
+ <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.server.default.dir.allowed</name>
<value>false</value>
diff --git a/docs/setup/operation/configuration.md
b/docs/setup/operation/configuration.md
index 1e994e0263..4215222c40 100644
--- a/docs/setup/operation/configuration.md
+++ b/docs/setup/operation/configuration.md
@@ -406,6 +406,18 @@ Sources descending by priority:
<td>1024000</td>
<td>Size(in characters) of the maximum text message that can be received
by websocket.</td>
</tr>
+ <tr>
+ <td><h6 class="properties">ZEPPELIN_WEBSOCKET_IDLE_TIMEOUT</h6></td>
+ <td><h6 class="properties">zeppelin.websocket.idle.timeout</h6></td>
+ <td>300000</td>
+ <td>Time(in milliseconds) before an idle websocket session is closed.</td>
+ </tr>
+ <tr>
+ <td><h6 class="properties">ZEPPELIN_WEBSOCKET_HEARTBEAT_INTERVAL</h6></td>
+ <td><h6 class="properties">zeppelin.websocket.heartbeat.interval</h6></td>
+ <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_SERVER_DEFAULT_DIR_ALLOWED</h6></td>
<td><h6 class="properties">zeppelin.server.default.dir.allowed</h6></td>
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 b2e15160b5..01ad388eba 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 class ZeppelinConfiguration {
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 class ZeppelinConfiguration {
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 b3f78816ae..6dece8a13d 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 @@ public class ZeppelinServer implements AutoCloseable {
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 85a552e7f4..cf4c91e593 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.Set;
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,12 @@ public class NotebookServer implements
AngularObjectRegistryListener,
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;
+ private boolean heartbeatInitialized;
+
// TODO(jl): This will be removed by handling session directly
private final Map<String, NotebookSocket> sessionIdNotebookSocketMap =
Metrics.gaugeMapSize("zeppelin_session_id_notebook_sockets", Tags.empty(), new
ConcurrentHashMap<>());
private ConnectionManager connectionManager;
@@ -265,6 +273,75 @@ public class NotebookServer implements
AngularObjectRegistryListener,
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 (heartbeatInitialized) {
+ return;
+ }
+ heartbeatInitialized = true;
+ 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;
+ }
+ heartbeatInitialized = false;
+ }
+
+ /**
+ * 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.
+ * Pong responses are not tracked; once the heartbeat is enabled, Jetty's
idle timeout no
+ * longer determines connection liveness (see ZEPPELIN-6694).
+ */
+ void sendHeartbeat() {
+ 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 1805ce456f..57edf1d79b 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.Logger;
import org.slf4j.LoggerFactory;
import java.io.IOException;
+import java.nio.ByteBuffer;
import java.util.Map;
import jakarta.websocket.Session;
@@ -32,6 +33,11 @@ import jakarta.websocket.Session;
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<String, Object> headers;
private String user;
@@ -55,6 +61,21 @@ public class NotebookSocket {
});
}
+ /**
+ * 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 a5cb0037fd..f1e4d0d4da 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 @@ class ZeppelinConfigurationTest {
// 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 0000000000..6f53668261
--- /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 0000000000..4382e0dafe
--- /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);
+ }
+}
