Merge pull request #91707 from Faless/ws/heartbeat

[WebSocket] Add optional heartbeat via "ping" control frames.

modules/websocket/doc_classes/WebSocketPeer.xml

@@ -155,6 +155,10 @@
 			The extra HTTP headers to be sent during the WebSocket handshake.
 			[b]Note:[/b] Not supported in Web exports due to browsers' restrictions.
 		</member>
+		<member name="heartbeat_interval" type="float" setter="set_heartbeat_interval" getter="get_heartbeat_interval" default="0.0">
+			The interval (in seconds) at which the peer will automatically send WebSocket "ping" control frames. When set to [code]0[/code], no "ping" control frames will be sent.
+			[b]Note:[/b] Has no effect in Web exports due to browser restrictions.
+		</member>
 		<member name="inbound_buffer_size" type="int" setter="set_inbound_buffer_size" getter="get_inbound_buffer_size" default="65535">
 			The size of the input buffer in bytes (roughly the maximum amount of memory that will be allocated for the inbound packets).
 		</member>

modules/websocket/websocket_peer.cpp

@@ -70,6 +70,9 @@ void WebSocketPeer::_bind_methods() {
 	ClassDB::bind_method(D_METHOD("set_max_queued_packets", "buffer_size"), &WebSocketPeer::set_max_queued_packets);
 	ClassDB::bind_method(D_METHOD("get_max_queued_packets"), &WebSocketPeer::get_max_queued_packets);
 
+	ClassDB::bind_method(D_METHOD("set_heartbeat_interval", "interval"), &WebSocketPeer::set_heartbeat_interval);
+	ClassDB::bind_method(D_METHOD("get_heartbeat_interval"), &WebSocketPeer::get_heartbeat_interval);
+
 	ADD_PROPERTY(PropertyInfo(Variant::PACKED_STRING_ARRAY, "supported_protocols"), "set_supported_protocols", "get_supported_protocols");
 	ADD_PROPERTY(PropertyInfo(Variant::PACKED_STRING_ARRAY, "handshake_headers"), "set_handshake_headers", "get_handshake_headers");
 
@@ -78,6 +81,8 @@ void WebSocketPeer::_bind_methods() {
 
 	ADD_PROPERTY(PropertyInfo(Variant::INT, "max_queued_packets"), "set_max_queued_packets", "get_max_queued_packets");
 
+	ADD_PROPERTY(PropertyInfo(Variant::INT, "heartbeat_interval"), "set_heartbeat_interval", "get_heartbeat_interval");
+
 	BIND_ENUM_CONSTANT(WRITE_MODE_TEXT);
 	BIND_ENUM_CONSTANT(WRITE_MODE_BINARY);
 
@@ -151,3 +156,12 @@ void WebSocketPeer::set_max_queued_packets(int p_max_queued_packets) {
 int WebSocketPeer::get_max_queued_packets() const {
 	return max_queued_packets;
 }
+
+double WebSocketPeer::get_heartbeat_interval() const {
+	return heartbeat_interval_msec / 1000.0;
+}
+
+void WebSocketPeer::set_heartbeat_interval(double p_interval) {
+	ERR_FAIL_COND(p_interval < 0);
+	heartbeat_interval_msec = p_interval * 1000.0;
+}

modules/websocket/websocket_peer.h

@@ -72,6 +72,7 @@ protected:
 	int outbound_buffer_size = DEFAULT_BUFFER_SIZE;
 	int inbound_buffer_size = DEFAULT_BUFFER_SIZE;
 	int max_queued_packets = 2048;
+	uint64_t heartbeat_interval_msec = 0;
 
 public:
 	static WebSocketPeer *create(bool p_notify_postinitialize = true) {
@@ -117,6 +118,9 @@ public:
 	void set_max_queued_packets(int p_max_queued_packets);
 	int get_max_queued_packets() const;
 
+	double get_heartbeat_interval() const;
+	void set_heartbeat_interval(double p_interval);
+
 	WebSocketPeer();
 	~WebSocketPeer();
 };

modules/websocket/wsl_peer.cpp

@@ -636,7 +636,10 @@ void WSLPeer::_wsl_msg_recv_callback(wslay_event_context_ptr ctx, const struct w
 		uint8_t is_string = arg->opcode == WSLAY_TEXT_FRAME ? 1 : 0;
 		peer->in_buffer.write_packet(arg->msg, arg->msg_length, &is_string);
 	}
-	// Ping or pong.
+	if (op == WSLAY_PONG) {
+		peer->heartbeat_waiting = false;
+	}
+	// Pong.
 }
 
 wslay_event_callbacks WSLPeer::_wsl_callbacks = {
@@ -680,7 +683,31 @@ void WSLPeer::poll() {
 
 	if (ready_state == STATE_OPEN || ready_state == STATE_CLOSING) {
 		ERR_FAIL_NULL(wsl_ctx);
+		uint64_t ticks = OS::get_singleton()->get_ticks_msec();
 		int err = 0;
+		if (heartbeat_interval_msec != 0 && ticks - last_heartbeat > heartbeat_interval_msec && ready_state == STATE_OPEN) {
+			if (heartbeat_waiting) {
+				wslay_event_context_free(wsl_ctx);
+				wsl_ctx = nullptr;
+				close(-1);
+				return;
+			}
+			heartbeat_waiting = true;
+			struct wslay_event_msg msg;
+			msg.opcode = WSLAY_PING;
+			msg.msg = nullptr;
+			msg.msg_length = 0;
+			err = wslay_event_queue_msg(wsl_ctx, &msg);
+			if (err == 0) {
+				last_heartbeat = ticks;
+			} else {
+				print_verbose("Websocket (wslay) failed to send ping: " + itos(err));
+				wslay_event_context_free(wsl_ctx);
+				wsl_ctx = nullptr;
+				close(-1);
+				return;
+			}
+		}
 		if ((err = wslay_event_recv(wsl_ctx)) != 0 || (err = wslay_event_send(wsl_ctx)) != 0) {
 			// Error close.
 			print_verbose("Websocket (wslay) poll error: " + itos(err));
@@ -781,6 +808,7 @@ void WSLPeer::close(int p_code, String p_reason) {
 		}
 	}
 
+	heartbeat_waiting = false;
 	in_buffer.clear();
 	packet_buffer.resize(0);
 }

modules/websocket/wsl_peer.h

@@ -99,6 +99,8 @@ private:
 	int close_code = -1;
 	String close_reason;
 	uint8_t was_string = 0;
+	uint64_t last_heartbeat = 0;
+	bool heartbeat_waiting = false;
 
 	// WebSocket configuration.
 	bool use_tls = true;