conn: avoid frame timeout and cleanup overhead

#565

Every frame read registered a context.AfterFunc callback, even when the
context could not be canceled, and returned a cleanup closure that
forced the caller's error and captured state onto the heap. Skip
timeout setup for contexts with a nil Done channel and move read
cleanup into a direct method call, while tracking whether a callback
was installed so cancelable operations retain the same
close-on-cancellation behavior. Writes use the same background-context
fast path, and BenchmarkConn now joins its writer goroutine so repeated
runs finish cleanly.

On an Apple M4 Max with GOMAXPROCS=1, a 10-sample in-package
benchmark against d099e16 reduced background header reads from
168.36 to 22.85 ns/op and 512-byte payload reads from 175.04 to
26.65 ns/op. Both fell from 216 B/op and 6 allocs/op to zero.
Cancelable header reads improved from 203.40 to 160.50 ns/op and
payload reads from 215.18 to 162.52 ns/op, with both falling to
152 B/op and 4 allocs/op. Comparing against the closure-only change
confirms that skipping the dead timeout registration accounts for
exactly 152 bytes and 4 allocations on background reads without
changing cancelable allocations.

BenchmarkConn/disabledCompress averaged 3.968 us/op, 1539 B/op, and 42
allocs/op before the change and 3.817 us/op, 1219 B/op, and 32 allocs/op
after it. This is a 3.8 percent time reduction, 20.8 percent fewer
bytes, and 23.8 percent fewer allocations per operation.

Author: mitchellh

conn.go

@@ -168,24 +168,34 @@ func (c *Conn) close() error {
 	return err
 }
 
-func (c *Conn) setupWriteTimeout(ctx context.Context) {
+func (c *Conn) setupWriteTimeout(ctx context.Context) bool {
+	if ctx.Done() == nil {
+		return false
+	}
+
 	stop := context.AfterFunc(ctx, func() {
 		c.clearWriteTimeout()
 		c.close()
 	})
 	swapTimeoutStop(&c.writeTimeoutStop, &stop)
+	return true
 }
 
 func (c *Conn) clearWriteTimeout() {
 	swapTimeoutStop(&c.writeTimeoutStop, nil)
 }
 
-func (c *Conn) setupReadTimeout(ctx context.Context) {
+func (c *Conn) setupReadTimeout(ctx context.Context) bool {
+	if ctx.Done() == nil {
+		return false
+	}
+
 	stop := context.AfterFunc(ctx, func() {
 		c.clearReadTimeout()
 		c.close()
 	})
 	swapTimeoutStop(&c.readTimeoutStop, &stop)
+	return true
 }
 
 func (c *Conn) clearReadTimeout() {

conn_test.go

@@ -592,10 +592,11 @@ func BenchmarkConn(b *testing.B) {
 			msg := []byte(strings.Repeat("1234", 128))
 			readBuf := make([]byte, len(msg))
 			writes := make(chan struct{})
-			defer close(writes)
 			werrs := make(chan error)
+			writerDone := make(chan struct{})
 
 			go func() {
+				defer close(writerDone)
 				for range writes {
 					select {
 					case werrs <- c1.Write(bb.ctx, websocket.MessageText, msg):
@@ -650,6 +651,13 @@ func BenchmarkConn(b *testing.B) {
 			}
 			b.StopTimer()
 
+			close(writes)
+			select {
+			case <-writerDone:
+			case <-bb.ctx.Done():
+				b.Fatal(bb.ctx.Err())
+			}
+
 			b.ReportMetric(float64(*bytesWritten/b.N), "written/op")
 			b.ReportMetric(float64(*bytesRead/b.N), "read/op")
 

read.go

@@ -217,51 +217,52 @@ 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) (bool, error) {
 	select {
 	case <-c.closed:
-		return nil, net.ErrClosed
+		return false, 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()
-		}
-	}
+	timeoutSet := c.setupReadTimeout(ctx)
 
 	c.closeStateMu.Lock()
 	closeReceivedErr := c.closeReceivedErr
 	c.closeStateMu.Unlock()
 	if closeReceivedErr != nil {
-		defer done()
-		return nil, closeReceivedErr
+		if timeoutSet {
+			c.clearReadTimeout()
+		}
+		return false, closeReceivedErr
 	}
 
-	return done, nil
+	return timeoutSet, 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, timeoutSet bool) {
+	if timeoutSet {
+		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)
+	timeoutSet, err := c.prepareRead(ctx)
 	if err != nil {
 		return header{}, err
 	}
-	defer readDone()
+	defer c.finishRead(ctx, &err, timeoutSet)
 
 	h, err := readFrameHeader(c.br, c.readHeaderBuf[:])
 	if err != nil {
@@ -272,11 +273,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)
+	timeoutSet, err := c.prepareRead(ctx)
 	if err != nil {
 		return 0, err
 	}
-	defer readDone()
+	defer c.finishRead(ctx, &err, timeoutSet)
 
 	n, err := io.ReadFull(c.br, p)
 	if err != nil {

write.go

@@ -318,8 +318,9 @@ func (c *Conn) writeFrame(ctx context.Context, fin bool, flate bool, opcode opco
 		return 0, net.ErrClosed
 	default:
 	}
-	c.setupWriteTimeout(ctx)
-	defer c.clearWriteTimeout()
+	if c.setupWriteTimeout(ctx) {
+		defer c.clearWriteTimeout()
+	}
 
 	c.writeHeader.fin = fin
 	c.writeHeader.opcode = opcode