1

diff --git a/io/io.go b/io/io.go
index b4fb4b1a97562b7838b749bedd758a7936392e3e..fa7babedcf519fae64ebddc701a322e8b150160c 100644 (file)
--- 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 9a35c668184fbc6b4f460af5ecacb1c78d530278..864001d6eb7a7ec274bbcd22891c88c035e8e870 100644 (file)
--- 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 360d858007201964066f02003821adf6af0fd638..eec5ece5cf7409bad6010b0a59456cc31956d583 100644 (file)
--- 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 {