im_server/server/kafka_message_persist.go

package main

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"time"

	"github.com/segmentio/kafka-go"
	"github.com/tinode/chat/server/logs"
	"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"`
}

var publishKafkaMessagePersist = publishKafkaEnvelope

func startKafkaMessagePersistConsumer() {
	startKafkaConsumer("message_persist", handleKafkaMessagePersist)
}

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
	}
	clientMsgID, err := store.NormalizeClientMsgID(msg.Head)
	if err != nil {
		return nil, err
	}
	msg.ClientMsgID = clientMsgID
	msg.InitTimes()
	if err = enqueueKafkaMessagePersist(ctx, tenantID, msg, attachmentURLs, readBySender); err != nil {
		return nil, err
	}
	if types.GetTopicCat(msg.Topic) == types.TopicCatGrp {
		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, MarkedReadBySender: readBySender && !types.ParseUid(msg.From).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,
	})
	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, "message_persist", msg.Topic, envelope)
}

func handleKafkaMessagePersist(ctx context.Context, envelope kafkaEnvelope, _ kafka.Message) error {
	tenantID := envelope.TenantID
	var err error
	if tenantID.IsZero() {
		tenantID, err = kafkaTenantID(envelope.TenantCode)
		if err != nil {
			return err
		}
	}
	consume, claimed, err := beginKafkaConsumeForTopic(tenantID, envelope, "message_persist")
	if err != nil || !claimed {
		return err
	}
	var payload kafkaMessagePersistPayload
	if err = json.Unmarshal(envelope.Payload, &payload); err != nil {
		return markKafkaDeadLetter(ctx, tenantID, envelope, consume, "MALFORMED_PAYLOAD", err.Error())
	}
	if payload.Message.TenantID == types.ZeroTenantID {
		payload.Message.TenantID = tenantID
	}
	if payload.Message.TenantID != tenantID {
		return markKafkaDeadLetter(ctx, tenantID, envelope, consume, "TENANT_MISMATCH", "message tenant does not match envelope tenant")
	}
	if err = validateKafkaMessageKind(payload); err != nil {
		return markKafkaDeadLetter(ctx, tenantID, envelope, consume, "MALFORMED_PAYLOAD", err.Error())
	}
	if _, err = store.PersistPreparedMessage(tenantID, &payload.Message, nil, payload.ReadBySender); err != nil && !errors.Is(err, types.ErrDuplicate) {
		failed, markErr := store.HTTPInfrastructure.MarkKafkaConsumeFailed(tenantID, envelope.EventID, kafkaBus.config.ConsumerGroup,
			kafkaCommandErrorReason(err), err.Error(), time.Now().UTC().Round(time.Millisecond))
		if markErr != nil {
			return markErr
		}
		if failed != nil && failed.RetryCount >= kafkaBus.config.MaxRetryCount {
			return markKafkaDeadLetter(ctx, tenantID, envelope, failed, kafkaCommandErrorReason(err), err.Error())
		}
		sleepKafkaRetryBackoff(failed.RetryCount)
		return err
	}
	if len(payload.AttachmentURLs) > 0 {
		if err = enqueueKafkaFileAttachmentsBind(ctx, tenantID, &payload.Message, payload.AttachmentURLs); err != nil {
			failed, markErr := store.HTTPInfrastructure.MarkKafkaConsumeFailed(tenantID, envelope.EventID, kafkaBus.config.ConsumerGroup,
				"KAFKA_FILE_LIFECYCLE_FAILED", err.Error(), time.Now().UTC().Round(time.Millisecond))
			if markErr != nil {
				return markErr
			}
			if failed != nil && failed.RetryCount >= kafkaBus.config.MaxRetryCount {
				return markKafkaDeadLetter(ctx, tenantID, envelope, failed, "KAFKA_FILE_LIFECYCLE_FAILED", err.Error())
			}
			sleepKafkaRetryBackoff(failed.RetryCount)
			return err
		}
	}
	return store.HTTPInfrastructure.MarkKafkaConsumeSucceeded(tenantID, envelope.EventID, kafkaBus.config.ConsumerGroup, time.Now().UTC().Round(time.Millisecond))
}

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
}