2024-01-06 01:50:20 +08:00
|
|
|
package pool
|
|
|
|
|
|
|
|
import (
|
|
|
|
"sync"
|
|
|
|
)
|
|
|
|
|
|
|
|
type TaskPool struct {
|
|
|
|
workers int
|
|
|
|
queue chan Task
|
|
|
|
wg sync.WaitGroup
|
|
|
|
|
|
|
|
lck sync.Mutex
|
|
|
|
curTaskID int
|
2024-01-28 18:16:04 +08:00
|
|
|
waitMap map[int]chan error
|
2024-01-06 01:50:20 +08:00
|
|
|
}
|
|
|
|
|
2024-01-28 22:02:26 +08:00
|
|
|
type ErrTaskNotFound struct{}
|
|
|
|
|
|
|
|
func (m *ErrTaskNotFound) Error() string {
|
|
|
|
return "task not found"
|
|
|
|
}
|
|
|
|
|
2024-01-06 01:50:20 +08:00
|
|
|
func NewTaskPool(maxWorkers, bufferSize int) *TaskPool {
|
|
|
|
return &TaskPool{
|
|
|
|
workers: maxWorkers,
|
|
|
|
queue: make(chan Task, bufferSize),
|
2024-01-28 18:16:04 +08:00
|
|
|
waitMap: make(map[int]chan error),
|
2024-01-06 01:50:20 +08:00
|
|
|
curTaskID: 1, // task id starts from 1
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
func (tp *TaskPool) Start() {
|
|
|
|
for i := 1; i <= tp.workers; i++ { // worker id starts from 1
|
|
|
|
worker := NewWorker(i, tp.queue, tp)
|
|
|
|
tp.wg.Add(1)
|
|
|
|
go worker.Start(&tp.wg)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2024-01-28 18:16:04 +08:00
|
|
|
func (tp *TaskPool) AddTask(f func() error) int {
|
2024-01-06 01:50:20 +08:00
|
|
|
tp.lck.Lock()
|
|
|
|
|
|
|
|
id := tp.curTaskID
|
|
|
|
tp.curTaskID++
|
|
|
|
|
2024-01-28 18:16:04 +08:00
|
|
|
waitChan := make(chan error, 1)
|
2024-01-06 01:50:20 +08:00
|
|
|
tp.waitMap[id] = waitChan
|
|
|
|
|
2024-01-06 21:03:30 +08:00
|
|
|
tp.lck.Unlock()
|
|
|
|
|
|
|
|
task := Task{id: id, f: f}
|
|
|
|
tp.queue <- task
|
|
|
|
|
2024-01-06 01:50:20 +08:00
|
|
|
return id
|
|
|
|
}
|
|
|
|
|
2024-01-28 18:16:04 +08:00
|
|
|
func (tp *TaskPool) WaitForTask(taskID int) error {
|
2024-01-06 01:50:20 +08:00
|
|
|
tp.lck.Lock()
|
|
|
|
waitChan, ok := tp.waitMap[taskID]
|
|
|
|
if !ok {
|
|
|
|
tp.lck.Unlock()
|
2024-01-28 22:02:26 +08:00
|
|
|
return &ErrTaskNotFound{}
|
2024-01-06 01:50:20 +08:00
|
|
|
}
|
|
|
|
tp.lck.Unlock()
|
|
|
|
|
2024-01-28 18:16:04 +08:00
|
|
|
ret := <-waitChan
|
|
|
|
close(waitChan)
|
2024-01-06 01:50:20 +08:00
|
|
|
|
|
|
|
tp.lck.Lock()
|
2024-01-28 18:16:04 +08:00
|
|
|
delete(tp.waitMap, taskID)
|
|
|
|
tp.lck.Unlock()
|
|
|
|
|
|
|
|
return ret
|
|
|
|
}
|
2024-01-06 01:50:20 +08:00
|
|
|
|
2024-01-28 18:16:04 +08:00
|
|
|
func (tp *TaskPool) markTaskComplete(taskID int, err error) {
|
|
|
|
tp.lck.Lock()
|
2024-01-06 01:50:20 +08:00
|
|
|
waitChan, ok := tp.waitMap[taskID]
|
|
|
|
if !ok {
|
2024-01-28 22:02:26 +08:00
|
|
|
// should never happen here
|
|
|
|
panic("worker: task destroyed before completion")
|
2024-01-06 01:50:20 +08:00
|
|
|
}
|
2024-01-28 18:16:04 +08:00
|
|
|
tp.lck.Unlock()
|
2024-01-06 01:50:20 +08:00
|
|
|
|
2024-01-28 18:16:04 +08:00
|
|
|
waitChan <- err
|
2024-01-06 01:50:20 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
func (tp *TaskPool) Stop() {
|
|
|
|
close(tp.queue)
|
|
|
|
tp.wg.Wait()
|
|
|
|
}
|