feat: complete shared database tenant sharding

2026-09-23 02:19:07 +08:00
parent 1ec2feee47
commit e37b52c3c0
37 changed files with 946 additions and 1420 deletions
@@ -1,6 +1,10 @@
 package store

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

 	adapter "github.com/tinode/chat/server/db"
@@ -8,9 +12,6 @@ import (
 )

 type HTTPInfrastructurePersistenceInterface interface {
-	BeginIdempotency(types.TenantID, *types.HTTPIdempotencyRecord) (*types.HTTPIdempotencyRecord, bool, error)
-	CompleteIdempotency(types.TenantID, *types.HTTPIdempotencyRecord) error
-	AbortIdempotency(types.TenantID, *types.HTTPIdempotencyRecord) error
 	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)
@@ -25,6 +26,49 @@ 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 {
@@ -37,52 +81,109 @@ func tenantHTTPInfrastructureAdapter(tenantID types.TenantID) (adapter.TenantHTT
 	return infrastructure, nil
 }

-func (httpInfrastructureMapper) BeginIdempotency(tenantID types.TenantID, record *types.HTTPIdempotencyRecord) (*types.HTTPIdempotencyRecord, bool, error) {
-	infrastructure, err := tenantHTTPInfrastructureAdapter(tenantID)
-	if err != nil {
-		return nil, false, err
-	}
-	return infrastructure.TenantHTTPIdempotencyBegin(tenantID, record)
-}
-
-func (httpInfrastructureMapper) CompleteIdempotency(tenantID types.TenantID, record *types.HTTPIdempotencyRecord) error {
-	infrastructure, err := tenantHTTPInfrastructureAdapter(tenantID)
-	if err != nil {
-		return err
-	}
-	return infrastructure.TenantHTTPIdempotencyComplete(tenantID, record)
-}
-
-func (httpInfrastructureMapper) AbortIdempotency(tenantID types.TenantID, record *types.HTTPIdempotencyRecord) error {
-	infrastructure, err := tenantHTTPInfrastructureAdapter(tenantID)
-	if err != nil {
-		return err
-	}
-	return infrastructure.TenantHTTPIdempotencyAbort(tenantID, record)
-}
-
 func (httpInfrastructureMapper) BeginKafkaConsume(tenantID types.TenantID, record *types.KafkaConsumeRecord, deadline time.Time) (*types.KafkaConsumeRecord, bool, error) {
-	infrastructure, err := tenantHTTPInfrastructureAdapter(tenantID)
-	if err != nil {
+	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
 	}
-	return infrastructure.TenantKafkaConsumeBegin(tenantID, record, deadline)
+	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 {
-	infrastructure, err := tenantHTTPInfrastructureAdapter(tenantID)
-	if err != nil {
+	record, found, enabled, err := kafkaConsumeRead(tenantID, eventID, consumerGroup)
+	if err != nil || !enabled {
+		if err == nil {
+			err = errKafkaConsumeRedisUnavailable
+		}
 		return err
 	}
-	return infrastructure.TenantKafkaConsumeSucceed(tenantID, eventID, consumerGroup, now)
+	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) {
-	infrastructure, err := tenantHTTPInfrastructureAdapter(tenantID)
-	if err != nil {
+	record, found, enabled, err := kafkaConsumeRead(tenantID, eventID, consumerGroup)
+	if err != nil || !enabled {
+		if err == nil {
+			err = errKafkaConsumeRedisUnavailable
+		}
 		return nil, err
 	}
-	return infrastructure.TenantKafkaConsumeFail(tenantID, eventID, consumerGroup, reason, message, now)
+	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 {
@@ -90,7 +191,18 @@ func (httpInfrastructureMapper) MarkKafkaConsumeDead(tenantID types.TenantID, de
 	if err != nil {
 		return err
 	}
-	return infrastructure.TenantKafkaConsumeDead(tenantID, dead, now)
+	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) {
@@ -122,5 +234,11 @@ func (httpInfrastructureMapper) MarkKafkaDeadLetterReplay(tenantID types.TenantI
 	if err != nil {
 		return nil, err
 	}
-	return infrastructure.TenantKafkaDeadLetterMarkReplay(tenantID, deadID, resolvedBy, note, now)
+	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
 }