[ZEPPELIN-6092] Send server-initiated websocket ping frames to keep connections alive

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 ####
 

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>

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>

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"),

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

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<String, NotebookSocket> 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

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<String, Object> 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;
   }

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

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