Websocket peer outbound buffer fixes. Expose outbound buffered amount.

10 files changed, 56 insertions, 4 deletions

diff --git a/modules/websocket/doc_classes/WebSocketPeer.xml b/modules/websocket/doc_classes/WebSocketPeer.xml
index 5125956416..9f2db37dd7 100644
--- a/modules/websocket/doc_classes/WebSocketPeer.xml
+++ b/modules/websocket/doc_classes/WebSocketPeer.xml
@@ -39,6 +39,12 @@
 				[b]Note:[/b] Not available in the HTML5 export.
 			</description>
 		</method>
+		<method name="get_current_outbound_buffered_amount" qualifiers="const">
+			<return type="int" />
+			<description>
+				Returns the current amount of data in the outbound websocket buffer. [b]Note:[/b] HTML5 exports use WebSocket.bufferedAmount, while other platforms use an internal buffer.
+			</description>
+		</method>
 		<method name="get_write_mode" qualifiers="const">
 			<return type="int" enum="WebSocketPeer.WriteMode">
 			</return>
diff --git a/modules/websocket/emws_client.cpp b/modules/websocket/emws_client.cpp
index d3d0066c12..5cd94e978f 100644
--- a/modules/websocket/emws_client.cpp
+++ b/modules/websocket/emws_client.cpp
@@ -95,7 +95,7 @@ Error EMWSClient::connect_to_host(String p_host, String p_path, uint16_t p_port,
 		return FAILED;
 	}
 
-	static_cast<Ref<EMWSPeer>>(_peer)->set_sock(_js_id, _in_buf_size, _in_pkt_size);
+	static_cast<Ref<EMWSPeer>>(_peer)->set_sock(_js_id, _in_buf_size, _in_pkt_size, _out_buf_size);
 
 	return OK;
 }
@@ -136,6 +136,7 @@ int EMWSClient::get_max_packet_size() const {
 Error EMWSClient::set_buffers(int p_in_buffer, int p_in_packets, int p_out_buffer, int p_out_packets) {
 	_in_buf_size = nearest_shift(p_in_buffer - 1) + 10;
 	_in_pkt_size = nearest_shift(p_in_packets - 1);
+	_out_buf_size = nearest_shift(p_out_buffer - 1) + 10;
 	return OK;
 }
 
diff --git a/modules/websocket/emws_client.h b/modules/websocket/emws_client.h
index ca2d7ed986..3b0b8395b9 100644
--- a/modules/websocket/emws_client.h
+++ b/modules/websocket/emws_client.h
@@ -45,6 +45,7 @@ private:
 	bool _is_connecting = false;
 	int _in_buf_size = DEF_BUF_SHIFT;
 	int _in_pkt_size = DEF_PKT_SHIFT;
+	int _out_buf_size = DEF_BUF_SHIFT;
 
 	static void _esws_on_connect(void *obj, char *proto);
 	static void _esws_on_message(void *obj, const uint8_t *p_data, int p_data_size, int p_is_string);
diff --git a/modules/websocket/emws_peer.cpp b/modules/websocket/emws_peer.cpp
index 05f9e12ae1..69822f6ff3 100644
--- a/modules/websocket/emws_peer.cpp
+++ b/modules/websocket/emws_peer.cpp
@@ -33,10 +33,11 @@
 #include "emws_peer.h"
 #include "core/io/ip.h"
 
-void EMWSPeer::set_sock(int p_sock, unsigned int p_in_buf_size, unsigned int p_in_pkt_size) {
+void EMWSPeer::set_sock(int p_sock, unsigned int p_in_buf_size, unsigned int p_in_pkt_size, unsigned int p_out_buf_size) {
 	peer_sock = p_sock;
 	_in_buffer.resize(p_in_pkt_size, p_in_buf_size);
 	_packet_buffer.resize((1 << p_in_buf_size));
+	_out_buf_size = p_out_buf_size;
 }
 
 void EMWSPeer::set_write_mode(WriteMode p_mode) {
@@ -53,7 +54,10 @@ Error EMWSPeer::read_msg(const uint8_t *p_data, uint32_t p_size, bool p_is_strin
 }
 
 Error EMWSPeer::put_packet(const uint8_t *p_buffer, int p_buffer_size) {
+	ERR_FAIL_COND_V(_out_buf_size && (godot_js_websocket_buffered_amount(peer_sock) >= (1ULL << _out_buf_size)), ERR_OUT_OF_MEMORY);
+
 	int is_bin = write_mode == WebSocketPeer::WRITE_MODE_BINARY ? 1 : 0;
+
 	godot_js_websocket_send(peer_sock, p_buffer, p_buffer_size, is_bin);
 	return OK;
 }
@@ -76,6 +80,13 @@ int EMWSPeer::get_available_packet_count() const {
 	return _in_buffer.packets_left();
 }
 
+int EMWSPeer::get_current_outbound_buffered_amount() const {
+	if (peer_sock != -1) {
+		return godot_js_websocket_buffered_amount(peer_sock);
+	}
+	return 0;
+}
+
 bool EMWSPeer::was_string_packet() const {
 	return _is_string;
 }
diff --git a/modules/websocket/emws_peer.h b/modules/websocket/emws_peer.h
index 73e701720b..6e93ea31a2 100644
--- a/modules/websocket/emws_peer.h
+++ b/modules/websocket/emws_peer.h
@@ -48,6 +48,7 @@ typedef void (*WSOnError)(void *p_ref);
 
 extern int godot_js_websocket_create(void *p_ref, const char *p_url, const char *p_proto, WSOnOpen p_on_open, WSOnMessage p_on_message, WSOnError p_on_error, WSOnClose p_on_close);
 extern int godot_js_websocket_send(int p_id, const uint8_t *p_buf, int p_buf_len, int p_raw);
+extern int godot_js_websocket_buffered_amount(int p_id);
 extern void godot_js_websocket_close(int p_id, int p_code, const char *p_reason);
 extern void godot_js_websocket_destroy(int p_id);
 }
@@ -62,14 +63,16 @@ private:
 	Vector<uint8_t> _packet_buffer;
 	PacketBuffer<uint8_t> _in_buffer;
 	uint8_t _is_string = 0;
+	int _out_buf_size = 0;
 
 public:
 	Error read_msg(const uint8_t *p_data, uint32_t p_size, bool p_is_string);
-	void set_sock(int p_sock, unsigned int p_in_buf_size, unsigned int p_in_pkt_size);
+	void set_sock(int p_sock, unsigned int p_in_buf_size, unsigned int p_in_pkt_size, unsigned int p_out_buf_size);
 	virtual int get_available_packet_count() const;
 	virtual Error get_packet(const uint8_t **r_buffer, int &r_buffer_size);
 	virtual Error put_packet(const uint8_t *p_buffer, int p_buffer_size);
 	virtual int get_max_packet_size() const { return _packet_buffer.size(); };
+	virtual int get_current_outbound_buffered_amount() const;
 
 	virtual void close(int p_code = 1000, String p_reason = "");
 	virtual bool is_connected_to_host() const;
diff --git a/modules/websocket/library_godot_websocket.js b/modules/websocket/library_godot_websocket.js
index b182d1ecde..dd2fd1e94f 100644
--- a/modules/websocket/library_godot_websocket.js
+++ b/modules/websocket/library_godot_websocket.js
@@ -101,6 +101,15 @@ const GodotWebSocket = {
 			return 0;
 		},
 
+		// Get current bufferedAmount
+		bufferedAmount: function (p_id) {
+			const ref = IDHandler.get(p_id);
+			if (!ref) {
+				return 0; // Godot object is gone.
+			}
+			return ref.bufferedAmount;
+		},
+
 		create: function (socket, p_on_open, p_on_message, p_on_error, p_on_close) {
 			const id = IDHandler.add(socket);
 			socket.onopen = GodotWebSocket._onopen.bind(null, id, p_on_open);
@@ -171,6 +180,11 @@ const GodotWebSocket = {
 		return GodotWebSocket.send(p_id, out);
 	},
 
+	godot_js_websocket_buffered_amount__sig: 'ii',
+	godot_js_websocket_buffered_amount: function (p_id) {
+		return GodotWebSocket.bufferedAmount(p_id);
+	},
+
 	godot_js_websocket_close__sig: 'viii',
 	godot_js_websocket_close: function (p_id, p_code, p_reason) {
 		const code = p_code;
diff --git a/modules/websocket/websocket_peer.cpp b/modules/websocket/websocket_peer.cpp
index e77fdcfed2..ee13040821 100644
--- a/modules/websocket/websocket_peer.cpp
+++ b/modules/websocket/websocket_peer.cpp
@@ -47,6 +47,7 @@ void WebSocketPeer::_bind_methods() {
 	ClassDB::bind_method(D_METHOD("get_connected_host"), &WebSocketPeer::get_connected_host);
 	ClassDB::bind_method(D_METHOD("get_connected_port"), &WebSocketPeer::get_connected_port);
 	ClassDB::bind_method(D_METHOD("set_no_delay", "enabled"), &WebSocketPeer::set_no_delay);
+	ClassDB::bind_method(D_METHOD("get_current_outbound_buffered_amount"), &WebSocketPeer::get_current_outbound_buffered_amount);
 
 	BIND_ENUM_CONSTANT(WRITE_MODE_TEXT);
 	BIND_ENUM_CONSTANT(WRITE_MODE_BINARY);
diff --git a/modules/websocket/websocket_peer.h b/modules/websocket/websocket_peer.h
index e9bb20f21f..517b8600d6 100644
--- a/modules/websocket/websocket_peer.h
+++ b/modules/websocket/websocket_peer.h
@@ -59,6 +59,7 @@ public:
 	virtual uint16_t get_connected_port() const = 0;
 	virtual bool was_string_packet() const = 0;
 	virtual void set_no_delay(bool p_enabled) = 0;
+	virtual int get_current_outbound_buffered_amount() const = 0;
 
 	WebSocketPeer();
 	~WebSocketPeer();
diff --git a/modules/websocket/wsl_peer.cpp b/modules/websocket/wsl_peer.cpp
index 1dbadfed74..7f027e1c0d 100644
--- a/modules/websocket/wsl_peer.cpp
+++ b/modules/websocket/wsl_peer.cpp
@@ -205,7 +205,9 @@ void WSLPeer::make_context(PeerData *p_data, unsigned int p_in_buf_size, unsigne
 	ERR_FAIL_COND(p_data == nullptr);
 
 	_in_buffer.resize(p_in_pkt_size, p_in_buf_size);
-	_packet_buffer.resize((1 << MAX(p_in_buf_size, p_out_buf_size)));
+	_packet_buffer.resize(1 << p_in_buf_size);
+	_out_buf_size = p_out_buf_size;
+	_out_pkt_size = p_out_pkt_size;
 
 	_data = p_data;
 	_data->peer = this;
@@ -239,6 +241,8 @@ void WSLPeer::poll() {
 
 Error WSLPeer::put_packet(const uint8_t *p_buffer, int p_buffer_size) {
 	ERR_FAIL_COND_V(!is_connected_to_host(), FAILED);
+	ERR_FAIL_COND_V(_out_pkt_size && (wslay_event_get_queued_msg_count(_data->ctx) >= (1ULL << _out_pkt_size)), ERR_OUT_OF_MEMORY);
+	ERR_FAIL_COND_V(_out_buf_size && (wslay_event_get_queued_msg_length(_data->ctx) >= (1ULL << _out_buf_size)), ERR_OUT_OF_MEMORY);
 
 	struct wslay_event_msg msg; // Should I use fragmented?
 	msg.opcode = write_mode == WRITE_MODE_TEXT ? WSLAY_TEXT_FRAME : WSLAY_BINARY_FRAME;
@@ -280,6 +284,12 @@ int WSLPeer::get_available_packet_count() const {
 	return _in_buffer.packets_left();
 }
 
+int WSLPeer::get_current_outbound_buffered_amount() const {
+	ERR_FAIL_COND_V(!_data, 0);
+
+	return wslay_event_get_queued_msg_length(_data->ctx);
+}
+
 bool WSLPeer::was_string_packet() const {
 	return _is_string;
 }
diff --git a/modules/websocket/wsl_peer.h b/modules/websocket/wsl_peer.h
index f1ea98d384..260d4b183d 100644
--- a/modules/websocket/wsl_peer.h
+++ b/modules/websocket/wsl_peer.h
@@ -77,6 +77,9 @@ private:
 
 	WriteMode write_mode = WRITE_MODE_BINARY;
 
+	int _out_buf_size = 0;
+	int _out_pkt_size = 0;
+
 public:
 	int close_code = -1;
 	String close_reason;
@@ -86,6 +89,7 @@ public:
 	virtual Error get_packet(const uint8_t **r_buffer, int &r_buffer_size);
 	virtual Error put_packet(const uint8_t *p_buffer, int p_buffer_size);
 	virtual int get_max_packet_size() const { return _packet_buffer.size(); };
+	virtual int get_current_outbound_buffered_amount() const;
 
 	virtual void close_now();
 	virtual void close(int p_code = 1000, String p_reason = "");