diff --git a/signal/client/grpc.go b/signal/client/grpc.go index 28df8f04f..0961a44ec 100644 --- a/signal/client/grpc.go +++ b/signal/client/grpc.go @@ -421,7 +421,7 @@ func (c *GrpcClient) receive(stream proto.SignalExchange_ConnectStreamClient) er continue } - if err := c.decryptionWorker.AddMsg(msg); err != nil { + if err := c.decryptionWorker.AddMsg(c.ctx, msg); err != nil { log.Errorf("failed to add message to decryption worker: %v", err) } } diff --git a/signal/client/worker.go b/signal/client/worker.go index 1a8656b1e..f9ec6b167 100644 --- a/signal/client/worker.go +++ b/signal/client/worker.go @@ -2,7 +2,6 @@ package client import ( "context" - "fmt" log "github.com/sirupsen/logrus" @@ -24,13 +23,12 @@ func NewWorker(decryptFn func(msg *proto.EncryptedMessage) (*proto.Message, erro } } -func (w *Worker) AddMsg(msg *proto.EncryptedMessage) error { - if w.encryptedMsgPool == nil { - return fmt.Errorf("worker is not initialized") - } - +func (w *Worker) AddMsg(ctx context.Context, msg *proto.EncryptedMessage) error { // this is blocker because do not want to drop messages here - w.encryptedMsgPool <- msg + select { + case w.encryptedMsgPool <- msg: + case <-ctx.Done(): + } return nil }