Files
im_server/server/kafka_runtime.go
T

625 lines
19 KiB
Go

package main
import (
"context"
"crypto/tls"
"encoding/json"
"errors"
"regexp"
"strings"
"sync"
"sync/atomic"
"time"
"github.com/segmentio/kafka-go"
"github.com/segmentio/kafka-go/sasl"
"github.com/segmentio/kafka-go/sasl/scram"
"github.com/tinode/chat/server/logs"
immetrics "github.com/tinode/chat/server/metrics"
"github.com/tinode/chat/server/store/types"
)
const (
kafkaTopicAdminCommand = "im.admin.command"
kafkaTopicAdminCommandResult = "im.admin.command.result"
kafkaTopicConfigEvent = "im.config.event"
kafkaTopicBusinessEvent = "im.business.event"
kafkaTopicMessagePersist = "im.message.persist"
kafkaTopicGroupMessagePersist = "im.group.message.persist"
kafkaTopicMessagePush = "im.message.push"
kafkaTopicFileSecurity = "im.file.security"
kafkaTopicWorkbenchEvent = "im.workbench.event"
kafkaTopicCacheInvalidate = "im.cache.invalidate"
kafkaTopicAuditEvent = "im.audit.event"
kafkaTopicDeadLetter = "im.dead.letter"
)
type kafkaConfig struct {
Enabled bool `json:"enabled"`
Brokers []string `json:"brokers"`
ClientID string `json:"client_id"`
ConsumerGroup string `json:"consumer_group"`
Topics map[string]string `json:"topics"`
DialTimeoutMS int `json:"dial_timeout_ms"`
WriteTimeoutMS int `json:"write_timeout_ms"`
ReadTimeoutMS int `json:"read_timeout_ms"`
BatchTimeoutMS int `json:"batch_timeout_ms"`
RequiredAcks int `json:"required_acks"`
AllowAutoTopic bool `json:"allow_auto_topic_creation"`
TLSEnabled bool `json:"tls_enabled"`
SASLEnabled bool `json:"sasl_enabled"`
SASLMechanism string `json:"sasl_mechanism"`
SASLUsername string `json:"sasl_username"`
SASLPassword string `json:"sasl_password"`
MaxRetryCount int `json:"max_retry_count"`
RetryBackoffMS []int `json:"retry_backoff_ms"`
MessageBatchSize int `json:"message_batch_size"`
MessageBatchWait int `json:"message_batch_wait_ms"`
MessageRetries int `json:"message_single_retry_count"`
DeadRetentionSec int `json:"dead_letter_retention_seconds"`
}
type kafkaEnvelope struct {
EventID string `json:"event_id"`
TraceID string `json:"trace_id,omitempty"`
TenantCode string `json:"tenant_code,omitempty"`
TenantID types.TenantID `json:"tenant_id,omitempty"`
SourceService string `json:"source_service,omitempty"`
TargetService string `json:"target_service,omitempty"`
EventType string `json:"event_type"`
BizType string `json:"biz_type,omitempty"`
BizID string `json:"biz_id,omitempty"`
IdempotencyKey string `json:"idempotency_key,omitempty"`
Version int64 `json:"version,omitempty"`
OccurredAt time.Time `json:"occurred_at"`
Actor *kafkaActor `json:"actor,omitempty"`
Payload json.RawMessage `json:"payload,omitempty"`
ParentEventID string `json:"parent_event_id,omitempty"`
}
type kafkaActor struct {
Type string `json:"type,omitempty"`
ID string `json:"id,omitempty"`
DisplayName string `json:"display_name,omitempty"`
}
type kafkaRuntime struct {
config kafkaConfig
topics map[string]string
writer *kafka.Writer
ctx context.Context
cancel context.CancelFunc
wg sync.WaitGroup
}
var kafkaBus *kafkaRuntime
var kafkaConsumerGroupUnsafeChars = regexp.MustCompile(`[^A-Za-z0-9._-]+`)
func initKafkaRuntime(raw json.RawMessage) error {
if len(raw) == 0 {
kafkaBus = nil
logs.Info.Println("Kafka disabled: no config")
return nil
}
var cfg kafkaConfig
if err := json.Unmarshal(raw, &cfg); err != nil {
return err
}
if !cfg.Enabled {
kafkaBus = nil
logs.Info.Println("Kafka disabled")
return nil
}
normalizeKafkaConfig(&cfg)
if len(cfg.Brokers) == 0 {
return errors.New("kafka: enabled but brokers are not configured")
}
saslMechanism, err := kafkaSASLMechanism(cfg)
if err != nil {
return err
}
ctx, cancel := context.WithCancel(context.Background())
rt := &kafkaRuntime{
config: cfg,
topics: kafkaTopics(cfg.Topics),
ctx: ctx,
cancel: cancel,
}
transport := &kafka.Transport{
ClientID: cfg.ClientID,
}
if cfg.TLSEnabled {
transport.TLS = &tls.Config{MinVersion: tls.VersionTLS12}
}
transport.SASL = saslMechanism
rt.writer = &kafka.Writer{
Addr: kafka.TCP(cfg.Brokers...),
Balancer: &kafka.Hash{},
Transport: transport,
RequiredAcks: kafka.RequiredAcks(cfg.RequiredAcks),
AllowAutoTopicCreation: cfg.AllowAutoTopic,
BatchTimeout: time.Duration(cfg.BatchTimeoutMS) * time.Millisecond,
WriteTimeout: time.Duration(cfg.WriteTimeoutMS) * time.Millisecond,
ReadTimeout: time.Duration(cfg.ReadTimeoutMS) * time.Millisecond,
MaxAttempts: cfg.MaxRetryCount,
}
if err := rt.ping(context.Background()); err != nil {
cancel()
_ = rt.writer.Close()
return err
}
kafkaBus = rt
logs.Info.Printf("Kafka enabled: brokers=%s client_id=%s group=%s", strings.Join(cfg.Brokers, ","), cfg.ClientID, cfg.ConsumerGroup)
return nil
}
func kafkaConfigured(raw json.RawMessage) bool {
var cfg kafkaConfig
return len(raw) > 0 && json.Unmarshal(raw, &cfg) == nil && cfg.Enabled
}
func normalizeKafkaConfig(cfg *kafkaConfig) {
if cfg.ClientID == "" {
cfg.ClientID = "im-server"
}
if cfg.ConsumerGroup == "" {
cfg.ConsumerGroup = "im-server"
}
if cfg.DialTimeoutMS <= 0 {
cfg.DialTimeoutMS = 3000
}
if cfg.WriteTimeoutMS <= 0 {
cfg.WriteTimeoutMS = 3000
}
if cfg.ReadTimeoutMS <= 0 {
cfg.ReadTimeoutMS = 3000
}
if cfg.BatchTimeoutMS <= 0 {
cfg.BatchTimeoutMS = 10
}
if cfg.RequiredAcks == 0 {
cfg.RequiredAcks = int(kafka.RequireAll)
}
if cfg.MaxRetryCount <= 0 {
cfg.MaxRetryCount = 3
}
if len(cfg.RetryBackoffMS) == 0 {
cfg.RetryBackoffMS = []int{10000, 30000, 60000, 300000, 900000}
}
if cfg.MessageBatchSize <= 0 {
cfg.MessageBatchSize = 100
} else if cfg.MessageBatchSize > 500 {
cfg.MessageBatchSize = 500
}
if cfg.MessageBatchWait <= 0 {
cfg.MessageBatchWait = 10
}
if cfg.MessageRetries <= 0 {
cfg.MessageRetries = 3
}
if cfg.DeadRetentionSec <= 0 {
cfg.DeadRetentionSec = 90 * 24 * 3600
}
if cfg.SASLMechanism == "" {
cfg.SASLMechanism = "SCRAM-SHA-512"
}
}
func kafkaSASLMechanism(cfg kafkaConfig) (sasl.Mechanism, error) {
if !cfg.SASLEnabled && cfg.SASLUsername == "" {
return nil, nil
}
if cfg.SASLUsername == "" || cfg.SASLPassword == "" {
return nil, errors.New("kafka: sasl enabled but username or password is empty")
}
switch strings.ToUpper(strings.ReplaceAll(cfg.SASLMechanism, "_", "-")) {
case "", "SCRAM-SHA-512":
return scram.Mechanism(scram.SHA512, cfg.SASLUsername, cfg.SASLPassword)
case "SCRAM-SHA-256":
return scram.Mechanism(scram.SHA256, cfg.SASLUsername, cfg.SASLPassword)
default:
return nil, errors.New("kafka: unsupported sasl mechanism " + cfg.SASLMechanism)
}
}
func kafkaTopics(overrides map[string]string) map[string]string {
topics := map[string]string{
"admin_command": kafkaTopicAdminCommand,
"admin_command_result": kafkaTopicAdminCommandResult,
"config_event": kafkaTopicConfigEvent,
"business_event": kafkaTopicBusinessEvent,
"message_persist": kafkaTopicMessagePersist,
"group_message_persist": kafkaTopicGroupMessagePersist,
"message_push": kafkaTopicMessagePush,
"file_security": kafkaTopicFileSecurity,
"workbench_external": kafkaTopicWorkbenchEvent,
"cache_invalidate": kafkaTopicCacheInvalidate,
"audit_event": kafkaTopicAuditEvent,
"dead_letter": kafkaTopicDeadLetter,
}
for key, value := range overrides {
if value != "" {
topics[key] = value
}
}
return topics
}
func (rt *kafkaRuntime) ping(ctx context.Context) error {
timeout := time.Duration(rt.config.DialTimeoutMS) * time.Millisecond
dialer := rt.dialer()
pingCtx, cancel := context.WithTimeout(ctx, timeout)
defer cancel()
conn, err := dialer.DialContext(pingCtx, "tcp", rt.config.Brokers[0])
if err != nil {
return err
}
return conn.Close()
}
func kafkaEnabled() bool {
return kafkaBus != nil
}
func kafkaTopic(name string) string {
if kafkaBus == nil {
return ""
}
return kafkaBus.topics[name]
}
func publishKafkaEnvelope(ctx context.Context, topicName string, key string, envelope kafkaEnvelope) error {
if kafkaBus == nil {
return nil
}
if envelope.OccurredAt.IsZero() {
envelope.OccurredAt = time.Now().UTC()
}
topic := kafkaBus.topics[topicName]
if topic == "" {
topic = topicName
}
value, err := json.Marshal(envelope)
if err != nil {
return err
}
return kafkaBus.writer.WriteMessages(ctx, kafka.Message{
Topic: topic,
Key: []byte(key),
Value: value,
Time: envelope.OccurredAt,
})
}
type kafkaMessageHandler func(context.Context, kafkaEnvelope, kafka.Message) error
type kafkaMessageCommitter func(context.Context, kafka.Message) error
type kafkaRetryWaiter func(context.Context, time.Duration) error
type kafkaBatchHandler func(context.Context, []kafka.Message) error
type kafkaBatchCommitter func(context.Context, ...kafka.Message) error
type kafkaMessageFetcher func(context.Context) (kafka.Message, error)
func startKafkaConsumer(topicName string, handler kafkaMessageHandler) {
startKafkaConsumerWithGroup(topicName, "", handler)
}
func startKafkaBroadcastConsumer(topicName string, handler kafkaMessageHandler) string {
if kafkaBus == nil || handler == nil {
return ""
}
groupID := kafkaBroadcastConsumerGroup(topicName)
startKafkaConsumerWithGroup(topicName, groupID, handler)
return groupID
}
func startKafkaConsumerWithGroup(topicName, groupID string, handler kafkaMessageHandler) {
if kafkaBus == nil || handler == nil {
return
}
topic := kafkaBus.topics[topicName]
if topic == "" {
topic = topicName
}
groupOverride := groupID != ""
if groupID == "" {
groupID = kafkaBus.config.ConsumerGroup
}
if topicName != "" && !groupOverride {
groupID += "." + topicName
}
reader := kafka.NewReader(kafka.ReaderConfig{
Brokers: kafkaBus.config.Brokers,
GroupID: groupID,
Topic: topic,
StartOffset: kafka.FirstOffset,
MinBytes: 1,
MaxBytes: 10e6,
MaxWait: time.Second,
Dialer: kafkaBus.dialer(),
})
readerMetrics := startKafkaReaderMetrics(kafkaBus.ctx, reader, topic, groupID)
kafkaBus.wg.Add(1)
go func() {
defer kafkaBus.wg.Done()
defer reader.Close()
for {
msg, err := reader.FetchMessage(kafkaBus.ctx)
if err != nil {
if errors.Is(err, context.Canceled) || errors.Is(kafkaBus.ctx.Err(), context.Canceled) {
return
}
logs.Err.Println("Kafka fetch failed:", err)
continue
}
readerMetrics.observe(msg.Time)
commit := func(ctx context.Context, message kafka.Message) error {
return reader.CommitMessages(ctx, message)
}
if err = consumeKafkaMessage(kafkaBus.ctx, msg, handler, commit, waitKafkaRetry); err != nil {
if errors.Is(err, context.Canceled) || errors.Is(kafkaBus.ctx.Err(), context.Canceled) {
return
}
logs.Err.Printf("Kafka message stopped without commit: topic=%s partition=%d offset=%d err=%v", msg.Topic, msg.Partition, msg.Offset, err)
}
readerMetrics.clear()
}
}()
logs.Info.Printf("Kafka consumer started: topic=%s group=%s", topic, groupID)
}
func startKafkaBatchConsumer(topicName string, handler kafkaBatchHandler) {
if kafkaBus == nil || handler == nil {
return
}
topic := kafkaBus.topics[topicName]
if topic == "" {
topic = topicName
}
groupID := kafkaBus.config.ConsumerGroup + "." + topicName
reader := kafka.NewReader(kafka.ReaderConfig{
Brokers: kafkaBus.config.Brokers, GroupID: groupID, Topic: topic, StartOffset: kafka.FirstOffset,
MinBytes: 1, MaxBytes: 10e6, MaxWait: time.Second, Dialer: kafkaBus.dialer(),
})
readerMetrics := startKafkaReaderMetrics(kafkaBus.ctx, reader, topic, groupID)
kafkaBus.wg.Add(1)
go func() {
defer kafkaBus.wg.Done()
defer reader.Close()
for {
messages, fetchErr := fetchKafkaMessageBatch(kafkaBus.ctx, reader, kafkaBus.config.MessageBatchSize,
time.Duration(kafkaBus.config.MessageBatchWait)*time.Millisecond)
if fetchErr != nil && len(messages) == 0 {
if errors.Is(fetchErr, context.Canceled) || errors.Is(kafkaBus.ctx.Err(), context.Canceled) {
return
}
logs.Err.Println("Kafka batch fetch failed:", fetchErr)
continue
}
if fetchErr != nil {
logs.Warn.Printf("Kafka batch fetch interrupted; processing retained messages: topic=%s size=%d err=%v",
topic, len(messages), fetchErr)
}
if len(messages) > 0 {
readerMetrics.observe(messages[0].Time)
}
if err := consumeKafkaMessageBatch(kafkaBus.ctx, messages, handler, reader.CommitMessages, waitKafkaRetry); err != nil {
if errors.Is(err, context.Canceled) || errors.Is(kafkaBus.ctx.Err(), context.Canceled) {
return
}
logs.Err.Printf("Kafka batch stopped without commit: topic=%s size=%d err=%v", topic, len(messages), err)
}
readerMetrics.clear()
}
}()
logs.Info.Printf("Kafka batch consumer started: topic=%s group=%s size=%d wait_ms=%d", topic, groupID, kafkaBus.config.MessageBatchSize, kafkaBus.config.MessageBatchWait)
}
type kafkaReaderMetrics struct {
reader *kafka.Reader
topic string
group string
oldestNano atomic.Int64
}
func startKafkaReaderMetrics(ctx context.Context, reader *kafka.Reader, topic, group string) *kafkaReaderMetrics {
metrics := &kafkaReaderMetrics{reader: reader, topic: topic, group: group}
kafkaBus.wg.Add(1)
go func() {
defer kafkaBus.wg.Done()
ticker := time.NewTicker(5 * time.Second)
defer ticker.Stop()
metrics.update()
for {
select {
case <-ctx.Done():
return
case <-ticker.C:
metrics.update()
}
}
}()
return metrics
}
func (m *kafkaReaderMetrics) observe(messageTime time.Time) {
if !messageTime.IsZero() {
m.oldestNano.Store(messageTime.UnixNano())
}
m.update()
}
func (m *kafkaReaderMetrics) clear() {
m.oldestNano.Store(0)
m.update()
}
func (m *kafkaReaderMetrics) update() {
stats := m.reader.Stats()
if stats.Lag >= 0 {
immetrics.KafkaConsumerLag.WithLabelValues(clusterPodName(), m.topic, m.group).Set(float64(stats.Lag))
}
age := float64(0)
if oldest := m.oldestNano.Load(); oldest > 0 {
age = time.Since(time.Unix(0, oldest)).Seconds()
if age < 0 {
age = 0
}
}
immetrics.KafkaBacklogOldestAge.WithLabelValues(clusterPodName(), m.topic, m.group).Set(age)
}
func fetchKafkaMessageBatch(ctx context.Context, reader *kafka.Reader, maxSize int, maxWait time.Duration) ([]kafka.Message, error) {
return fetchKafkaMessageBatchWith(ctx, maxSize, maxWait, reader.FetchMessage)
}
func fetchKafkaMessageBatchWith(ctx context.Context, maxSize int, maxWait time.Duration, fetch kafkaMessageFetcher) ([]kafka.Message, error) {
first, err := fetch(ctx)
if err != nil {
return nil, err
}
messages := []kafka.Message{first}
deadline := time.Now().Add(maxWait)
for len(messages) < maxSize {
remaining := time.Until(deadline)
if remaining <= 0 {
break
}
fetchCtx, cancel := context.WithTimeout(ctx, remaining)
message, fetchErr := fetch(fetchCtx)
cancel()
if errors.Is(fetchErr, context.DeadlineExceeded) && ctx.Err() == nil {
break
}
if fetchErr != nil {
// FetchMessage advances the reader's internal position before offsets are
// committed. Preserve a partial batch so a later higher-offset commit cannot
// skip messages already fetched into this batch.
return messages, fetchErr
}
messages = append(messages, message)
}
return messages, nil
}
func consumeKafkaMessageBatch(ctx context.Context, messages []kafka.Message, handler kafkaBatchHandler, commit kafkaBatchCommitter, wait kafkaRetryWaiter) error {
for attempt := 1; ; attempt++ {
if err := handler(ctx, messages); err == nil {
return commitKafkaMessageBatch(ctx, messages, commit, wait)
}
if err := wait(ctx, kafkaRetryBackoff(attempt)); err != nil {
return err
}
}
}
func consumeKafkaMessage(ctx context.Context, msg kafka.Message, handler kafkaMessageHandler, commit kafkaMessageCommitter, wait kafkaRetryWaiter) error {
var envelope kafkaEnvelope
if err := json.Unmarshal(msg.Value, &envelope); err != nil {
return commitKafkaMessage(ctx, msg, commit, wait)
}
for attempt := 1; ; attempt++ {
if err := handler(ctx, envelope, msg); err == nil {
return commitKafkaMessage(ctx, msg, commit, wait)
}
if err := wait(ctx, kafkaRetryBackoff(attempt)); err != nil {
return err
}
}
}
func commitKafkaMessage(ctx context.Context, msg kafka.Message, commit kafkaMessageCommitter, wait kafkaRetryWaiter) error {
for attempt := 1; ; attempt++ {
if err := commit(ctx, msg); err == nil {
return nil
}
if err := wait(ctx, kafkaRetryBackoff(attempt)); err != nil {
return err
}
}
}
func commitKafkaMessageBatch(ctx context.Context, messages []kafka.Message, commit kafkaBatchCommitter, wait kafkaRetryWaiter) error {
for attempt := 1; ; attempt++ {
if err := commit(ctx, messages...); err == nil {
return nil
}
if err := wait(ctx, kafkaRetryBackoff(attempt)); err != nil {
return err
}
}
}
func kafkaRetryBackoff(attempt int) time.Duration {
if kafkaBus == nil || len(kafkaBus.config.RetryBackoffMS) == 0 {
return time.Second
}
if attempt < 1 {
attempt = 1
}
index := attempt - 1
if index >= len(kafkaBus.config.RetryBackoffMS) {
index = len(kafkaBus.config.RetryBackoffMS) - 1
}
return time.Duration(kafkaBus.config.RetryBackoffMS[index]) * time.Millisecond
}
func waitKafkaRetry(ctx context.Context, delay time.Duration) error {
timer := time.NewTimer(delay)
defer timer.Stop()
select {
case <-ctx.Done():
return ctx.Err()
case <-timer.C:
return nil
}
}
func kafkaBroadcastConsumerGroup(topicName string) string {
if kafkaBus == nil {
return ""
}
base := strings.TrimSpace(kafkaBus.config.ConsumerGroup)
if base == "" {
base = "im-server"
}
nodeID := strings.TrimSpace(currentNodeID())
if nodeID == "" {
nodeID = "local"
}
nodeID = kafkaConsumerGroupUnsafeChars.ReplaceAllString(nodeID, "-")
return base + "." + topicName + "." + nodeID
}
func (rt *kafkaRuntime) dialer() *kafka.Dialer {
dialer := &kafka.Dialer{
ClientID: rt.config.ClientID,
Timeout: time.Duration(rt.config.DialTimeoutMS) * time.Millisecond,
}
if rt.config.TLSEnabled {
dialer.TLS = &tls.Config{MinVersion: tls.VersionTLS12}
}
if mechanism, err := kafkaSASLMechanism(rt.config); err != nil {
logs.Err.Println("Kafka SASL config failed:", err)
} else {
dialer.SASLMechanism = mechanism
}
return dialer
}
func stopKafkaRuntime() {
if kafkaBus == nil {
return
}
kafkaBus.cancel()
kafkaBus.wg.Wait()
if kafkaBus.writer != nil {
_ = kafkaBus.writer.Close()
}
logs.Info.Println("Stopped Kafka runtime")
kafkaBus = nil
}