wsjs: Ensure no goroutines leak after Close

Closes#330

Author: nhooyr

close.go

@@ -99,14 +99,14 @@ func CloseStatus(err error) StatusCode{
 // Close will unblock all goroutines interacting with the connection once
 // complete.
 func (c*Conn) Close(codeStatusCode, reasonstring) error{
-deferc.wgWait()
+deferc.wg.Wait()
 returnc.closeHandshake(code, reason)
 }
 
 // CloseNow closes the WebSocket connection without attempting a close handshake.
 // Use when you do not want the overhead of the close handshake.
 func (c*Conn) CloseNow() (errerror){
-deferc.wgWait()
+deferc.wg.Wait()
 defererrd.Wrap(&err, "failed to close WebSocket")
 
 ifc.isClosed(){

conn.go

@@ -45,8 +45,6 @@ const (
 typeConnstruct{
 noCopy
 
-wg sync.WaitGroup
-
 subprotocolstring
 rwc io.ReadWriteCloser
 clientbool
@@ -72,6 +70,7 @@ type Conn struct{
 writeHeaderBuf [8]byte
 writeHeaderheader
 
+wg sync.WaitGroup
 closedchanstruct{}
 closeMu sync.Mutex
 closeErrerror
@@ -132,7 +131,11 @@ func newConn(cfg connConfig) *Conn{
 c.close(errors.New("connection garbage collected"))
  })
 
-c.wgGo(c.timeoutLoop)
+c.wg.Add(1)
+gofunc(){
+deferc.wg.Done()
+c.timeoutLoop
+ }()
 
 returnc
 }
@@ -163,10 +166,12 @@ func (c *Conn) close(err error){
 // closeErr.
 c.rwc.Close()
 
-c.wgGo(func(){
+c.wg.Add(1)
+gofunc(){
+deferc.wg.Done()
 c.msgWriter.close()
 c.msgReader.close()
- })
+ }()
 }
 
 func (c*Conn) timeoutLoop(){
@@ -183,9 +188,11 @@ func (c *Conn) timeoutLoop(){
 
 case<-readCtx.Done():
 c.setCloseErr(fmt.Errorf("read timed out: %w", readCtx.Err()))
-c.wgGo(func(){
+c.wg.Add(1)
+gofunc(){
+deferc.wg.Done()
 c.writeError(StatusPolicyViolation, errors.New("read timed out"))
- })
+ }()
 case<-writeCtx.Done():
 c.close(fmt.Errorf("write timed out: %w", writeCtx.Err()))
 return
@@ -302,15 +309,3 @@ func (m *mu) unlock(){
 typenoCopystruct{}
 
 func (*noCopy) Lock(){}
-
-func (c*Conn) wgGo(fnfunc()){
-c.wg.Add(1)
-gofunc(){
-deferc.wg.Done()
-fn()
- }()
-}
-
-func (c*Conn) wgWait(){
-c.wg.Wait()
-}

ws_js.go

@@ -47,6 +47,7 @@ type Conn struct{
 // read limit for a message in bytes.
 msgReadLimit xsync.Int64
 
+wg sync.WaitGroup
 closingMu sync.Mutex
 isReadClosed xsync.Int64
 closeOnce sync.Once
@@ -223,6 +224,7 @@ func (c *Conn) write(ctx context.Context, typ MessageType, p []byte) error{
 // or the connection is closed.
 // It thus performs the full WebSocket close handshake.
 func (c*Conn) Close(codeStatusCode, reasonstring) error{
+deferc.wg.Wait()
 err:=c.exportedClose(code, reason)
 iferr!=nil{
 returnfmt.Errorf("failed to close WebSocket: %w", err)
@@ -236,6 +238,7 @@ func (c *Conn) Close(code StatusCode, reason string) error{
 // note: No different from Close(StatusGoingAway, "") in WASM as there is no way to close
 // a WebSocket without the close handshake.
 func (c*Conn) CloseNow() error{
+deferc.wg.Wait()
 returnc.Close(StatusGoingAway, "")
 }
 
@@ -388,10 +391,15 @@ func (c *Conn) CloseRead(ctx context.Context) context.Context{
 c.isReadClosed.Store(1)
 
 ctx, cancel:=context.WithCancel(ctx)
+c.wg.Add(1)
 gofunc(){
+deferc.CloseNow()
+deferc.wg.Done()
 defercancel()
-c.read(ctx)
-c.Close(StatusPolicyViolation, "unexpected data message")
+_, _, err:=c.read(ctx)
+iferr!=nil{
+c.Close(StatusPolicyViolation, "unexpected data message")
+ }
  }()
 returnctx
 }