go协程池设计优化

    前些时间因项目需求，需使用协程池处理事件，所以在网上看好很多相关文章，做了总结，然后根据业务需求做了进一步的优化。项目部署在生产环境中运行了很长时间， 协程池处理事务比较稳定，目前还没有出现问题。所以分享给同行，能够相互学习。咋们就不多说，已代码为主。需要优化的可以多提提意见。

       一、协程池

package goroutine_pool

import (
    "runtime/debug"
    "sync"
    "sync/atomic"
    "time"
    "***/golib/log"
)

const (
    pool_status_ready    int32 = 0 // 就绪
    pool_status_running  int32 = 1 // 运行中
    pool_status_pre_exit int32 = 2 // 退出前
    pool_status_exit     int32 = 3 // 已退出
)

type GRPool struct {
    capacity       int           // 协程数量上限
    workerNum      int           // 已创建的工作者数
    expireDuration time.Duration // 空闲工作者过期时间

    
    lock     *sync.Mutex   // 同步锁
    condLock *sync.Mutex   // 繁忙时通知有工作者空闲
    cond     *sync.Cond

    cacheSize int         // 任务缓存channel大小
    taskCache chan GRTask // 任务缓存

    idleWorkers []*GRWorker // 空闲工作者队列
    workerCache []*GRWorker // 空闲工作者临时缓存，定期清理

    status int32           // 状态
    wg     *sync.WaitGroup // 控制等待所有工作正常结束
}

func NewGRPool(capacity int, cache_size int, expire_duration time.Duration) *GRPool {
    if 0 == capacity {
        return nil
    }

    condLock := new(sync.Mutex)
    return &GRPool{
        capacity:       capacity,
        workerNum:      0,
        expireDuration: expire_duration,
        lock:           new(sync.Mutex),
        condLock:       condLock,
        cond:           sync.NewCond(condLock),
        cacheSize:      cache_size,
        taskCache:      make(chan GRTask, cache_size),
        idleWorkers:    make([]*GRWorker, 0),
        workerCache:    make([]*GRWorker, 0),
        status:         pool_status_ready,
        wg:             &sync.WaitGroup{},
    }
}

func (this *GRPool) AddTask(task GRTask) {
    defer func() {
        if err := recover(); err != nil {
            log.Error("AddTask panic. err:%+v", err)
            debug.PrintStack()
        }
    }()

    this.taskCache <- task
}

func (this *GRPool) Run() {
    atomic.StoreInt32(&this.status, pool_status_running)

    this.wg.Add(1)
    go func() {
        defer this.wg.Done()

        for task := range this.taskCache {
            if nil == task {
                break
            }

            worker := this.fetchWorker()
            if worker != nil {
                worker.SetTask(task)
            } else {
                // 未获取到可用工作者--等待有工作者可用
                this.cond.L.Lock()
                this.cond.Wait()

                worker = this.fetchWorker()
                worker.SetTask(task)

                this.cond.L.Unlock()
            }
        }

        log.Info("GRPool run done")
    }()

    this.wg.Add(1)
    go this.purge()
}

func (this *GRPool) Stop() {
    if !atomic.CompareAndSwapInt32(&this.status, pool_status_running, pool_status_pre_exit) {
        return
    }

    log.Info("GRPool stop")
    this.taskCache <- nil
    this.wg.Wait()
    atomic.StoreInt32(&this.status, pool_status_exit)
    log.Info("GRPool exit")
}

func (this *GRPool) fetchWorker() *GRWorker {
    this.lock.Lock()
    defer this.lock.Unlock()

    n := len(this.idleWorkers) - 1
    if n >= 0 {
        // 空闲队列不为空，取出队列尾的worker
        worker := this.idleWorkers[n]
        this.idleWorkers[n] = nil
        this.idleWorkers = this.idleWorkers[:n]
        return worker
    } else {
        // 空闲队列为空
        if this.workerNum < this.capacity {
            // 协程数未达上限--新建worker
            worker := newGRWorker(this.expireDuration, this)
            worker.Run()
            this.workerNum++
            return worker
        }
        return nil
    }
}

func (this *GRPool) onWorkerDone(worker *GRWorker) {
    this.lock.Lock()
    defer this.lock.Unlock()
    this.idleWorkers = append(this.idleWorkers, worker)
    this.cond.Signal()
}

func (this *GRPool) purge() {
    defer this.wg.Done()

    ticker := time.Tick(time.Second)
    for _ = range ticker {
        if pool_status_pre_exit == atomic.LoadInt32(&this.status) {
            break
        }

        this.lock.Lock()

        idx := -1
        for i, worker := range this.idleWorkers {
            if worker.isExpired() {
                // 停止过期worker
                idx = i
                worker.Stop()
                this.workerNum--
                this.idleWorkers[i] = nil
            }
        }

        if idx > -1 {
            if idx >= len(this.idleWorkers)-1 {
                // 全部过期
                this.idleWorkers = this.idleWorkers[:0]
            } else {
                // 部分过期
                this.idleWorkers = this.idleWorkers[idx+1:]
            }
        }

        this.lock.Unlock()
    }
}

func (this *GRPool) Resize(capacity int) {
    this.lock.Lock()
    defer this.lock.Unlock()
    this.capacity = capacity
}

       二、  工作对象

package goroutine_pool

import (
    "runtime/debug"
    "sync"
    "time"
    "***/golib/log"
)

type GRWorker struct {
    pool           *GRPool
    taskCh         chan GRTask
    wg             *sync.WaitGroup
    expireDuration time.Duration
    recycleTime    time.Time
}

func newGRWorker(expire_duration time.Duration, pool *GRPool) *GRWorker {
    return &GRWorker{
        pool:           pool,
        taskCh:         make(chan GRTask),
        wg:             &sync.WaitGroup{},
        expireDuration: expire_duration,
        recycleTime:    time.Now(),
    }
}

func (this *GRWorker) SetTask(task GRTask) {
    defer func() {
        if err := recover(); err != nil {
            log.Debug("SetTask panic err:%+v", err)
            debug.PrintStack()
        }
    }()
    this.taskCh <- task
}

func (this *GRWorker) Run() {
    this.wg.Add(1)
    go func() {
        defer this.wg.Done()
        for task := range this.taskCh {
            if nil == task {
                break
            }

            // 执行
            task.Do()

            this.recycleTime = time.Now()
            this.pool.onWorkerDone(this)
        }
        log.Info("worker done")
    }()
}

func (this *GRWorker) Stop() {
    this.taskCh<- nil
    this.wg.Wait()
}

func (this *GRWorker) isExpired() bool {
    return time.Now().Sub(this.recycleTime) > this.expireDuration
}

三、任务调度

package goroutine_pool

type GRTask interface {
    Do()
}

四、测试

package main

import (
	"fmt"
	"goroutine_pool"
	"time"
)

type MyTask struct {
	my_task goroutine_pool.GRTask
}

func NewMyTask(msgPusher goroutine_pool.GRTask) *MyTask {
	return &MyTask{
		my_task: msgPusher,
	}
}

func (this *MyTask) Do() {
	fmt.Printf("协程调度执行任务\n")
}

func main() {
	gRPool := goroutine_pool.NewGRPool(10, 6, 10*time.Second)
	gRPool.Run()

	var msgPusher goroutine_pool.GRTask
	pushTask := NewMyTask(msgPusher)

	gRPool.AddTask(pushTask)
	gRPool.AddTask(pushTask)

	time.Sleep(5 * time.Second)
}

 

参考资料：

1、https://segmentfault.com/a/1190000018193161

2、https://www.jianshu.com/p/73624f96d286

3、https://blog.csdn.net/wyfhist/article/details/80288903

5、https://studygolang.com/articles/15477

来源：网络
所涉及知识点归原创文章所有，如有任何授权并注明出处，会立马删除处理。本文仅做学习交流。