Repository: zeromicro/go-zero
Pull request #3896 (2024-02-06 05:05:04)

feat: support ScheduleImmediately in TaskRunner

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/zeromicro/go-zero/pull/3896?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zeromicro) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`10f1d93`)](https://app.codecov.io/gh/zeromicro/go-zero/commit/10f1d93e2aa053f443f24e6307c02d7c9bbe47e1?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zeromicro) 94.42% compared to head [(`1332f45`)](https://app.codecov.io/gh/zeromicro/go-zero/pull/3896?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zeromicro) 94.43%.

> :exclamation: Current head 1332f45 differs from pull request most recent head fa2b858. Consider uploading reports for the commit fa2b858 to get more accurate results

<details><summary>Additional details and impacted files</summary>



| [Files](https://app.codecov.io/gh/zeromicro/go-zero/pull/3896?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zeromicro) | Coverage Δ | |
|---|---|---|
| [core/threading/taskrunner.go](https://app.codecov.io/gh/zeromicro/go-zero/pull/3896?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=zeromicro#diff-Y29yZS90aHJlYWRpbmcvdGFza3J1bm5lci5nbw==) | `100.00% <100.00%> (ø)` | |

</details>

Patch:
diff --git a/core/threading/taskrunner.go b/core/threading/taskrunner.go
index 955beeab..18ac4f6c 100644
--- a/core/threading/taskrunner.go
+++ b/core/threading/taskrunner.go
@@ -1,13 +1,20 @@
 package threading
 
 import (
+	"errors"
+	"sync"
+
 	"github.com/zeromicro/go-zero/core/lang"
 	"github.com/zeromicro/go-zero/core/rescue"
 )
 
+// ErrTaskRunnerBusy is the error that indicates the runner is busy.
+var ErrTaskRunnerBusy = errors.New("task runner is busy")
+
 // A TaskRunner is used to control the concurrency of goroutines.
 type TaskRunner struct {
 	limitChan chan lang.PlaceholderType
+	waitGroup sync.WaitGroup
 }
 
 // NewTaskRunner returns a TaskRunner.
@@ -19,13 +26,47 @@ func NewTaskRunner(concurrency int) *TaskRunner {
 
 // Schedule schedules a task to run under concurrency control.
 func (rp *TaskRunner) Schedule(task func()) {
+	// Why we add waitGroup first, in case of race condition on starting a task and wait returns.
+	// For example, limitChan is full, and the task is scheduled to run, but the waitGroup is not added,
+	// then the wait returns, and the task is then scheduled to run, but caller thinks all tasks are done.
+	// the same reason for ScheduleImmediately.
+	rp.waitGroup.Add(1)
 	rp.limitChan <- lang.Placeholder
 
 	go func() {
 		defer rescue.Recover(func() {
 			<-rp.limitChan
+			rp.waitGroup.Done()
 		})
 
 		task()
 	}()
 }
+
+// ScheduleImmediately schedules a task to run immediately under concurrency control.
+// It returns ErrTaskRunnerBusy if the runner is busy.
+func (rp *TaskRunner) ScheduleImmediately(task func()) error {
+	// Why we add waitGroup first, check the comment in Schedule.
+	rp.waitGroup.Add(1)
+	select {
+	case rp.limitChan <- lang.Placeholder:
+	default:
+		rp.waitGroup.Done()
+		return ErrTaskRunnerBusy
+	}
+
+	go func() {
+		defer rescue.Recover(func() {
+			<-rp.limitChan
+			rp.waitGroup.Done()
+		})
+		task()
+	}()
+
+	return nil
+}
+
+// Wait waits all running tasks to be done.
+func (rp *TaskRunner) Wait() {
+	rp.waitGroup.Wait()
+}


Test patch:
diff --git a/core/threading/taskrunner_test.go b/core/threading/taskrunner_test.go
index 81cefc82..7771760a 100644
--- a/core/threading/taskrunner_test.go
+++ b/core/threading/taskrunner_test.go
@@ -2,32 +2,52 @@ package threading
 
 import (
 	"runtime"
-	"sync"
 	"sync/atomic"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 )
 
-func TestRoutinePool(t *testing.T) {
+func TestTaskRunner_Schedule(t *testing.T) {
 	times := 100
 	pool := NewTaskRunner(runtime.NumCPU())
 
 	var counter int32
-	var waitGroup sync.WaitGroup
 	for i := 0; i < times; i++ {
-		waitGroup.Add(1)
 		pool.Schedule(func() {
 			atomic.AddInt32(&counter, 1)
-			waitGroup.Done()
 		})
 	}
 
-	waitGroup.Wait()
+	pool.Wait()
 
 	assert.Equal(t, times, int(counter))
 }
 
+func TestTaskRunner_ScheduleImmediately(t *testing.T) {
+	cpus := runtime.NumCPU()
+	times := cpus * 2
+	pool := NewTaskRunner(cpus)
+
+	var counter int32
+	for i := 0; i < times; i++ {
+		err := pool.ScheduleImmediately(func() {
+			atomic.AddInt32(&counter, 1)
+			time.Sleep(time.Millisecond * 100)
+		})
+		if i < cpus {
+			assert.Nil(t, err)
+		} else {
+			assert.ErrorIs(t, err, ErrTaskRunnerBusy)
+		}
+	}
+
+	pool.Wait()
+
+	assert.Equal(t, cpus, int(counter))
+}
+
 func BenchmarkRoutinePool(b *testing.B) {
 	queue := NewTaskRunner(runtime.NumCPU())
 	for i := 0; i < b.N; i++ {