open-im-server/pkg/common/storage/database/mgo/seq_conversation.go

package mgo

import (
	"context"
	"errors"
	"github.com/openimsdk/open-im-server/v3/pkg/common/storage/database"
	"github.com/openimsdk/open-im-server/v3/pkg/common/storage/model"
	"github.com/openimsdk/tools/db/mongoutil"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

func NewSeqConversationMongo(db *mongo.Database) (database.SeqConversation, error) {
	coll := db.Collection(database.SeqConversationName)
	_, err := coll.Indexes().CreateOne(context.Background(), mongo.IndexModel{
		Keys: bson.D{
			{Key: "conversation_id", Value: 1},
		},
	})
	if err != nil {
		return nil, err
	}
	return &seqConversationMongo{coll: coll}, nil
}

type seqConversationMongo struct {
	coll *mongo.Collection
}

func (s *seqConversationMongo) setSeq(ctx context.Context, conversationID string, seq int64, field string) error {
	filter := map[string]any{
		"conversation_id": conversationID,
	}
	insert := bson.M{
		"conversation_id": conversationID,
		"min_seq":         0,
		"max_seq":         0,
	}
	delete(insert, field)
	update := map[string]any{
		"$set": bson.M{
			field: seq,
		},
		"$setOnInsert": insert,
	}
	opt := options.Update().SetUpsert(true)
	return mongoutil.UpdateOne(ctx, s.coll, filter, update, false, opt)
}

func (s *seqConversationMongo) Malloc(ctx context.Context, conversationID string, size int64) (int64, error) {
	if size < 0 {
		return 0, errors.New("size must be greater than 0")
	}
	if size == 0 {
		return s.GetMaxSeq(ctx, conversationID)
	}
	filter := map[string]any{"conversation_id": conversationID}
	update := map[string]any{
		"$inc": map[string]any{"max_seq": size},
		"$set": map[string]any{"min_seq": int64(0)},
	}
	opt := options.FindOneAndUpdate().SetUpsert(true).SetReturnDocument(options.After).SetProjection(map[string]any{"_id": 0, "max_seq": 1})
	lastSeq, err := mongoutil.FindOneAndUpdate[int64](ctx, s.coll, filter, update, opt)
	if err != nil {
		return 0, err
	}
	return lastSeq - size, nil
}

func (s *seqConversationMongo) SetMaxSeq(ctx context.Context, conversationID string, seq int64) error {
	return s.setSeq(ctx, conversationID, seq, "max_seq")
}

func (s *seqConversationMongo) GetMaxSeq(ctx context.Context, conversationID string) (int64, error) {
	seq, err := mongoutil.FindOne[int64](ctx, s.coll, bson.M{"conversation_id": conversationID}, options.FindOne().SetProjection(map[string]any{"_id": 0, "max_seq": 1}))
	if err == nil {
		return seq, nil
	} else if IsNotFound(err) {
		return 0, nil
	} else {
		return 0, err
	}
}

func (s *seqConversationMongo) GetMinSeq(ctx context.Context, conversationID string) (int64, error) {
	seq, err := mongoutil.FindOne[int64](ctx, s.coll, bson.M{"conversation_id": conversationID}, options.FindOne().SetProjection(map[string]any{"_id": 0, "min_seq": 1}))
	if err == nil {
		return seq, nil
	} else if IsNotFound(err) {
		return 0, nil
	} else {
		return 0, err
	}
}

func (s *seqConversationMongo) SetMinSeq(ctx context.Context, conversationID string, seq int64) error {
	return s.setSeq(ctx, conversationID, seq, "min_seq")
}

func (s *seqConversationMongo) GetConversation(ctx context.Context, conversationID string) (*model.SeqConversation, error) {
	return mongoutil.FindOne[*model.SeqConversation](ctx, s.coll, bson.M{"conversation_id": conversationID})
}