Merge pull request #3 from boreq/closed-mutex

Fix a data race in the close method
czeslavo · Feb 5, 2020 · 3059fcb · 3059fcb
2 parents a50d137 + 8436e6a
commit 3059fcb
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/pkg/firestore/subscriber.go b/pkg/firestore/subscriber.go
@@ -77,6 +77,8 @@ type Subscriber struct {
 	closed                       bool
 	closing                      chan struct{}
 	allSubscriptionsWaitingGroup sync.WaitGroup
+
+	closedMutex sync.Mutex
 }
 
 func NewSubscriber(config SubscriberConfig, logger watermill.LoggerAdapter) (*Subscriber, error) {
@@ -97,10 +99,15 @@ func NewSubscriber(config SubscriberConfig, logger watermill.LoggerAdapter) (*Su
 }
 
 func (s *Subscriber) Subscribe(ctx context.Context, topic string) (<-chan *message.Message, error) {
+	s.closedMutex.Lock()
+
 	if s.closed {
+		s.closedMutex.Unlock()
 		return nil, errors.New("subscriber is closed")
 	}
 
+	s.closedMutex.Unlock()
+
 	ctx, cancel := context.WithCancel(ctx)
 
 	subscriptionName := s.config.GenerateSubscriptionName(topic)
@@ -155,6 +162,9 @@ func (s *Subscriber) subscriptionCollection(topic string) *firestore.CollectionR
 }
 
 func (s *Subscriber) Close() error {
+	s.closedMutex.Lock()
+	defer s.closedMutex.Unlock()
+
 	if s.closed {
 		return nil
 	}