avoid per-frame cleanup closures on read to remove 2 allocs per frame read (#565)

read: avoid per-frame cleanup closures

Every frame header and payload read called prepareRead, which returned a
cleanup function to clear the timeout and translate close or cancellation
errors. That function captured the context, connection, and the address of
the caller's named error result. Because prepareRead returned the function,
the closure outlived its stack frame and Go allocated its captured state on
the heap for every frame read.

Move the cleanup logic to a normal finishRead method and defer a direct
method call instead. This preserves timeout cleanup and error translation
without returning a closure. Compiler escape analysis with -gcflags=-m=2
confirms that the old function literal escaped and forced the named error
result in readFrameHeader and readFramePayload onto the heap; neither escape
remains after this change.

The results below compare parent d099e16 with this commit on an Apple
M4 Max using GOMAXPROCS=1 and benchstat over 10 samples. A temporary
in-package harness, not included in this commit, repeatedly called one
internal frame read. Header reads parse a minimal frame header; payload reads
copy 512 bytes from a buffered repeating reader. The background cases use
context.Background, while the cancelable cases use an uncanceled
context.WithCancel.

Removing the escaping closure eliminates 2 allocations and 64 bytes from
every frame read: one allocation for the closure environment and one for the
named error result retained by that closure. Header reads with a background
context improve from 170.5 to 118.5 ns/op, 216 to 152 B/op, and 6 to
4 allocs/op. Header reads with a cancelable context improve from 211.1 to
158.7 ns/op with the same allocation reduction. Payload reads remove the
same fixed overhead; they remain slower because the benchmark also copies
512 bytes.

Both context types benefit because this commit does not change timeout
registration. It only removes cleanup allocations made after every
prepareRead call. An interleaved 12-sample
BenchmarkConn/disabledCompress run, which exercises complete message reads,
improves from 3.875 to 3.657 us/op (-5.63%), 42 to 32 allocs/op (-23.81%),
and 1536 to 1216 B/op (-20.83%).

Author: mitchellh

read.go

@@ -217,51 +217,48 @@ func (c *Conn) readLoop(ctx context.Context) (header, error) {
 	}
 }
 
-// prepareRead sets the readTimeout context and returns a done function
-// to be called after the read is done. It also returns an error if the
-// connection is closed. The reference to the error is used to assign
-// an error depending on if the connection closed or the context timed
-// out during use. Typically, the referenced error is a named return
-// variable of the function calling this method.
-func (c *Conn) prepareRead(ctx context.Context, err *error) (func(), error) {
+// prepareRead sets the read timeout and checks whether the connection is closed.
+func (c *Conn) prepareRead(ctx context.Context) error {
 	select {
 	case <-c.closed:
-		return nil, net.ErrClosed
+		return net.ErrClosed
 	default:
 	}
 	c.setupReadTimeout(ctx)
 
-	done := func() {
-		c.clearReadTimeout()
-		select {
-		case <-c.closed:
-			if *err != nil {
-				*err = net.ErrClosed
-			}
-		default:
-		}
-		if *err != nil && ctx.Err() != nil {
-			*err = ctx.Err()
-		}
-	}
-
 	c.closeStateMu.Lock()
 	closeReceivedErr := c.closeReceivedErr
 	c.closeStateMu.Unlock()
 	if closeReceivedErr != nil {
-		defer done()
-		return nil, closeReceivedErr
+		c.clearReadTimeout()
+		return closeReceivedErr
 	}
 
-	return done, nil
+	return nil
+}
+
+// finishRead clears the read timeout and reports whether the connection or
+// operation context ended while the read was in progress.
+func (c *Conn) finishRead(ctx context.Context, err *error) {
+	c.clearReadTimeout()
+	select {
+	case <-c.closed:
+		if *err != nil {
+			*err = net.ErrClosed
+		}
+	default:
+	}
+	if *err != nil && ctx.Err() != nil {
+		*err = ctx.Err()
+	}
 }
 
 func (c *Conn) readFrameHeader(ctx context.Context) (_ header, err error) {
-	readDone, err := c.prepareRead(ctx, &err)
+	err = c.prepareRead(ctx)
 	if err != nil {
 		return header{}, err
 	}
-	defer readDone()
+	defer c.finishRead(ctx, &err)
 
 	h, err := readFrameHeader(c.br, c.readHeaderBuf[:])
 	if err != nil {
@@ -272,11 +269,11 @@ func (c *Conn) readFrameHeader(ctx context.Context) (_ header, err error) {
 }
 
 func (c *Conn) readFramePayload(ctx context.Context, p []byte) (_ int, err error) {
-	readDone, err := c.prepareRead(ctx, &err)
+	err = c.prepareRead(ctx)
 	if err != nil {
 		return 0, err
 	}
-	defer readDone()
+	defer c.finishRead(ctx, &err)
 
 	n, err := io.ReadFull(c.br, p)
 	if err != nil {