From e4437d6241d8dc69dab4c27e2a21897fd8e222ec Mon Sep 17 00:00:00 2001
From: Zoltan Papp <zoltan.pmail@gmail.com>
Date: Mon, 29 Jun 2026 02:15:59 +0200
Subject: [PATCH] [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          | 19 +++++++++++++------
 shared/signal/client/watchdog_test.go |  2 +-
 2 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/shared/signal/client/grpc.go b/shared/signal/client/grpc.go
index d8dfe4337..7e8e551b6 100644
--- a/shared/signal/client/grpc.go
+++ b/shared/signal/client/grpc.go
@@ -409,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")
 	}
@@ -425,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)
 
@@ -495,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()
@@ -504,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: