1 Star 0 Fork 15

曾华萍/goreplay

加入 Gitee
与超过 1200万 开发者一起发现、参与优秀开源项目,私有仓库也完全免费 :)
免费加入
文件
克隆/下载
input_kafka.go 2.71 KB
一键复制 编辑 原始数据 按行查看 历史
Erik Schweller 提交于 2021-06-08 14:05 . Address issues flagged by (#938)
package main
import (
"encoding/json"
"log"
"strings"
"github.com/Shopify/sarama"
"github.com/Shopify/sarama/mocks"
)
// KafkaInput is used for receiving Kafka messages and
// transforming them into HTTP payloads.
type KafkaInput struct {
config *InputKafkaConfig
consumers []sarama.PartitionConsumer
messages chan *sarama.ConsumerMessage
quit chan struct{}
}
// NewKafkaInput creates instance of kafka consumer client with TLS config
func NewKafkaInput(address string, config *InputKafkaConfig, tlsConfig *KafkaTLSConfig) *KafkaInput {
c := NewKafkaConfig(tlsConfig)
var con sarama.Consumer
if mock, ok := config.consumer.(*mocks.Consumer); ok && mock != nil {
con = config.consumer
} else {
var err error
con, err = sarama.NewConsumer(strings.Split(config.Host, ","), c)
if err != nil {
log.Fatalln("Failed to start Sarama(Kafka) consumer:", err)
}
}
partitions, err := con.Partitions(config.Topic)
if err != nil {
log.Fatalln("Failed to collect Sarama(Kafka) partitions:", err)
}
i := &KafkaInput{
config: config,
consumers: make([]sarama.PartitionConsumer, len(partitions)),
messages: make(chan *sarama.ConsumerMessage, 256),
quit: make(chan struct{}),
}
for index, partition := range partitions {
consumer, err := con.ConsumePartition(config.Topic, partition, sarama.OffsetNewest)
if err != nil {
log.Fatalln("Failed to start Sarama(Kafka) partition consumer:", err)
}
go func(consumer sarama.PartitionConsumer) {
defer consumer.Close()
for message := range consumer.Messages() {
i.messages <- message
}
}(consumer)
go i.ErrorHandler(consumer)
i.consumers[index] = consumer
}
return i
}
// ErrorHandler should receive errors
func (i *KafkaInput) ErrorHandler(consumer sarama.PartitionConsumer) {
for err := range consumer.Errors() {
Debug(1, "Failed to read access log entry:", err)
}
}
// PluginRead a reads message from this plugin
func (i *KafkaInput) PluginRead() (*Message, error) {
var message *sarama.ConsumerMessage
var msg Message
select {
case <-i.quit:
return nil, ErrorStopped
case message = <-i.messages:
}
msg.Data = message.Value
if i.config.UseJSON {
var kafkaMessage KafkaMessage
json.Unmarshal(message.Value, &kafkaMessage)
var err error
msg.Data, err = kafkaMessage.Dump()
if err != nil {
Debug(1, "[INPUT-KAFKA] failed to decode access log entry:", err)
return nil, err
}
}
// does it have meta
if isOriginPayload(msg.Data) {
msg.Meta, msg.Data = payloadMetaWithBody(msg.Data)
}
return &msg, nil
}
func (i *KafkaInput) String() string {
return "Kafka Input: " + i.config.Host + "/" + i.config.Topic
}
// Close closes this plugin
func (i *KafkaInput) Close() error {
close(i.quit)
return nil
}
Loading...
马建仓 AI 助手
尝试更多
代码解读
代码找茬
代码优化
Go
1
https://gitee.com/zeng_huaping/goreplay.git
git@gitee.com:zeng_huaping/goreplay.git
zeng_huaping
goreplay
goreplay
master

搜索帮助