iotlcm/admin-svr/server/internal/tasks/job_lock.go

package tasks

import (
	"context"
	"fmt"
	"os"
	"strconv"
	"strings"
	"time"

	"git.5iot.com/iotlcm-public/iotlcm-admin-svr/utility/uuid32"

	"github.com/gogf/gf/v2/frame/g"
)

const jobLockKeyPrefix = "iotlcm:job-lock"

type redisJobLockOptions struct {
	InvokeTarget string
	Period       time.Duration
	TTL          time.Duration
}

type redisJobLock struct {
	key   string
	value string
	ttl   time.Duration
}

func (t Task) runWithRedisJobLock(ctx context.Context, options redisJobLockOptions, run func(context.Context)) {
	if ctx == nil {
		ctx = context.Background()
	}
	if run == nil {
		return
	}
	invokeTarget := strings.TrimSpace(options.InvokeTarget)
	if strings.TrimSpace(invokeTarget) == "" {
		invokeTarget = taskInvokeTarget(t.ID, "")
	}
	if options.TTL <= 0 {
		options.TTL = 30 * time.Minute
	}

	lock := redisJobLock{
		key:   jobLockKey(invokeTarget, time.Now(), options.Period),
		value: jobLockValue(),
		ttl:   options.TTL,
	}
	acquired, err := lock.acquire(ctx)
	if err != nil {
		g.Log().Errorf(ctx, "job lock acquire failed, invokeTarget:%s, key:%s, err:%+v", invokeTarget, lock.key, err)
		return
	}
	if !acquired {
		g.Log().Infof(ctx, "job skipped by redis lock, invokeTarget:%s, key:%s", invokeTarget, lock.key)
		return
	}

	runCtx, cancel := context.WithCancel(ctx)
	stopRenew := lock.startRenew(runCtx, cancel, invokeTarget)
	defer stopRenew()
	defer func() {
		if err := lock.release(context.Background()); err != nil {
			g.Log().Warningf(ctx, "job lock release failed, invokeTarget:%s, key:%s, err:%+v", invokeTarget, lock.key, err)
		}
	}()

	g.Log().Infof(ctx, "job lock acquired, invokeTarget:%s, key:%s, ttl:%s", invokeTarget, lock.key, lock.ttl)
	run(runCtx)
}

func (l redisJobLock) acquire(ctx context.Context) (bool, error) {
	res, err := g.Redis().Do(ctx, "SET", l.key, l.value, "NX", "PX", l.ttl.Milliseconds())
	if err != nil {
		return false, err
	}
	return strings.EqualFold(strings.TrimSpace(res.String()), "OK"), nil
}

func (l redisJobLock) startRenew(ctx context.Context, cancel context.CancelFunc, invokeTarget string) func() {
	interval := l.ttl / 3
	if interval <= 0 {
		interval = time.Minute
	}
	stop := make(chan struct{})
	done := make(chan struct{})
	go func() {
		defer close(done)
		ticker := time.NewTicker(interval)
		defer ticker.Stop()
		for {
			select {
			case <-ctx.Done():
				return
			case <-stop:
				return
			case <-ticker.C:
				renewCtx, renewCancel := context.WithTimeout(context.Background(), 5*time.Second)
				ok, err := l.renew(renewCtx)
				renewCancel()
				if err != nil {
					g.Log().Warningf(ctx, "job lock renew failed, invokeTarget:%s, key:%s, err:%+v", invokeTarget, l.key, err)
					continue
				}
				if !ok {
					g.Log().Warningf(ctx, "job lock lost, cancel current run, invokeTarget:%s, key:%s", invokeTarget, l.key)
					cancel()
					return
				}
			}
		}
	}()
	return func() {
		close(stop)
		<-done
		cancel()
	}
}

func (l redisJobLock) renew(ctx context.Context) (bool, error) {
	const script = `if redis.call("GET", KEYS[1]) == ARGV[1] then return redis.call("PEXPIRE", KEYS[1], ARGV[2]) else return 0 end`
	res, err := g.Redis().Do(ctx, "EVAL", script, 1, l.key, l.value, strconv.FormatInt(l.ttl.Milliseconds(), 10))
	if err != nil {
		return false, err
	}
	return res.Int() == 1, nil
}

func (l redisJobLock) release(ctx context.Context) error {
	const script = `if redis.call("GET", KEYS[1]) == ARGV[1] then return redis.call("DEL", KEYS[1]) else return 0 end`
	_, err := g.Redis().Do(ctx, "EVAL", script, 1, l.key, l.value)
	return err
}

func jobLockKey(invokeTarget string, now time.Time, period time.Duration) string {
	return fmt.Sprintf("%s:%s:%s", jobLockKeyPrefix, normalizeLockPart(invokeTarget), jobLockPeriodKey(now, period))
}

func jobLockPeriodKey(now time.Time, period time.Duration) string {
	if now.IsZero() {
		now = time.Now()
	}
	if period < 0 {
		return "active"
	}
	if period >= 24*time.Hour {
		return now.Format("20060102")
	}
	if period > 0 {
		return strconv.FormatInt(now.Truncate(period).Unix(), 10)
	}
	return now.Format("20060102150405")
}

func taskInvokeTarget(taskID string, fallback string) string {
	taskID = strings.TrimSpace(taskID)
	if taskID == "" {
		return fallback
	}
	if index := strings.LastIndex(taskID, "-job-"); index > 0 {
		return taskID[:index]
	}
	return fallback
}

func normalizeLockPart(value string) string {
	value = strings.TrimSpace(value)
	value = strings.ReplaceAll(value, ":", "_")
	value = strings.ReplaceAll(value, " ", "_")
	if value == "" {
		return "unknown"
	}
	return value
}

func jobLockValue() string {
	host, _ := os.Hostname()
	if strings.TrimSpace(host) == "" {
		host = "unknown"
	}
	return host + ":" + uuid32.S()
}