Update

vendor/github.com/coder/websocket/conn.go

@@ -1,5 +1,4 @@
 //go:build !js
-// +build !js
 
 package websocket
 
@@ -52,9 +51,8 @@ type Conn struct {
 	br             *bufio.Reader
 	bw             *bufio.Writer
 
-	readTimeout     chan context.Context
-	writeTimeout    chan context.Context
-	timeoutLoopDone chan struct{}
+	readTimeoutStop  atomic.Pointer[func() bool]
+	writeTimeoutStop atomic.Pointer[func() bool]
 
 	// Read state.
 	readMu         *mu
@@ -69,17 +67,25 @@ type Conn struct {
 	writeHeaderBuf [8]byte
 	writeHeader    header
 
+	// Close handshake state.
+	closeStateMu     sync.RWMutex
+	closeReceivedErr error
+	closeSentErr     error
+
+	// CloseRead state.
 	closeReadMu   sync.Mutex
 	closeReadCtx  context.Context
 	closeReadDone chan struct{}
 
+	closing atomic.Bool
+	closeMu sync.Mutex // Protects following.
 	closed  chan struct{}
-	closeMu sync.Mutex
-	closing bool
 
-	pingCounter   int32
-	activePingsMu sync.Mutex
-	activePings   map[string]chan<- struct{}
+	pingCounter    atomic.Int64
+	activePingsMu  sync.Mutex
+	activePings    map[string]chan<- struct{}
+	onPingReceived func(context.Context, []byte) bool
+	onPongReceived func(context.Context, []byte)
 }
 
 type connConfig struct {
@@ -88,6 +94,8 @@ type connConfig struct {
 	client         bool
 	copts          *compressionOptions
 	flateThreshold int
+	onPingReceived func(context.Context, []byte) bool
+	onPongReceived func(context.Context, []byte)
 
 	br *bufio.Reader
 	bw *bufio.Writer
@@ -104,12 +112,10 @@ func newConn(cfg connConfig) *Conn {
 		br: cfg.br,
 		bw: cfg.bw,
 
-		readTimeout:     make(chan context.Context),
-		writeTimeout:    make(chan context.Context),
-		timeoutLoopDone: make(chan struct{}),
-
-		closed:      make(chan struct{}),
-		activePings: make(map[string]chan<- struct{}),
+		closed:         make(chan struct{}),
+		activePings:    make(map[string]chan<- struct{}),
+		onPingReceived: cfg.onPingReceived,
+		onPongReceived: cfg.onPongReceived,
 	}
 
 	c.readMu = newMu(c)
@@ -133,8 +139,6 @@ func newConn(cfg connConfig) *Conn {
 		c.close()
 	})
 
-	go c.timeoutLoop()
-
 	return c
 }
 
@@ -164,27 +168,44 @@ func (c *Conn) close() error {
 	return err
 }
 
-func (c *Conn) timeoutLoop() {
-	defer close(c.timeoutLoopDone)
+func (c *Conn) setupWriteTimeout(ctx context.Context) bool {
+	if ctx.Done() == nil {
+		return false
+	}
 
-	readCtx := context.Background()
-	writeCtx := context.Background()
+	stop := context.AfterFunc(ctx, func() {
+		c.clearWriteTimeout()
+		c.close()
+	})
+	swapTimeoutStop(&c.writeTimeoutStop, &stop)
+	return true
+}
 
-	for {
-		select {
-		case <-c.closed:
-			return
+func (c *Conn) clearWriteTimeout() {
+	swapTimeoutStop(&c.writeTimeoutStop, nil)
+}
 
-		case writeCtx = <-c.writeTimeout:
-		case readCtx = <-c.readTimeout:
+func (c *Conn) setupReadTimeout(ctx context.Context) bool {
+	if ctx.Done() == nil {
+		return false
+	}
 
-		case <-readCtx.Done():
-			c.close()
-			return
-		case <-writeCtx.Done():
-			c.close()
-			return
-		}
+	stop := context.AfterFunc(ctx, func() {
+		c.clearReadTimeout()
+		c.close()
+	})
+	swapTimeoutStop(&c.readTimeoutStop, &stop)
+	return true
+}
+
+func (c *Conn) clearReadTimeout() {
+	swapTimeoutStop(&c.readTimeoutStop, nil)
+}
+
+func swapTimeoutStop(p *atomic.Pointer[func() bool], newStop *func() bool) {
+	oldStop := p.Swap(newStop)
+	if oldStop != nil {
+		(*oldStop)()
 	}
 }
 
@@ -200,9 +221,9 @@ func (c *Conn) flate() bool {
 //
 // TCP Keepalives should suffice for most use cases.
 func (c *Conn) Ping(ctx context.Context) error {
-	p := atomic.AddInt32(&c.pingCounter, 1)
+	p := c.pingCounter.Add(1)
 
-	err := c.ping(ctx, strconv.Itoa(int(p)))
+	err := c.ping(ctx, strconv.FormatInt(p, 10))
 	if err != nil {
 		return fmt.Errorf("failed to ping: %w", err)
 	}