naza/pkg/taskpool/pool.go

// Copyright 2019, Chef.  All rights reserved.
// https://github.com/q191201771/naza
//
// Use of this source code is governed by a MIT-style license
// that can be found in the License file.
//
// Author: Chef (191201771@qq.com)

package taskpool

import (
	"sync"
)

type taskWrapper struct {
	taskFn TaskFn
	param  []interface{}
}

type pool struct {
	maxWorkerNum int

	m              sync.Mutex
	totalWorkerNum int
	idleWorkerList []*worker
	blockTaskList  []taskWrapper
}

func newPool(option Option) *pool {
	p := pool{
		maxWorkerNum: option.MaxWorkerNum,
	}
	for i := 0; i < option.InitWorkerNum; i++ {
		p.newWorker()
	}
	return &p
}

func (p *pool) Go(task TaskFn, param ...interface{}) {
	tw := taskWrapper{
		taskFn: task,
		param:  param,
	}
	var w *worker
	p.m.Lock()
	if len(p.idleWorkerList) != 0 {
		// 还有空闲worker

		w = p.idleWorkerList[len(p.idleWorkerList)-1]
		p.idleWorkerList = p.idleWorkerList[0 : len(p.idleWorkerList)-1]
		w.Go(tw)
	} else {
		// 无空闲worker

		if p.maxWorkerNum == 0 ||
			(p.maxWorkerNum > 0 && p.totalWorkerNum < p.maxWorkerNum) {
			// 无最大worker限制，或还未达到限制

			p.newWorkerWithTask(tw)
		} else {
			// 已达到限制

			p.blockTaskList = append(p.blockTaskList, tw)
		}
	}
	p.m.Unlock()
}

func (p *pool) KillIdleWorkers() {
	p.m.Lock()
	p.totalWorkerNum = p.totalWorkerNum - len(p.idleWorkerList)
	for i := range p.idleWorkerList {
		p.idleWorkerList[i].Stop()
	}
	p.idleWorkerList = p.idleWorkerList[0:0]
	p.m.Unlock()
}

func (p *pool) GetCurrentStatus() Status {
	p.m.Lock()
	defer p.m.Unlock()
	return Status{
		TotalWorkerNum: p.totalWorkerNum,
		IdleWorkerNum:  len(p.idleWorkerList),
		BlockTaskNum:   len(p.blockTaskList),
	}
}

func (p *pool) newWorker() *worker {
	w := NewWorker(p)
	w.Start()
	p.idleWorkerList = append(p.idleWorkerList, w)
	p.totalWorkerNum++
	return w
}

func (p *pool) newWorkerWithTask(task taskWrapper) {
	w := NewWorker(p)
	w.Start()
	w.Go(task)
	p.totalWorkerNum++
}

func (p *pool) onIdle(w *worker) {
	p.m.Lock()
	if len(p.blockTaskList) == 0 {
		// 没有等待执行的任务

		p.idleWorkerList = append(p.idleWorkerList, w)
	} else {
		t := p.blockTaskList[0]
		p.blockTaskList = p.blockTaskList[1:]
		w.Go(t)
	}
	p.m.Unlock()
}