github.com/johnnyeven/libtools@v0.0.0-20191126065708-61829c1adf46/task/kafka/producer.go (about)

     1  package kafka
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"github.com/johnnyeven/libtools/task/constants"
     7  	"github.com/segmentio/kafka-go"
     8  	"github.com/sirupsen/logrus"
     9  )
    10  
    11  type KafkaProducer struct {
    12  	w *kafka.Writer
    13  }
    14  
    15  func NewKafkaProducer(info constants.ConnectionInfo) *KafkaProducer {
    16  	w := kafka.NewWriter(kafka.WriterConfig{
    17  		Brokers:           []string{fmt.Sprintf("%s:%d", info.Host, info.Port)},
    18  		Topic:             "test",
    19  		Balancer:          &kafka.LeastBytes{},
    20  	})
    21  	return &KafkaProducer{
    22  		w: w,
    23  	}
    24  }
    25  
    26  func (p *KafkaProducer) SendTask(task *constants.Task) error {
    27  	data, err := constants.MarshalData(task)
    28  	if err != nil {
    29  		logrus.Errorf("GearmanProducer.SendTask err: %v", err)
    30  		return err
    31  	}
    32  	err = p.w.WriteMessages(context.Background(), kafka.Message{
    33  		Topic: task.Channel,
    34  		Key:   []byte(task.Subject),
    35  		Value: data,
    36  	})
    37  	if err != nil {
    38  		return err
    39  	}
    40  
    41  	return nil
    42  }
    43  
    44  func (p *KafkaProducer) Stop() {
    45  	p.w.Close()
    46  }