From b0a2994b11b1835f74ccf0b64e5c8e59fa1f6f91 Mon Sep 17 00:00:00 2001
From: qydysky <qydysky@foxmail.com>
Date: Wed, 7 Aug 2024 01:14:37 +0800
Subject: [PATCH] 1

---
 io/io.go            | 19 +++++++++++++++++++
 io/io_test.go       | 33 +++++++++++++++++++++++++++++++++
 websocket/Client.go |  9 ++++++++-
 3 files changed, 60 insertions(+), 1 deletion(-)

diff --git a/io/io.go b/io/io.go
index b4fb4b1..fa7babe 100644
--- a/io/io.go
+++ b/io/io.go
@@ -414,3 +414,22 @@ func Copy(r io.Reader, w io.Writer, c CopyConfig) (e error) {
 		}
 	}
 }
+
+func ReadAll(r io.Reader, b []byte) ([]byte, error) {
+	b = b[0:cap(b)]
+	for {
+		n, err := r.Read(b[len(b):cap(b)])
+		b = b[:len(b)+n]
+		if err != nil {
+			if err == io.EOF {
+				err = nil
+			}
+			return b, err
+		}
+
+		if len(b) == cap(b) {
+			// Add more capacity (let append pick how much).
+			b = append(b, 0)[:len(b)]
+		}
+	}
+}
diff --git a/io/io_test.go b/io/io_test.go
index 9a35c66..864001d 100644
--- a/io/io_test.go
+++ b/io/io_test.go
@@ -63,3 +63,36 @@ func Test_RW2Chan(t *testing.T) {
 		}
 	}
 }
+
+func Test_readall(t *testing.T) {
+	var buf = []byte{}
+	result, e := ReadAll(bytes.NewReader([]byte{0x01, 0x02, 0x03}), buf)
+	if e != nil || !bytes.Equal(result, []byte{0x01, 0x02, 0x03}) {
+		t.Fatal()
+	}
+}
+
+// 4248350               281.0 ns/op            16 B/op          1 allocs/op
+func Benchmark_readall(b *testing.B) {
+	var buf = []byte{}
+	var data = []byte{0x01, 0x02, 0x03}
+	r := bytes.NewReader(data)
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		ReadAll(r, buf)
+		r.Reset(data)
+	}
+}
+
+// 2806576               424.2 ns/op           512 B/op          1 allocs/op
+func Benchmark_readall1(b *testing.B) {
+	var data = []byte{0x01, 0x02, 0x03}
+	r := bytes.NewReader(data)
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		io.ReadAll(r)
+		r.Reset(data)
+	}
+}
diff --git a/websocket/Client.go b/websocket/Client.go
index 360d858..eec5ece 100644
--- a/websocket/Client.go
+++ b/websocket/Client.go
@@ -8,8 +8,10 @@ import (
 	"sync"
 	"time"
 
+	"github.com/dustin/go-humanize"
 	"github.com/gorilla/websocket"
 
+	pio "github.com/qydysky/part/io"
 	msgq "github.com/qydysky/part/msgq"
 )
 
@@ -145,12 +147,17 @@ func (o *Client) Handle() (*msgq.MsgType[*WsMsg], error) {
 			o.l.Unlock()
 		}()
 
+		buf := make([]byte, humanize.KByte)
+		var message []byte
 		for {
 			if err := c.SetReadDeadline(time.Now().Add(time.Duration(o.RTOMs * int(time.Millisecond)))); err != nil {
 				o.error(err)
 				return
 			}
-			msg_type, message, err := c.ReadMessage()
+			msg_type, r, err := c.NextReader()
+			if err == nil {
+				message, err = pio.ReadAll(r, buf)
+			}
 			if err != nil {
 				if e, ok := err.(*websocket.CloseError); ok {
 					switch e.Code {
-- 
2.39.2