refactor: improve task management with xsync for concurrent access and enhance callback and subtasks handling as well as memory allocation

internal/task/utils.go

@@ -2,7 +2,6 @@ package task
 
 import (
 	"context"
-	"encoding/json"
 	"errors"
 	"os"
 	"os/signal"
@@ -10,73 +9,34 @@ import (
 	"time"
 
 	"github.com/rs/zerolog/log"
-	F "github.com/yusing/go-proxy/internal/utils/functional"
 )
 
 var ErrProgramExiting = errors.New("program exiting")
 
-var (
-	root     = newRoot()
-	allTasks = F.NewSet[*Task]()
-)
-
-func testCleanup() {
-	root = newRoot()
-	allTasks.Clear()
-}
-
 // RootTask returns a new Task with the given name, derived from the root context.
-func RootTask(name string, needFinish ...bool) *Task {
-	return root.Subtask(name, needFinish...)
-}
-
-func newRoot() *Task {
-	t := &Task{
-		name:         "root",
-		childrenDone: make(chan struct{}),
-		finished:     make(chan struct{}),
-	}
-	t.ctx, t.cancel = context.WithCancelCause(context.Background())
-	return t
+//
+//go:inline
+func RootTask(name string, needFinish bool) *Task {
+	return root.Subtask(name, needFinish)
 }
 
 func RootContext() context.Context {
-	return root.ctx
+	return root
 }
 
 func RootContextCanceled() <-chan struct{} {
-	return root.ctx.Done()
+	return root.Done()
 }
 
 func OnProgramExit(about string, fn func()) {
 	root.OnFinished(about, fn)
 }
 
-// GracefulShutdown waits for all tasks to finish, up to the given timeout.
+// WaitExit waits for a signal to shutdown the program, and then 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 GracefulShutdown(timeout time.Duration) (err error) {
-	go root.Finish(ErrProgramExiting)
-
-	after := time.After(timeout)
-	for {
-		select {
-		case <-root.finished:
-			return
-		case <-after:
-			b, err := json.Marshal(DebugTaskList())
-			if err != nil {
-				log.Warn().Err(err).Msg("failed to marshal tasks")
-				return context.DeadlineExceeded
-			}
-			log.Warn().RawJSON("tasks", b).Msgf("Timeout waiting for these %d tasks to finish", allTasks.Size())
-			return context.DeadlineExceeded
-		}
-	}
-}
-
 func WaitExit(shutdownTimeout int) {
 	sig := make(chan os.Signal, 1)
 	signal.Notify(sig, syscall.SIGINT)
@@ -88,5 +48,22 @@ func WaitExit(shutdownTimeout int) {
 
 	// gracefully shutdown
 	log.Info().Msg("shutting down")
-	_ = GracefulShutdown(time.Second * time.Duration(shutdownTimeout))
+	if err := gracefulShutdown(time.Second * time.Duration(shutdownTimeout)); err != nil {
+		root.reportStucked()
+	}
+}
+
+// gracefulShutdown 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 gracefulShutdown(timeout time.Duration) error {
+	root.Finish(ErrProgramExiting)
+	root.finishChildren()
+	root.runCallbacks()
+	if !root.waitFinish(timeout) {
+		return context.DeadlineExceeded
+	}
+	return nil
 }