Golang goroutine pool 的实现

Golang 语言很好的支持高并发场景，goroutine相比java的thread开销更小。但是大量的goroutine会带来内存开销，如果无限的创建goroutine则会出现内存溢出的灾难，所以萌生出了goroutine pool 的想法，仿照java中的ThreadPoolExecutor实现一个简单的Golang 版本的协程池。代码如下：

package goroutine_pool

import (
	"fmt"
	"sync"
	"time"
)
/*任务的接口*/
type Task interface {
	Run()
}
//任务的具体结构
type TestTask struct {
	I int
}

//实现接口
func (t TestTask) Run() {
	time.Sleep(300 * time.Millisecond)
	fmt.Printf("No: %d\n",t.I)
}

//协程池结构
type GroutinePool struct {
	lock sync.Mutex  //处理并发的锁
	coreGoroutineSize int32 //核心工作协程数量
	maxGoroutineSize  int32 //最大工作协程数量
	queue chan Task //任务队列
	duration time.Duration //最大等待时间
	status   int32   //线程池状态0:初始化；1:运行；2:shutdown;3:stop
	size     int32   //当前线程池中的工作线程数
	maxQueueSize int32 //工作队列中的最大数量
	reject func ()  //丢弃任务之后回调该方法
}
//原子操作函数，保证方法f对全局变量的操作是原子的
func (pool *GroutinePool) atomicOper(f func() bool) bool {
	pool.lock.Lock()
	defer pool.lock.Unlock()
	return f()
}
//构造一个协程池
func NewGroutinePool(core int32,max int32,queueSize int32,duration time.Duration,rejectHandler func()) *GroutinePool{
	pool := GroutinePool{}
	pool.status = 0
	pool.coreGoroutineSize = core
	pool.maxGoroutineSize = max
	pool.size = 0
	pool.maxQueueSize = queueSize
	pool.queue = make(chan Task,pool.maxQueueSize)
	pool.duration = duration
	pool.lock = sync.Mutex{}
	pool.reject = rejectHandler
	return &pool
}
//启动协程池
func (pool *GroutinePool) Start() {
	if pool.status != 0 {
		panic("GroutinePool has started!")
	}
	pool.status = 1
}

//向协程池放入任务，该逻辑参考java ThreadPoolExecutor
func (pool *GroutinePool)Put(t Task) bool {
     if pool.status != 1 {
		mes := fmt.Sprintf("GroutinePool status: %d\n",pool.status)
		panic(mes)
	}
	return pool.atomicOper(func() bool{
	//如果当前协程数小于核心协程数则创建新协程处理任务
		if pool.size < pool.coreGoroutineSize {
			pool.size++
			go pool.run(t)
			return true
		} else if len(pool.queue) < int(pool.maxQueueSize) {
		//当前协程数大于等于核心协程数并且队列没有满，则将任务放入队列
			pool.queue <- t
			return true
		} else if pool.size < pool.maxGoroutineSize {
		//如果协程数小于最大协程数则创建新协程处理任务
			pool.size++
			go pool.run(t)
			return true
		} else {
		//协程数已经是最大协程数则拒绝该任务
			pool.reject()
			return false
		}
	})
}

//工作协程处理任务，不断从队列中获取任务，直到超时
func (pool *GroutinePool) run(t Task) {
    if t != nil {//正在处理的任务不为nil,则处理任务t
		t.Run()
	}
	loop:
	for pool.status == 1 || pool.status == 2 {
		select {
			case t := <-pool.queue:
				t.Run()
			case <-time.After(pool.duration):
			//如果超时还没有拿到任务，说明该协程已经空闲了一段时间，该协程的生命周期可以结束
				if pool.atomicOper(func() bool {
				//当协程池中的协程数量大于核心数量的时候，将协程结束
					if pool.size > pool.coreGoroutineSize {
						pool.size--
						return true
					} else if pool.status == 2 {
					//如果协程池状态为shutdown,需要查看队列中是否还有任务，如果有任务则保留协程，否则结束协程
						if len(pool.queue) > 0 {
							return false
						} else {
							pool.size--
							return true
						}
					} else {
						return false
					}
				}) {
					break loop
				}
		}
	}
	pool.atomicOper(func() bool {//如果协程池中协程数为0，则关闭协程池
		if pool.status == 2 && pool.size == 0 {
			pool.status = 3
			return true
		} else {
			return false
		}
	})
}

func (pool *GroutinePool) Shutdown() {
	pool.status = 2
}

func (pool *GroutinePool) ShutdownNow() {
	pool.status = 3
}

func (pool *GroutinePool) GetStatus() int32 {
	return pool.status
}

func (pool *GroutinePool) GetQueueSize() int32 {
	return int32(len(pool.queue))
}

func (pool *GroutinePool) GetPoolSize() int32 {
	return pool.size
}

测试代码

func main() {
//构造一个协程池
	pool := goroutine_pool.NewGroutinePool(1,2,5,500 * time.Millisecond,func(){
		fmt.Printf("reject task !\n")
	})
	pool.Start()//启动
//放任务
	for i := 0;i < 20;i++ {
		fmt.Printf("i : %d\n",i)
		t := &goroutine_pool.TestTask{}
		t.I = i
		if !pool.Put(t) {
			time.Sleep(100 * time.Millisecond)
		}
	}

	time.Sleep(5 * time.Second)

	fmt.Printf("pool status : %d;\n queue size : %d;\n goroutine size : %d\n",
		pool.GetStatus(),pool.GetQueueSize(),pool.GetPoolSize())
//分别测试shutdown,shutdownNow
	//pool.Shutdown()
	pool.ShutdownNow()

	for pool.GetStatus() != 3 {
		/*fmt.Printf("pool status : %d;\n queue size : %d;\n goroutine size : %d\n",
			pool.GetStatus(),pool.GetQueueSize(),pool.GetPoolSize())
		time.Sleep(time.Millisecond * 10)*/
		fmt.Printf("wait pool shutdown ! \n")
	}

	fmt.Printf("pool status : %d;\n queue size : %d;\n goroutine size : %d\n",
		pool.GetStatus(),pool.GetQueueSize(),pool.GetPoolSize())

}