Files
im_server/server/kafka_message_persist.go
T

342 lines
13 KiB
Go

package main
import (
"context"
"encoding/json"
"errors"
"fmt"
"strings"
"time"
"github.com/segmentio/kafka-go"
"github.com/tinode/chat/server/logs"
immetrics "github.com/tinode/chat/server/metrics"
"github.com/tinode/chat/server/store"
"github.com/tinode/chat/server/store/types"
)
type kafkaMessagePersistPayload struct {
Message types.Message `json:"message"`
AttachmentURLs []string `json:"attachment_urls,omitempty"`
ReadBySender bool `json:"read_by_sender,omitempty"`
MessageKind string `json:"message_kind"`
ShardMonth string `json:"shard_month,omitempty"`
}
var publishKafkaMessagePersist = publishKafkaEnvelope
func startKafkaMessagePersistConsumer() {
startKafkaBatchConsumer("message_persist", func(ctx context.Context, messages []kafka.Message) error {
return handleKafkaMessagePersistBatch(ctx, messages, "message_persist")
})
startKafkaBatchConsumer("group_message_persist", func(ctx context.Context, messages []kafka.Message) error {
return handleKafkaMessagePersistBatch(ctx, messages, "group_message_persist")
})
}
func saveMessageMaybeKafka(ctx context.Context, tenantID types.TenantID, msg *types.Message, attachmentURLs []string, readBySender bool) (*types.MessageSaveResult, error) {
if !kafkaEnabled() {
return store.Messages.Save(tenantID, msg, attachmentURLs, readBySender)
}
if msg == nil || msg.TenantID != tenantID || msg.Topic == "" || msg.SeqId <= 0 {
return nil, types.ErrMalformed
}
msg.InitTimes()
if msg.Uid().IsZero() {
msg.SetUid(store.Store.GetUid())
}
isGroupMessage := types.GetTopicCat(msg.Topic) == types.TopicCatGrp
clientMsgID, err := store.NormalizeClientMsgID(msg.Head)
if err != nil {
return nil, err
}
msg.ClientMsgID = clientMsgID
fromUid := types.ParseUid(msg.From)
if clientMsgID != "" {
claim, claimed, runtimeReady, claimErr := store.ClaimMessageUnique(tenantID, msg.Topic, fromUid, clientMsgID)
if claimErr != nil {
return nil, claimErr
}
if runtimeReady && !claimed {
if claim != nil && claim.Seq > 0 {
msg.SeqId = claim.Seq
return &types.MessageSaveResult{SeqId: claim.Seq, Duplicate: true, Pending: claim.Status == "pending"}, nil
}
if claim != nil && claim.Status == "pending" {
return &types.MessageSaveResult{Duplicate: true, Pending: true}, types.ErrDuplicate
}
return nil, types.ErrDuplicate
}
}
candidateMonth := msg.CreatedAt.UTC().Format("200601")
if seq, shardMonth, runtimeReady, seqErr := store.ReserveMessageSeq(tenantID, msg.Topic, msg.SeqId-1, candidateMonth); seqErr != nil {
if clientMsgID != "" {
_ = store.MarkMessageUniqueFailed(tenantID, msg.Topic, fromUid, clientMsgID)
}
return nil, seqErr
} else if runtimeReady {
msg.SeqId = seq
if isGroupMessage {
msg.ShardMonth = shardMonth
}
}
if msg.ShardMonth == "" && isGroupMessage {
msg.ShardMonth = candidateMonth
}
if err = enqueueKafkaMessagePersist(ctx, tenantID, msg, attachmentURLs, readBySender); err != nil {
if clientMsgID != "" {
if markErr := store.MarkMessageUniquePending(tenantID, msg.Topic, fromUid, clientMsgID, msg.SeqId); markErr != nil {
logs.Warn.Printf("kafka message persist: failed to mark message unique key pending topic=%s seq=%d: %v", msg.Topic, msg.SeqId, markErr)
}
}
logs.Warn.Printf("kafka message persist: publish state unknown topic=%s seq=%d client_mid=%s err=%v", msg.Topic, msg.SeqId, clientMsgID, err)
return &types.MessageSaveResult{SeqId: msg.SeqId, Pending: true, MarkedReadBySender: readBySender && !fromUid.IsZero()}, err
}
if clientMsgID != "" {
if err = store.MarkMessageUniquePending(tenantID, msg.Topic, fromUid, clientMsgID, msg.SeqId); err != nil {
logs.Warn.Printf("kafka message persist: failed to mark message unique key pending topic=%s seq=%d: %v", msg.Topic, msg.SeqId, err)
}
}
if isGroupMessage {
if cacheErr := store.CachePendingHotMessage(tenantID, msg); cacheErr != nil {
logs.Warn.Printf("kafka message persist: failed to cache pending hot message after publish topic=%s seq=%d: %v", msg.Topic, msg.SeqId, cacheErr)
}
}
return &types.MessageSaveResult{SeqId: msg.SeqId, Pending: true, MarkedReadBySender: readBySender && !fromUid.IsZero()}, nil
}
func enqueueKafkaMessagePersist(ctx context.Context, tenantID types.TenantID, msg *types.Message, attachmentURLs []string, readBySender bool) error {
kind := "p2p"
if types.GetTopicCat(msg.Topic) == types.TopicCatGrp {
kind = "group"
}
payload, err := json.Marshal(kafkaMessagePersistPayload{
Message: *msg,
AttachmentURLs: attachmentURLs,
ReadBySender: readBySender,
MessageKind: kind,
ShardMonth: msg.ShardMonth,
})
if err != nil {
return err
}
eventID := fmt.Sprintf("msg.persist.%d.%s.%d", tenantID, msg.Topic, msg.SeqId)
envelope := kafkaEnvelope{
EventID: eventID,
TenantID: tenantID,
SourceService: "im-server",
TargetService: "im-server",
EventType: "message.persist." + kind,
BizType: "message",
BizID: fmt.Sprintf("%s:%d", msg.Topic, msg.SeqId),
IdempotencyKey: msg.ClientMsgID,
Version: msg.SeqId,
OccurredAt: msg.CreatedAt,
Actor: &kafkaActor{Type: "user", ID: types.ParseUid(msg.From).UserId()},
Payload: payload,
}
return publishKafkaMessagePersist(ctx, kafkaMessagePersistTopicName(msg.Topic), msg.Topic, envelope)
}
type kafkaMessagePersistItem struct {
envelope kafkaEnvelope
tenantID types.TenantID
payload kafkaMessagePersistPayload
}
func handleKafkaMessagePersistBatch(ctx context.Context, messages []kafka.Message, sourceTopicName string) error {
groups := make(map[string][]*kafkaMessagePersistItem)
for _, message := range messages {
item, terminal, err := prepareKafkaMessagePersistItem(ctx, message, sourceTopicName)
if err != nil {
return err
}
if terminal {
continue
}
key := fmt.Sprintf("%d", item.tenantID)
if item.payload.MessageKind == "group" {
key += ":" + item.payload.Message.ShardMonth
}
groups[key] = append(groups[key], item)
}
for _, items := range groups {
kind := items[0].payload.MessageKind
if kind == "" {
kind = "p2p"
}
started := time.Now()
batch := make([]store.PreparedMessageBatchItem, 0, len(items))
for _, item := range items {
batch = append(batch, store.PreparedMessageBatchItem{Message: &item.payload.Message, ReadBySender: item.payload.ReadBySender})
}
if err := store.PersistPreparedMessageBatch(items[0].tenantID, batch); err == nil {
immetrics.KafkaMessagePersistTotal.WithLabelValues(kind, "batch_ok").Add(float64(len(items)))
immetrics.MessagePersistDuration.WithLabelValues(kind, "batch_ok").Observe(time.Since(started).Seconds())
for _, item := range items {
if err = finishKafkaMessagePersist(ctx, item); err != nil {
if err = persistKafkaMessageIndividually(ctx, item); err != nil {
return err
}
}
}
continue
} else {
immetrics.KafkaMessagePersistTotal.WithLabelValues(kind, "batch_failed").Add(float64(len(items)))
immetrics.MessagePersistDuration.WithLabelValues(kind, "batch_failed").Observe(time.Since(started).Seconds())
logs.Warn.Printf("kafka message batch persist failed, degrading to singles: tenant=%d kind=%s shard=%s size=%d err=%v",
items[0].tenantID, items[0].payload.MessageKind, items[0].payload.Message.ShardMonth, len(items), err)
}
for _, item := range items {
if err := persistKafkaMessageIndividually(ctx, item); err != nil {
return err
}
}
}
return nil
}
func prepareKafkaMessagePersistItem(ctx context.Context, message kafka.Message, sourceTopicName string) (*kafkaMessagePersistItem, bool, error) {
var envelope kafkaEnvelope
if err := json.Unmarshal(message.Value, &envelope); err != nil {
logs.Err.Printf("kafka message malformed envelope skipped: partition=%d offset=%d err=%v", message.Partition, message.Offset, err)
return nil, true, nil
}
tenantID := envelope.TenantID
var err error
if tenantID.IsZero() {
tenantID, err = kafkaTenantID(envelope.TenantCode)
if err != nil {
return nil, false, err
}
}
consume, claimed, err := beginKafkaConsumeForTopic(tenantID, envelope, sourceTopicName)
if err != nil {
return nil, false, err
}
if !claimed {
return nil, true, nil
}
item := &kafkaMessagePersistItem{envelope: envelope, tenantID: tenantID}
if err = json.Unmarshal(envelope.Payload, &item.payload); err != nil {
return nil, true, markKafkaDeadLetter(ctx, tenantID, envelope, consume, "MALFORMED_PAYLOAD", err.Error())
}
if item.payload.Message.TenantID == types.ZeroTenantID {
item.payload.Message.TenantID = tenantID
}
if item.payload.Message.ShardMonth == "" {
item.payload.Message.ShardMonth = item.payload.ShardMonth
}
if item.payload.Message.TenantID != tenantID {
return nil, true, markKafkaDeadLetter(ctx, tenantID, envelope, consume, "TENANT_MISMATCH", "message tenant does not match envelope tenant")
}
if err = validateKafkaMessageKind(item.payload); err != nil {
return nil, true, markKafkaDeadLetter(ctx, tenantID, envelope, consume, "MALFORMED_PAYLOAD", err.Error())
}
return item, false, nil
}
func persistKafkaMessageIndividually(ctx context.Context, item *kafkaMessagePersistItem) error {
kind := item.payload.MessageKind
if kind == "" {
kind = "p2p"
}
started := time.Now()
operation := func() error {
err := store.PersistPreparedMessageBatch(item.tenantID, []store.PreparedMessageBatchItem{{
Message: &item.payload.Message, ReadBySender: item.payload.ReadBySender,
}})
if errors.Is(err, types.ErrUnsupported) {
_, err = store.PersistPreparedMessage(item.tenantID, &item.payload.Message, nil, item.payload.ReadBySender)
}
if err != nil && !errors.Is(err, types.ErrDuplicate) {
return err
}
if err = finishKafkaMessagePersist(ctx, item); err != nil {
return err
}
immetrics.KafkaMessagePersistTotal.WithLabelValues(kind, "single_ok").Inc()
immetrics.MessagePersistDuration.WithLabelValues(kind, "single_ok").Observe(time.Since(started).Seconds())
return nil
}
markFailed := func(err error) (*types.KafkaConsumeRecord, error) {
immetrics.KafkaMessagePersistTotal.WithLabelValues(kind, "retry").Inc()
return store.HTTPInfrastructure.MarkKafkaConsumeFailed(item.tenantID, item.envelope.EventID, kafkaBus.config.ConsumerGroup,
kafkaMessagePersistFailureReason(err), err.Error(), time.Now().UTC().Round(time.Millisecond))
}
markDead := func(failed *types.KafkaConsumeRecord, err error) error {
immetrics.KafkaMessagePersistTotal.WithLabelValues(kind, "dead").Inc()
immetrics.MessagePersistDuration.WithLabelValues(kind, "dead").Observe(time.Since(started).Seconds())
reason := kafkaMessagePersistFailureReason(err)
immetrics.KafkaDeadLettersTotal.WithLabelValues(kind, reason).Inc()
return markKafkaDeadLetter(ctx, item.tenantID, item.envelope, failed, reason, err.Error())
}
return runKafkaMessageSingleRetries(ctx, kafkaBus.config.MessageRetries, operation, markFailed, markDead, waitKafkaRetry)
}
func runKafkaMessageSingleRetries(ctx context.Context, retryCount int, operation func() error,
markFailed func(error) (*types.KafkaConsumeRecord, error), markDead func(*types.KafkaConsumeRecord, error) error,
wait kafkaRetryWaiter) error {
var lastErr error
var failed *types.KafkaConsumeRecord
for attempt := 1; attempt <= retryCount+1; attempt++ {
if lastErr = operation(); lastErr == nil {
return nil
}
var err error
if failed, err = markFailed(lastErr); err != nil {
return err
}
if attempt <= retryCount {
if err = wait(ctx, kafkaRetryBackoff(attempt)); err != nil {
return err
}
}
}
return markDead(failed, lastErr)
}
func finishKafkaMessagePersist(ctx context.Context, item *kafkaMessagePersistItem) error {
if len(item.payload.AttachmentURLs) > 0 {
if err := enqueueKafkaFileAttachmentsBind(ctx, item.tenantID, &item.payload.Message, item.payload.AttachmentURLs); err != nil {
return fmt.Errorf("file lifecycle enqueue failed: %w", err)
}
}
return store.HTTPInfrastructure.MarkKafkaConsumeSucceeded(item.tenantID, item.envelope.EventID,
kafkaBus.config.ConsumerGroup, time.Now().UTC().Round(time.Millisecond))
}
func kafkaMessagePersistFailureReason(err error) string {
if err != nil && strings.Contains(err.Error(), "file lifecycle") {
return "KAFKA_FILE_LIFECYCLE_FAILED"
}
return kafkaCommandErrorReason(err)
}
func kafkaMessagePersistTopicName(topic string) string {
if types.GetTopicCat(topic) == types.TopicCatGrp {
return "group_message_persist"
}
return "message_persist"
}
func validateKafkaMessageKind(payload kafkaMessagePersistPayload) error {
if payload.Message.Topic == "" || payload.Message.SeqId <= 0 {
return types.ErrMalformed
}
isGroup := types.GetTopicCat(payload.Message.Topic) == types.TopicCatGrp
switch payload.MessageKind {
case "group":
if !isGroup {
return types.ErrMalformed
}
case "p2p", "":
if isGroup {
return types.ErrMalformed
}
default:
return types.ErrMalformed
}
return nil
}