Merge pull request #2 from zeromicro/master

merge

Author: lvphps

kq/config.go

@@ -12,6 +12,7 @@ type KqConf struct {
 	Brokers     []string
 	Group       string
 	Topic       string
+	CaFile      string `json:",optional"`
 	Offset      string `json:",options=first|last,default=last"`
 	Conns       int    `json:",default=1"`
 	Consumers   int    `json:",default=8"`

kq/pusher.go

@@ -11,92 +11,111 @@ import (
 )
 
 type (
-	PushOption func(options *chunkOptions)
+	PushOption func(options *pushOptions)
 
 	Pusher struct {
-		produer  *kafka.Writer
+		producer *kafka.Writer
 		topic    string
 		executor *executors.ChunkExecutor
 	}
 
-	chunkOptions struct {
+	pushOptions struct {
+		// kafka.Writer options
+		allowAutoTopicCreation bool
+
+		// executors.ChunkExecutor options
 		chunkSize     int
 		flushInterval time.Duration
 	}
 )
 
+// NewPusher returns a Pusher with the given Kafka addresses and topic.
 func NewPusher(addrs []string, topic string, opts ...PushOption) *Pusher {
 	producer := &kafka.Writer{
 		Addr:        kafka.TCP(addrs...),
 		Topic:       topic,
 		Balancer:    &kafka.LeastBytes{},
 		Compression: kafka.Snappy,
 	}
+
+	var options pushOptions
+	for _, opt := range opts {
+		opt(&options)
+	}
+
+	// apply kafka.Writer options
+	producer.AllowAutoTopicCreation = options.allowAutoTopicCreation
+
+	// apply ChunkExecutor options
+	var chunkOpts []executors.ChunkOption
+	if options.chunkSize > 0 {
+		chunkOpts = append(chunkOpts, executors.WithChunkBytes(options.chunkSize))
+	}
+	if options.flushInterval > 0 {
+		chunkOpts = append(chunkOpts, executors.WithFlushInterval(options.flushInterval))
+	}
+
 	pusher := &Pusher{
-		produer: producer,
-		topic:   topic,
+		producer: producer,
+		topic:    topic,
 	}
 	pusher.executor = executors.NewChunkExecutor(func(tasks []interface{}) {
 		chunk := make([]kafka.Message, len(tasks))
 		for i := range tasks {
 			chunk[i] = tasks[i].(kafka.Message)
 		}
-		if err := pusher.produer.WriteMessages(context.Background(), chunk...); err != nil {
+		if err := pusher.producer.WriteMessages(context.Background(), chunk...); err != nil {
 			logx.Error(err)
 		}
-	}, newOptions(opts)...)
+	}, chunkOpts...)
 
 	return pusher
 }
 
+// Close closes the Pusher and releases any resources used by it.
 func (p *Pusher) Close() error {
 	if p.executor != nil {
 		p.executor.Flush()
 	}
-	
-	return p.produer.Close()
+
+	return p.producer.Close()
 }
 
+// Name returns the name of the Kafka topic that the Pusher is sending messages to.
 func (p *Pusher) Name() string {
 	return p.topic
 }
 
+// Push sends a message to the Kafka topic.
 func (p *Pusher) Push(v string) error {
 	msg := kafka.Message{
-		Key:   []byte(strconv.FormatInt(time.Now().UnixNano(), 10)),
+		Key:   []byte(strconv.FormatInt(time.Now().UnixNano(), 10)), // current timestamp
 		Value: []byte(v),
 	}
 	if p.executor != nil {
 		return p.executor.Add(msg, len(v))
 	} else {
-		return p.produer.WriteMessages(context.Background(), msg)
+		return p.producer.WriteMessages(context.Background(), msg)
 	}
 }
 
+// WithChunkSize customizes the Pusher with the given chunk size.
 func WithChunkSize(chunkSize int) PushOption {
-	return func(options *chunkOptions) {
+	return func(options *pushOptions) {
 		options.chunkSize = chunkSize
 	}
 }
 
+// WithFlushInterval customizes the Pusher with the given flush interval.
 func WithFlushInterval(interval time.Duration) PushOption {
-	return func(options *chunkOptions) {
+	return func(options *pushOptions) {
 		options.flushInterval = interval
 	}
 }
 
-func newOptions(opts []PushOption) []executors.ChunkOption {
-	var options chunkOptions
-	for _, opt := range opts {
-		opt(&options)
-	}
-
-	var chunkOpts []executors.ChunkOption
-	if options.chunkSize > 0 {
-		chunkOpts = append(chunkOpts, executors.WithChunkBytes(options.chunkSize))
-	}
-	if options.flushInterval > 0 {
-		chunkOpts = append(chunkOpts, executors.WithFlushInterval(options.flushInterval))
+// WithAllowAutoTopicCreation allows the Pusher to create the given topic if it does not exist.
+func WithAllowAutoTopicCreation() PushOption {
+	return func(options *pushOptions) {
+		options.allowAutoTopicCreation = true
 	}
-	return chunkOpts
 }

kq/queue.go

@@ -2,8 +2,11 @@ package kq
 
 import (
 	"context"
+	"crypto/tls"
+	"crypto/x509"
 	"io"
 	"log"
+	"os"
 	"time"
 
 	"github.com/segmentio/kafka-go"
@@ -121,6 +124,23 @@ func newKafkaQueue(c KqConf, handler ConsumeHandler, options queueOptions) queue
 			},
 		}
 	}
+	if len(c.CaFile) > 0 {
+		caCert, err := os.ReadFile(c.CaFile)
+		if err != nil {
+			log.Fatal(err)
+		}
+
+		caCertPool := x509.NewCertPool()
+		ok := caCertPool.AppendCertsFromPEM(caCert)
+		if !ok {
+			log.Fatal(err)
+		}
+
+		readerConfig.Dialer.TLS = &tls.Config{
+			RootCAs:            caCertPool,
+			InsecureSkipVerify: true,
+		}
+	}
 	consumer := kafka.NewReader(readerConfig)
 
 	return &kafkaQueue{