You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
106 lines
2.0 KiB
106 lines
2.0 KiB
package mq
|
|
|
|
import (
|
|
"sync"
|
|
|
|
"github.com/IBM/sarama"
|
|
)
|
|
|
|
type KafkaConfig struct {
|
|
topic string
|
|
addr []string
|
|
config *sarama.Config
|
|
}
|
|
|
|
type KafkaProducer struct {
|
|
config *KafkaConfig
|
|
client sarama.SyncProducer
|
|
}
|
|
|
|
type KafkaConsumer struct {
|
|
config *KafkaConfig
|
|
client sarama.Consumer
|
|
}
|
|
|
|
type ProducerResp struct {
|
|
Partition int32
|
|
Offset int64
|
|
}
|
|
|
|
func NewKafkaConfig(topic string, config *sarama.Config) *KafkaConfig {
|
|
return &KafkaConfig{
|
|
topic: topic,
|
|
config: config,
|
|
}
|
|
}
|
|
|
|
func NewKafkaProducer(config *KafkaConfig) (*KafkaProducer, func(), error) {
|
|
client, err := sarama.NewSyncProducer(config.addr, config.config)
|
|
shutdown := func() {
|
|
client.Close() //nolint
|
|
}
|
|
if err != nil {
|
|
return nil, shutdown, err
|
|
}
|
|
|
|
return &KafkaProducer{
|
|
config: config,
|
|
client: client,
|
|
}, shutdown, nil
|
|
}
|
|
|
|
func NewKafkaConsumer(config *KafkaConfig) (*KafkaConsumer, func(), error) {
|
|
consumer, err := sarama.NewConsumer(config.addr, nil)
|
|
shutdown := func() {
|
|
consumer.Close() //nolint
|
|
}
|
|
if err != nil {
|
|
return nil, shutdown, err
|
|
}
|
|
return &KafkaConsumer{
|
|
config: config,
|
|
client: consumer,
|
|
}, shutdown, nil
|
|
}
|
|
|
|
func (producer *KafkaProducer) Produce(data []byte) (*ProducerResp, error) {
|
|
partition, offset, err := producer.client.SendMessage(&sarama.ProducerMessage{
|
|
Topic: producer.config.topic,
|
|
Value: sarama.ByteEncoder(data),
|
|
})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return &ProducerResp{
|
|
Partition: partition,
|
|
Offset: offset,
|
|
}, nil
|
|
}
|
|
|
|
func (consumer *KafkaConsumer) Consume(op func(b []byte)) error {
|
|
var wg sync.WaitGroup
|
|
|
|
partitionList, err := consumer.client.Partitions(consumer.config.topic)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
for partition := range partitionList {
|
|
pc, err := consumer.client.ConsumePartition(consumer.config.topic, int32(partition), sarama.OffsetNewest)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
wg.Add(1)
|
|
|
|
go func(sarama.PartitionConsumer) {
|
|
defer pc.AsyncClose()
|
|
for message := range pc.Messages() {
|
|
op(message.Value)
|
|
}
|
|
}(pc)
|
|
}
|
|
|
|
wg.Wait()
|
|
return nil
|
|
}
|
|
|