[signal] Serialize concurrent sends to a peer signal stream (#6463)

signal/peer/peer.go

@@ -26,6 +26,10 @@ type Peer struct {
 
 	// a gRpc connection stream to the Peer
 	Stream proto.SignalExchange_ConnectStreamServer
+	// sendMu serializes writes to Stream. gRPC forbids concurrent SendMsg on
+	// the same ServerStream, and a peer can be the target of many senders at
+	// once.
+	sendMu sync.Mutex
 
 	// registration time
 	RegisteredAt time.Time
@@ -33,6 +37,13 @@ type Peer struct {
 	Cancel context.CancelFunc
 }
 
+// Send writes a message to the peer's stream, serializing concurrent senders.
+func (p *Peer) Send(msg *proto.EncryptedMessage) error {
+	p.sendMu.Lock()
+	defer p.sendMu.Unlock()
+	return p.Stream.Send(msg)
+}
+
 // NewPeer creates a new instance of a connected Peer
 func NewPeer(id string, stream proto.SignalExchange_ConnectStreamServer, cancel context.CancelFunc) *Peer {
 	return &Peer{

signal/server/concurrent_send_test.go

@@ -0,0 +1,67 @@
+package server
+
+import (
+	"context"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	"go.opentelemetry.io/otel"
+
+	"github.com/netbirdio/netbird/shared/signal/proto"
+	"github.com/netbirdio/netbird/signal/peer"
+)
+
+// concurrencyCheckStream records the maximum number of Send calls in flight at
+// once. gRPC forbids concurrent SendMsg on the same ServerStream, so a correct
+// server must never have more than one in flight per peer.
+type concurrencyCheckStream struct {
+	proto.SignalExchange_ConnectStreamServer
+	ctx      context.Context
+	inflight atomic.Int32
+	maxSeen  atomic.Int32
+}
+
+func (s *concurrencyCheckStream) Send(*proto.EncryptedMessage) error {
+	n := s.inflight.Add(1)
+	for {
+		old := s.maxSeen.Load()
+		if n <= old || s.maxSeen.CompareAndSwap(old, n) {
+			break
+		}
+	}
+	// Widen the window so overlapping callers are reliably observed.
+	time.Sleep(time.Millisecond)
+	s.inflight.Add(-1)
+	return nil
+}
+
+func (s *concurrencyCheckStream) Context() context.Context { return s.ctx }
+
+// TestForwardMessageToPeerSerializesSend verifies that concurrent forwards to the
+// same peer never call Stream.Send concurrently, which would violate the gRPC
+// ServerStream contract.
+func TestForwardMessageToPeerSerializesSend(t *testing.T) {
+	s, err := NewServer(context.Background(), otel.Meter(""))
+	require.NoError(t, err)
+
+	const peerID = "peerX"
+	stream := &concurrencyCheckStream{ctx: context.Background()}
+	_, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+	require.NoError(t, s.registry.Register(peer.NewPeer(peerID, stream, cancel)))
+
+	var wg sync.WaitGroup
+	for i := 0; i < 50; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			s.forwardMessageToPeer(context.Background(), &proto.EncryptedMessage{Key: "sender", RemoteKey: peerID})
+		}()
+	}
+	wg.Wait()
+
+	require.Equal(t, int32(1), stream.maxSeen.Load(), "Stream.Send must never run concurrently on the same peer stream")
+}

signal/server/signal.go

@@ -179,7 +179,7 @@ func (s *Server) forwardMessageToPeer(ctx context.Context, msg *proto.EncryptedM
 	sendResultChan := make(chan error, 1)
 	go func() {
 		select {
-		case sendResultChan <- dstPeer.Stream.Send(msg):
+		case sendResultChan <- dstPeer.Send(msg):
 			return
 		case <-dstPeer.Stream.Context().Done():
 			return