Add some tests for workerPool.ScheduleWithTimeout()

Author: johnteee

worker/pool.go

@@ -78,8 +78,9 @@ type DefaultWorkerPool struct {
 	spawnWorkerCh fpgo.ChannelQueue[int]
 	lastAliveTime time.Time
 
-	scheduleWaitCh    fpgo.ChannelQueue[int]
-	scheduleWaitCount int
+	scheduleWaitQueue  fpgo.Queue[fpgo.ChannelQueue[int]]
+	scheduleWaitChPool sync.Pool
+	scheduleWaitCount  int
 
 	// Settings
 	DefaultWorkerPoolSettings
@@ -93,12 +94,15 @@ func NewDefaultWorkerPool(jobQueue *fpgo.BufferedChannelQueue[func()], settings
 	workerPool := &DefaultWorkerPool{
 		jobQueue: jobQueue,
 
-		spawnWorkerCh:  fpgo.NewChannelQueue[int](1),
-		scheduleWaitCh: fpgo.NewChannelQueue[int](1),
+		spawnWorkerCh:     fpgo.NewChannelQueue[int](1),
+		scheduleWaitQueue: fpgo.NewConcurrentQueue[fpgo.ChannelQueue[int]](fpgo.NewLinkedListQueue[fpgo.ChannelQueue[int]]()),
 
 		// Settings
 		DefaultWorkerPoolSettings: *settings,
 	}
+	workerPool.scheduleWaitChPool.New = func() interface{} {
+		return fpgo.NewChannelQueue[int](1)
+	}
 	go workerPool.spawnLoop()
 
 	return workerPool
@@ -108,12 +112,12 @@ func NewDefaultWorkerPool(jobQueue *fpgo.BufferedChannelQueue[func()], settings
 func (workerPoolSelf *DefaultWorkerPool) trySpawn() {
 	workerPoolSelf.lock.RLock()
 	batchSize := workerPoolSelf.workerBatchSize
-	if batchSize < 1 {
-		batchSize = 1
-	}
-	expectedWorkerCount := workerPoolSelf.jobQueue.Count() / batchSize
-	if workerPoolSelf.jobQueue.Count()%batchSize > 0 {
-		expectedWorkerCount++
+	var expectedWorkerCount int
+	if batchSize > 0 {
+		expectedWorkerCount = workerPoolSelf.jobQueue.Count() / batchSize
+		if workerPoolSelf.jobQueue.Count()%batchSize > 0 {
+			expectedWorkerCount++
+		}
 	}
 	if workerPoolSelf.workerSizeStandBy > expectedWorkerCount {
 		expectedWorkerCount = workerPoolSelf.workerSizeStandBy
@@ -198,8 +202,18 @@ func (workerPoolSelf *DefaultWorkerPool) generateWorkerWithMaximum(maximum int)
 	loopLabel:
 		for {
 			workerPoolSelf.lastAliveTime = time.Now()
+			workerPoolSelf.lock.RLock()
 			if workerPoolSelf.scheduleWaitCount > 0 {
-				workerPoolSelf.scheduleWaitCh.Offer(1)
+				ch, _ := workerPoolSelf.scheduleWaitQueue.Poll()
+				if ch != nil {
+					ch.Offer(1)
+					defer workerPoolSelf.scheduleWaitChPool.Put(ch)
+				}
+			}
+			workerPoolSelf.lock.RUnlock()
+
+			if workerPoolSelf.IsClosed() {
+				return
 			}
 
 			select {
@@ -229,6 +243,12 @@ func (workerPoolSelf *DefaultWorkerPool) SetJobQueue(jobQueue *fpgo.BufferedChan
 	return workerPoolSelf
 }
 
+// SetScheduleWaitQueue Set the ScheduleWaitQueue(WARNING: if the pool has started to use, doing this is not safe)
+func (workerPoolSelf *DefaultWorkerPool) SetScheduleWaitQueue(scheduleWaitQueue fpgo.Queue[fpgo.ChannelQueue[int]]) *DefaultWorkerPool {
+	workerPoolSelf.scheduleWaitQueue = scheduleWaitQueue
+	return workerPoolSelf
+}
+
 // SetIsJobQueueClosedWhenClose Set is the JobQueue closed when the WorkerPool.Close()
 func (workerPoolSelf *DefaultWorkerPool) SetIsJobQueueClosedWhenClose(isJobQueueClosedWhenClose bool) *DefaultWorkerPool {
 	workerPoolSelf.isJobQueueClosedWhenClose = isJobQueueClosedWhenClose
@@ -328,9 +348,20 @@ func (workerPoolSelf *DefaultWorkerPool) ScheduleWithTimeout(fn func(), timeout
 		return err
 	}
 
+	// if retry < 1 {
+	// 	retry = 1
+	// }
+	// retryInterval := timeout / time.Duration(retry+1)
 	deadline := time.Now().Add(timeout)
 	workerPoolSelf.addScheduleWaitCount(1)
-	go workerPoolSelf.addScheduleWaitCount(-1)
+	defer workerPoolSelf.addScheduleWaitCount(-1)
+
+	scheduleWaitCh := workerPoolSelf.scheduleWaitChPool.Get().(fpgo.ChannelQueue[int])
+	defer workerPoolSelf.scheduleWaitChPool.Put(scheduleWaitCh)
+	err = workerPoolSelf.scheduleWaitQueue.Offer(scheduleWaitCh)
+	if err != nil {
+		return err
+	}
 
 	for {
 		if workerPoolSelf.IsClosed() {
@@ -342,8 +373,15 @@ func (workerPoolSelf *DefaultWorkerPool) ScheduleWithTimeout(fn func(), timeout
 			return err
 		}
 
+		// fmt.Println(retryInterval)
+		// fmt.Println(deadline.Sub(now))
+		// if time.Now().After(deadline) {
+		// 	return ErrWorkerPoolScheduleTimeout
+		// }
+		// time.Sleep(retryInterval)
+
 		select {
-		case <-workerPoolSelf.scheduleWaitCh:
+		case <-scheduleWaitCh:
 			continue
 		case <-time.After(deadline.Sub(time.Now())):
 			return ErrWorkerPoolScheduleTimeout

worker/pool_test.go

@@ -54,9 +54,42 @@ func TestWorkerPool(t *testing.T) {
 	}
 	time.Sleep(1 * time.Millisecond)
 	// BatchSize: 3, Jobs: 4 -> ceil(4/3) = 2 workers
-	// assert.Equal(t, 2, defaultWorkerPool.workerCount)
+	assert.GreaterOrEqual(t, defaultWorkerPool.workerCount, 2)
 	defaultWorkerPool.SetWorkerSizeStandBy(0)
 	time.Sleep(10 * time.Millisecond)
 	// workerSizeStandBy: 1
 	assert.Equal(t, 0, defaultWorkerPool.workerCount)
 }
+
+func TestScheduleWithTimeout(t *testing.T) {
+	var workerPool WorkerPool
+	var err error
+	defaultWorkerPool := NewDefaultWorkerPool(fpgo.NewBufferedChannelQueue[func()](3, 1, 3), nil).
+		SetSpawnWorkerDuration(1 * time.Millisecond / 10).
+		SetWorkerExpiryDuration(2 * time.Millisecond).
+		SetWorkerSizeMaximum(3).
+		SetWorkerSizeStandBy(3).
+		SetWorkerBatchSize(0)
+	// defaultWorkerPool.PreAllocWorkerSize(5)
+	workerPool = defaultWorkerPool
+
+	// Test ScheduleWithTimeout
+	// channel: 3 positions, buffered 1 => 4 positions
+	for i := 0; i < 4; i++ {
+		v := i
+		err = workerPool.ScheduleWithTimeout(func() {
+			// Nothing to do
+			time.Sleep(3 * time.Millisecond)
+			t.Log(v)
+		}, 1*time.Millisecond)
+		assert.NoError(t, err)
+	}
+	err = workerPool.Schedule(func() {})
+	assert.Equal(t, ErrWorkerPoolJobQueueIsFull, err)
+	err = workerPool.ScheduleWithTimeout(func() {}, 1*time.Millisecond)
+	assert.Equal(t, ErrWorkerPoolScheduleTimeout, err)
+
+	// defaultWorkerPool.SetWorkerSizeMaximum(3)
+	err = workerPool.ScheduleWithTimeout(func() {}, 10*time.Millisecond)
+	assert.Equal(t, nil, err)
+}