Skip to content
This repository has been archived by the owner on Jan 30, 2025. It is now read-only.

Fix recvloop so it doens't use as much memory #1482

Closed
wants to merge 2 commits into from
Closed
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 16 additions & 1 deletion common/connection.go
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
package common

import (
"bytes"
"context"
"crypto/tls"
"errors"
Expand Down Expand Up @@ -318,14 +319,28 @@ func (c *Connection) findTargetIDForLog(id target.SessionID) target.ID {
}

func (c *Connection) recvLoop() {
var buffer bytes.Buffer
Copy link
Collaborator Author

@ankur22 ankur22 Oct 14, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's nice that we find a solution for this tedious issue 🎉

I have suggestions.

This PR adds this:

// 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 {
_, 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)

I believe keeping this already tricky and central part of the code simpler will serve us better in the long run.

Also, correctly using io.Reader directly is risky and highly likely to create unforeseen issues. Since our goal is to reuse the same buffer instead of allocating a new one, we can use bytes.Buffer to let it handle the io.Reader. It reuses the same internal bytes buffer. Go has bytes.Buffer for these kind of situations:

var buffer bytes.Buffer

// loop:
if _, err := buffer.ReadFrom(reader); err != nil {
    // ...
}

// Clone so that we can share the underlying bytes
buf := bytes.Clone(buffer.Bytes())

// Reuse the underlying byte slice
buffer.Reset()

// Use buf here
...

I ran this suggested change through my test and it all checks off Perfectly 👍 a631b38

If your profiling shows that bufio.Reader is helpful, you can wrap reader with it: buffer.ReadFrom(bufio.NewReader(reader)).

This didn't make a difference, so I've left it out.

💡 It also makes sense to add a Benchmark to see how much memory usage we reduce and keep it at that level.

I've added a task in the original issue tracking all the memory related issues. I can't create a benchmark test yet as they aren't behaving as we would expect since the memory keeps going up due to the other memory leaks (investigating in #1484). Once memory leaking/usage has come under a bit more control then I think it would be best to complete the benchmarks.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I ran this suggested change through my test and it all checks off Perfectly 👍 a631b38

Nice 🎉

I've added a task in the original #1480 tracking all the memory related issues. I can't create a benchmark test yet as they aren't behaving as we would expect since the memory keeps going up due to the other memory leaks (investigating in #1484). Once memory leaking/usage has come under a bit more control then I think it would be best to complete the benchmarks.

This is fine, but benchmarks can also help with pprofing.


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
}

if _, err := buffer.ReadFrom(reader); err != nil {
c.handleIOError(err)
return
}

// Clone the bytes from the recvBuffer to buf before unmarshaling
// the message. Without this, the unmarshaling will fail as it
// reuses the buf as the underlying bytes buffer.
buf := bytes.Clone(buffer.Bytes())

buffer.Reset()

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

var msg cdproto.Message
Expand Down
Loading