From 5f0831789f31dcd9466d1909b5ac28ebbeed2ece Mon Sep 17 00:00:00 2001
From: Anmol Sethi <hi@nhooyr.io>
Date: Sat, 9 Jan 2021 08:22:31 -0500
Subject: [PATCH] netconn: Cleanup contexts on close

Updates #255
---
 netconn.go | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/netconn.go b/netconn.go
index c6f8dc13..4046137f 100644
--- 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, "")
 }