fix(task): refactor task module and fix reload stuck/error, fix some logic

internal/task/utils.go

@@ -3,6 +3,7 @@ package task
 import (
 	"context"
 	"errors"
+	"fmt"
 	"os"
 	"os/signal"
 	"syscall"
@@ -13,6 +14,31 @@ import (
 
 var ErrProgramExiting = errors.New("program exiting")
 
+var root *Task
+
+var closedCh = make(chan struct{})
+
+func init() {
+	close(closedCh)
+	initRoot()
+}
+
+func initRoot() {
+	ctx, cancel := context.WithCancelCause(context.Background())
+	root = &Task{
+		name:   "root",
+		ctx:    ctx,
+		cancel: cancel,
+		done:   closedCh,
+	}
+	root.parent = root
+}
+
+func testCleanup() {
+	root.cancel(nil)
+	initRoot()
+}
+
 // RootTask returns a new Task with the given name, derived from the root context.
 //
 //go:inline
@@ -29,7 +55,7 @@ func RootContextCanceled() <-chan struct{} {
 }
 
 func OnProgramExit(about string, fn func()) {
-	root.OnFinished(about, fn)
+	root.OnCancel(about, fn)
 }
 
 // WaitExit waits for a signal to shutdown the program, and then waits for all tasks to finish, up to the given timeout.
@@ -59,19 +85,33 @@ func WaitExit(shutdownTimeout int) {
 // still running when the timeout was reached, and their current tree
 // of subtasks.
 func gracefulShutdown(timeout time.Duration) error {
-	root.mu.Lock()
-	if root.isCanceled() {
-		cause := context.Cause(root.ctx)
-		root.mu.Unlock()
-		return cause
-	}
-	root.mu.Unlock()
-
-	root.cancel(ErrProgramExiting)
-	ok := waitEmpty(root.children, timeout)
-	root.runOnFinishCallbacks()
-	if !ok || !root.waitFinish(timeout) {
+	go root.Finish(ErrProgramExiting)
+	if !root.waitFinish(timeout) {
 		return context.DeadlineExceeded
 	}
 	return nil
 }
+
+func invokeWithRecover(cb *Callback) {
+	defer func() {
+		if err := recover(); err != nil {
+			log.Err(fmtCause(err)).Str("callback", cb.about).Msg("panic")
+			panicWithDebugStack()
+		}
+	}()
+	cb.fn()
+}
+
+//go:inline
+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)
+	}
+}