116 lines
2.0 KiB
Go
116 lines
2.0 KiB
Go
package pool
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/unistack-org/micro/v3/network/transport"
|
|
)
|
|
|
|
type pool struct {
|
|
size int
|
|
ttl time.Duration
|
|
tr transport.Transport
|
|
|
|
sync.Mutex
|
|
conns map[string][]*poolConn
|
|
}
|
|
|
|
type poolConn struct {
|
|
transport.Client
|
|
id string
|
|
created time.Time
|
|
}
|
|
|
|
func newPool(options Options) *pool {
|
|
return &pool{
|
|
size: options.Size,
|
|
tr: options.Transport,
|
|
ttl: options.TTL,
|
|
conns: make(map[string][]*poolConn),
|
|
}
|
|
}
|
|
|
|
func (p *pool) Close() error {
|
|
p.Lock()
|
|
for k, c := range p.conns {
|
|
for _, conn := range c {
|
|
conn.Client.Close()
|
|
}
|
|
delete(p.conns, k)
|
|
}
|
|
p.Unlock()
|
|
return nil
|
|
}
|
|
|
|
// NoOp the Close since we manage it
|
|
func (p *poolConn) Close() error {
|
|
return nil
|
|
}
|
|
|
|
func (p *poolConn) Id() string {
|
|
return p.id
|
|
}
|
|
|
|
func (p *poolConn) Created() time.Time {
|
|
return p.created
|
|
}
|
|
|
|
func (p *pool) Get(ctx context.Context, addr string, opts ...transport.DialOption) (Conn, error) {
|
|
p.Lock()
|
|
conns := p.conns[addr]
|
|
|
|
// while we have conns check age and then return one
|
|
// otherwise we'll create a new conn
|
|
for len(conns) > 0 {
|
|
conn := conns[len(conns)-1]
|
|
conns = conns[:len(conns)-1]
|
|
p.conns[addr] = conns
|
|
|
|
// if conn is old kill it and move on
|
|
if d := time.Since(conn.Created()); d > p.ttl {
|
|
conn.Client.Close()
|
|
continue
|
|
}
|
|
|
|
// we got a good conn, lets unlock and return it
|
|
p.Unlock()
|
|
|
|
return conn, nil
|
|
}
|
|
|
|
p.Unlock()
|
|
|
|
// create new conn
|
|
c, err := p.tr.Dial(ctx, addr, opts...)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &poolConn{
|
|
Client: c,
|
|
id: uuid.New().String(),
|
|
created: time.Now(),
|
|
}, nil
|
|
}
|
|
|
|
func (p *pool) Release(conn Conn, err error) error {
|
|
// don't store the conn if it has errored
|
|
if err != nil {
|
|
return conn.(*poolConn).Client.Close()
|
|
}
|
|
|
|
// otherwise put it back for reuse
|
|
p.Lock()
|
|
conns := p.conns[conn.Remote()]
|
|
if len(conns) >= p.size {
|
|
p.Unlock()
|
|
return conn.(*poolConn).Client.Close()
|
|
}
|
|
p.conns[conn.Remote()] = append(conns, conn.(*poolConn))
|
|
p.Unlock()
|
|
|
|
return nil
|
|
}
|