github.com/cloudreve/Cloudreve/v3@v3.0.0-20240224133659-3edb00a6484c/pkg/task/pool.go (about)

     1  package task
     2  
     3  import (
     4  	model "github.com/cloudreve/Cloudreve/v3/models"
     5  	"github.com/cloudreve/Cloudreve/v3/pkg/conf"
     6  	"github.com/cloudreve/Cloudreve/v3/pkg/util"
     7  )
     8  
     9  // TaskPoll 要使用的任务池
    10  var TaskPoll Pool
    11  
    12  type Pool interface {
    13  	Add(num int)
    14  	Submit(job Job)
    15  }
    16  
    17  // AsyncPool 带有最大配额的任务池
    18  type AsyncPool struct {
    19  	// 容量
    20  	idleWorker chan int
    21  }
    22  
    23  // Add 增加可用Worker数量
    24  func (pool *AsyncPool) Add(num int) {
    25  	for i := 0; i < num; i++ {
    26  		pool.idleWorker <- 1
    27  	}
    28  }
    29  
    30  // ObtainWorker 阻塞直到获取新的Worker
    31  func (pool *AsyncPool) obtainWorker() Worker {
    32  	select {
    33  	case <-pool.idleWorker:
    34  		// 有空闲Worker名额时，返回新Worker
    35  		return &GeneralWorker{}
    36  	}
    37  }
    38  
    39  // FreeWorker 添加空闲Worker
    40  func (pool *AsyncPool) freeWorker() {
    41  	pool.Add(1)
    42  }
    43  
    44  // Submit 开始提交任务
    45  func (pool *AsyncPool) Submit(job Job) {
    46  	go func() {
    47  		util.Log().Debug("Waiting for Worker.")
    48  		worker := pool.obtainWorker()
    49  		util.Log().Debug("Worker obtained.")
    50  		worker.Do(job)
    51  		util.Log().Debug("Worker released.")
    52  		pool.freeWorker()
    53  	}()
    54  }
    55  
    56  // Init 初始化任务池
    57  func Init() {
    58  	maxWorker := model.GetIntSetting("max_worker_num", 10)
    59  	TaskPoll = &AsyncPool{
    60  		idleWorker: make(chan int, maxWorker),
    61  	}
    62  	TaskPoll.Add(maxWorker)
    63  	util.Log().Info("Initialize task queue with WorkerNum = %d", maxWorker)
    64  
    65  	if conf.SystemConfig.Mode == "master" {
    66  		Resume(TaskPoll)
    67  	}
    68  }