github.com/aaabigfish/gopkg@v1.1.0/mq/kafka/writer.go (about)

     1  package kafka
     2  
     3  import (
     4  	"context"
     5  	"time"
     6  
     7  	"github.com/rs/xid"
     8  	kf "github.com/segmentio/kafka-go"
     9  	"github.com/segmentio/kafka-go/snappy"
    10  )
    11  
    12  type Writer interface {
    13  	WriteMessage(tx context.Context, msg Message) error
    14  	WriteMessages(tx context.Context, msgs []Message) error
    15  	GetWriter() *kf.Writer
    16  	Close() error
    17  }
    18  type writer struct {
    19  	kfWriter *kf.Writer
    20  	c        *kf.WriterConfig
    21  }
    22  
    23  func NewWriter(brokers []string, topic ...string) Writer {
    24  	dialer := &kf.Dialer{
    25  		Timeout:  10 * time.Second,
    26  		ClientID: xid.New().String(),
    27  	}
    28  
    29  	conf := kf.WriterConfig{
    30  		Brokers:          brokers,
    31  		Async:            true,
    32  		Balancer:         &kf.Hash{},
    33  		Dialer:           dialer,
    34  		WriteTimeout:     10 * time.Second,
    35  		ReadTimeout:      10 * time.Second,
    36  		CompressionCodec: snappy.NewCompressionCodec(),
    37  	}
    38  
    39  	if len(topic) > 0 {
    40  		conf.Topic = topic[0]
    41  	}
    42  
    43  	return &writer{kfWriter: kf.NewWriter(conf), c: &conf}
    44  }
    45  
    46  func (w *writer) GetWriter() *kf.Writer {
    47  	return w.kfWriter
    48  }
    49  
    50  func (w *writer) WriteMessage(tx context.Context, msg Message) error {
    51  	return w.kfWriter.WriteMessages(tx, msg)
    52  }
    53  
    54  func (w *writer) WriteMessages(tx context.Context, msgs []Message) error {
    55  	return w.kfWriter.WriteMessages(tx, msgs...)
    56  }
    57  
    58  func (w *writer) Close() error {
    59  	return w.kfWriter.Close()
    60  }