sendnrw/netbird
2
0
Fork 0
mirror of https://github.com/netbirdio/netbird.git synced 2026-07-18 12:39:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

[client] Bind signal receive probe to the stream context

Browse Source 
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.
This commit is contained in:
Zoltan Papp
2026-06-29 02:15:59 +02:00
parent 12e4bf698b
commit e4437d6241
2 changed files with 14 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
shared/signal/client/grpc.go
View File

@@ -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},

2
shared/signal/client/watchdog_test.go
View File

@@ -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: