基于 golang 从零到一实现时间轮算法 (三)

引言

本文参考小徐先生的相关博客整理，项目地址为：
https://github.com/xiaoxuxiansheng/timewheel/blob/main/redis_time_wheel.go。主要是完善流程以及记录个人学习笔记。

分布式版实现

本章我们讨论一下，如何基于 redis 实现分布式版本的时间轮，以贴合实际生产环境对分布式定时任务调度系统的诉求.

redis 版时间轮的实现思路是使用 redis 中的有序集合 sorted set（简称 zset）进行定时任务的存储管理，其中以每个定时任务执行时间对应的时间戳作为 zset 中的 score，完成定时任务的有序排列组合.

zset 数据结构的 redis 官方文档链接：https://redis.io/docs/data-type，这里简单看一下使用。

在这里插入图片描述

Redis 的 ZSET（有序集合）是 Redis 数据类型之一，它是字符串元素的集合，且不允许重复的成员。不同的是，每个元素都会关联一个 double 类型的分数。Redis 正是通过分数来为集合中的成员进行从小到大的排序。ZSET的成员是唯一的，但分数（score）却可以重复。

基本操作包括添加元素、删除元素、修改元素的分数、查询元素的分数等。以下是一些常用的 ZSET 操作命令：

ZADD key score member：向 ZSET 中添加一个元素，如果元素已存在则更新其分数。
ZSCORE key member：获取 ZSET 中元素的分数。
ZRANGE key start stop [WITHSCORES]：按照分数从低到高的顺序返回 ZSET 中指定区间内的元素，如果使用了 WITHSCORES 选项，则结果中会包含元素的分数。
ZREVRANGE key start stop [WITHSCORES]：功能与 ZRANGE 相同，但是元素是按分数从高到低返回的。
ZRANGEBYSCORE key min max [WITHSCORES] [LIMIT offset count]：返回 ZSET 中分数在 min 和 max 之间的元素。
ZREMRANGEBYRANK key start stop：删除 ZSET 中排名在给定区间内的所有成员。
ZREMRANGEBYSCORE key min max：删除 ZSET 中分数在给定区间内的所有成员。
ZINCRBY key increment member：增加或减少 ZSET 中指定成员的分数。
ZCARD key：获取 ZSET 的成员数。
ZCOUNT key min max：计算 ZSET 中分数在 min 和 max 之间的成员数量。
ZREM key member [member …]：删除 ZSET 中的一个或多个成员。

docker安装

取最新版的 Redis 镜像

docker pull redis:latest

在这里插入图片描述

运行容器
安装完成后，我们可以使用以下命令来运行 redis 容器：

$ docker run -itd --name redis-test -p 6379:6379 redis

接着我们通过 redis-cli 连接测试使用 redis 服务。

$ docker exec -it redis-test /bin/bash

使用示例：

# 向名为 myzset 的 ZSET 中添加三个元素
ZADD myzset 1 "one" 2 "two" 3 "three"

# 获取 myzset 中的所有元素和它们的分数
ZRANGE myzset 0 -1 WITHSCORES

# 获取 myzset 中分数为 2 的成员的数量
ZCOUNT myzset 2 2

# 增加元素 "one" 的分数
ZINCRBY myzset 10 "one"

在这里插入图片描述

代码使用

代码可以参考仓库https://github.com/xiaoxuxiansheng/timewheel/blob/main/redis_time_wheel.go。这里输出启动过程。

在这里插入图片描述

数据结构

redis 时间轮

在 redis 版时间轮中有两个核心类，第一个是关于时间轮的类定义：

redisClient：定时任务的存储是基于 redis zset 实现的，因此需要内置一个 redis 客户端，这部分在 3.2 小节展开；
httpClient：定时任务执行时，是通过请求使用方预留回调地址的方式实现的，因此需要内置一个 http 客户端
channel × 2：ticker 和 stopc 对应为 golang 标准库定时器以及停止 goroutine 的控制器

// 基于 redis 实现的分布式版时间轮
type RTimeWheel struct {
    // 内置的单例工具，用于保证 stopc 只被关闭一次
    sync.Once
    // redis 客户端
    redisClient *redis.Client
    // http 客户端. 在执行定时任务时需要使用到.
    httpClient  *thttp.Client
    // 用于停止时间轮的控制器 channel
    stopc       chan struct{
    // 触发定时扫描任务的定时器 
    ticker      *time.Ticker
}

在这里插入图片描述

定时任务

定时任务的类型定义如下，其中包括定时任务的唯一键 key，以及执行定时任务回调时需要使用到的 http 协议参数.

// 使用方提交的每一笔定时任务
type RTaskElement struct {
    // 定时任务全局唯一 key
    Key         string            `json:"key"`
    // 定时任务执行时，回调的 http url
    CallbackURL string            `json:"callback_url"`
    // 回调时使用的 http 方法
    Method      string            `json:"method"`
    // 回调时传递的请求参数
    Req         interface{}       `json:"req"`
    // 回调时使用的 http 请求头
    Header      map[string]string `json:"header"`
}

在这里插入图片描述

构造器

在构造时间轮实例时，使用方需要注入 redis 客户端以及 http 客户端.

在初始化流程中，ticker 为 golang 标准库实现的定时器，定时器的执行时间间隔固定为 1 s. 此外会异步运行 run 方法，启动一个常驻 goroutine，生命周期会通过 stopc channel 进行控制.

func NewRTimeWheel(redisClient *redis.Client, httpClient *thttp.Client) *RTimeWheel {
	r := RTimeWheel{
		ticker:      time.NewTicker(time.Second),
		redisClient: redisClient,
		httpClient:  httpClient,
		stopc:       make(chan struct{}),
	}
	go r.run()
	return &r
}

启动与停止

时间轮常驻 goroutine 运行流程同样通过 for + select 的形式运行：

接收到 stopc 信号时，goroutine 退出，时间轮停止运行
接收到 ticker 信号时，开启一个异步 goroutine 用于执行当前批次的定时任务

// 运行时间轮
func (r *RTimeWheel) run() {
    // 通过 for + select 的代码结构运行一个常驻 goroutine 是常规操作
    for {
        select {
        // 接收到终止信号，则退出 goroutine
        case <-r.stopc:
            return
        // 每次接收到来自定时器的信号，则批量扫描并执行定时任务
        case <-r.ticker.C:
            // 每次 tick 获取任务
            go r.executeTasks()
        }
    }
}

停止时间轮的 Stop 方法通过关闭 stopc 保证常驻 goroutine 能够及时退出.

// 停止时间轮
func (r *RTimeWheel) Stop() {
    // 基于单例工具，保证 stopc 只能被关闭一次
    r.Do(func() {
        // 关闭 stopc，使得常驻 goroutine 停止运行
        close(r.stopc)
        // 终止定时器 ticker
        r.ticker.Stop()
    })
}

创建任务

在创建定时任务时，每笔定时任务需要根据其执行的时间找到从属的分钟时间片.

定时任务真正的存储逻辑定义在一段 lua 脚本中，通过 redis 客户端的 Eval 方法执行.

// 添加定时任务
func (r *RTimeWheel) AddTask(ctx context.Context, key string, task *RTaskElement, executeAt time.Time) error {
    // 前置对定时任务的参数进行校验
    if err := r.addTaskPrecheck(task); err != nil {
        return err
    }


    task.Key = key
    // 将定时任务序列化成字节数组
    taskBody, _ := json.Marshal(task)
    // 通过执行 lua 脚本，实现将定时任务添加 redis zset 中. 本质上底层使用的是 zadd 指令.
    _, err := r.redisClient.Eval(ctx, LuaAddTasks, 2, []interface{}{
        // 分钟级 zset 时间片
        r.getMinuteSlice(executeAt),
        // 标识任务删除的集合
        r.getDeleteSetKey(executeAt),
        // 以执行时刻的秒级时间戳作为 zset 中的 score
        executeAt.Unix(),
        // 任务明细
        string(taskBody),
        // 任务 key，用于存放在删除集合中
        key,
    })
    return err
}

//使用示例
if err := rTimeWheel.AddTask(ctx, "test1", &RTaskElement{
		CallbackURL: callbackURL,
		Method:      callbackMethod,
		Req:         callbackReq,
		Header:      callbackHeader,
}, time.Now().Add(time.Second)); err != nil {
		t.Error(err)
		return
}

// 1 添加任务时，如果存在删除 key 的标识，则将其删除
// 添加任务时，根据时间（所属的 min）决定数据从属于哪个分片{}
LuaAddTasks = `
   local zsetKey = KEYS[1]
   local deleteSetKey = KEYS[2]
   local score = ARGV[1]
   local task = ARGV[2]
   local taskKey = ARGV[3]
   redis.call('srem',deleteSetKey,taskKey)
   return redis.call('zadd',zsetKey,score,task)
`

下面展示的是获取分钟级定时任务有序表 minuteSlice 以及已删除任务集合 deleteSet 的细节.
我们首先看一下addTaskPrecheck这个函数，是对task参数对校验。

func (r *RTimeWheel) addTaskPrecheck(task *RTaskElement) error {
	if task.Method != http.MethodGet && task.Method != http.MethodPost {
		return fmt.Errorf("invalid method: %s", task.Method)
	}
	if !strings.HasPrefix(task.CallbackURL, "http://") && !strings.HasPrefix(task.CallbackURL, "https://") {
		return fmt.Errorf("invalid url: %s", task.CallbackURL)
	}
	return nil
}

现在看一下getMinuteSlice，获取定时任务有序表 key 的方法：

func (r *RTimeWheel) getMinuteSlice(executeAt time.Time) string {
	return fmt.Sprintf("xiaoxu_timewheel_task_{%s}", util.GetTimeMinuteStr(executeAt))
}

func GetTimeMinuteStr(t time.Time) string {
	return t.Format(YYYY_MM_DD_HH_MM)
}

例如生成的key是xiaoxu_timewheel_task_{2023-11-05-15:08}；

获取删除任务集合 key 的方法：

func (r *RTimeWheel) getDeleteSetKey(executeAt time.Time) string {
    return fmt.Sprintf("xiaoxu_timewheel_delset_{%s}", util.GetTimeMinuteStr(executeAt))
}

现在我们看一下Lua脚本

type Client struct {
	opts *ClientOptions
	pool *redis.Pool
}
// Eval 支持使用 lua 脚本.
func (c *Client) Eval(ctx context.Context, src string, keyCount int, keysAndArgs []interface{}) (interface{}, error) {
	args := make([]interface{}, 2+len(keysAndArgs))
	args[0] = src
	args[1] = keyCount
	copy(args[2:], keysAndArgs)

	conn, err := c.pool.GetContext(ctx)
	if err != nil {
		return -1, err
	}
	defer conn.Close()

	return conn.Do("EVAL", args...)
}


// 1 添加任务时，如果存在删除 key 的标识，则将其删除
// 添加任务时，根据时间（所属的 min）决定数据从属于哪个分片{}
LuaAddTasks = `
   local zsetKey = KEYS[1]
   local deleteSetKey = KEYS[2]
   local score = ARGV[1]
   local task = ARGV[2]
   local taskKey = ARGV[3]
   redis.call('srem',deleteSetKey,taskKey)
   return redis.call('zadd',zsetKey,score,task)
`

这段Go代码定义了一个名为 Eval 的方法，这个方法使得 Go 客户端能够通过 Redis 连接执行 Lua 脚本。Eval 方法是如何工作的，以及它与 Lua 脚本 LuaAddTasks 是如何配合使用的，我们可以逐步解析如下：

Eval 方法：

这个方法属于 Client 类型，接受一个上下文（context.Context），Lua 脚本的源代码（src 字符串），键的数量（keyCount 整数），以及一个包含键和参数的切片（keysAndArgs []interface{}）。
方法的开始，首先初始化一个足够大的切片 args 来存储 Lua 脚本的源码、键的数量以及所有的键和参数。
然后尝试从连接池 c.pool 中获取一个连接，并处理可能出现的错误。如果无法获取连接，则返回错误。
在连接使用完毕后，通过 defer 声明确保连接最终会关闭。
使用获得的连接执行 Redis 的 EVAL 命令，传入前面构造的 args 切片。

Redis 的 EVAL 命令有什么用?

Redis 的 EVAL 命令用于执行 Lua 脚本。Lua 脚本在 Redis 中的执行是原子性的，意味着脚本运行期间，Redis 服务器不会执行任何其他命令，直到该脚本完成。这为用户提供了在一个执行步骤中执行多个命令的能力，这些命令要么全部执行，要么全部不执行，这类似于数据库的事务。
EVAL 命令的基本用法是：

EVAL script numkeys key [key ...] arg [arg ...]

script 是要执行的 Lua 脚本代码。
numkeys 是键的数量，这个参数告诉 Redis 哪些是键参数，哪些是普通参数，以便它可以正确地处理数据分片和脚本缓存。
key [key …] 是传递给脚本的键名，这些键名由 numkeys 参数指定。
arg [arg …] 是传递给脚本的其他参数，这些参数不会被 Redis 当作键来处理。

下面是每部分的详细说明：

“local current = redis.call(‘get’, KEYS[1]) if current then current = redis.call(‘incr’, KEYS[1]) else current = redis.call(‘set’, KEYS[1], 1) end return current” 是 Lua 脚本。

首先，我们使用 redis.call(‘get’, KEYS[1]) 获取 counter 的当前值。
- 如果 counter 存在（即 current 不为 nil），我们执行自增操作 redis.call(‘incr’, KEYS[1])。
- 如果 counter 不存在（即 current 为 nil），我们使用 redis.call(‘set’, KEYS[1], 1) 设置 counter 的值为 1。
- 最后，脚本返回 counter 的新值。
1 是 numkeys 参数，指示给 Lua 脚本的键参数数量。
counter 是键名，这是我们要自增的键。

LuaAddTasks 脚本：

这个 Lua 脚本预计接收两个键和三个参数。
KEYS[1] 是一个有序集合的键（zsetKey），用来存储需要添加的任务。
KEYS[2] 是一个集合的键（deleteSetKey），其中包含需要删除的任务键名。
ARGV[1] 是一个分数（score），在有序集合中用来排序任务。
ARGV[2] 是任务内容（task），这是要添加到有序集合的值。
ARGV[3] 是任务的键名（taskKey），在删除集合中用来指定要删除的任务。
Lua 脚本首先调用 redis.call(‘srem’, deleteSetKey, taskKey) 来从 deleteSetKey 集合中移除指定的 taskKey。
然后，脚本通过 redis.call(‘zadd’, zsetKey, score, task) 将任务 task 与它的分数 score 添加到 zsetKey 的有序集合中，并返回该操作的结果。
当客户端想要添加一个新任务时，它可以使用 Eval 方法执行 LuaAddTasks 脚本。如果添加任务时存在要删除的键，那么 Lua 脚本首先会处理这个删除操作，接着再添加新任务到对应的有序集合中。这种方式是原子性的，也就是说，删除和添加操作要么都发生，要么都不发生，这是利用 Lua 脚本操作 Redis 的一大优势。

下面展示一下创建定时任务流程中 lua 脚本的执行逻辑：

在这里插入图片描述

删除任务

删除定时任务的方式是将定时任务追加到分钟级的已删除任务 set 中. 之后在检索定时任务时，会根据这个 set 对定时任务进行过滤，实现惰性删除机制.

// 从 redis 时间轮中删除一个定时任务
func (r *RTimeWheel) RemoveTask(ctx context.Context, key string, executeAt time.Time) error {
    // 执行 lua 脚本，将被删除的任务追加到 set 中.
    _, err := r.redisClient.Eval(ctx, LuaDeleteTask, 1, []interface{}{
        r.getDeleteSetKey(executeAt),
        key,
    })
    return err
}

const(    
    // 删除定时任务 lua 脚本
    LuaDeleteTask = `
       -- 获取标识删除任务的 set 集合的 key
       local deleteSetKey = KEYS[1]
       -- 获取定时任务的唯一键
       local taskKey = ARGV[1]
       -- 将定时任务唯一键添加到 set 中
       redis.call('sadd',deleteSetKey,taskKey)
       -- 倘若是 set 中的首个元素，则对 set 设置 120 s 的过期时间
       local scnt = redis.call('scard',deleteSetKey)
       if (tonumber(scnt) == 1)
       then
           redis.call('expire',deleteSetKey,120)
       end
       return scnt
)    `

在这里插入图片描述

执行定时任务

在执行定时任务时，会通过 getExecutableTasks 方法批量获取到满足执行条件的定时任务 list，然后并发调用 execute 方法完成定时任务的回调执行.

// 批量执行定时任务
func (r *RTimeWheel) executeTasks() {
    defer func() {
        if err := recover(); err != nil {
            // log
        }
    }()
    // 并发控制，保证 30 s 之内完成该批次全量任务的执行，及时回收 goroutine，避免发生 goroutine 泄漏
    tctx, cancel := context.WithTimeout(context.Background(), time.Second*30)
    defer cancel()
    // 根据当前时间条件扫描 redis zset，获取所有满足执行条件的定时任务
    tasks, err := r.getExecutableTasks(tctx)
    if err != nil {
        // log
        return
    }
    // 并发执行任务，通过 waitGroup 进行聚合收口
    var wg sync.WaitGroup
    for _, task := range tasks {
        wg.Add(1)
        // shadow
        task := task
        go func() {
            defer func() {
                if err := recover(); err != nil {
                }
                wg.Done()
            }()
            // 执行定时任务
            if err := r.executeTask(tctx, task); err != nil {
                // log
            }
        }()
    }
    wg.Wait()
}

这个 Lua 脚本是为了在 Redis 中处理和管理集合（Set）类型的数据，并且带有某种形式的过期时间管理。具体步骤如下：

local deleteSetKey = KEYS[1]: 将 Lua 脚本中的第一个键参数赋值给变量 deleteSetKey。这里 KEYS 是一个从 EVAL 命令传入的参数数组，代表了键的名称。在 Redis 的 Lua 脚本中，KEYS 数组用于传递键名参数。
local taskKey = ARGV[1]: 将脚本的第一个非键参数赋值给变量 taskKey。在 EVAL 命令中，ARGV 数组用于传递除了键之外的其他参数。
redis.call(‘sadd’, deleteSetKey, taskKey): 使用 sadd 命令将 taskKey 添加到名为 deleteSetKey 的集合中。如果 taskKey 已经是集合的成员，则该命令不做任何操作。如果成功添加了新元素，它会返回 1。
local scnt = redis.call(‘scard’, deleteSetKey): 获取名为 deleteSetKey 的集合的成员数量，并将这个数量赋值给变量 scnt。
if (tonumber(scnt) == 1) then: 判断 deleteSetKey 集合中元素的数量是否为 1。Lua 脚本中，tonumber 函数用于确保 scnt 的值被当作数字处理。
redis.call(‘expire’, deleteSetKey, 120): 如果 deleteSetKey 的集合只有一个成员（即刚添加的 taskKey），则设置该集合的过期时间为 120 秒。expire 命令用于设置键的生存时间（TTL）。
return scnt: 脚本返回 deleteSetKey 集合的成员数量。

检索定时任务

最后介绍一下，如何根据当前时间获取到满足执行条件的定时任务列表：

每次检索时，首先根据当前时刻，推算出所从属的分钟级时间片
然后获得当前的秒级时间戳，作为 zrange 指令检索的 score 范围
调用 lua 脚本，同时获取到已删除任务 set 以及 score 范围内的定时任务 list.
通过 set过滤掉被删除的任务，然后返回满足执行条件的定时任务

func (r *RTimeWheel) getExecutableTasks(ctx context.Context) ([]*RTaskElement, error) {
    now := time.Now()
    // 根据当前时间，推算出其从属的分钟级时间片
    minuteSlice := r.getMinuteSlice(now)
    // 推算出其对应的分钟级已删除任务集合
    deleteSetKey := r.getDeleteSetKey(now)
    nowSecond := util.GetTimeSecond(now)
    // 以秒级时间戳作为 score 进行 zset 检索
    score1 := nowSecond.Unix()
    score2 := nowSecond.Add(time.Second).Unix()
    // 执行 lua 脚本，本质上是通过 zrange 指令结合秒级时间戳对应的 score 进行定时任务检索
    rawReply, err := r.redisClient.Eval(ctx, LuaZrangeTasks, 2, []interface{}{
        minuteSlice, deleteSetKey, score1, score2,
    })
    if err != nil {
        return nil, err
    }
    // 结果中，首个元素对应为已删除任务的 key 集合，后续元素对应为各笔定时任务
    replies := gocast.ToInterfaceSlice(rawReply)
    if len(replies) == 0 {
        return nil, fmt.Errorf("invalid replies: %v", replies)
    }
    deleteds := gocast.ToStringSlice(replies[0])
    //获取删除元素集合
    deletedSet := make(map[string]struct{}, len(deleteds))
    for _, deleted := range deleteds {
        deletedSet[deleted] = struct{}{}
    }
    // 遍历各笔定时任务，倘若其存在于删除集合中，则跳过，否则追加到 list 中返回，用于后续执行
    tasks := make([]*RTaskElement, 0, len(replies)-1)
    for i := 1; i < len(replies); i++ {
        var task RTaskElement
        if err := json.Unmarshal([]byte(gocast.ToString(replies[i])), &task); err != nil {
            // log
            continue
        }


        if _, ok := deletedSet[task.Key]; ok {
            continue
        }
        tasks = append(tasks, &task)
    }
    return tasks, nil
}

lua 脚本的执行逻辑如下：

(    
    // 扫描 redis 时间轮. 获取分钟范围内,已删除任务集合 以及在时间上达到执行条件的定时任务进行返回
    LuaZrangeTasks = `
       -- 第一个 key 为存储定时任务的 zset key
       local zsetKey = KEYS[1]
       -- 第二个 key 为已删除任务 set 的 key
       local deleteSetKey = KEYS[2]
       -- 第一个 arg 为 zrange 检索的 score 左边界
       local score1 = ARGV[1]
       -- 第二个 arg 为 zrange 检索的 score 右边界
       local score2 = ARGV[2]
       -- 获取到已删除任务的集合
       local deleteSet = redis.call('smembers',deleteSetKey)
       -- 根据秒级时间戳对 zset 进行 zrange 检索，获取到满足时间条件的定时任务
       local targets = redis.call('zrange',zsetKey,score1,score2,'byscore')
       -- 检索到的定时任务直接从时间轮中移除，保证分布式场景下定时任务不被重复获取
       redis.call('zremrangebyscore',zsetKey,score1,score2)
       -- 返回的结果是一个 table
       local reply = {}
       -- table 的首个元素为已删除任务集合
       reply[1] = deleteSet
       -- 依次将检索到的定时任务追加到 table 中
       for i, v in ipairs(targets) do
           reply[#reply+1]=v
       end
       return reply
    `
)

在这里插入图片描述