conn.go

// Package conn provides a safe MySQL connection pool with human-friendly error
// and metrics.
package conn

import (
	"context"
	"database/sql"
	"database/sql/driver"
	"errors"
	"net"
	"sync"
	"time"

	"github.com/go-sql-driver/mysql"
)

var (
	// ErrConnCannotConnect is returned when a connection cannot be made because
	// MySQL is down or unreachable for any reason. This is usually because
	// MySQL is down, but it could also indicate a network issue. If this error
	// occurs frequently, verify the network connection and MySQL address. If it
	// occurs infrequently, it could indicate a transient state that will
	// recover automatically; for example: a failover. Pool.Error returns true
	// with this error.
	ErrConnCannotConnect = errors.New("cannot connect")

	// ErrConnLost is returned when the connection is lost, closed, or killed
	// for any reason. This could mean MySQL crashed, or the KILL command was used
	// to kill the connection. Lost is distinct from down: a connection can only
	// be lost if it was previously connected. If MySQL is up and ok, this could
	// indicate a transient state that will recover automatically. If not,
	// ErrConnCannotConnect will probably be returned next when the driver tries
	// but fails to reestablish the connection. Pool.Error returns true with this error.
	ErrConnLost = errors.New("connection lost")

	// ErrQueryKilled is returned when the KILL QUERY command is used. This only
	// kills the currently active query; the connection is still ok. This error
	// is not connection-related, so Pool.Error returns false.
	ErrQueryKilled = errors.New("query killed")

	// ErrReadOnly is returned when MySQL read-only is enabled. This error is not
	// connection-related, so Pool.Error returns false.
	ErrReadOnly = errors.New("server is read-only")

	// ErrDupeKey is returned when a unique index prevents a value from being
	// inserted or updated. This error is not connection-related, so Pool.Error
	// returns false.
	ErrDupeKey = errors.New("duplicate key value")

	// ErrTimeout is returned when a context timeout happens. This error is not
	// connection-related, so Pool.Error returns false.
	ErrTimeout = errors.New("timeout")
)

type Connector interface {
	Open(context.Context) (*sql.Conn, error)
	Close(*sql.Conn)
	Error(error) (bool, error)
	Stats() Stats
}

type Stats struct {
	*sync.Mutex
	Ts        int64
	Open      int  // sql.DBStats.OpenConnections, not always accurate (gauge)
	OpenCalls int  // All calls to Open (counter)
	Opened    uint // Successful calls to Open (counter)
	Closed    uint // All calls to close (counter)
	Timeout   uint // Timeouts during calls to Open (counter)
	Lost      uint // Lost connections (counter)
	Down      uint // ErrConnCannotConnect errors (counter)
}

type Pool struct {
	db    *sql.DB
	stats *Stats
}

// NewPool creates a new pool using the given, pre-configured sql.DB.
func NewPool(db *sql.DB) *Pool {
	return &Pool{
		db: db,
		stats: &Stats{
			Mutex: &sync.Mutex{},
		},
	}
}

// Open opens a sql.Conn from the pool. The caller should always call this
// function to open a connection because it saves stats.
func (p *Pool) Open(ctx context.Context) (*sql.Conn, error) {
	p.stats.Lock()
	p.stats.OpenCalls++
	p.stats.Unlock()

	conn, err := p.db.Conn(ctx)
	if err != nil {
		return conn, err
	}

	p.stats.Lock()
	p.stats.Opened++
	p.stats.Unlock()

	return conn, err
}

// Close closes a sql.Conn obtained from the pool by calling Open. The caller
// should always call this function to close the sql.Conn because it saves stats.
func (p *Pool) Close(conn *sql.Conn) {
	conn.Close()

	p.stats.Lock()
	p.stats.Closed++
	p.stats.Unlock()
}

// Error returns true if the given error is a connection error (ErrConn*), else
// it returns false. If true, the caller should wait and retry the operation;
// the driver will attempt to reconnect. However, it usually take one or two
// attempts after the first error before the driver reestablishes the connection.
// This is a driver implementation detail that cannot be changed by the caller.
// If false, the returned error could be one provided by this package, like
// ErrReadOnly. The caller should check and handle accordingly.
//
// If the given error is not a MySQL error or nil, it returns false and the given
// error.
//
// The caller should always call this function on error from any code that uses
// a sql.Conn obtained from the pool because it saves stats based on the error.
// See packages docs for the correct, typical workflow.
func (p *Pool) Error(err error) (bool, error) {
	if err == nil {
		return false, nil
	}

	p.stats.Lock()
	defer p.stats.Unlock()

	if Lost(err) {
		p.stats.Lost++
		return true, ErrConnLost
	}

	if Down(err) {
		p.stats.Down++
		return true, ErrConnCannotConnect
	}

	// Not connection issues (return false)
	if errCode := MySQLErrorCode(err); errCode != 0 {
		switch errCode {
		case 1317: // ER_QUERY_INTERRUPTED
			return false, ErrQueryKilled
		case 1290, 1836: // ER_OPTION_PREVENTS_STATEMENT, ER_READ_ONLY_MODE
			return false, ErrReadOnly
		case 1062: // ER_DUP_ENTRY
			return false, ErrDupeKey
		}
	}

	if err == context.DeadlineExceeded {
		p.stats.Timeout++
		return false, ErrTimeout
	}

	// Not a MySQL error we handle, or some other type of error
	return false, err
}

// Stats returns the stats since the last call to Stats. Stats are never
// reset. The Open stat is not always accurate. The driver updates it lazily,
// so it can be higher than the actual number of currently open connections.
func (p *Pool) Stats() Stats {
	// Set ts and copy stats
	p.stats.Lock()
	defer p.stats.Unlock()
	dbstats := p.db.Stats()
	return Stats{
		Ts:        time.Now().Unix(),
		Open:      dbstats.OpenConnections,
		OpenCalls: p.stats.OpenCalls,
		Opened:    p.stats.Opened,
		Closed:    p.stats.Closed,
		Lost:      p.stats.Lost,
		Timeout:   p.stats.Timeout,
		Down:      p.stats.Down,
	}
}

// //////////////////////////////////////////////////////////////////////////
// Helper functions
// //////////////////////////////////////////////////////////////////////////

// Lost returns true if the error indicates the MySQL connection was lost for
// any reason. Lost is distinct from down: a connection can only be lost if it
// was previously connected. A connection can be lost for many reasons. MySQL
// could be up and ok, but particular connections were lost. For example, the
// KILL command causes a lost connection.
func Lost(err error) bool {
	// mysql.ErrInvalidConn is returned for sql.DB functions. driver.ErrBadConn
	// is returned for sql.Conn functions. These are the normal errors when
	// MySQL is lost.
	if err == mysql.ErrInvalidConn || err == driver.ErrBadConn {
		return true
	}

	// Server shutdown in progress is a special case: the conn will be lost
	// soon. The next call will most likely return in the block above ^.
	if errCode := MySQLErrorCode(err); errCode == 1053 { // ER_SERVER_SHUTDOWN
		return true
	}

	return false
}

// Down returns true if the error indicates MySQL cannot be reached for any
// reason, probably because it's not running. Down is distinct from lost: it's
// a network-level error that means MySQL cannot be reached. MySQL could be up
// and ok, but particular connections are failing, so from the program's point
// of view MySQL is down.
func Down(err error) bool {
	// Being unable to reach MySQL is a network issue, so we get a net.OpError.
	// If MySQL is reachable, then we'd get a mysql.* or driver.* error instead.
	_, ok := err.(*net.OpError)
	return ok
}

func MySQLErrorCode(err error) uint16 {
	if val, ok := err.(*mysql.MySQLError); ok {
		return val.Number
	}
	return 0 // not a mysql error
}