enqueue.go

package work

import (
	"sync"
	"time"

	"github.com/garyburd/redigo/redis"
)

// Enqueuer can enqueue jobs.
type Enqueuer struct {
	Namespace string // eg, "myapp-work"
	Pool      *redis.Pool

	queuePrefix           string // eg, "myapp-work:jobs:"
	knownJobs             map[string]int64
	enqueueUniqueScript   *redis.Script
	enqueueUniqueInScript *redis.Script
	mtx                   sync.RWMutex
}

// NewEnqueuer creates a new enqueuer with the specified Redis namespace and Redis pool.
func NewEnqueuer(namespace string, pool *redis.Pool) *Enqueuer {
	if pool == nil {
		panic("NewEnqueuer needs a non-nil *redis.Pool")
	}

	return &Enqueuer{
		Namespace:             namespace,
		Pool:                  pool,
		queuePrefix:           redisKeyJobsPrefix(namespace),
		knownJobs:             make(map[string]int64),
		enqueueUniqueScript:   redis.NewScript(3, redisLuaEnqueueUnique),
		enqueueUniqueInScript: redis.NewScript(3, redisLuaEnqueueUniqueIn),
	}
}

// Enqueue will enqueue the specified job name and arguments. The args param can be nil if no args ar needed.
// Example: e.Enqueue("send_email", work.Q{"addr": "test@example.com"})
func (e *Enqueuer) Enqueue(jobName string, args map[string]interface{}) (*Job, error) {
	job := &Job{
		Name:       jobName,
		ID:         makeIdentifier(),
		EnqueuedAt: nowEpochSeconds(),
		Args:       args,
	}

	rawJSON, err := job.serialize()
	if err != nil {
		return nil, err
	}

	conn := e.Pool.Get()
	defer conn.Close()

	if _, err := conn.Do("LPUSH", e.queuePrefix+jobName, rawJSON); err != nil {
		return nil, err
	}

	if err := e.addToKnownJobs(conn, jobName); err != nil {
		return job, err
	}

	return job, nil
}

// EnqueueIn enqueues a job in the scheduled job queue for execution in secondsFromNow seconds.
func (e *Enqueuer) EnqueueIn(jobName string, secondsFromNow int64, args map[string]interface{}) (*ScheduledJob, error) {
	job := &Job{
		Name:        jobName,
		ID:          makeIdentifier(),
		EnqueuedAt:  nowEpochSeconds(),
		Args:        args,
		ScheduledAt: nowEpochSeconds() + secondsFromNow,
	}

	rawJSON, err := job.serialize()
	if err != nil {
		return nil, err
	}

	conn := e.Pool.Get()
	defer conn.Close()

	scheduledJob := &ScheduledJob{
		RunAt: nowEpochSeconds() + secondsFromNow,
		Job:   job,
	}

	_, err = conn.Do("ZADD", redisKeyScheduled(e.Namespace), scheduledJob.RunAt, rawJSON)
	if err != nil {
		return nil, err
	}

	if err := e.addToKnownJobs(conn, jobName); err != nil {
		return scheduledJob, err
	}

	return scheduledJob, nil
}

type EnqueueOp struct {
	Expire    int
	UniqueKey string
}

type EnqueueOption func(*EnqueueOp)

// WithExpireTime sets the expire time(in seconds) of the uniqueness of an unique job.
// The default expire time is 10 minutes.
// It is ignored if expire time <= 0.
func WithExpireTime(expire int) EnqueueOption {
	return func(op *EnqueueOp) {
		op.Expire = expire
	}
}

// WithUniqueKey sets customized unique keys used to distinguish unique jobs.
// By default, an unique is generated by all args of a job.
func WithUniqueKey(k string) EnqueueOption {
	return func(op *EnqueueOp) {
		op.UniqueKey = k
	}
}

const expireTime = 600

// EnqueueUnique enqueues a job unless a job is already enqueued with the same name and arguments.
// The already-enqueued job can be in the normal work queue or in the scheduled job queue.
// Only if a job is processed, another job with the same name and arguments can be enqueued again.
// Any failed jobs in the retry queue or dead queue don't count against the uniqueness -- so if a job fails and is retried, two unique jobs with the same name and arguments can be enqueued at once.
// In order to add robustness to the system, jobs are only unique for 10 minutes after they're enqueued. This is mostly relevant for scheduled jobs.
// EnqueueUnique returns the job if it was enqueued and nil if it wasn't
func (e *Enqueuer) EnqueueUnique(jobName string, args map[string]interface{}, opts ...EnqueueOption) (*Job, error) {
	op := &EnqueueOp{}
	for _, opt := range opts {
		opt(op)
	}
	expire := expireTime
	if op.Expire > 0 {
		expire = op.Expire
	}
	args, uniqueKey, err := uniqueKey(op, e.Namespace, jobName, args)
	if err != nil {
		return nil, err
	}

	job := &Job{
		Name:       jobName,
		ID:         makeIdentifier(),
		EnqueuedAt: nowEpochSeconds(),
		Args:       args,
		Unique:     true,
	}

	rawJSON, err := job.serialize()
	if err != nil {
		return nil, err
	}

	conn := e.Pool.Get()
	defer conn.Close()

	if err := e.addToKnownJobs(conn, jobName); err != nil {
		return nil, err
	}

	scriptArgs := make([]interface{}, 0, 3)
	scriptArgs = append(scriptArgs, e.queuePrefix+jobName) // KEY[1]
	scriptArgs = append(scriptArgs, uniqueKey)             // KEY[2]
	scriptArgs = append(scriptArgs, expire)                // KEY[3]
	scriptArgs = append(scriptArgs, rawJSON)               // ARGV[1]

	res, err := redis.String(e.enqueueUniqueScript.Do(conn, scriptArgs...))
	if res == "ok" && err == nil {
		return job, nil
	}
	return nil, err
}

// EnqueueUniqueIn enqueues a unique job in the scheduled job queue for execution in secondsFromNow seconds. See EnqueueUnique for the semantics of unique jobs.
func (e *Enqueuer) EnqueueUniqueIn(jobName string, secondsFromNow int64, args map[string]interface{}, opts ...EnqueueOption) (*ScheduledJob, error) {
	op := &EnqueueOp{}
	for _, opt := range opts {
		opt(op)
	}
	expire := expireTime
	if op.Expire > 0 {
		expire = op.Expire
	}

	args, uniqueKey, err := uniqueKey(op, e.Namespace, jobName, args)
	if err != nil {
		return nil, err
	}

	job := &Job{
		Name:        jobName,
		ID:          makeIdentifier(),
		EnqueuedAt:  nowEpochSeconds(),
		Args:        args,
		Unique:      true,
		ScheduledAt: nowEpochSeconds() + secondsFromNow,
	}

	rawJSON, err := job.serialize()
	if err != nil {
		return nil, err
	}

	conn := e.Pool.Get()
	defer conn.Close()

	if err := e.addToKnownJobs(conn, jobName); err != nil {
		return nil, err
	}

	scheduledJob := &ScheduledJob{
		RunAt: nowEpochSeconds() + secondsFromNow,
		Job:   job,
	}

	scriptArgs := make([]interface{}, 0, 4)
	scriptArgs = append(scriptArgs, redisKeyScheduled(e.Namespace)) // KEY[1]
	scriptArgs = append(scriptArgs, uniqueKey)                      // KEY[2]
	scriptArgs = append(scriptArgs, expire)                         // KEY[3]
	scriptArgs = append(scriptArgs, rawJSON)                        // ARGV[1]
	scriptArgs = append(scriptArgs, scheduledJob.RunAt)             // ARGV[2]

	res, err := redis.String(e.enqueueUniqueInScript.Do(conn, scriptArgs...))

	if res == "ok" && err == nil {
		return scheduledJob, nil
	}
	return nil, err
}

const UniqueKeyArg = "unique_job_key"

func uniqueKey(op *EnqueueOp, namespace, jobName string, args map[string]interface{}) (map[string]interface{}, string, error) {
	if len(op.UniqueKey) > 0 {
		if args == nil {
			args = make(map[string]interface{}, 1)
		}
		args[UniqueKeyArg] = op.UniqueKey
	}
	k, err := redisKeyUniqueJob(namespace, jobName, args)
	return args, k, err
}

func (e *Enqueuer) addToKnownJobs(conn redis.Conn, jobName string) error {
	needSadd := true
	now := time.Now().Unix()

	e.mtx.RLock()
	t, ok := e.knownJobs[jobName]
	e.mtx.RUnlock()

	if ok {
		if now < t {
			needSadd = false
		}
	}
	if needSadd {
		if _, err := conn.Do("SADD", redisKeyKnownJobs(e.Namespace), jobName); err != nil {
			return err
		}

		e.mtx.Lock()
		e.knownJobs[jobName] = now + 300
		e.mtx.Unlock()
	}

	return nil
}