diff --git a/pkg/wsclient/wsclient.go b/pkg/wsclient/wsclient.go index d0f3789..af28604 100644 --- a/pkg/wsclient/wsclient.go +++ b/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 } diff --git a/pkg/wsclient/wsclient_cycle_test.go b/pkg/wsclient/wsclient_cycle_test.go index 9ef0fc1..00e4317 100644 --- a/pkg/wsclient/wsclient_cycle_test.go +++ b/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) +}