Refactor Timer

common/signal/timer.go

@@ -5,9 +5,6 @@ import (
 	"sync"
 	"sync/atomic"
 	"time"
-
-	"github.com/xtls/xray-core/common"
-	"github.com/xtls/xray-core/common/task"
 )
 
 type ActivityUpdater interface {
@@ -15,45 +12,28 @@ type ActivityUpdater interface {
 }
 
 type ActivityTimer struct {
-	mu        sync.RWMutex
-	updated   chan struct{}
-	checkTask *task.Periodic
+	mu        sync.Mutex
+	timer     atomic.Pointer[time.Timer]
+	timeout   time.Duration
 	onTimeout func()
-	consumed  atomic.Bool
-	once      sync.Once
 }
 
 func (t *ActivityTimer) Update() {
-	select {
-	case t.updated <- struct{}{}:
-	default:
+	t.mu.Lock()
+	defer t.mu.Unlock()
+	if timer := t.timer.Load(); timer != nil {
+		timer.Reset(t.timeout)
 	}
 }
 
-func (t *ActivityTimer) check() error {
-	select {
-	case <-t.updated:
-	default:
-		t.finish()
-	}
-	return nil
-}
-
 func (t *ActivityTimer) finish() {
-	t.once.Do(func() {
-		t.consumed.Store(true)
-		t.mu.Lock()
-		defer t.mu.Unlock()
-
-		common.CloseIfExists(t.checkTask)
+	if timer := t.timer.Swap(nil); timer != nil {
+		timer.Stop()
 		t.onTimeout()
-	})
+	}
 }
 
 func (t *ActivityTimer) SetTimeout(timeout time.Duration) {
-	if t.consumed.Load() {
-		return
-	}
 	if timeout == 0 {
 		t.finish()
 		return
@@ -61,25 +41,21 @@ func (t *ActivityTimer) SetTimeout(timeout time.Duration) {
 
 	t.mu.Lock()
 	defer t.mu.Unlock()
-	// double check, just in case
-	if t.consumed.Load() {
-		return
+	if timer := t.timer.Load(); timer != nil {
+		t.timeout = timeout
+		timer.Reset(timeout)
 	}
-	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 {
-	timer := &ActivityTimer{
-		updated:   make(chan struct{}, 1),
+	activityTimer := &ActivityTimer{
+		timeout:   timeout,
 		onTimeout: cancel,
 	}
-	timer.SetTimeout(timeout)
-	return timer
+	if timeout == 0 {
+		cancel()
+		return activityTimer
+	}
+	activityTimer.timer.Store(time.AfterFunc(timeout, activityTimer.finish))
+	return activityTimer
 }