golang 协程池 动态扩缩容

参考 github.com/panjf2000/ants

go 复制代码
package main

import (
	"fmt"
	"sync"
	"sync/atomic"
	"time"

	_ "github.com/panjf2000/ants"
)

type pool struct {
	// 协程池最大容量
	cap int32
	// 当前运行的协程个数
	run   int32
	block bool

	// 空闲goroutinue
	idleWorkers []*worker
	// 最大空闲时间,超过就挥手
	idleTimeoutSec uint32

	lock  sync.Mutex
	cond  *sync.Cond
	close chan struct{}
	wg    sync.WaitGroup
}

func NewPool(cap int32) *pool {
	c := &pool{
		cap: cap,
		//idleWorkers: make([]*worker, 0, cap),
		close: make(chan struct{}),
	}
	c.cond = sync.NewCond(&c.lock)
	go c.checkIdleWork()
	return c
}

func (p *pool) checkIdleWork() {
	for {
		select {
		case <-time.After(time.Second * 5):
			{
				p.lock.Lock()
				defer p.lock.Unlock()

				var i = 0
				for ; i < len(p.idleWorkers) &&
					uint32(time.Now().Unix())-p.idleWorkers[i].lastWorkTs >= p.idleTimeoutSec; i++ {
				}

				fmt.Printf("checkIdleWork del=%d\n", i)
				if i > 0 {
					for j := 0; j < i; j++ {
						atomic.AddInt32(&p.run, -1)
						p.idleWorkers[j].task <- nil
					}

					p.idleWorkers = p.idleWorkers[i:]
					fmt.Printf("checkIdleWork now run=%d, del=%d\n", len(p.idleWorkers), i)
				}
			}
		}
	}
}

func (p *pool) Done() {
	p.wg.Done()
}

func (p *pool) Close() {
	fmt.Printf("pool Close\n")
	close(p.close)
	p.wg.Wait()
}

func (p *pool) backPool(w *worker) {
	p.lock.Lock()
	p.idleWorkers = append(p.idleWorkers, w)
	p.lock.Unlock()
	p.cond.Signal()
	fmt.Printf("backPool\n")
}

func (p *pool) Submit(task func()) error {

	p.lock.Lock()
	if len(p.idleWorkers) > 0 {
		work := p.idleWorkers[len(p.idleWorkers)-1]
		p.idleWorkers = p.idleWorkers[:len(p.idleWorkers)-1]
		work.task <- task
		p.lock.Unlock()
		fmt.Printf("use idle\n")
		return nil
	}
	p.lock.Unlock()

	if p.run < p.cap {
		fmt.Printf("use new\n")
		work := &worker{
			pool: p,
			task: make(chan func()),
		}

		p.wg.Add(1)
		work.run()
		atomic.AddInt32(&p.run, 1)

		work.task <- task
		return nil

	} else {

	loop:
		if p.block {
			return fmt.Errorf("pool max")
		}

		fmt.Printf("pool max, wait\n")
		p.lock.Lock()
		p.cond.Wait()

		if len(p.idleWorkers) > 0 {
			work := p.idleWorkers[len(p.idleWorkers)-1]
			p.idleWorkers = p.idleWorkers[:len(p.idleWorkers)-1]
			work.task <- task
			p.lock.Unlock()
			fmt.Printf("use idle\n")
			return nil
		}
		p.lock.Unlock()

		if p.run < p.cap {
			work := &worker{
				pool: p,
				task: make(chan func()),
			}
			p.wg.Add(1)
			work.run()
			atomic.AddInt32(&p.run, 1)
			work.task <- task

			fmt.Printf("use new\n")
			return nil
		}

		goto loop

		return nil
	}
}

type worker struct {
	pool       *pool
	task       chan func()
	lastWorkTs uint32
}

func (w *worker) run() {

	go func() {
		defer w.pool.Done()

		for {
			select {
			case <-w.pool.close:
				return
			case taskFunc := <-w.task:
				if taskFunc == nil {
					return
				}
				taskFunc()
				w.lastWorkTs = uint32(time.Now().Unix())
				w.pool.backPool(w)
			}
		}
	}()
}

func main() {

	p := NewPool(20)

	for i := 0; i < 5; i++ {

		p.Submit(func() {
			func(idx int) {
				fmt.Printf("work %d ..\n", idx)
				time.Sleep(time.Second * 4)
			}(i)
		})

	}

	time.Sleep(time.Second * 10)
	p.Close()
}
相关推荐
Ai 编码助手41 分钟前
Go 语言 API 限流实战:保障系统稳定性的护盾
开发语言·后端·golang
玩大数据的龙威2 小时前
【ArcGIS Pro】完整的nc文件整理表格模型构建流程及工具练习数据分享
开发语言·python
唐棣棣3 小时前
期末速成C++【知识点汇总完】
开发语言·c++
yannan201903133 小时前
【数据结构】(Python)差分数组。差分数组与树状数组结合
开发语言·python·算法
中國移动丶移不动3 小时前
Java List 源码解析——从基础到深度剖析
java·后端·list
javaTodo3 小时前
消息队列kafka详解:Kafka架构介绍
后端
m0_748248234 小时前
Springboot项目:使用MockMvc测试get和post接口(含单个和多个请求参数场景)
java·spring boot·后端
WongKyunban4 小时前
Bash Shell知识合集
开发语言·chrome·bash
努力的小雨4 小时前
KES的执行计划分析与索引优化
数据库·后端