Go

Go の Prepared Statement は Connection を気にせず使える

More than 3 years have passed since last update.

コネクションプールと prepared statement

database/sql は、トランザクションを除いて基本的にコネクションをユーザーに見せず、全ての操作をコネクションプール sql.DB を通して行う設計になっています。

コネクションプールと言えば、 prepared statement の実装が気になります。
prepared statement は基本的にコネクションに紐付いていて、

  1. プレースホルダ付きのクエリを投げてコネクションローカルの「ハンドル」を貰う
  2. 「ハンドル」にパラメータを送ってクエリを実行する
  3. 「ハンドル」を閉じる(あるいはコネクション自体を閉じる)

という流れになるので、コネクションプールと組み合わせて使う場合には、

  1. 毎回ハンドルを取得&開放する (1クエリに3回通信が発生)
  2. ハンドルを開放しない (DBサーバー側のリソースを食いつぶす)
  3. コネクションごとに、クエリに対する「ハンドル」を管理して、まだ prepare してないなら prepare する

といった方法が考えられます。
prepared statement のメリットの1つは、クエリのパースを毎回しないで済むからパフォーマンスが上がるというものなのに (1) は効率が悪くて残念すぎます。
(2) は問題外です。そんなDBドライバ窓から投げ捨てましょう。
(3) が一番賢くて、 Go の prepared statement はこれになっていました。

Prepared Statement の使い方

func (db *DB) Prepare(query string) (*Stmt, error) を呼ぶと、 prepare された Stmt オブジェクトが返ります。

この Stmt が、コネクションと prepared statement の対応を管理する役目も持っているので、アプリケーション内でクエリごとに singleton になるようにしてください。
簡単なアプリケーションの場合、アプリケーションの起動直後に DB に接続して Prepare を済ませておき、 *DB*Stmt をグローバル変数にしてしまうのが楽でしょう。

package main

import (
    "database/sql"
    "flag"

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

var db *sql.DB
var stmt *sql.Stmt

func run() {
    // ...
}

func main() {
    var dsn string
    flag.StringVar(&dsn, "dsn", "", "DSN (see https://github.com/go-sql-driver/mysql#dsn-data-source-name)")
    flag.Parse()

    var err error
    db, err = sql.Open("mysql", dsn)
    if err != nil {
        log.Fatal(err)
    }
    stmt, err = db.Prepare("SELECT 1+1")
    if err != nil {
        log.Fatal(err)
    }
    run()
}

実装を覗いてみる

DB.Prepare()

DB から1つだけ接続 (dc) を取り出して、 dc.prepareLocked(query) を実行します。
prepareLocked()driver.Stmt を返すので、 Stmt の中に、 dc とペア (connStmt) を、 Stmt.css という名前で格納します。

この段階では1接続だけで prepare されている状態です。

// connStmt is a prepared statement on a particular connection.
type connStmt struct {
    dc *driverConn
    si driver.Stmt
}

// Prepare creates a prepared statement for later queries or executions.
// Multiple queries or executions may be run concurrently from the
// returned statement.
func (db *DB) Prepare(query string) (*Stmt, error) {
    var stmt *Stmt
    var err error
    for i := 0; i < maxBadConnRetries; i++ {
        stmt, err = db.prepare(query)
        if err != driver.ErrBadConn {
            break
        }
    }
    return stmt, err
}

func (db *DB) prepare(query string) (*Stmt, error) {
    // TODO: check if db.driver supports an optional
    // driver.Preparer interface and call that instead, if so,
    // otherwise we make a prepared statement that's bound
    // to a connection, and to execute this prepared statement
    // we either need to use this connection (if it's free), else
    // get a new connection + re-prepare + execute on that one.
    dc, err := db.conn()
    if err != nil {
        return nil, err
    }
    dc.Lock()
    si, err := dc.prepareLocked(query)
    dc.Unlock()
    if err != nil {
        db.putConn(dc, err)
        return nil, err
    }
    stmt := &Stmt{
        db:    db,
        query: query,
        css:   []connStmt{{dc, si}},
    }
    db.addDep(stmt, stmt)
    db.putConn(dc, nil)
    return stmt, nil
}

Stmt.Exec()

Stmt.Exec()Stmt.connStmt() から、 prepare 済みのコネクションを受け取って実行します。
Stmt.connStmt() は、まず Stmt.css の中に入ってる prepare 済みのコネクションを舐めて、もしそのコネクションが開いていたらそれを返します。そのコネクションが既に閉じられていたら Stmt.css から削除します。
Stmt.css の中にあるコネクションのなかで空きがなかった場合は、新しい接続に対して prepare() し、 Stmt.css に追加しつつ返します。

// Exec executes a prepared statement with the given arguments and
// returns a Result summarizing the effect of the statement.
func (s *Stmt) Exec(args ...interface{}) (Result, error) {
    s.closemu.RLock()
    defer s.closemu.RUnlock()

    var res Result
    for i := 0; i < maxBadConnRetries; i++ {
        dc, releaseConn, si, err := s.connStmt()
        if err != nil {
            if err == driver.ErrBadConn {
                continue
            }
            return nil, err
        }

        res, err = resultFromStatement(driverStmt{dc, si}, args...)
        releaseConn(err)
        if err != driver.ErrBadConn {
            return res, err
        }
    }
    return nil, driver.ErrBadConn
}

func resultFromStatement(ds driverStmt, args ...interface{}) (Result, error) {
    ds.Lock()
    want := ds.si.NumInput()
    ds.Unlock()

    // -1 means the driver doesn't know how to count the number of
    // placeholders, so we won't sanity check input here and instead let the
    // driver deal with errors.
    if want != -1 && len(args) != want {
        return nil, fmt.Errorf("sql: expected %d arguments, got %d", want, len(args))
    }

    dargs, err := driverArgs(&ds, args)
    if err != nil {
        return nil, err
    }

    ds.Lock()
    resi, err := ds.si.Exec(dargs)
    ds.Unlock()
    if err != nil {
        return nil, err
    }
    return driverResult{ds.Locker, resi}, nil
}

// connStmt returns a free driver connection on which to execute the
// statement, a function to call to release the connection, and a
// statement bound to that connection.
func (s *Stmt) connStmt() (ci *driverConn, releaseConn func(error), si driver.Stmt, err error) {
    if err = s.stickyErr; err != nil {
        return
    }
    s.mu.Lock()
    if s.closed {
        s.mu.Unlock()
        err = errors.New("sql: statement is closed")
        return
    }

    // In a transaction, we always use the connection that the
    // transaction was created on.
    if s.tx != nil {
        s.mu.Unlock()
        ci, err = s.tx.grabConn() // blocks, waiting for the connection.
        if err != nil {
            return
        }
        releaseConn = func(error) {}
        return ci, releaseConn, s.txsi.si, nil
    }

    var cs connStmt
    match := false
    for i := 0; i < len(s.css); i++ {
        v := s.css[i]
        _, err := s.db.connIfFree(v.dc)
        if err == nil {
            match = true
            cs = v
            break
        }
        if err == errConnClosed {
            // Lazily remove dead conn from our freelist.
            s.css[i] = s.css[len(s.css)-1]
            s.css = s.css[:len(s.css)-1]
            i--
        }

    }
    s.mu.Unlock()

    // Make a new conn if all are busy.
    // TODO(bradfitz): or wait for one? make configurable later?
    if !match {
        dc, err := s.db.conn()
        if err != nil {
            return nil, nil, nil, err
        }
        dc.Lock()
        si, err := dc.prepareLocked(s.query)
        dc.Unlock()
        if err != nil {
            s.db.putConn(dc, err)
            return nil, nil, nil, err
        }
        s.mu.Lock()
        cs = connStmt{dc, si}
        s.css = append(s.css, cs)
        s.mu.Unlock()
    }

    conn := cs.dc
    return conn, conn.releaseConn, cs.si, nil
}

追記

    s.mu.Unlock()

    // Make a new conn if all are busy.
    // TODO(bradfitz): or wait for one? make configurable later?
    if !match {
        dc, err := s.db.conn()

の部分で、他の goroutine がちょうど返したばかりの prepared 済みのコネクションに対して重ねて prepare しないのかなーというのがコードを読んでいた時に気になっていたのですが、報告済みのバグでした。

https://code.google.com/p/go/issues/detail?id=8376&q=Prepare&colspec=ID%20Status%20Stars%20Release%20Owner%20Repo%20Summary