Merge pull request #55 from testerxiaodong/master

feat(kq): support custom producer balancer and message key

Author: kevwan

kq/pusher.go

@@ -24,6 +24,7 @@ type (
 	pushOptions struct {
 		// kafka.Writer options
 		allowAutoTopicCreation bool
+		balancer               kafka.Balancer
 
 		// executors.ChunkExecutor options
 		chunkSize     int
@@ -50,6 +51,9 @@ func NewPusher(addrs []string, topic string, opts ...PushOption) *Pusher {
 
 	// apply kafka.Writer options
 	producer.AllowAutoTopicCreation = options.allowAutoTopicCreation
+	if options.balancer != nil {
+		producer.Balancer = options.balancer
+	}
 
 	pusher := &Pusher{
 		producer: producer,
@@ -134,6 +138,26 @@ func (p *Pusher) PushWithKey(ctx context.Context, key, v string) error {
 	}
 }
 
+// SetWriterBalancer set kafka-go custom writer balancer.
+func (p *Pusher) SetWriterBalancer(balancer kafka.Balancer) {
+	if p.producer != nil {
+		p.producer.Balancer = balancer
+	}
+}
+
+// PushWithKey sends a message to the Kafka topic with custom message key.
+func (p *Pusher) PushWithKey(k, v string) error {
+	msg := kafka.Message{
+		Key:   []byte(k), // custom message key
+		Value: []byte(v),
+	}
+	if p.executor != nil {
+		return p.executor.Add(msg, len(v))
+	} else {
+		return p.producer.WriteMessages(context.Background(), msg)
+	}
+}
+
 // WithChunkSize customizes the Pusher with the given chunk size.
 func WithChunkSize(chunkSize int) PushOption {
 	return func(options *pushOptions) {
@@ -155,6 +179,13 @@ func WithAllowAutoTopicCreation() PushOption {
 	}
 }
 
+// WithBalancer customizes the Pusher with the given balancer.
+func WithBalancer(balancer kafka.Balancer) PushOption {
+	return func(options *pushOptions) {
+		options.balancer = balancer
+  }
+}
+    
 // WithSyncPush enables the Pusher to push messages synchronously.
 func WithSyncPush() PushOption {
 	return func(options *pushOptions) {