socket: use more atomic in FrameSocket

Closes #1150
Fixes #1085

socket/framesocket.go

@@ -12,6 +12,7 @@ import (
 	"fmt"
 	"net/http"
 	"sync"
+	"sync/atomic"
 
 	"github.com/coder/websocket"
 
@@ -22,7 +23,7 @@ type FrameSocket struct {
 	parentCtx context.Context
 	cancelCtx context.Context
 	cancel    context.CancelFunc
-	conn      *websocket.Conn
+	conn      atomic.Pointer[websocket.Conn]
 	log       waLog.Logger
 	lock      sync.Mutex
 
@@ -35,7 +36,7 @@ type FrameSocket struct {
 
 	Header []byte
 
-	closed bool
+	closed atomic.Bool
 
 	incomingLength int
 	receivedLength int
@@ -56,30 +57,30 @@ func NewFrameSocket(log waLog.Logger, client *http.Client) *FrameSocket {
 }
 
 func (fs *FrameSocket) IsConnected() bool {
-	return fs.conn != nil
+	return fs.conn.Load() != nil
 }
 
 func (fs *FrameSocket) Close(code websocket.StatusCode) {
 	fs.lock.Lock()
 	defer fs.lock.Unlock()
 
-	if fs.conn == nil {
+	conn := fs.conn.Swap(nil)
+	if conn == nil {
 		return
 	}
 
-	fs.closed = true
+	fs.closed.Store(true)
 	if code > 0 {
-		err := fs.conn.Close(code, "")
+		err := conn.Close(code, "")
 		if err != nil {
 			fs.log.Warnf("Error sending close to websocket: %v", err)
 		}
 	} else {
-		err := fs.conn.CloseNow()
+		err := conn.CloseNow()
 		if err != nil {
 			fs.log.Debugf("Error force closing websocket: %v", err)
 		}
 	}
-	fs.conn = nil
 	fs.cancel()
 	fs.cancel = nil
 	if fs.OnDisconnect != nil {
@@ -90,7 +91,7 @@ func (fs *FrameSocket) Close(code websocket.StatusCode) {
 func (fs *FrameSocket) Connect(ctx context.Context) error {
 	fs.lock.Lock()
 	defer fs.lock.Unlock()
-	if fs.conn != nil {
+	if fs.conn.Load() != nil {
 		return ErrSocketAlreadyOpen
 	}
 	fs.parentCtx = ctx
@@ -107,7 +108,7 @@ func (fs *FrameSocket) Connect(ctx context.Context) error {
 	}
 	conn.SetReadLimit(FrameMaxSize)
 
-	fs.conn = conn
+	fs.conn.Store(conn)
 
 	go fs.readPump(conn, ctx)
 	return nil
@@ -118,7 +119,7 @@ func (fs *FrameSocket) Context() context.Context {
 }
 
 func (fs *FrameSocket) SendFrame(data []byte) error {
-	conn := fs.conn
+	conn := fs.conn.Load()
 	if conn == nil {
 		return ErrSocketClosed
 	}
@@ -209,7 +210,7 @@ func (fs *FrameSocket) readPump(conn *websocket.Conn, ctx context.Context) {
 		msgType, data, err := conn.Read(ctx)
 		if err != nil {
 			// Ignore the error if the context has been closed
-			if !fs.closed && !errors.Is(ctx.Err(), context.Canceled) {
+			if !fs.closed.Load() && !errors.Is(ctx.Err(), context.Canceled) {
 				fs.log.Errorf("Error reading from websocket: %v", err)
 			}
 			return