close: Implement CloseNow

Closes #384

Author: nhooyr

close.go

@@ -102,6 +102,19 @@ func (c *Conn) Close(code StatusCode, reason string) error {
 	return c.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() (err error) {
+	defer errd.Wrap(&err, "failed to close WebSocket")
+
+	if c.isClosed() {
+		return errClosed
+	}
+
+	c.close(nil)
+	return c.closeErr
+}
+
 func (c *Conn) closeHandshake(code StatusCode, reason string) (err error) {
 	defer errd.Wrap(&err, "failed to close WebSocket")
 
@@ -265,7 +278,7 @@ func (c *Conn) setCloseErr(err error) {
 }
 
 func (c *Conn) setCloseErrLocked(err error) {
-	if c.closeErr == nil {
+	if c.closeErr == nil && err != nil {
 		c.closeErr = fmt.Errorf("WebSocket closed: %w", err)
 	}
 }

conn.go

@@ -147,6 +147,9 @@ func (c *Conn) close(err error) {
 	if c.isClosed() {
 		return
 	}
+	if err == nil {
+		err = c.rwc.Close()
+	}
 	c.setCloseErrLocked(err)
 	close(c.closed)
 	runtime.SetFinalizer(c, nil)

conn_test.go

@@ -295,6 +295,19 @@ func TestConn(t *testing.T) {
 		err = c1.Close(websocket.StatusNormalClosure, "")
 		assert.Success(t, err)
 	})
+
+	t.Run("CloseNow", func(t *testing.T) {
+		_, c1, c2 := newConnTest(t, nil, nil)
+
+		err1 := c1.CloseNow()
+		err2 := c2.CloseNow()
+		assert.Success(t, err1)
+		assert.Success(t, err2)
+		err1 = c1.CloseNow()
+		err2 = c2.CloseNow()
+		assert.ErrorIs(t, websocket.ErrClosed, err1)
+		assert.ErrorIs(t, websocket.ErrClosed, err2)
+	})
 }
 
 func TestWasm(t *testing.T) {

export_test.go

@@ -23,3 +23,5 @@ func (c *Conn) RecordBytesRead() *int {
 	}))
 	return &bytesRead
 }
+
+var ErrClosed = errClosed