// Package db provides both a generic interface to the go db layer, and a
// database connection pool abstraction.  The generic interface can be used
// without the pooling layer if pooling is not desired.  The things to check
// out are:
//
// * Pool and PoolImpl - Pool defines an interface to db pools, and PoolImpl is
// the default implementation.
//
// * PoolConfig - Configure a Pool how you have always wanted.
//
// * ResultSet and ResultMap - ResultSet is a slice of ResultMaps, each of
// which holds a row of results from the db.
package db

import (
    "fmt"
    "time"
    "database/sql"
    "sync"

    "code.justin.tv/common/golibs/log"
)

const (
    // Time we wait for a "stalled" db query before discarding the associated
    // connection from the pool.  Not really sure what a good value would be.
    // Too short, and we risk putting more load on a db that is struggling or
    // deadlocked; too slow and we shrink the connection pool, possibly without
    // a good reason.  For now, let's make sure we cause some back pressure if
    // there is high demand, but try again in the future if the situation has
    // resolved.  10 seconds seems to be a reasonable time, but it's pulled out
    // of my ass: above four TCP retries/backoffs, and long enough to force
    // lots of failures if the whole pool gets stuck, but short enough to throw
    // everything out and retry if the db gets restarted or clears, and to do
    // so without needed manual intervention.
    STALL_WAIT_TIME = 10 * time.Second
)

// Pool provides an interface to a database pool.  The standard usage of a
// pool follows the following pattern:
//
//    if conn, err := pool.Claim(); err == nil {
//        // do things with connection
//        pool.Release(conn)
//    }
//
// or, even more simply:
//
//    if res, err := pool.Query("select ... ", x, y, z); err == nil {
//        // use results in res
//    }
//
// In brief, Claim() returns a connection from the pool; connections must be
// returned to the pool via Release(), else they will be lost.  Connections can
// be evicted from the pool using Discard(), which will close them if possible
// and then remove them.  The most useful function is Query(), which will
// Claim(), run a query, Release(), and return the results.  Query() will wait
// up to PoolConfig.PoolWaitTimeout for a free connection if all connections
// are currently in use, and up to PoolConfig.QueryTimeout for the query to
// complete.
type Pool interface {
    // Claim a database connection from the pool.
    Claim() (*sql.DB, error)

    // Release a connection back to the pool.
    Release(conn *sql.DB)

    // Discard a connection from the pool.
    Discard(conn *sql.DB)

    // Query the database using a free connection from the pool.
    Query(query string, params ...interface{}) (ResultSet, error)
}

// PoolImpl implements Pool.
type PoolImpl struct {
    userConf  PoolConfig

    sizeMutex sync.Mutex
    curSize   uint32
    pool      chan *sql.DB
    waitQueue chan chan *sql.DB
}

type NoFreeDBConnectionError struct { }

func (e *NoFreeDBConnectionError) Error() string {
    return fmt.Sprintf("No free DB connections available.")
}

// PoolConfig provides the configuration parameters for a specific pool.
// Passed to CreatePool.
type PoolConfig struct {
    ConnectString   string        // Specifies the db connection parameters
    ConnectTimeout  time.Duration // Connection timeout
    PoolSize        uint32        // Maximum size of pool
    PoolWaitTimeout time.Duration // Maximum time to wait for a free connection during Claim()
    QueryTimeout    time.Duration // Maximum time to wait for a Query() to return

    StatsLogger     func(event string, t time.Duration, err error)  // Logging callback for an external service to capture db event timings; if null, do no callbacks
}

// CreatePool creates a new db pool with the settings given in conf.  May return
// any error db.NewConnection returns, as it attempts to initialize at least
// one connection to the database.  The underlying type of the returned Pool is
// PoolImpl.
func CreatePool(conf PoolConfig) (Pool, error) {
    log.Logln("Creating DB pool.")

    pool, err := newPool(conf)
    if err != nil {
        pool = nil
        return nil, err
    }

    // force NewConnection to be called; also create the first connection to
    // the DB
    conn, err := pool.Claim()
    if err == nil {
        pool.Release(conn)
        log.Logln("DB pool created.")
        return pool, nil
    } else if conn != nil {
        // XXX: thought to be unused code-path
        // if we had an error but got a connection back, Discard(), as we don't
        // know the state; this /should/ be dead code with the current flow of
        // the db driver, but the current flow is not contractual.
        pool.Discard(conn)
    }

    return nil, err
}

func newPool(conf PoolConfig) (Pool, error) {
    pool := new(PoolImpl)
    pool.userConf = conf

    pool.curSize = 0
    pool.pool = make(chan *sql.DB, pool.userConf.PoolSize)
    pool.waitQueue = make(chan chan *sql.DB, pool.userConf.PoolSize * 10)

    return pool, nil
}

// Claim attempts to grab a free database connection, or create a new
// connection if all current connection are busy and there is room left in the
// pool.  If it can neither grab a free connection or has space to create a new
// one, it will wait for PoolConfig.PoolWaitTimeout for a active connection to
// be Released().  Can return the same errors as db.QueryWithConnection(), or:
//
//    * NoFreeDBConnectionError: Could not claim a connection within PoolConfig.PoolWaitTimeout.
func (p *PoolImpl) Claim() (*sql.DB, error) {
    var conn *sql.DB
    var err error

    startTime := time.Now()

    defer func() {
        p.logClaim(startTime, err)
    }()

    select {
    case conn = <-p.pool:
        log.Debugln("Got a connection from the pool.")
    default:
        conn, err = p.tryCreatePooledConnection()
        if err != nil {
            return nil, err
        }
    }

    if conn != nil {
        return conn, err
    }

    // else we could not create a new pooled connection (pool's full); wait for
    // one to be returned
    conn, err = p.waitForRelease()

    return conn, err
}

// Release returns a database connection back to the pool.
func (p *PoolImpl) Release(conn *sql.DB) {
    select {
    case wait_chn := <-p.waitQueue:
        // pass connection directly to waiting process
        wait_chn <- conn
    default:
        // put connection back into the main pool
        p.pool <- conn
    }
}

// Discard closes and removes a database connection from the pool.
func (p *PoolImpl) Discard(conn *sql.DB) {
    if conn == nil {
        return
    }

    tryClose(conn)

    p.sizeMutex.Lock()
    p.curSize -= 1
    p.sizeMutex.Unlock()
}

type queryResult struct {
    res ResultSet
    err error
}

func asyncQuery(ret_c chan queryResult, db *sql.DB, query string, params []interface{}) {
    res, err := QueryWithConnection(db, query, params...)
    ret_c <- queryResult{res, err}
}

// Query executes the given query using db.QueryWithConnection and a database
// connection returned by Claim().  It can return the same errors as
// db.QueryWithConnection().  An example usage would be:
//    res, err := pool.Query("SELECT * FROM table WHERE name = $1 AND other_thing = $2", "a_name", int_val)
//    for obj := range res {
//        fmt.Printf("%s: %d", obj["name"].(string), obj["some_other_property"].(int64))
//    }
//
// It's the callers responsibility to know the underlying type of the objects
// returned.  The typed used is the one that maps to the db type in the
// underlying db driver.  This situation is equivalent to having to know the
// type at the time one call Rows.Scan() in the go db interface.
func (p *PoolImpl) Query(query string, params ...interface{}) (ResultSet, error) {
    startTime := time.Now()

    var err error
    var db *sql.DB

    defer func() {
        p.logQuery(startTime, err)
    }()

    db, err = p.Claim()
    if err != nil {
        return nil, QueryFailedError{Query: query, UnderlyingError: err}
    }

    timer := time.NewTimer(p.userConf.QueryTimeout)
    defer timer.Stop()

    query_c := make(chan queryResult)
    go asyncQuery(query_c, db, query, params)

    var res ResultSet

    select {
    case r := <-query_c:
        p.Release(db)
        res, err = r.res, r.err

    case <-timer.C:
        res = nil
        err = QueryTimedOutError{QueryString: query, Timeout: p.userConf.QueryTimeout}
        go p.cleanupStalledQuery(db, query_c)
    }

    return res, err
}

// Maybe we just had an unusual network hiccup - wait a retry timeout and then
// reap the connection if we still haven't heard back.
func (p *PoolImpl) cleanupStalledQuery(conn *sql.DB, recv_c chan queryResult) {
    // wait a TCP timeout window of time for the query, then give up, and throw
    // the connection away.
    timer := time.NewTimer(STALL_WAIT_TIME)
    defer timer.Stop()

    select {
    case <-recv_c:
        log.Warnln("Late DB query returned; releasing.")
        p.Release(conn)

    case <-timer.C:
        log.Errln("DB never returned; discarding connection.")
        p.Discard(conn)
    }
}

func (p *PoolImpl) tryCreatePooledConnection() (*sql.DB, error) {
    var conn *sql.DB
    var err error

    // Note that this is /not/ an locked check - we may continue past this
    // point and then, after the counter lock is grabbed at point 1 below, fail
    // the next conditional.  It's performed here in an attempt to
    // optimistically avoid lock contention and waiting when the pool is empty
    // and multiple go procs need to wait for connections to be released in
    // waitForRelease().
    if p.curSize >= p.userConf.PoolSize {
        return nil, nil
    }

    // NOTE: It could be argued that we should defer p.sizeMutex.Unlock()
    // immediately after Lock(); however, NewConnection() can, under some error
    // conditions, block while connecting, while  other, simultaneous attempts
    // will succeed.  As such, we do not want to wait until NewConnection() has
    // returned to unlock the mutex.  I believe there is no condition under
    // which the code will not reach one of the two Unlock() calls and that
    // this should be a safe optimization.  (Also, it's invalid to unlock a
    // mutex twice.) The Lock/Unlock calls are label with numbers as to their
    // pairings.
    p.sizeMutex.Lock() // Lock 1
    if p.curSize < p.userConf.PoolSize {
        p.curSize += 1
        p.sizeMutex.Unlock()  // Unlock 1

        // if NewConnection panics, we want to make sure we keep the pool count
        // accurate; something above us may catch the panic and continue
        // running, and we want to leave the db pool count in a good state.
        defer func() {
            if conn == nil || err != nil {
                p.sizeMutex.Lock()    // Lock 2
                p.curSize -= 1
                p.sizeMutex.Unlock()  // Unlock 2
            }
        }()

        log.Logln("Creating new DB connection (pool %d/%d).", p.curSize, p.userConf.PoolSize)
        conn, err = NewConnection(p.userConf.ConnectString, p.userConf.ConnectTimeout)
        if conn == nil || err != nil {
            // note, the defered func above will trigger here and decrement the
            // pool count size
            return nil, err
        }
    } else {
        p.sizeMutex.Unlock()  // Unlock 1
    }

    return conn, nil
}

func (p *PoolImpl) waitForRelease() (*sql.DB, error) {
    var conn *sql.DB

    wait_chnl := make(chan *sql.DB)
    p.waitQueue <- wait_chnl

    log.Debugln("Waiting for a DB connection to be released.")
    timer := time.NewTimer(p.userConf.PoolWaitTimeout)
    defer timer.Stop()

    select {
    case conn = <-wait_chnl:
        log.Debugln("Got a DB connection from release.")
        return conn, nil
    case <-timer.C:
        // we still need to read the connection off the socket once it is ready
        // and then try and release again, else we'll block on the send in
        // Release, causing the gofunc calling Release to hang and the db
        // connection to become lost
        go func() {
            p.Release(<-wait_chnl)
        }()

        // return error below
    }

    log.Warnln("Timeout; no free DB connections within %d milliseconds.", p.userConf.PoolWaitTimeout / (1 * time.Millisecond))
    return nil, &NoFreeDBConnectionError{}
}

func (p *PoolImpl) logClaim(startTime time.Time, err error) {
    p.logStat("claim", startTime, err)
}

func (p *PoolImpl) logQuery(startTime time.Time, err error) {
    p.logStat("query", startTime, err)
}

func (p *PoolImpl) logStat(event string, startTime time.Time, err error) {
    if p.userConf.StatsLogger != nil {
        duration := time.Now().Sub(startTime)
        p.userConf.StatsLogger(event, duration, err)
    }
}
