graceful shutdown and ref count related

internal/common/task.go

@@ -13,8 +13,8 @@ import (
 
 var (
 	globalCtx, globalCtxCancel = context.WithCancel(context.Background())
-	globalCtxWg                sync.WaitGroup
-	globalCtxTraceMap          = xsync.NewMapOf[*task, struct{}]()
+	taskWg                     sync.WaitGroup
+	tasksMap                   = xsync.NewMapOf[*task, struct{}]()
 )
 
 type (
@@ -38,10 +38,14 @@ func (t *task) Name() string {
 	return t.name
 }
 
+// Context returns the context associated with the task. This context is
+// canceled when the task is finished.
 func (t *task) Context() context.Context {
 	return t.ctx
 }
 
+// Finished marks the task as finished and notifies the global wait group.
+// Finished is thread-safe and idempotent.
 func (t *task) Finished() {
 	t.mu.Lock()
 	defer t.mu.Unlock()
@@ -50,12 +54,22 @@ func (t *task) Finished() {
 		return
 	}
 	t.finished = true
-	if _, ok := globalCtxTraceMap.Load(t); ok {
-		globalCtxWg.Done()
-		globalCtxTraceMap.Delete(t)
+	if _, ok := tasksMap.Load(t); ok {
+		taskWg.Done()
+		tasksMap.Delete(t)
 	}
+	logrus.Debugf("task %q finished", t.Name())
 }
 
+// Subtask returns a new subtask with the given name, derived from the receiver's context.
+//
+// The returned subtask is associated with the receiver's context and will be
+// automatically registered and deregistered from the global task wait group.
+//
+// If the receiver's context is already canceled, the returned subtask will be
+// canceled immediately.
+//
+// The returned subtask is safe for concurrent use.
 func (t *task) Subtask(format string, args ...interface{}) Task {
 	if len(args) > 0 {
 		format = fmt.Sprintf(format, args...)
@@ -67,6 +81,14 @@ func (t *task) Subtask(format string, args ...interface{}) Task {
 	return sub
 }
 
+// SubtaskWithCancel returns a new subtask with the given name, derived from the receiver's context,
+// and a cancel function. The returned subtask is associated with the receiver's context and will be
+// automatically registered and deregistered from the global task wait group.
+//
+// If the receiver's context is already canceled, the returned subtask will be canceled immediately.
+//
+// The returned cancel function is safe for concurrent use, and can be used to cancel the returned
+// subtask at any time.
 func (t *task) SubtaskWithCancel(format string, args ...interface{}) (Task, context.CancelFunc) {
 	if len(args) > 0 {
 		format = fmt.Sprintf(format, args...)
@@ -79,7 +101,7 @@ func (t *task) SubtaskWithCancel(format string, args ...interface{}) (Task, cont
 	return sub, cancel
 }
 
-func (t *task) Tree(prefix ...string) string {
+func (t *task) tree(prefix ...string) string {
 	var sb strings.Builder
 	var pre string
 	if len(prefix) > 0 {
@@ -91,7 +113,7 @@ func (t *task) Tree(prefix ...string) string {
 		if sub.finished {
 			continue
 		}
-		sb.WriteString(sub.Tree(pre + "  "))
+		sb.WriteString(sub.tree(pre + "  "))
 	}
 	return sb.String()
 }
@@ -101,11 +123,22 @@ func newSubTask(ctx context.Context, name string) *task {
 		ctx:  ctx,
 		name: name,
 	}
-	globalCtxTraceMap.Store(t, struct{}{})
-	globalCtxWg.Add(1)
+	tasksMap.Store(t, struct{}{})
+	taskWg.Add(1)
 	return t
 }
 
+// NewTask returns a new Task with the given name, derived from the global
+// context.
+//
+// The returned Task is associated with the global context and will be
+// automatically registered and deregistered from the global context's wait
+// group.
+//
+// If the global context is already canceled, the returned Task will be
+// canceled immediately.
+//
+// The returned Task is not safe for concurrent use.
 func NewTask(format string, args ...interface{}) Task {
 	if len(args) > 0 {
 		format = fmt.Sprintf(format, args...)
@@ -113,6 +146,18 @@ func NewTask(format string, args ...interface{}) Task {
 	return newSubTask(globalCtx, format)
 }
 
+// NewTaskWithCancel returns a new Task with the given name, derived from the
+// global context, and a cancel function. The returned Task is associated with
+// the global context and will be automatically registered and deregistered
+// from the global task wait group.
+//
+// If the global context is already canceled, the returned Task will be
+// canceled immediately.
+//
+// The returned Task is safe for concurrent use.
+//
+// The returned cancel function is safe for concurrent use, and can be used
+// to cancel the returned Task at any time.
 func NewTaskWithCancel(format string, args ...interface{}) (Task, context.CancelFunc) {
 	subCtx, cancel := context.WithCancel(globalCtx)
 	if len(args) > 0 {
@@ -121,6 +166,17 @@ func NewTaskWithCancel(format string, args ...interface{}) (Task, context.Cancel
 	return newSubTask(subCtx, format), cancel
 }
 
+// GlobalTask returns a new Task with the given name, associated with the
+// global context.
+//
+// Unlike NewTask, GlobalTask does not automatically register or deregister
+// the Task with the global task wait group. The returned Task is not
+// started, but the name is formatted immediately.
+//
+// This is best used for main task that do not need to wait and
+// will create a bunch of subtasks.
+//
+// The returned Task is safe for concurrent use.
 func GlobalTask(format string, args ...interface{}) Task {
 	if len(args) > 0 {
 		format = fmt.Sprintf(format, args...)
@@ -131,15 +187,24 @@ func GlobalTask(format string, args ...interface{}) Task {
 	}
 }
 
+// CancelGlobalContext cancels the global context, which will cause all tasks
+// created by GlobalTask or NewTask to be canceled. This should be called
+// before exiting the program to ensure that all tasks are properly cleaned
+// up.
 func CancelGlobalContext() {
 	globalCtxCancel()
 }
 
+// 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) {
 	done := make(chan struct{})
 	after := time.After(timeout)
 	go func() {
-		globalCtxWg.Wait()
+		taskWg.Wait()
 		close(done)
 	}()
 	for {
@@ -148,8 +213,8 @@ func GlobalContextWait(timeout time.Duration) {
 			return
 		case <-after:
 			logrus.Println("Timeout waiting for these tasks to finish:")
-			globalCtxTraceMap.Range(func(t *task, _ struct{}) bool {
-				logrus.Println(t.Tree())
+			tasksMap.Range(func(t *task, _ struct{}) bool {
+				logrus.Println(t.tree())
 				return true
 			})
 			return