once使用封装的leader和heartbeat

leader/leader.go

@@ -0,0 +1,163 @@
+package leader
+
+// 竞选Leader
+
+import (
+	"context"
+	"errors"
+	"sync"
+	"time"
+
+	"github.com/go-redis/redis/v8"
+	"github.com/yuninks/lockx"
+	"github.com/yuninks/timerx/logger"
+	"github.com/yuninks/timerx/priority"
+)
+
+type Leader struct {
+	ctx    context.Context
+	cancel context.CancelFunc
+
+	isLeader         bool         // 是否是领导
+	leaderLock       sync.RWMutex // 领导锁
+	leaderUniLockKey string       // 领导唯一锁
+	leaderKey        string       // 上报当前的Leader
+
+	redis  redis.UniversalClient // redis
+	logger logger.Logger
+
+	priority *priority.Priority // 全局优先级
+	wg       sync.WaitGroup
+
+	instanceId string // 实例ID
+}
+
+// leader负责选举
+
+func InitLeader(ctx context.Context, ref redis.UniversalClient, keyPrefix string, opts ...Option) (*Leader, error) {
+
+	if ref == nil {
+		return nil, errors.New("redis is nil")
+	}
+
+	op := newOptions(opts...)
+
+	ctx, cancel := context.WithCancel(ctx)
+
+	l := &Leader{
+		ctx:              ctx,
+		cancel:           cancel,
+		redis:            ref,
+		leaderUniLockKey: "timer:leader_lockKey" + keyPrefix,
+		priority:         op.priority,
+		instanceId:       op.instanceId,
+		logger:           op.logger,
+	}
+
+	l.wg.Add(1)
+	go l.leaderElection()
+
+	l.logger.Infof(l.ctx, "InitLeader InstanceId %s lockKey:%s", l.instanceId, l.leaderUniLockKey)
+
+	return l, nil
+}
+
+func (l *Leader) Close() {
+	l.cancel()
+	l.wg.Wait()
+}
+
+// 领导选举
+// 领导作用：全局推选一个人计算执行时间&移入队列，避免每个都进行计算浪费资源
+func (l *Leader) leaderElection() {
+	defer l.wg.Done()
+
+	// 先执行一次
+	l.getLeaderLock()
+
+	ticker := time.NewTicker(5 * time.Second)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ticker.C:
+			l.getLeaderLock()
+		case <-l.ctx.Done():
+			return
+		}
+	}
+}
+
+// 成为领导
+func (l *Leader) getLeaderLock() error {
+
+	// 非当前优先级不用抢leader
+	if l.priority != nil && !l.priority.IsLatest(l.ctx) {
+		return nil
+	}
+
+	ctx, cancel := context.WithCancel(l.ctx)
+	defer cancel()
+
+	// 尝试加锁
+	lock, err := lockx.NewGlobalLock(ctx, l.redis, l.leaderUniLockKey)
+	if err != nil {
+		l.logger.Errorf(l.ctx, "getLeaderLock err:%+v", err)
+		return err
+	}
+	if b, _ := lock.Lock(); !b {
+		// 加锁失败 非Reader
+		l.leaderLock.Lock()
+		l.isLeader = false
+		l.leaderLock.Unlock()
+		return nil
+	}
+	defer lock.Unlock()
+
+	// 加锁成功
+	l.leaderLock.Lock()
+	l.isLeader = true
+	l.leaderLock.Unlock()
+
+	// 上报当前的Leader实例
+	l.redis.Set(l.ctx, l.leaderKey, l.instanceId, time.Hour*24)
+
+	l.logger.Infof(l.ctx, "getLeaderLock Instance %s became leader", lock.GetValue())
+
+	go func() {
+		if l.priority == nil {
+			return
+		}
+
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			default:
+				if !l.priority.IsLatest(l.ctx) {
+					cancel()
+					return
+				}
+				time.Sleep(100 * time.Millisecond)
+			}
+		}
+	}()
+
+	// 等待超时退出
+	<-lock.GetCtx().Done()
+
+	// 已过期
+	// l.leaderLock.Lock()
+	// l.isLeader = false
+	// l.leaderLock.Unlock()
+
+	return nil
+
+}
+
+// isCurrentLeader 检查当前实例是否是leader
+func (c *Leader) IsLeader() bool {
+	c.leaderLock.RLock()
+	defer c.leaderLock.RUnlock()
+	return c.isLeader
+}