Ensure message order with a buffer

nhooyr · nhooyr · commit ae56de60e4fc · 2019-09-23T16:53:37.000-05:00
diff --git a/websocket_js.go b/websocket_js.go
@@ -31,7 +31,9 @@ type Conn struct {
 	releaseOnClose   func()
 	releaseOnMessage func()
 
-	readch chan wsjs.MessageEvent
+	readSignal chan struct{}
+	readBufMu  sync.Mutex
+	readBuf    []wsjs.MessageEvent
 }
 
 func (c *Conn) close(err error) {
@@ -45,7 +47,7 @@ func (c *Conn) close(err error) {
 
 func (c *Conn) init() {
 	c.closed = make(chan struct{})
-	c.readch = make(chan wsjs.MessageEvent, 1)
+	c.readSignal = make(chan struct{}, 1)
 	c.msgReadLimit = 32768
 
 	c.releaseOnClose = c.ws.OnClose(func(e wsjs.CloseEvent) {
@@ -61,7 +63,16 @@ func (c *Conn) init() {
 	})
 
 	c.releaseOnMessage = c.ws.OnMessage(func(e wsjs.MessageEvent) {
-		c.readch <- e
+		c.readBufMu.Lock()
+		defer c.readBufMu.Unlock()
+
+		c.readBuf = append(c.readBuf, e)
+
+		// Lets the read goroutine know there is definitely something in readBuf.
+		select {
+		case c.readSignal <- struct{}{}:
+		default:
+		}
 	})
 
 	runtime.SetFinalizer(c, func(c *Conn) {
@@ -89,16 +100,29 @@ func (c *Conn) Read(ctx context.Context) (MessageType, []byte, error) {
 }
 
 func (c *Conn) read(ctx context.Context) (MessageType, []byte, error) {
-	var me wsjs.MessageEvent
 	select {
 	case <-ctx.Done():
 		c.Close(StatusPolicyViolation, "read timed out")
 		return 0, nil, ctx.Err()
-	case me = <-c.readch:
+	case <-c.readSignal:
 	case <-c.closed:
 		return 0, nil, c.closeErr
 	}
 
+	c.readBufMu.Lock()
+	defer c.readBufMu.Unlock()
+
+	me := c.readBuf[0]
+	c.readBuf = c.readBuf[1:]
+
+	if len(c.readBuf) > 0 {
+		// Next time we read, we'll grab the message.
+		select {
+		case c.readSignal <- struct{}{}:
+		default:
+		}
+	}
+
 	switch p := me.Data.(type) {
 	case string:
 		return MessageText, []byte(p), nil
