80 lines
1.7 KiB
Go
80 lines
1.7 KiB
Go
package server
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
|
|
"ky-go-kratos/app/websocket/internal/conf"
|
|
"ky-go-kratos/app/websocket/internal/service"
|
|
"ky-go-kratos/pkg/kafka"
|
|
|
|
"github.com/go-kratos/kratos/v2/log"
|
|
)
|
|
|
|
type KafkaConsumerServer struct {
|
|
consumer *kafka.KafkaConsumer
|
|
messageHandler *service.MessageHandler
|
|
wg sync.WaitGroup
|
|
}
|
|
|
|
func NewKafkaConsumerServer(c *conf.Data, messageHandler *service.MessageHandler) (*KafkaConsumerServer, func()) {
|
|
consumer := kafka.NewKafkaReader(c.Kafka.Brokers, c.Kafka.Topic, int(c.Kafka.Partition))
|
|
server := &KafkaConsumerServer{
|
|
consumer: consumer,
|
|
messageHandler: messageHandler,
|
|
}
|
|
|
|
ctx := context.Background()
|
|
if err := server.Start(ctx); err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
return server, func() {
|
|
if err := server.Stop(ctx); err != nil {
|
|
log.Errorf("failed to stop kafka consumer: %v", err)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (s *KafkaConsumerServer) Start(ctx context.Context) error {
|
|
if err := s.messageHandler.Start(); err != nil {
|
|
return err
|
|
}
|
|
|
|
s.wg.Add(1)
|
|
go s.consumeMessages(ctx)
|
|
return nil
|
|
}
|
|
|
|
func (s *KafkaConsumerServer) Stop(ctx context.Context) error {
|
|
s.wg.Wait()
|
|
if err := s.messageHandler.Stop(); err != nil {
|
|
return err
|
|
}
|
|
return s.consumer.Close()
|
|
}
|
|
|
|
func (s *KafkaConsumerServer) consumeMessages(ctx context.Context) {
|
|
defer s.wg.Done()
|
|
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
return
|
|
default:
|
|
message, err := s.consumer.Reader.ReadMessage(ctx)
|
|
log.Infof("message data: %v", message)
|
|
if err != nil {
|
|
if s.consumer.IsTransientNetworkError(err) {
|
|
continue
|
|
}
|
|
return
|
|
}
|
|
|
|
if err := s.messageHandler.HandleMessage(&message); err != nil {
|
|
log.Errorf("failed to handle message: %v", err)
|
|
}
|
|
}
|
|
}
|
|
}
|