open-im-server/pkg/common/storage/cache/redis/msg.go

package redis

import (
	"context"
	"github.com/dtm-labs/rockscache"
	"github.com/openimsdk/open-im-server/v3/pkg/common/storage/cache"
	"github.com/openimsdk/open-im-server/v3/pkg/common/storage/cache/cachekey"
	"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/open-im-server/v3/pkg/msgprocessor"
	"github.com/openimsdk/protocol/sdkws"
	"github.com/openimsdk/tools/errs"
	"github.com/openimsdk/tools/utils/datautil"
	"github.com/redis/go-redis/v9"
	"time"
) //

// msgCacheTimeout is  expiration time of message cache, 86400 seconds
const msgCacheTimeout = time.Hour * 24

func NewMsgCache(client redis.UniversalClient) cache.MsgCache {
	return &msgCache{rdb: client}
}

type msgCache struct {
	rdb            redis.UniversalClient
	rcClient       *rockscache.Client
	msgDocDatabase database.Msg
	msgTable       model.MsgDocModel
}

func (c *msgCache) getMessageCacheKey(conversationID string, seq int64) string {
	return cachekey.GetMessageCacheKey(conversationID, seq)
}

func (c *msgCache) getSendMsgKey(id string) string {
	return cachekey.GetSendMsgKey(id)
}

func (c *msgCache) SetMessagesToCache(ctx context.Context, conversationID string, msgs []*sdkws.MsgData) (int, error) {
	msgMap := datautil.SliceToMap(msgs, func(msg *sdkws.MsgData) string {
		return c.getMessageCacheKey(conversationID, msg.Seq)
	})
	keys := datautil.Slice(msgs, func(msg *sdkws.MsgData) string {
		return c.getMessageCacheKey(conversationID, msg.Seq)
	})
	err := ProcessKeysBySlot(ctx, c.rdb, keys, func(ctx context.Context, slot int64, keys []string) error {
		var values []string
		for _, key := range keys {
			if msg, ok := msgMap[key]; ok {
				s, err := msgprocessor.Pb2String(msg)
				if err != nil {
					return err
				}
				values = append(values, s)
			}
		}
		return LuaSetBatchWithCommonExpire(ctx, c.rdb, keys, values, int(msgCacheTimeout/time.Second))
	})
	if err != nil {
		return 0, err
	}
	return len(msgs), nil
}

func (c *msgCache) DeleteMessagesFromCache(ctx context.Context, conversationID string, seqs []int64) error {
	var keys []string
	for _, seq := range seqs {
		keys = append(keys, c.getMessageCacheKey(conversationID, seq))
	}
	return ProcessKeysBySlot(ctx, c.rdb, keys, func(ctx context.Context, slot int64, keys []string) error {
		return LuaDeleteBatch(ctx, c.rdb, keys)
	})
}

func (c *msgCache) SetSendMsgStatus(ctx context.Context, id string, status int32) error {
	return errs.Wrap(c.rdb.Set(ctx, c.getSendMsgKey(id), status, time.Hour*24).Err())
}

func (c *msgCache) GetSendMsgStatus(ctx context.Context, id string) (int32, error) {
	result, err := c.rdb.Get(ctx, c.getSendMsgKey(id)).Int()
	return int32(result), errs.Wrap(err)
}

func (c *msgCache) GetMessagesBySeq(ctx context.Context, conversationID string, seqs []int64) (seqMsgs []*sdkws.MsgData, failedSeqs []int64, err error) {
	var keys []string
	keySeqMap := make(map[string]int64, 10)
	for _, seq := range seqs {
		key := c.getMessageCacheKey(conversationID, seq)
		keys = append(keys, key)
		keySeqMap[key] = seq
	}
	err = ProcessKeysBySlot(ctx, c.rdb, keys, func(ctx context.Context, slot int64, keys []string) error {
		result, err := LuaGetBatch(ctx, c.rdb, keys)
		if err != nil {
			return err
		}
		for i, value := range result {
			seq := keySeqMap[keys[i]]
			if value == nil {
				failedSeqs = append(failedSeqs, seq)
				continue
			}

			msg := &sdkws.MsgData{}
			msgString, ok := value.(string)
			if !ok || msgprocessor.String2Pb(msgString, msg) != nil {
				failedSeqs = append(failedSeqs, seq)
				continue
			}
			seqMsgs = append(seqMsgs, msg)

		}
		return nil
	})
	if err != nil {
		return nil, nil, err
	}
	return seqMsgs, failedSeqs, nil
}

func (c *msgCache) GetMessageBySeqs(ctx context.Context, conversationID string, seqs []int64) ([]*model.MsgInfoModel, error) {
	if len(seqs) == 0 {
		return nil, nil
	}
	getKey := func(seq int64) string {
		return cachekey.GetMessageCacheKeyV2(conversationID, seq)
	}
	getMsgID := func(msg *model.MsgInfoModel) int64 {
		return msg.Msg.Seq
	}
	find := func(ctx context.Context, seqs []int64) ([]*model.MsgInfoModel, error) {
		return c.msgDocDatabase.FindSeqs(ctx, conversationID, seqs)
	}
	return batchGetCache2(ctx, c.rcClient, msgCacheTimeout, seqs, getKey, getMsgID, find)
}

func (c *msgCache) DelMessageBySeqs(ctx context.Context, conversationID string, seqs []int64) error {
	if len(seqs) == 0 {
		return nil
	}
	keys := datautil.Slice(seqs, func(seq int64) string {
		return cachekey.GetMessageCacheKeyV2(conversationID, seq)
	})
	slotKeys, err := groupKeysBySlot(ctx, getRocksCacheRedisClient(c.rcClient), keys)
	if err != nil {
		return err
	}
	for _, keys := range slotKeys {
		if err := c.rcClient.TagAsDeletedBatch2(ctx, keys); err != nil {
			return err
		}
	}
	return nil
}