Print trace if close take too long

common/signal/timer.go

@@ -3,7 +3,11 @@ package signal
 import (
 	"context"
 	"sync"
+	"sync/atomic"
 	"time"
+
+	"github.com/xtls/xray-core/common"
+	"github.com/xtls/xray-core/common/task"
 )
 
 type ActivityUpdater interface {
@@ -11,35 +15,45 @@ type ActivityUpdater interface {
 }
 
 type ActivityTimer struct {
-	mu sync.Mutex
-	// timer will be nil if this timer is already finished
-	timer     *time.Timer
-	timeout   time.Duration
+	mu        sync.RWMutex
+	updated   chan struct{}
+	checkTask *task.Periodic
 	onTimeout func()
+	consumed  atomic.Bool
+	once      sync.Once
 }
 
 func (t *ActivityTimer) Update() {
-	// someone already called Update or closing, just return
-	if !t.mu.TryLock() {
-		return
+	select {
+	case t.updated <- struct{}{}:
+	default:
 	}
-	defer t.mu.Unlock()
-	if t.timer != nil {
-		t.timer.Reset(t.timeout)
+}
+
+func (t *ActivityTimer) check() error {
+	select {
+	case <-t.updated:
+	default:
+		t.finish()
 	}
+	return nil
 }
 
 func (t *ActivityTimer) finish() {
-	t.mu.Lock()
-	defer t.mu.Unlock()
-	if t.timer != nil {
-		t.timer.Stop()
+	t.once.Do(func() {
+		t.consumed.Store(true)
+		t.mu.Lock()
+		defer t.mu.Unlock()
+
+		common.CloseIfExists(t.checkTask)
 		t.onTimeout()
-		t.timer = nil
-	}
+	})
 }
 
 func (t *ActivityTimer) SetTimeout(timeout time.Duration) {
+	if t.consumed.Load() {
+		return
+	}
 	if timeout == 0 {
 		t.finish()
 		return
@@ -47,22 +61,25 @@ func (t *ActivityTimer) SetTimeout(timeout time.Duration) {
 
 	t.mu.Lock()
 	defer t.mu.Unlock()
-	if t.timer != nil {
-		t.timeout = timeout
-		t.timer.Reset(timeout)
+	// double check, just in case
+	if t.consumed.Load() {
+		return
 	}
+	newCheckTask := &task.Periodic{
+		Interval: timeout,
+		Execute:  t.check,
+	}
+	common.CloseIfExists(t.checkTask)
+	t.checkTask = newCheckTask
+	t.Update()
+	common.Must(newCheckTask.Start())
 }
 
 func CancelAfterInactivity(ctx context.Context, cancel context.CancelFunc, timeout time.Duration) *ActivityTimer {
-	activityTimer := &ActivityTimer{
-		timeout:   timeout,
+	timer := &ActivityTimer{
+		updated:   make(chan struct{}, 1),
 		onTimeout: cancel,
 	}
-	// strange situation
-	if timeout == 0 {
-		cancel()
-		return activityTimer
-	}
-	activityTimer.timer = time.AfterFunc(timeout, activityTimer.finish)
-	return activityTimer
+	timer.SetTimeout(timeout)
+	return timer
 }

main/run.go

@@ -91,7 +91,28 @@ func executeRun(cmd *base.Command, args []string) {
 		fmt.Println("Failed to start:", err)
 		os.Exit(-1)
 	}
-	defer server.Close()
+	defer func() {
+		closeErrCh := make(chan error, 1)
+		go func() {
+			closeErrCh <- server.Close()
+		}()
+		select {
+		case err := <-closeErrCh:
+			if err != nil {
+				fmt.Println("Failed to close server:", err)
+			}
+		case <-time.After(10 * time.Second):
+			fmt.Println("Timeout when closing, printing traces:")
+			buf := make([]byte, 1<<20)
+			n := runtime.Stack(buf, true)
+			blocks := strings.Split(string(buf[:n]), "\n\n")
+			for _, block := range blocks {
+				if strings.Contains(block, "github.com/xtls/xray-core/core.(*Instance).Close") {
+					fmt.Println(block)
+				}
+			}
+		}
+	}()
 
 	// Explicitly triggering GC to remove garbage from config loading.
 	runtime.GC()