Build a Delay Queue with Redis

type Job struct {
	Id string `json:"id"`
	Topic string `json:"topic"`
	Delay int64 `json:"delay"`
	TTR int64 `json:"ttr"`  // time to run, guarantee reliability
	Body string `json:"body"` // service data
}

// poolAdd add the job hash
func poolAdd(jobId string, data []byte) error {
	return rdb.HSet(background, PoolKey, jobId, data).Err()
}

// poolRem remove the job hash by job's Id
func poolRem(jobId string) error {
	return rdb.HDel(background, PoolKey, jobId).Err()
}

// poolGet get job meta data by job's Id
// the meta data is transformed into Job struct with unmarshal
func poolGet(jobId string) *Job{
	jobStr, err := rdb.HGet(background, PoolKey, jobId).Result()
	if err != nil {
		return nil
	}
	job := new(Job)
	if err = json.Unmarshal([]byte(jobStr), job); err != nil {
		return nil
	}
	return job
}

func queuePop(ctx context.Context, topic string) (string, error) {
	key := fmt.Sprintf(QueueKeyPrefix, topic)
	vals, err := rdb.BRPop(ctx, 24*time.Hour, key).Result()
	if err != nil {
		return "", err
	}
	return vals[1], nil
}

func queueAdd(topic, jobId string) error {
	queueKey := fmt.Sprintf(QueueKeyPrefix, topic)
	return rdb.LPush(background, queueKey, jobId).Err()
}

func bucketAdd(jobId string, jobDelay int64) error {
	bucketNum := <- bucketNumChan
	key := fmt.Sprintf(BucketKeyPrefix, bucketNum)
	return rdb.ZAdd(background, key,
		&redis.Z{Score: float64(time.Now().Unix()+jobDelay), Member: jobId}).Err()
}

func bucketRem(bucketNum int, jobId string) error {
	key := fmt.Sprintf(BucketKeyPrefix, bucketNum)
	return rdb.ZRem(background, key, jobId).Err()
}

func bucketGetFirst(bucketNum int) *redis.Z {
	key := fmt.Sprintf(BucketKeyPrefix, bucketNum)
	bucketItems := rdb.ZRangeWithScores(background, key, 0, 0).Val()
	if len(bucketItems) == 0 {
		return nil
	}
	return &bucketItems[0]
}

func initTimers() {
	timers = make([]*time.Ticker, BucketSize)
	for i := 0; i < BucketSize; i++ {
		timers[i] = time.NewTicker(1 * time.Second)
		go poolingBucket(timers[i], i)
	}
}

func poolingBucket(ticker *time.Ticker, bucketNum int) {
	for t := range ticker.C {
		item := bucketGetFirst(bucketNum)
		if item == nil || int64(item.Score) > t.Unix() { // does not come the time to run
			continue
		}

		jobId := item.Member.(string)
		job := poolGet(jobId)
		if job == nil { // job is not found in the pool (usually is deleted)
			bucketRem(bucketNum, jobId)
			continue
		}
		if queueAdd(job.Topic, job.Id) != nil {
			continue
		}
		bucketRem(bucketNum, job.Id)
	}
}

syntax = "proto3";

option go_package = "delayqueue/pb";

service DelayQueue {
  rpc Push(PushReq) returns (PushResp);
  rpc Pop(PopReq) returns (PopResp);
  rpc Finish(FinishReq) returns (FinishResp);
}

message PushReq {
  string topic = 1;
  int64 delay = 2;
  int64 ttr = 3;
  string body = 4;
}

message PushResp {
  string id = 1;
}

message PopReq {
  string topic = 1;
}

message PopResp {
  string body = 1;
}

message FinishReq {
  string job_id = 1;
}

message FinishResp {}

package main

import (
	"context"
	"encoding/json"
	"errors"
	"net"
	"time"

	"github.com/go-basic/uuid"
	"github.com/go-redis/redis/v8"
	"google.golang.org/grpc"

	pb "delayqueue/pb"
)

const (
	PoolKey         = "dq:pool"
	QueueKeyPrefix  = "dq:queue:%s"
	BucketSize      = 5
	BucketKeyPrefix = "dq:bucket:%d"
)

var rdb *redis.Client

var timers []*time.Ticker

var bucketNumChan <-chan int

var background = context.Background()

type server struct {
	pb.UnimplementedDelayQueueServer
}

func (s *server) Push(ctx context.Context, req *pb.PushReq) (*pb.PushResp, error) {
	if req.Topic == "" {
		return nil, errors.New("topic is empty")
	}
	if req.Delay < 0 {
		req.Delay = 0
	}
	if req.Ttr < 0 {
		req.Ttr = 0
	}

	job := &Job{
		Id:    uuid.New(),
		Topic: req.Topic,
		Delay: req.Delay,
		TTR:   req.Ttr,
		Body:  req.Body,
	}
	jobBytes, err := json.Marshal(job)
	if err != nil {
		return nil, err
	}

	// 用 lua 脚本写成一个命令，保证操作的原子性
	if err = poolAdd(job.Id, jobBytes); err != nil {
		return nil, err
	}

	if err = bucketAdd(job.Id, job.Delay); err != nil {
		return nil, err
	}
	return &pb.PushResp{Id: job.Id}, nil
}

func (s *server) Finish(ctx context.Context, req *pb.FinishReq) (*pb.FinishResp, error) {
	poolRem(req.JobId)
	return &pb.FinishResp{}, nil
}

func (s *server) Pop(ctx context.Context, req *pb.PopReq) (*pb.PopResp, error) {
	jobId, err := queuePop(ctx, req.Topic)
	if err != nil {
		return nil, err
	}

	job := poolGet(jobId)
	if job == nil {
		return nil, errors.New("no job at present")
	}

	if job.TTR > 0 {
		if err = bucketAdd(job.Id, job.TTR); err != nil {
			return nil, err
		}
	}

	return &pb.PopResp{Body: job.Body}, nil
}

func initBucketNumChan() {
	c := make(chan int)
	go func() {
		i := 1
		for {
			c <- i
			i++
			if i > BucketSize {
				i = 1
			}
		}
	}()
	bucketNumChan = c
}

func main() {
	// connect to redis
	rdb = redis.NewClient(&redis.Options{
		Addr:     "localhost:6379",
		Password: "",
		DB:       0,
	})

	// pooling the buckets
	initTimers()
	initBucketNumChan()

	lis, err := net.Listen("tcp", ":8080")
	if err != nil {
		panic(err)
	}
	s := grpc.NewServer()
	pb.RegisterDelayQueueServer(s, &server{})
	if err = s.Serve(lis); err != nil {
		panic(err)
	}
}