package main

import (
	"fmt"
	"time"

	"github.com/gammazero/workerpool"
)

func main() {
	wp := workerpool.New(2)
	requests := []string{"alpha", "beta", "gamma", "delta", "epsilon"}

	for _, r := range requests {
		wp.Submit(func() {
			fmt.Printf("%s: Handling request: %s\n", time.Now().Format(time.RFC3339), r)
			time.Sleep(1 * time.Second)
		})
	}

	wp.StopWait()
}

// WorkerPool 是 Go 協(xié)程的集合池，用于確保同時處理請求的協(xié)程數(shù)量嚴格受控于預設的上限值
type WorkerPool struct {
	maxWorkers   int                 // 最大工作協(xié)程數(shù)
	taskQueue    chan func()         // 任務提交隊列
	workerQueue  chan func()         // 工作協(xié)程消費隊列
	stoppedChan  chan struct{}       // 停止完成通知通道
	stopSignal   chan struct{}       // 停止信號通道
	waitingQueue deque.Deque[func()] // 等待隊列（雙端隊列）
	stopLock     sync.Mutex          // 停止操作互斥鎖
	stopOnce     sync.Once           // 控制只停止一次
	stopped      bool                // 是否已經(jīng)停止
	waiting      int32               // 等待隊列中任務計數(shù)
	wait         bool                // 協(xié)程池退出時是否等待已入隊任務執(zhí)行完成
}

// New 創(chuàng)建并啟動協(xié)程池
// maxWorkers 參數(shù)指定可以并發(fā)執(zhí)行任務的最大工作協(xié)程數(shù)。
func New(maxWorkers int) *WorkerPool {
	// 至少有一個 worker
	if maxWorkers < 1 {
		maxWorkers = 1
	}

	// 實例化協(xié)程池對象
	pool := &WorkerPool{
		maxWorkers:  maxWorkers,
		taskQueue:   make(chan func()),
		workerQueue: make(chan func()),
		stopSignal:  make(chan struct{}),
		stoppedChan: make(chan struct{}),
	}

	// 啟動任務調度器
	go pool.dispatch()

	return pool
}

// Submit 將任務函數(shù)提交到工作池隊列等待執(zhí)行，不會等待任務執(zhí)行完成
func (p *WorkerPool) Submit(task func()) {
	if task != nil {
		p.taskQueue <- task
	}
}

// 任務派發(fā)，循環(huán)的將下一個排隊中的任務發(fā)送給可用的工作協(xié)程（worker）執(zhí)行
func (p *WorkerPool) dispatch() {
	defer close(p.stoppedChan)            // 保證調度器退出時關閉停止通知通道
	timeout := time.NewTimer(idleTimeout) // 創(chuàng)建 2 秒周期的空閑檢測定時器
	var workerCount int                   // 當前活躍 worker 計數(shù)器
	var idle bool                         // 空閑狀態(tài)標識
	var wg sync.WaitGroup                 // 用于等待所有 worker 完成

Loop:
	for { // 主循環(huán)處理任務分發(fā)

		// 隊列優(yōu)先模式：優(yōu)先檢測等待隊列
		if p.waitingQueue.Len() != 0 {
			if !p.processWaitingQueue() {
				break Loop // 等待隊列為空，退出循環(huán)
			}
			continue
		}

		// 直通模式：開始處理提交上來的新任務
		select {
		case task, ok := <-p.taskQueue: // 接收到新任務
			if !ok { // 協(xié)程池停止時會關閉任務通道，如果 !ok 說明協(xié)程池已停止，退出循環(huán)
				break Loop
			}

			select {
			case p.workerQueue <- task: // 嘗試派發(fā)任務
			default: // 沒有空閑的 worker，無法立即派發(fā)任務
				if workerCount < p.maxWorkers { // 如果協(xié)程池中的活躍協(xié)程數(shù)量小于最大值，那么創(chuàng)建一個新的協(xié)程（worker）來執(zhí)行任務
					wg.Add(1)
					go worker(task, p.workerQueue, &wg) // 創(chuàng)建新的 worker 執(zhí)行任務
					workerCount++                       // worker 記數(shù)加 1
				} else { // 已達協(xié)程池容量上限
					p.waitingQueue.PushBack(task)                              // 將任務提交到等待隊列
					atomic.StoreInt32(&p.waiting, int32(p.waitingQueue.Len())) // 原子更新等待計數(shù)
				}
			}
			idle = false // 標記為非空閑
		case <-timeout.C: // 空閑超時處理
			// 在一個空閑超時周期內，存在空閑的 workers，那么停止一個 worker
			if idle && workerCount > 0 {
				if p.killIdleWorker() { // 回收一個 worker
					workerCount-- // worker 計數(shù)減 1
				}
			}
			idle = true                // 標記為空閑
			timeout.Reset(idleTimeout) // 復用定時器
		}
	}

	if p.wait { // 調用了 StopWait() 方法，需要運行等待隊列中的任務，直至隊列清空
		p.runQueuedTasks()
	}

	// 終止所有 worker
	for workerCount > 0 {
		p.workerQueue <- nil // 發(fā)送終止信號給 worker
		workerCount--        // worker 計數(shù)減 1，直至為 0 退出循環(huán)
	}
	wg.Wait() // 阻塞等待所有 worker 完成

	timeout.Stop() // 停止定時器
}

// 隊列優(yōu)先模式：優(yōu)先檢測等待隊列
if p.waitingQueue.Len() != 0 {
    if !p.processWaitingQueue() {
        break Loop // 協(xié)程池已經(jīng)停止
    }
    continue // 隊列不為空則繼續(xù)下一輪循環(huán)
}

// 處理等待隊列
func (p *WorkerPool) processWaitingQueue() bool {
	select {
	case task, ok := <-p.taskQueue: // 接收到新任務
		if !ok { // 協(xié)程池停止時會關閉任務通道，如果 !ok 說明協(xié)程池已停止，返回 false，不再繼續(xù)處理
			return false
		}
		p.waitingQueue.PushBack(task) // 將新任務加入等待隊列隊尾
	case p.workerQueue <- p.waitingQueue.Front(): // 從等待隊列隊頭獲取任務并放入工作隊列
		p.waitingQueue.PopFront() // 任務已經(jīng)開始處理，所以要從等待隊列中移除任務
	}
	atomic.StoreInt32(&p.waiting, int32(p.waitingQueue.Len())) // 原子修改等待隊列中任務計數(shù)
	return true
}

// 直通模式：開始處理提交上來的新任務
select {
case task, ok := <-p.taskQueue: // 接收到新任務
    ...
case <-timeout.C: // 空閑超時處理
    ...
}

case task, ok := <-p.taskQueue: // 接收到新任務
    if !ok { // 協(xié)程池停止時會關閉任務通道，如果 !ok 說明協(xié)程池已停止，退出循環(huán)
        break Loop
    }

    select {
    case p.workerQueue <- task: // 嘗試派發(fā)任務
    default: // 沒有空閑的 worker，無法立即派發(fā)任務
        if workerCount < p.maxWorkers { // 如果協(xié)程池中的活躍協(xié)程數(shù)量小于最大值，那么創(chuàng)建一個新的協(xié)程（worker）來執(zhí)行任務
            wg.Add(1)
            go worker(task, p.workerQueue, &wg) // 創(chuàng)建新的 worker 執(zhí)行任務
            workerCount++                       // worker 記數(shù)加 1
        } else { // 已達協(xié)程池容量上限
            p.waitingQueue.PushBack(task)                              // 將任務提交到等待隊列
            atomic.StoreInt32(&p.waiting, int32(p.waitingQueue.Len())) // 原子更新等待計數(shù)
        }
    }
    idle = false // 標記為非空閑

// 工作協(xié)程，執(zhí)行任務并在收到 nil 信號時停止
func worker(task func(), workerQueue chan func(), wg *sync.WaitGroup) {
	for task != nil { // 循環(huán)執(zhí)行任務，直至接收到終止信號 nil
		task()               // 執(zhí)行任務
		task = <-workerQueue // 接收新任務
	}
	wg.Done() // 標記 worker 完成
}

case <-timeout.C: // 空閑超時處理
    // 在一個空閑超時周期內，存在空閑的 workers，那么停止一個 worker
    if idle && workerCount > 0 {
        if p.killIdleWorker() { // 回收一個 worker
            workerCount-- // worker 計數(shù)減 1
        }
    }
    idle = true                // 標記為空閑
    timeout.Reset(idleTimeout) // 復用定時器

// 停止一個空閑 worker
func (p *WorkerPool) killIdleWorker() bool {
	select {
	case p.workerQueue <- nil: // 發(fā)送終止信號給工作協(xié)程（worker）
		// Sent kill signal to worker.
		return true
	default:
		// No ready workers. All, if any, workers are busy.
		return false
	}
}