From 5882597ee1a3cf31dd38a2fff08378a3b30bf195 Mon Sep 17 00:00:00 2001 From: withchao <993506633@qq.com> Date: Fri, 3 Jul 2026 17:51:42 +0800 Subject: [PATCH] feat(redis): implement Redis-based locking mechanism for cron tasks --- internal/tools/cron/cron_task.go | 34 ++++------ .../cron/{dist_look.go => etcd_locker.go} | 7 +- internal/tools/cron/locker.go | 14 ++++ internal/tools/cron/redis_locker.go | 68 +++++++++++++++++++ pkg/common/cmd/cron_task.go | 4 +- .../cache/redis/standalone_gateway_test.go | 2 +- 6 files changed, 103 insertions(+), 26 deletions(-) rename internal/tools/cron/{dist_look.go => etcd_locker.go} (97%) create mode 100644 internal/tools/cron/locker.go create mode 100644 internal/tools/cron/redis_locker.go diff --git a/internal/tools/cron/cron_task.go b/internal/tools/cron/cron_task.go index 2c8655d4a..e4644435b 100644 --- a/internal/tools/cron/cron_task.go +++ b/internal/tools/cron/cron_task.go @@ -3,8 +3,12 @@ package cron import ( "context" + "github.com/robfig/cron/v3" + "google.golang.org/grpc" + "github.com/openimsdk/open-im-server/v3/pkg/common/config" disetcd "github.com/openimsdk/open-im-server/v3/pkg/common/discovery/etcd" + "github.com/openimsdk/open-im-server/v3/pkg/dbbuild" pbconversation "github.com/openimsdk/protocol/conversation" "github.com/openimsdk/protocol/msg" "github.com/openimsdk/protocol/third" @@ -14,14 +18,13 @@ import ( "github.com/openimsdk/tools/log" "github.com/openimsdk/tools/mcontext" "github.com/openimsdk/tools/utils/runtimeenv" - "github.com/robfig/cron/v3" - "google.golang.org/grpc" ) type Config struct { - CronTask config.CronTask - Share config.Share - Discovery config.Discovery + CronTask config.CronTask + Share config.Share + Discovery config.Discovery + RedisConfig config.Redis } func Start(ctx context.Context, conf *Config, client discovery.SvcDiscoveryRegistry, service grpc.ServiceRegistrar) error { @@ -60,10 +63,13 @@ func Start(ctx context.Context, conf *Config, client discovery.SvcDiscoveryRegis if err != nil { return err } - } - - if locker == nil { - locker = emptyLocker{} + } else { + builder := dbbuild.NewBuilder(nil, &conf.RedisConfig) + rdb, err := builder.Redis(ctx) + if err != nil { + return err + } + locker = NewRedisLocker(rdb) } srv := &cronServer{ @@ -95,16 +101,6 @@ func Start(ctx context.Context, conf *Config, client discovery.SvcDiscoveryRegis return nil } -type Locker interface { - ExecuteWithLock(ctx context.Context, taskName string, task func()) -} - -type emptyLocker struct{} - -func (emptyLocker) ExecuteWithLock(ctx context.Context, taskName string, task func()) { - task() -} - type cronServer struct { ctx context.Context config *Config diff --git a/internal/tools/cron/dist_look.go b/internal/tools/cron/etcd_locker.go similarity index 97% rename from internal/tools/cron/dist_look.go rename to internal/tools/cron/etcd_locker.go index e46b9206c..959110fea 100644 --- a/internal/tools/cron/dist_look.go +++ b/internal/tools/cron/etcd_locker.go @@ -6,13 +6,10 @@ import ( "os" "time" - "github.com/openimsdk/tools/log" clientv3 "go.etcd.io/etcd/client/v3" "go.etcd.io/etcd/client/v3/concurrency" -) -const ( - lockLeaseTTL = 300 + "github.com/openimsdk/tools/log" ) type EtcdLocker struct { @@ -35,7 +32,7 @@ func NewEtcdLocker(client *clientv3.Client) (*EtcdLocker, error) { } func (e *EtcdLocker) ExecuteWithLock(ctx context.Context, taskName string, task func()) { - session, err := concurrency.NewSession(e.client, concurrency.WithTTL(lockLeaseTTL)) + session, err := concurrency.NewSession(e.client, concurrency.WithTTL(int(lockLeaseTTL/time.Second))) if err != nil { log.ZWarn(ctx, "Failed to create etcd session", err, "taskName", taskName, diff --git a/internal/tools/cron/locker.go b/internal/tools/cron/locker.go new file mode 100644 index 000000000..4a1115486 --- /dev/null +++ b/internal/tools/cron/locker.go @@ -0,0 +1,14 @@ +package cron + +import ( + "context" + "time" +) + +const ( + lockLeaseTTL = time.Second * 300 +) + +type Locker interface { + ExecuteWithLock(ctx context.Context, taskName string, task func()) +} diff --git a/internal/tools/cron/redis_locker.go b/internal/tools/cron/redis_locker.go new file mode 100644 index 000000000..329e9cb77 --- /dev/null +++ b/internal/tools/cron/redis_locker.go @@ -0,0 +1,68 @@ +package cron + +import ( + "context" + "strings" + "time" + + "github.com/google/uuid" + "github.com/redis/go-redis/v9" + + "github.com/openimsdk/tools/log" +) + +func NewRedisLocker(client redis.UniversalClient) *RedisLocker { + return &RedisLocker{ + client: client, + script: redis.NewScript(strings.TrimSpace(` +if redis.call("get", KEYS[1]) == ARGV[1] then + return redis.call("del", KEYS[1]) +else + return 0 +end +`)), + } +} + +type RedisLocker struct { + client redis.UniversalClient + script *redis.Script +} + +func (e *RedisLocker) getKey(name string) string { + return "CRON_LOCKED:" + name +} + +func (e *RedisLocker) lock(ctx context.Context, name string, owner string) (bool, error) { + ctx, cancel := context.WithTimeout(ctx, time.Second) + defer cancel() + return e.client.SetNX(ctx, e.getKey(name), owner, lockLeaseTTL).Result() +} + +func (e *RedisLocker) unlock(ctx context.Context, name string, owner string) error { + ctx, cancel := context.WithTimeout(ctx, time.Second) + defer cancel() + return e.script.Run(ctx, e.client, []string{e.getKey(name)}, owner).Err() +} + +func (e *RedisLocker) ExecuteWithLock(ctx context.Context, taskName string, task func()) { + owner := uuid.New().String() + ok, err := e.lock(ctx, taskName, owner) + if err != nil { + log.ZWarn(ctx, "cron lock get lock", err, "taskName", taskName) + return + } + log.ZDebug(ctx, "cron lock get lock", "taskName", taskName, "ok", ok, "owner", owner) + if !ok { + return + } + defer func() { + err := e.unlock(ctx, taskName, owner) + if err == nil { + log.ZDebug(ctx, "cron lock unlock", "taskName", taskName, "owner", owner) + } else { + log.ZWarn(ctx, "cron lock unlock", err, "taskName", taskName, "owner", owner) + } + }() + task() +} diff --git a/pkg/common/cmd/cron_task.go b/pkg/common/cmd/cron_task.go index c666bd021..041deea67 100644 --- a/pkg/common/cmd/cron_task.go +++ b/pkg/common/cmd/cron_task.go @@ -17,12 +17,13 @@ package cmd import ( "context" + "github.com/spf13/cobra" + "github.com/openimsdk/open-im-server/v3/internal/tools/cron" "github.com/openimsdk/open-im-server/v3/pkg/common/config" "github.com/openimsdk/open-im-server/v3/pkg/common/startrpc" "github.com/openimsdk/open-im-server/v3/version" "github.com/openimsdk/tools/system/program" - "github.com/spf13/cobra" ) type CronTaskCmd struct { @@ -39,6 +40,7 @@ func NewCronTaskCmd() *CronTaskCmd { config.OpenIMCronTaskCfgFileName: &cronTaskConfig.CronTask, config.ShareFileName: &cronTaskConfig.Share, config.DiscoveryConfigFilename: &cronTaskConfig.Discovery, + config.RedisConfigFileName: &cronTaskConfig.RedisConfig, } ret.RootCmd = NewRootCmd(program.GetProcessName(), WithConfigMap(ret.configMap)) ret.ctx = context.WithValue(context.Background(), "version", version.Version) diff --git a/pkg/common/storage/cache/redis/standalone_gateway_test.go b/pkg/common/storage/cache/redis/standalone_gateway_test.go index b4f2c025f..3b88df104 100644 --- a/pkg/common/storage/cache/redis/standalone_gateway_test.go +++ b/pkg/common/storage/cache/redis/standalone_gateway_test.go @@ -47,6 +47,6 @@ func TestStandaloneGatewayRedisGetGatewayAddrs(t *testing.T) { addrs, err := cache.GetGatewayAddrs(context.Background()) require.NoError(t, err) - assert.Equal(t, []string{"127.0.0.1:10001", "127.0.0.1:10003"}, addrs) + assert.ElementsMatch(t, []string{"127.0.0.1:10001", "127.0.0.1:10003"}, addrs) assert.NoError(t, mock.ExpectationsWereMet()) }