Merge remote-tracking branch 'origin/main' into dmitri-event-aggregation

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

shared/signal/client/client.go

@@ -33,7 +33,7 @@ type Client interface {
 	Receive(ctx context.Context, msgHandler func(msg *proto.Message) error) error
 	Ready() bool
 	IsHealthy() bool
-	WaitStreamConnected()
+	WaitStreamConnected(context.Context)
 	SendToStream(msg *proto.EncryptedMessage) error
 	Send(msg *proto.Message) error
 	SetOnReconnectedListener(func())

shared/signal/client/client_test.go

@@ -65,7 +65,10 @@ var _ = Describe("GrpcClient", func() {
 						return
 					}
 				}()
-				clientA.WaitStreamConnected()
+				ctxA, cancelA := context.WithTimeout(context.Background(), 5*time.Second)
+				defer cancelA()
+				clientA.WaitStreamConnected(ctxA)
+				Expect(clientA.StreamConnected()).To(BeTrue())
 
 				// connect PeerB to Signal
 				keyB, _ := wgtypes.GenerateKey()
@@ -91,7 +94,10 @@ var _ = Describe("GrpcClient", func() {
 					}
 				}()
 
-				clientB.WaitStreamConnected()
+				ctxB, cancelB := context.WithTimeout(context.Background(), 5*time.Second)
+				defer cancelB()
+				clientB.WaitStreamConnected(ctxB)
+				Expect(clientB.StreamConnected()).To(BeTrue())
 
 				// PeerA initiates ping-pong
 				err := clientA.Send(&sigProto.Message{
@@ -129,8 +135,10 @@ var _ = Describe("GrpcClient", func() {
 						return
 					}
 				}()
-				client.WaitStreamConnected()
-				Expect(client).NotTo(BeNil())
+				ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+				defer cancel()
+				client.WaitStreamConnected(ctx)
+				Expect(client.StreamConnected()).To(BeTrue())
 			})
 		})
 

shared/signal/client/grpc.go

@@ -246,15 +246,6 @@ func (c *GrpcClient) notifyStreamConnected() {
 	}
 }
 
-func (c *GrpcClient) getStreamStatusChan() <-chan struct{} {
-	c.mux.Lock()
-	defer c.mux.Unlock()
-	if c.connectedCh == nil {
-		c.connectedCh = make(chan struct{})
-	}
-	return c.connectedCh
-}
-
 func (c *GrpcClient) connect(ctx context.Context, key string) (proto.SignalExchange_ConnectStreamClient, error) {
 	c.stream = nil
 
@@ -310,14 +301,24 @@ func (c *GrpcClient) IsHealthy() bool {
 }
 
 // WaitStreamConnected waits until the client is connected to the Signal stream
-func (c *GrpcClient) WaitStreamConnected() {
-
+func (c *GrpcClient) WaitStreamConnected(ctx context.Context) {
+	// Check the status and obtain the wait channel atomically: otherwise
+	// notifyStreamConnected could flip the status and close/clear the channel
+	// between the check and the channel creation, leaving us waiting forever on
+	// a stale channel.
+	c.mux.Lock()
 	if c.status == StreamConnected {
+		c.mux.Unlock()
 		return
 	}
+	if c.connectedCh == nil {
+		c.connectedCh = make(chan struct{})
+	}
+	ch := c.connectedCh
+	c.mux.Unlock()
 
-	ch := c.getStreamStatusChan()
 	select {
+	case <-ctx.Done():
 	case <-c.ctx.Done():
 	case <-ch:
 	}

shared/signal/client/mock.go

@@ -55,7 +55,7 @@ func (sm *MockClient) Ready() bool {
 	return sm.ReadyFunc()
 }
 
-func (sm *MockClient) WaitStreamConnected() {
+func (sm *MockClient) WaitStreamConnected(context.Context) {
 	if sm.WaitStreamConnectedFunc == nil {
 		return
 	}

shared/signal/client/watchdog_test.go

@@ -65,7 +65,7 @@ func TestReceiveProbeRoundTrips(t *testing.T) {
 
 	streamReady := make(chan struct{})
 	go func() {
-		client.WaitStreamConnected()
+		client.WaitStreamConnected(ctx)
 		close(streamReady)
 	}()
 	select {