From 1cad1a4c42edab7c742b093406becc381856d5de Mon Sep 17 00:00:00 2001 From: Peter Broadhurst Date: Mon, 5 Oct 2026 13:58:22 -0400 Subject: [PATCH 1/4] Document health-warnings for stashing conn handles, and try to protect against misuse Signed-off-by: Peter Broadhurst --- pkg/wsclient/wsclient.go | 29 ++++++++++++++-- pkg/wsclient/wsclient_cycle_test.go | 53 +++++++++++++++++++++++++++++ 2 files changed, 79 insertions(+), 3 deletions(-) diff --git a/pkg/wsclient/wsclient.go b/pkg/wsclient/wsclient.go index d0f3789..d394c43 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 to 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 to 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.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) +} From c37a844ab46d2b565d04f498a8eefc647f9d3279 Mon Sep 17 00:00:00 2001 From: Peter Broadhurst Date: Tue, 6 Oct 2026 07:47:29 -0400 Subject: [PATCH 2/4] Update pkg/wsclient/wsclient.go Co-authored-by: Enrique Lacal Signed-off-by: Peter Broadhurst --- pkg/wsclient/wsclient.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pkg/wsclient/wsclient.go b/pkg/wsclient/wsclient.go index d394c43..a4f12e1 100644 --- a/pkg/wsclient/wsclient.go +++ b/pkg/wsclient/wsclient.go @@ -160,7 +160,7 @@ type WSPostConnectHandler func(ctx context.Context, establishing WSClient) error // - When cycling the connection (after the new connection is established, before post-connect is called) // // Passed a disconnecting WSConn handle that will route to the OLD connection. -// You should not to use this connection beyond the scope of the callback. +// 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. From 9d0f7b1d7c025e801e8cbcedefdd1b68a5257592 Mon Sep 17 00:00:00 2001 From: Peter Broadhurst Date: Tue, 6 Oct 2026 07:47:37 -0400 Subject: [PATCH 3/4] Update pkg/wsclient/wsclient.go Co-authored-by: Enrique Lacal Signed-off-by: Peter Broadhurst --- pkg/wsclient/wsclient.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pkg/wsclient/wsclient.go b/pkg/wsclient/wsclient.go index a4f12e1..748ff22 100644 --- a/pkg/wsclient/wsclient.go +++ b/pkg/wsclient/wsclient.go @@ -152,7 +152,7 @@ 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. // During auto-cycle it is passed an establishing WSConn handle that will route to the NEW connection. -// You should not to use this connection beyond the scope of the callback. +// 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): From a9acd5a9a495b5fa84cc4785d277c3f7a50b94a5 Mon Sep 17 00:00:00 2001 From: Peter Broadhurst Date: Tue, 6 Oct 2026 07:47:44 -0400 Subject: [PATCH 4/4] Update pkg/wsclient/wsclient.go Co-authored-by: Enrique Lacal Signed-off-by: Peter Broadhurst --- pkg/wsclient/wsclient.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pkg/wsclient/wsclient.go b/pkg/wsclient/wsclient.go index 748ff22..af28604 100644 --- a/pkg/wsclient/wsclient.go +++ b/pkg/wsclient/wsclient.go @@ -500,7 +500,7 @@ func (bc *connBoundClient) isRetired() bool { } func (bc *connBoundClient) Send(ctx context.Context, message []byte) error { - if bc.isRetired() { + 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) }