go-kratos/pkg/kafka/producer.go

44 lines
843 B
Go

package kafka
import (
"context"
"fmt"
"github.com/segmentio/kafka-go"
)
type KafkaProducer struct {
writer *kafka.Writer
}
func NewKafkaWriter(address []string, topic string) *KafkaProducer {
return &KafkaProducer{
writer: &kafka.Writer{
Addr: kafka.TCP(address...),
Topic: topic,
Balancer: &kafka.LeastBytes{},
RequiredAcks: kafka.RequireOne,
Async: false,
BatchSize: 1,
BatchTimeout: 0,
},
}
}
func (p *KafkaProducer) SendToKafkaMessage(ctx context.Context, key string, message []byte) error {
msg := kafka.Message{
Key: []byte(key),
Value: []byte(message),
}
err := p.writer.WriteMessages(ctx, msg)
if err != nil {
fmt.Printf("send kafka messages error: %v", err)
return err
}
return nil
}
func (p *KafkaProducer) Close() error {
return p.writer.Close()
}