routes in loadbalance pool no longer listed in ls-route and its API, the loadbalancer is listed instead. improved context handling and grateful shutdown

internal/common/task.go

@@ -0,0 +1,158 @@
+package common
+
+import (
+	"context"
+	"fmt"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/puzpuzpuz/xsync/v3"
+	"github.com/sirupsen/logrus"
+)
+
+var (
+	globalCtx, globalCtxCancel = context.WithCancel(context.Background())
+	globalCtxWg                sync.WaitGroup
+	globalCtxTraceMap          = xsync.NewMapOf[*task, struct{}]()
+)
+
+type (
+	Task interface {
+		Name() string
+		Context() context.Context
+		Subtask(usageFmt string, args ...interface{}) Task
+		SubtaskWithCancel(usageFmt string, args ...interface{}) (Task, context.CancelFunc)
+		Finished()
+	}
+	task struct {
+		ctx      context.Context
+		subtasks []*task
+		name     string
+		finished bool
+		mu       sync.Mutex
+	}
+)
+
+func (t *task) Name() string {
+	return t.name
+}
+
+func (t *task) Context() context.Context {
+	return t.ctx
+}
+
+func (t *task) Finished() {
+	t.mu.Lock()
+	defer t.mu.Unlock()
+
+	if t.finished {
+		return
+	}
+	t.finished = true
+	if _, ok := globalCtxTraceMap.Load(t); ok {
+		globalCtxWg.Done()
+		globalCtxTraceMap.Delete(t)
+	}
+}
+
+func (t *task) Subtask(format string, args ...interface{}) Task {
+	if len(args) > 0 {
+		format = fmt.Sprintf(format, args...)
+	}
+	t.mu.Lock()
+	defer t.mu.Unlock()
+	sub := newSubTask(t.ctx, format)
+	t.subtasks = append(t.subtasks, sub)
+	return sub
+}
+
+func (t *task) SubtaskWithCancel(format string, args ...interface{}) (Task, context.CancelFunc) {
+	if len(args) > 0 {
+		format = fmt.Sprintf(format, args...)
+	}
+	t.mu.Lock()
+	defer t.mu.Unlock()
+	ctx, cancel := context.WithCancel(t.ctx)
+	sub := newSubTask(ctx, format)
+	t.subtasks = append(t.subtasks, sub)
+	return sub, cancel
+}
+
+func (t *task) Tree(prefix ...string) string {
+	var sb strings.Builder
+	var pre string
+	if len(prefix) > 0 {
+		pre = prefix[0]
+	}
+	sb.WriteString(pre)
+	sb.WriteString(t.Name() + "\n")
+	for _, sub := range t.subtasks {
+		if sub.finished {
+			continue
+		}
+		sb.WriteString(sub.Tree(pre + "  "))
+	}
+	return sb.String()
+}
+
+func newSubTask(ctx context.Context, name string) *task {
+	t := &task{
+		ctx:  ctx,
+		name: name,
+	}
+	globalCtxTraceMap.Store(t, struct{}{})
+	globalCtxWg.Add(1)
+	return t
+}
+
+func NewTask(format string, args ...interface{}) Task {
+	if len(args) > 0 {
+		format = fmt.Sprintf(format, args...)
+	}
+	return newSubTask(globalCtx, format)
+}
+
+func NewTaskWithCancel(format string, args ...interface{}) (Task, context.CancelFunc) {
+	subCtx, cancel := context.WithCancel(globalCtx)
+	if len(args) > 0 {
+		format = fmt.Sprintf(format, args...)
+	}
+	return newSubTask(subCtx, format), cancel
+}
+
+func GlobalTask(format string, args ...interface{}) Task {
+	if len(args) > 0 {
+		format = fmt.Sprintf(format, args...)
+	}
+	return &task{
+		ctx:  globalCtx,
+		name: format,
+	}
+}
+
+func CancelGlobalContext() {
+	globalCtxCancel()
+}
+
+func GlobalContextWait(timeout time.Duration) {
+	done := make(chan struct{})
+	after := time.After(timeout)
+	go func() {
+		globalCtxWg.Wait()
+		close(done)
+	}()
+	for {
+		select {
+		case <-done:
+			return
+		case <-after:
+			logrus.Println("Timeout waiting for these tasks to finish:")
+			globalCtxTraceMap.Range(func(t *task, _ struct{}) bool {
+				logrus.Println(t.Tree())
+				return true
+			})
+			return
+		}
+	}
+}