github.com/kaydxh/golang@v0.0.131/pkg/mq/kafka/kafka_producer.go (about)

     1  package kafka
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  
     7  	kafka "github.com/segmentio/kafka-go"
     8  )
     9  
    10  type Producer struct {
    11  	*kafka.Writer
    12  	config    kafka.WriterConfig
    13  	topic     string
    14  	closeOnce sync.Once
    15  	closeCh   chan struct{}
    16  }
    17  
    18  func NewProducer(config kafka.WriterConfig) (*Producer, error) {
    19  	p := &Producer{
    20  		config:  config,
    21  		closeCh: make(chan struct{}),
    22  	}
    23  	w := kafka.NewWriter(config)
    24  	p.Writer = w
    25  	return p, nil
    26  }
    27  
    28  func (p *Producer) Send(ctx context.Context, msgs ...kafka.Message) error {
    29  	return p.Writer.WriteMessages(ctx, msgs...)
    30  }
    31  
    32  func (p *Producer) Close() {
    33  	p.closeOnce.Do(func() {
    34  		if p.closeCh != nil {
    35  			close(p.closeCh)
    36  		}
    37  	})
    38  }