ピン留めされたアイテム

概要

SaasのリアーキテクトでApartment gemによるテナント分離方式からRLS(Row Level Security)を用いたテナント分離方式に変える方向で動いている
Apartment gemはpostgresqlのschemaを切り替えることによって分離を実現するマルチテナント方式だが
RLSは各行のカラムを利用してテナント分離を実現するシングルテナント方式なので
本番データを全てシングルスキーマへと移行する必要がある
その場合既存データのプライマリキーの連番がシングルスキーマ化によって競合してしまうので対策を考える必要がある。

これまで
テナントAスキーマ.users

id	name	email
1	hoge	hoge@a.com
2	fuga	fuga@a.com
3	piyo	piyo@a.com

テナントBスキーマ.users

id	name	email
1	hoge	hoge@b.com
2	fuga	fuga@b.com
3	piyo	piyo@b.com

↓
これから
共通スキーマ.users

id	name	email
1	hoge	hoge@a.com
2	fuga	fuga@a.com
3	piyo	piyo@a.com
1	hoge	hoge@b.com
2	fuga	fuga@b.com
3	piyo	piyo@b.com

↑プライマリキーの連番が重複している！

連番のIDに依存している実装や要件が数多く存在しており、連番をアプリケーションから排除することはできない、辛いと結論づいたためテナント/テーブル単位で連番を採番できるようにする必要がある
※urlの連番を維持したい、他サービスとの結合キーに利用など→ストレージパスに利用してるのが致命的だった。。。

詳細

IDを採番する上の要件としては

番号の重複があってはならない
並列で採番を行なっても重複が起きることなく採番ができる
テナント/テーブル単位でUniqueな連番
ある程度高速に(パフォーマンスを極端に損なわない)
保守性や拡張性に問題がない

これを達成するための案として

ID生成機
Redisでシーケンスオブジェクト作成
DBでシーケンスオブジェクト作成
が候補に上がった。

今回は最もシンプルかつコストの少ないDBで行う方式について検討している

DBでシーケンスを実装する場合、シーケンス保持用のテーブルを作成し、
シーケンスの更新/番号の払い出しをアプリケーションサーバー側で実装するのが良さそうではあるが、
ネットワークI/Oによるボトルネックが懸念されるので
postgresqlのユーザー定義関数を利用することとした。

具体的には

最新の連番をシーケンス用テーブルから取得。
取得できればincrementして次の番号として返却
取得できなかった場合、連番取得対象のテーブル/テナントのidの最終値を取得
idの最終値が取れればそれにincrementして次の番号として返却
idの最終値が取れなければ1からstart(新規テーブル/テナントの組み合わせとみなす)

で、かつIDの競合が起きないように、行ロックやアドバイザリロックを利用する

そうして定義したpostgresqlのユーザー定義関数をRailsのbefore_create側で呼び出させる。
before_createやfunctionを呼び出すロジックは全てconcernに詰め込んで共通化しておく。

上記前提でぶつかった問題を挙げて整理していく

take_webengineer

連番を発行するためのカスタムシーケンス実装でトラブル発生

    def set_next_id_for_tenant
      tenant_name = Apartment::Tenant.current
      tenant_id = Tenant.find_by(name: tenant_name)&.id

      # connectionを分離する(トランザクションを採番処理のみ分離することで採番の重複を防ぐ)
      ActiveRecord::Base.connected_to database: :primary_sequence do
        # テナント指定(RLS or schema switch)
        # connected_toもapartmentでpatchされているため厳密には以下不要だがパフォーマンスの差分を計測するためにあえて記載
        if $test_patterns[:tenant_switch] === :schema_switch
          Apartment::Tenant.switch! tenant_name
        else
          ActiveRecord::Base.connection.execute("SELECT set_config('app.current_tenant_id', '#{tenant_id}', false)")
        end
        tenant_id = "NULL" if tenant_id.nil?
        get_next_id(tenant_id)
      end
    end

  def get_next_id(tenant_id)
    # max_retry_count = ENV.fetch("MAX_SEQ_ATTEMPTS",3).to_i
    max_retry_count = 5
    retry_count = 0
    self.tenant_id = tenant_id
    begin
      ActiveRecord::Base.transaction do
        self.sample_id = ActiveRecord::Base.connection.execute("SELECT public.custom_next_val('#{self.class.table_name}',#{tenant_id})")[0]["custom_next_val"]
      end
    rescue ActiveRecord::LockWaitTimeout,ActiveRecord::StatementInvalid => e
      if retry_count < max_retry_count
        # rubocop:disable Rails/Output
        print "\n\e[33mretry #{retry_count} times. | thread_object_id: #{Thread.current.object_id} | time: #{Time.now}\e[0m"
        puts "\e[33mClass: #{e.class}  |   Message: #{e.message}\e[0m"
        # rubocop:enable Rails/Output
        retry_count += 1
        sleep((2 ** retry_count) + rand(0.0..1.0))
        retry
      else
        # raise
        puts "\e[31mretry failed. #{Thread.current.object_id}\e[0m"
        puts "\e[31m#{e.message}\e[0m"
      end
    end

get_next_idでrescueしたとき、transaction側で先にrescueとrollbackが走っているはずなのだが、
rollbackされておらず、abortされた状態のtransactionが生き残り続ける問題が発生
ERROR: current transaction is aborted, commands ignored until end of transaction block

コード例のように雑にROLLBACKをEXECUTEするとトランザクションがないケースも発生している
WARNING: there is no transaction in progress

何か見落としてそう。。。
connected_to/transaction/create!とかtransactionに絡む実装を読んでみることにする

take_webengineer

postgresqlの実装としてはユーザー定義関数(function)は上位のトランザクション内部で実行され、exceptionが関数内部で発生すると上位のtransactionをabortするとのこと。

2024-02-21 16:56:39.170 GMT [802] LOG:  duration: 0.048 ms
2024-02-21 16:56:39.177 GMT [802] LOG:  statement: COMMIT
2024-02-21 16:56:39.177 GMT [802] LOG:  duration: 0.769 ms
2024-02-21 16:56:38.293 GMT [802] STATEMENT:  SELECT public.custom_next_val('hoges',1)
2024-02-21 16:56:38.305 GMT [802] ERROR:  current transaction is aborted, commands ignored until end of transaction block

postgresqlログ側では実行とエラーしか吐き出してくれないんだよなあ。。。
直前のログでcommitしてその後クエリ走ってないしabortされるタイミングがないように見えるのだが。。。
視点として何が抜けてるのか

take_webengineer

ログと実装眺めててちょっと気づいたけど

before_createより前にbegin(create!側の処理)
採番function側のtransactionブロックはすでにActiveのトランザクションがあるのでスルー
採番functionの実行にコケる
トランザクションブロック内でExceptionが出てるのでRollbackが走るはずだが生きてるトランザクションは1のトランザクションなのでRollbackしない
2,3でfunctionコケてabortしたままrollbackしないので後続のクエリが全てコケる(無視される)

って流れかも。推測の範囲。実装読まねば

# どこで始まってる?
# before_createより前にbegin
2024-02-21 16:56:38.361 GMT [802] LOG:  statement: BEGIN
2024-02-21 16:56:38.361 GMT [802] LOG:  duration: 0.127 ms
2024-02-21 16:56:38.366 GMT [802] LOG:  duration: 0.123 ms  parse <unnamed>: SELECT "public"."tenants".* FROM "public"."tenants" WHERE "public"."tenants"."name" = 'n2i' LIMIT 1
2024-02-21 16:56:38.366 GMT [802] LOG:  duration: 0.072 ms  bind <unnamed>: SELECT "public"."tenants".* FROM "public"."tenants" WHERE "public"."tenants"."name" = 'n2i' LIMIT 1
2024-02-21 16:56:38.366 GMT [802] LOG:  execute <unnamed>: SELECT "public"."tenants".* FROM "public"."tenants" WHERE "public"."tenants"."name" = 'n2i' LIMIT 1
2024-02-21 16:56:38.366 GMT [802] LOG:  duration: 0.020 ms
2024-02-21 16:56:38.372 GMT [802] LOG:  statement: SELECT COUNT(*) FROM pg_namespace WHERE nspname = 'hoge_tenant'
2024-02-21 16:56:38.373 GMT [802] LOG:  duration: 0.490 ms
2024-02-21 16:56:38.375 GMT [802] LOG:  statement: SET search_path TO "hoge_tenant"
2024-02-21 16:56:38.375 GMT [802] LOG:  duration: 0.178 ms
2024-02-21 16:56:38.377 GMT [802] LOG:  statement: SELECT set_config('app.current_tenant_id', '1', false)
2024-02-21 16:56:38.377 GMT [802] LOG:  duration: 0.085 ms
2024-02-21 16:56:38.382 GMT [802] LOG:  statement: SELECT public.custom_next_val('hoges',1)

    included do
      before_create :set_next_id_for_tenant
    end

    def set_next_id_for_tenant
      tenant_name = Apartment::Tenant.current
      tenant_id = Tenant.find_by(name: tenant_name)&.id

      # connectionを分離する(トランザクションを採番処理のみ分離することで採番の重複を防ぐ)
      ActiveRecord::Base.connected_to database: :primary_sequence do
        # テナント指定(RLS or schema switch)
        # connected_toもapartmentでpatchされているため厳密には以下不要だがパフォーマンスの差分を計測するためにあえて記載
        if $test_patterns[:tenant_switch] === :schema_switch
          Apartment::Tenant.switch! tenant_name
        else
          ActiveRecord::Base.connection.execute("SELECT set_config('app.current_tenant_id', '#{tenant_id}', false)")
        end
        tenant_id = "NULL" if tenant_id.nil?
        self.sample_id = get_next_id(tenant_id)
      end
    end

take_webengineer

じゃあやっつけでrescue側にROLLBACK仕込むとどうなるかというと
WARNING: there is no transaction in progress
が発生する。

この推測が正しいと仮定すると、

差し込んだROLLBACKの後にcreate!がexceptionを拾ったことによるROLLBACKが別で実行されて
余剰なROLLBACKとなってしまっている説がありそう

現象的に推測の通りだとしか考えられないけど
とりあえず、create! , transaction絡みの実装読まないと判断できないな