如何实现一个协程池

1 协程的概念

2 协程池的作用

3 手写实现协程池

import (
	"context"
	"errors"
	"math/rand"
	"sync"
	"sync/atomic"
	"time"
)

type Pool struct {
	ctx        context.Context
	expiration time.Duration
	workers    chan *worker
	capacity   int32
	alives     int32
	cancel     func()
	blocking   bool
	sync.Mutex
	sync.Once
}

func NewPool(capacity int32, blocking bool) *Pool {
	p := Pool{
		capacity: capacity,
		workers:  make(chan *worker, capacity),
		blocking: blocking,
	}
	p.ctx, p.cancel = context.WithCancel(context.Background())
	go p.releaseExpiration()
	return &p
}

func (p *Pool) Submit(task func()) error {
	select {
	case w, ok := <-p.workers:
		if !ok {
			return errors.New("closed pool")
		}
		w.submit(task)
		return nil
	default:
	}

	// 没有现成可用的 worker，尝试创建一个新的 worker
	for {
		alives := p.alives
		if alives == p.capacity && !p.blocking {
			return errors.New("overload")
		}

		if alives == p.capacity {
			return p.submitInBlockingMode(task)
		}

		if atomic.CompareAndSwapInt32(&p.alives, alives, alives+1) {
			break
		}
	}
	w := newWorker(p)
	w.submit(task)
	return nil
}

func (p *Pool) submitInBlockingMode(task func()) error {
	// 阻塞直到有 worker 可用
	w, ok := <-p.workers
	if !ok {
		return errors.New("closed pool")
	}
	w.submit(task)
	return nil
}

func (p *Pool) Close() {
	p.Do(p.close)
}

func (p *Pool) close() {
	close(p.workers)
	p.cancel()
}

func (p *Pool) put(w *worker) {
	select {
	case <-p.ctx.Done():
	default:
		p.workers <- w
	}
}

// 每个 worker 记录一个时间，记录最后一次被使用的时间
// pool 开启一个清理协程，间隔一定的时间对过期的 worker 进行清理
func (p *Pool) releaseExpiration() {
	for {
		time.Sleep(time.Duration(rand.Intn(100)) * time.Millisecond)
		select {
		case <-p.ctx.Done():
			return
		case worker := <-p.workers:
			// 未过期则归还
			if !worker.expired(p.expiration) {
				p.put(worker)
				continue
			}
			// 过期了，则回收
			atomic.AddInt32(&p.alives, -1)
			worker.close()
		}
	}
}

// 构造函数制定一个 size，确定好 worker 容量上限
// 每当 submit 的时候，查看是否有可用的，有则复用，没有则创建一个新的
// 每个 worker 有一个过期时间，启动一个协程持续对过期的 worker 进行回收
type worker struct {
	sync.Once
	tasks      chan func()
	lastUsedAt time.Time
	pool       *Pool
}

func newWorker(p *Pool) *worker {
	w := worker{
		tasks:      make(chan func()),
		lastUsedAt: time.Now(),
		pool:       p,
	}
	go w.start()
	return &w
}

func (w *worker) start() {
	for t := range w.tasks {
		t()
		w.lastUsedAt = time.Now()
		// 每次执行完任务后将自己归还回池子当中
		w.pool.put(w)
	}
}

func (w *worker) submit(task func()) {
	w.tasks <- task
}

func (w *worker) expired(duration time.Duration) bool {
	return time.Now().Sub(w.lastUsedAt) >= duration
}

func (w *worker) close() {
	w.Do(func() {
		close(w.tasks)
	})
}

4 对标 ants 的差异

5 性能测试