Skip to content

Commit

Permalink
reuse buffer when reading messages from websocket.
Browse files Browse the repository at this point in the history
  • Loading branch information
fiatjaf committed Jul 30, 2023
1 parent cfcd195 commit 35faff8
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 12 deletions.
17 changes: 8 additions & 9 deletions connection.go
Original file line number Diff line number Diff line change
Expand Up @@ -121,47 +121,46 @@ func (c *Connection) WriteMessage(data []byte) error {
return nil
}

func (c *Connection) ReadMessage(ctx context.Context) ([]byte, error) {
func (c *Connection) ReadMessage(ctx context.Context, buf io.Writer) error {
for {
select {
case <-ctx.Done():
return nil, errors.New("context canceled")
return errors.New("context canceled")
default:
}

h, err := c.reader.NextFrame()
if err != nil {
c.conn.Close()
return nil, fmt.Errorf("failed to advance frame: %w", err)
return fmt.Errorf("failed to advance frame: %w", err)
}

if h.OpCode.IsControl() {
if err := c.controlHandler(h, c.reader); err != nil {
return nil, fmt.Errorf("failed to handle control frame: %w", err)
return fmt.Errorf("failed to handle control frame: %w", err)
}
} else if h.OpCode == ws.OpBinary ||
h.OpCode == ws.OpText {
break
}

if err := c.reader.Discard(); err != nil {
return nil, fmt.Errorf("failed to discard: %w", err)
return fmt.Errorf("failed to discard: %w", err)
}
}

buf := new(bytes.Buffer)
if c.msgState.IsCompressed() && c.enableCompression {
c.flateReader.Reset(c.reader)
if _, err := io.Copy(buf, c.flateReader); err != nil {
return nil, fmt.Errorf("failed to read message: %w", err)
return fmt.Errorf("failed to read message: %w", err)
}
} else {
if _, err := io.Copy(buf, c.reader); err != nil {
return nil, fmt.Errorf("failed to read message: %w", err)
return fmt.Errorf("failed to read message: %w", err)
}
}

return buf.Bytes(), nil
return nil
}

func (c *Connection) Close() error {
Expand Down
9 changes: 6 additions & 3 deletions relay.go
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
package nostr

import (
"bytes"
"context"
"fmt"
"log"
Expand Down Expand Up @@ -232,16 +233,18 @@ func (r *Relay) Connect(ctx context.Context) error {

// general message reader loop
go func() {
buf := new(bytes.Buffer)

for {
message, err := conn.ReadMessage(r.connectionContext)
if err != nil {
buf.Reset()
if err := conn.ReadMessage(r.connectionContext, buf); err != nil {
r.ConnectionError = err
r.Close()
break
}

message := buf.Bytes()
debugLogf("{%s} %v\n", r.URL, message)

envelope := ParseMessage(message)
if envelope == nil {
continue
Expand Down

0 comments on commit 35faff8

Please sign in to comment.