Skip to content

Commit

Permalink
close: Implement CloseNow
Browse files Browse the repository at this point in the history
Closes #384
  • Loading branch information
nhooyr committed Oct 14, 2023
1 parent 1dbc141 commit a94999f
Show file tree
Hide file tree
Showing 4 changed files with 32 additions and 1 deletion.
15 changes: 14 additions & 1 deletion close.go
Expand Up @@ -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")

Expand Down Expand Up @@ -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)
}
}
Expand Down
3 changes: 3 additions & 0 deletions conn.go
Expand Up @@ -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)
Expand Down
13 changes: 13 additions & 0 deletions conn_test.go
Expand Up @@ -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) {
Expand Down
2 changes: 2 additions & 0 deletions export_test.go
Expand Up @@ -23,3 +23,5 @@ func (c *Conn) RecordBytesRead() *int {
}))
return &bytesRead
}

var ErrClosed = errClosed

0 comments on commit a94999f

Please sign in to comment.