refactor: db cache batch refactor and batch consume message. (#2325)
* refactor: cmd update. * refactor: msg transfer refactor. * refactor: msg transfer refactor. * refactor: msg transfer refactor. * fix: read prometheus port when flag set to enable and prevent failure during startup. * fix: notification has counted unread counts bug fix. * fix: merge opensource code into local. * refactor: delete message and message batch use lua. * refactor: delete message and message batch use lua. * refactor: delete message and message batch use lua. * refactor: delete message and message batch use lua. * refactor: delete message and message batch use lua. * refactor: delete message and message batch use lua. * refactor: delete message and message batch use lua. * refactor: delete message and message batch use lua. * refactor: delete message and message batch use lua. * refactor: delete message and message batch use lua. * refactor: delete message and message batch use lua. * refactor: delete message and message batch use lua. * refactor: delete message and message batch use lua. * refactor: delete message and message batch use lua. * refactor: delete message and message batch use lua. * refactor: delete message and message batch use lua. * refactor: delete message and message batch use lua. * refactor: delete message and message batch use lua. * refactor: delete message and message batch use lua. * refactor: delete message and message batch use lua. * refactor: delete message and message batch use lua. * refactor: delete message and message batch use lua. * refactor: delete message and message batch use lua. * refactor: delete message and message batch use lua. * refactor: delete message and message batch use lua. * refactor: delete message and message batch use lua. * refactor: delete message and message batch use lua. * refactor: delete message and message batch use lua. * refactor: delete message and message batch use lua. * refactor: delete message and message batch use lua. * fix: add protective measures against memory overflow.
This commit is contained in:
OpenIM-Gordon
@@ -323,13 +323,12 @@ type User struct {
|
||||
}
|
||||
|
||||
type Redis struct {
|
||||
Address []string `mapstructure:"address"`
|
||||
Username string `mapstructure:"username"`
|
||||
Password string `mapstructure:"password"`
|
||||
EnablePipeline bool `mapstructure:"enablePipeline"`
|
||||
ClusterMode bool `mapstructure:"clusterMode"`
|
||||
DB int `mapstructure:"storage"`
|
||||
MaxRetry int `mapstructure:"MaxRetry"`
|
||||
Address []string `mapstructure:"address"`
|
||||
Username string `mapstructure:"username"`
|
||||
Password string `mapstructure:"password"`
|
||||
ClusterMode bool `mapstructure:"clusterMode"`
|
||||
DB int `mapstructure:"storage"`
|
||||
MaxRetry int `mapstructure:"MaxRetry"`
|
||||
}
|
||||
|
||||
type BeforeConfig struct {
|
||||
|
||||
@@ -52,12 +52,9 @@ func Start[T any](ctx context.Context, discovery *config2.Discovery, prometheusC
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
prometheusPort, err := datautil.GetElemByIndex(prometheusConfig.Ports, index)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
log.CInfo(ctx, "RPC server is initializing", "rpcRegisterName", rpcRegisterName, "rpcPort", rpcPort,
|
||||
"prometheusPort", prometheusPort)
|
||||
"prometheusPorts", prometheusConfig.Ports)
|
||||
rpcTcpAddr := net.JoinHostPort(network.GetListenIP(listenIP), strconv.Itoa(rpcPort))
|
||||
listener, err := net.Listen(
|
||||
"tcp",
|
||||
@@ -117,9 +114,14 @@ func Start[T any](ctx context.Context, discovery *config2.Discovery, prometheusC
|
||||
netErr error
|
||||
httpServer *http.Server
|
||||
)
|
||||
|
||||
go func() {
|
||||
if prometheusConfig.Enable && prometheusPort != 0 {
|
||||
if prometheusConfig.Enable {
|
||||
go func() {
|
||||
prometheusPort, err := datautil.GetElemByIndex(prometheusConfig.Ports, index)
|
||||
if err != nil {
|
||||
netErr = err
|
||||
netDone <- struct{}{}
|
||||
return
|
||||
}
|
||||
metric.InitializeMetrics(srv)
|
||||
// Create a HTTP server for prometheus.
|
||||
httpServer = &http.Server{Handler: promhttp.HandlerFor(reg, promhttp.HandlerOpts{}), Addr: fmt.Sprintf("0.0.0.0:%d", prometheusPort)}
|
||||
@@ -127,8 +129,8 @@ func Start[T any](ctx context.Context, discovery *config2.Discovery, prometheusC
|
||||
netErr = errs.WrapMsg(err, "prometheus start err", httpServer.Addr)
|
||||
netDone <- struct{}{}
|
||||
}
|
||||
}
|
||||
}()
|
||||
}()
|
||||
}
|
||||
|
||||
go func() {
|
||||
err := srv.Serve(listener)
|
||||
|
||||
-4
@@ -31,10 +31,6 @@ const (
|
||||
reactionNotification = "EX_NOTIFICATION_"
|
||||
)
|
||||
|
||||
func GetAllMessageCacheKey(conversationID string) string {
|
||||
return messageCache + conversationID + "_*"
|
||||
}
|
||||
|
||||
func GetMessageCacheKey(conversationID string, seq int64) string {
|
||||
return messageCache + conversationID + "_" + strconv.Itoa(int(seq))
|
||||
}
|
||||
|
||||
-3
@@ -52,9 +52,6 @@ type ConversationCache interface {
|
||||
// GetUserAllHasReadSeqs(ctx context.Context, ownerUserID string) (map[string]int64, error)
|
||||
DelUserAllHasReadSeqs(ownerUserID string, conversationIDs ...string) ConversationCache
|
||||
|
||||
GetConversationsByConversationID(ctx context.Context,
|
||||
conversationIDs []string) ([]*relationtb.Conversation, error)
|
||||
DelConversationByConversationID(conversationIDs ...string) ConversationCache
|
||||
GetConversationNotReceiveMessageUserIDs(ctx context.Context, conversationID string) ([]string, error)
|
||||
DelConversationNotReceiveMessageUserIDs(conversationIDs ...string) ConversationCache
|
||||
}
|
||||
|
||||
Vendored
+2
-7
@@ -23,13 +23,8 @@ import (
|
||||
|
||||
type MsgCache interface {
|
||||
GetMessagesBySeq(ctx context.Context, conversationID string, seqs []int64) (seqMsg []*sdkws.MsgData, failedSeqList []int64, err error)
|
||||
SetMessageToCache(ctx context.Context, conversationID string, msgs []*sdkws.MsgData) (int, error)
|
||||
UserDeleteMsgs(ctx context.Context, conversationID string, seqs []int64, userID string) error
|
||||
DelUserDeleteMsgsList(ctx context.Context, conversationID string, seqs []int64)
|
||||
DeleteMessages(ctx context.Context, conversationID string, seqs []int64) error
|
||||
GetUserDelList(ctx context.Context, userID, conversationID string) (seqs []int64, err error)
|
||||
CleanUpOneConversationAllMsg(ctx context.Context, conversationID string) error
|
||||
DelMsgFromCache(ctx context.Context, userID string, seqList []int64) error
|
||||
SetMessagesToCache(ctx context.Context, conversationID string, msgs []*sdkws.MsgData) (int, error)
|
||||
DeleteMessagesFromCache(ctx context.Context, conversationID string, seqs []int64) error
|
||||
SetSendMsgStatus(ctx context.Context, id string, status int32) error
|
||||
GetSendMsgStatus(ctx context.Context, id string) (int32, error)
|
||||
JudgeMessageReactionExist(ctx context.Context, clientMsgID string, sessionType int32) (bool, error)
|
||||
|
||||
+4
-39
@@ -62,17 +62,13 @@ func (c *BatchDeleterRedis) ChainExecDel(ctx context.Context) error {
|
||||
func (c *BatchDeleterRedis) execDel(ctx context.Context, keys []string) error {
|
||||
if len(keys) > 0 {
|
||||
log.ZDebug(ctx, "delete cache", "topic", c.redisPubTopics, "keys", keys)
|
||||
slotMapKeys, err := groupKeysBySlot(ctx, c.redisClient, keys)
|
||||
// Batch delete keys
|
||||
err := ProcessKeysBySlot(ctx, c.redisClient, keys, func(ctx context.Context, slot int64, keys []string) error {
|
||||
return c.rocksClient.TagAsDeletedBatch2(ctx, keys)
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// Batch delete keys
|
||||
for slot, singleSlotKeys := range slotMapKeys {
|
||||
if err := c.rocksClient.TagAsDeletedBatch2(ctx, singleSlotKeys); err != nil {
|
||||
log.ZWarn(ctx, "Batch delete cache failed", err, "slot", slot, "keys", singleSlotKeys)
|
||||
continue
|
||||
}
|
||||
}
|
||||
// Publish the keys that have been deleted to Redis to update the local cache information of other nodes
|
||||
if len(c.redisPubTopics) > 0 && len(keys) > 0 {
|
||||
keysByTopic := localcache.GetPublishKeysByTopic(c.redisPubTopics, keys)
|
||||
@@ -117,37 +113,6 @@ func GetRocksCacheOptions() *rockscache.Options {
|
||||
return &opts
|
||||
}
|
||||
|
||||
// groupKeysBySlot groups keys by their Redis cluster hash slots.
|
||||
func groupKeysBySlot(ctx context.Context, redisClient redis.UniversalClient, keys []string) (map[int64][]string, error) {
|
||||
slots := make(map[int64][]string)
|
||||
clusterClient, isCluster := redisClient.(*redis.ClusterClient)
|
||||
if isCluster {
|
||||
pipe := clusterClient.Pipeline()
|
||||
cmds := make([]*redis.IntCmd, len(keys))
|
||||
for i, key := range keys {
|
||||
cmds[i] = pipe.ClusterKeySlot(ctx, key)
|
||||
}
|
||||
_, err := pipe.Exec(ctx)
|
||||
if err != nil {
|
||||
return nil, errs.WrapMsg(err, "get slot err")
|
||||
}
|
||||
|
||||
for i, cmd := range cmds {
|
||||
slot, err := cmd.Result()
|
||||
if err != nil {
|
||||
log.ZWarn(ctx, "some key get slot err", err, "key", keys[i])
|
||||
continue
|
||||
}
|
||||
slots[slot] = append(slots[slot], keys[i])
|
||||
}
|
||||
} else {
|
||||
// If not a cluster client, put all keys in the same slot (0)
|
||||
slots[0] = keys
|
||||
}
|
||||
|
||||
return slots, nil
|
||||
}
|
||||
|
||||
func getCache[T any](ctx context.Context, rcClient *rockscache.Client, key string, expire time.Duration, fn func(ctx context.Context) (T, error)) (T, error) {
|
||||
var t T
|
||||
var write bool
|
||||
|
||||
@@ -222,14 +222,6 @@ func (c *ConversationRedisCache) DelUserAllHasReadSeqs(ownerUserID string, conve
|
||||
return cache
|
||||
}
|
||||
|
||||
func (c *ConversationRedisCache) GetConversationsByConversationID(ctx context.Context, conversationIDs []string) ([]*model.Conversation, error) {
|
||||
panic("implement me")
|
||||
}
|
||||
|
||||
func (c *ConversationRedisCache) DelConversationByConversationID(conversationIDs ...string) cache.ConversationCache {
|
||||
panic("implement me")
|
||||
}
|
||||
|
||||
func (c *ConversationRedisCache) GetConversationNotReceiveMessageUserIDs(ctx context.Context, conversationID string) ([]string, error) {
|
||||
return getCache(ctx, c.rcClient, c.getConversationNotReceiveMessageUserIDsKey(conversationID), c.expireTime, func(ctx context.Context) ([]string, error) {
|
||||
return c.conversationDB.GetConversationNotReceiveMessageUserIDs(ctx, conversationID)
|
||||
|
||||
+125
@@ -0,0 +1,125 @@
|
||||
package redis
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"github.com/openimsdk/open-im-server/v3/pkg/common/servererrs"
|
||||
"github.com/openimsdk/tools/errs"
|
||||
"github.com/openimsdk/tools/log"
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
var (
|
||||
setBatchWithCommonExpireScript = redis.NewScript(`
|
||||
local expire = tonumber(ARGV[1])
|
||||
for i, key in ipairs(KEYS) do
|
||||
redis.call('SET', key, ARGV[i + 1])
|
||||
redis.call('EXPIRE', key, expire)
|
||||
end
|
||||
return #KEYS
|
||||
`)
|
||||
|
||||
setBatchWithIndividualExpireScript = redis.NewScript(`
|
||||
local n = #KEYS
|
||||
for i = 1, n do
|
||||
redis.call('SET', KEYS[i], ARGV[i])
|
||||
redis.call('EXPIRE', KEYS[i], ARGV[i + n])
|
||||
end
|
||||
return n
|
||||
`)
|
||||
|
||||
deleteBatchScript = redis.NewScript(`
|
||||
for i, key in ipairs(KEYS) do
|
||||
redis.call('DEL', key)
|
||||
end
|
||||
return #KEYS
|
||||
`)
|
||||
|
||||
getBatchScript = redis.NewScript(`
|
||||
local values = {}
|
||||
for i, key in ipairs(KEYS) do
|
||||
local value = redis.call('GET', key)
|
||||
table.insert(values, value)
|
||||
end
|
||||
return values
|
||||
`)
|
||||
)
|
||||
|
||||
func callLua(ctx context.Context, rdb redis.Scripter, script *redis.Script, keys []string, args []any) (any, error) {
|
||||
log.ZDebug(ctx, "callLua args", "scriptHash", script.Hash(), "keys", keys, "args", args)
|
||||
r := script.EvalSha(ctx, rdb, keys, args)
|
||||
if redis.HasErrorPrefix(r.Err(), "NOSCRIPT") {
|
||||
if err := script.Load(ctx, rdb).Err(); err != nil {
|
||||
r = script.Eval(ctx, rdb, keys, args)
|
||||
} else {
|
||||
r = script.EvalSha(ctx, rdb, keys, args)
|
||||
}
|
||||
}
|
||||
v, err := r.Result()
|
||||
if err == redis.Nil {
|
||||
err = nil
|
||||
}
|
||||
return v, errs.WrapMsg(err, "call lua err", "scriptHash", script.Hash(), "keys", keys, "args", args)
|
||||
}
|
||||
|
||||
func LuaSetBatchWithCommonExpire(ctx context.Context, rdb redis.Scripter, keys []string, values []string, expire int) error {
|
||||
// Check if the lengths of keys and values match
|
||||
if len(keys) != len(values) {
|
||||
return errs.New("keys and values length mismatch").Wrap()
|
||||
}
|
||||
|
||||
// Ensure allocation size does not overflow
|
||||
maxAllowedLen := (1 << 31) - 1 // 2GB limit (maximum address space for 32-bit systems)
|
||||
|
||||
if len(values) > maxAllowedLen-1 {
|
||||
return fmt.Errorf("values length is too large, causing overflow")
|
||||
}
|
||||
var vals = make([]any, 0, 1+len(values))
|
||||
vals = append(vals, expire)
|
||||
for _, v := range values {
|
||||
vals = append(vals, v)
|
||||
}
|
||||
_, err := callLua(ctx, rdb, setBatchWithCommonExpireScript, keys, vals)
|
||||
return err
|
||||
}
|
||||
|
||||
func LuaSetBatchWithIndividualExpire(ctx context.Context, rdb redis.Scripter, keys []string, values []string, expires []int) error {
|
||||
// Check if the lengths of keys, values, and expires match
|
||||
if len(keys) != len(values) || len(keys) != len(expires) {
|
||||
return errs.New("keys and values length mismatch").Wrap()
|
||||
}
|
||||
|
||||
// Ensure the allocation size does not overflow
|
||||
maxAllowedLen := (1 << 31) - 1 // 2GB limit (maximum address space for 32-bit systems)
|
||||
|
||||
if len(values) > maxAllowedLen-1 {
|
||||
return errs.New(fmt.Sprintf("values length %d exceeds the maximum allowed length %d", len(values), maxAllowedLen-1)).Wrap()
|
||||
}
|
||||
var vals = make([]any, 0, len(values)+len(expires))
|
||||
for _, v := range values {
|
||||
vals = append(vals, v)
|
||||
}
|
||||
for _, ex := range expires {
|
||||
vals = append(vals, ex)
|
||||
}
|
||||
_, err := callLua(ctx, rdb, setBatchWithIndividualExpireScript, keys, vals)
|
||||
return err
|
||||
}
|
||||
|
||||
func LuaDeleteBatch(ctx context.Context, rdb redis.Scripter, keys []string) error {
|
||||
_, err := callLua(ctx, rdb, deleteBatchScript, keys, nil)
|
||||
return err
|
||||
}
|
||||
|
||||
func LuaGetBatch(ctx context.Context, rdb redis.Scripter, keys []string) ([]any, error) {
|
||||
v, err := callLua(ctx, rdb, getBatchScript, keys, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
values, ok := v.([]any)
|
||||
if !ok {
|
||||
return nil, servererrs.ErrArgs.WrapMsg("invalid lua get batch result")
|
||||
}
|
||||
return values, nil
|
||||
|
||||
}
|
||||
@@ -0,0 +1,75 @@
|
||||
package redis
|
||||
|
||||
import (
|
||||
"context"
|
||||
"github.com/go-redis/redismock/v9"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestLuaSetBatchWithCommonExpire(t *testing.T) {
|
||||
rdb, mock := redismock.NewClientMock()
|
||||
ctx := context.Background()
|
||||
|
||||
keys := []string{"key1", "key2"}
|
||||
values := []string{"value1", "value2"}
|
||||
expire := 10
|
||||
|
||||
mock.ExpectEvalSha(setBatchWithCommonExpireScript.Hash(), keys, []any{expire, "value1", "value2"}).SetVal(int64(len(keys)))
|
||||
|
||||
err := LuaSetBatchWithCommonExpire(ctx, rdb, keys, values, expire)
|
||||
require.NoError(t, err)
|
||||
assert.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestLuaSetBatchWithIndividualExpire(t *testing.T) {
|
||||
rdb, mock := redismock.NewClientMock()
|
||||
ctx := context.Background()
|
||||
|
||||
keys := []string{"key1", "key2"}
|
||||
values := []string{"value1", "value2"}
|
||||
expires := []int{10, 20}
|
||||
|
||||
args := make([]any, 0, len(values)+len(expires))
|
||||
for _, v := range values {
|
||||
args = append(args, v)
|
||||
}
|
||||
for _, ex := range expires {
|
||||
args = append(args, ex)
|
||||
}
|
||||
|
||||
mock.ExpectEvalSha(setBatchWithIndividualExpireScript.Hash(), keys, args).SetVal(int64(len(keys)))
|
||||
|
||||
err := LuaSetBatchWithIndividualExpire(ctx, rdb, keys, values, expires)
|
||||
require.NoError(t, err)
|
||||
assert.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestLuaDeleteBatch(t *testing.T) {
|
||||
rdb, mock := redismock.NewClientMock()
|
||||
ctx := context.Background()
|
||||
|
||||
keys := []string{"key1", "key2"}
|
||||
|
||||
mock.ExpectEvalSha(deleteBatchScript.Hash(), keys, []any{}).SetVal(int64(len(keys)))
|
||||
|
||||
err := LuaDeleteBatch(ctx, rdb, keys)
|
||||
require.NoError(t, err)
|
||||
assert.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestLuaGetBatch(t *testing.T) {
|
||||
rdb, mock := redismock.NewClientMock()
|
||||
ctx := context.Background()
|
||||
|
||||
keys := []string{"key1", "key2"}
|
||||
expectedValues := []any{"value1", "value2"}
|
||||
|
||||
mock.ExpectEvalSha(getBatchScript.Hash(), keys, []any{}).SetVal(expectedValues)
|
||||
|
||||
values, err := LuaGetBatch(ctx, rdb, keys)
|
||||
require.NoError(t, err)
|
||||
assert.NoError(t, mock.ExpectationsWereMet())
|
||||
assert.Equal(t, expectedValues, values)
|
||||
}
|
||||
-15
@@ -1,15 +0,0 @@
|
||||
// Copyright © 2023 OpenIM. All rights reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package redis
|
||||
+61
-313
@@ -16,37 +16,25 @@ package redis
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"github.com/gogo/protobuf/jsonpb"
|
||||
"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/msgprocessor"
|
||||
"github.com/openimsdk/protocol/constant"
|
||||
"github.com/openimsdk/protocol/sdkws"
|
||||
"github.com/openimsdk/tools/errs"
|
||||
"github.com/openimsdk/tools/log"
|
||||
"github.com/openimsdk/tools/utils/stringutil"
|
||||
"github.com/openimsdk/tools/utils/datautil"
|
||||
"github.com/redis/go-redis/v9"
|
||||
"golang.org/x/sync/errgroup"
|
||||
"time"
|
||||
)
|
||||
) //
|
||||
|
||||
const msgCacheTimeout = 86400 * time.Second
|
||||
// msgCacheTimeout is expiration time of message cache, 86400 seconds
|
||||
const msgCacheTimeout = 86400
|
||||
|
||||
var concurrentLimit = 3
|
||||
|
||||
func NewMsgCache(client redis.UniversalClient, redisEnablePipeline bool) cache.MsgCache {
|
||||
return &msgCache{rdb: client, msgCacheTimeout: msgCacheTimeout, redisEnablePipeline: redisEnablePipeline}
|
||||
func NewMsgCache(client redis.UniversalClient) cache.MsgCache {
|
||||
return &msgCache{rdb: client}
|
||||
}
|
||||
|
||||
type msgCache struct {
|
||||
rdb redis.UniversalClient
|
||||
msgCacheTimeout time.Duration
|
||||
redisEnablePipeline bool
|
||||
}
|
||||
|
||||
func (c *msgCache) getAllMessageCacheKey(conversationID string) string {
|
||||
return cachekey.GetAllMessageCacheKey(conversationID)
|
||||
rdb redis.UniversalClient
|
||||
}
|
||||
|
||||
func (c *msgCache) getMessageCacheKey(conversationID string, seq int64) string {
|
||||
@@ -72,218 +60,41 @@ func (c *msgCache) getMessageReactionExPrefix(clientMsgID string, sessionType in
|
||||
return cachekey.GetMessageReactionExKey(clientMsgID, sessionType)
|
||||
}
|
||||
|
||||
func (c *msgCache) SetMessageToCache(ctx context.Context, conversationID string, msgs []*sdkws.MsgData) (int, error) {
|
||||
if c.redisEnablePipeline {
|
||||
return c.PipeSetMessageToCache(ctx, conversationID, msgs)
|
||||
}
|
||||
return c.ParallelSetMessageToCache(ctx, conversationID, msgs)
|
||||
}
|
||||
|
||||
func (c *msgCache) PipeSetMessageToCache(ctx context.Context, conversationID string, msgs []*sdkws.MsgData) (int, error) {
|
||||
pipe := c.rdb.Pipeline()
|
||||
for _, msg := range msgs {
|
||||
s, err := msgprocessor.Pb2String(msg)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
key := c.getMessageCacheKey(conversationID, msg.Seq)
|
||||
_ = pipe.Set(ctx, key, s, c.msgCacheTimeout)
|
||||
}
|
||||
|
||||
results, err := pipe.Exec(ctx)
|
||||
if err != nil {
|
||||
return 0, errs.Wrap(err)
|
||||
}
|
||||
|
||||
for _, res := range results {
|
||||
if res.Err() != nil {
|
||||
return 0, errs.Wrap(err)
|
||||
}
|
||||
}
|
||||
|
||||
return len(msgs), nil
|
||||
}
|
||||
|
||||
func (c *msgCache) ParallelSetMessageToCache(ctx context.Context, conversationID string, msgs []*sdkws.MsgData) (int, error) {
|
||||
wg := errgroup.Group{}
|
||||
wg.SetLimit(concurrentLimit)
|
||||
|
||||
for _, msg := range msgs {
|
||||
msg := msg // closure safe var
|
||||
wg.Go(func() error {
|
||||
s, err := msgprocessor.Pb2String(msg)
|
||||
if err != nil {
|
||||
return errs.Wrap(err)
|
||||
}
|
||||
|
||||
key := c.getMessageCacheKey(conversationID, msg.Seq)
|
||||
if err := c.rdb.Set(ctx, key, s, c.msgCacheTimeout).Err(); err != nil {
|
||||
return errs.Wrap(err)
|
||||
}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
err := wg.Wait()
|
||||
if err != nil {
|
||||
return 0, errs.WrapMsg(err, "wg.Wait failed")
|
||||
}
|
||||
|
||||
return len(msgs), nil
|
||||
}
|
||||
|
||||
func (c *msgCache) UserDeleteMsgs(ctx context.Context, conversationID string, seqs []int64, userID string) error {
|
||||
for _, seq := range seqs {
|
||||
delUserListKey := c.getMessageDelUserListKey(conversationID, seq)
|
||||
userDelListKey := c.getUserDelList(conversationID, userID)
|
||||
err := c.rdb.SAdd(ctx, delUserListKey, userID).Err()
|
||||
if err != nil {
|
||||
return errs.Wrap(err)
|
||||
}
|
||||
err = c.rdb.SAdd(ctx, userDelListKey, seq).Err()
|
||||
if err != nil {
|
||||
return errs.Wrap(err)
|
||||
}
|
||||
if err := c.rdb.Expire(ctx, delUserListKey, c.msgCacheTimeout).Err(); err != nil {
|
||||
return errs.Wrap(err)
|
||||
}
|
||||
if err := c.rdb.Expire(ctx, userDelListKey, c.msgCacheTimeout).Err(); err != nil {
|
||||
return errs.Wrap(err)
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *msgCache) GetUserDelList(ctx context.Context, userID, conversationID string) (seqs []int64, err error) {
|
||||
result, err := c.rdb.SMembers(ctx, c.getUserDelList(conversationID, userID)).Result()
|
||||
if err != nil {
|
||||
return nil, errs.Wrap(err)
|
||||
}
|
||||
seqs = make([]int64, len(result))
|
||||
for i, v := range result {
|
||||
seqs[i] = stringutil.StringToInt64(v)
|
||||
}
|
||||
|
||||
return seqs, nil
|
||||
}
|
||||
|
||||
func (c *msgCache) DelUserDeleteMsgsList(ctx context.Context, conversationID string, seqs []int64) {
|
||||
for _, seq := range seqs {
|
||||
delUsers, err := c.rdb.SMembers(ctx, c.getMessageDelUserListKey(conversationID, seq)).Result()
|
||||
if err != nil {
|
||||
log.ZWarn(ctx, "DelUserDeleteMsgsList failed", err, "conversationID", conversationID, "seq", seq)
|
||||
|
||||
continue
|
||||
}
|
||||
if len(delUsers) > 0 {
|
||||
var failedFlag bool
|
||||
for _, userID := range delUsers {
|
||||
err = c.rdb.SRem(ctx, c.getUserDelList(conversationID, userID), seq).Err()
|
||||
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 {
|
||||
failedFlag = true
|
||||
log.ZWarn(ctx, "DelUserDeleteMsgsList failed", err, "conversationID", conversationID, "seq", seq, "userID", userID)
|
||||
}
|
||||
}
|
||||
if !failedFlag {
|
||||
if err := c.rdb.Del(ctx, c.getMessageDelUserListKey(conversationID, seq)).Err(); err != nil {
|
||||
log.ZWarn(ctx, "DelUserDeleteMsgsList failed", err, "conversationID", conversationID, "seq", seq)
|
||||
return err
|
||||
}
|
||||
values = append(values, s)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *msgCache) DeleteMessages(ctx context.Context, conversationID string, seqs []int64) error {
|
||||
if c.redisEnablePipeline {
|
||||
return c.PipeDeleteMessages(ctx, conversationID, seqs)
|
||||
}
|
||||
|
||||
return c.ParallelDeleteMessages(ctx, conversationID, seqs)
|
||||
}
|
||||
|
||||
func (c *msgCache) ParallelDeleteMessages(ctx context.Context, conversationID string, seqs []int64) error {
|
||||
wg := errgroup.Group{}
|
||||
wg.SetLimit(concurrentLimit)
|
||||
|
||||
for _, seq := range seqs {
|
||||
seq := seq
|
||||
wg.Go(func() error {
|
||||
err := c.rdb.Del(ctx, c.getMessageCacheKey(conversationID, seq)).Err()
|
||||
if err != nil {
|
||||
return errs.Wrap(err)
|
||||
}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
return wg.Wait()
|
||||
}
|
||||
|
||||
func (c *msgCache) PipeDeleteMessages(ctx context.Context, conversationID string, seqs []int64) error {
|
||||
pipe := c.rdb.Pipeline()
|
||||
for _, seq := range seqs {
|
||||
_ = pipe.Del(ctx, c.getMessageCacheKey(conversationID, seq))
|
||||
}
|
||||
|
||||
results, err := pipe.Exec(ctx)
|
||||
return LuaSetBatchWithCommonExpire(ctx, c.rdb, keys, values, msgCacheTimeout)
|
||||
})
|
||||
if err != nil {
|
||||
return errs.WrapMsg(err, "pipe.del")
|
||||
return 0, err
|
||||
}
|
||||
|
||||
for _, res := range results {
|
||||
if res.Err() != nil {
|
||||
return errs.Wrap(err)
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
return len(msgs), nil
|
||||
}
|
||||
|
||||
func (c *msgCache) CleanUpOneConversationAllMsg(ctx context.Context, conversationID string) error {
|
||||
vals, err := c.rdb.Keys(ctx, c.getAllMessageCacheKey(conversationID)).Result()
|
||||
if errors.Is(err, redis.Nil) {
|
||||
return nil
|
||||
}
|
||||
if err != nil {
|
||||
return errs.Wrap(err)
|
||||
}
|
||||
for _, v := range vals {
|
||||
if err := c.rdb.Del(ctx, v).Err(); err != nil {
|
||||
return errs.Wrap(err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *msgCache) DelMsgFromCache(ctx context.Context, userID string, seqs []int64) error {
|
||||
func (c *msgCache) DeleteMessagesFromCache(ctx context.Context, conversationID string, seqs []int64) error {
|
||||
var keys []string
|
||||
for _, seq := range seqs {
|
||||
key := c.getMessageCacheKey(userID, seq)
|
||||
result, err := c.rdb.Get(ctx, key).Result()
|
||||
if err != nil {
|
||||
if errors.Is(err, redis.Nil) {
|
||||
continue
|
||||
}
|
||||
|
||||
return errs.Wrap(err)
|
||||
}
|
||||
var msg sdkws.MsgData
|
||||
err = jsonpb.UnmarshalString(result, &msg)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
msg.Status = constant.MsgDeleted
|
||||
s, err := msgprocessor.Pb2String(&msg)
|
||||
if err != nil {
|
||||
return errs.Wrap(err)
|
||||
}
|
||||
if err := c.rdb.Set(ctx, key, s, c.msgCacheTimeout).Err(); err != nil {
|
||||
return errs.Wrap(err)
|
||||
}
|
||||
keys = append(keys, c.getMessageCacheKey(conversationID, seq))
|
||||
}
|
||||
|
||||
return nil
|
||||
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 {
|
||||
@@ -338,102 +149,39 @@ func (c *msgCache) DeleteOneMessageKey(ctx context.Context, clientMsgID string,
|
||||
}
|
||||
|
||||
func (c *msgCache) GetMessagesBySeq(ctx context.Context, conversationID string, seqs []int64) (seqMsgs []*sdkws.MsgData, failedSeqs []int64, err error) {
|
||||
if c.redisEnablePipeline {
|
||||
return c.PipeGetMessagesBySeq(ctx, conversationID, seqs)
|
||||
}
|
||||
|
||||
return c.ParallelGetMessagesBySeq(ctx, conversationID, seqs)
|
||||
}
|
||||
|
||||
func (c *msgCache) PipeGetMessagesBySeq(ctx context.Context, conversationID string, seqs []int64) (seqMsgs []*sdkws.MsgData, failedSeqs []int64, err error) {
|
||||
pipe := c.rdb.Pipeline()
|
||||
|
||||
results := []*redis.StringCmd{}
|
||||
var keys []string
|
||||
keySeqMap := make(map[string]int64, 10)
|
||||
for _, seq := range seqs {
|
||||
results = append(results, pipe.Get(ctx, c.getMessageCacheKey(conversationID, seq)))
|
||||
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
|
||||
}
|
||||
|
||||
_, err = pipe.Exec(ctx)
|
||||
if err != nil && err != redis.Nil {
|
||||
return seqMsgs, failedSeqs, errs.WrapMsg(err, "pipe.get")
|
||||
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
|
||||
|
||||
for idx, res := range results {
|
||||
seq := seqs[idx]
|
||||
if res.Err() != nil {
|
||||
log.ZError(ctx, "GetMessagesBySeq failed", err, "conversationID", conversationID, "seq", seq, "err", res.Err())
|
||||
failedSeqs = append(failedSeqs, seq)
|
||||
continue
|
||||
}
|
||||
|
||||
msg := sdkws.MsgData{}
|
||||
if err = msgprocessor.String2Pb(res.Val(), &msg); err != nil {
|
||||
log.ZError(ctx, "GetMessagesBySeq Unmarshal failed", err, "res", res, "conversationID", conversationID, "seq", seq)
|
||||
failedSeqs = append(failedSeqs, seq)
|
||||
continue
|
||||
}
|
||||
|
||||
if msg.Status == constant.MsgDeleted {
|
||||
failedSeqs = append(failedSeqs, seq)
|
||||
continue
|
||||
}
|
||||
|
||||
seqMsgs = append(seqMsgs, &msg)
|
||||
}
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
func (c *msgCache) ParallelGetMessagesBySeq(ctx context.Context, conversationID string, seqs []int64) (seqMsgs []*sdkws.MsgData, failedSeqs []int64, err error) {
|
||||
type entry struct {
|
||||
err error
|
||||
msg *sdkws.MsgData
|
||||
}
|
||||
|
||||
wg := errgroup.Group{}
|
||||
wg.SetLimit(concurrentLimit)
|
||||
|
||||
results := make([]entry, len(seqs)) // set slice len/cap to length of seqs.
|
||||
for idx, seq := range seqs {
|
||||
// closure safe var
|
||||
idx := idx
|
||||
seq := seq
|
||||
|
||||
wg.Go(func() error {
|
||||
res, err := c.rdb.Get(ctx, c.getMessageCacheKey(conversationID, seq)).Result()
|
||||
if err != nil {
|
||||
log.ZError(ctx, "GetMessagesBySeq failed", err, "conversationID", conversationID, "seq", seq)
|
||||
results[idx] = entry{err: err}
|
||||
return nil
|
||||
}
|
||||
|
||||
msg := sdkws.MsgData{}
|
||||
if err = msgprocessor.String2Pb(res, &msg); err != nil {
|
||||
log.ZError(ctx, "GetMessagesBySeq Unmarshal failed", err, "res", res, "conversationID", conversationID, "seq", seq)
|
||||
results[idx] = entry{err: err}
|
||||
return nil
|
||||
}
|
||||
|
||||
if msg.Status == constant.MsgDeleted {
|
||||
results[idx] = entry{err: err}
|
||||
return nil
|
||||
}
|
||||
|
||||
results[idx] = entry{msg: &msg}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
_ = wg.Wait()
|
||||
|
||||
for idx, res := range results {
|
||||
if res.err != nil {
|
||||
failedSeqs = append(failedSeqs, seqs[idx])
|
||||
continue
|
||||
}
|
||||
|
||||
seqMsgs = append(seqMsgs, res.msg)
|
||||
}
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
+95
-362
@@ -4,14 +4,13 @@
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package redis
|
||||
|
||||
import (
|
||||
@@ -20,381 +19,115 @@ import (
|
||||
"github.com/openimsdk/protocol/sdkws"
|
||||
"github.com/redis/go-redis/v9"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"math/rand"
|
||||
"google.golang.org/protobuf/proto"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestParallelSetMessageToCache(t *testing.T) {
|
||||
var (
|
||||
cid = fmt.Sprintf("cid-%v", rand.Int63())
|
||||
seqFirst = rand.Int63()
|
||||
msgs = []*sdkws.MsgData{}
|
||||
)
|
||||
|
||||
for i := 0; i < 100; i++ {
|
||||
msgs = append(msgs, &sdkws.MsgData{
|
||||
Seq: seqFirst + int64(i),
|
||||
func Test_msgCache_SetMessagesToCache(t *testing.T) {
|
||||
type fields struct {
|
||||
rdb redis.UniversalClient
|
||||
}
|
||||
type args struct {
|
||||
ctx context.Context
|
||||
conversationID string
|
||||
msgs []*sdkws.MsgData
|
||||
}
|
||||
tests := []struct {
|
||||
name string
|
||||
fields fields
|
||||
args args
|
||||
want int
|
||||
wantErr assert.ErrorAssertionFunc
|
||||
}{
|
||||
{"test1", fields{rdb: redis.NewClient(&redis.Options{Addr: "localhost:16379", Username: "", Password: "openIM123", DB: 0})}, args{context.Background(),
|
||||
"cid", []*sdkws.MsgData{{Seq: 1}, {Seq: 2}, {Seq: 3}}}, 3, assert.NoError},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
c := &msgCache{
|
||||
rdb: tt.fields.rdb,
|
||||
}
|
||||
got, err := c.SetMessagesToCache(tt.args.ctx, tt.args.conversationID, tt.args.msgs)
|
||||
if !tt.wantErr(t, err, fmt.Sprintf("SetMessagesToCache(%v, %v, %v)", tt.args.ctx, tt.args.conversationID, tt.args.msgs)) {
|
||||
return
|
||||
}
|
||||
assert.Equalf(t, tt.want, got, "SetMessagesToCache(%v, %v, %v)", tt.args.ctx, tt.args.conversationID, tt.args.msgs)
|
||||
})
|
||||
}
|
||||
|
||||
testParallelSetMessageToCache(t, cid, msgs)
|
||||
}
|
||||
|
||||
func testParallelSetMessageToCache(t *testing.T, cid string, msgs []*sdkws.MsgData) {
|
||||
rdb := redis.NewClient(&redis.Options{})
|
||||
defer rdb.Close()
|
||||
|
||||
cacher := msgCache{rdb: rdb}
|
||||
|
||||
ret, err := cacher.ParallelSetMessageToCache(context.Background(), cid, msgs)
|
||||
assert.Nil(t, err)
|
||||
assert.Equal(t, len(msgs), ret)
|
||||
|
||||
// validate
|
||||
for _, msg := range msgs {
|
||||
key := cacher.getMessageCacheKey(cid, msg.Seq)
|
||||
val, err := rdb.Exists(context.Background(), key).Result()
|
||||
assert.Nil(t, err)
|
||||
assert.EqualValues(t, 1, val)
|
||||
func Test_msgCache_GetMessagesBySeq(t *testing.T) {
|
||||
type fields struct {
|
||||
rdb redis.UniversalClient
|
||||
}
|
||||
}
|
||||
|
||||
func TestPipeSetMessageToCache(t *testing.T) {
|
||||
var (
|
||||
cid = fmt.Sprintf("cid-%v", rand.Int63())
|
||||
seqFirst = rand.Int63()
|
||||
msgs = []*sdkws.MsgData{}
|
||||
)
|
||||
|
||||
for i := 0; i < 100; i++ {
|
||||
msgs = append(msgs, &sdkws.MsgData{
|
||||
Seq: seqFirst + int64(i),
|
||||
type args struct {
|
||||
ctx context.Context
|
||||
conversationID string
|
||||
seqs []int64
|
||||
}
|
||||
var failedSeq []int64
|
||||
tests := []struct {
|
||||
name string
|
||||
fields fields
|
||||
args args
|
||||
wantSeqMsgs []*sdkws.MsgData
|
||||
wantFailedSeqs []int64
|
||||
wantErr assert.ErrorAssertionFunc
|
||||
}{
|
||||
{"test1", fields{rdb: redis.NewClient(&redis.Options{Addr: "localhost:16379", Password: "openIM123", DB: 0})},
|
||||
args{context.Background(), "cid", []int64{1, 2, 3}},
|
||||
[]*sdkws.MsgData{{Seq: 1}, {Seq: 2}, {Seq: 3}}, failedSeq, assert.NoError},
|
||||
{"test2", fields{rdb: redis.NewClient(&redis.Options{Addr: "localhost:16379", Password: "openIM123", DB: 0})},
|
||||
args{context.Background(), "cid", []int64{4, 5, 6}},
|
||||
nil, []int64{4, 5, 6}, assert.NoError},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
c := &msgCache{
|
||||
rdb: tt.fields.rdb,
|
||||
}
|
||||
gotSeqMsgs, gotFailedSeqs, err := c.GetMessagesBySeq(tt.args.ctx, tt.args.conversationID, tt.args.seqs)
|
||||
if !tt.wantErr(t, err, fmt.Sprintf("GetMessagesBySeq(%v, %v, %v)", tt.args.ctx, tt.args.conversationID, tt.args.seqs)) {
|
||||
return
|
||||
}
|
||||
equalMsgDataSlices(t, tt.wantSeqMsgs, gotSeqMsgs)
|
||||
assert.Equalf(t, tt.wantFailedSeqs, gotFailedSeqs, "GetMessagesBySeq(%v, %v, %v)", tt.args.ctx, tt.args.conversationID, tt.args.seqs)
|
||||
})
|
||||
}
|
||||
|
||||
testPipeSetMessageToCache(t, cid, msgs)
|
||||
}
|
||||
|
||||
func testPipeSetMessageToCache(t *testing.T, cid string, msgs []*sdkws.MsgData) {
|
||||
rdb := redis.NewClient(&redis.Options{})
|
||||
defer rdb.Close()
|
||||
|
||||
cacher := msgCache{rdb: rdb}
|
||||
|
||||
ret, err := cacher.PipeSetMessageToCache(context.Background(), cid, msgs)
|
||||
assert.Nil(t, err)
|
||||
assert.Equal(t, len(msgs), ret)
|
||||
|
||||
// validate
|
||||
for _, msg := range msgs {
|
||||
key := cacher.getMessageCacheKey(cid, msg.Seq)
|
||||
val, err := rdb.Exists(context.Background(), key).Result()
|
||||
assert.Nil(t, err)
|
||||
assert.EqualValues(t, 1, val)
|
||||
func equalMsgDataSlices(t *testing.T, expected, actual []*sdkws.MsgData) {
|
||||
assert.Equal(t, len(expected), len(actual), "Slices have different lengths")
|
||||
for i := range expected {
|
||||
assert.True(t, proto.Equal(expected[i], actual[i]), "Element %d not equal: expected %v, got %v", i, expected[i], actual[i])
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetMessagesBySeq(t *testing.T) {
|
||||
var (
|
||||
cid = fmt.Sprintf("cid-%v", rand.Int63())
|
||||
seqFirst = rand.Int63()
|
||||
msgs = []*sdkws.MsgData{}
|
||||
)
|
||||
|
||||
seqs := []int64{}
|
||||
for i := 0; i < 100; i++ {
|
||||
msgs = append(msgs, &sdkws.MsgData{
|
||||
Seq: seqFirst + int64(i),
|
||||
SendID: fmt.Sprintf("fake-sendid-%v", i),
|
||||
func Test_msgCache_DeleteMessagesFromCache(t *testing.T) {
|
||||
type fields struct {
|
||||
rdb redis.UniversalClient
|
||||
}
|
||||
type args struct {
|
||||
ctx context.Context
|
||||
conversationID string
|
||||
seqs []int64
|
||||
}
|
||||
tests := []struct {
|
||||
name string
|
||||
fields fields
|
||||
args args
|
||||
wantErr assert.ErrorAssertionFunc
|
||||
}{
|
||||
{"test1", fields{rdb: redis.NewClient(&redis.Options{Addr: "localhost:16379", Password: "openIM123"})},
|
||||
args{context.Background(), "cid", []int64{1, 2, 3}}, assert.NoError},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
c := &msgCache{
|
||||
rdb: tt.fields.rdb,
|
||||
}
|
||||
tt.wantErr(t, c.DeleteMessagesFromCache(tt.args.ctx, tt.args.conversationID, tt.args.seqs),
|
||||
fmt.Sprintf("DeleteMessagesFromCache(%v, %v, %v)", tt.args.ctx, tt.args.conversationID, tt.args.seqs))
|
||||
})
|
||||
seqs = append(seqs, seqFirst+int64(i))
|
||||
}
|
||||
|
||||
// set data to cache
|
||||
testPipeSetMessageToCache(t, cid, msgs)
|
||||
|
||||
// get data from cache with parallet mode
|
||||
testParallelGetMessagesBySeq(t, cid, seqs, msgs)
|
||||
|
||||
// get data from cache with pipeline mode
|
||||
testPipeGetMessagesBySeq(t, cid, seqs, msgs)
|
||||
}
|
||||
|
||||
func testParallelGetMessagesBySeq(t *testing.T, cid string, seqs []int64, inputMsgs []*sdkws.MsgData) {
|
||||
rdb := redis.NewClient(&redis.Options{})
|
||||
defer rdb.Close()
|
||||
|
||||
cacher := msgCache{rdb: rdb}
|
||||
|
||||
respMsgs, failedSeqs, err := cacher.ParallelGetMessagesBySeq(context.Background(), cid, seqs)
|
||||
assert.Nil(t, err)
|
||||
assert.Equal(t, 0, len(failedSeqs))
|
||||
assert.Equal(t, len(respMsgs), len(seqs))
|
||||
|
||||
// validate
|
||||
for idx, msg := range respMsgs {
|
||||
assert.Equal(t, msg.Seq, inputMsgs[idx].Seq)
|
||||
assert.Equal(t, msg.SendID, inputMsgs[idx].SendID)
|
||||
}
|
||||
}
|
||||
|
||||
func testPipeGetMessagesBySeq(t *testing.T, cid string, seqs []int64, inputMsgs []*sdkws.MsgData) {
|
||||
rdb := redis.NewClient(&redis.Options{})
|
||||
defer rdb.Close()
|
||||
|
||||
cacher := msgCache{rdb: rdb}
|
||||
|
||||
respMsgs, failedSeqs, err := cacher.PipeGetMessagesBySeq(context.Background(), cid, seqs)
|
||||
assert.Nil(t, err)
|
||||
assert.Equal(t, 0, len(failedSeqs))
|
||||
assert.Equal(t, len(respMsgs), len(seqs))
|
||||
|
||||
// validate
|
||||
for idx, msg := range respMsgs {
|
||||
assert.Equal(t, msg.Seq, inputMsgs[idx].Seq)
|
||||
assert.Equal(t, msg.SendID, inputMsgs[idx].SendID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetMessagesBySeqWithEmptySeqs(t *testing.T) {
|
||||
var (
|
||||
cid = fmt.Sprintf("cid-%v", rand.Int63())
|
||||
seqFirst int64 = 0
|
||||
msgs = []*sdkws.MsgData{}
|
||||
)
|
||||
|
||||
seqs := []int64{}
|
||||
for i := 0; i < 100; i++ {
|
||||
msgs = append(msgs, &sdkws.MsgData{
|
||||
Seq: seqFirst + int64(i),
|
||||
SendID: fmt.Sprintf("fake-sendid-%v", i),
|
||||
})
|
||||
seqs = append(seqs, seqFirst+int64(i))
|
||||
}
|
||||
|
||||
// don't set cache, only get data from cache.
|
||||
|
||||
// get data from cache with parallet mode
|
||||
testParallelGetMessagesBySeqWithEmptry(t, cid, seqs, msgs)
|
||||
|
||||
// get data from cache with pipeline mode
|
||||
testPipeGetMessagesBySeqWithEmptry(t, cid, seqs, msgs)
|
||||
}
|
||||
|
||||
func testParallelGetMessagesBySeqWithEmptry(t *testing.T, cid string, seqs []int64, inputMsgs []*sdkws.MsgData) {
|
||||
rdb := redis.NewClient(&redis.Options{})
|
||||
defer rdb.Close()
|
||||
|
||||
cacher := msgCache{rdb: rdb}
|
||||
|
||||
respMsgs, failedSeqs, err := cacher.ParallelGetMessagesBySeq(context.Background(), cid, seqs)
|
||||
assert.Nil(t, err)
|
||||
assert.Equal(t, len(seqs), len(failedSeqs))
|
||||
assert.Equal(t, 0, len(respMsgs))
|
||||
}
|
||||
|
||||
func testPipeGetMessagesBySeqWithEmptry(t *testing.T, cid string, seqs []int64, inputMsgs []*sdkws.MsgData) {
|
||||
rdb := redis.NewClient(&redis.Options{})
|
||||
defer rdb.Close()
|
||||
|
||||
cacher := msgCache{rdb: rdb}
|
||||
|
||||
respMsgs, failedSeqs, err := cacher.PipeGetMessagesBySeq(context.Background(), cid, seqs)
|
||||
assert.Equal(t, err, redis.Nil)
|
||||
assert.Equal(t, len(seqs), len(failedSeqs))
|
||||
assert.Equal(t, 0, len(respMsgs))
|
||||
}
|
||||
|
||||
func TestGetMessagesBySeqWithLostHalfSeqs(t *testing.T) {
|
||||
var (
|
||||
cid = fmt.Sprintf("cid-%v", rand.Int63())
|
||||
seqFirst int64 = 0
|
||||
msgs = []*sdkws.MsgData{}
|
||||
)
|
||||
|
||||
seqs := []int64{}
|
||||
for i := 0; i < 100; i++ {
|
||||
msgs = append(msgs, &sdkws.MsgData{
|
||||
Seq: seqFirst + int64(i),
|
||||
SendID: fmt.Sprintf("fake-sendid-%v", i),
|
||||
})
|
||||
seqs = append(seqs, seqFirst+int64(i))
|
||||
}
|
||||
|
||||
// Only set half the number of messages.
|
||||
testParallelSetMessageToCache(t, cid, msgs[:50])
|
||||
|
||||
// get data from cache with parallet mode
|
||||
testParallelGetMessagesBySeqWithLostHalfSeqs(t, cid, seqs, msgs)
|
||||
|
||||
// get data from cache with pipeline mode
|
||||
testPipeGetMessagesBySeqWithLostHalfSeqs(t, cid, seqs, msgs)
|
||||
}
|
||||
|
||||
func testParallelGetMessagesBySeqWithLostHalfSeqs(t *testing.T, cid string, seqs []int64, inputMsgs []*sdkws.MsgData) {
|
||||
rdb := redis.NewClient(&redis.Options{})
|
||||
defer rdb.Close()
|
||||
|
||||
cacher := msgCache{rdb: rdb}
|
||||
|
||||
respMsgs, failedSeqs, err := cacher.ParallelGetMessagesBySeq(context.Background(), cid, seqs)
|
||||
assert.Nil(t, err)
|
||||
assert.Equal(t, len(seqs)/2, len(failedSeqs))
|
||||
assert.Equal(t, len(seqs)/2, len(respMsgs))
|
||||
|
||||
for idx, msg := range respMsgs {
|
||||
assert.Equal(t, msg.Seq, seqs[idx])
|
||||
}
|
||||
}
|
||||
|
||||
func testPipeGetMessagesBySeqWithLostHalfSeqs(t *testing.T, cid string, seqs []int64, inputMsgs []*sdkws.MsgData) {
|
||||
rdb := redis.NewClient(&redis.Options{})
|
||||
defer rdb.Close()
|
||||
|
||||
cacher := msgCache{rdb: rdb}
|
||||
|
||||
respMsgs, failedSeqs, err := cacher.PipeGetMessagesBySeq(context.Background(), cid, seqs)
|
||||
assert.Nil(t, err)
|
||||
assert.Equal(t, len(seqs)/2, len(failedSeqs))
|
||||
assert.Equal(t, len(seqs)/2, len(respMsgs))
|
||||
|
||||
for idx, msg := range respMsgs {
|
||||
assert.Equal(t, msg.Seq, seqs[idx])
|
||||
}
|
||||
}
|
||||
|
||||
func TestPipeDeleteMessages(t *testing.T) {
|
||||
var (
|
||||
cid = fmt.Sprintf("cid-%v", rand.Int63())
|
||||
seqFirst = rand.Int63()
|
||||
msgs = []*sdkws.MsgData{}
|
||||
)
|
||||
|
||||
var seqs []int64
|
||||
for i := 0; i < 100; i++ {
|
||||
msgs = append(msgs, &sdkws.MsgData{
|
||||
Seq: seqFirst + int64(i),
|
||||
})
|
||||
seqs = append(seqs, msgs[i].Seq)
|
||||
}
|
||||
|
||||
testPipeSetMessageToCache(t, cid, msgs)
|
||||
testPipeDeleteMessagesOK(t, cid, seqs, msgs)
|
||||
|
||||
// set again
|
||||
testPipeSetMessageToCache(t, cid, msgs)
|
||||
testPipeDeleteMessagesMix(t, cid, seqs[:90], msgs)
|
||||
}
|
||||
|
||||
func testPipeDeleteMessagesOK(t *testing.T, cid string, seqs []int64, inputMsgs []*sdkws.MsgData) {
|
||||
rdb := redis.NewClient(&redis.Options{})
|
||||
defer rdb.Close()
|
||||
|
||||
cacher := msgCache{rdb: rdb}
|
||||
|
||||
err := cacher.PipeDeleteMessages(context.Background(), cid, seqs)
|
||||
assert.Nil(t, err)
|
||||
|
||||
// validate
|
||||
for _, msg := range inputMsgs {
|
||||
key := cacher.getMessageCacheKey(cid, msg.Seq)
|
||||
val := rdb.Exists(context.Background(), key).Val()
|
||||
assert.EqualValues(t, 0, val)
|
||||
}
|
||||
}
|
||||
|
||||
func testPipeDeleteMessagesMix(t *testing.T, cid string, seqs []int64, inputMsgs []*sdkws.MsgData) {
|
||||
rdb := redis.NewClient(&redis.Options{})
|
||||
defer rdb.Close()
|
||||
|
||||
cacher := msgCache{rdb: rdb}
|
||||
|
||||
err := cacher.PipeDeleteMessages(context.Background(), cid, seqs)
|
||||
assert.Nil(t, err)
|
||||
|
||||
// validate
|
||||
for idx, msg := range inputMsgs {
|
||||
key := cacher.getMessageCacheKey(cid, msg.Seq)
|
||||
val, err := rdb.Exists(context.Background(), key).Result()
|
||||
assert.Nil(t, err)
|
||||
if idx < 90 {
|
||||
assert.EqualValues(t, 0, val) // not exists
|
||||
continue
|
||||
}
|
||||
|
||||
assert.EqualValues(t, 1, val) // exists
|
||||
}
|
||||
}
|
||||
|
||||
func TestParallelDeleteMessages(t *testing.T) {
|
||||
var (
|
||||
cid = fmt.Sprintf("cid-%v", rand.Int63())
|
||||
seqFirst = rand.Int63()
|
||||
msgs = []*sdkws.MsgData{}
|
||||
)
|
||||
|
||||
var seqs []int64
|
||||
for i := 0; i < 100; i++ {
|
||||
msgs = append(msgs, &sdkws.MsgData{
|
||||
Seq: seqFirst + int64(i),
|
||||
})
|
||||
seqs = append(seqs, msgs[i].Seq)
|
||||
}
|
||||
|
||||
randSeqs := []int64{}
|
||||
for i := seqFirst + 100; i < seqFirst+200; i++ {
|
||||
randSeqs = append(randSeqs, i)
|
||||
}
|
||||
|
||||
testParallelSetMessageToCache(t, cid, msgs)
|
||||
testParallelDeleteMessagesOK(t, cid, seqs, msgs)
|
||||
|
||||
// set again
|
||||
testParallelSetMessageToCache(t, cid, msgs)
|
||||
testParallelDeleteMessagesMix(t, cid, seqs[:90], msgs, 90)
|
||||
testParallelDeleteMessagesOK(t, cid, seqs[90:], msgs[:90])
|
||||
|
||||
// set again
|
||||
testParallelSetMessageToCache(t, cid, msgs)
|
||||
testParallelDeleteMessagesMix(t, cid, randSeqs, msgs, 0)
|
||||
}
|
||||
|
||||
func testParallelDeleteMessagesOK(t *testing.T, cid string, seqs []int64, inputMsgs []*sdkws.MsgData) {
|
||||
rdb := redis.NewClient(&redis.Options{})
|
||||
defer rdb.Close()
|
||||
|
||||
cacher := msgCache{rdb: rdb}
|
||||
|
||||
err := cacher.PipeDeleteMessages(context.Background(), cid, seqs)
|
||||
assert.Nil(t, err)
|
||||
|
||||
// validate
|
||||
for _, msg := range inputMsgs {
|
||||
key := cacher.getMessageCacheKey(cid, msg.Seq)
|
||||
val := rdb.Exists(context.Background(), key).Val()
|
||||
assert.EqualValues(t, 0, val)
|
||||
}
|
||||
}
|
||||
|
||||
func testParallelDeleteMessagesMix(t *testing.T, cid string, seqs []int64, inputMsgs []*sdkws.MsgData, lessValNonExists int) {
|
||||
rdb := redis.NewClient(&redis.Options{})
|
||||
defer rdb.Close()
|
||||
|
||||
cacher := msgCache{rdb: rdb}
|
||||
|
||||
err := cacher.PipeDeleteMessages(context.Background(), cid, seqs)
|
||||
assert.Nil(t, err)
|
||||
|
||||
// validate
|
||||
for idx, msg := range inputMsgs {
|
||||
key := cacher.getMessageCacheKey(cid, msg.Seq)
|
||||
val, err := rdb.Exists(context.Background(), key).Result()
|
||||
assert.Nil(t, err)
|
||||
if idx < lessValNonExists {
|
||||
assert.EqualValues(t, 0, val) // not exists
|
||||
continue
|
||||
}
|
||||
|
||||
assert.EqualValues(t, 1, val) // exists
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,197 @@
|
||||
package redis
|
||||
|
||||
import (
|
||||
"context"
|
||||
"github.com/openimsdk/tools/errs"
|
||||
"github.com/openimsdk/tools/log"
|
||||
"github.com/redis/go-redis/v9"
|
||||
"golang.org/x/sync/errgroup"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultBatchSize = 50
|
||||
defaultConcurrentLimit = 3
|
||||
)
|
||||
|
||||
// RedisShardManager is a class for sharding and processing keys
|
||||
type RedisShardManager struct {
|
||||
redisClient redis.UniversalClient
|
||||
config *Config
|
||||
}
|
||||
type Config struct {
|
||||
batchSize int
|
||||
continueOnError bool
|
||||
concurrentLimit int
|
||||
}
|
||||
|
||||
// Option is a function type for configuring Config
|
||||
type Option func(c *Config)
|
||||
|
||||
// NewRedisShardManager creates a new RedisShardManager instance
|
||||
func NewRedisShardManager(redisClient redis.UniversalClient, opts ...Option) *RedisShardManager {
|
||||
config := &Config{
|
||||
batchSize: defaultBatchSize, // Default batch size is 50 keys
|
||||
continueOnError: false,
|
||||
concurrentLimit: defaultConcurrentLimit, // Default concurrent limit is 3
|
||||
}
|
||||
for _, opt := range opts {
|
||||
opt(config)
|
||||
}
|
||||
rsm := &RedisShardManager{
|
||||
redisClient: redisClient,
|
||||
config: config,
|
||||
}
|
||||
return rsm
|
||||
}
|
||||
|
||||
// WithBatchSize sets the number of keys to process per batch
|
||||
func WithBatchSize(size int) Option {
|
||||
return func(c *Config) {
|
||||
c.batchSize = size
|
||||
}
|
||||
}
|
||||
|
||||
// WithContinueOnError sets whether to continue processing on error
|
||||
func WithContinueOnError(continueOnError bool) Option {
|
||||
return func(c *Config) {
|
||||
c.continueOnError = continueOnError
|
||||
}
|
||||
}
|
||||
|
||||
// WithConcurrentLimit sets the concurrency limit
|
||||
func WithConcurrentLimit(limit int) Option {
|
||||
return func(c *Config) {
|
||||
c.concurrentLimit = limit
|
||||
}
|
||||
}
|
||||
|
||||
// ProcessKeysBySlot groups keys by their Redis cluster hash slots and processes them using the provided function.
|
||||
func (rsm *RedisShardManager) ProcessKeysBySlot(
|
||||
ctx context.Context,
|
||||
keys []string,
|
||||
processFunc func(ctx context.Context, slot int64, keys []string) error,
|
||||
) error {
|
||||
|
||||
// Group keys by slot
|
||||
slots, err := groupKeysBySlot(ctx, rsm.redisClient, keys)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
g, ctx := errgroup.WithContext(ctx)
|
||||
g.SetLimit(rsm.config.concurrentLimit)
|
||||
|
||||
// Process keys in each slot using the provided function
|
||||
for slot, singleSlotKeys := range slots {
|
||||
batches := splitIntoBatches(singleSlotKeys, rsm.config.batchSize)
|
||||
for _, batch := range batches {
|
||||
slot, batch := slot, batch // Avoid closure capture issue
|
||||
g.Go(func() error {
|
||||
err := processFunc(ctx, slot, batch)
|
||||
if err != nil {
|
||||
log.ZWarn(ctx, "Batch processFunc failed", err, "slot", slot, "keys", batch)
|
||||
if !rsm.config.continueOnError {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
if err := g.Wait(); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// groupKeysBySlot groups keys by their Redis cluster hash slots.
|
||||
func groupKeysBySlot(ctx context.Context, redisClient redis.UniversalClient, keys []string) (map[int64][]string, error) {
|
||||
slots := make(map[int64][]string)
|
||||
clusterClient, isCluster := redisClient.(*redis.ClusterClient)
|
||||
if isCluster {
|
||||
pipe := clusterClient.Pipeline()
|
||||
cmds := make([]*redis.IntCmd, len(keys))
|
||||
for i, key := range keys {
|
||||
cmds[i] = pipe.ClusterKeySlot(ctx, key)
|
||||
}
|
||||
_, err := pipe.Exec(ctx)
|
||||
if err != nil {
|
||||
return nil, errs.WrapMsg(err, "get slot err")
|
||||
}
|
||||
|
||||
for i, cmd := range cmds {
|
||||
slot, err := cmd.Result()
|
||||
if err != nil {
|
||||
log.ZWarn(ctx, "some key get slot err", err, "key", keys[i])
|
||||
return nil, errs.WrapMsg(err, "get slot err", "key", keys[i])
|
||||
}
|
||||
slots[slot] = append(slots[slot], keys[i])
|
||||
}
|
||||
} else {
|
||||
// If not a cluster client, put all keys in the same slot (0)
|
||||
slots[0] = keys
|
||||
}
|
||||
|
||||
return slots, nil
|
||||
}
|
||||
|
||||
// splitIntoBatches splits keys into batches of the specified size
|
||||
func splitIntoBatches(keys []string, batchSize int) [][]string {
|
||||
var batches [][]string
|
||||
for batchSize < len(keys) {
|
||||
keys, batches = keys[batchSize:], append(batches, keys[0:batchSize:batchSize])
|
||||
}
|
||||
return append(batches, keys)
|
||||
}
|
||||
|
||||
// ProcessKeysBySlot groups keys by their Redis cluster hash slots and processes them using the provided function.
|
||||
func ProcessKeysBySlot(
|
||||
ctx context.Context,
|
||||
redisClient redis.UniversalClient,
|
||||
keys []string,
|
||||
processFunc func(ctx context.Context, slot int64, keys []string) error,
|
||||
opts ...Option,
|
||||
) error {
|
||||
|
||||
config := &Config{
|
||||
batchSize: defaultBatchSize,
|
||||
continueOnError: false,
|
||||
concurrentLimit: defaultConcurrentLimit,
|
||||
}
|
||||
for _, opt := range opts {
|
||||
opt(config)
|
||||
}
|
||||
|
||||
// Group keys by slot
|
||||
slots, err := groupKeysBySlot(ctx, redisClient, keys)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
g, ctx := errgroup.WithContext(ctx)
|
||||
g.SetLimit(config.concurrentLimit)
|
||||
|
||||
// Process keys in each slot using the provided function
|
||||
for slot, singleSlotKeys := range slots {
|
||||
batches := splitIntoBatches(singleSlotKeys, config.batchSize)
|
||||
for _, batch := range batches {
|
||||
slot, batch := slot, batch // Avoid closure capture issue
|
||||
g.Go(func() error {
|
||||
err := processFunc(ctx, slot, batch)
|
||||
if err != nil {
|
||||
log.ZWarn(ctx, "Batch processFunc failed", err, "slot", slot, "keys", batch)
|
||||
if !config.continueOnError {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
if err := g.Wait(); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
Vendored
+3
-3
@@ -16,15 +16,15 @@ package cache
|
||||
|
||||
import (
|
||||
"context"
|
||||
relationtb "github.com/openimsdk/open-im-server/v3/pkg/common/storage/model"
|
||||
"github.com/openimsdk/open-im-server/v3/pkg/common/storage/model"
|
||||
"github.com/openimsdk/protocol/user"
|
||||
)
|
||||
|
||||
type UserCache interface {
|
||||
BatchDeleter
|
||||
CloneUserCache() UserCache
|
||||
GetUserInfo(ctx context.Context, userID string) (userInfo *relationtb.User, err error)
|
||||
GetUsersInfo(ctx context.Context, userIDs []string) ([]*relationtb.User, error)
|
||||
GetUserInfo(ctx context.Context, userID string) (userInfo *model.User, err error)
|
||||
GetUsersInfo(ctx context.Context, userIDs []string) ([]*model.User, error)
|
||||
DelUsersInfo(userIDs ...string) UserCache
|
||||
GetUserGlobalRecvMsgOpt(ctx context.Context, userID string) (opt int, err error)
|
||||
DelUsersGlobalRecvMsgOpt(userIDs ...string) UserCache
|
||||
|
||||
@@ -54,8 +54,6 @@ type CommonMsgDatabase interface {
|
||||
MarkSingleChatMsgsAsRead(ctx context.Context, userID string, conversationID string, seqs []int64) error
|
||||
// DeleteMessagesFromCache deletes message caches from Redis by sequence numbers.
|
||||
DeleteMessagesFromCache(ctx context.Context, conversationID string, seqs []int64) error
|
||||
// DelUserDeleteMsgsList deletes user's message deletion list.
|
||||
DelUserDeleteMsgsList(ctx context.Context, conversationID string, seqs []int64)
|
||||
// BatchInsertChat2Cache increments the sequence number and then batch inserts messages into the cache.
|
||||
BatchInsertChat2Cache(ctx context.Context, conversationID string, msgs []*sdkws.MsgData) (seq int64, isNewConversation bool, err error)
|
||||
// GetMsgBySeqsRange retrieves messages from MongoDB by a range of sequence numbers.
|
||||
@@ -98,7 +96,6 @@ type CommonMsgDatabase interface {
|
||||
|
||||
// to mq
|
||||
MsgToMQ(ctx context.Context, key string, msg2mq *sdkws.MsgData) error
|
||||
MsgToModifyMQ(ctx context.Context, key, conversarionID string, msgs []*sdkws.MsgData) error
|
||||
MsgToPushMQ(ctx context.Context, key, conversarionID string, msg2mq *sdkws.MsgData) (int32, int64, error)
|
||||
MsgToMongoMQ(ctx context.Context, key, conversarionID string, msgs []*sdkws.MsgData, lastSeq int64) error
|
||||
|
||||
@@ -150,14 +147,13 @@ func NewCommonMsgDatabase(msgDocModel database.Msg, msg cache.MsgCache, seq cach
|
||||
//}
|
||||
|
||||
type commonMsgDatabase struct {
|
||||
msgDocDatabase database.Msg
|
||||
msgTable model.MsgDocModel
|
||||
msg cache.MsgCache
|
||||
seq cache.SeqCache
|
||||
producer *kafka.Producer
|
||||
producerToMongo *kafka.Producer
|
||||
producerToModify *kafka.Producer
|
||||
producerToPush *kafka.Producer
|
||||
msgDocDatabase database.Msg
|
||||
msgTable model.MsgDocModel
|
||||
msg cache.MsgCache
|
||||
seq cache.SeqCache
|
||||
producer *kafka.Producer
|
||||
producerToMongo *kafka.Producer
|
||||
producerToPush *kafka.Producer
|
||||
}
|
||||
|
||||
func (db *commonMsgDatabase) MsgToMQ(ctx context.Context, key string, msg2mq *sdkws.MsgData) error {
|
||||
@@ -165,14 +161,6 @@ func (db *commonMsgDatabase) MsgToMQ(ctx context.Context, key string, msg2mq *sd
|
||||
return err
|
||||
}
|
||||
|
||||
func (db *commonMsgDatabase) MsgToModifyMQ(ctx context.Context, key, conversationID string, messages []*sdkws.MsgData) error {
|
||||
if len(messages) > 0 {
|
||||
_, _, err := db.producerToModify.SendMessage(ctx, key, &pbmsg.MsgDataToModifyByMQ{ConversationID: conversationID, Messages: messages})
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (db *commonMsgDatabase) MsgToPushMQ(ctx context.Context, key, conversationID string, msg2mq *sdkws.MsgData) (int32, int64, error) {
|
||||
partition, offset, err := db.producerToPush.SendMessage(ctx, key, &pbmsg.PushMsgDataToMQ{MsgData: msg2mq, ConversationID: conversationID})
|
||||
if err != nil {
|
||||
@@ -357,11 +345,7 @@ func (db *commonMsgDatabase) MarkSingleChatMsgsAsRead(ctx context.Context, userI
|
||||
}
|
||||
|
||||
func (db *commonMsgDatabase) DeleteMessagesFromCache(ctx context.Context, conversationID string, seqs []int64) error {
|
||||
return db.msg.DeleteMessages(ctx, conversationID, seqs)
|
||||
}
|
||||
|
||||
func (db *commonMsgDatabase) DelUserDeleteMsgsList(ctx context.Context, conversationID string, seqs []int64) {
|
||||
db.msg.DelUserDeleteMsgsList(ctx, conversationID, seqs)
|
||||
return db.msg.DeleteMessagesFromCache(ctx, conversationID, seqs)
|
||||
}
|
||||
|
||||
func (db *commonMsgDatabase) BatchInsertChat2Cache(ctx context.Context, conversationID string, msgs []*sdkws.MsgData) (seq int64, isNew bool, err error) {
|
||||
@@ -388,7 +372,7 @@ func (db *commonMsgDatabase) BatchInsertChat2Cache(ctx context.Context, conversa
|
||||
userSeqMap[m.SendID] = m.Seq
|
||||
}
|
||||
|
||||
failedNum, err := db.msg.SetMessageToCache(ctx, conversationID, msgs)
|
||||
failedNum, err := db.msg.SetMessagesToCache(ctx, conversationID, msgs)
|
||||
if err != nil {
|
||||
prommetrics.MsgInsertRedisFailedCounter.Add(float64(failedNum))
|
||||
log.ZError(ctx, "setMessageToCache error", err, "len", len(msgs), "conversationID", conversationID)
|
||||
@@ -584,6 +568,7 @@ func (db *commonMsgDatabase) GetMsgBySeqsRange(ctx context.Context, userID strin
|
||||
}
|
||||
newBegin := seqs[0]
|
||||
newEnd := seqs[len(seqs)-1]
|
||||
var successMsgs []*sdkws.MsgData
|
||||
log.ZDebug(ctx, "GetMsgBySeqsRange", "first seqs", seqs, "newBegin", newBegin, "newEnd", newEnd)
|
||||
cachedMsgs, failedSeqs, err := db.msg.GetMessagesBySeq(ctx, conversationID, seqs)
|
||||
if err != nil {
|
||||
@@ -592,54 +577,12 @@ func (db *commonMsgDatabase) GetMsgBySeqsRange(ctx context.Context, userID strin
|
||||
log.ZError(ctx, "get message from redis exception", err, "conversationID", conversationID, "seqs", seqs)
|
||||
}
|
||||
}
|
||||
var successMsgs []*sdkws.MsgData
|
||||
if len(cachedMsgs) > 0 {
|
||||
delSeqs, err := db.msg.GetUserDelList(ctx, userID, conversationID)
|
||||
if err != nil && errs.Unwrap(err) != redis.Nil {
|
||||
return 0, 0, nil, err
|
||||
}
|
||||
var cacheDelNum int
|
||||
for _, msg := range cachedMsgs {
|
||||
if !datautil.Contain(msg.Seq, delSeqs...) {
|
||||
successMsgs = append(successMsgs, msg)
|
||||
} else {
|
||||
cacheDelNum += 1
|
||||
}
|
||||
}
|
||||
log.ZDebug(ctx, "get delSeqs from redis", "delSeqs", delSeqs, "userID", userID, "conversationID", conversationID, "cacheDelNum", cacheDelNum)
|
||||
var reGetSeqsCache []int64
|
||||
for i := 1; i <= cacheDelNum; {
|
||||
newSeq := newBegin - int64(i)
|
||||
if newSeq >= begin {
|
||||
if !datautil.Contain(newSeq, delSeqs...) {
|
||||
log.ZDebug(ctx, "seq del in cache, a new seq in range append", "new seq", newSeq)
|
||||
reGetSeqsCache = append(reGetSeqsCache, newSeq)
|
||||
i++
|
||||
}
|
||||
} else {
|
||||
break
|
||||
}
|
||||
}
|
||||
if len(reGetSeqsCache) > 0 {
|
||||
log.ZDebug(ctx, "reGetSeqsCache", "reGetSeqsCache", reGetSeqsCache)
|
||||
cachedMsgs, failedSeqs2, err := db.msg.GetMessagesBySeq(ctx, conversationID, reGetSeqsCache)
|
||||
if err != nil {
|
||||
if err != redis.Nil {
|
||||
|
||||
log.ZError(ctx, "get message from redis exception", err, "conversationID", conversationID, "seqs", reGetSeqsCache)
|
||||
}
|
||||
}
|
||||
failedSeqs = append(failedSeqs, failedSeqs2...)
|
||||
successMsgs = append(successMsgs, cachedMsgs...)
|
||||
}
|
||||
}
|
||||
log.ZDebug(ctx, "get msgs from cache", "successMsgs", successMsgs)
|
||||
if len(failedSeqs) != 0 {
|
||||
log.ZDebug(ctx, "msgs not exist in redis", "seqs", failedSeqs)
|
||||
}
|
||||
// get from cache or storage
|
||||
successMsgs = append(successMsgs, cachedMsgs...)
|
||||
log.ZDebug(ctx, "get msgs from cache", "cachedMsgs", cachedMsgs)
|
||||
// get from cache or db
|
||||
|
||||
if len(failedSeqs) > 0 {
|
||||
log.ZDebug(ctx, "msgs not exist in redis", "seqs", failedSeqs)
|
||||
mongoMsgs, err := db.getMsgBySeqsRange(ctx, userID, conversationID, failedSeqs, begin, end)
|
||||
if err != nil {
|
||||
|
||||
@@ -679,7 +622,7 @@ func (db *commonMsgDatabase) GetMsgBySeqs(ctx context.Context, userID string, co
|
||||
log.ZError(ctx, "get message from redis exception", err, "failedSeqs", failedSeqs, "conversationID", conversationID)
|
||||
}
|
||||
}
|
||||
log.ZDebug(ctx, "storage.seq.GetMessagesBySeq", "userID", userID, "conversationID", conversationID, "seqs",
|
||||
log.ZDebug(ctx, "db.seq.GetMessagesBySeq", "userID", userID, "conversationID", conversationID, "seqs",
|
||||
seqs, "len(successMsgs)", len(successMsgs), "failedSeqs", failedSeqs)
|
||||
|
||||
if len(failedSeqs) > 0 {
|
||||
@@ -705,12 +648,6 @@ func (db *commonMsgDatabase) DeleteConversationMsgsAndSetMinSeq(ctx context.Cont
|
||||
if minSeq == 0 {
|
||||
return nil
|
||||
}
|
||||
if remainTime == 0 {
|
||||
err = db.msg.CleanUpOneConversationAllMsg(ctx, conversationID)
|
||||
if err != nil {
|
||||
log.ZWarn(ctx, "CleanUpOneUserAllMsg", err, "conversationID", conversationID)
|
||||
}
|
||||
}
|
||||
return db.seq.SetMinSeq(ctx, conversationID, minSeq)
|
||||
}
|
||||
|
||||
@@ -830,7 +767,7 @@ func (db *commonMsgDatabase) deleteMsgRecursion(ctx context.Context, conversatio
|
||||
}
|
||||
|
||||
func (db *commonMsgDatabase) DeleteMsgsPhysicalBySeqs(ctx context.Context, conversationID string, allSeqs []int64) error {
|
||||
if err := db.msg.DeleteMessages(ctx, conversationID, allSeqs); err != nil {
|
||||
if err := db.msg.DeleteMessagesFromCache(ctx, conversationID, allSeqs); err != nil {
|
||||
return err
|
||||
}
|
||||
for docID, seqs := range db.msgTable.GetDocIDSeqsMap(conversationID, allSeqs) {
|
||||
@@ -846,21 +783,9 @@ func (db *commonMsgDatabase) DeleteMsgsPhysicalBySeqs(ctx context.Context, conve
|
||||
}
|
||||
|
||||
func (db *commonMsgDatabase) DeleteUserMsgsBySeqs(ctx context.Context, userID string, conversationID string, seqs []int64) error {
|
||||
cachedMsgs, _, err := db.msg.GetMessagesBySeq(ctx, conversationID, seqs)
|
||||
if err != nil && errs.Unwrap(err) != redis.Nil {
|
||||
log.ZWarn(ctx, "DeleteUserMsgsBySeqs", err, "conversationID", conversationID, "seqs", seqs)
|
||||
if err := db.msg.DeleteMessagesFromCache(ctx, conversationID, seqs); err != nil {
|
||||
return err
|
||||
}
|
||||
if len(cachedMsgs) > 0 {
|
||||
var cacheSeqs []int64
|
||||
for _, msg := range cachedMsgs {
|
||||
cacheSeqs = append(cacheSeqs, msg.Seq)
|
||||
}
|
||||
if err := db.msg.UserDeleteMsgs(ctx, conversationID, cacheSeqs, userID); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
for docID, seqs := range db.msgTable.GetDocIDSeqsMap(conversationID, seqs) {
|
||||
for _, seq := range seqs {
|
||||
if _, err := db.msgDocDatabase.PushUnique(ctx, docID, db.msgTable.GetMsgIndex(seq), "del_list", []string{userID}); err != nil {
|
||||
@@ -1085,14 +1010,14 @@ func (db *commonMsgDatabase) DeleteDocMsgBefore(ctx context.Context, ts int64, d
|
||||
}
|
||||
}
|
||||
|
||||
//func (storage *commonMsgDatabase) ClearMsg(ctx context.Context, ts int64) (err error) {
|
||||
//func (db *commonMsgDatabase) ClearMsg(ctx context.Context, ts int64) (err error) {
|
||||
// var (
|
||||
// docNum int
|
||||
// msgNum int
|
||||
// start = time.Now()
|
||||
// )
|
||||
// for {
|
||||
// msgs, err := storage.msgDocDatabase.GetBeforeMsg(ctx, ts, 100)
|
||||
// msgs, err := db.msgDocDatabase.GetBeforeMsg(ctx, ts, 100)
|
||||
// if err != nil {
|
||||
// return err
|
||||
// }
|
||||
@@ -1100,7 +1025,7 @@ func (db *commonMsgDatabase) DeleteDocMsgBefore(ctx context.Context, ts int64, d
|
||||
// return nil
|
||||
// }
|
||||
// for _, msg := range msgs {
|
||||
// num, err := storage.deleteOneMsg(ctx, ts, msg)
|
||||
// num, err := db.deleteOneMsg(ctx, ts, msg)
|
||||
// if err != nil {
|
||||
// return err
|
||||
// }
|
||||
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
"fmt"
|
||||
"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/utils/datautil"
|
||||
"time"
|
||||
|
||||
"github.com/openimsdk/protocol/constant"
|
||||
@@ -108,29 +109,11 @@ func (m *MsgMgo) GetMsgBySeqIndexIn1Doc(ctx context.Context, userID, docID strin
|
||||
{Key: "input", Value: indexs},
|
||||
{Key: "as", Value: "index"},
|
||||
{Key: "in", Value: bson.D{
|
||||
{Key: "$let", Value: bson.D{
|
||||
{Key: "vars", Value: bson.D{
|
||||
{Key: "currentMsg", Value: bson.D{
|
||||
{Key: "$arrayElemAt", Value: bson.A{"$msgs", "$$index"}},
|
||||
}},
|
||||
}},
|
||||
{Key: "in", Value: bson.D{
|
||||
{Key: "$cond", Value: bson.D{
|
||||
{Key: "if", Value: bson.D{
|
||||
{Key: "$in", Value: bson.A{userID, "$$currentMsg.del_list"}},
|
||||
}},
|
||||
{Key: "then", Value: nil},
|
||||
{Key: "else", Value: "$$currentMsg"},
|
||||
}},
|
||||
}},
|
||||
}},
|
||||
{Key: "$arrayElemAt", Value: bson.A{"$msgs", "$$index"}},
|
||||
}},
|
||||
}},
|
||||
}},
|
||||
}}},
|
||||
bson.D{{Key: "$project", Value: bson.D{
|
||||
{Key: "msgs.del_list", Value: 0},
|
||||
}}},
|
||||
}
|
||||
msgDocModel, err := mongoutil.Aggregate[*model.MsgDocModel](ctx, m.coll, pipeline)
|
||||
if err != nil {
|
||||
@@ -145,6 +128,10 @@ func (m *MsgMgo) GetMsgBySeqIndexIn1Doc(ctx context.Context, userID, docID strin
|
||||
if msg == nil || msg.Msg == nil {
|
||||
continue
|
||||
}
|
||||
if datautil.Contain(userID, msg.DelList...) {
|
||||
msg.Msg.Content = ""
|
||||
msg.Msg.Status = constant.MsgDeleted
|
||||
}
|
||||
if msg.Revoke != nil {
|
||||
revokeContent := sdkws.MessageRevokedContent{
|
||||
RevokerID: msg.Revoke.UserID,
|
||||
|
||||
@@ -0,0 +1,272 @@
|
||||
package batcher
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"github.com/openimsdk/tools/errs"
|
||||
"github.com/openimsdk/tools/utils/idutil"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
var (
|
||||
DefaultDataChanSize = 1000
|
||||
DefaultSize = 100
|
||||
DefaultBuffer = 100
|
||||
DefaultWorker = 5
|
||||
DefaultInterval = time.Second
|
||||
)
|
||||
|
||||
type Config struct {
|
||||
size int // Number of message aggregations
|
||||
buffer int // The number of caches running in a single coroutine
|
||||
dataBuffer int // The size of the main data channel
|
||||
worker int // Number of coroutines processed in parallel
|
||||
interval time.Duration // Time of message aggregations
|
||||
syncWait bool // Whether to wait synchronously after distributing messages have been consumed
|
||||
}
|
||||
|
||||
type Option func(c *Config)
|
||||
|
||||
func WithSize(s int) Option {
|
||||
return func(c *Config) {
|
||||
c.size = s
|
||||
}
|
||||
}
|
||||
|
||||
func WithBuffer(b int) Option {
|
||||
return func(c *Config) {
|
||||
c.buffer = b
|
||||
}
|
||||
}
|
||||
|
||||
func WithWorker(w int) Option {
|
||||
return func(c *Config) {
|
||||
c.worker = w
|
||||
}
|
||||
}
|
||||
|
||||
func WithInterval(i time.Duration) Option {
|
||||
return func(c *Config) {
|
||||
c.interval = i
|
||||
}
|
||||
}
|
||||
|
||||
func WithSyncWait(wait bool) Option {
|
||||
return func(c *Config) {
|
||||
c.syncWait = wait
|
||||
}
|
||||
}
|
||||
|
||||
func WithDataBuffer(size int) Option {
|
||||
return func(c *Config) {
|
||||
c.dataBuffer = size
|
||||
}
|
||||
}
|
||||
|
||||
type Batcher[T any] struct {
|
||||
config *Config
|
||||
|
||||
globalCtx context.Context
|
||||
cancel context.CancelFunc
|
||||
Do func(ctx context.Context, channelID int, val *Msg[T])
|
||||
OnComplete func(lastMessage *T, totalCount int)
|
||||
Sharding func(key string) int
|
||||
Key func(data *T) string
|
||||
HookFunc func(triggerID string, messages map[string][]*T, totalCount int, lastMessage *T)
|
||||
data chan *T
|
||||
chArrays []chan *Msg[T]
|
||||
wait sync.WaitGroup
|
||||
counter sync.WaitGroup
|
||||
}
|
||||
|
||||
func emptyOnComplete[T any](*T, int) {}
|
||||
func emptyHookFunc[T any](string, map[string][]*T, int, *T) {
|
||||
}
|
||||
|
||||
func New[T any](opts ...Option) *Batcher[T] {
|
||||
b := &Batcher[T]{
|
||||
OnComplete: emptyOnComplete[T],
|
||||
HookFunc: emptyHookFunc[T],
|
||||
}
|
||||
config := &Config{
|
||||
size: DefaultSize,
|
||||
buffer: DefaultBuffer,
|
||||
worker: DefaultWorker,
|
||||
interval: DefaultInterval,
|
||||
}
|
||||
for _, opt := range opts {
|
||||
opt(config)
|
||||
}
|
||||
b.config = config
|
||||
b.data = make(chan *T, DefaultDataChanSize)
|
||||
b.globalCtx, b.cancel = context.WithCancel(context.Background())
|
||||
|
||||
b.chArrays = make([]chan *Msg[T], b.config.worker)
|
||||
for i := 0; i < b.config.worker; i++ {
|
||||
b.chArrays[i] = make(chan *Msg[T], b.config.buffer)
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
func (b *Batcher[T]) Worker() int {
|
||||
return b.config.worker
|
||||
}
|
||||
|
||||
func (b *Batcher[T]) Start() error {
|
||||
if b.Sharding == nil {
|
||||
return errs.New("Sharding function is required").Wrap()
|
||||
}
|
||||
if b.Do == nil {
|
||||
return errs.New("Do function is required").Wrap()
|
||||
}
|
||||
if b.Key == nil {
|
||||
return errs.New("Key function is required").Wrap()
|
||||
}
|
||||
b.wait.Add(b.config.worker)
|
||||
for i := 0; i < b.config.worker; i++ {
|
||||
go b.run(i, b.chArrays[i])
|
||||
}
|
||||
b.wait.Add(1)
|
||||
go b.scheduler()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (b *Batcher[T]) Put(ctx context.Context, data *T) error {
|
||||
if data == nil {
|
||||
return errs.New("data can not be nil").Wrap()
|
||||
}
|
||||
select {
|
||||
case <-b.globalCtx.Done():
|
||||
return errs.New("data channel is closed").Wrap()
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case b.data <- data:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func (b *Batcher[T]) scheduler() {
|
||||
ticker := time.NewTicker(b.config.interval)
|
||||
defer func() {
|
||||
ticker.Stop()
|
||||
for _, ch := range b.chArrays {
|
||||
close(ch)
|
||||
}
|
||||
close(b.data)
|
||||
b.wait.Done()
|
||||
}()
|
||||
|
||||
vals := make(map[string][]*T)
|
||||
count := 0
|
||||
var lastAny *T
|
||||
|
||||
for {
|
||||
select {
|
||||
case data, ok := <-b.data:
|
||||
if !ok {
|
||||
// If the data channel is closed unexpectedly
|
||||
return
|
||||
}
|
||||
if data == nil {
|
||||
if count > 0 {
|
||||
b.distributeMessage(vals, count, lastAny)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
key := b.Key(data)
|
||||
vals[key] = append(vals[key], data)
|
||||
lastAny = data
|
||||
|
||||
count++
|
||||
if count >= b.config.size {
|
||||
|
||||
b.distributeMessage(vals, count, lastAny)
|
||||
vals = make(map[string][]*T)
|
||||
count = 0
|
||||
}
|
||||
|
||||
case <-ticker.C:
|
||||
if count > 0 {
|
||||
|
||||
b.distributeMessage(vals, count, lastAny)
|
||||
vals = make(map[string][]*T)
|
||||
count = 0
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type Msg[T any] struct {
|
||||
key string
|
||||
triggerID string
|
||||
val []*T
|
||||
}
|
||||
|
||||
func (m Msg[T]) Key() string {
|
||||
return m.key
|
||||
}
|
||||
|
||||
func (m Msg[T]) TriggerID() string {
|
||||
return m.triggerID
|
||||
}
|
||||
|
||||
func (m Msg[T]) Val() []*T {
|
||||
return m.val
|
||||
}
|
||||
|
||||
func (m Msg[T]) String() string {
|
||||
var sb strings.Builder
|
||||
sb.WriteString("Key: ")
|
||||
sb.WriteString(m.key)
|
||||
sb.WriteString(", Values: [")
|
||||
for i, v := range m.val {
|
||||
if i > 0 {
|
||||
sb.WriteString(", ")
|
||||
}
|
||||
sb.WriteString(fmt.Sprintf("%v", *v))
|
||||
}
|
||||
sb.WriteString("]")
|
||||
return sb.String()
|
||||
}
|
||||
|
||||
func (b *Batcher[T]) distributeMessage(messages map[string][]*T, totalCount int, lastMessage *T) {
|
||||
triggerID := idutil.OperationIDGenerator()
|
||||
b.HookFunc(triggerID, messages, totalCount, lastMessage)
|
||||
for key, data := range messages {
|
||||
if b.config.syncWait {
|
||||
b.counter.Add(1)
|
||||
}
|
||||
channelID := b.Sharding(key)
|
||||
b.chArrays[channelID] <- &Msg[T]{key: key, triggerID: triggerID, val: data}
|
||||
}
|
||||
if b.config.syncWait {
|
||||
b.counter.Wait()
|
||||
}
|
||||
b.OnComplete(lastMessage, totalCount)
|
||||
}
|
||||
|
||||
func (b *Batcher[T]) run(channelID int, ch <-chan *Msg[T]) {
|
||||
defer b.wait.Done()
|
||||
for {
|
||||
select {
|
||||
case messages, ok := <-ch:
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
b.Do(context.Background(), channelID, messages)
|
||||
if b.config.syncWait {
|
||||
b.counter.Done()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (b *Batcher[T]) Close() {
|
||||
b.cancel() // Signal to stop put data
|
||||
b.data <- nil
|
||||
//wait all goroutines exit
|
||||
b.wait.Wait()
|
||||
}
|
||||
@@ -0,0 +1,66 @@
|
||||
package batcher
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"github.com/openimsdk/tools/utils/stringutil"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestBatcher(t *testing.T) {
|
||||
config := Config{
|
||||
size: 1000,
|
||||
buffer: 10,
|
||||
worker: 10,
|
||||
interval: 5 * time.Millisecond,
|
||||
}
|
||||
|
||||
b := New[string](
|
||||
WithSize(config.size),
|
||||
WithBuffer(config.buffer),
|
||||
WithWorker(config.worker),
|
||||
WithInterval(config.interval),
|
||||
WithSyncWait(true),
|
||||
)
|
||||
|
||||
// Mock Do function to simply print values for demonstration
|
||||
b.Do = func(ctx context.Context, channelID int, vals *Msg[string]) {
|
||||
t.Logf("Channel %d Processed batch: %v", channelID, vals)
|
||||
}
|
||||
b.OnComplete = func(lastMessage *string, totalCount int) {
|
||||
t.Logf("Completed processing with last message: %v, total count: %d", *lastMessage, totalCount)
|
||||
}
|
||||
b.Sharding = func(key string) int {
|
||||
hashCode := stringutil.GetHashCode(key)
|
||||
return int(hashCode) % config.worker
|
||||
}
|
||||
b.Key = func(data *string) string {
|
||||
return *data
|
||||
}
|
||||
|
||||
err := b.Start()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// Test normal data processing
|
||||
for i := 0; i < 10000; i++ {
|
||||
data := "data" + fmt.Sprintf("%d", i)
|
||||
if err := b.Put(context.Background(), &data); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
time.Sleep(time.Duration(1) * time.Second)
|
||||
start := time.Now()
|
||||
// Wait for all processing to finish
|
||||
b.Close()
|
||||
|
||||
elapsed := time.Since(start)
|
||||
t.Logf("Close took %s", elapsed)
|
||||
|
||||
if len(b.data) != 0 {
|
||||
t.Error("Data channel should be empty after closing")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user