88 lines
1.4 KiB
Go
88 lines
1.4 KiB
Go
|
package pool
|
||
|
|
|
||
|
|
import (
|
||
|
|
"net"
|
||
|
|
"sync"
|
||
|
|
"sync/atomic"
|
||
|
|
)
|
||
|
|
|
||
|
|
var (
|
||
|
|
defaultMaxConns = 20
|
||
|
|
defaultStrategy = &RoundrobinStrategy{
|
||
|
|
lastSelected: initialRoundrobinAtomic(),
|
||
|
|
}
|
||
|
|
)
|
||
|
|
|
||
|
|
func initialRoundrobinAtomic() atomic.Int64 {
|
||
|
|
a := atomic.Int64{}
|
||
|
|
a.Store(-1)
|
||
|
|
return a
|
||
|
|
}
|
||
|
|
|
||
|
|
type ConnSelectionStrategy interface {
|
||
|
|
Select() int
|
||
|
|
}
|
||
|
|
|
||
|
|
type RoundrobinStrategy struct {
|
||
|
|
lastSelected atomic.Int64
|
||
|
|
}
|
||
|
|
|
||
|
|
func (r *RoundrobinStrategy) Select() int {
|
||
|
|
return int(r.lastSelected.Add(1))
|
||
|
|
}
|
||
|
|
|
||
|
|
type ConnPool struct {
|
||
|
|
m sync.RWMutex
|
||
|
|
strategy ConnSelectionStrategy
|
||
|
|
conns []net.Conn
|
||
|
|
}
|
||
|
|
|
||
|
|
type newConnPoolOpts struct {
|
||
|
|
strategy ConnSelectionStrategy
|
||
|
|
maxConns int
|
||
|
|
}
|
||
|
|
|
||
|
|
func newNewConnPoolOpts() newConnPoolOpts {
|
||
|
|
return newConnPoolOpts{
|
||
|
|
strategy: defaultStrategy,
|
||
|
|
maxConns: defaultMaxConns,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
type NewConnPoolOpt func(p *newConnPoolOpts)
|
||
|
|
|
||
|
|
func WithStrategy(s ConnSelectionStrategy) NewConnPoolOpt {
|
||
|
|
return func(p *newConnPoolOpts) {
|
||
|
|
p.strategy = s
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func WithMaxConns(mc int) NewConnPoolOpt {
|
||
|
|
return func(p *newConnPoolOpts) {
|
||
|
|
p.maxConns = mc
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func NewConnPool(opts ...NewConnPoolOpt) *ConnPool {
|
||
|
|
o := newNewConnPoolOpts()
|
||
|
|
for _, opt := range opts {
|
||
|
|
opt(&o)
|
||
|
|
}
|
||
|
|
return &ConnPool{
|
||
|
|
conns: make([]net.Conn, 0),
|
||
|
|
strategy: o.strategy,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (p *ConnPool) SelectConn() net.Conn {
|
||
|
|
p.m.RLock()
|
||
|
|
defer p.m.RUnlock()
|
||
|
|
return p.conns[p.strategy.Select()]
|
||
|
|
}
|
||
|
|
|
||
|
|
func (p *ConnPool) AddConn(conn net.Conn) {
|
||
|
|
p.m.Lock()
|
||
|
|
defer p.m.Unlock()
|
||
|
|
p.conns = append(p.conns, conn)
|
||
|
|
}
|