From deff8af59f13222d245abfdfc7e2e700c74dd8c7 Mon Sep 17 00:00:00 2001 From: Zoltan Papp Date: Mon, 29 Jun 2026 11:24:25 +0200 Subject: [PATCH] [client] Wait for signal receive watchdog to stop before reconnect (#6574) * [client] Wait for signal receive watchdog to stop before reconnect The per-stream watchReceiveStream goroutine was started fire-and-forget and never joined. On reconnect a lingering watchdog could still flip shared client state (receiveStalled, the disconnect notifier) on the freshly established stream, since cancelStream only cancels its own stream context. Track the watchdog with a WaitGroup and wait for it to exit (after cancelling its stream) before the operation returns, so each reconnect starts with no stale watchdog. * [client] Bind signal receive probe to the stream context The watchdog probe reused the generic Send, which derives its per-attempt timeouts from the long-lived client context, so cancelStream could not interrupt an in-flight probe. After joining the watchdog on reconnect, watchdogWg.Wait() could then block for the full send-attempt chain. Split Send into a context-aware send and pass the stream context down through sendReceiveProbe, so cancelStream aborts any in-flight probe and the watchdog exits promptly. --- shared/signal/client/grpc.go | 30 ++++++++++++++++++++------- shared/signal/client/watchdog_test.go | 2 +- 2 files changed, 24 insertions(+), 8 deletions(-) diff --git a/shared/signal/client/grpc.go b/shared/signal/client/grpc.go index 611ab0c45..7e8e551b6 100644 --- a/shared/signal/client/grpc.go +++ b/shared/signal/client/grpc.go @@ -85,6 +85,7 @@ type GrpcClient struct { // receive backpressure as a dead stream: reconnecting cannot help, since the // new stream feeds the same worker, and only triggers a reconnect storm. receiveHandoffBlocked atomic.Bool + watchdogWg sync.WaitGroup } // NewClient creates a new Signal client @@ -200,10 +201,18 @@ func (c *GrpcClient) Receive(ctx context.Context, msgHandler func(msg *proto.Mes // Guard the receive direction: the transport can stay healthy while the // server stops delivering messages. The watchdog reconnects via cancelStream. c.markReceived() - go c.watchReceiveStream(streamCtx, cancelStream) + c.watchdogWg.Add(1) + go func() { + defer c.watchdogWg.Done() + c.watchReceiveStream(streamCtx, cancelStream) + }() // start receiving messages from the Signal stream (from other peers through signal) err = c.receive(stream) + + cancelStream() + c.watchdogWg.Wait() + if err != nil { // Check the parent context, not streamCtx: a watchdog-triggered // cancelStream must reconnect, only a parent cancel is shutdown. @@ -400,7 +409,12 @@ func (c *GrpcClient) encryptMessage(msg *proto.Message) (*proto.EncryptedMessage // Send sends a message to the remote Peer through the Signal Exchange. func (c *GrpcClient) Send(msg *proto.Message) error { + return c.send(c.ctx, msg) +} +// send delivers a message deriving per-attempt timeouts from parentCtx, so a +// caller can abort an in-flight send by cancelling that context. +func (c *GrpcClient) send(parentCtx context.Context, msg *proto.Message) error { if !c.Ready() { return fmt.Errorf("no connection to signal") } @@ -416,7 +430,7 @@ func (c *GrpcClient) Send(msg *proto.Message) error { if attempt > 1 { attemptTimeout = time.Duration(attempt) * 5 * time.Second } - ctx, cancel := context.WithTimeout(c.ctx, attemptTimeout) + ctx, cancel := context.WithTimeout(parentCtx, attemptTimeout) _, err = c.realClient.Send(ctx, encryptedMessage) @@ -486,7 +500,7 @@ func (c *GrpcClient) watchReceiveStream(ctx context.Context, cancelStream contex } if probeSentAt.IsZero() { - if err := c.sendReceiveProbe(); err != nil { + if err := c.sendReceiveProbe(ctx); err != nil { log.Debugf("failed to send signal receive probe: %v", err) } probeSentAt = time.Now() @@ -495,11 +509,13 @@ func (c *GrpcClient) watchReceiveStream(ctx context.Context, cancelStream contex } } -// sendReceiveProbe sends a self-addressed heartbeat. The Signal server routes it -// back to this client, exercising the exact receive path the watchdog guards. -func (c *GrpcClient) sendReceiveProbe() error { +// sendReceiveProbe sends a self-addressed heartbeat bound to ctx, so cancelStream +// aborts an in-flight probe instead of leaving the watchdog blocked on send timeouts. +// The Signal server routes it back to this client, exercising the exact receive +// path the watchdog guards. +func (c *GrpcClient) sendReceiveProbe(ctx context.Context) error { self := c.key.PublicKey().String() - return c.Send(&proto.Message{ + return c.send(ctx, &proto.Message{ Key: self, RemoteKey: self, Body: &proto.Body{Type: proto.Body_HEARTBEAT}, diff --git a/shared/signal/client/watchdog_test.go b/shared/signal/client/watchdog_test.go index bc6b5520b..eeb9aec30 100644 --- a/shared/signal/client/watchdog_test.go +++ b/shared/signal/client/watchdog_test.go @@ -74,7 +74,7 @@ func TestReceiveProbeRoundTrips(t *testing.T) { t.Fatal("signal stream did not connect within timeout") } - require.NoError(t, client.sendReceiveProbe()) + require.NoError(t, client.sendReceiveProbe(ctx)) select { case <-received: