`*sql.DB` を観察する #golang

Photo by Abo Ngalonkulu on Unsplash

People tribe / Backend squad の @izumin5210 です。もう12月ですね。自分は Advent Calendar に登録しすぎて後悔するのが得意です。

この記事は Go3 Advent Calendar 2019 の4日目です。

TL;DR

• SQL のメトリクス・トレースを収集したいは driver.Driver をラップするのが常套手段
• コネクション取得までの待ち時間まで見たい場合は、DBStats を見るのがよさそう

ことの発端

Wantedly では Application Performance Monitoring に New Relic を利用しています。New Relic APM には様々な機能が存在しますが、例えばエンドポイントごとに「どの処理でどれくらいの時間がかかっているか」などを良い感じに可視化してくれたりします。

最近、あるマイクロサービスでレイテンシ大爆発が起きました。結局その原因は DB に負荷をかけすぎていたことだったのですが、調査の過程で「そもそもこの Postgres とはどこからどこまでを含むんだろう」という疑問が湧いてきました。

database/sql/driver.Driver

Go で RDB を扱うコードを書くとき、大体の場合は次のようなコードから始めるでしょう。

import "database/sql"
import _ "github.com/lib/pq" // どこかで driver を import しておく


db, err := sql.Open("postgres", databaseURL)

この第一引数 "postgres" は Driver の名前です。PostgreSQL でも MySQL でもだいたい最初に driver package の blank import を求められますが、これは driver の名前と実装を紐付けるための記述です。

// https://github.com/lib/pq/blob/v1.2.0/conn.go#L51-L53

func init() {
sql.Register("postgres", &Driver{})
}

ここで sql.Register の第2引数には database/sql/driver.Driver という interface を受けるようになっており、後々どんな DB が生まれても接続できるようになっています。

New Relic など、Go で DB のパフォーマンスを測るための instrumentation の多くは、この driver.Driver をラップする形で実装されています。いくつか眺めてみましょう。

newrelic/go-agent

New Relic の database/sql 用の instrumentation はオフィシャルの go-agent v2.8.0 から同梱されるようになりました。（それ以前は Wantedly では拙作の izumin5210/nrsql を利用していました。実装はほぼ同じ。）

_integrations の下に有名な DBMS ごとに、nr* という名前でパッケージが用意されています。たとえば PostgreSQL であれば nrpq で、内部では pq.Driver をラップした何かを nrpostgres という名前で登録しています。

// https://github.com/newrelic/go-agent/blob/v2.16.0/_integrations/nrpq/nrpq.go#L92-L95

func init() {
sql.Register("nrpostgres", newrelic.InstrumentSQLDriver(&pq.Driver{}, baseBuilder))
internal.TrackUsage("integration", "driver", "postgres")
}

実装は比較的筋肉で、 driver.Driver はもちろん、そこから返りうる driver.Conn や driver.Stmt に対応する wrapper をひとつひとつ定義し、メソッドに計測用のコードを挟んでいます。

// https://github.com/newrelic/go-agent/blob/v2.16.0/sql_driver.go#L181-L193

func (w *wrapConn) Query(query string, args []driver.Value) (driver.Rows, error) {
return w.original.(driver.Queryer).Query(query, args)
}


// QueryContext implements QueryerContext.
func (w *wrapConn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
segment := w.bld.useQuery(query).startSegment(ctx)
rows, err := w.original.(driver.QueryerContext).QueryContext(ctx, query, args)
if err != driver.ErrSkip {
segment.End()
}
return rows, err
}

これはひとつ面白いポイントですが、New Relic は追跡用の ID みたいなものを context.Context に入れて伝搬させるため、context.Context を渡せない (*sql.DB).Query などを利用していると正しく計測が行われません。これは他の多くの計測ツールでも同様です。

このようにして PostgreSQL や MySQL にクエリを投げる前後に計測処理を挟むことで、あの便利なグラフが描かれています。

opencensus-integrations/ocsql

もうひとつ、分散トレーシングのフレームワークである OpenCensus の Go 実装も見てみましょう。

こちらも Register の API は用意されていますが、Register なしでも使えるよう Wrap という関数が公開されています。これは名前の通りラップされた driver.Driver を返してくれます。

// https://github.com/opencensus-integrations/ocsql/blob/v0.1.4/driver.go#L89-L99

// Wrap takes a SQL driver and wraps it with OpenCensus instrumentation.
func Wrap(d driver.Driver, options ...TraceOption) driver.Driver {
o := TraceOptions{}
for _, option := range options {
option(&o)
}
if o.QueryParams && !o.Query {
o.QueryParams = false
}
return wrapDriver(d, o)
}

こちらも New Relic のものとおなじく、database/sq/driver のインタフェースひとつひとつに wrapped struct を定義しています。大変そう。

// https://github.com/opencensus-integrations/ocsql/blob/v0.1.4/driver.go#L275-L313
// 一部省略

func (c ocConn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (rows driver.Rows, err error) {
defer recordCallStats(ctx, "go.sql.query")(err)


if queryerCtx, ok := c.parent.(driver.QueryerContext); ok {
parentSpan := trace.FromContext(ctx)

// ...

defer func() {
setSpanStatus(span, err)
span.End()
}()


rows, err = queryerCtx.QueryContext(ctx, query, args)
if err != nil {
return nil, err
}


return wrapRows(ctx, rows, c.options), nil
}


return nil, driver.ErrSkip
}

Driver でどこまで計測できるのか

さて、ここまでで New Relic や OpenCensus の SQL instrument が「クエリが行って返ってくるまで」を観測していることがわかりました。ところで、New Relic にはたまに次のようなメトリクスが記録されることがあります。

このタイミングではちょうど DB の負荷が高かったことがわかっています。黄色の Postgres はまだわかりますが、水色の Go はなぜボコッと出っ張ってしまったのでしょうか。この謎を解明するため、毎週水曜に開催しているコードリーディング会で *sql.DB の内部を読んでみました。

一般的な SELECT クエリを発行する際に利用する (*sql.DB).QueryContext をエントリポイントとして、コードを読んでいきます。

ここでは driver.ErrBadConn が返ってきたときにリトライしつつ、 (*sql.DB).query を実行しています。

// https://github.com/golang/go/blob/go1.13.4/src/database/sql/sql.go#L1541-L1556

// QueryContext executes a query that returns rows, typically a SELECT.
// The args are for any placeholder parameters in the query.
func (db *DB) QueryContext(ctx context.Context, query string, args ...interface{}) (*Rows, error) {
var rows *Rows
var err error
for i := 0; i < maxBadConnRetries; i++ {
rows, err = db.query(ctx, query, args, cachedOrNewConn)
if err != driver.ErrBadConn {
break
}
}
if err == driver.ErrBadConn {
return db.query(ctx, query, args, alwaysNewConn)
}
return rows, err
}

(*sql.DB).query は conn と queryDC の2つの関数を呼び出します。

// https://github.com/golang/go/blob/go1.13.4/src/database/sql/sql.go#L1564-L1571

func (db *DB) query(ctx context.Context, query string, args []interface{}, strategy connReuseStrategy) (*Rows, error) {
dc, err := db.conn(ctx, strategy)
if err != nil {
return nil, err
}


return db.queryDC(ctx, nil, dc, dc.releaseConn, query, args)
}

この2つの関数について、シグネチャを見てみます。 conn のほうは、 driverConn を返しています。これは名のとおりですが、 driver.Conn をラップしたものです。実際、そのあとの queryDC の実装を流し見すると、この driverConn から driver.QueryerContext を取り出していることが見て取れます。

// https://github.com/golang/go/blob/go1.13.4/src/database/sql/sql.go#L1129-L1246

// conn returns a newly-opened or cached *driverConn.
func (db *DB) conn(ctx context.Context, strategy connReuseStrategy) (*driverConn, error) {

// https://github.com/golang/go/blob/go1.13.4/src/database/sql/sql.go#L1573-L1639

// queryDC executes a query on the given connection.
// The connection gets released by the releaseConn function.
// The ctx context is from a query method and the txctx context is from an
// optional transaction context.
func (db *DB) queryDC(ctx, txctx context.Context, dc *driverConn, releaseConn func(error), query string, args []interface{}) (*Rows, error) {

では、 conn の実装はどうでしょう。なかなかに長い関数なのですべては引用しませんが、関数の最後で (driver.Connector).Connect を実行していることがわかります。そして、関数の手前の部分では maxOpen や numOpen の値を見つつ、プールからコネクションの取り出し・待ちを行っています。

// https://github.com/golang/go/blob/go1.13.4/src/database/sql/sql.go#L1228-L1245
// 一部省略

ci, err := db.connector.Connect(ctx)
if err != nil {
// ...
}
db.mu.Lock()
dc := &driverConn{
db: db,
createdAt: nowFunc(),
ci: ci,
inUse: true,
}
db.addDepLocked(dc, dc)
db.mu.Unlock()
return dc, nil

numOpen >= maxOpen であった場合、 Idle なコネクションができるまでは処理をブロックします。そして、ここは driver.Conn を取り出す前なので、driver.Driver をラップした instrumentation の計測には引っかからないことになります。もしかしたら、これが Go レイヤーでの謎のレイテンシ増加の原因かもしれません。

DBStats

この「Idle なコネクション待ち」がどれくらいアプリケーションのレイテンシに影響しているか、どうすれば計測できるでしょう。

単純な例だと、 sql.DB のラッパー型を作るというのがあるでしょう。しかし、sql.DB は interface ではなく、sqlx など一部のパッケージでは明確に *sql.DB を要求してくることがあります。ラッパー型だとこのような場合に対応できないため、悩ましいところです（sqlboiler などは、sql.DB を interface 化した型を独自定義していたりします）。

ところで、 (*sql.DB).conn に目を戻してみると、コネクション待ちになるタイミングで waitCount と waitDuration なる値を変更しているコードが存在します。

// https://github.com/golang/go/blob/go1.13.4/src/database/sql/sql.go#L1177-L1191

db.waitCount++
db.mu.Unlock()

waitStart := time.Now()

// Timeout the connection request with the context.
select {
case <-ctx.Done():
// ...

atomic.AddInt64(&db.waitDuration, int64(time.Since(waitStart)))

これらは (*sql.DB).Stats という関数で取得できる sql.DBStats というオブジェクト経由で取得することができます。この sql.DBStats に含まれる値を何かしらのツールで監視できれば、コネクションが取得できない問題にも簡単に気付くことができるかもしれません。

まとめ

TL;DR に書いたとおりですが、 Go で発行した SQL のメトリクスを監視したい場合は driver.Driver のラッパー型を用意するのが常套手段です。独自の Instrumentation を実装したい場合は New Relic や OpenCensus の実装が参考になるでしょう。

ただし、Driver だけを監視しているとコネクションの確立自体に問題が起きた場合や、コネクションプールが枯れた場合の原因究明がすこし難しいかもしれません。もし謎のレイテンシ増加に悩まされているのなら、 (*sql.DB).Stats のことを思い出すと良いことがあるかもしれません。