mirror of
https://github.com/yusing/godoxy.git
synced 2026-04-09 18:33:36 +02:00
simplify task package implementation
This commit is contained in:
yusing
@@ -4,355 +4,194 @@ import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"runtime/debug"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/yusing/go-proxy/internal/common"
|
||||
E "github.com/yusing/go-proxy/internal/error"
|
||||
"github.com/yusing/go-proxy/internal/logging"
|
||||
F "github.com/yusing/go-proxy/internal/utils/functional"
|
||||
"github.com/yusing/go-proxy/internal/utils/strutils"
|
||||
)
|
||||
|
||||
var globalTask = createGlobalTask()
|
||||
|
||||
func createGlobalTask() (t *Task) {
|
||||
t = new(Task)
|
||||
t.name = "root"
|
||||
t.ctx, t.cancel = context.WithCancelCause(context.Background())
|
||||
t.subtasks = F.NewSet[*Task]()
|
||||
return
|
||||
}
|
||||
|
||||
func testResetGlobalTask() {
|
||||
globalTask = createGlobalTask()
|
||||
}
|
||||
|
||||
type (
|
||||
TaskStarter interface {
|
||||
// Start starts the object that implements TaskStarter,
|
||||
// and returns an error if it fails to start.
|
||||
//
|
||||
// The task passed must be a subtask of the caller task.
|
||||
//
|
||||
// callerSubtask.Finish must be called when start fails or the object is finished.
|
||||
Start(callerSubtask *Task) E.Error
|
||||
Start(parent Parent) E.Error
|
||||
Task() *Task
|
||||
}
|
||||
TaskFinisher interface {
|
||||
// Finish marks the task as finished and cancel its context.
|
||||
//
|
||||
// Then call Wait to wait for all subtasks, OnFinished and OnSubtasksFinished
|
||||
// of the task to finish.
|
||||
//
|
||||
// Note that it will also cancel all subtasks.
|
||||
Finish(reason any)
|
||||
}
|
||||
// Task controls objects' lifetime.
|
||||
//
|
||||
// Objects that uses a Task should implement the TaskStarter and the TaskFinisher interface.
|
||||
//
|
||||
// When passing a Task object to another function,
|
||||
// it must be a sub-Task of the current Task,
|
||||
// in name of "`currentTaskName`Subtask"
|
||||
//
|
||||
// Use Task.Finish to stop all subtasks of the Task.
|
||||
Task struct {
|
||||
name string
|
||||
|
||||
children sync.WaitGroup
|
||||
|
||||
onFinished sync.WaitGroup
|
||||
finished chan struct{}
|
||||
|
||||
ctx context.Context
|
||||
cancel context.CancelCauseFunc
|
||||
|
||||
parent *Task
|
||||
subtasks F.Set[*Task]
|
||||
subTasksWg sync.WaitGroup
|
||||
|
||||
name string
|
||||
|
||||
OnFinishedFuncs []func()
|
||||
OnFinishedMu sync.Mutex
|
||||
onFinishedWg sync.WaitGroup
|
||||
|
||||
finishOnce sync.Once
|
||||
once sync.Once
|
||||
}
|
||||
Parent interface {
|
||||
Context() context.Context
|
||||
Subtask(name string, needFinish ...bool) *Task
|
||||
Name() string
|
||||
Finish(reason any)
|
||||
OnCancel(name string, f func())
|
||||
}
|
||||
)
|
||||
|
||||
var (
|
||||
ErrProgramExiting = errors.New("program exiting")
|
||||
ErrTaskCanceled = errors.New("task canceled")
|
||||
|
||||
logger = logging.With().Str("module", "task").Logger()
|
||||
)
|
||||
|
||||
// GlobalTask returns a new Task with the given name, derived from the global context.
|
||||
func GlobalTask(format string, args ...any) *Task {
|
||||
if len(args) > 0 {
|
||||
format = fmt.Sprintf(format, args...)
|
||||
}
|
||||
return globalTask.Subtask(format)
|
||||
}
|
||||
|
||||
// DebugTaskMap returns a map[string]any representation of the global task tree.
|
||||
//
|
||||
// The returned map is suitable for encoding to JSON, and can be used
|
||||
// to debug the task tree.
|
||||
//
|
||||
// The returned map is not guaranteed to be stable, and may change
|
||||
// between runs of the program. It is intended for debugging purposes
|
||||
// only.
|
||||
func DebugTaskMap() map[string]any {
|
||||
return globalTask.serialize()
|
||||
}
|
||||
|
||||
// CancelGlobalContext cancels the global task context, which will cause all tasks
|
||||
// created to be canceled. This should be called before exiting the program
|
||||
// to ensure that all tasks are properly cleaned up.
|
||||
func CancelGlobalContext() {
|
||||
globalTask.cancel(ErrProgramExiting)
|
||||
}
|
||||
|
||||
// GlobalContextWait waits for all tasks to finish, up to the given timeout.
|
||||
//
|
||||
// If the timeout is exceeded, it prints a list of all tasks that were
|
||||
// still running when the timeout was reached, and their current tree
|
||||
// of subtasks.
|
||||
func GlobalContextWait(timeout time.Duration) (err error) {
|
||||
done := make(chan struct{})
|
||||
after := time.After(timeout)
|
||||
go func() {
|
||||
globalTask.Wait()
|
||||
close(done)
|
||||
}()
|
||||
for {
|
||||
select {
|
||||
case <-done:
|
||||
return
|
||||
case <-after:
|
||||
logger.Warn().Msg("Timeout waiting for these tasks to finish:\n" + globalTask.tree())
|
||||
return context.DeadlineExceeded
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (t *Task) trace(msg string) {
|
||||
logger.Trace().Str("name", t.name).Msg(msg)
|
||||
}
|
||||
|
||||
// Name returns the name of the task.
|
||||
func (t *Task) Name() string {
|
||||
if !common.IsTrace {
|
||||
return t.name
|
||||
}
|
||||
parts := strings.Split(t.name, " > ")
|
||||
return parts[len(parts)-1]
|
||||
}
|
||||
|
||||
// String returns the name of the task.
|
||||
func (t *Task) String() string {
|
||||
return t.name
|
||||
}
|
||||
|
||||
// Context returns the context associated with the task. This context is
|
||||
// canceled when Finish of the task is called, or parent task is canceled.
|
||||
func (t *Task) Context() context.Context {
|
||||
return t.ctx
|
||||
}
|
||||
|
||||
func (t *Task) Finished() <-chan struct{} {
|
||||
return t.finished
|
||||
}
|
||||
|
||||
// FinishCause returns the reason / error that caused the task to be finished.
|
||||
func (t *Task) FinishCause() error {
|
||||
cause := context.Cause(t.ctx)
|
||||
if cause == nil {
|
||||
return t.ctx.Err()
|
||||
}
|
||||
return cause
|
||||
return context.Cause(t.ctx)
|
||||
}
|
||||
|
||||
// Parent returns the parent task of the current task.
|
||||
func (t *Task) Parent() *Task {
|
||||
return t.parent
|
||||
}
|
||||
|
||||
func (t *Task) runAllOnFinished(onCompTask *Task) {
|
||||
<-t.ctx.Done()
|
||||
t.WaitSubTasks()
|
||||
for _, OnFinishedFunc := range t.OnFinishedFuncs {
|
||||
OnFinishedFunc()
|
||||
t.onFinishedWg.Done()
|
||||
}
|
||||
onCompTask.Finish(fmt.Errorf("%w: %s, reason: %s", ErrTaskCanceled, t.name, "done"))
|
||||
}
|
||||
|
||||
// OnFinished calls fn when all subtasks are finished.
|
||||
// OnFinished calls fn when the task is canceled and all subtasks are finished.
|
||||
//
|
||||
// It cannot be called after Finish or Wait is called.
|
||||
// It should not be called after Finish is called.
|
||||
func (t *Task) OnFinished(about string, fn func()) {
|
||||
if t.parent == globalTask {
|
||||
t.OnCancel(about, fn)
|
||||
return
|
||||
}
|
||||
t.onFinishedWg.Add(1)
|
||||
t.OnFinishedMu.Lock()
|
||||
defer t.OnFinishedMu.Unlock()
|
||||
|
||||
if t.OnFinishedFuncs == nil {
|
||||
onCompTask := GlobalTask(t.name + " > OnFinished > " + about)
|
||||
go t.runAllOnFinished(onCompTask)
|
||||
}
|
||||
idx := len(t.OnFinishedFuncs)
|
||||
wrapped := func() {
|
||||
defer func() {
|
||||
if err := recover(); err != nil {
|
||||
logger.Error().
|
||||
Str("name", t.name).
|
||||
Interface("err", err).
|
||||
Msg("panic in " + about)
|
||||
}
|
||||
}()
|
||||
fn()
|
||||
logger.Trace().Str("name", t.name).Msgf("OnFinished[%d] done: %s", idx, about)
|
||||
}
|
||||
t.OnFinishedFuncs = append(t.OnFinishedFuncs, wrapped)
|
||||
t.onCancel(about, fn, true)
|
||||
}
|
||||
|
||||
// OnCancel calls fn when the task is canceled.
|
||||
//
|
||||
// It cannot be called after Finish or Wait is called.
|
||||
// It should not be called after Finish is called.
|
||||
func (t *Task) OnCancel(about string, fn func()) {
|
||||
onCompTask := GlobalTask(t.name + " > OnFinished")
|
||||
t.onCancel(about, fn, false)
|
||||
}
|
||||
|
||||
func (t *Task) onCancel(about string, fn func(), waitSubTasks bool) {
|
||||
t.onFinished.Add(1)
|
||||
go func() {
|
||||
<-t.ctx.Done()
|
||||
fn()
|
||||
onCompTask.Finish("done")
|
||||
t.trace("onCancel done: " + about)
|
||||
if waitSubTasks {
|
||||
t.children.Wait()
|
||||
}
|
||||
t.invokeWithRecover(fn, about)
|
||||
t.onFinished.Done()
|
||||
}()
|
||||
}
|
||||
|
||||
// Finish marks the task as finished and cancel its context.
|
||||
//
|
||||
// Then call Wait to wait for all subtasks, OnFinished and OnSubtasksFinished
|
||||
// of the task to finish.
|
||||
//
|
||||
// Note that it will also cancel all subtasks.
|
||||
// Finish cancel all subtasks and wait for them to finish,
|
||||
// then marks the task as finished, with the given reason (if any).
|
||||
func (t *Task) Finish(reason any) {
|
||||
var format string
|
||||
switch reason.(type) {
|
||||
case error:
|
||||
format = "%w"
|
||||
case string, fmt.Stringer:
|
||||
format = "%s"
|
||||
select {
|
||||
case <-t.finished:
|
||||
return
|
||||
default:
|
||||
format = "%v"
|
||||
t.once.Do(func() {
|
||||
t.finish(reason)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func (t *Task) finish(reason any) {
|
||||
t.cancel(fmtCause(reason))
|
||||
t.children.Wait()
|
||||
t.onFinished.Wait()
|
||||
if t.finished != nil {
|
||||
close(t.finished)
|
||||
}
|
||||
logger.Trace().Msg("task " + t.name + " finished")
|
||||
}
|
||||
|
||||
func fmtCause(cause any) error {
|
||||
switch cause := cause.(type) {
|
||||
case nil:
|
||||
return nil
|
||||
case error:
|
||||
return cause
|
||||
case string:
|
||||
return errors.New(cause)
|
||||
default:
|
||||
return fmt.Errorf("%v", cause)
|
||||
}
|
||||
t.finishOnce.Do(func() {
|
||||
t.cancel(fmt.Errorf("%w: %s, reason: "+format, ErrTaskCanceled, t.name, reason))
|
||||
})
|
||||
t.Wait()
|
||||
}
|
||||
|
||||
// Subtask returns a new subtask with the given name, derived from the parent's context.
|
||||
//
|
||||
// If the parent's context is already canceled, the returned subtask will be canceled immediately.
|
||||
//
|
||||
// This should not be called after Finish, Wait, or WaitSubTasks is called.
|
||||
func (t *Task) Subtask(name string) *Task {
|
||||
ctx, cancel := context.WithCancelCause(t.ctx)
|
||||
return t.newSubTask(ctx, cancel, name)
|
||||
}
|
||||
// This should not be called after Finish is called.
|
||||
func (t *Task) Subtask(name string, needFinish ...bool) *Task {
|
||||
nf := len(needFinish) == 0 || needFinish[0]
|
||||
|
||||
func (t *Task) newSubTask(ctx context.Context, cancel context.CancelCauseFunc, name string) *Task {
|
||||
parent := t
|
||||
if common.IsTrace {
|
||||
name = parent.name + " > " + name
|
||||
}
|
||||
subtask := &Task{
|
||||
ctx, cancel := context.WithCancelCause(t.ctx)
|
||||
child := &Task{
|
||||
finished: make(chan struct{}, 1),
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
name: name,
|
||||
parent: parent,
|
||||
subtasks: F.NewSet[*Task](),
|
||||
}
|
||||
parent.subTasksWg.Add(1)
|
||||
parent.subtasks.Add(subtask)
|
||||
if common.IsTrace {
|
||||
subtask.trace("started")
|
||||
if t != root {
|
||||
child.name = t.name + "." + name
|
||||
allTasks.Add(child)
|
||||
} else {
|
||||
child.name = name
|
||||
}
|
||||
|
||||
allTasksWg.Add(1)
|
||||
t.children.Add(1)
|
||||
|
||||
if !nf {
|
||||
go func() {
|
||||
subtask.Wait()
|
||||
subtask.trace("finished: " + subtask.FinishCause().Error())
|
||||
<-child.ctx.Done()
|
||||
child.Finish(nil)
|
||||
}()
|
||||
}
|
||||
|
||||
go func() {
|
||||
subtask.Wait()
|
||||
parent.subtasks.Remove(subtask)
|
||||
parent.subTasksWg.Done()
|
||||
<-child.finished
|
||||
allTasksWg.Done()
|
||||
t.children.Done()
|
||||
allTasks.Remove(child)
|
||||
}()
|
||||
return subtask
|
||||
|
||||
logger.Trace().Msg("task " + child.name + " started")
|
||||
return child
|
||||
}
|
||||
|
||||
// Wait waits for all subtasks, itself, OnFinished and OnSubtasksFinished to finish.
|
||||
//
|
||||
// It must be called only after Finish is called.
|
||||
func (t *Task) Wait() {
|
||||
<-t.ctx.Done()
|
||||
t.WaitSubTasks()
|
||||
t.onFinishedWg.Wait()
|
||||
// Name returns the name of the task without parent names.
|
||||
func (t *Task) Name() string {
|
||||
parts := strutils.SplitRune(t.name, '.')
|
||||
return parts[len(parts)-1]
|
||||
}
|
||||
|
||||
// WaitSubTasks waits for all subtasks of the task to finish.
|
||||
//
|
||||
// No more subtasks can be added after this call.
|
||||
//
|
||||
// It can be called before Finish is called.
|
||||
func (t *Task) WaitSubTasks() {
|
||||
t.subTasksWg.Wait()
|
||||
// String returns the full name of the task.
|
||||
func (t *Task) String() string {
|
||||
return t.name
|
||||
}
|
||||
|
||||
// tree returns a string representation of the task tree, with the given
|
||||
// prefix prepended to each line. The prefix is used to indent the tree,
|
||||
// and should be a string of spaces or a similar separator.
|
||||
//
|
||||
// The resulting string is suitable for printing to the console, and can be
|
||||
// used to debug the task tree.
|
||||
//
|
||||
// The tree is traversed in a depth-first manner, with each task's name and
|
||||
// line number (if available) printed on a separate line. The line number is
|
||||
// only printed if the task was created with a non-empty line argument.
|
||||
//
|
||||
// The returned string is not guaranteed to be stable, and may change between
|
||||
// runs of the program. It is intended for debugging purposes only.
|
||||
func (t *Task) tree(prefix ...string) string {
|
||||
var sb strings.Builder
|
||||
var pre string
|
||||
if len(prefix) > 0 {
|
||||
pre = prefix[0]
|
||||
sb.WriteString(pre + "- ")
|
||||
}
|
||||
sb.WriteString(t.Name() + "\n")
|
||||
t.subtasks.RangeAll(func(subtask *Task) {
|
||||
sb.WriteString(subtask.tree(pre + " "))
|
||||
})
|
||||
return sb.String()
|
||||
// MarshalText implements encoding.TextMarshaler.
|
||||
func (t *Task) MarshalText() ([]byte, error) {
|
||||
return []byte(t.name), nil
|
||||
}
|
||||
|
||||
// serialize returns a map[string]any representation of the task tree.
|
||||
//
|
||||
// The map contains the following keys:
|
||||
// - name: the name of the task
|
||||
// - subtasks: a slice of maps, each representing a subtask
|
||||
//
|
||||
// The subtask maps contain the same keys, recursively.
|
||||
//
|
||||
// The returned map is suitable for encoding to JSON, and can be used
|
||||
// to debug the task tree.
|
||||
//
|
||||
// The returned map is not guaranteed to be stable, and may change
|
||||
// between runs of the program. It is intended for debugging purposes
|
||||
// only.
|
||||
func (t *Task) serialize() map[string]any {
|
||||
m := make(map[string]any)
|
||||
parts := strings.Split(t.name, " > ")
|
||||
m["name"] = parts[len(parts)-1]
|
||||
if t.subtasks.Size() > 0 {
|
||||
m["subtasks"] = make([]map[string]any, 0, t.subtasks.Size())
|
||||
t.subtasks.RangeAll(func(subtask *Task) {
|
||||
m["subtasks"] = append(m["subtasks"].([]map[string]any), subtask.serialize())
|
||||
})
|
||||
}
|
||||
return m
|
||||
func (t *Task) invokeWithRecover(fn func(), caller string) {
|
||||
defer func() {
|
||||
if err := recover(); err != nil {
|
||||
logger.Error().
|
||||
Interface("err", err).
|
||||
Msg("panic in task " + t.name + "." + caller)
|
||||
if common.IsDebug {
|
||||
panic(string(debug.Stack()))
|
||||
}
|
||||
}
|
||||
}()
|
||||
fn()
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user