Golang实现自定义协程池

type ConnPool struct {
    opt *Options

    dialErrorsNum uint32 // atomic

    lastDialErrorMu sync.RWMutex
    lastDialError   error

    queue chan struct{}

    connsMu      sync.Mutex
    conns        []*Conn
    idleConns    []*Conn
    poolSize     int
    idleConnsLen int

    stats Stats

    _closed uint32 // atomic
}

type Options struct {
    Dialer  func() (net.Conn, error)
    OnClose func(*Conn) error

    PoolSize           int
    MinIdleConns       int
    MaxConnAge         time.Duration
    PoolTimeout        time.Duration
    IdleTimeout        time.Duration
    IdleCheckFrequency time.Duration
}

func NewConnPool(opt *Options) *ConnPool {
    p := &ConnPool{
        opt: opt,

        queue:     make(chan struct{}, opt.PoolSize),
        conns:     make([]*Conn, 0, opt.PoolSize),
        idleConns: make([]*Conn, 0, opt.PoolSize),
    }

    for i := 0; i < opt.MinIdleConns; i++ {
        p.checkMinIdleConns()
    }

    if opt.IdleTimeout > 0 && opt.IdleCheckFrequency > 0 {
        go p.reaper(opt.IdleCheckFrequency)
    }

    return p
}

if p.closed() {
    return nil, ErrClosed
}

err := p.waitTurn()
if err != nil {
    return nil, err
}

func (p *ConnPool) waitTurn() error {
    select {
    case p.queue <- struct{}{}:
        return nil
    default:
        timer := timers.Get().(*time.Timer)
        timer.Reset(p.opt.PoolTimeout)

        select {
        case p.queue <- struct{}{}:
            if !timer.Stop() {
                <-timer.C
            }
            timers.Put(timer)
            return nil
        case <-timer.C:
            timers.Put(timer)
            return ErrPoolTimeout
        }
    }
}

for {
    p.connsMu.Lock()
    cn := p.popIdle()
    p.connsMu.Unlock()

    if cn == nil {
        break
    }

    if p.isStaleConn(cn) {
        _ = p.CloseConn(cn)
        continue
    }

    return cn, nil
}

newcn, err := p._NewConn(true)
if err != nil {
    p.freeTurn()
    return nil, err
}

return newcn, nil

cn, err := p.newConn(pooled)
if err != nil {
    return nil, err
}

p.connsMu.Lock()
p.conns = append(p.conns, cn)
if pooled {
    if p.poolSize < p.opt.PoolSize {
        p.poolSize++
    } else {
        cn.pooled = false
    }
}
p.connsMu.Unlock()
return cn, nil

func (p *ConnPool) Put(cn *Conn) {
    if !cn.pooled {
        p.Remove(cn)
        return
    }

    p.connsMu.Lock()
    p.idleConns = append(p.idleConns, cn)
    p.idleConnsLen++
    p.connsMu.Unlock()
    p.freeTurn()
}

func (p *ConnPool) removeConn(cn *Conn) {
    p.connsMu.Lock()
    for i, c := range p.conns {
        if c == cn {
            p.conns = append(p.conns[:i], p.conns[i+1:]...)
            if cn.pooled {
                p.poolSize--
                p.checkMinIdleConns()
            }
            break
        }
    }
    p.connsMu.Unlock()
}

var n int
for {
    p.getTurn()

    p.connsMu.Lock()
    cn := p.reapStaleConn()
    p.connsMu.Unlock()

    if cn != nil {
        p.removeConn(cn)
    }

    p.freeTurn()

    if cn != nil {
        p.closeConn(cn)
        n++
    } else {
        break
    }
}
return n, nil

type Pool struct {
    Dial func() (Conn, error)

    TestOnBorrow func(c Conn, t time.Time) error

    MaxIdle int
    MaxActive int

    IdleTimeout time.Duration

    Wait bool

    MaxConnLifetime time.Duration

    chInitialized uint32

    mu     sync.Mutex
    closed bool
    active int
    ch     chan struct{}
    idle   idleList
}

func newPool(addr string) *redis.Pool {
   return &redis.Pool{
     MaxIdle: 3,
     IdleTimeout: 240 * time.Second,
     Dial: func () (redis.Conn, error) { return redis.Dial("tcp", addr)},
   }
 }

// Handle limit for p.Wait == true.
if p.Wait && p.MaxActive > 0 {
    p.lazyInit()
    if ctx == nil {
        <-p.ch
    } else {
        select {
        case <-p.ch:
        case <-ctx.Done():
            return nil, ctx.Err()
        }
    }
}

p.mu.Lock()

// Prune stale connections at the back of the idle list.
if p.IdleTimeout > 0 {
    n := p.idle.count
    for i := 0; i < n && p.idle.back != nil && p.idle.back.t.Add(p.IdleTimeout).Before(nowFunc()); i++ {
        pc := p.idle.back
        p.idle.popBack()
        p.mu.Unlock()
        pc.c.Close()
        p.mu.Lock()
        p.active--
    }
}

// Get idle connection from the front of idle list.
for p.idle.front != nil {
    pc := p.idle.front
    p.idle.popFront()
    p.mu.Unlock()
    if (p.TestOnBorrow == nil || p.TestOnBorrow(pc.c, pc.t) == nil) &&
        (p.MaxConnLifetime == 0 || nowFunc().Sub(pc.created) < p.MaxConnLifetime) {
        return pc, nil
    }
    pc.c.Close()
    p.mu.Lock()
    p.active--
}

// Check for pool closed before dialing a new connection.
if p.closed {
    p.mu.Unlock()
    return nil, errors.New("redigo: get on closed pool")
}

// Handle limit for p.Wait == false.
if !p.Wait && p.MaxActive > 0 && p.active >= p.MaxActive {
    p.mu.Unlock()
    return nil, ErrPoolExhausted
}

p.active++
p.mu.Unlock()

c, err := p.Dial()
if err != nil {
    c = nil
    p.mu.Lock()
    p.active--
    if p.ch != nil && !p.closed {
        p.ch <- struct{}{}
    }
    p.mu.Unlock()
}
return &poolConn{c: c, created: nowFunc()}, err

func (p *ConnPool) Get() (*Conn, error) // go-redis
func (p *Pool) Get() Conn // redigo

p.mu.Lock()
if !p.closed && !forceClose {
    pc.t = nowFunc()
    p.idle.pushFront(pc)
    if p.idle.count > p.MaxIdle {
        pc = p.idle.back
        p.idle.popBack()
    } else {
        pc = nil
    }
}

if pc != nil {
    p.mu.Unlock()
    pc.c.Close()
    p.mu.Lock()
    p.active--
}

if p.ch != nil && !p.closed {
    p.ch <- struct{}{}
}
p.mu.Unlock()
return nil

func serveHome(w http.ResponseWriter, r *http.Request) {
     conn := pool.Get()
     defer conn.Close()
     //...
 }

type Pool struct {
    pool        chan *redis.Client
    reservePool chan *redis.Client
    df          DialFunc

    po opts

    limited chan bool

    initDoneCh chan bool // used for tests
    stopCh     chan bool

    Network, Addr string
}

var defaultPoolOpts []Opt
// if pool size is 0 don't do any pinging, cause there'd be no point
if size > 0 {
    defaultPoolOpts = append(defaultPoolOpts, PingInterval(10*time.Second/time.Duration(size)))
}

var po opts
for _, opt := range append(defaultPoolOpts, os...) {
    opt(&po)
}

p := Pool{
    Network:     network,
    Addr:        addr,
    po:          po,
    pool:        make(chan *redis.Client, size),
    reservePool: make(chan *redis.Client, po.overflowSize),
    limited:     make(chan bool, po.createLimitBuffer),
    df:          df,
    initDoneCh:  make(chan bool),
    stopCh:      make(chan bool),
}

if po.pingInterval > 0 {
    doEvery(po.pingInterval, func() {
        // instead of using Cmd/Get, which might make a new connection,
        // we only check from the pool
        select {
        case conn := <-p.pool:
            // we don't care if PING errors since Put will handle that
            conn.Cmd("PING")
            p.Put(conn)
        default:
        }
    })
}

if po.overflowSize > 0 {
    doEvery(po.overflowDrainInterval, func() {
        // remove one from the reservePool, if there is any, and try putting it
        // into the main pool
        select {
        case conn := <-p.reservePool:
            select {
            case p.pool <- conn:
            default:
                // if the main pool is full then just close it
                conn.Close()
            }
        default:
        }
    })
}

mkConn := func() error {
    client, err := df(network, addr)
    if err == nil {
        p.pool <- client
    }
    return err
}

if size > 0 {
    // make one connection to make sure the redis instance is actually there
    if err := mkConn(); err != nil {
        return &p, err
    }
}

// make the rest of the connections in the background, if any fail it's fine
go func() {
    for i := 0; i < size-1; i++ {
        mkConn()
    }
    close(p.initDoneCh)
}()

return &p, nil

select {
case conn := <-p.pool:
    return conn, nil
case conn := <-p.reservePool:
    return conn, nil
case <-p.stopCh:
    return nil, errors.New("pool emptied")
default:
    var timeoutCh <-chan time.Time
    if p.po.getTimeout > 0 {
        timer := time.NewTimer(p.po.getTimeout)
        defer timer.Stop()
        timeoutCh = timer.C
    }

    select {
    case conn := <-p.pool:
        return conn, nil
    case conn := <-p.reservePool:
        return conn, nil
    case <-timeoutCh:
        return nil, ErrGetTimeout
    case <-p.limited:
        return p.df(p.Network, p.Addr)
    }
}

select {
case <-p.stopCh:
    conn.Close()
    return
default:
}

select {
case p.pool <- conn:
default:
    if p.po.overflowSize == 0 {
        conn.Close()
        return
    }

    // we need a separate select here since it's indeterminate which case go
    // will select and we want to always prefer the main pool over the reserve
    select {
    case p.reservePool <- conn:
    default:
        conn.Close()
    }
}

package channel_pool

// The pool interface
type Pool interface {
	// Get returns a new connection from the pool. Closing the connections puts
	// it back to the Pool. Closing it when the pool is destroyed or full will
	// be counted as an error.
	Get() (interface{}, error)

	// Put puts the connection into the pool instead of closing it.
	Put(interface{}) error

	// CloseConn directly close the connection
	CloseConn(interface{}) error

	// ShutDown closes the pool and all its connections.
	// After ShutDown() the pool is no longer usable.
	ShutDown() error

	// Len returns the current number of connections of the pool.
	Len() int
}

// generalConn is a wrapper around the connection
type generalConn struct {
	conn     interface{}
	t        time.Time
	unusable bool
}

// the pool
type channelPool struct {
	mu           sync.RWMutex
	conns        chan *generalConn
	factory      func() (interface{}, error)
	close        func(interface{}) error
	ping         func(interface{}) error
	idleTimeout  time.Duration
	waitTimeOut  time.Duration
	maxActive    int
	openingConns int
	connReqs     []chan *generalConn
}

// Configs for pool
type Options struct {
	// The number of the connections when initiate the pool
	// Also, the least connection number of the pool
	InitialCap int

	// Max connection number in the pool
	MaxCap int

	// Max idle number in the pool
	MaxIdle int

	// The method the build the connection
	Factory func() (interface{}, error)

	// The method to close the connection
	Close func(interface{}) error

	// Check connection health
	Ping func(interface{}) error

	// Max life time for idle connection
	IdleTimeout time.Duration

	// Max time to get a connection from pool
	// Else this will return a errs.MaxActiveConnectionErr
	WaitTimeout time.Duration
}

func newChannelPool() (Pool, error) {
	ops := Options{
		InitialCap:  InitialCap,
		MaxCap:      MaximumCap,
		MaxIdle:     MaxIdleCap,
		Factory:     factory,
		Ping:        pingFunc,
		Close:       closeFunc,
		IdleTimeout: IdleTimeout,
		WaitTimeout: WaitTimeout,
	}
	return NewChannelPool(&ops)
}

func xxx() {
    conn, _ := p.Get()
    cli := conn.(*rpc.Client)
    defer p.Put(conn)
	
    // 使用cli连接处理逻辑
    ...
}

func TestPool_Get(t *testing.T) {
	p, err := newChannelPool()
	if err != nil {
		t.Errorf("create pool error: %s", err)
	}
	defer p.ShutDown()

	_, err = p.Get()
	if err != nil {
		t.Errorf("Get error: %s", err)
	}

	// after one get, current capacity should be lowered by one.
	if p.Len() != (InitialCap - 1) {
		t.Errorf("Get error. Expecting %d, got %d", InitialCap-1, p.Len())
	}
}

type Number int

type Args struct {
	A, B int
}

func rpcServer() {
	number := new(Number)
	_ = rpc.Register(number)
	rpc.HandleHTTP()

	l, e := net.Listen("tcp", address)
	if e != nil {
		panic(e)
	}
	go http.Serve(l, nil)
}

func (n *Number) Multiply(args *Args, reply *int) error {
	*reply = args.A * args.B
	return nil
}