[client] Detect and recover from a stalled signal receive stream (#6459)

shared/signal/client/grpc.go

@@ -2,9 +2,11 @@ package client
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/cenkalti/backoff/v4"
@@ -23,7 +25,23 @@ import (
 	"github.com/netbirdio/netbird/util/wsproxy"
 )
 
-const healthCheckTimeout = 5 * time.Second
+const (
+	// receiveInactivityThreshold is how long the receive stream may be silent
+	// before the watchdog actively probes it. The gRPC transport can stay
+	// healthy (keepalive satisfied) while the server stops delivering messages,
+	// which the transport layer cannot detect.
+	receiveInactivityThreshold = 30 * time.Second
+	// receiveProbeTimeout is how long the watchdog waits for its self-addressed
+	// probe to round-trip back on the stream before declaring the receive
+	// direction dead.
+	receiveProbeTimeout = 10 * time.Second
+	// receiveWatchdogInterval is how often the watchdog evaluates the stream.
+	receiveWatchdogInterval = 10 * time.Second
+)
+
+// errReceiveStreamStalled is reported when the receive stream is transport-alive
+// but no longer delivering messages, so the stream is torn down to reconnect.
+var errReceiveStreamStalled = errors.New("signal receive stream stalled")
 
 // ConnStateNotifier is a wrapper interface of the status recorder
 type ConnStateNotifier interface {
@@ -52,6 +70,14 @@ type GrpcClient struct {
 	decryptionWorker       *Worker
 	decryptionWorkerCancel context.CancelFunc
 	decryptionWg           sync.WaitGroup
+
+	// lastReceived holds the Unix-nano timestamp of the last message read from
+	// the receive stream, used by the receive watchdog.
+	lastReceived atomic.Int64
+	// receiveStalled is set by the receive watchdog when the stream is
+	// transport-alive but no longer delivering messages. It is the source of
+	// truth IsHealthy reads, and is cleared once any frame is received again.
+	receiveStalled atomic.Bool
 }
 
 // NewClient creates a new Signal client
@@ -148,9 +174,9 @@ func (c *GrpcClient) Receive(ctx context.Context, msgHandler func(msg *proto.Mes
 
 		// connect to Signal stream identifying ourselves with a public WireGuard key
 		// todo once the key rotation logic has been implemented, consider changing to some other identifier (received from management)
-		ctx, cancelStream := context.WithCancel(ctx)
+		streamCtx, cancelStream := context.WithCancel(ctx)
 		defer cancelStream()
-		stream, err := c.connect(ctx, c.key.PublicKey().String())
+		stream, err := c.connect(streamCtx, c.key.PublicKey().String())
 		if err != nil {
 			log.Warnf("disconnected from the Signal Exchange due to an error: %v", err)
 			return err
@@ -164,9 +190,16 @@ func (c *GrpcClient) Receive(ctx context.Context, msgHandler func(msg *proto.Mes
 		// Start worker pool if not already started
 		c.startEncryptionWorker(msgHandler)
 
+		// 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)
+
 		// start receiving messages from the Signal stream (from other peers through signal)
 		err = c.receive(stream)
 		if err != nil {
+			// Check the parent context, not streamCtx: a watchdog-triggered
+			// cancelStream must reconnect, only a parent cancel is shutdown.
 			if ctx.Err() != nil {
 				log.Debugf("signal connection context has been canceled, this usually indicates shutdown")
 				return nil
@@ -252,7 +285,10 @@ func (c *GrpcClient) Ready() bool {
 	return c.signalConn.GetState() == connectivity.Ready || c.signalConn.GetState() == connectivity.Idle
 }
 
-// IsHealthy probes the gRPC connection and returns false on errors
+// IsHealthy reports whether the Signal connection is usable, based on the
+// transport state plus the receive watchdog's verdict, and updates the status
+// recorder accordingly. It does not actively probe: the watchdog
+// (watchReceiveStream) owns probing the receive path and reconnecting.
 func (c *GrpcClient) IsHealthy() bool {
 	switch c.signalConn.GetState() {
 	case connectivity.TransientFailure:
@@ -265,16 +301,8 @@ func (c *GrpcClient) IsHealthy() bool {
 	case connectivity.Ready:
 	}
 
-	ctx, cancel := context.WithTimeout(c.ctx, healthCheckTimeout)
-	defer cancel()
-	_, err := c.realClient.Send(ctx, &proto.EncryptedMessage{
-		Key:       c.key.PublicKey().String(),
-		RemoteKey: "dummy",
-		Body:      nil,
-	})
-	if err != nil {
-		c.notifyDisconnected(err)
-		log.Warnf("health check returned: %s", err)
+	if c.receiveStalled.Load() {
+		c.notifyDisconnected(errReceiveStreamStalled)
 		return false
 	}
 	c.notifyConnected()
@@ -398,6 +426,68 @@ func (c *GrpcClient) Send(msg *proto.Message) error {
 	return err
 }
 
+// markReceived records that a frame was just read from the receive stream and
+// clears the stalled flag.
+func (c *GrpcClient) markReceived() {
+	c.lastReceived.Store(time.Now().UnixNano())
+	c.receiveStalled.Store(false)
+}
+
+// idleSinceReceive returns how long the receive stream has been silent.
+func (c *GrpcClient) idleSinceReceive() time.Duration {
+	return time.Since(time.Unix(0, c.lastReceived.Load()))
+}
+
+// watchReceiveStream guards against a receive stream that is transport-alive but
+// no longer delivering messages. While the stream is idle past
+// receiveInactivityThreshold it sends a self-addressed probe that the Signal
+// server routes back to this client. If the probe does not round-trip within
+// receiveProbeTimeout the receive direction is considered dead and cancelStream
+// is called so the retry loop reconnects.
+func (c *GrpcClient) watchReceiveStream(ctx context.Context, cancelStream context.CancelFunc) {
+	ticker := time.NewTicker(receiveWatchdogInterval)
+	defer ticker.Stop()
+
+	var probeSentAt time.Time
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+			if c.idleSinceReceive() < receiveInactivityThreshold {
+				probeSentAt = time.Time{}
+				continue
+			}
+
+			if !probeSentAt.IsZero() && time.Since(probeSentAt) >= receiveProbeTimeout {
+				log.Warnf("signal receive stream stalled: no messages for %s and probe did not return, reconnecting", c.idleSinceReceive().Round(time.Second))
+				c.receiveStalled.Store(true)
+				c.notifyDisconnected(errReceiveStreamStalled)
+				cancelStream()
+				return
+			}
+
+			if probeSentAt.IsZero() {
+				if err := c.sendReceiveProbe(); err != nil {
+					log.Debugf("failed to send signal receive probe: %v", err)
+				}
+				probeSentAt = time.Now()
+			}
+		}
+	}
+}
+
+// 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 {
+	self := c.key.PublicKey().String()
+	return c.Send(&proto.Message{
+		Key:       self,
+		RemoteKey: self,
+		Body:      &proto.Body{Type: proto.Body_HEARTBEAT},
+	})
+}
+
 // receive receives messages from other peers coming through the Signal Exchange
 // and distributes them to worker threads for processing
 func (c *GrpcClient) receive(stream proto.SignalExchange_ConnectStreamClient) error {
@@ -419,6 +509,9 @@ func (c *GrpcClient) receive(stream proto.SignalExchange_ConnectStreamClient) er
 			return err
 		}
 
+		// Any frame from the server proves the receive direction is alive.
+		c.markReceived()
+
 		if msg == nil {
 			continue
 		}