netconn.go: Cleanup contexts on close

nhooyr · nhooyr · commit 5070c1479d4d · 2021-01-09T08:22:52.000-05:00
Updates #255
diff --git a/netconn.go b/netconn.go
@@ -50,10 +50,8 @@ func NetConn(ctx context.Context, c *Conn, msgType MessageType) net.Conn {
 		writeMu: newMu(c),
 	}
 
-	var writeCancel context.CancelFunc
-	nc.writeCtx, writeCancel = context.WithCancel(ctx)
-	var readCancel context.CancelFunc
-	nc.readCtx, readCancel = context.WithCancel(ctx)
+	nc.writeCtx, nc.writeCancel = context.WithCancel(ctx)
+	nc.readCtx, nc.readCancel = context.WithCancel(ctx)
 
 	nc.writeTimer = time.AfterFunc(math.MaxInt64, func() {
 		if !nc.writeMu.tryLock() {
@@ -98,11 +96,13 @@ type netConn struct {
 	writeMu      *mu
 	writeExpired int64
 	writeCtx     context.Context
+	writeCancel  context.CancelFunc
 
 	readTimer   *time.Timer
 	readMu      *mu
 	readExpired int64
 	readCtx     context.Context
+	readCancel  context.CancelFunc
 	readEOFed   bool
 	reader      io.Reader
 }
@@ -111,7 +111,9 @@ var _ net.Conn = &netConn{}
 
 func (nc *netConn) Close() error {
 	nc.writeTimer.Stop()
+	nc.writeCancel()
 	nc.readTimer.Stop()
+	nc.readCancel()
 	return nc.c.Close(StatusNormalClosure, "")
 }
 
