sink: refine close in kafka sink (#586)

pingcap · May 20, 2020 · 13ccbe6 · 13ccbe6
1 parent 557a397
commit 13ccbe6
Showing 1 changed file with 19 additions and 4 deletions.
diff --git a/cdc/sink/mqProducer/kafka.go b/cdc/sink/mqProducer/kafka.go
@@ -90,7 +90,12 @@ func (k *kafkaSaramaProducer) SyncBroadcastMessage(ctx context.Context, key []by
 			Partition: int32(i),
 		}
 	}
-	return errors.Trace(k.syncClient.SendMessages(msgs))
+	select {
+	case <-k.closeCh:
+		return nil
+	default:
+		return errors.Trace(k.syncClient.SendMessages(msgs))
+	}
 }
 
 func (k *kafkaSaramaProducer) Flush(ctx context.Context) error {
@@ -134,7 +139,12 @@ func (k *kafkaSaramaProducer) GetPartitionNum() int32 {
 }
 
 func (k *kafkaSaramaProducer) Close() error {
-	close(k.closeCh)
+	select {
+	case <-k.closeCh:
+		return nil
+	default:
+		close(k.closeCh)
+	}
 	err1 := k.syncClient.Close()
 	err2 := k.asyncClient.Close()
 	if err1 != nil {
@@ -147,7 +157,13 @@ func (k *kafkaSaramaProducer) Close() error {
 }
 
 func (k *kafkaSaramaProducer) run(ctx context.Context) error {
-	defer k.flushedReceiver.Stop()
+	defer func() {
+		k.flushedReceiver.Stop()
+		err := k.Close()
+		if err != nil {
+			log.Error("close kafkaSaramaProducer with error", zap.Error(err))
+		}
+	}()
 	for {
 		select {
 		case <-ctx.Done():
@@ -162,7 +178,6 @@ func (k *kafkaSaramaProducer) run(ctx context.Context) error {
 			atomic.StoreUint64(&k.partitionOffset[msg.Partition].flushed, flushedOffset)
 			k.flushedNotifier.Notify()
 		case err := <-k.asyncClient.Errors():
-			close(k.closeCh)
 			return errors.Annotate(err, "write kafka error")
 		}
 	}