Browse Source

Fix data race in consumer mock

Willem van Bergen 10 years ago
parent
commit
c85d3d0d49
2 changed files with 21 additions and 10 deletions
  1. 18 8
      mocks/consumer.go
  2. 3 2
      mocks/mocks.go

+ 18 - 8
mocks/consumer.go

@@ -97,30 +97,40 @@ func (pc *PartitionConsumer) AsyncClose() {
 func (pc *PartitionConsumer) Close() error {
 	if !pc.consumed {
 		pc.t.Errorf("Expectations set on %s/%d, but no partition consumer was started.", pc.topic, pc.partition)
+		return errPartitionConsumerNotStarted
 	}
 
 	pc.AsyncClose()
 
-	var errs = make(sarama.ConsumerErrors, 0)
+	var (
+		closeErr error
+		wg       sync.WaitGroup
+	)
+
+	wg.Add(1)
 	go func() {
+		defer wg.Done()
+
+		var errs = make(sarama.ConsumerErrors, 0)
 		for err := range pc.errors {
 			errs = append(errs, err)
 		}
+
+		if len(errs) > 0 {
+			closeErr = errs
+		}
 	}()
 
+	wg.Add(1)
 	go func() {
+		defer wg.Done()
 		for _ = range pc.messages {
 			// drain
 		}
 	}()
 
-	pc.l.Lock()
-	pc.l.Unlock()
-
-	if len(errs) > 0 {
-		return errs
-	}
-	return nil
+	wg.Wait()
+	return closeErr
 }
 
 // Errors implements the Errors method from the sarama.PartitionConsumer interface.

+ 3 - 2
mocks/mocks.go

@@ -26,8 +26,9 @@ type ErrorReporter interface {
 }
 
 var (
-	errProduceSuccess    error = nil
-	errOutOfExpectations       = errors.New("No more expectations set on mock")
+	errProduceSuccess              error = nil
+	errOutOfExpectations                 = errors.New("No more expectations set on mock")
+	errPartitionConsumerNotStarted       = errors.New("The partition consumer was never started")
 )
 
 type producerExpectation struct {