micro/util/pool/default.go

115 lines
2.0 KiB
Go
Raw Normal View History

2019-07-28 20:56:18 +03:00
package pool
import (
"context"
2019-07-28 20:56:18 +03:00
"sync"
"time"
"github.com/google/uuid"
"github.com/unistack-org/micro/v3/network/transport"
2019-07-28 20:56:18 +03:00
)
type pool struct {
tr transport.Transport
2019-07-28 20:56:18 +03:00
conns map[string][]*poolConn
size int
ttl time.Duration
sync.Mutex
2019-07-28 20:56:18 +03:00
}
type poolConn struct {
created time.Time
transport.Client
id string
2019-07-28 20:56:18 +03:00
}
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) {
2019-07-28 20:56:18 +03:00
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...)
2019-07-28 20:56:18 +03:00
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
}