From 546826d74e1575cc5b86a57f6fbe1a0e908ec9d5 Mon Sep 17 00:00:00 2001
From: Lee Min Jea <minjea.lee@naverlabs.com>
Date: Wed, 14 Feb 2024 15:39:33 +0900
Subject: [PATCH] separate goroutine invoking given functions internally from
 user's goroutine which calls throttled function and purge function.

---
 retry.go      | 14 ++++++++------
 retry_test.go | 42 ++++++++++++++++++++----------------------
 2 files changed, 28 insertions(+), 28 deletions(-)

diff --git a/retry.go b/retry.go
index a3b4bc3c..002ecd47 100644
--- a/retry.go
+++ b/retry.go
@@ -306,21 +306,22 @@ func (th *throttle) throttledFunc() {
 	}
 }
 
+func (th *throttle) invokeFunctions() {
+	for _, f := range th.callbacks {
+		go f()
+	}
+}
+
 func (th *throttle) purge(forcePurge bool) {
 	th.mu.Lock()
 	defer th.mu.Unlock()
-	defer func() {
-		th.needInvoke = false
-	}()
 
 	if th.timer != nil {
 		th.timer.Stop()
 	}
 
 	if th.needInvoke || forcePurge {
-		for _, f := range th.callbacks {
-			f()
-		}
+		th.invokeFunctions()
 		th.timer = time.AfterFunc(th.interval, func() {
 			th.purge(false)
 		})
@@ -328,6 +329,7 @@ func (th *throttle) purge(forcePurge bool) {
 		th.timer = nil
 	}
 
+	th.needInvoke = false
 }
 
 // NewThrottle creates a throttled instance that invokes given functions only once in every interval.
diff --git a/retry_test.go b/retry_test.go
index ed935a98..246b5403 100644
--- a/retry_test.go
+++ b/retry_test.go
@@ -503,36 +503,34 @@ func TestNewThrottle(t *testing.T) {
 	t.Parallel()
 	is := assert.New(t)
 	callCount := 0
-	f1 := func() { callCount++ }
+	f1 := func() {
+		callCount++
+	}
 	th, purge := NewThrottle(10*time.Millisecond, f1)
 
 	is.Equal(0, callCount)
-	for i := 0; i < 10; i++ {
-		th()
-	}
-	is.Equal(0, callCount)
-
-	time.Sleep(11 * time.Millisecond)
-	is.Equal(1, callCount)
-	for i := 0; i < 10; i++ {
-		th()
+	for i := 0; i < 7; i++ {
+		var wg sync.WaitGroup
+		for j := 0; j < 100; j++ {
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+				th()
+			}()
+		}
+		wg.Wait()
+		time.Sleep(5 * time.Millisecond)
 	}
-	is.Equal(1, callCount)
+	// 35 ms passed
+	is.Equal(3, callCount)
 
 	purge()
-	is.Equal(2, callCount)
+	// awaits go routine which invokes given functions to be scheduled
+	time.Sleep(1 * time.Millisecond)
+	is.Equal(4, callCount)
 
 	// pause a little bit without calling
-	time.Sleep(11 * time.Millisecond)
-	is.Equal(2, callCount)
-	for i := 0; i < 10; i++ {
-		th()
-	}
-	is.Equal(2, callCount)
-
-	time.Sleep(11 * time.Millisecond)
-	is.Equal(3, callCount)
-	purge()
+	time.Sleep(20 * time.Millisecond)
 	is.Equal(4, callCount)
 
 }