Skip to content

Commit

Permalink
Fix recvloop so it doens't use as much memory
Browse files Browse the repository at this point in the history
recvloop was using ws.ReadMessage, which was in turn working with
io.ReadAll. It turns out for the current use, this isn't a good way to
read the message of the web socket connection. Instead we should buffer
the read, store it in a buffer that can increase with the test, which
is then cloned before being unmarshaled. We have more control over the
memory by working in this way.
  • Loading branch information
ankur22 committed Oct 11, 2024
1 parent 75647e9 commit c8ee6f4
Showing 1 changed file with 34 additions and 1 deletion.
35 changes: 34 additions & 1 deletion common/connection.go
Original file line number Diff line number Diff line change
@@ -1,10 +1,13 @@
package common

import (
"bufio"
"bytes"
"context"
"crypto/tls"
"errors"
"fmt"
"io"
"net/http"
"sync"
"sync/atomic"
Expand Down Expand Up @@ -318,14 +321,44 @@ func (c *Connection) findTargetIDForLog(id target.SessionID) target.ID {
}

func (c *Connection) recvLoop() {
// Setting the wsBuffer as the same size as bufio.NewReader. We should look at
// optimizing it depending on the average website, or allow it to be
// configurable, but it depends on how much this affects the user.
wsBuffer := make([]byte, 4096)
// This buffer grows as the test progresses.
var recvBuffer []byte

c.logger.Debugf("Connection:recvLoop", "wsURL:%q", c.wsURL)
for {
_, buf, err := c.conn.ReadMessage()
_, reader, err := c.conn.NextReader()
if err != nil {
c.handleIOError(err)
return
}

// Buffered reads from the websocket connection.
bufReader := bufio.NewReader(reader)
for {
n, err := bufReader.Read(wsBuffer)
if err == io.EOF {
break
}
if err != nil {
c.handleIOError(err)
return
}
recvBuffer = append(recvBuffer, bytes.Clone(wsBuffer[:n])...)
}

// Clone the bytes from the recvBuffer to buf before unmarshaling
// the message. Without this test the unmarshaling will fail as
// it reuses the buf as the underlying bytes buffer.
buf := make([]byte, len(recvBuffer))
copy(buf, recvBuffer)

// Reset recvBuffer so it can be reused.
recvBuffer = recvBuffer[:0]

c.logger.Tracef("cdp:recv", "<- %s", buf)

var msg cdproto.Message
Expand Down

0 comments on commit c8ee6f4

Please sign in to comment.