taskq.go

package taskq

import (
	"context"
	"errors"
	"runtime"
	"sync/atomic"
	"time"
)

var (
	ErrStarted = errors.New("taskq started")
	ErrClosed  = errors.New("taskq closed")
	ErrNilTask = errors.New("nil task")
)

type worker struct {
	id uint64
}

type TaskQ struct {
	queue Queue

	isRunning int32
	isClosed  int32
	isStopped int32

	workerCount int32
	workers     chan worker

	// TaskContextFunc is a function that returns a context for a task.Do execution function.
	TaskContextFunc func(t Task) context.Context
	// OnDequeueContextFunc is a function that returns a context for Queue.Dequeue function.
	DequeueContextFunc func() context.Context
	// OnEnqueueError is a function that is called when an error occurs during the Dequeue operation.
	OnDequeueError func(ctx context.Context, workerID uint64, err error)
}

func New(limit int) *TaskQ {
	return NewWithQueue(limit, NewConcurrentQueue())
}

func NewWithQueue(limit int, q Queue) *TaskQ {
	if limit <= 0 {
		limit = runtime.NumCPU()
	}
	// init worker pool
	workers := make(chan worker, limit)
	for i := 1; i < limit+1; i++ {
		workers <- worker{
			id: uint64(i),
		}
	}
	return &TaskQ{
		queue:          q,
		isRunning:      0,
		isClosed:       0,
		isStopped:      0,
		workers:        workers,
		OnDequeueError: nil,
	}
}

func (t *TaskQ) triggerDequeue() bool {
	if atomic.LoadInt32(&t.isRunning) != 1 {
		return false
	}
	select {
	// trying to fetch free worker for immediate task execution
	case w, ok := <-t.workers:
		if !ok {
			return false
		}
		atomic.AddInt32(&t.workerCount, 1)
		go func(w worker) {
			var qCtx context.Context
			for atomic.LoadInt32(&t.isStopped) != 1 {
				if t.DequeueContextFunc != nil {
					qCtx = t.DequeueContextFunc()
				} else {
					qCtx = context.Background()
				}
				task, err := t.queue.Dequeue(qCtx)
				if err != nil {
					if err == EmptyQueue {
						break
					}
					if t.OnDequeueError == nil {
						panic(err)
					}
					t.OnDequeueError(qCtx, w.id, err)
					break
				}
				var tCtx context.Context
				if t.TaskContextFunc != nil {
					tCtx = t.TaskContextFunc(task)
				} else {
					tCtx = context.Background()
				}
				processTask(tCtx, task)
			}
			t.workers <- w // return worker to pool
			atomic.AddInt32(&t.workerCount, -1)
		}(w)
		return true
	default:
		return false
	}
}

func (t *TaskQ) Enqueue(ctx context.Context, task Task) (int64, error) {
	if task == nil {
		return -1, ErrNilTask
	}

	if atomic.LoadInt32(&t.isClosed) != 0 {
		return -1, ErrClosed
	}

	id, err := t.queue.Enqueue(ctx, task)
	if err != nil {
		return -1, err
	}

	t.triggerDequeue()
	return id, nil
}

func (t *TaskQ) triggerFreeWorkers() {
	count := len(t.workers)
	for i := 0; i < count; i++ {
		if !t.triggerDequeue() {
			return
		}
	}
}

func (t *TaskQ) Start() error {
	if atomic.LoadInt32(&t.isClosed) != 0 {
		return ErrClosed
	}
	if !atomic.CompareAndSwapInt32(&t.isRunning, 0, 1) {
		return ErrStarted
	}
	t.triggerFreeWorkers()
	return nil
}

func (t *TaskQ) Shutdown(ctx context.Context) error {
	if !atomic.CompareAndSwapInt32(&t.isClosed, 0, 1) {
		return ErrClosed
	}
	defer close(t.workers)
	if wait, _ := ctx.Value(ctxWaitKey{}).(bool); !wait {
		atomic.StoreInt32(&t.isStopped, 1)
	} else {
		t.triggerFreeWorkers()
	}

	var pollDuration = time.Millisecond * 500
	timer := time.NewTimer(pollDuration)
	for atomic.LoadInt32(&t.workerCount) > 0 {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-timer.C:
			const delta = 1.1
			timer.Reset(time.Duration(float64(pollDuration) * delta))
		}
	}
	return nil
}

func (t *TaskQ) Close() error {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	return t.Shutdown(ctx)
}