From e78b01b296e5eec63e50ec727e9801905b3d8d1e Mon Sep 17 00:00:00 2001 From: Monet Lee Date: Tue, 3 Jun 2025 16:30:21 +0800 Subject: [PATCH 1/6] feat: support distributed lock in crontask. --- internal/tools/cron/cron_task.go | 27 +++- internal/tools/cron/dist_look.go | 260 +++++++++++++++++++++++++++++++ start-config.yml | 2 +- 3 files changed, 285 insertions(+), 4 deletions(-) create mode 100644 internal/tools/cron/dist_look.go diff --git a/internal/tools/cron/cron_task.go b/internal/tools/cron/cron_task.go index 7ae314193..b57baa81f 100644 --- a/internal/tools/cron/cron_task.go +++ b/internal/tools/cron/cron_task.go @@ -58,6 +58,15 @@ func Start(ctx context.Context, conf *Config, client discovery.Conn, service grp cm.Watch(ctx) } + locker, err := NewEtcdLocker(client.(*etcd.SvcDiscoveryRegistryImpl).GetClient()) + if err != nil { + return err + } + + if err := locker.Start(ctx); err != nil { + return err + } + srv := &cronServer{ ctx: ctx, config: conf, @@ -65,6 +74,7 @@ func Start(ctx context.Context, conf *Config, client discovery.Conn, service grp msgClient: msg.NewMsgClient(msgConn), conversationClient: pbconversation.NewConversationClient(conversationConn), thirdClient: third.NewThirdClient(thirdConn), + locker: locker, } if err := srv.registerClearS3(); err != nil { @@ -81,6 +91,9 @@ func Start(ctx context.Context, conf *Config, client discovery.Conn, service grp log.ZDebug(ctx, "cron task server is running") <-ctx.Done() log.ZDebug(ctx, "cron task server is shutting down") + + locker.Stop() // release distributed lock + return nil } @@ -91,6 +104,7 @@ type cronServer struct { msgClient msg.MsgClient conversationClient pbconversation.ConversationClient thirdClient third.ThirdClient + locker *EtcdLocker } func (c *cronServer) registerClearS3() error { @@ -98,7 +112,10 @@ func (c *cronServer) registerClearS3() error { log.ZInfo(c.ctx, "disable scheduled cleanup of s3", "fileExpireTime", c.config.CronTask.FileExpireTime, "deleteObjectType", c.config.CronTask.DeleteObjectType) return nil } - _, err := c.cron.AddFunc(c.config.CronTask.CronExecuteTime, c.clearS3) + _, err := c.cron.AddFunc(c.config.CronTask.CronExecuteTime, func() { + c.locker.ExecuteWithLock(c.ctx, c.clearS3) + }, + ) return errs.WrapMsg(err, "failed to register clear s3 cron task") } @@ -107,11 +124,15 @@ func (c *cronServer) registerDeleteMsg() error { log.ZInfo(c.ctx, "disable scheduled cleanup of chat records", "retainChatRecords", c.config.CronTask.RetainChatRecords) return nil } - _, err := c.cron.AddFunc(c.config.CronTask.CronExecuteTime, c.deleteMsg) + _, err := c.cron.AddFunc(c.config.CronTask.CronExecuteTime, func() { + c.locker.ExecuteWithLock(c.ctx, c.deleteMsg) + }) return errs.WrapMsg(err, "failed to register delete msg cron task") } func (c *cronServer) registerClearUserMsg() error { - _, err := c.cron.AddFunc(c.config.CronTask.CronExecuteTime, c.clearUserMsg) + _, err := c.cron.AddFunc(c.config.CronTask.CronExecuteTime, func() { + c.locker.ExecuteWithLock(c.ctx, c.clearUserMsg) + }) return errs.WrapMsg(err, "failed to register clear user msg cron task") } diff --git a/internal/tools/cron/dist_look.go b/internal/tools/cron/dist_look.go new file mode 100644 index 000000000..57da8254f --- /dev/null +++ b/internal/tools/cron/dist_look.go @@ -0,0 +1,260 @@ +package cron + +import ( + "context" + "fmt" + "os" + "sync/atomic" + "time" + + "github.com/openimsdk/tools/log" + clientv3 "go.etcd.io/etcd/client/v3" +) + +const ( + lockKey = "openim/crontask/dist-lock" + lockLeaseTTL = 15 // Lease TTL in seconds + acquireRetryDelay = 500 * time.Millisecond +) + +type EtcdLocker struct { + client *clientv3.Client + instanceID string + leaseID clientv3.LeaseID + isLockOwner int32 // Using atomic for lock ownership check + watchCh clientv3.WatchChan + watchCancel context.CancelFunc + leaseTTL int64 + stopCh chan struct{} + stoppedCh chan struct{} + acquireDelay time.Duration +} + +// NewEtcdLocker creates a new etcd distributed lock +func NewEtcdLocker(client *clientv3.Client) (*EtcdLocker, error) { + hostname, _ := os.Hostname() + pid := os.Getpid() + instanceID := fmt.Sprintf("%s-pid-%d-%d", hostname, pid, time.Now().UnixNano()) + + locker := &EtcdLocker{ + client: client, + instanceID: instanceID, + leaseTTL: lockLeaseTTL, + stopCh: make(chan struct{}), + stoppedCh: make(chan struct{}), + acquireDelay: acquireRetryDelay, + } + + return locker, nil +} + +func (e *EtcdLocker) Start(ctx context.Context) error { + log.ZInfo(ctx, "Starting etcd distributed lock", "instanceID", e.instanceID) + go e.runLockLoop(ctx) + return nil +} + +func (e *EtcdLocker) runLockLoop(ctx context.Context) { + defer close(e.stoppedCh) + + for { + select { + case <-e.stopCh: + e.releaseLock(ctx) + return + case <-ctx.Done(): + e.releaseLock(ctx) + return + default: + acquired, err := e.tryAcquireLock(ctx) + if err != nil { + log.ZWarn(ctx, "Failed to acquire lock", err, "instanceID", e.instanceID) + time.Sleep(e.acquireDelay) + continue + } + + if acquired { + e.runKeepAlive(ctx) + + time.Sleep(e.acquireDelay) + } else { + e.watchLock(ctx) + } + } + } +} + +func (e *EtcdLocker) tryAcquireLock(ctx context.Context) (bool, error) { + lease, err := e.client.Grant(ctx, e.leaseTTL) + if err != nil { + return false, fmt.Errorf("failed to create lease: %w", err) + } + + txnResp, err := e.client.Txn(ctx). + If(clientv3.Compare(clientv3.CreateRevision(lockKey), "=", 0)). + Then(clientv3.OpPut(lockKey, e.instanceID, clientv3.WithLease(lease.ID))). + Else(clientv3.OpGet(lockKey)). + Commit() + + if err != nil { + e.client.Revoke(ctx, lease.ID) + return false, fmt.Errorf("transaction failed: %w", err) + } + + if !txnResp.Succeeded { + rangeResp := txnResp.Responses[0].GetResponseRange() + if len(rangeResp.Kvs) > 0 { + currentOwner := string(rangeResp.Kvs[0].Value) + log.ZInfo(ctx, "Lock already owned", "instanceID", e.instanceID, "owner", currentOwner) + } + + e.client.Revoke(ctx, lease.ID) + return false, nil + } + + e.leaseID = lease.ID + atomic.StoreInt32(&e.isLockOwner, 1) + log.ZInfo(ctx, "Successfully acquired lock", "instanceID", e.instanceID, "leaseID", lease.ID) + return true, nil +} + +func (e *EtcdLocker) runKeepAlive(ctx context.Context) { + keepAliveCh, err := e.client.KeepAlive(ctx, e.leaseID) + if err != nil { + log.ZError(ctx, "Failed to start lease keepalive", err, "instanceID", e.instanceID) + e.releaseLock(ctx) + return + } + + for { + select { + case _, ok := <-keepAliveCh: + if !ok { + log.ZWarn(ctx, "Keepalive channel closed, lock lost", nil, "instanceID", e.instanceID) + atomic.StoreInt32(&e.isLockOwner, 0) // Set to false atomically + return + } + + case <-ctx.Done(): + log.ZInfo(ctx, "Context canceled, releasing lock", "instanceID", e.instanceID) + e.releaseLock(ctx) + return + + case <-e.stopCh: + log.ZInfo(ctx, "Stop signal received, releasing lock", "instanceID", e.instanceID) + e.releaseLock(ctx) + return + } + } +} + +// Watch lock status directly in etcd +func (e *EtcdLocker) watchLock(ctx context.Context) { + log.ZInfo(ctx, "Starting to watch lock status", "instanceID", e.instanceID) + + watchCtx, cancel := context.WithCancel(ctx) + e.watchCancel = cancel + + defer e.cancelWatch() + + // Watch for changes to the lock key + e.watchCh = e.client.Watch(watchCtx, lockKey) + + for { + select { + case resp, ok := <-e.watchCh: + if !ok { + log.ZWarn(ctx, "Watch channel closed", nil, "instanceID", e.instanceID) + return + } + + for _, event := range resp.Events { + if event.Type == clientv3.EventTypeDelete { + log.ZInfo(ctx, "Lock released, attempting to acquire", "instanceID", e.instanceID) + return + } + } + case <-ctx.Done(): + return + case <-e.stopCh: + return + } + } +} + +// Release the lock +func (e *EtcdLocker) releaseLock(ctx context.Context) { + if atomic.LoadInt32(&e.isLockOwner) == 0 { + return + } + + leaseID := e.leaseID + atomic.StoreInt32(&e.isLockOwner, 0) + e.leaseID = 0 + + if leaseID != 0 { + _, err := e.client.Revoke(context.Background(), leaseID) + if err != nil { + log.ZWarn(ctx, "Failed to revoke lease", err, "instanceID", e.instanceID, "error", err) + } else { + log.ZInfo(ctx, "Successfully released lock", "instanceID", e.instanceID) + } + } +} + +func (e *EtcdLocker) CheckLockOwnership(ctx context.Context) (bool, error) { + if atomic.LoadInt32(&e.isLockOwner) == 0 { + return false, nil + } + + resp, err := e.client.Get(ctx, lockKey) + if err != nil { + return false, fmt.Errorf("failed to check lock status: %w", err) + } + if len(resp.Kvs) > 0 && string(resp.Kvs[0].Value) == e.instanceID { + return true, nil + } + if atomic.LoadInt32(&e.isLockOwner) == 1 { + log.ZWarn(ctx, "Lock ownership lost unexpectedly", nil, "instanceID", e.instanceID) + atomic.StoreInt32(&e.isLockOwner, 0) + } + + return false, nil +} + +func (e *EtcdLocker) cancelWatch() { + if e.watchCancel != nil { + e.watchCancel() + e.watchCancel = nil + } +} + +func (e *EtcdLocker) Stop() { + e.cancelWatch() + close(e.stopCh) + <-e.stoppedCh +} + +func (e *EtcdLocker) IsLockOwner() bool { + return atomic.LoadInt32(&e.isLockOwner) == 1 +} + +func (e *EtcdLocker) ExecuteWithLock(ctx context.Context, task func()) { + if atomic.LoadInt32(&e.isLockOwner) == 0 { + log.ZDebug(ctx, "Instance does not own lock (local check), skipping task execution", "instanceID", e.instanceID) + return + } + isOwner, err := e.CheckLockOwnership(ctx) + if err != nil { + log.ZWarn(ctx, "Failed to verify lock ownership", err, "instanceID", e.instanceID) + return + } + if !isOwner { + log.ZDebug(ctx, "Instance does not own lock (etcd verification), skipping task execution", "instanceID", e.instanceID) + return + } + + log.ZInfo(ctx, "Starting lock-protected task execution", "instanceID", e.instanceID) + task() + log.ZInfo(ctx, "Lock-protected task execution completed", "instanceID", e.instanceID) +} diff --git a/start-config.yml b/start-config.yml index 1231b5d0d..da959044d 100644 --- a/start-config.yml +++ b/start-config.yml @@ -1,6 +1,6 @@ serviceBinaries: openim-api: 1 - openim-crontask: 1 + openim-crontask: 4 openim-rpc-user: 1 openim-msggateway: 1 openim-push: 8 From e42126687c84b3f6dbcd6dc45ec6bf80c2845457 Mon Sep 17 00:00:00 2001 From: Monet Lee Date: Tue, 3 Jun 2025 16:44:44 +0800 Subject: [PATCH 2/6] remove space. --- internal/tools/cron/cron_task.go | 3 +-- internal/tools/cron/dist_look.go | 9 --------- 2 files changed, 1 insertion(+), 11 deletions(-) diff --git a/internal/tools/cron/cron_task.go b/internal/tools/cron/cron_task.go index b57baa81f..24f0762eb 100644 --- a/internal/tools/cron/cron_task.go +++ b/internal/tools/cron/cron_task.go @@ -114,8 +114,7 @@ func (c *cronServer) registerClearS3() error { } _, err := c.cron.AddFunc(c.config.CronTask.CronExecuteTime, func() { c.locker.ExecuteWithLock(c.ctx, c.clearS3) - }, - ) + }) return errs.WrapMsg(err, "failed to register clear s3 cron task") } diff --git a/internal/tools/cron/dist_look.go b/internal/tools/cron/dist_look.go index 57da8254f..9edfbc794 100644 --- a/internal/tools/cron/dist_look.go +++ b/internal/tools/cron/dist_look.go @@ -56,7 +56,6 @@ func (e *EtcdLocker) Start(ctx context.Context) error { func (e *EtcdLocker) runLockLoop(ctx context.Context) { defer close(e.stoppedCh) - for { select { case <-e.stopCh: @@ -75,7 +74,6 @@ func (e *EtcdLocker) runLockLoop(ctx context.Context) { if acquired { e.runKeepAlive(ctx) - time.Sleep(e.acquireDelay) } else { e.watchLock(ctx) @@ -134,12 +132,10 @@ func (e *EtcdLocker) runKeepAlive(ctx context.Context) { atomic.StoreInt32(&e.isLockOwner, 0) // Set to false atomically return } - case <-ctx.Done(): log.ZInfo(ctx, "Context canceled, releasing lock", "instanceID", e.instanceID) e.releaseLock(ctx) return - case <-e.stopCh: log.ZInfo(ctx, "Stop signal received, releasing lock", "instanceID", e.instanceID) e.releaseLock(ctx) @@ -151,15 +147,12 @@ func (e *EtcdLocker) runKeepAlive(ctx context.Context) { // Watch lock status directly in etcd func (e *EtcdLocker) watchLock(ctx context.Context) { log.ZInfo(ctx, "Starting to watch lock status", "instanceID", e.instanceID) - watchCtx, cancel := context.WithCancel(ctx) e.watchCancel = cancel - defer e.cancelWatch() // Watch for changes to the lock key e.watchCh = e.client.Watch(watchCtx, lockKey) - for { select { case resp, ok := <-e.watchCh: @@ -167,7 +160,6 @@ func (e *EtcdLocker) watchLock(ctx context.Context) { log.ZWarn(ctx, "Watch channel closed", nil, "instanceID", e.instanceID) return } - for _, event := range resp.Events { if event.Type == clientv3.EventTypeDelete { log.ZInfo(ctx, "Lock released, attempting to acquire", "instanceID", e.instanceID) @@ -191,7 +183,6 @@ func (e *EtcdLocker) releaseLock(ctx context.Context) { leaseID := e.leaseID atomic.StoreInt32(&e.isLockOwner, 0) e.leaseID = 0 - if leaseID != 0 { _, err := e.client.Revoke(context.Background(), leaseID) if err != nil { From cdf3e9cdb3d479208863e38510282baf9acfe86a Mon Sep 17 00:00:00 2001 From: Monet Lee Date: Tue, 3 Jun 2025 16:49:47 +0800 Subject: [PATCH 3/6] remove comment. --- internal/tools/cron/dist_look.go | 3 --- 1 file changed, 3 deletions(-) diff --git a/internal/tools/cron/dist_look.go b/internal/tools/cron/dist_look.go index 9edfbc794..2bcf7f951 100644 --- a/internal/tools/cron/dist_look.go +++ b/internal/tools/cron/dist_look.go @@ -144,14 +144,12 @@ func (e *EtcdLocker) runKeepAlive(ctx context.Context) { } } -// Watch lock status directly in etcd func (e *EtcdLocker) watchLock(ctx context.Context) { log.ZInfo(ctx, "Starting to watch lock status", "instanceID", e.instanceID) watchCtx, cancel := context.WithCancel(ctx) e.watchCancel = cancel defer e.cancelWatch() - // Watch for changes to the lock key e.watchCh = e.client.Watch(watchCtx, lockKey) for { select { @@ -174,7 +172,6 @@ func (e *EtcdLocker) watchLock(ctx context.Context) { } } -// Release the lock func (e *EtcdLocker) releaseLock(ctx context.Context) { if atomic.LoadInt32(&e.isLockOwner) == 0 { return From dae7c072ed895c40344c17ff786f2f695b9d1085 Mon Sep 17 00:00:00 2001 From: Monet Lee Date: Tue, 3 Jun 2025 16:52:34 +0800 Subject: [PATCH 4/6] remove log. --- internal/tools/cron/dist_look.go | 2 -- 1 file changed, 2 deletions(-) diff --git a/internal/tools/cron/dist_look.go b/internal/tools/cron/dist_look.go index 2bcf7f951..356edb6ef 100644 --- a/internal/tools/cron/dist_look.go +++ b/internal/tools/cron/dist_look.go @@ -229,7 +229,6 @@ func (e *EtcdLocker) IsLockOwner() bool { func (e *EtcdLocker) ExecuteWithLock(ctx context.Context, task func()) { if atomic.LoadInt32(&e.isLockOwner) == 0 { - log.ZDebug(ctx, "Instance does not own lock (local check), skipping task execution", "instanceID", e.instanceID) return } isOwner, err := e.CheckLockOwnership(ctx) @@ -238,7 +237,6 @@ func (e *EtcdLocker) ExecuteWithLock(ctx context.Context, task func()) { return } if !isOwner { - log.ZDebug(ctx, "Instance does not own lock (etcd verification), skipping task execution", "instanceID", e.instanceID) return } From 1389c8b80d305956e8ce8e1c084acc87790d1f84 Mon Sep 17 00:00:00 2001 From: Monet Lee Date: Tue, 3 Jun 2025 18:00:51 +0800 Subject: [PATCH 5/6] Update logic. --- internal/tools/cron/cron_task.go | 12 +- internal/tools/cron/dist_look.go | 261 +++++++++++-------------------- 2 files changed, 94 insertions(+), 179 deletions(-) diff --git a/internal/tools/cron/cron_task.go b/internal/tools/cron/cron_task.go index 24f0762eb..33ebda57d 100644 --- a/internal/tools/cron/cron_task.go +++ b/internal/tools/cron/cron_task.go @@ -63,10 +63,6 @@ func Start(ctx context.Context, conf *Config, client discovery.Conn, service grp return err } - if err := locker.Start(ctx); err != nil { - return err - } - srv := &cronServer{ ctx: ctx, config: conf, @@ -92,8 +88,6 @@ func Start(ctx context.Context, conf *Config, client discovery.Conn, service grp <-ctx.Done() log.ZDebug(ctx, "cron task server is shutting down") - locker.Stop() // release distributed lock - return nil } @@ -113,7 +107,7 @@ func (c *cronServer) registerClearS3() error { return nil } _, err := c.cron.AddFunc(c.config.CronTask.CronExecuteTime, func() { - c.locker.ExecuteWithLock(c.ctx, c.clearS3) + c.locker.ExecuteWithLock(c.ctx, "clearS3", c.clearS3) }) return errs.WrapMsg(err, "failed to register clear s3 cron task") } @@ -124,14 +118,14 @@ func (c *cronServer) registerDeleteMsg() error { return nil } _, err := c.cron.AddFunc(c.config.CronTask.CronExecuteTime, func() { - c.locker.ExecuteWithLock(c.ctx, c.deleteMsg) + c.locker.ExecuteWithLock(c.ctx, "deleteMsg", c.deleteMsg) }) return errs.WrapMsg(err, "failed to register delete msg cron task") } func (c *cronServer) registerClearUserMsg() error { _, err := c.cron.AddFunc(c.config.CronTask.CronExecuteTime, func() { - c.locker.ExecuteWithLock(c.ctx, c.clearUserMsg) + c.locker.ExecuteWithLock(c.ctx, "clearUserMsg", c.clearUserMsg) }) return errs.WrapMsg(err, "failed to register clear user msg cron task") } diff --git a/internal/tools/cron/dist_look.go b/internal/tools/cron/dist_look.go index 356edb6ef..f7911951e 100644 --- a/internal/tools/cron/dist_look.go +++ b/internal/tools/cron/dist_look.go @@ -4,7 +4,6 @@ import ( "context" "fmt" "os" - "sync/atomic" "time" "github.com/openimsdk/tools/log" @@ -12,22 +11,12 @@ import ( ) const ( - lockKey = "openim/crontask/dist-lock" - lockLeaseTTL = 15 // Lease TTL in seconds - acquireRetryDelay = 500 * time.Millisecond + lockLeaseTTL = 3000 ) type EtcdLocker struct { - client *clientv3.Client - instanceID string - leaseID clientv3.LeaseID - isLockOwner int32 // Using atomic for lock ownership check - watchCh clientv3.WatchChan - watchCancel context.CancelFunc - leaseTTL int64 - stopCh chan struct{} - stoppedCh chan struct{} - acquireDelay time.Duration + client *clientv3.Client + instanceID string } // NewEtcdLocker creates a new etcd distributed lock @@ -37,55 +26,19 @@ func NewEtcdLocker(client *clientv3.Client) (*EtcdLocker, error) { instanceID := fmt.Sprintf("%s-pid-%d-%d", hostname, pid, time.Now().UnixNano()) locker := &EtcdLocker{ - client: client, - instanceID: instanceID, - leaseTTL: lockLeaseTTL, - stopCh: make(chan struct{}), - stoppedCh: make(chan struct{}), - acquireDelay: acquireRetryDelay, + client: client, + instanceID: instanceID, } return locker, nil } -func (e *EtcdLocker) Start(ctx context.Context) error { - log.ZInfo(ctx, "Starting etcd distributed lock", "instanceID", e.instanceID) - go e.runLockLoop(ctx) - return nil -} +func (e *EtcdLocker) tryAcquireTaskLock(ctx context.Context, taskName string) (clientv3.LeaseID, bool, error) { + lockKey := fmt.Sprintf("openim/crontask/%s-lock", taskName) -func (e *EtcdLocker) runLockLoop(ctx context.Context) { - defer close(e.stoppedCh) - for { - select { - case <-e.stopCh: - e.releaseLock(ctx) - return - case <-ctx.Done(): - e.releaseLock(ctx) - return - default: - acquired, err := e.tryAcquireLock(ctx) - if err != nil { - log.ZWarn(ctx, "Failed to acquire lock", err, "instanceID", e.instanceID) - time.Sleep(e.acquireDelay) - continue - } - - if acquired { - e.runKeepAlive(ctx) - time.Sleep(e.acquireDelay) - } else { - e.watchLock(ctx) - } - } - } -} - -func (e *EtcdLocker) tryAcquireLock(ctx context.Context) (bool, error) { - lease, err := e.client.Grant(ctx, e.leaseTTL) + lease, err := e.client.Grant(ctx, lockLeaseTTL) if err != nil { - return false, fmt.Errorf("failed to create lease: %w", err) + return 0, false, fmt.Errorf("failed to create lease: %w", err) } txnResp, err := e.client.Txn(ctx). @@ -96,151 +49,119 @@ func (e *EtcdLocker) tryAcquireLock(ctx context.Context) (bool, error) { if err != nil { e.client.Revoke(ctx, lease.ID) - return false, fmt.Errorf("transaction failed: %w", err) + return 0, false, fmt.Errorf("transaction failed: %w", err) } if !txnResp.Succeeded { rangeResp := txnResp.Responses[0].GetResponseRange() if len(rangeResp.Kvs) > 0 { currentOwner := string(rangeResp.Kvs[0].Value) - log.ZInfo(ctx, "Lock already owned", "instanceID", e.instanceID, "owner", currentOwner) + log.ZInfo(ctx, "Task lock already owned, skipping execution", + "taskName", taskName, + "instanceID", e.instanceID, + "currentOwner", currentOwner) } - e.client.Revoke(ctx, lease.ID) - return false, nil + return 0, false, nil } - e.leaseID = lease.ID - atomic.StoreInt32(&e.isLockOwner, 1) - log.ZInfo(ctx, "Successfully acquired lock", "instanceID", e.instanceID, "leaseID", lease.ID) - return true, nil + log.ZInfo(ctx, "Successfully acquired task lock", + "taskName", taskName, + "instanceID", e.instanceID, + "leaseID", lease.ID) + return lease.ID, true, nil } -func (e *EtcdLocker) runKeepAlive(ctx context.Context) { - keepAliveCh, err := e.client.KeepAlive(ctx, e.leaseID) - if err != nil { - log.ZError(ctx, "Failed to start lease keepalive", err, "instanceID", e.instanceID) - e.releaseLock(ctx) +func (e *EtcdLocker) releaseTaskLock(ctx context.Context, taskName string, leaseID clientv3.LeaseID) { + if leaseID == 0 { return } - for { - select { - case _, ok := <-keepAliveCh: - if !ok { - log.ZWarn(ctx, "Keepalive channel closed, lock lost", nil, "instanceID", e.instanceID) - atomic.StoreInt32(&e.isLockOwner, 0) // Set to false atomically - return - } - case <-ctx.Done(): - log.ZInfo(ctx, "Context canceled, releasing lock", "instanceID", e.instanceID) - e.releaseLock(ctx) - return - case <-e.stopCh: - log.ZInfo(ctx, "Stop signal received, releasing lock", "instanceID", e.instanceID) - e.releaseLock(ctx) - return - } + _, err := e.client.Revoke(ctx, leaseID) + if err != nil { + log.ZWarn(ctx, "Failed to revoke task lease", err, + "taskName", taskName, + "instanceID", e.instanceID, + "leaseID", leaseID) + } else { + log.ZInfo(ctx, "Successfully released task lock", + "taskName", taskName, + "instanceID", e.instanceID) } } -func (e *EtcdLocker) watchLock(ctx context.Context) { - log.ZInfo(ctx, "Starting to watch lock status", "instanceID", e.instanceID) - watchCtx, cancel := context.WithCancel(ctx) - e.watchCancel = cancel - defer e.cancelWatch() +func (e *EtcdLocker) startLeaseKeepAlive(ctx context.Context, taskName string, leaseID clientv3.LeaseID) (context.CancelFunc, error) { + keepAliveCtx, cancel := context.WithCancel(ctx) - e.watchCh = e.client.Watch(watchCtx, lockKey) - for { - select { - case resp, ok := <-e.watchCh: - if !ok { - log.ZWarn(ctx, "Watch channel closed", nil, "instanceID", e.instanceID) - return - } - for _, event := range resp.Events { - if event.Type == clientv3.EventTypeDelete { - log.ZInfo(ctx, "Lock released, attempting to acquire", "instanceID", e.instanceID) + keepAliveCh, err := e.client.KeepAlive(keepAliveCtx, leaseID) + if err != nil { + cancel() + return nil, fmt.Errorf("failed to start keepalive: %w", err) + } + + go func() { + defer cancel() + for { + select { + case _, ok := <-keepAliveCh: + if !ok { + log.ZWarn(keepAliveCtx, "KeepAlive channel closed, lease may have expired", nil, + "taskName", taskName, + "instanceID", e.instanceID, + "leaseID", leaseID) return } + + case <-keepAliveCtx.Done(): + log.ZDebug(keepAliveCtx, "KeepAlive stopped", + "taskName", taskName, + "instanceID", e.instanceID) + return } - case <-ctx.Done(): - return - case <-e.stopCh: - return } - } + }() + + return cancel, nil } -func (e *EtcdLocker) releaseLock(ctx context.Context) { - if atomic.LoadInt32(&e.isLockOwner) == 0 { - return - } - - leaseID := e.leaseID - atomic.StoreInt32(&e.isLockOwner, 0) - e.leaseID = 0 - if leaseID != 0 { - _, err := e.client.Revoke(context.Background(), leaseID) - if err != nil { - log.ZWarn(ctx, "Failed to revoke lease", err, "instanceID", e.instanceID, "error", err) - } else { - log.ZInfo(ctx, "Successfully released lock", "instanceID", e.instanceID) - } - } -} - -func (e *EtcdLocker) CheckLockOwnership(ctx context.Context) (bool, error) { - if atomic.LoadInt32(&e.isLockOwner) == 0 { - return false, nil - } - - resp, err := e.client.Get(ctx, lockKey) +func (e *EtcdLocker) ExecuteWithLock(ctx context.Context, taskName string, task func()) { + leaseID, acquired, err := e.tryAcquireTaskLock(ctx, taskName) if err != nil { - return false, fmt.Errorf("failed to check lock status: %w", err) - } - if len(resp.Kvs) > 0 && string(resp.Kvs[0].Value) == e.instanceID { - return true, nil - } - if atomic.LoadInt32(&e.isLockOwner) == 1 { - log.ZWarn(ctx, "Lock ownership lost unexpectedly", nil, "instanceID", e.instanceID) - atomic.StoreInt32(&e.isLockOwner, 0) - } - - return false, nil -} - -func (e *EtcdLocker) cancelWatch() { - if e.watchCancel != nil { - e.watchCancel() - e.watchCancel = nil - } -} - -func (e *EtcdLocker) Stop() { - e.cancelWatch() - close(e.stopCh) - <-e.stoppedCh -} - -func (e *EtcdLocker) IsLockOwner() bool { - return atomic.LoadInt32(&e.isLockOwner) == 1 -} - -func (e *EtcdLocker) ExecuteWithLock(ctx context.Context, task func()) { - if atomic.LoadInt32(&e.isLockOwner) == 0 { + log.ZWarn(ctx, "Failed to acquire task lock", err, + "taskName", taskName, + "instanceID", e.instanceID) return } - isOwner, err := e.CheckLockOwnership(ctx) + + if !acquired { + log.ZDebug(ctx, "Task is being executed by another instance, skipping", + "taskName", taskName, + "instanceID", e.instanceID) + return + } + + cancelKeepAlive, err := e.startLeaseKeepAlive(ctx, taskName, leaseID) if err != nil { - log.ZWarn(ctx, "Failed to verify lock ownership", err, "instanceID", e.instanceID) - return - } - if !isOwner { + log.ZWarn(ctx, "Failed to start lease keepalive", err, + "taskName", taskName, + "instanceID", e.instanceID) + e.releaseTaskLock(ctx, taskName, leaseID) return } - log.ZInfo(ctx, "Starting lock-protected task execution", "instanceID", e.instanceID) + defer func() { + cancelKeepAlive() + e.releaseTaskLock(ctx, taskName, leaseID) + }() + + log.ZInfo(ctx, "Starting task execution with lease keepalive", + "taskName", taskName, + "instanceID", e.instanceID, + "leaseID", leaseID) + task() - log.ZInfo(ctx, "Lock-protected task execution completed", "instanceID", e.instanceID) + + log.ZInfo(ctx, "Task execution completed", + "taskName", taskName, + "instanceID", e.instanceID) } From f21bc47d7173885cb47ef9d00d2e234e025e31d2 Mon Sep 17 00:00:00 2001 From: Monet Lee Date: Tue, 3 Jun 2025 18:27:10 +0800 Subject: [PATCH 6/6] Update contents. --- internal/tools/cron/cron_task.go | 1 + internal/tools/cron/dist_look.go | 140 +++++++------------------------ 2 files changed, 32 insertions(+), 109 deletions(-) diff --git a/internal/tools/cron/cron_task.go b/internal/tools/cron/cron_task.go index 33ebda57d..a4de309d4 100644 --- a/internal/tools/cron/cron_task.go +++ b/internal/tools/cron/cron_task.go @@ -87,6 +87,7 @@ func Start(ctx context.Context, conf *Config, client discovery.Conn, service grp log.ZDebug(ctx, "cron task server is running") <-ctx.Done() log.ZDebug(ctx, "cron task server is shutting down") + srv.cron.Stop() return nil } diff --git a/internal/tools/cron/dist_look.go b/internal/tools/cron/dist_look.go index f7911951e..d1d2d1cb0 100644 --- a/internal/tools/cron/dist_look.go +++ b/internal/tools/cron/dist_look.go @@ -8,10 +8,11 @@ import ( "github.com/openimsdk/tools/log" clientv3 "go.etcd.io/etcd/client/v3" + "go.etcd.io/etcd/client/v3/concurrency" ) const ( - lockLeaseTTL = 3000 + lockLeaseTTL = 300 ) type EtcdLocker struct { @@ -33,131 +34,52 @@ func NewEtcdLocker(client *clientv3.Client) (*EtcdLocker, error) { return locker, nil } -func (e *EtcdLocker) tryAcquireTaskLock(ctx context.Context, taskName string) (clientv3.LeaseID, bool, error) { - lockKey := fmt.Sprintf("openim/crontask/%s-lock", taskName) - - lease, err := e.client.Grant(ctx, lockLeaseTTL) - if err != nil { - return 0, false, fmt.Errorf("failed to create lease: %w", err) - } - - txnResp, err := e.client.Txn(ctx). - If(clientv3.Compare(clientv3.CreateRevision(lockKey), "=", 0)). - Then(clientv3.OpPut(lockKey, e.instanceID, clientv3.WithLease(lease.ID))). - Else(clientv3.OpGet(lockKey)). - Commit() - - if err != nil { - e.client.Revoke(ctx, lease.ID) - return 0, false, fmt.Errorf("transaction failed: %w", err) - } - - if !txnResp.Succeeded { - rangeResp := txnResp.Responses[0].GetResponseRange() - if len(rangeResp.Kvs) > 0 { - currentOwner := string(rangeResp.Kvs[0].Value) - log.ZInfo(ctx, "Task lock already owned, skipping execution", - "taskName", taskName, - "instanceID", e.instanceID, - "currentOwner", currentOwner) - } - e.client.Revoke(ctx, lease.ID) - return 0, false, nil - } - - log.ZInfo(ctx, "Successfully acquired task lock", - "taskName", taskName, - "instanceID", e.instanceID, - "leaseID", lease.ID) - return lease.ID, true, nil -} - -func (e *EtcdLocker) releaseTaskLock(ctx context.Context, taskName string, leaseID clientv3.LeaseID) { - if leaseID == 0 { - return - } - - _, err := e.client.Revoke(ctx, leaseID) - if err != nil { - log.ZWarn(ctx, "Failed to revoke task lease", err, - "taskName", taskName, - "instanceID", e.instanceID, - "leaseID", leaseID) - } else { - log.ZInfo(ctx, "Successfully released task lock", - "taskName", taskName, - "instanceID", e.instanceID) - } -} - -func (e *EtcdLocker) startLeaseKeepAlive(ctx context.Context, taskName string, leaseID clientv3.LeaseID) (context.CancelFunc, error) { - keepAliveCtx, cancel := context.WithCancel(ctx) - - keepAliveCh, err := e.client.KeepAlive(keepAliveCtx, leaseID) - if err != nil { - cancel() - return nil, fmt.Errorf("failed to start keepalive: %w", err) - } - - go func() { - defer cancel() - for { - select { - case _, ok := <-keepAliveCh: - if !ok { - log.ZWarn(keepAliveCtx, "KeepAlive channel closed, lease may have expired", nil, - "taskName", taskName, - "instanceID", e.instanceID, - "leaseID", leaseID) - return - } - - case <-keepAliveCtx.Done(): - log.ZDebug(keepAliveCtx, "KeepAlive stopped", - "taskName", taskName, - "instanceID", e.instanceID) - return - } - } - }() - - return cancel, nil -} - func (e *EtcdLocker) ExecuteWithLock(ctx context.Context, taskName string, task func()) { - leaseID, acquired, err := e.tryAcquireTaskLock(ctx, taskName) + session, err := concurrency.NewSession(e.client, concurrency.WithTTL(lockLeaseTTL)) if err != nil { - log.ZWarn(ctx, "Failed to acquire task lock", err, + log.ZWarn(ctx, "Failed to create etcd session", err, "taskName", taskName, "instanceID", e.instanceID) return } + defer session.Close() - if !acquired { - log.ZDebug(ctx, "Task is being executed by another instance, skipping", - "taskName", taskName, - "instanceID", e.instanceID) - return - } + lockKey := fmt.Sprintf("openim/crontask/%s", taskName) + mutex := concurrency.NewMutex(session, lockKey) - cancelKeepAlive, err := e.startLeaseKeepAlive(ctx, taskName, leaseID) + ctxWithTimeout, cancel := context.WithTimeout(ctx, 100*time.Millisecond) + defer cancel() + + err = mutex.TryLock(ctxWithTimeout) if err != nil { - log.ZWarn(ctx, "Failed to start lease keepalive", err, - "taskName", taskName, - "instanceID", e.instanceID) - e.releaseTaskLock(ctx, taskName, leaseID) + if err == context.DeadlineExceeded { + log.ZDebug(ctx, "Task is being executed by another instance, skipping", + "taskName", taskName, + "instanceID", e.instanceID) + } else { + log.ZWarn(ctx, "Failed to acquire task lock", err, + "taskName", taskName, + "instanceID", e.instanceID) + } return } defer func() { - cancelKeepAlive() - e.releaseTaskLock(ctx, taskName, leaseID) + if err := mutex.Unlock(ctx); err != nil { + log.ZWarn(ctx, "Failed to release task lock", err, + "taskName", taskName, + "instanceID", e.instanceID) + } else { + log.ZInfo(ctx, "Successfully released task lock", + "taskName", taskName, + "instanceID", e.instanceID) + } }() - log.ZInfo(ctx, "Starting task execution with lease keepalive", + log.ZInfo(ctx, "Successfully acquired task lock, starting execution", "taskName", taskName, "instanceID", e.instanceID, - "leaseID", leaseID) + "sessionID", session.Lease()) task()