You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
219 lines
5.3 KiB
Go
219 lines
5.3 KiB
Go
// Source: https://github.com/eternnoir/gncp/blob/master/pool.go
|
|
|
|
package gncp
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"net"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
type ConnPool interface {
|
|
Get() (net.Conn, error)
|
|
GetWithTimeout(timeout time.Duration) (net.Conn, error)
|
|
Close() error
|
|
Remove(conn net.Conn) error
|
|
}
|
|
|
|
// GncpPool implements ConnPool interface. Use channel buffer connections.
|
|
type GncpPool struct {
|
|
lock sync.Mutex
|
|
conns chan net.Conn
|
|
minConnNum int
|
|
maxConnNum int
|
|
totalConnNum int
|
|
closed bool
|
|
connCreator func() (net.Conn, error)
|
|
}
|
|
|
|
var (
|
|
errPoolIsClose = errors.New("Connection pool has been closed")
|
|
// Error for get connection time out.
|
|
errTimeOut = errors.New("Get Connection timeout")
|
|
errContextClose = errors.New("Get Connection close by context")
|
|
)
|
|
|
|
// NewPool return new ConnPool. It base on channel. It will init minConn connections in channel first.
|
|
// When Get()/GetWithTimeout called, if channel still has connection it will get connection from channel.
|
|
// Otherwise GncpPool check number of connection which had already created as the number are less than maxConn,
|
|
// it use connCreator function to create new connection.
|
|
func NewPool(minConn, maxConn int, connCreator func() (net.Conn, error)) (*GncpPool, error) {
|
|
if minConn > maxConn || minConn < 0 || maxConn <= 0 {
|
|
return nil, errors.New("Number of connection bound error")
|
|
}
|
|
|
|
pool := &GncpPool{}
|
|
pool.minConnNum = minConn
|
|
pool.maxConnNum = maxConn
|
|
pool.connCreator = connCreator
|
|
pool.conns = make(chan net.Conn, maxConn)
|
|
pool.closed = false
|
|
pool.totalConnNum = 0
|
|
err := pool.init()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return pool, nil
|
|
}
|
|
|
|
func (p *GncpPool) init() error {
|
|
for i := 0; i < p.minConnNum; i++ {
|
|
conn, err := p.createConn()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
p.conns <- conn
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// GetMaxConns returns the maximum number of connections configured
|
|
func (p *GncpPool) GetMaxConns() int {
|
|
return p.maxConnNum
|
|
}
|
|
|
|
// Get get connection from connection pool. If connection poll is empty and alreay created connection number less than Max number of connection
|
|
// it will create new one. Otherwise it wil wait someone put connection back.
|
|
func (p *GncpPool) Get() (net.Conn, error) {
|
|
if p.isClosed() == true {
|
|
return nil, errPoolIsClose
|
|
}
|
|
go func() {
|
|
conn, err := p.createConn()
|
|
if err != nil {
|
|
return
|
|
}
|
|
p.conns <- conn
|
|
}()
|
|
select {
|
|
case conn := <-p.conns:
|
|
return p.packConn(conn), nil
|
|
}
|
|
}
|
|
|
|
// GetWithTimeout can let you get connection wait for a time duration. If cannot get connection in this time.
|
|
// It will return TimeOutError.
|
|
func (p *GncpPool) GetWithTimeout(timeout time.Duration) (net.Conn, error) {
|
|
if p.isClosed() == true {
|
|
return nil, errPoolIsClose
|
|
}
|
|
go func() {
|
|
conn, err := p.createConn()
|
|
if err != nil {
|
|
return
|
|
}
|
|
p.conns <- conn
|
|
}()
|
|
select {
|
|
case conn := <-p.conns:
|
|
return p.packConn(conn), nil
|
|
case <-time.After(timeout):
|
|
return nil, errTimeOut
|
|
}
|
|
}
|
|
|
|
func (p *GncpPool) GetWithContext(ctx context.Context) (net.Conn, error) {
|
|
if p.isClosed() == true {
|
|
return nil, errPoolIsClose
|
|
}
|
|
go func() {
|
|
conn, err := p.createConn()
|
|
if err != nil {
|
|
return
|
|
}
|
|
p.conns <- conn
|
|
}()
|
|
select {
|
|
case conn := <-p.conns:
|
|
return p.packConn(conn), nil
|
|
case <-ctx.Done():
|
|
return nil, errContextClose
|
|
}
|
|
}
|
|
|
|
// Close close the connection pool. When close the connection pool it also close all connection already in connection pool.
|
|
// If connection not put back in connection it will not close. But it will close when it put back.
|
|
func (p *GncpPool) Close() error {
|
|
if p.isClosed() == true {
|
|
return errPoolIsClose
|
|
}
|
|
|
|
p.lock.Lock()
|
|
p.closed = true
|
|
p.lock.Unlock()
|
|
|
|
close(p.conns)
|
|
for conn := range p.conns {
|
|
conn.Close()
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// Put can put connection back in connection pool. If connection has been closed, the connection will be close too.
|
|
func (p *GncpPool) Put(conn net.Conn) error {
|
|
if p.isClosed() == true {
|
|
return errPoolIsClose
|
|
}
|
|
if conn == nil {
|
|
p.lock.Lock()
|
|
p.totalConnNum = p.totalConnNum - 1
|
|
p.lock.Unlock()
|
|
return errors.New("Cannot put nil to connection pool")
|
|
}
|
|
|
|
select {
|
|
case p.conns <- conn:
|
|
return nil
|
|
default:
|
|
return conn.Close()
|
|
}
|
|
}
|
|
|
|
func (p *GncpPool) isClosed() bool {
|
|
p.lock.Lock()
|
|
defer p.lock.Unlock()
|
|
return p.closed
|
|
}
|
|
|
|
// Remove let connection not belong connection pool. And it will close connection.
|
|
func (p *GncpPool) Remove(conn net.Conn) error {
|
|
if p.isClosed() == true {
|
|
return errPoolIsClose
|
|
}
|
|
|
|
switch conn.(type) {
|
|
case *CpConn:
|
|
// Destroy calls pool.Remove, so do not decrease the number of connections here
|
|
return conn.(*CpConn).Destroy()
|
|
default:
|
|
p.lock.Lock()
|
|
p.totalConnNum = p.totalConnNum - 1
|
|
p.lock.Unlock()
|
|
return conn.Close()
|
|
}
|
|
}
|
|
|
|
// createConn will create one connection from connCreator. And increase connection counter.
|
|
func (p *GncpPool) createConn() (net.Conn, error) {
|
|
p.lock.Lock()
|
|
defer p.lock.Unlock()
|
|
if p.totalConnNum >= p.maxConnNum {
|
|
return nil, fmt.Errorf("Connot Create a new connection. Pool now has %d conns. Max is %d", p.totalConnNum, p.maxConnNum)
|
|
}
|
|
conn, err := p.connCreator()
|
|
if err != nil {
|
|
return nil, fmt.Errorf("Cannot create a new connection: %s", err)
|
|
}
|
|
p.totalConnNum = p.totalConnNum + 1
|
|
return conn, nil
|
|
}
|
|
|
|
func (p *GncpPool) packConn(conn net.Conn) net.Conn {
|
|
ret := &CpConn{pool: p}
|
|
ret.Conn = conn
|
|
return ret
|
|
}
|