69 lines
1.4 KiB
Go
69 lines
1.4 KiB
Go
package task
|
||
|
||
import (
|
||
model "github.com/cloudreve/Cloudreve/v3/models"
|
||
"github.com/cloudreve/Cloudreve/v3/pkg/conf"
|
||
"github.com/cloudreve/Cloudreve/v3/pkg/util"
|
||
)
|
||
|
||
// TaskPoll 要使用的任务池
|
||
var TaskPoll Pool
|
||
|
||
type Pool interface {
|
||
Add(num int)
|
||
Submit(job Job)
|
||
}
|
||
|
||
// AsyncPool 带有最大配额的任务池
|
||
type AsyncPool struct {
|
||
// 容量
|
||
idleWorker chan int
|
||
}
|
||
|
||
// Add 增加可用Worker数量
|
||
func (pool *AsyncPool) Add(num int) {
|
||
for i := 0; i < num; i++ {
|
||
pool.idleWorker <- 1
|
||
}
|
||
}
|
||
|
||
// ObtainWorker 阻塞直到获取新的Worker
|
||
func (pool *AsyncPool) obtainWorker() Worker {
|
||
select {
|
||
case <-pool.idleWorker:
|
||
// 有空闲Worker名额时，返回新Worker
|
||
return &GeneralWorker{}
|
||
}
|
||
}
|
||
|
||
// FreeWorker 添加空闲Worker
|
||
func (pool *AsyncPool) freeWorker() {
|
||
pool.Add(1)
|
||
}
|
||
|
||
// Submit 开始提交任务
|
||
func (pool *AsyncPool) Submit(job Job) {
|
||
go func() {
|
||
util.Log().Debug("Waiting for Worker.")
|
||
worker := pool.obtainWorker()
|
||
util.Log().Debug("Worker obtained.")
|
||
worker.Do(job)
|
||
util.Log().Debug("Worker released.")
|
||
pool.freeWorker()
|
||
}()
|
||
}
|
||
|
||
// Init 初始化任务池
|
||
func Init() {
|
||
maxWorker := model.GetIntSetting("max_worker_num", 10)
|
||
TaskPoll = &AsyncPool{
|
||
idleWorker: make(chan int, maxWorker),
|
||
}
|
||
TaskPoll.Add(maxWorker)
|
||
util.Log().Info("Initialize task queue with WorkerNum = %d", maxWorker)
|
||
|
||
if conf.SystemConfig.Mode == "master" {
|
||
Resume(TaskPoll)
|
||
}
|
||
}
|