go-omni/task/task.go

// Package task provides a bounded in-process worker pool and interval jobs.
// Jobs are not durable and are not coordinated across service replicas.
package task

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"sync/atomic"
	"time"

	"git.yuntuo.hk/k2safe/go-omni/logx"
)

var ErrClosed = errors.New("task: closed")
var ErrQueueFull = errors.New("task: queue full")

type Func func(context.Context) error
type Config struct {
	Workers         int                 `json:"workers"`
	QueueSize       int                 `json:"queueSize"`
	MaxSchedules    int                 `json:"maxSchedules"`
	TaskTimeout     time.Duration       `json:"taskTimeout"`
	ShutdownTimeout time.Duration       `json:"shutdownTimeout"`
	OnError         func(string, error) `json:"-"`
}

type Future struct {
	done chan struct{}
	err  error
}

func (f *Future) Wait(ctx context.Context) error {
	if ctx == nil {
		return errors.New("task: nil context")
	}
	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-f.done:
		return f.err
	}
}
func (f *Future) Done() <-chan struct{} { return f.done }

type job struct {
	ctx    context.Context
	name   string
	fn     Func
	future *Future
}

type Runner struct {
	config    Config
	mu        sync.Mutex
	queue     chan job
	closed    bool
	schedules map[string]struct{}
	stop      chan struct{}
	done      chan struct{}
	wg        sync.WaitGroup
	ctx       context.Context
	cancel    context.CancelFunc
	active    atomic.Int64
	completed atomic.Uint64
	failed    atomic.Uint64
	skipped   atomic.Uint64
}

type Stats struct {
	Active    int64  `json:"active"`
	Queued    int    `json:"queued"`
	Completed uint64 `json:"completed"`
	Failed    uint64 `json:"failed"`
	Skipped   uint64 `json:"skipped"`
}

func New(cfg Config) (*Runner, error) {
	if cfg.Workers < 0 || cfg.QueueSize < 0 || cfg.MaxSchedules < 0 || cfg.TaskTimeout < 0 || cfg.ShutdownTimeout < 0 {
		return nil, errors.New("task: invalid configuration")
	}
	if cfg.Workers == 0 {
		cfg.Workers = 4
	}
	if cfg.QueueSize == 0 {
		cfg.QueueSize = 128
	}
	if cfg.MaxSchedules == 0 {
		cfg.MaxSchedules = 128
	}
	if cfg.TaskTimeout == 0 {
		cfg.TaskTimeout = 30 * time.Second
	}
	if cfg.ShutdownTimeout == 0 {
		cfg.ShutdownTimeout = 10 * time.Second
	}
	ctx, cancel := context.WithCancel(context.Background())
	r := &Runner{config: cfg, queue: make(chan job, cfg.QueueSize), schedules: make(map[string]struct{}), stop: make(chan struct{}), done: make(chan struct{}), ctx: ctx, cancel: cancel}
	r.wg.Add(cfg.Workers)
	for i := 0; i < cfg.Workers; i++ {
		go r.worker()
	}
	return r, nil
}

// Submit never blocks on a full queue. The job retains the caller's context;
// use a service-owned context for work intended to outlive an HTTP request.
func (r *Runner) Submit(ctx context.Context, name string, fn Func) (*Future, error) {
	if ctx == nil || fn == nil || name == "" || len(name) > 256 {
		return nil, errors.New("task: context, function and bounded name are required")
	}
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	if r.closed {
		return nil, ErrClosed
	}
	f := &Future{done: make(chan struct{})}
	select {
	case r.queue <- job{ctx: ctx, name: name, fn: fn, future: f}:
		return f, nil
	default:
		return nil, ErrQueueFull
	}
}

// Every schedules at a fixed interval, with no initial execution. A tick is
// skipped if this schedule's previous job is queued/running or the pool is full.
func (r *Runner) Every(name string, interval time.Duration, fn Func) error {
	if name == "" || len(name) > 256 || interval <= 0 || fn == nil {
		return errors.New("task: invalid schedule")
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	if r.closed {
		return ErrClosed
	}
	if _, ok := r.schedules[name]; ok {
		return errors.New("task: duplicate schedule name")
	}
	if len(r.schedules) >= r.config.MaxSchedules {
		return errors.New("task: schedule capacity reached")
	}
	r.schedules[name] = struct{}{}
	r.wg.Add(1)
	go func() {
		defer r.wg.Done()
		ticker := time.NewTicker(interval)
		defer ticker.Stop()
		var previous *Future
		for {
			select {
			case <-r.stop:
				return
			case <-ticker.C:
				if previous != nil {
					select {
					case <-previous.Done():
						previous = nil
					default:
						r.skipped.Add(1)
						continue
					}
				}
				var err error
				previous, err = r.Submit(r.ctx, name, fn)
				if err != nil {
					r.skipped.Add(1)
				}
			}
		}
	}()
	return nil
}

func (r *Runner) worker() {
	defer r.wg.Done()
	for item := range r.queue {
		r.active.Add(1)
		ctx, cancel := context.WithTimeout(item.ctx, r.config.TaskTimeout)
		stop := context.AfterFunc(r.ctx, cancel)
		if r.ctx.Err() != nil {
			cancel()
		}
		err := run(ctx, item.fn)
		// AfterFunc cancellation may still be queued when the job returns.
		// Check the pool directly before publishing a successful outcome.
		if err == nil {
			err = r.ctx.Err()
		}
		stop()
		cancel()
		if err != nil {
			r.failed.Add(1)
			func() {
				defer func() { _ = recover() }()
				if r.config.OnError != nil {
					r.config.OnError(item.name, err)
				} else {
					logx.ErrorContext(item.ctx, "task.failed", "task", item.name, "error", err)
				}
			}()
		}
		r.active.Add(-1)
		r.completed.Add(1)
		item.future.err = err
		close(item.future.done)
	}
}

func run(ctx context.Context, fn Func) (err error) {
	defer func() {
		if value := recover(); value != nil {
			err = fmt.Errorf("task: panic: %v", value)
		}
	}()
	if err = ctx.Err(); err != nil {
		return err
	}
	err = fn(ctx)
	if err == nil {
		err = ctx.Err()
	}
	return err
}

// StopAccepting stops schedules and submission without canceling accepted jobs.
// Pair it with Wait to establish completion before releasing dependencies.
func (r *Runner) StopAccepting(ctx context.Context) error {
	if ctx == nil {
		return errors.New("task: nil context")
	}
	r.mu.Lock()
	if !r.closed {
		r.closed = true
		close(r.stop)
		close(r.queue)
		go func() { r.wg.Wait(); r.cancel(); close(r.done) }()
	}
	r.mu.Unlock()
	return nil
}

// Wait joins all workers and schedules after StopAccepting. A canceled context
// cancels work, but callers can wait again to establish actual completion.
func (r *Runner) Wait(ctx context.Context) error {
	if ctx == nil {
		return errors.New("task: nil context")
	}
	select {
	case <-r.done:
		return nil
	default:
	}
	select {
	case <-r.done:
		return nil
	case <-ctx.Done():
		r.cancel()
		return ctx.Err()
	}
}

// Shutdown stops admission and waits with the caller's deadline.
func (r *Runner) Shutdown(ctx context.Context) error {
	if err := r.StopAccepting(ctx); err != nil {
		return err
	}
	return r.Wait(ctx)
}

func (r *Runner) Close() error {
	ctx, cancel := context.WithTimeout(context.Background(), r.config.ShutdownTimeout)
	defer cancel()
	return r.Shutdown(ctx)
}

func (r *Runner) Stats() Stats {
	return Stats{Active: r.active.Load(), Queued: len(r.queue), Completed: r.completed.Load(), Failed: r.failed.Load(), Skipped: r.skipped.Load()}
}