connection_pool.go

package pool

import (
	"container/list"
	"errors"
	"sync"

	"fmt"
)

// ------------- basic connection pool implementation -----------------

type Connection interface {
	// Close the connection.
	Close() error
	// IsBad return true/false when a connection is bad/good.
	IsBad() bool
}

// connectFunc: factory function for new connection.
type connectFunc func() (Connection, error)

// Pool: basic connection pool.
type Pool struct {
	size         int64
	mutex        sync.RWMutex
	pool         *list.List
	currentCount int64
	connecting   int64
	connect      connectFunc
}

// New: create a new connection pool.
func New(size int64, hook connectFunc) *Pool {
	return &Pool{
		size:    size,
		pool:    list.New(),
		connect: hook,
	}
}

// Checkout: checkout connection from pool or create one.
func (p *Pool) Checkout() (Connection, error) {
	if conn := p.peek(); conn != nil {
		fmt.Println("connection obtained from pool")
		return conn, nil
	}
	// create a new connection
	return p.createConnection()
}

func (p *Pool) createConnection() (Connection, error) {
	fmt.Println("creating a new connection")
	newConnection := false
	p.mutex.Lock()

	if p.currentCount+p.connecting < p.size {
		p.connecting++
		newConnection = true
	}
	p.mutex.Unlock()

	if newConnection {
		conn, err := p.connect()
		p.mutex.Lock()
		p.connecting--
		if conn != nil {
			p.currentCount++
			fmt.Println("new connection setup")
		}
		p.mutex.Unlock()

		return conn, err
	}
	fmt.Println("pool size exceed")
	return nil, errors.New("pool size exceed")
}

// Checkin: checkin the connection onto pool
func (p *Pool) Checkin(conn Connection) {
	fmt.Println("checkin connection")
	if conn.IsBad() {
		p.mutex.Lock()
		defer p.mutex.Unlock()
		p.currentCount--
		return
	}
	p.mutex.Lock()
	defer p.mutex.Unlock()
	p.pool.PushBack(conn)
}

// peek: return the removed connection at the front.
func (p *Pool) peek() Connection {
	p.mutex.Lock()
	if p.pool.Len() == 0 {
		p.mutex.Unlock()
		return nil
	}
	// if we are here we know we have connection
	conn := p.pool.Front()
	p.pool.Remove(conn)
	p.mutex.Unlock()
	c, _ := conn.Value.(Connection)
	return c
}

// TearDown: teardown the connection and close them.
func (p *Pool) TearDown() {
	p.mutex.Lock()
	defer p.mutex.Unlock()

	// tear down all connections
	for p.pool.Len() != 0 {
		c := p.pool.Front()
		conn, _ := c.Value.(Connection)
		conn.Close()
		p.pool.Remove(c)
	}
}