feat: add cache, outbound HTTP, request protection, tasks and module scaffolding
This commit is contained in:
262
task/task.go
Normal file
262
task/task.go
Normal file
@@ -0,0 +1,262 @@
|
||||
// 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
|
||||
}
|
||||
|
||||
// Shutdown stops schedules and submission, then drains accepted jobs. Deadline
|
||||
// expiry cancels running/queued jobs and returns without waiting for code that
|
||||
// ignores cancellation. Such code cannot be forcibly stopped by Go.
|
||||
func (r *Runner) Shutdown(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()
|
||||
select {
|
||||
case <-r.done:
|
||||
return nil
|
||||
case <-ctx.Done():
|
||||
r.cancel()
|
||||
return ctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
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()}
|
||||
}
|
||||
102
task/task_test.go
Normal file
102
task/task_test.go
Normal file
@@ -0,0 +1,102 @@
|
||||
package task
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestCapacityDrainPanicAndClose(t *testing.T) {
|
||||
r, err := New(Config{Workers: 1, QueueSize: 1, OnError: func(string, error) { panic("callback failure") }})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer r.Close()
|
||||
started, release := make(chan struct{}), make(chan struct{})
|
||||
first, _ := r.Submit(context.Background(), "first", func(context.Context) error { close(started); <-release; return nil })
|
||||
<-started
|
||||
second, _ := r.Submit(context.Background(), "panic", func(context.Context) error { panic("job failure") })
|
||||
if _, err := r.Submit(context.Background(), "overflow", func(context.Context) error { return nil }); !errors.Is(err, ErrQueueFull) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
done := make(chan error, 1)
|
||||
go func() { done <- r.Close() }()
|
||||
close(release)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
if err := first.Wait(ctx); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := second.Wait(ctx); err == nil {
|
||||
t.Fatal("panic lost")
|
||||
}
|
||||
if err := <-done; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if s := r.Stats(); s.Completed != 2 || s.Failed != 1 || s.Active != 0 || s.Queued != 0 {
|
||||
t.Fatal(s)
|
||||
}
|
||||
if _, err := r.Submit(ctx, "late", func(context.Context) error { return nil }); !errors.Is(err, ErrClosed) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestShutdownDeadlineAndQueuedCancellation(t *testing.T) {
|
||||
r, _ := New(Config{Workers: 1, QueueSize: 1, ShutdownTimeout: 20 * time.Millisecond, OnError: func(string, error) {}})
|
||||
started, release := make(chan struct{}), make(chan struct{})
|
||||
var queued atomic.Bool
|
||||
first, _ := r.Submit(context.Background(), "ignoresCancel", func(context.Context) error { close(started); <-release; return nil })
|
||||
<-started
|
||||
second, _ := r.Submit(context.Background(), "queued", func(context.Context) error { queued.Store(true); return nil })
|
||||
if err := r.Close(); !errors.Is(err, context.DeadlineExceeded) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
close(release)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
if err := first.Wait(ctx); !errors.Is(err, context.Canceled) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := second.Wait(ctx); !errors.Is(err, context.Canceled) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if queued.Load() {
|
||||
t.Fatal("ran queued work after forced shutdown")
|
||||
}
|
||||
if err := r.Shutdown(ctx); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestIntervalsDoNotOverlapAndTaskDeadline(t *testing.T) {
|
||||
r, _ := New(Config{Workers: 2, TaskTimeout: 25 * time.Millisecond, OnError: func(string, error) {}})
|
||||
defer r.Close()
|
||||
var active, maxActive atomic.Int32
|
||||
job := func(ctx context.Context) error {
|
||||
n := active.Add(1)
|
||||
if n > maxActive.Load() {
|
||||
maxActive.Store(n)
|
||||
}
|
||||
defer active.Add(-1)
|
||||
<-ctx.Done()
|
||||
return ctx.Err()
|
||||
}
|
||||
if err := r.Every("periodic", 3*time.Millisecond, job); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := r.Every("periodic", time.Second, job); err == nil {
|
||||
t.Fatal("duplicate schedule accepted")
|
||||
}
|
||||
deadline := time.Now().Add(time.Second)
|
||||
for r.Stats().Completed < 2 && time.Now().Before(deadline) {
|
||||
time.Sleep(2 * time.Millisecond)
|
||||
}
|
||||
if err := r.Close(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if maxActive.Load() != 1 || r.Stats().Skipped == 0 || r.Stats().Failed < 2 {
|
||||
t.Fatal(maxActive.Load(), r.Stats())
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user