Health-warnings for stashing callback conn handles, and try to safeguard use

pkg/wsclient/wsclient.go

@@ -151,13 +151,17 @@ type wsClient struct {
 type WSPreConnectHandler func(ctx context.Context, w WSClient) error
 
 // WSPostConnectHandler will be called after every connect/reconnect. Can send data over ws, but must not block listening for data on the ws.
-// Note: During auto-cycle this is called on the new connection, after the WSPreDisconnectHandler is called on the old one, but before the old connection is closed.
-type WSPostConnectHandler func(ctx context.Context, w WSClient) error
+// During auto-cycle it is passed an establishing WSConn handle that will route to the NEW connection.
+// You should not use this connection beyond the scope of the callback.
+type WSPostConnectHandler func(ctx context.Context, establishing WSClient) error
 
 // WSPreDisconnectHandler is called before a graceful close, to allow cleanup (such as unsubscribe):
 //   - When closed explicitly
 //   - When cycling the connection (after the new connection is established, before post-connect is called)
-type WSPreDisconnectHandler func(ctx context.Context, w WSClient) error
+//
+// Passed a disconnecting WSConn handle that will route to the OLD connection.
+// You should not use this connection beyond the scope of the callback.
+type WSPreDisconnectHandler func(ctx context.Context, disconnecting WSClient) error
 
 // New creates a new outbound client that can be connected to a remote server.
 // ** Recommend using NewWithConfig directly **
@@ -480,7 +484,26 @@ type connBoundClient struct {
 	c *wsConnection
 }
 
+func (bc *connBoundClient) isRetired() bool {
+	// Selecting on sendDone/default gets compiled down to an efficient `runtime.selectnbrecv`,
+	// so it's acceptable overhead on the Send() critical path.
+	select {
+	case <-bc.c.sendDone:
+		select {
+		case <-bc.c.promoted:
+			return true
+		default:
+		}
+	default:
+	}
+	return false
+}
+
 func (bc *connBoundClient) Send(ctx context.Context, message []byte) error {
+	if !bc.disableReconnect && bc.isRetired() {
+		log.L(ctx).Debugf("WS %s send on retired connection routed to current connection", bc.url)
+		return bc.wsClient.Send(ctx, message)
+	}
 	if err := bc.waitRateLimiter(ctx); err != nil {
 		return err
 	}

pkg/wsclient/wsclient_cycle_test.go

@@ -734,3 +734,56 @@ func TestConnBoundClientSendWriteTimeout(t *testing.T) {
 	err := w.boundTo(c).Send(ctx, []byte(`never written`))
 	assert.Regexp(t, "FF00146", err)
 }
+
+// A hook facade retained after its connection is gone falls back to the current connection,
+// so consumers that captured the afterConnect WSClient keep working across reconnects
+func TestConnBoundClientRetainedAcrossReconnect(t *testing.T) {
+	connections, url, _, done := NewTestWSServerMulti(nil)
+	defer done()
+
+	var retained atomic.Pointer[WSClient]
+	wsConfig := &WSConfig{
+		HTTPURL:      url,
+		InitialDelay: 5 * time.Millisecond,
+		MaximumDelay: 20 * time.Millisecond,
+		PostConnectHandler: func(ctx context.Context, w WSClient) error {
+			retained.CompareAndSwap(nil, &w) // keep the first connection's facade
+			return w.Send(ctx, []byte(`subscribe`))
+		},
+	}
+	wsc, err := NewWithConfig(context.Background(), wsConfig)
+	assert.NoError(t, err)
+	assert.NoError(t, wsc.Connect())
+	defer wsc.Close()
+
+	conn1 := nextConn(t, connections)
+	expectMsg(t, conn1, `subscribe`)
+	facade := *retained.Load()
+
+	// Still bound to conn1 while it is live
+	assert.NoError(t, facade.Send(context.Background(), []byte(`on conn1`)))
+	expectMsg(t, conn1, `on conn1`)
+
+	// Server drops conn1 - the client reconnects on conn2
+	conn1.CloseConn()
+	waitDone(t, conn1)
+	conn2 := nextConn(t, connections)
+	expectMsg(t, conn2, `subscribe`)
+
+	// The retained facade now routes via the shared client to conn2
+	assert.NoError(t, facade.Send(context.Background(), []byte(`after reconnect`)))
+	expectMsg(t, conn2, `after reconnect`)
+}
+
+func TestConnBoundClientUnpromotedStaysStrict(t *testing.T) {
+	w := &wsClient{
+		ctx:     context.Background(),
+		closing: make(chan struct{}),
+	}
+	// A connection that was torn down without ever being promoted (e.g. afterConnect failed
+	// during a cycle) must not leak sends onto the current connection
+	c := &wsConnection{w: w, send: make(chan *trackedSend), sendDone: make(chan []byte), promoted: make(chan struct{})}
+	close(c.sendDone)
+	err := w.boundTo(c).Send(context.Background(), []byte(`a`))
+	assert.Regexp(t, "FF00147", err)
+}