367 lines
9.9 KiB
Go
367 lines
9.9 KiB
Go
// Package queue
|
|
// @Link https://github.com/bufanyun/hotgo
|
|
// @Copyright Copyright (c) 2023 HotGo CLI
|
|
// @Author Ms <133814250@qq.com>
|
|
// @License https://github.com/bufanyun/hotgo/blob/master/LICENSE
|
|
package queue
|
|
|
|
import (
|
|
"context"
|
|
"crypto/sha256"
|
|
"crypto/sha512"
|
|
"errors"
|
|
"hash"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/IBM/sarama"
|
|
"github.com/gogf/gf/v2/errors/gerror"
|
|
"github.com/gogf/gf/v2/frame/g"
|
|
"github.com/xdg-go/scram"
|
|
|
|
"hotgo/internal/consts"
|
|
"hotgo/utility/simple"
|
|
)
|
|
|
|
type KafkaMq struct {
|
|
Partitions int32
|
|
producerIns sarama.AsyncProducer
|
|
consumerIns sarama.ConsumerGroup
|
|
}
|
|
|
|
type KafkaConfig struct {
|
|
ClientId string
|
|
Brokers []string
|
|
GroupID string
|
|
Partitions int32
|
|
Replication int16
|
|
Version string
|
|
UserName string
|
|
Password string
|
|
Mechanism string
|
|
}
|
|
|
|
type xdgSCRAMClient struct {
|
|
*scram.ClientConversation
|
|
HashGeneratorFcn scram.HashGeneratorFcn
|
|
}
|
|
|
|
func (x *xdgSCRAMClient) Begin(userName, password, authzID string) error {
|
|
client, err := x.HashGeneratorFcn.NewClient(userName, password, authzID)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
x.ClientConversation = client.NewConversation()
|
|
return nil
|
|
}
|
|
|
|
func (x *xdgSCRAMClient) Step(challenge string) (string, error) {
|
|
return x.ClientConversation.Step(challenge)
|
|
}
|
|
|
|
func (x *xdgSCRAMClient) Done() bool {
|
|
return x.ClientConversation.Done()
|
|
}
|
|
|
|
// SendMsg 按字符串类型生产数据
|
|
func (r *KafkaMq) SendMsg(topic string, body string) (mqMsg MqMsg, err error) {
|
|
return r.SendByteMsg(topic, []byte(body))
|
|
}
|
|
|
|
// SendByteMsg 生产数据
|
|
func (r *KafkaMq) SendByteMsg(topic string, body []byte) (mqMsg MqMsg, err error) {
|
|
msg := &sarama.ProducerMessage{
|
|
Topic: topic,
|
|
Value: sarama.ByteEncoder(body),
|
|
Timestamp: time.Now(),
|
|
}
|
|
|
|
if r.producerIns == nil {
|
|
err = gerror.New("queue kafka producerIns is nil")
|
|
return
|
|
}
|
|
|
|
r.producerIns.Input() <- msg
|
|
sendCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
defer cancel()
|
|
|
|
select {
|
|
case info := <-r.producerIns.Successes():
|
|
return MqMsg{
|
|
RunType: SendMsg,
|
|
Topic: info.Topic,
|
|
Offset: info.Offset,
|
|
Partition: info.Partition,
|
|
Timestamp: info.Timestamp,
|
|
}, nil
|
|
case fail := <-r.producerIns.Errors():
|
|
if nil != fail {
|
|
return mqMsg, fail.Err
|
|
}
|
|
case <-sendCtx.Done():
|
|
return mqMsg, gerror.New("send mqMst timeout")
|
|
}
|
|
return mqMsg, nil
|
|
}
|
|
|
|
func (r *KafkaMq) SendDelayMsg(topic string, body string, delaySecond int64) (mqMsg MqMsg, err error) {
|
|
err = gerror.New("implement me")
|
|
return
|
|
}
|
|
|
|
// ListenReceiveMsgDo 消费数据
|
|
func (r *KafkaMq) ListenReceiveMsgDo(topic string, receiveDo func(mqMsg MqMsg)) (err error) {
|
|
if r.consumerIns == nil {
|
|
return gerror.New("queue kafka consumer not register")
|
|
}
|
|
|
|
consumer := KaConsumer{
|
|
ready: make(chan bool),
|
|
receiveDoFun: receiveDo,
|
|
}
|
|
|
|
consumerCtx, cancel := context.WithCancel(context.Background())
|
|
go func(consumerCtx context.Context) {
|
|
for {
|
|
if err = r.consumerIns.Consume(consumerCtx, []string{topic}, &consumer); err != nil {
|
|
Logger().Errorf(ctx, "kafka error from consumer, topic:%s err:%+v", topic, err)
|
|
select {
|
|
case <-consumerCtx.Done():
|
|
return
|
|
case <-time.After(3 * time.Second):
|
|
}
|
|
}
|
|
|
|
if consumerCtx.Err() != nil {
|
|
Logger().Debugf(ctx, "kafka consumer stop : %v", consumerCtx.Err())
|
|
return
|
|
}
|
|
consumer.ready = make(chan bool)
|
|
}
|
|
}(consumerCtx)
|
|
|
|
// await till the consumer has been set up
|
|
select {
|
|
case <-consumer.ready:
|
|
case <-time.After(10 * time.Second):
|
|
Logger().Warningf(ctx, "kafka consumer setup timeout, topic:%s", topic)
|
|
}
|
|
Logger().Debug(ctx, "kafka consumer up and running!...")
|
|
|
|
simple.Event().Register(consts.EventServerClose, func(ctx context.Context, args ...interface{}) {
|
|
Logger().Debug(ctx, "kafka consumer close...")
|
|
cancel()
|
|
if err = r.consumerIns.Close(); err != nil {
|
|
Logger().Errorf(ctx, "kafka error closing client, err:%+v", err)
|
|
}
|
|
})
|
|
return
|
|
}
|
|
|
|
// RegisterKafkaMqConsumer 注册消费者
|
|
func RegisterKafkaMqConsumer(connOpt KafkaConfig) (client MqConsumer, err error) {
|
|
mqIns := &KafkaMq{}
|
|
kfkVersion, err := parseKafkaVersion(connOpt.Version)
|
|
if err != nil {
|
|
return
|
|
}
|
|
if !validateVersion(kfkVersion) {
|
|
kfkVersion = sarama.V2_4_0_0
|
|
}
|
|
|
|
brokers := connOpt.Brokers
|
|
conf := sarama.NewConfig()
|
|
conf.Consumer.Return.Errors = true
|
|
conf.Version = kfkVersion
|
|
configureKafkaSASL(conf, connOpt)
|
|
|
|
// 默认按随机方式消费
|
|
conf.Consumer.Group.Rebalance.Strategy = sarama.BalanceStrategyRange
|
|
conf.Consumer.Offsets.Initial = sarama.OffsetNewest
|
|
conf.Consumer.Offsets.AutoCommit.Interval = 10 * time.Millisecond
|
|
conf.ClientID = connOpt.ClientId
|
|
|
|
consumerClient, err := sarama.NewConsumerGroup(brokers, connOpt.GroupID, conf)
|
|
if err != nil {
|
|
return
|
|
}
|
|
mqIns.consumerIns = consumerClient
|
|
return mqIns, err
|
|
}
|
|
|
|
// RegisterKafkaProducer 注册并启动生产者接口实现
|
|
func RegisterKafkaProducer(connOpt KafkaConfig) (client MqProducer, err error) {
|
|
mqIns := &KafkaMq{}
|
|
connOpt.ClientId = "HOTGO-Producer"
|
|
|
|
// 这里如果使用go程需要处理chan同步问题
|
|
if err = doRegisterKafkaProducer(connOpt, mqIns); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return mqIns, nil
|
|
}
|
|
|
|
// doRegisterKafkaProducer 注册同步类型实例
|
|
func doRegisterKafkaProducer(connOpt KafkaConfig, mqIns *KafkaMq) (err error) {
|
|
kfkVersion, err := parseKafkaVersion(connOpt.Version)
|
|
if err != nil {
|
|
return
|
|
}
|
|
if !validateVersion(kfkVersion) {
|
|
kfkVersion = sarama.V2_4_0_0
|
|
}
|
|
|
|
brokers := connOpt.Brokers
|
|
conf := sarama.NewConfig()
|
|
// 等待服务器所有副本都保存成功后的响应
|
|
conf.Producer.RequiredAcks = sarama.WaitForAll
|
|
// 随机向partition发送消息
|
|
conf.Producer.Partitioner = sarama.NewRandomPartitioner
|
|
// 是否等待成功和失败后的响应,只有上面的RequireAcks设置不是NoReponse这里才有用.
|
|
conf.Producer.Return.Successes = true
|
|
|
|
conf.Producer.Return.Errors = true
|
|
conf.Producer.Compression = sarama.CompressionNone
|
|
conf.ClientID = connOpt.ClientId
|
|
|
|
conf.Version = kfkVersion
|
|
configureKafkaSASL(conf, connOpt)
|
|
|
|
mqIns.producerIns, err = sarama.NewAsyncProducer(brokers, conf)
|
|
if err != nil {
|
|
return
|
|
}
|
|
|
|
simple.Event().Register(consts.EventServerClose, func(ctx context.Context, args ...interface{}) {
|
|
g.Log().Debug(ctx, "kafka producer AsyncClose...")
|
|
mqIns.producerIns.AsyncClose()
|
|
})
|
|
return
|
|
}
|
|
|
|
func parseKafkaVersion(version string) (sarama.KafkaVersion, error) {
|
|
parts := strings.Split(version, ".")
|
|
if len(parts) == 4 && parts[0] != "0" && parts[3] == "0" {
|
|
version = strings.Join(parts[:3], ".")
|
|
}
|
|
return sarama.ParseKafkaVersion(version)
|
|
}
|
|
|
|
func configureKafkaSASL(conf *sarama.Config, connOpt KafkaConfig) {
|
|
if connOpt.UserName == "" {
|
|
return
|
|
}
|
|
conf.Net.SASL.Enable = true
|
|
conf.Net.SASL.User = connOpt.UserName
|
|
conf.Net.SASL.Password = connOpt.Password
|
|
mechanism := strings.ToUpper(strings.ReplaceAll(connOpt.Mechanism, "_", "-"))
|
|
if mechanism == "" {
|
|
mechanism = sarama.SASLTypeSCRAMSHA512
|
|
}
|
|
switch mechanism {
|
|
case sarama.SASLTypeSCRAMSHA512:
|
|
conf.Net.SASL.Mechanism = sarama.SASLTypeSCRAMSHA512
|
|
conf.Net.SASL.SCRAMClientGeneratorFunc = scramClientGenerator(sha512.New)
|
|
case sarama.SASLTypeSCRAMSHA256:
|
|
conf.Net.SASL.Mechanism = sarama.SASLTypeSCRAMSHA256
|
|
conf.Net.SASL.SCRAMClientGeneratorFunc = scramClientGenerator(sha256.New)
|
|
default:
|
|
conf.Net.SASL.Mechanism = sarama.SASLMechanism(mechanism)
|
|
}
|
|
}
|
|
|
|
// KafkaHealthCheck verifies broker connectivity by creating a short-lived Sarama client.
|
|
func KafkaHealthCheck(ctx context.Context) error {
|
|
if len(config.Kafka.Address) == 0 {
|
|
return errors.New("queue kafka address is empty")
|
|
}
|
|
kfkVersion, err := parseKafkaVersion(config.Kafka.Version)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if !validateVersion(kfkVersion) {
|
|
kfkVersion = sarama.V2_4_0_0
|
|
}
|
|
conf := sarama.NewConfig()
|
|
conf.Version = kfkVersion
|
|
conf.ClientID = "hotgo-healthz"
|
|
conf.Net.DialTimeout = 2 * time.Second
|
|
conf.Net.ReadTimeout = 2 * time.Second
|
|
conf.Net.WriteTimeout = 2 * time.Second
|
|
configureKafkaSASL(conf, KafkaConfig{
|
|
UserName: config.Kafka.Username,
|
|
Password: config.Kafka.Password,
|
|
Mechanism: config.Kafka.Mechanism,
|
|
})
|
|
|
|
select {
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
default:
|
|
}
|
|
client, err := sarama.NewClient(config.Kafka.Address, conf)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer func() { _ = client.Close() }()
|
|
if len(client.Brokers()) == 0 {
|
|
return errors.New("no kafka brokers available")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func scramClientGenerator(hashFunc func() hash.Hash) func() sarama.SCRAMClient {
|
|
return func() sarama.SCRAMClient {
|
|
return &xdgSCRAMClient{HashGeneratorFcn: scram.HashGeneratorFcn(hashFunc)}
|
|
}
|
|
}
|
|
|
|
// validateVersion 验证版本是否有效
|
|
func validateVersion(version sarama.KafkaVersion) bool {
|
|
for _, item := range sarama.SupportedVersions {
|
|
if version.String() == item.String() {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
type KaConsumer struct {
|
|
ready chan bool
|
|
receiveDoFun func(mqMsg MqMsg)
|
|
}
|
|
|
|
// Setup is run at the beginning of a new session, before ConsumeClaim
|
|
func (consumer *KaConsumer) Setup(sarama.ConsumerGroupSession) error {
|
|
// Mark the consumer as ready
|
|
close(consumer.ready)
|
|
return nil
|
|
}
|
|
|
|
// Cleanup is run at the end of a session, once all ConsumeClaim goroutines have exited
|
|
func (consumer *KaConsumer) Cleanup(sarama.ConsumerGroupSession) error {
|
|
return nil
|
|
}
|
|
|
|
// ConsumeClaim must start a consumer loop of ConsumerGroupClaim's Messages().
|
|
func (consumer *KaConsumer) ConsumeClaim(session sarama.ConsumerGroupSession, claim sarama.ConsumerGroupClaim) error {
|
|
// NOTE:
|
|
// Do not move the code below to a goroutine.
|
|
// The `ConsumeClaim` itself is called within a goroutine, see:
|
|
// https://github.com/Shopify/sarama/blob/master/consumer_group.go#L27-L29
|
|
// `ConsumeClaim` 方法已经是 goroutine 调用 不要在该方法内进行 goroutine
|
|
for message := range claim.Messages() {
|
|
consumer.receiveDoFun(MqMsg{
|
|
RunType: ReceiveMsg,
|
|
Topic: message.Topic,
|
|
Body: message.Value,
|
|
Offset: message.Offset,
|
|
Timestamp: message.Timestamp,
|
|
Partition: message.Partition,
|
|
})
|
|
session.MarkMessage(message, "")
|
|
}
|
|
return nil
|
|
}
|