-
Notifications
You must be signed in to change notification settings - Fork 151
/
tasks.go
163 lines (133 loc) · 3.75 KB
/
tasks.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
//go:build !wasm
package tasks
import (
"context"
"fmt"
"os"
"os/signal"
"reflect"
)
type TaskFunc[T any, U any] func(t *Task, args T) (nextArgs U, err error)
type CleanupFunc func(ctx context.Context) error
type Task struct {
Name string
Ctx context.Context
Logs *os.File
taskFunction TaskFunc[any, any]
argType reflect.Type
returnType reflect.Type
cleanFunctions []CleanupFunc
}
type Tasks struct {
tasks []*Task
LoggerMode LoggerMode
}
func Begin() *Tasks {
return &Tasks{
LoggerMode: PrinterModeAuto,
}
}
func (ts *Tasks) SetLoggerMode(mode LoggerMode) {
ts.LoggerMode = mode
}
func Add[TaskArg any, TaskReturn any](ts *Tasks, name string, taskFunc TaskFunc[TaskArg, TaskReturn]) {
var argValue TaskArg
var returnValue TaskReturn
argType := reflect.TypeOf(argValue)
returnType := reflect.TypeOf(returnValue)
tasksAmount := len(ts.tasks)
if tasksAmount > 0 {
lastTask := ts.tasks[tasksAmount-1]
if argType != lastTask.returnType {
panic(fmt.Errorf("invalid task declared, wait for %s, previous task returns %s", argType.Name(), lastTask.returnType.Name()))
}
}
ts.tasks = append(ts.tasks, &Task{
Name: name,
argType: argType,
returnType: returnType,
taskFunction: func(t *Task, i interface{}) (passedData interface{}, err error) {
if i == nil {
var zero TaskArg
passedData, err = taskFunc(t, zero)
} else {
passedData, err = taskFunc(t, i.(TaskArg))
}
return
},
})
}
func (t *Task) AddToCleanUp(cleanupFunc CleanupFunc) {
t.cleanFunctions = append(t.cleanFunctions, cleanupFunc)
}
// setupContext return a contextWithCancel that will cancel on os interrupt (Ctrl-C)
func setupContext(ctx context.Context) (context.Context, func()) {
return signal.NotifyContext(ctx, os.Interrupt)
}
// Cleanup execute all tasks cleanup function before failed one in reverse order
func (ts *Tasks) Cleanup(ctx context.Context, logger *Logger, failed int) {
cancelableCtx, cleanCtx := setupContext(ctx)
defer cleanCtx()
for i := failed; i >= 0; i-- {
task := ts.tasks[i]
select {
case <-cancelableCtx.Done():
fmt.Println("cleanup has been cancelled, there may be dangling resources")
return
default:
}
if len(task.cleanFunctions) != 0 {
var err error
for i, cleanUpFunc := range task.cleanFunctions {
loggerEntry := logger.AddEntry(fmt.Sprintf("Cleaning task %q %d/%d", task.Name, i+1, len(task.cleanFunctions)))
task.Logs = loggerEntry.Logs
loggerEntry.Start()
err = cleanUpFunc(cancelableCtx)
if err != nil {
loggerEntry.Complete(err)
break
}
loggerEntry.Complete(nil)
}
}
}
}
// Execute tasks with interactive display and cleanup on fail
func (ts *Tasks) Execute(ctx context.Context, data interface{}) (interface{}, error) {
cancelableCtx, cleanCtx := setupContext(ctx)
defer cleanCtx()
logger, err := NewTasksLogger(context.Background(), ts.LoggerMode)
if err != nil {
return nil, err
}
defer func() {
err := logger.CloseAndWait()
if err != nil {
fmt.Println("failed to close logger:", err)
}
}()
for i := range ts.tasks {
task := ts.tasks[i]
loggerEntry := logger.AddEntry(task.Name)
task.Logs = loggerEntry.Logs
loggerEntry.Start()
// Add context and reset cleanup functions, allows to execute multiple times
task.Ctx = cancelableCtx
task.cleanFunctions = []CleanupFunc(nil)
data, err = task.taskFunction(task, data)
if err != nil {
loggerEntry.Complete(err)
ts.Cleanup(ctx, logger, i)
return nil, fmt.Errorf("task %d %q failed: %w", i+1, task.Name, err)
}
select {
case <-ctx.Done():
loggerEntry.Complete(context.Canceled)
ts.Cleanup(ctx, logger, i)
return nil, fmt.Errorf("task %d %q failed: context canceled", i+1, task.Name)
default:
}
loggerEntry.Complete(nil)
}
return data, nil
}