Fix Kafka message seq idempotency

server/kafka_message_persist.go

@@ -39,16 +39,46 @@ func saveMessageMaybeKafka(ctx context.Context, tenantID types.TenantID, msg *ty
 		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}, nil
+			}
+			return nil, types.ErrDuplicate
+		}
+	}
+	if seq, runtimeReady, seqErr := store.ReserveMessageSeq(tenantID, msg.Topic, msg.SeqId-1); seqErr != nil {
+		if clientMsgID != "" {
+			_ = store.MarkMessageUniqueFailed(tenantID, msg.Topic, fromUid, clientMsgID)
+		}
+		return nil, seqErr
+	} else if runtimeReady {
+		msg.SeqId = seq
+	}
 	msg.InitTimes()
 	if err = enqueueKafkaMessagePersist(ctx, tenantID, msg, attachmentURLs, readBySender); err != nil {
+		if clientMsgID != "" {
+			_ = store.MarkMessageUniqueFailed(tenantID, msg.Topic, fromUid, clientMsgID)
+		}
 		return nil, err
 	}
+	if clientMsgID != "" {
+		if err = store.MarkMessageUniqueDone(tenantID, msg.Topic, fromUid, clientMsgID, msg.SeqId, ""); err != nil {
+			logs.Warn.Printf("kafka message persist: failed to mark message unique key done topic=%s seq=%d: %v", msg.Topic, msg.SeqId, 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
+	return &types.MessageSaveResult{SeqId: msg.SeqId, MarkedReadBySender: readBySender && !fromUid.IsZero()}, nil
 }
 
 func enqueueKafkaMessagePersist(ctx context.Context, tenantID types.TenantID, msg *types.Message, attachmentURLs []string, readBySender bool) error {

server/store/group_message_runtime.go

@@ -39,6 +39,11 @@ type groupMessageUniqueRecord struct {
 	CreatedAt time.Time `json:"created_at"`
 }
 
+type MessageUniqueClaim struct {
+	Status string
+	Seq    int64
+}
+
 type groupFanoutEnvelope struct {
 	OriginNodeID string          `json:"origin_node_id"`
 	Message      json.RawMessage `json:"message"`
@@ -76,6 +81,39 @@ type groupMessageRuntime interface {
 
 var groupMsgRuntime groupMessageRuntime
 
+func ClaimMessageUnique(tenantID types.TenantID, topic string, from types.Uid, clientMsgID string) (*MessageUniqueClaim, bool, bool, error) {
+	if groupMsgRuntime == nil || !tenantID.IsValid() || topic == "" || from.IsZero() || clientMsgID == "" {
+		return nil, true, false, nil
+	}
+	rec, claimed, err := groupMsgRuntime.Claim(tenantID, topic, from, clientMsgID)
+	if rec == nil {
+		return nil, claimed, true, err
+	}
+	return &MessageUniqueClaim{Status: rec.Status, Seq: rec.Seq}, claimed, true, err
+}
+
+func ReserveMessageSeq(tenantID types.TenantID, topic string, baseline int64) (int64, bool, error) {
+	if groupMsgRuntime == nil || !tenantID.IsValid() || topic == "" {
+		return 0, false, nil
+	}
+	seq, err := groupMsgRuntime.NextSeq(tenantID, topic, baseline)
+	return seq, true, err
+}
+
+func MarkMessageUniqueDone(tenantID types.TenantID, topic string, from types.Uid, clientMsgID string, seq int64, msgID string) error {
+	if groupMsgRuntime == nil || !tenantID.IsValid() || topic == "" || from.IsZero() || clientMsgID == "" {
+		return nil
+	}
+	return groupMsgRuntime.MarkDone(tenantID, topic, from, clientMsgID, seq, msgID)
+}
+
+func MarkMessageUniqueFailed(tenantID types.TenantID, topic string, from types.Uid, clientMsgID string) error {
+	if groupMsgRuntime == nil || !tenantID.IsValid() || topic == "" || from.IsZero() || clientMsgID == "" {
+		return nil
+	}
+	return groupMsgRuntime.MarkFailed(tenantID, topic, from, clientMsgID)
+}
+
 type UserPresence struct {
 	Online       bool      `json:"online"`
 	NodeID       string    `json:"node_id"`

server/store/group_message_runtime_test.go

@@ -0,0 +1,178 @@
+package store
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/tinode/chat/server/store/types"
+)
+
+type testGroupMessageRuntime struct {
+	claims map[string]*groupMessageUniqueRecord
+	seqs   map[string]int64
+}
+
+func newTestGroupMessageRuntime() *testGroupMessageRuntime {
+	return &testGroupMessageRuntime{
+		claims: make(map[string]*groupMessageUniqueRecord),
+		seqs:   make(map[string]int64),
+	}
+}
+
+func (r *testGroupMessageRuntime) Close() error { return nil }
+
+func (r *testGroupMessageRuntime) Claim(tenantID types.TenantID, topic string, from types.Uid, clientMsgID string) (*groupMessageUniqueRecord, bool, error) {
+	key := topic + ":" + from.UserId() + ":" + clientMsgID
+	if rec := r.claims[key]; rec != nil {
+		return rec, false, nil
+	}
+	rec := &groupMessageUniqueRecord{Status: "pending", CreatedAt: time.Now().UTC()}
+	r.claims[key] = rec
+	return rec, true, nil
+}
+
+func (r *testGroupMessageRuntime) NextSeq(_ types.TenantID, topic string, baseline int64) (int64, error) {
+	if r.seqs[topic] < baseline {
+		r.seqs[topic] = baseline
+	}
+	r.seqs[topic]++
+	return r.seqs[topic], nil
+}
+
+func (r *testGroupMessageRuntime) MarkDone(_ types.TenantID, topic string, from types.Uid, clientMsgID string, seq int64, msgID string) error {
+	key := topic + ":" + from.UserId() + ":" + clientMsgID
+	r.claims[key] = &groupMessageUniqueRecord{Status: "done", Seq: seq, MessageID: msgID, CreatedAt: time.Now().UTC()}
+	return nil
+}
+
+func (r *testGroupMessageRuntime) MarkFailed(_ types.TenantID, topic string, from types.Uid, clientMsgID string) error {
+	key := topic + ":" + from.UserId() + ":" + clientMsgID
+	r.claims[key] = &groupMessageUniqueRecord{Status: "failed", CreatedAt: time.Now().UTC()}
+	return nil
+}
+
+func (r *testGroupMessageRuntime) CacheHotMessage(types.TenantID, *types.Message) error {
+	return nil
+}
+
+func (r *testGroupMessageRuntime) CachePendingHotMessage(types.TenantID, *types.Message) error {
+	return nil
+}
+
+func (r *testGroupMessageRuntime) GetHotMessages(types.TenantID, string, *types.QueryOpt, int) ([]types.Message, bool, error) {
+	return nil, false, nil
+}
+
+func (r *testGroupMessageRuntime) InvalidateHotMessages(types.TenantID, string, []types.Range) error {
+	return nil
+}
+
+func (r *testGroupMessageRuntime) MarkUserOnline(types.TenantID, types.Uid, string, string, string, string) error {
+	return nil
+}
+
+func (r *testGroupMessageRuntime) RefreshUserPresence(types.TenantID, types.Uid, string, string) error {
+	return nil
+}
+
+func (r *testGroupMessageRuntime) MarkUserOffline(types.TenantID, types.Uid, string, string) error {
+	return nil
+}
+
+func (r *testGroupMessageRuntime) GetUserPresence(types.TenantID, types.Uid) (*UserPresence, error) {
+	return nil, nil
+}
+
+func (r *testGroupMessageRuntime) MarkTopicMemberOnline(types.TenantID, string, types.Uid, string) error {
+	return nil
+}
+
+func (r *testGroupMessageRuntime) MarkTopicMemberOffline(types.TenantID, string, types.Uid, string) error {
+	return nil
+}
+
+func (r *testGroupMessageRuntime) GetTopicNodeOnlineMembers(types.TenantID, string, string) ([]types.Uid, error) {
+	return nil, nil
+}
+
+func (r *testGroupMessageRuntime) PublishGroupFanout(string, any) error {
+	return nil
+}
+
+func (r *testGroupMessageRuntime) StartGroupFanoutSubscriber(context.Context, string, func([]byte)) error {
+	return nil
+}
+
+func (r *testGroupMessageRuntime) CacheObject(types.TenantID, string, string, any) error {
+	return nil
+}
+
+func (r *testGroupMessageRuntime) GetCachedObject(types.TenantID, string, string, any) (bool, error) {
+	return false, nil
+}
+
+func (r *testGroupMessageRuntime) DeleteCachedObject(types.TenantID, string, string) error {
+	return nil
+}
+
+func (r *testGroupMessageRuntime) AddCachedSetMember(types.TenantID, string, string, string) error {
+	return nil
+}
+
+func (r *testGroupMessageRuntime) RemoveCachedSetMember(types.TenantID, string, string, string) error {
+	return nil
+}
+
+func (r *testGroupMessageRuntime) CachedSetContains(types.TenantID, string, string, string) (bool, error) {
+	return false, nil
+}
+
+func (r *testGroupMessageRuntime) AllowRate(types.TenantID, string, int, time.Duration) (bool, int, time.Duration, error) {
+	return true, 0, 0, nil
+}
+
+func (r *testGroupMessageRuntime) ClaimShortState(string, string, time.Duration) (bool, error) {
+	return true, nil
+}
+
+func (r *testGroupMessageRuntime) ClaimShortStateValue(string, string, string, time.Duration) (bool, bool, error) {
+	return true, true, nil
+}
+
+func TestMessageRuntimeSeqAndClientMidIdempotency(t *testing.T) {
+	oldRuntime := groupMsgRuntime
+	runtime := newTestGroupMessageRuntime()
+	groupMsgRuntime = runtime
+	t.Cleanup(func() { groupMsgRuntime = oldRuntime })
+
+	tenantID := types.TenantID(7)
+	topic := types.Uid(1).P2PName(types.Uid(2))
+	from := types.Uid(1)
+
+	claim, claimed, ready, err := ClaimMessageUnique(tenantID, topic, from, "mid-1")
+	if err != nil || !ready || !claimed || claim == nil || claim.Status != "pending" {
+		t.Fatalf("unexpected first claim: claim=%#v claimed=%v ready=%v err=%v", claim, claimed, ready, err)
+	}
+	seq, ready, err := ReserveMessageSeq(tenantID, topic, 184)
+	if err != nil || !ready || seq != 185 {
+		t.Fatalf("unexpected first seq: seq=%d ready=%v err=%v", seq, ready, err)
+	}
+	if err = MarkMessageUniqueDone(tenantID, topic, from, "mid-1", seq, ""); err != nil {
+		t.Fatalf("mark done: %v", err)
+	}
+
+	claim, claimed, ready, err = ClaimMessageUnique(tenantID, topic, from, "mid-1")
+	if err != nil || !ready || claimed || claim == nil || claim.Seq != 185 {
+		t.Fatalf("unexpected duplicate claim: claim=%#v claimed=%v ready=%v err=%v", claim, claimed, ready, err)
+	}
+
+	claim, claimed, ready, err = ClaimMessageUnique(tenantID, topic, from, "mid-2")
+	if err != nil || !ready || !claimed || claim == nil {
+		t.Fatalf("unexpected second claim: claim=%#v claimed=%v ready=%v err=%v", claim, claimed, ready, err)
+	}
+	seq, ready, err = ReserveMessageSeq(tenantID, topic, 184)
+	if err != nil || !ready || seq != 186 {
+		t.Fatalf("unexpected second seq: seq=%d ready=%v err=%v", seq, ready, err)
+	}
+}