From 659ad2987519a6e18b74b9641ae306c3ef168311 Mon Sep 17 00:00:00 2001
From: yusing <yusing@6uo.me>
Date: Wed, 1 Jan 2025 16:51:45 +0800
Subject: [PATCH] add timeout on task wait, temporary fix task stuck

---
 internal/task/task.go | 47 ++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 44 insertions(+), 3 deletions(-)

diff --git a/internal/task/task.go b/internal/task/task.go
index b431b18..a8b6cde 100644
--- a/internal/task/task.go
+++ b/internal/task/task.go
@@ -5,7 +5,9 @@ import (
 	"errors"
 	"fmt"
 	"runtime/debug"
+	"strings"
 	"sync"
+	"time"
 
 	"github.com/yusing/go-proxy/internal/common"
 	E "github.com/yusing/go-proxy/internal/error"
@@ -51,6 +53,8 @@ type (
 	}
 )
 
+const taskTimeout = 3 * time.Second
+
 func (t *Task) Context() context.Context {
 	return t.ctx
 }
@@ -83,7 +87,7 @@ func (t *Task) onCancel(about string, fn func(), waitSubTasks bool) {
 	go func() {
 		<-t.ctx.Done()
 		if waitSubTasks {
-			t.children.Wait()
+			waitWithTimeout(&t.children)
 		}
 		t.invokeWithRecover(fn, about)
 		t.onFinished.Done()
@@ -100,14 +104,51 @@ func (t *Task) Finish(reason any) {
 
 func (t *Task) finish(reason any) {
 	t.cancel(fmtCause(reason))
-	t.children.Wait()
-	t.onFinished.Wait()
+	if !waitWithTimeout(&t.children) {
+		logger.Debug().
+			Strs("subtasks", t.listChildren()).
+			Msg("Timeout waiting for these subtasks to finish")
+	}
+	if !waitWithTimeout(&t.onFinished) {
+		logger.Debug().
+			Str("task", t.name).
+			Msg("Timeout waiting for callbacks to finish")
+	}
 	if t.finished != nil {
 		close(t.finished)
 	}
 	logger.Trace().Msg("task " + t.name + " finished")
 }
 
+// debug only.
+func (t *Task) listChildren() []string {
+	var children []string
+	allTasks.Range(func(child *Task) bool {
+		if strings.HasPrefix(child.name, t.name+".") {
+			children = append(children, child.name)
+		}
+		return true
+	})
+	return children
+}
+
+func waitWithTimeout(wg *sync.WaitGroup) bool {
+	done := make(chan struct{})
+	timeout := time.After(taskTimeout)
+
+	go func() {
+		wg.Wait()
+		close(done)
+	}()
+
+	select {
+	case <-done:
+		return true
+	case <-timeout:
+		return false
+	}
+}
+
 func fmtCause(cause any) error {
 	switch cause := cause.(type) {
 	case nil: