【HRBrain初のアドベントカレンダー!!】PostgreSQLのRow Level SecurityをGoで単一DBコネクションプールにて使いこなす

こんにちは。HRBrainでSREをやらせてもらっている@mano_0307です。

今年も早いものでアドベントカレンダーの時期になってしまいました。あっという間の年末ですね。
昨年まで弊社ではアドベントカレンダーをやる文化はなかったのですが、先月終了間際にダメ元で「やろうよ!」とチームの方々にお声がけしたところ、予想よりも多い方が「書く!」と集まってくれました。僕としては赤井秀一(名探偵コナンの登場人物)を真似して

「・・・嫌だ、と言ったら?」

と、全員から断られてしまうことも覚悟していたので、とても嬉しく思います。

まぁでも急すぎたのもあって、さすがに25人は集まらなかった・・・。なので穴開きカレンダーです。構うものですか!アドベントカレンダーに穴が開いているからとて、別に犯罪でもあるまいに!(居直り)

でも、このようなツイートをしたのですが、

誰かが書かないから、その分Aさんが複数記事を書いた。Aさんは偉い!
書かない人、けしからんぞ!僕はこんなに寝る間を惜しんで書いたのに・・・
・・・なんか書いてないからなんか肩身狭いな・・・

みたいな分断は絶対に、絶対に避けたい。そういう考えもあって無理はせず、穴は穴のまま残します。穴だって僕らの建設的な選択の結果です。


前置きが長くなりましたが、そういうわけでこの記事はHRBrainアドベントカレンダーの第一日目の記事です。

qiita.com

HRBrainはマルチテナント・アーキテクチャ

弊社が運用しているHRBrainは、クラウド人材サービスです。複数の企業様にご利用いただいている「マルチテナント型SaaS」です。

マルチテナント・アーキテクチャのサービスを運用する中で、非常に悩ましいのは

「DBをどう管理するか」

ということではないかと思います。

DBはテナント毎に独立させる?それとも全テナント情報を一つのDBにまとめる?

DBを分けようか一つにしようか、悩むポイントを表にしてみました。

テナント毎にDBを分ける 全テナント情報を一つのDBで管理する
メリット WHERE文ミスなどで別テナントの情報が閲覧できてしまう可能性がない ・DBコネクション管理が楽
・メンテが楽
・低コスト(スケール度合いによってはそうとは言い切れないが)
デメリット ・アプリケーションのDBコネクション管理が煩雑化
・マイグレーションやその他メンテナンスが大変
・インスタンス数が増えればコストが上がる
WHERE文を間違えたらもう大変

弊社では、「全テナント情報を一つのDBで管理する」という方法が採用されています。 でもそのまま採用するとWHERE文誤りが致命傷になってしまうため、PostgreSQLのRow Level Securityを使ってそのリスクを回避しています。

マルチテナント・アーキテクチャと、PostgreSQLのRow Level Securityについては、弊社CTOが書いたこちらを是非読んでいただきたいと思います。

speakerdeck.com

また、弊社インフラエンジニアのユキチさんが書いたこちらの記事も具体的な部分が書かれていてとても面白いです。

times.hrbrain.co.jp

当記事は、これらの「続編」という形式を取らせていただきます。

どうすればRow Level Securityを使ってもパフォーマンスが出るようになる?

テナント毎にコネクションプールを持つのは何かとしんどい

上の2記事を読んでいただくとわかるように、Row Level Security利用時に要となるのが、PostgreSQLの「current userもしくはcurrent setting」です。 A社のユーザー様による利用であればA社用のroleでログインした(もしくはA社用のcurrent settingを設定した)コネクションを、B社であればB社用のroleでログインした(もしくはB社用のcurrent settingを設定した)コネクションを使う必要があります。そのため、当初はDBコネクションプールごと会社ごとに分ける形で実装していました。

しかしここで問題になるのが、コネクション管理の煩雑さ(先ほどの表ではメリット側に挙げていたはず)もさることながら、まずは「アクティブなコネクション数の上限」です。
弊社が利用しているCloud SQLでは、使っているマシンタイプ(メモリーサイズ)によって上限値が変わります。

cloud.google.com

アプリケーションはKubernetesで動いており、Podの数を増やすことでスケールアウトできる設計ではありますが、このコネクション数上限がネックとなってきます。
さらに、コネクションの無駄が多くなります。A社からのアクセスは多く、B社からのアクセスは少ないというとき、

  • A社用のコネクションプールは枯渇!猫の手だって借りたい!😹😹😹
  • B社用のコネクションプールにはアイドル状態のコネクションがすやすやと眠っている😪😴💤

という理不尽な状態になるのです。このままではいつかA社用コネクションプールは蜂起します!!(しない)

これを避けるために例えば

アイドル状態のコネクションなんかいらないかも!IdleTimeoutを1秒くらいまで短くして、さっさとCloseするようにしよう!!

ということもできます。ただし、そうすると使用するたびに毎回毎回

  • アプリケーションPodとPostgreSQLとのTCPハンドシェイク(「やぁ初めまして。」「つい1ミリ秒前にきたばっか!寝ぼけてるよね!!」)
  • PostgreSQLにログイン
  • Goのコネクション実体を生成

をやることになります。うーむ。

なんとか一つのコネクションプールで全テナントを管理できないか

どうにかこうにか一つのコネクションプールでRow Level Securityを使えないか。 つまり、

コネクションを使う直前に、所属テナント用のroleに(もしくは設定に)切り替えられないか?

という問題の解を見出すことができればいい、ということでした。

そこで僕は、Goの標準DBライブラリであるdatabase/sqlのコードを、穴が開くほど読みまくりました。もともとこちらのGo製memcachedクライアントライブラリを実装したときにもコネクションプール実装を真似するために読みまくったのですが、それ以来でした。

github.com

いずれにせよ、僕はまるでルーペが黒い紙を燃やしてしまうようにGoのコードを見つめ続けました。そしてついに、たどり 着いた (参照: HUNTER×HUNTER ネテロ vs メルエム)。

GoのSessionResetterインタフェースを使う

ここで簡単に、database/sqlのコネクションプールがどのような実装になっているかを見てみます。
Goのコネクションプール実装で骨格となっているのはこのconn関数です。

何をやっているかというと、このようなことをやっています。

  • アイドル状態のコネクションがあるとき、そこから一つ拝借する
  • アイドル状態のコネクションがなく、かつすでにactiveなコネクションの数がMaxOpenConnsに到達しているとき、もうコネクションを新規作成することはできないため、誰かがコネクションを使い終わるのを待つ。使い終わったコネクションはchannelを介して渡される
  • アイドル状態のコネクションがなく、でもまだactiveなコネクション数がMaxOpenConnsに到達していないなら、ここぞとばかりに「よっしゃ!」とコネクションを新規作成する

さて、僕が辿り着いたのは、database/sqlに用意されているSessionResetterインタフェースでした。

https://github.com/golang/go/blob/e5da18df52e3f81534d7cdb6920cf993b5f079d2/src/database/sql/driver/driver.go#L289-L296

このインタフェースの ResetSession を実装すれば、アイドル状態または他の処理にて使用されたコネクションを使う直前に行う処理を実装できます。

でもそれだけでは不十分です。「新規コネクション生成時」にも行わないと手落ちになります。危ない危ない。

さらにConnectorインタフェースも使う

新規コネクション生成処理を実装するにはConnectorインタフェースを実装します。

https://github.com/golang/go/blob/e5da18df52e3f81534d7cdb6920cf993b5f079d2/src/database/sql/driver/driver.go#L109-L138

この二つ、SessionResetterインタフェースとConnectorインタフェースを実装すれば、既存コネクションを再利用するときも、新規作成直後のコネクションを使うときも網羅できます。

なお、これらのインタフェースを実装したdriverを使って*sql.DBの実体を生成するには、OpenDB関数をコールします。

https://github.com/golang/go/blob/e5da18df52e3f81534d7cdb6920cf993b5f079d2/src/database/sql/sql.go#L714-L743

下記の例では、このSessionResetterインタフェースを実装したconnと、Connectorインタフェースを実装したconnectorのサンプルを書いてみました。

import (
    "context"
    "database/sql/driver"
    "fmt"
    "os"
)

type connector struct {
    dsn string
    d   driver.Driver
}

// Connect database/sqlにて、新しいconnectionが生成されるときに呼ばれる。
// https://github.com/golang/go/blob/689a7a1378/src/database/sql/sql.go#L1301
func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
    cn, err := c.d.Open(c.dsn)
    if err != nil {
        return nil, fmt.Errorf("failed c.d.Open: %w", err)
    }
    if err = setApplicationName(ctx, cn); err != nil {
        return nil, fmt.Errorf("failed setApplicationName: %w", err)
    }
    return &conn{cn}, nil
}

func (c *connector) Driver() driver.Driver {
    return c.d
}

type conn struct {
    driver.Conn
}

// ResetSession database/sqlにて、ConnectionPoolからconnが取り出されるときに呼ばれる。
// https://github.com/golang/go/blob/689a7a1378/src/database/sql/sql.go#L1227
// https://github.com/golang/go/blob/689a7a1378/src/database/sql/sql.go#L1291
func (c *conn) ResetSession(ctx context.Context) error {
    if err := setApplicationName(ctx, c); err != nil {
        return fmt.Errorf("failed setApplicationName: %w", err)
    }
    return nil
}

func setApplicationName(ctx context.Context, cn driver.Conn) error {
    xxxID := GetXXXIDFromContext(ctx) // ctxからapplication_name切り替えに必要な情報を取得する。当サンプルでは未実装。

    stmt, err := cn.Prepare(fmt.Sprintf("SET application_name TO %d", xxxID))
    if err != nil {
        return fmt.Errorf("failed Prepare: %w", err)
    }
    if _, err = stmt.Exec(nil); err != nil {
        return fmt.Errorf("failed stmt.Exec: %w", err)
    }

    return nil
}

var (
    _ driver.Connector       = &connector{}
    _ driver.SessionResetter = &conn{}
)

このように用意したconnectorOpenDB()関数の引数に下記のように渡せば、*sql.DBが生成されます。ちなみに下記の例はpqというPostgreSQLドライバーを使用しています。

github.com

db := sql.OpenDB(&connector{dsn: source, d: &pq.Driver{}})

補足: ORMライブラリを使用するとき

弊社ではマイクロサービス・アーキテクチャを採用しており、サービスによって使用しているORマッパーライブラリも異なります。 弊社で現在メインで使用しているのは、gormxormです。

github.com

gitea.com

gormを使うとき

Config.ConnPoolに、↑のように自分で用意した*sql.DBをセットしてOpenをコールします。

xormを使うとき

当初、xormにはgormのように、自分で作った*sql.DBを渡すような方法は用意されていませんでした。なので下記PRを出したところ、すぐにマージしてくれました。

gitea.com

↑で実装されているNewEngineWithDialectAndDBに、自分で用意した*sql.DBを用意してやればいいです。

その他ORマッパーライブラリを使うときも、「このライブラリは自分で用意した*sql.DBをそのまま使えるか!?」という観点で選ぶのが重要です。

終わりに

こんな感じに、一つのDBでコネクションの無駄を発生させることなく、Row Level Securityを実現させることができています。
あと小並感ですが、database/sqlの実装を読むのはchannelの勉強にもなってとてもおすすめです。

な、なんとか書けました・・・!自分で言い出しっぺとなったアドベントカレンダー、さすがに初日に穴を開けるわけにはいかなかった・・・でも書くテーマもさっき(11/30深夜)決めたし、なぜかこういう日に限って妻と「世界の成り立ち」に関して長時間話し合った結果時間がどんどんなくなるし・・・よくやったよmanoくん!!😭

HRBrainでは、トライドリブンという文化の中、一緒に様々な興味深い難題を解決し続けてくれるエンジニアを募集しています!もう面白い仕事が山ほどあります。いささか手が余るほどにね。是非一緒に働きましょう!

www.wantedly.com