[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

@@ -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
+// sendReceiveProbe sends a self-addressed heartbeat bound to ctx, so cancelStream
-// back to this client, exercising the exact receive path the watchdog guards.
+// aborts an in-flight probe instead of leaving the watchdog blocked on send timeouts.
-func (c *GrpcClient) sendReceiveProbe() error {
+// 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},

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: