2017-08-28 09:37:10 +00:00
|
|
|
package overflow_grpc_pool
|
|
|
|
|
|
|
|
import (
|
2017-08-30 06:32:53 +00:00
|
|
|
"context"
|
2017-08-28 09:37:10 +00:00
|
|
|
"errors"
|
|
|
|
"time"
|
|
|
|
|
|
|
|
"google.golang.org/grpc"
|
|
|
|
)
|
|
|
|
|
|
|
|
var (
|
|
|
|
// ErrFullPool is the error when the pool is already full
|
|
|
|
ErrPoolFull = errors.New("grpc pool: Pool is reached maximum capacity.")
|
|
|
|
)
|
|
|
|
|
|
|
|
type pooledInstance struct {
|
|
|
|
clientConn *grpc.ClientConn
|
|
|
|
idleTime time.Time
|
|
|
|
instance interface{}
|
|
|
|
inUse bool
|
|
|
|
}
|
|
|
|
|
|
|
|
type Pool interface {
|
|
|
|
Get() (interface{}, error)
|
|
|
|
Put(i interface{})
|
|
|
|
Capacity() int
|
|
|
|
Available() int
|
|
|
|
Destroy()
|
|
|
|
}
|
|
|
|
|
|
|
|
type pool struct {
|
2017-08-30 06:32:53 +00:00
|
|
|
ctx context.Context
|
2017-09-01 03:16:45 +00:00
|
|
|
handler PoolHandler
|
2017-08-28 09:37:10 +00:00
|
|
|
instances map[interface{}]*pooledInstance
|
|
|
|
instancesCh chan *pooledInstance
|
|
|
|
}
|
|
|
|
|
2017-09-01 03:16:45 +00:00
|
|
|
func New(ctx context.Context, handler PoolHandler) (Pool, error) {
|
2017-08-28 09:37:10 +00:00
|
|
|
var err error
|
|
|
|
p := &pool{
|
2017-08-30 06:32:53 +00:00
|
|
|
ctx: ctx,
|
2017-09-01 03:16:45 +00:00
|
|
|
handler: handler,
|
|
|
|
instances: make(map[interface{}]*pooledInstance, handler.GetMaxCapacity()),
|
|
|
|
instancesCh: make(chan *pooledInstance, handler.GetMaxCapacity()),
|
2017-08-28 09:37:10 +00:00
|
|
|
}
|
|
|
|
|
2017-09-01 03:16:45 +00:00
|
|
|
if 0 < handler.GetMaxIdle() {
|
|
|
|
for i := 0; i < handler.GetMaxIdle(); i++ {
|
2017-08-28 09:37:10 +00:00
|
|
|
err = p.createInstance()
|
|
|
|
if nil != err {
|
|
|
|
return nil, err
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
return p, nil
|
|
|
|
}
|
|
|
|
|
|
|
|
func (p *pool) createInstance() error {
|
2017-09-01 03:16:45 +00:00
|
|
|
maxCapacity := p.handler.GetMaxCapacity()
|
2017-08-28 09:37:10 +00:00
|
|
|
currentInstanceCount := len(p.instances)
|
|
|
|
if currentInstanceCount >= maxCapacity {
|
|
|
|
return ErrPoolFull
|
|
|
|
}
|
|
|
|
|
2017-09-01 03:16:45 +00:00
|
|
|
conn, i, err := p.handler.OnCreate()
|
2017-08-28 09:37:10 +00:00
|
|
|
if err != nil {
|
|
|
|
return err
|
|
|
|
}
|
|
|
|
pi := &pooledInstance{
|
|
|
|
clientConn: conn,
|
|
|
|
idleTime: time.Now(),
|
|
|
|
instance: i,
|
|
|
|
inUse: false,
|
|
|
|
}
|
|
|
|
p.instances[i] = pi
|
|
|
|
p.instancesCh <- pi
|
|
|
|
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
2017-08-30 06:32:53 +00:00
|
|
|
func (p *pool) destroyInstance(i interface{}) error {
|
2017-08-28 09:37:10 +00:00
|
|
|
var err error
|
|
|
|
pi, ok := p.instances[i]
|
|
|
|
if !ok {
|
2017-08-30 06:32:53 +00:00
|
|
|
return nil
|
2017-08-28 09:37:10 +00:00
|
|
|
}
|
|
|
|
err = pi.clientConn.Close()
|
|
|
|
if nil != err {
|
2017-08-30 06:32:53 +00:00
|
|
|
return err
|
2017-08-28 09:37:10 +00:00
|
|
|
}
|
|
|
|
delete(p.instances, i)
|
2017-08-30 06:32:53 +00:00
|
|
|
|
|
|
|
return nil
|
2017-08-28 09:37:10 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
func (p *pool) get() (*pooledInstance, error) {
|
|
|
|
var pi *pooledInstance
|
|
|
|
var err error
|
|
|
|
|
|
|
|
avail := len(p.instancesCh)
|
|
|
|
if 0 == avail {
|
|
|
|
if err = p.createInstance(); nil != err {
|
|
|
|
return nil, err
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
for {
|
|
|
|
select {
|
|
|
|
case pi = <-p.instancesCh:
|
|
|
|
// All good
|
|
|
|
default:
|
|
|
|
}
|
2017-09-01 03:16:45 +00:00
|
|
|
idleTimeout := p.handler.GetIdleTimeout()
|
2017-08-29 08:54:51 +00:00
|
|
|
if 1 < len(p.instancesCh) && 0 < idleTimeout && pi.idleTime.Add(idleTimeout).Before(time.Now()) {
|
2017-08-28 09:37:10 +00:00
|
|
|
go p.destroyInstance(pi.instance)
|
|
|
|
continue
|
|
|
|
}
|
|
|
|
break
|
|
|
|
}
|
|
|
|
|
|
|
|
return pi, nil
|
|
|
|
}
|
|
|
|
|
|
|
|
func (p *pool) Get() (interface{}, error) {
|
|
|
|
var err error
|
|
|
|
var pi *pooledInstance
|
|
|
|
if pi, err = p.get(); nil != err {
|
|
|
|
return nil, err
|
|
|
|
}
|
|
|
|
|
|
|
|
pi.inUse = true
|
|
|
|
|
|
|
|
return pi.instance, nil
|
|
|
|
}
|
|
|
|
|
|
|
|
func (p *pool) Put(i interface{}) {
|
|
|
|
pi, ok := p.instances[i]
|
|
|
|
if !ok {
|
|
|
|
return
|
|
|
|
}
|
|
|
|
|
|
|
|
pi.idleTime = time.Now()
|
|
|
|
pi.inUse = false
|
|
|
|
|
|
|
|
select {
|
|
|
|
case p.instancesCh <- pi:
|
|
|
|
// All good
|
|
|
|
default:
|
|
|
|
// channel is full
|
|
|
|
return
|
|
|
|
}
|
|
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
func (p *pool) Capacity() int {
|
|
|
|
return cap(p.instancesCh)
|
|
|
|
}
|
|
|
|
|
|
|
|
func (p *pool) Available() int {
|
|
|
|
return len(p.instancesCh)
|
|
|
|
}
|
|
|
|
|
|
|
|
func (p *pool) Destroy() {
|
|
|
|
close(p.instancesCh)
|
|
|
|
|
|
|
|
for k, _ := range p.instances {
|
|
|
|
p.destroyInstance(k)
|
|
|
|
}
|
|
|
|
|
|
|
|
}
|