conn: skip timeout callbacks for background contexts (#566)

Every frame read and write registered a context.AfterFunc callback so a
context cancellation could close the connection. context.Background
cannot be canceled because its Done channel is nil, but the code still
constructed, stored, stopped, and cleared a callback registration for
each operation. Skip that work when ctx.Done() is nil, and only clear a
timeout when one was actually installed. BenchmarkConn also joins its
writer goroutine after timing so repeated benchmark runs finish cleanly.

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

Before this change, setupReadTimeout called context.AfterFunc even for
context.Background. Avoiding that dead registration reduces background
header reads from 121.95 to 29.94 ns/op and payload reads from 124.17 to
29.15 ns/op. Both fall from 152 B/op and 4 allocs/op to zero.

Cancelable contexts still need an AfterFunc registration, so they
remain at 152 B/op and 4 allocs/op. BenchmarkConn/disabledCompress uses
a cancelable context as well, so it remains at 1219 B/op and 32
allocs/op.

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

@@ -218,29 +218,33 @@ func (c *Conn) readLoop(ctx context.Context) (header, error) {
 }
 
 // prepareRead sets the read timeout and checks whether the connection is closed.
-func (c *Conn) prepareRead(ctx context.Context) error {
+func (c *Conn) prepareRead(ctx context.Context) (bool, error) {
 	select {
 	case <-c.closed:
-		return net.ErrClosed
+		return false, net.ErrClosed
 	default:
 	}
-	c.setupReadTimeout(ctx)
+	timeoutSet := c.setupReadTimeout(ctx)
 
 	c.closeStateMu.Lock()
 	closeReceivedErr := c.closeReceivedErr
 	c.closeStateMu.Unlock()
 	if closeReceivedErr != nil {
-		c.clearReadTimeout()
-		return closeReceivedErr
+		if timeoutSet {
+			c.clearReadTimeout()
+		}
+		return false, closeReceivedErr
 	}
 
-	return 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) {
-	c.clearReadTimeout()
+func (c *Conn) finishRead(ctx context.Context, err *error, timeoutSet bool) {
+	if timeoutSet {
+		c.clearReadTimeout()
+	}
 	select {
 	case <-c.closed:
 		if *err != nil {
@@ -254,11 +258,11 @@ func (c *Conn) finishRead(ctx context.Context, err *error) {
 }
 
 func (c *Conn) readFrameHeader(ctx context.Context) (_ header, err error) {
-	err = c.prepareRead(ctx)
+	timeoutSet, err := c.prepareRead(ctx)
 	if err != nil {
 		return header{}, err
 	}
-	defer c.finishRead(ctx, &err)
+	defer c.finishRead(ctx, &err, timeoutSet)
 
 	h, err := readFrameHeader(c.br, c.readHeaderBuf[:])
 	if err != nil {
@@ -269,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) {
-	err = c.prepareRead(ctx)
+	timeoutSet, err := c.prepareRead(ctx)
 	if err != nil {
 		return 0, err
 	}
-	defer c.finishRead(ctx, &err)
+	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