im_server/server/store/infrastructure.go

package store

import (
	"encoding/json"
	"errors"
	"strconv"
	"strings"
	"time"

	adapter "github.com/tinode/chat/server/db"
	"github.com/tinode/chat/server/store/types"
)

type HTTPInfrastructurePersistenceInterface interface {
	BeginKafkaConsume(types.TenantID, *types.KafkaConsumeRecord, time.Time) (*types.KafkaConsumeRecord, bool, error)
	MarkKafkaConsumeSucceeded(types.TenantID, string, string, time.Time) error
	MarkKafkaConsumeFailed(types.TenantID, string, string, string, string, time.Time) (*types.KafkaConsumeRecord, error)
	MarkKafkaConsumeDead(types.TenantID, *types.KafkaDeadLetter, time.Time) error
	ListKafkaDeadLetters(types.TenantID, types.KafkaDeadLetterListOptions) (*types.KafkaDeadLetterPage, error)
	GetKafkaDeadLetter(types.TenantID, string) (*types.KafkaDeadLetter, error)
	MarkKafkaDeadLetter(types.TenantID, string, string, string, string, time.Time) error
	MarkKafkaDeadLetterReplay(types.TenantID, string, string, string, time.Time) (*types.KafkaDeadLetter, error)
}

type httpInfrastructureMapper struct{}

var HTTPInfrastructure HTTPInfrastructurePersistenceInterface = httpInfrastructureMapper{}

const (
	kafkaConsumeNamespace = "kafka_consume"
	kafkaConsumeRecordTTL = 7 * 24 * time.Hour
)

var errKafkaConsumeRedisUnavailable = errors.New("kafka consume redis unavailable")

func kafkaConsumeKey(tenantID types.TenantID, eventID, consumerGroup string) string {
	return strconv.FormatInt(int64(tenantID), 10) + ":" + strings.TrimSpace(eventID) + ":" + strings.TrimSpace(consumerGroup)
}

func kafkaConsumeRecordKey(tenantID types.TenantID, eventID, consumerGroup string) string {
	return kafkaConsumeKey(tenantID, eventID, consumerGroup) + ":record"
}

func kafkaConsumeRead(tenantID types.TenantID, eventID, consumerGroup string) (*types.KafkaConsumeRecord, bool, bool, error) {
	raw, found, enabled, err := GetShortStateValue(kafkaConsumeNamespace, kafkaConsumeRecordKey(tenantID, eventID, consumerGroup))
	if err != nil || !enabled || !found {
		return nil, found, enabled, err
	}
	var record types.KafkaConsumeRecord
	if err = json.Unmarshal([]byte(raw), &record); err != nil {
		_, _ = DeleteShortState(kafkaConsumeNamespace, kafkaConsumeRecordKey(tenantID, eventID, consumerGroup))
		return nil, false, true, err
	}
	return &record, true, true, nil
}

func kafkaConsumeWrite(tenantID types.TenantID, record *types.KafkaConsumeRecord) error {
	if record == nil {
		return types.ErrMalformed
	}
	payload, err := json.Marshal(record)
	if err != nil {
		return err
	}
	enabled, err := SetShortStateValue(kafkaConsumeNamespace, kafkaConsumeRecordKey(tenantID, record.EventID, record.ConsumerGroup), string(payload), kafkaConsumeRecordTTL)
	if !enabled {
		return errKafkaConsumeRedisUnavailable
	}
	return err
}

func tenantHTTPInfrastructureAdapter(tenantID types.TenantID) (adapter.TenantHTTPInfrastructureAdapter, error) {
	business, err := tenantBusinessAdapter(tenantID)
	if err != nil {
		return nil, err
	}
	infrastructure, ok := business.(adapter.TenantHTTPInfrastructureAdapter)
	if !ok {
		return nil, types.ErrUnsupported
	}
	return infrastructure, nil
}

func (httpInfrastructureMapper) BeginKafkaConsume(tenantID types.TenantID, record *types.KafkaConsumeRecord, deadline time.Time) (*types.KafkaConsumeRecord, bool, error) {
	if record == nil || strings.TrimSpace(record.EventID) == "" || strings.TrimSpace(record.ConsumerGroup) == "" || strings.TrimSpace(record.SourceTopic) == "" {
		return nil, false, types.ErrMalformed
	}
	if existing, found, enabled, err := kafkaConsumeRead(tenantID, record.EventID, record.ConsumerGroup); err != nil || !enabled {
		if err == nil {
			err = errKafkaConsumeRedisUnavailable
		}
		return nil, false, err
	} else if found && (existing.Status == types.KafkaConsumeSucceeded || existing.Status == types.KafkaConsumeDead) {
		return existing, false, nil
	} else if found {
		record.RetryCount = existing.RetryCount
		record.LastReason = existing.LastReason
		record.LastError = existing.LastError
		record.CreatedAt = existing.CreatedAt
	}
	ttl := time.Until(deadline)
	if ttl <= 0 {
		ttl = 5 * time.Minute
	}
	claimed, _, enabled, err := ClaimShortStateValue(kafkaConsumeNamespace, kafkaConsumeKey(tenantID, record.EventID, record.ConsumerGroup), types.KafkaConsumeProcessing, ttl)
	if err != nil || !enabled {
		if err == nil {
			err = errKafkaConsumeRedisUnavailable
		}
		return nil, false, err
	}
	if !claimed {
		existing, found, enabled, err := kafkaConsumeRead(tenantID, record.EventID, record.ConsumerGroup)
		if err != nil || !enabled {
			if err == nil {
				err = errKafkaConsumeRedisUnavailable
			}
			return nil, false, err
		}
		if found {
			return existing, false, nil
		}
		copy := *record
		copy.Status = types.KafkaConsumeProcessing
		copy.ProcessingDeadline = &deadline
		return &copy, false, nil
	}
	now := time.Now().UTC().Round(time.Millisecond)
	copy := *record
	copy.Status = types.KafkaConsumeProcessing
	copy.ProcessingDeadline = &deadline
	if copy.CreatedAt.IsZero() {
		copy.CreatedAt = now
	}
	copy.UpdatedAt = now
	if err = kafkaConsumeWrite(tenantID, &copy); err != nil {
		_, _ = DeleteShortState(kafkaConsumeNamespace, kafkaConsumeKey(tenantID, record.EventID, record.ConsumerGroup))
		return nil, false, err
	}
	return &copy, true, nil
}

func (httpInfrastructureMapper) MarkKafkaConsumeSucceeded(tenantID types.TenantID, eventID, consumerGroup string, now time.Time) error {
	record, found, enabled, err := kafkaConsumeRead(tenantID, eventID, consumerGroup)
	if err != nil || !enabled {
		if err == nil {
			err = errKafkaConsumeRedisUnavailable
		}
		return err
	}
	if !found {
		return types.ErrNotFound
	}
	record.Status = types.KafkaConsumeSucceeded
	record.SucceededAt = &now
	record.ProcessingDeadline = nil
	record.UpdatedAt = now
	if err = kafkaConsumeWrite(tenantID, record); err != nil {
		return err
	}
	_, err = DeleteShortState(kafkaConsumeNamespace, kafkaConsumeKey(tenantID, eventID, consumerGroup))
	return err
}

func (httpInfrastructureMapper) MarkKafkaConsumeFailed(tenantID types.TenantID, eventID, consumerGroup, reason, message string, now time.Time) (*types.KafkaConsumeRecord, error) {
	record, found, enabled, err := kafkaConsumeRead(tenantID, eventID, consumerGroup)
	if err != nil || !enabled {
		if err == nil {
			err = errKafkaConsumeRedisUnavailable
		}
		return nil, err
	}
	if !found {
		return nil, types.ErrNotFound
	}
	record.Status = types.KafkaConsumeFailed
	record.RetryCount++
	record.LastReason = reason
	record.LastError = message
	record.ProcessingDeadline = nil
	record.UpdatedAt = now
	if err = kafkaConsumeWrite(tenantID, record); err != nil {
		return nil, err
	}
	_, err = DeleteShortState(kafkaConsumeNamespace, kafkaConsumeKey(tenantID, eventID, consumerGroup))
	return record, err
}

func (httpInfrastructureMapper) MarkKafkaConsumeDead(tenantID types.TenantID, dead *types.KafkaDeadLetter, now time.Time) error {
	infrastructure, err := tenantHTTPInfrastructureAdapter(tenantID)
	if err != nil {
		return err
	}
	if err = infrastructure.TenantKafkaConsumeDead(tenantID, dead, now); err != nil {
		return err
	}
	record := &types.KafkaConsumeRecord{
		EventID: dead.EventID, ConsumerGroup: dead.ConsumerGroup, SourceTopic: dead.SourceTopic,
		EventType: dead.EventType, BizType: dead.BizType, BizID: dead.BizID,
		Status: types.KafkaConsumeDead, RetryCount: dead.RetryCount, LastReason: dead.LastReason,
		LastError: dead.LastError, DeadAt: &now, CreatedAt: now, UpdatedAt: now,
	}
	_ = kafkaConsumeWrite(tenantID, record)
	_, _ = DeleteShortState(kafkaConsumeNamespace, kafkaConsumeKey(tenantID, dead.EventID, dead.ConsumerGroup))
	return nil
}

func (httpInfrastructureMapper) ListKafkaDeadLetters(tenantID types.TenantID, opts types.KafkaDeadLetterListOptions) (*types.KafkaDeadLetterPage, error) {
	infrastructure, err := tenantHTTPInfrastructureAdapter(tenantID)
	if err != nil {
		return nil, err
	}
	return infrastructure.TenantKafkaDeadLetterList(tenantID, opts)
}

func (httpInfrastructureMapper) GetKafkaDeadLetter(tenantID types.TenantID, deadID string) (*types.KafkaDeadLetter, error) {
	infrastructure, err := tenantHTTPInfrastructureAdapter(tenantID)
	if err != nil {
		return nil, err
	}
	return infrastructure.TenantKafkaDeadLetterGet(tenantID, deadID)
}

func (httpInfrastructureMapper) MarkKafkaDeadLetter(tenantID types.TenantID, deadID, status, resolvedBy, note string, now time.Time) error {
	infrastructure, err := tenantHTTPInfrastructureAdapter(tenantID)
	if err != nil {
		return err
	}
	return infrastructure.TenantKafkaDeadLetterMark(tenantID, deadID, status, resolvedBy, note, now)
}

func (httpInfrastructureMapper) MarkKafkaDeadLetterReplay(tenantID types.TenantID, deadID, resolvedBy, note string, now time.Time) (*types.KafkaDeadLetter, error) {
	infrastructure, err := tenantHTTPInfrastructureAdapter(tenantID)
	if err != nil {
		return nil, err
	}
	dead, err := infrastructure.TenantKafkaDeadLetterMarkReplay(tenantID, deadID, resolvedBy, note, now)
	if err != nil {
		return nil, err
	}
	_, _ = DeleteShortState(kafkaConsumeNamespace, kafkaConsumeKey(tenantID, dead.EventID, dead.ConsumerGroup))
	_, _ = DeleteShortState(kafkaConsumeNamespace, kafkaConsumeRecordKey(tenantID, dead.EventID, dead.ConsumerGroup))
	return dead, nil
}