Improve task

common/task/task.go

@@ -7,48 +7,99 @@ import (
 	E "github.com/sagernet/sing/common/exceptions"
 )
 
-func Run(ctx context.Context, tasks ...func() error) error {
-	runtimeCtx, cancel := context.WithCancel(ctx)
-	wg := &sync.WaitGroup{}
-	wg.Add(len(tasks))
-	var retErr []error
-	for _, task := range tasks {
-		currentTask := task
-		go func() {
-			if err := currentTask(); err != nil {
-				retErr = append(retErr, err)
-			}
-			wg.Done()
-		}()
-	}
-	go func() {
-		wg.Wait()
-		cancel()
-	}()
-	select {
-	case <-ctx.Done():
-	case <-runtimeCtx.Done():
-	}
-	retErr = append(retErr, ctx.Err())
-	return E.Errors(retErr...)
+type taskItem struct {
+	Name string
+	Run  func(ctx context.Context) error
 }
 
-func Any(ctx context.Context, tasks ...func(ctx context.Context) error) error {
-	runtimeCtx, cancel := context.WithCancel(ctx)
-	defer cancel()
+type Group struct {
+	tasks    []taskItem
+	cleanup  func()
+	fastFail bool
+}
+
+func (g *Group) Append(name string, f func(ctx context.Context) error) {
+	g.tasks = append(g.tasks, taskItem{
+		Name: name,
+		Run:  f,
+	})
+}
+
+func (g *Group) Append0(f func(ctx context.Context) error) {
+	g.tasks = append(g.tasks, taskItem{
+		Run: f,
+	})
+}
+
+func (g *Group) Cleanup(f func()) {
+	g.cleanup = f
+}
+
+func (g *Group) FastFail() {
+	g.fastFail = true
+}
+
+func (g *Group) Run(ctx context.Context) error {
+	var retAccess sync.Mutex
 	var retErr error
-	for _, task := range tasks {
+
+	taskCount := int8(len(g.tasks))
+	taskCtx, taskFinish := context.WithCancel(context.Background())
+	var mixedCtx context.Context
+	var mixedFinish context.CancelFunc
+	if ctx.Done() != nil || g.fastFail {
+		mixedCtx, mixedFinish = context.WithCancel(ctx)
+	} else {
+		mixedCtx, mixedFinish = taskCtx, taskFinish
+	}
+
+	for _, task := range g.tasks {
 		currentTask := task
 		go func() {
-			if err := currentTask(runtimeCtx); err != nil {
-				retErr = err
+			err := currentTask.Run(mixedCtx)
+			retAccess.Lock()
+			if err != nil {
+				retErr = E.Append(retErr, err, func(err error) error {
+					if currentTask.Name == "" {
+						return err
+					}
+					return E.Cause(err, currentTask.Name)
+				})
+				if g.fastFail {
+					mixedFinish()
+				}
+			}
+			taskCount--
+			currentCount := taskCount
+			retAccess.Unlock()
+			if currentCount == 0 {
+				taskFinish()
 			}
-			cancel()
 		}()
 	}
+
+	var upstreamErr error
+
 	select {
 	case <-ctx.Done():
-	case <-runtimeCtx.Done():
+		upstreamErr = ctx.Err()
+	case <-taskCtx.Done():
+		mixedFinish()
+	case <-mixedCtx.Done():
 	}
-	return E.Errors(retErr, ctx.Err())
+
+	if g.cleanup != nil {
+		g.cleanup()
+	}
+
+	<-taskCtx.Done()
+
+	taskFinish()
+	mixedFinish()
+
+	retErr = E.Append(retErr, upstreamErr, func(err error) error {
+		return E.Cause(err, "upstream")
+	})
+
+	return retErr
 }