From 4bcdd7f9646f5845d3186e97e7349f22d51b1245 Mon Sep 17 00:00:00 2001
From: "maye.chy" <chenhongyu@come-future.com>
Date: Tue, 26 Nov 2024 18:07:55 +0800
Subject: [PATCH] change Timer to TimerTask interface

---
 bucket.go                     |  4 ++++
 timingwheel.go                | 12 ++++++++----
 timingwheel_benchmark_test.go |  2 +-
 3 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/bucket.go b/bucket.go
index 5c297bf..093b03a 100644
--- a/bucket.go
+++ b/bucket.go
@@ -7,6 +7,10 @@ import (
 	"unsafe"
 )
 
+type TimerTask interface {
+	Stop() bool
+}
+
 // Timer represents a single event. When the Timer expires, the given
 // task will be executed.
 type Timer struct {
diff --git a/timingwheel.go b/timingwheel.go
index e2c64e0..d08de3a 100644
--- a/timingwheel.go
+++ b/timingwheel.go
@@ -166,7 +166,7 @@ func (tw *TimingWheel) Stop() {
 
 // AfterFunc waits for the duration to elapse and then calls f in its own goroutine.
 // It returns a Timer that can be used to cancel the call using its Stop method.
-func (tw *TimingWheel) AfterFunc(d time.Duration, f func()) *Timer {
+func (tw *TimingWheel) AfterFunc(d time.Duration, f func()) TimerTask {
 	t := &Timer{
 		expiration: timeToMs(time.Now().UTC().Add(d)),
 		task:       f,
@@ -199,11 +199,15 @@ type Scheduler interface {
 // Afterwards, it will ask the next execution time each time f is about to
 // be executed, and f will be called at the next execution time if the time
 // is non-zero.
-func (tw *TimingWheel) ScheduleFunc(s Scheduler, f func()) (t *Timer) {
+func (tw *TimingWheel) ScheduleFunc(s Scheduler, f func()) TimerTask {
+	var (
+		t = new(Timer)
+	)
+
 	expiration := s.Next(time.Now().UTC())
 	if expiration.IsZero() {
 		// No time is scheduled, return nil.
-		return
+		return t
 	}
 
 	t = &Timer{
@@ -222,5 +226,5 @@ func (tw *TimingWheel) ScheduleFunc(s Scheduler, f func()) (t *Timer) {
 	}
 	tw.addOrRun(t)
 
-	return
+	return t
 }
diff --git a/timingwheel_benchmark_test.go b/timingwheel_benchmark_test.go
index f9bd7f1..deab8e6 100644
--- a/timingwheel_benchmark_test.go
+++ b/timingwheel_benchmark_test.go
@@ -26,7 +26,7 @@ func BenchmarkTimingWheel_StartStop(b *testing.B) {
 	}
 	for _, c := range cases {
 		b.Run(c.name, func(b *testing.B) {
-			base := make([]*timingwheel.Timer, c.N)
+			base := make([]timingwheel.TimerTask, c.N)
 			for i := 0; i < len(base); i++ {
 				base[i] = tw.AfterFunc(genD(i), func() {})
 			}