From 3b345c3f35b1ac7717fc9d628d646173888b20f0 Mon Sep 17 00:00:00 2001
From: Andreas Linde <mail@andreaslinde.de>
Date: Tue, 26 Nov 2024 10:52:52 +0100
Subject: [PATCH] Improve websocket write handling
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Increase the buffer size of the write channel to 1024. This way adding messages will never block if e.g. at the same time the connection is closed. If the buffer is full, return an error, as sending isn’t possible.

This fixes https://github.com/enbility/ship-go/issues/42
---
 ws/websocket.go      | 12 +++++++++---
 ws/websocket_test.go | 17 +++++++++++++++++
 2 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/ws/websocket.go b/ws/websocket.go
index 65f9a97..27b0160 100644
--- a/ws/websocket.go
+++ b/ws/websocket.go
@@ -80,8 +80,8 @@ func (w *WebsocketConnection) isConnClosed() bool {
 }
 
 func (w *WebsocketConnection) run() {
-	w.shipWriteChannel = make(chan []byte, 1) // Send outgoing ship messages
-	w.closeChannel = make(chan struct{}, 1)   // Listen to close events
+	w.shipWriteChannel = make(chan []byte, 1024) // Send outgoing ship messages
+	w.closeChannel = make(chan struct{}, 1)      // Listen to close events
 
 	go w.readShipPump()
 	go w.writeShipPump()
@@ -258,7 +258,13 @@ func (w *WebsocketConnection) WriteMessageToWebsocketConnection(message []byte)
 		return errors.New(connIsClosedError)
 	}
 
-	w.shipWriteChannel <- message
+	select {
+	case w.shipWriteChannel <- message:
+	default:
+		// too many messages are pending, this doesn't look good
+		return errors.New("could not send message, buffer is full")
+	}
+
 	return nil
 }
 
diff --git a/ws/websocket_test.go b/ws/websocket_test.go
index c352576..05a2431 100644
--- a/ws/websocket_test.go
+++ b/ws/websocket_test.go
@@ -128,6 +128,23 @@ func (s *WebsocketSuite) TestConnectionClose() {
 	assert.NotNil(s.T(), err)
 }
 
+func (s *WebsocketSuite) TestWriteBufferFull() {
+	amountNil := 0
+	amountNotNil := 0
+	for i := 0; i < 10000; i++ {
+		msg := []byte{1}
+		msg = append(msg, []byte("message")...)
+		err := s.sut.WriteMessageToWebsocketConnection(msg)
+		if err == nil {
+			amountNil++
+		} else {
+			amountNotNil++
+		}
+	}
+	assert.Greater(s.T(), amountNotNil, 0)
+	assert.Greater(s.T(), amountNil, 0)
+}
+
 func (s *WebsocketSuite) TestPingPeriod() {
 	isClosed, err := s.sut.IsDataConnectionClosed()
 	assert.Equal(s.T(), false, isClosed)