diff --git a/client/internal/engine.go b/client/internal/engine.go index 11dccc036..68b80f7ab 100644 --- a/client/internal/engine.go +++ b/client/internal/engine.go @@ -1776,7 +1776,8 @@ func (e *Engine) receiveSignalEvents() { } }() - e.signal.WaitStreamConnected() + // todo: consider to remove this blocker. I do not see benefit to block the Start operations + e.signal.WaitStreamConnected(e.ctx) } func (e *Engine) parseNATExternalIPMappings() []string { diff --git a/shared/signal/client/client.go b/shared/signal/client/client.go index 9dc6ccd37..fb77cb90f 100644 --- a/shared/signal/client/client.go +++ b/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()) diff --git a/shared/signal/client/client_test.go b/shared/signal/client/client_test.go index 1af34e37a..b4ea2bee1 100644 --- a/shared/signal/client/client_test.go +++ b/shared/signal/client/client_test.go @@ -65,7 +65,7 @@ var _ = Describe("GrpcClient", func() { return } }() - clientA.WaitStreamConnected() + clientA.WaitStreamConnected(context.Background()) // connect PeerB to Signal keyB, _ := wgtypes.GenerateKey() @@ -91,7 +91,7 @@ var _ = Describe("GrpcClient", func() { } }() - clientB.WaitStreamConnected() + clientB.WaitStreamConnected(context.Background()) // PeerA initiates ping-pong err := clientA.Send(&sigProto.Message{ @@ -129,7 +129,7 @@ var _ = Describe("GrpcClient", func() { return } }() - client.WaitStreamConnected() + client.WaitStreamConnected(context.Background()) Expect(client).NotTo(BeNil()) }) }) diff --git a/shared/signal/client/grpc.go b/shared/signal/client/grpc.go index b245b2296..db23ed75b 100644 --- a/shared/signal/client/grpc.go +++ b/shared/signal/client/grpc.go @@ -282,7 +282,7 @@ 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) { if c.status == StreamConnected { return @@ -290,6 +290,7 @@ func (c *GrpcClient) WaitStreamConnected() { ch := c.getStreamStatusChan() select { + case <-ctx.Done(): case <-c.ctx.Done(): case <-ch: } diff --git a/shared/signal/client/mock.go b/shared/signal/client/mock.go index 95381a5b0..0c8a083c5 100644 --- a/shared/signal/client/mock.go +++ b/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 }