给 GPT 的 Prompts
User
gopool.go
package gopoolimport ( "sync" "time" )
// task represents a function that will be executed by a worker. // It returns a result and an error. type task func() (interface{}, error)
// goPool represents a pool of workers. type goPool struct { workers []*worker workerStack []int maxWorkers int // Set by WithMinWorkers(), used to adjust the number of workers. Default equals to maxWorkers. minWorkers int // tasks are added to this channel first, then dispatched to workers. Default buffer size is 1 million. taskQueue chan task // Set by WithRetryCount(), used to retry a task when it fails. Default is 0. retryCount int lock sync.Locker cond *sync.Cond // Set by WithTimeout(), used to set a timeout for a task. Default is 0, which means no timeout. timeout time.Duration // Set by WithResultCallback(), used to handle the result of a task. Default is nil. resultCallback func(interface{}) // Set by WithErrorCallback(), used to handle the error of a task. Default is nil. errorCallback func(error) // adjustInterval is the interval to adjust the number of workers. Default is 1 second. adjustInterval time.Duration }
// NewGoPool creates a new pool of workers. func NewGoPool(maxWorkers int, opts ...Option) *goPool { pool := &goPool{ maxWorkers: maxWorkers, // Set minWorkers to maxWorkers by default minWorkers: maxWorkers, workers: make([]*worker, maxWorkers), workerStack: make([]int, maxWorkers), taskQueue: make(chan task, 1e6), retryCount: 0, lock: new(sync.Mutex), timeout: 0, adjustInterval: 1 * time.Second, } // Apply options for _, opt := range opts { opt(pool) } if pool.cond == nil { pool.cond = sync.NewCond(pool.lock) } // Create workers with the minimum number. Don't use pushWorker() here. for i := 0; i < pool.minWorkers; i++ { worker := newWorker() pool.workers[i] = worker pool.workerStack[i] = i worker.start(pool, i) } go pool.adjustWorkers() go pool.dispatch() return pool }
// AddTask adds a task to the pool. func (p *goPool) AddTask(t task) { p.taskQueue <- t }
// Release stops all workers and releases resources. func (p *goPool) Release() { close(p.taskQueue) p.cond.L.Lock() for len(p.workerStack) != p.minWorkers { p.cond.Wait() } p.cond.L.Unlock() for _, worker := range p.workers { close(worker.taskQueue) } p.workers = nil p.workerStack = nil }
func (p *goPool) popWorker() int { p.lock.Lock() workerIndex := p.workerStack[len(p.workerStack)-1] p.workerStack = p.workerStack[:len(p.workerStack)-1] p.lock.Unlock() return workerIndex }
func (p *goPool) pushWorker(workerIndex int) { p.lock.Lock() p.workerStack = append(p.workerStack, workerIndex) p.lock.Unlock() p.cond.Signal() }
// adjustWorkers adjusts the number of workers according to the number of tasks in the queue. func (p *goPool) adjustWorkers() { ticker := time.NewTicker(p.adjustInterval) defer ticker.Stop()
for range ticker.C {
p.cond.L.Lock()
if len(p.taskQueue) > len(p.workerStack)*3/4 && len(p.workerStack) < p.maxWorkers {
// Double the number of workers until it reaches the maximum
newWorkers := min(len(p.workerStack)*2, p.maxWorkers) - len(p.workerStack)
for i := 0; i < newWorkers; i++ {
worker := newWorker()
p.workers = append(p.workers, worker)
p.workerStack = append(p.workerStack, len(p.workers)-1)
worker.start(p, len(p.workers)-1)
}
} else if len(p.taskQueue) == 0 && len(p.workerStack) > p.minWorkers {
// Halve the number of workers until it reaches the minimum
removeWorkers := max((len(p.workerStack)-p.minWorkers)/2, p.minWorkers)
p.workers = p.workers[:len(p.workers)-removeWorkers]
p.workerStack = p.workerStack[:len(p.workerStack)-removeWorkers]
}
p.cond.L.Unlock()
}
}
// dispatch dispatches tasks to workers. func (p *goPool) dispatch() { for t := range p.taskQueue { p.cond.L.Lock() for len(p.workerStack) == 0 { p.cond.Wait() } p.cond.L.Unlock() workerIndex := p.popWorker() p.workers[workerIndex].taskQueue <- t } }
func min(a, b int) int { if a < b { return a } return b }
func max(a, b int) int { if a > b { return a } return b }
调用了 Release() 方法之后,如果 goPool.taskQueue 不为空,则会导致 dispatch() 方法往已经 close 的 Channel 中发送数据。我想要给 goPool 添加一个 Wait 方法,方便用户在调用 Release 方法之前先通过 Wait 方法等待已经添加的任务完成执行。
DevChat
你的建议很好,我们可以添加一个 Wait
方法来等待所有已添加的任务完成执行。这个方法可以通过检查 taskQueue
是否为空来实现。如果 taskQueue
不为空,我们就让当前的 goroutine 等待,直到 taskQueue
为空。
以下是添加 Wait
方法的代码:
// Wait waits for all tasks to finish.
func (p *goPool) Wait() {
for len(p.taskQueue) > 0 {
time.Sleep(100 * time.Millisecond)
}
}
在这个方法中,我们使用了一个无限循环来检查 taskQueue
是否为空。如果 taskQueue
不为空,我们就让当前的 goroutine 等待一段时间(这里我们选择了100毫秒),然后再次检查。这个过程会一直持续,直到 taskQueue
为空。
然后,你可以在调用 Release
方法之前先调用 Wait
方法,以确保所有已添加的任务都已完成执行:
pool.Wait()
pool.Release()
这样,你就可以确保在释放资源之前,所有的任务都已经完成执行了。