diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index c02a405..8150cd7 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -2,11 +2,21 @@ name: Plugin on: [push, pull_request] permissions: contents: read +concurrency: + group: plugin-${{ github.event_name }}-${{ github.ref }} + cancel-in-progress: ${{ !startsWith(github.ref, 'refs/tags/') }} jobs: test: runs-on: ubuntu-24.04 + strategy: + fail-fast: false + matrix: + postgres: ['16', '17', '18'] + env: + POSTGRES_IMAGE: postgres:${{ matrix.postgres }}-alpine steps: - uses: actions/checkout@v4 + - run: python3 -m unittest discover -s tests/release -v - uses: dtolnay/rust-toolchain@stable with: components: rustfmt, clippy @@ -14,6 +24,7 @@ jobs: - run: cargo clippy --all-targets --locked -- -D warnings - run: cargo test --locked - run: bash scripts/test-postgres.sh + - run: python3 scripts/bench-postgres.py --plans-only --image "$POSTGRES_IMAGE" --output artifacts/query-plans.json - run: cargo build --release --locked - uses: actions/checkout@v4 with: @@ -22,7 +33,29 @@ jobs: path: .packaging persist-credentials: false - run: python3 .packaging/scripts/package-plugin.py target/release/amitoki-plugin-postgres dist + - run: cp schema.sql dist/schema.sql - uses: actions/upload-artifact@v4 with: - name: amitoki-plugin-postgres-x86_64-linux + name: amitoki-plugin-postgres-x86_64-linux-pg${{ matrix.postgres }} path: dist/ + publish: + if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags/v') + needs: test + runs-on: ubuntu-24.04 + timeout-minutes: 10 + permissions: + contents: write + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + - uses: actions/download-artifact@v4 + with: + name: amitoki-plugin-postgres-x86_64-linux-pg17 + path: dist + - name: Verify the package and publish the version tag + env: + GH_TOKEN: ${{ github.token }} + RELEASE_TAG: ${{ github.ref_name }} + RELEASE_REPOSITORY: ${{ github.repository }} + run: python3 scripts/publish-release.py --directory dist --tag "$RELEASE_TAG" --repository "$RELEASE_REPOSITORY" diff --git a/.gitignore b/.gitignore index 20f7978..27d3722 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,7 @@ /target/ /dist/ +/artifacts/ +__pycache__/ .env* *.der *.pem diff --git a/Cargo.lock b/Cargo.lock index cdc9062..fbe2931 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -54,7 +54,7 @@ dependencies = [ [[package]] name = "amitoki-relay-postgres" -version = "0.1.2" +version = "0.2.0" dependencies = [ "amitoki-plugin-sdk", "amitoki-relay", diff --git a/Cargo.toml b/Cargo.toml index 5c067b2..f70f680 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "amitoki-relay-postgres" -version = "0.1.2" +version = "0.2.0" edition = "2021" [dependencies] diff --git a/README.md b/README.md index 8c684bc..ef00bcf 100644 --- a/README.md +++ b/README.md @@ -4,12 +4,12 @@ PostgreSQLを使うamitokiの外部プロセス型プラグイン。本体の再 ## 利用する -amitoki 0.3以降を使う。公開リポジトリなので、公式配布物の取得にGitHubトークンは不要。 +amitoki 0.2以降を使う。公開リポジトリなので、公式配布物の取得にGitHubトークンは不要。 ```bash -amitoki plugin add postgres -amitoki plugin configure postgres --set connection_env=AMITOKI_POSTGRES_URL --set max_connections=4 -amitoki plugin validate postgres +amitoki plugin relay add https://github.com/amitoki/amitoki-plugin-postgres +amitoki plugin relay configure postgres --set connection_env=AMITOKI_POSTGRES_URL --set max_connections=4 +amitoki plugin relay validate postgres read -r -s -p 'PostgreSQL接続文字列: ' AMITOKI_POSTGRES_URL printf '\n' export AMITOKI_POSTGRES_URL @@ -35,13 +35,19 @@ max_connections = 4 replay_window_ms = 4000 ``` -`replay_window_ms`は初回登録前のフレームを何ミリ秒分受け取るか。再起動時は値に関係なく既存の未ACKキューを引き継ぐ。 +`replay_window_ms`は初回登録前のフレームを何ミリ秒分受け取るか。基準は登録ロック取得後の再生SQL開始時刻。0では履歴を取り込まない。再起動時は値に関係なく既存の未ACKキューを引き継ぐ。 ## 配送の扱い -フレームと各ノードの未処理キューを同じトランザクションで保存する。receiveは非破壊で、NICへの注入後にACKする。同じUUIDの再送とACKは冪等。登録・送信の競合はchannel単位のadvisory lockで防ぎ、同一node_idの多重起動は専用DB接続で拒否する。 +フレームと各ノードの未処理キューを同じトランザクションで保存する。receiveは非破壊で、NICへの注入後にACKする。同じUUIDの再送とACKは冪等。送信同士はchannelの共有advisory lockで並行実行し、ノード登録とは排他にする。DBの既定に関係なく登録・送信はREAD COMMITTEDで実行する。同一node_idの多重起動は専用DB接続で拒否する。 -NIC注入とACKの間のクラッシュは重複し得る。DBの保存済みデータは再起動後も残るが、収集直後の本体メモリは永続化しない。DBの保存期限・掃除は自動化していない。 +単一バッチ内の入力順を維持する。別の送信が先にcommitすることはあり、全送信元を通じたcommit順は保証しない。重なるUUIDを逆順で送った場合などのDBデッドロックは、再試行可能エラーで返す。本体は同じUUIDのバッチを再送できる。 + +NIC注入とACKの間のクラッシュは重複し得る。DBの保存済みデータは再起動後も残るが、収集直後の本体メモリは永続化しない。保持期限を設定すると、削除済みUUIDの再送は新規フレームになる。 + +## 保持期限とDB更新 + +自動削除は既定で無効。利用する場合は、channel内の全ノードで`retention_ms`を指定する。設定例とDB更新手順は[保持期限と運用](docs/retention.md)を参照。既存DBでは、更新した実行ファイルから取り出した`schema.sql`をプラグインの起動前に再適用する。ノード設定用の列と未ACK確認用の索引を追加し、既存フレーム・未ACKは維持する。 ## 開発・試験 @@ -60,8 +66,24 @@ cargo build --release --locked Dockerを実行できるユーザで試験する。DB試験は一時コンテナを終了時に削除する。公開SDKのGit revisionとCargo.lockを固定している。 +送信・受信・ACKの性能とSQL実行計画は[コンテナでの検証手順](docs/benchmark.md)で測定する。検証用の一時DBで、ノード数・パケット長・バッチサイズを変えて全フレームを照合する。 + 配布物は本体の`scripts/package-plugin.py target/release/amitoki-plugin-postgres dist`で生成する。CIも同じスクリプトを使用する。Linux x86_64向けの初回配布はUbuntu 24.04でビルド・試験した。 +## リリースする + +`Cargo.toml`と`Cargo.lock`の版を更新し、`docs/releases/v<版>.md`へ更新手順を追加する。mainへマージしてCIが成功したcommitに注釈タグ`v<版>`を付けてpushする。タグのCIがPostgreSQL 16・17・18の試験を完了すると、実行ファイル・manifest・SQLの一致を検証し、SHA256SUMSとともにGitHub Releasesへ公開する。正式公開前にアップロード済み全ファイルのSHA256を照合する。公開済み版を上書きしない。 + +```bash +git switch main +git pull --ff-only +# 作業ツリーがクリーンで、対象commitのCIが成功していることを確認する。 +git tag -a v0.2.0 -m 'PostgreSQLプラグイン0.2.0を公開' +git push origin v0.2.0 +``` + +公開前のdraftで止まった場合は、原因を確認して同じタグのCIを再実行できる。コード修正が必要な場合は新しい版を作り、既存タグは変更しない。 + ## stegrdb版から移行する 旧タグとリリースは保持している。amitoki版は実行ファイル・crate・環境変数の接頭辞をamitokiへ変更した。本体の[移行手順](https://github.com/amitoki/amitoki/blob/main/docs/amitoki-migration.md)に従って設定とプラグインを配置する。 diff --git a/docs/benchmark.md b/docs/benchmark.md new file mode 100644 index 0000000..fa423d3 --- /dev/null +++ b/docs/benchmark.md @@ -0,0 +1,66 @@ +# コンテナで配送性能を検証する + +検証用のPostgreSQLだけを作り、全フレームの内容とACK完了を照合する。既存DBの接続情報は使わない。プラグインのRust実装とDBの性能を測るため、NIC・本体のパイプライン・外部プラグインのMessagePack通信は測定に含まない。 + +## 実行 + +Rust・Python 3・Dockerを準備する。Rust/Dockerの導入は[README](../README.md#開発試験)を参照。 + +```bash +cargo build --release --example relay_bench --locked +python3 scripts/bench-postgres.py \ + --output artifacts/postgres-bench/result.json +``` + +各条件を3回測る。コンテナは4CPU・2GiB、ホストのloopbackのランダムポートを使う。既定イメージは`postgres:17-alpine`で、結果には実際のイメージIDとPostgreSQL版を記録する。比較を再現する場合は`--image postgres@sha256:...`で同じイメージを指定する。 + +`fsync`と`synchronous_commit`を無効化しない。コンテナの通常ストレージを使い、tmpfsへDBを置かない。ホストのディスクや同時負荷によって速度は変わる。例外・中断時も、このコマンドが作ったコンテナと匿名volumeを削除する。 + +## 測定条件と結果 + +- 2・3ノードすべてが同時に送信し、他ノードの全フレームを受信・照合・ACKする。 +- フレーム長64・1400バイト、送信バッチ1・32・128件。受信は最大128件。 +- 本測定前に各ノード1024件を送信し、統計を更新してから受信する。ウォームアップは速度に含めない。 +- `--warmup-frames 0`では空テーブルの統計から始める。投入直後の実行計画の影響を含むため、通常の測定値と混ぜない。 +- 速度は「元の送信フレーム数 ÷ 全ノードのACK完了時間」。複数宛先への配送数は`verified_deliveries`に別記する。 +- バッチ送信のp50・p95・p99、ACK後にDBに残るフレーム数、未ACK数を保存する。保存フレーム数にはウォームアップを含む。 + +各条件の後、検証用データを削除する。この削除時間は配送速度に含めない。通常の比較では保持期限による自動削除を無効にする。 + +## 改善候補との比較 + +比較用ビルドを別の出力先で作り、両方の実行ファイルを指定する。 + +```bash +python3 scripts/bench-postgres.py \ + --binary /absolute/path/current/relay_bench \ + --compare-binary /absolute/path/candidate/relay_bench \ + --output artifacts/postgres-bench/comparison.json +``` + +比較ビルドでは`CARGO_TARGET_DIR`を分ける。同じpackage名・版のworktree同士で出力先を共有すると、成果物を取り違える場合がある。スクリプトは実行ファイルのSHA256を記録し、同一ファイル同士の比較を拒否する。実行順は繰り返しごとに交替する。 + +## SQL実行計画だけを調べる + +```bash +python3 scripts/bench-postgres.py --plans-only \ + --baseline-queries /absolute/path/before/src/queries.rs \ + --output artifacts/postgres-bench/query-plans.json +``` + +受信SQLは1万件のフレームと128件の未ACKで、統計更新前後の計画を比較する。初期条件では、以前のchannelの統計を残したままデータを入れ替え、新channelの件数が少なく見積もられる状況を再現する。時間の閾値ではなく、履歴を繰り返し全走査していないことを実行計画の実走査件数で検証し、PostgreSQL 16/17/18のCIでも実行する。初回再生SQLは25万件の古い履歴で測定する。`--baseline-queries`には比較元の`src/queries.rs`を指定し、省略時は現行SQLだけを測る。比較元のファイルは`git show :src/queries.rs > before-queries.rs`で取得できる。配送するUUID・本文・順序の一致も確認する。 + +`EXPLAIN ANALYZE`の結果には実行時間・走査件数・バッファ参照を保存する。実際の稼働中はパラメータ、prepared statementの計画、データ分布、統計の鮮度で計画が変わるため、この固定データでの結果だけで常時高速化を保証しない。 + +初回再生の基準時刻にはSQL開始時刻で固定される`statement_timestamp()`を使う。[PostgreSQLの時刻関数](https://www.postgresql.org/docs/17/functions-datetime.html#FUNCTIONS-DATETIME-CURRENT)の違いに従い、インデックスで期間を絞れるようにしている。 + +## 接続を維持した耐久試験 + +```bash +python3 scripts/bench-postgres.py --soak-seconds 300 --retention-ms 5000 \ + --output artifacts/postgres-bench/soak.json +``` + +同じ測定プロセス内で3つのRelayとDB接続を維持し、1400バイト・バッチ128件の送受信を繰り返す。すべての内容とACKを照合し、約10秒ごとのRSS、残存フレーム数、未ACK数、表と索引の使用量を記録する。上の5秒保持は掃除を短時間で試験するための値で、本番の推奨保持期間ではない。`--retention-ms 0`では履歴を残したまま測る。耐久試験では掃除間隔を100ミリ秒にする。 + +各ラウンドの速度と、生成・観測を含む全体時間を別に記録する。10秒のサンプルはラウンド完了後なので、負荷が高いと記録間隔は延びる。RSSには生成・照合用メモリを含み、DBの常駐メモリを含まない。5分の成功だけで日単位のメモリ安定性やDB容量上限は保証しない。 diff --git a/docs/releases/v0.2.0.md b/docs/releases/v0.2.0.md new file mode 100644 index 0000000..1791621 --- /dev/null +++ b/docs/releases/v0.2.0.md @@ -0,0 +1,27 @@ +# PostgreSQLプラグイン 0.2.0 + +同じchannelへの送信を並行処理し、履歴が増えた場合や統計が古い場合の受信・再生SQLを改善しました。ACK済みフレームの保持期限も設定できます。 + +- ノード登録との排他を維持したまま送信を共有ロックへ変更し、登録・送信の分離レベルをREAD COMMITTEDに固定しました。 +- 受信は配送対象の主キーを参照し、古い統計による履歴の全走査を避けます。初回再生には期間索引を使います。 +- `retention_ms`・`cleanup_interval_ms`・`cleanup_batch_size`を追加しました。自動削除は既定で無効です。全登録ノードが有効にした場合だけ、最長の保持・再生期間を過ぎたACK済みフレームを削除します。未ACKは保持します。 +- 古いACKが、期限切れ削除後に同じUUIDで登録された新しい配送を消さないようにしました。 +- PostgreSQL 16・17・18の実DB試験と実行計画の回帰検証をCIで実行します。タグのCI成功後、実行ファイル・manifest・SQL・SHA256SUMSを公開します。 + +3ノード・1400バイト・128件バッチの比較では、合計約54.7Mbpsから179.1Mbpsへ改善しました。PostgreSQL 17.11、4CPU/2GiBのコンテナ、loopback・TLSなしで各条件3回の中央値です。本体・IPC・NICは含みません。72試行で884,736配送、保持を有効にした約5分の試験で3,366,912配送の内容とACKを照合しました。 + +## v0.1.2から更新する + +本体はamitoki 0.2以降を使えます。本体の再ビルドは不要です。対象DBを使うamitokiを停止してから、プラグインを更新します。 + +```bash +amitoki plugin relay update postgres +~/.local/share/amitoki/plugins/postgres/amitoki-plugin-postgres --schema > schema.sql +psql "$AMITOKI_POSTGRES_URL" --set ON_ERROR_STOP=1 -f schema.sql +``` + +`AMITOKI_PLUGIN_DIR`・`XDG_DATA_HOME`を使う場合は、その配置先からSQLを取り出してください。接続変数名を変更している場合も読み替えてください。スキーマ更新後にamitokiを起動します。既存データと未ACKは維持します。大きな既存DBでの索引作成と、以前の版へ戻す手順は[保持期限と運用](https://github.com/amitoki/amitoki-plugin-postgres/blob/v0.2.0/docs/retention.md)に記載しています。 + +保持期間は容量上限ではありません。停止したノードの未ACKは残ります。DB再起動で占有ロックを失った場合には本体の明示再起動が必要です。 + +配布対象はLinux x86_64です。P2Pプラグインと中継IPC v1は変更していません。 diff --git a/docs/retention.md b/docs/retention.md new file mode 100644 index 0000000..a4abf17 --- /dev/null +++ b/docs/retention.md @@ -0,0 +1,59 @@ +# 保持期限と運用 + +既定の`retention_ms = 0`は自動削除を無効にする。フレームの保存期間を決めた場合だけ、channel内の全ノードで有効にする。期間は保存時刻から数え、ACK時刻からは数えない。 + +```toml +[relay.options] +connection_env = "AMITOKI_POSTGRES_URL" +max_connections = 4 +replay_window_ms = 4000 +retention_ms = 86400000 # 24時間 +cleanup_interval_ms = 1000 +cleanup_batch_size = 4096 +``` + +CLIの設定も使える。保存した設定を稼働中のプラグインへ反映するには本体を再起動する。 + +```bash +amitoki plugin relay configure postgres --set retention_ms=86400000 +``` + +## 削除する条件 + +登録済みノードが1つでも`retention_ms = 0`なら、そのchannelでは削除しない。全ノードが有効にした場合、最長の保持期間と再生期間を適用する。停止中ノードの設定もDBに残して保護する。設定変更は接続時の登録トランザクションで反映する。 + +期限を過ぎても、いずれかのノードに未ACKが残るフレームは削除しない。従って、停止したノードのキューや削除を無効にしたノードがあれば、DBの使用量は増え続ける。容量上限を保証する設定ではない。ノードの自動削除や未ACKの破棄は行わない。 + +保持期間内はUUIDの再送を重複排除する。実際に削除した後の同一UUIDは新規フレームとして配送する。保持期間を後から延ばしても、削除した履歴は復元されない。 + +受領情報にはUUIDと配送の世代（DBのposition）を含める。削除前のACKが遅れて届いても、新しい世代の未ACKを消さない。受領情報はプラグインが返した値をそのままACKに渡し、呼び出し側でUUIDから作らない。旧版のUUIDだけの受領情報は新しいプロセスでは受け付けない。 + +| 設定 | 範囲 | 既定 | +| --- | --- | --- | +| `retention_ms` | 0、または自ノードの再生期間以上〜2147483647ミリ秒 | 0 | +| `cleanup_interval_ms` | 100〜3600000ミリ秒 | 1000 | +| `cleanup_batch_size` | 1〜65536件 | 4096 | + +掃除は送受信と同じ接続プールを利用する。登録とは排他、送信とは共有のロックで動き、登録中は待たずに次回へ回す。複数の掃除処理が重なった場合はロック済みの行をスキップする。1回の削除件数と各SQLの250ミリ秒の期限で処理を制限する。期限超過では削除件数を半減し、設定間隔と1秒の長い方で再試行する。50ミリ秒未満で上限件数を16回連続で削除できたら25%ずつ元の上限へ戻す。その他の失敗は60秒後に再試行する。stderrには認証情報を含まない診断を記録する。削除件数の上限は走査件数の上限ではない。 + +掃除能力より多く保存すると期限切れの行が残る。間隔とバッチ件数は実際の負荷で調整する。DELETE後の領域はPostgreSQLのVACUUMで再利用され、DBファイルが直ちに縮むわけではない。通常のautovacuum/autoanalyzeは無効にしない。 + +`schema.sql`は中継の2表だけにvacuum/analyzeのscale factorを0.02、vacuum cost limitを1000に設定する。末尾ページの縮小に伴う排他ロックを避けるため`vacuum_truncate=false`を指定し、領域は次のINSERTで再利用する。DB全体の設定は変更しない。[PostgreSQLの表単位の設定](https://www.postgresql.org/docs/17/sql-createtable.html#SQL-CREATETABLE-STORAGE-PARAMETERS)で調整できる。高負荷でVACUUMが追いつかない場合はDB管理者が起動間隔やI/O予算も確認する。 + +## 既存DBの更新 + +新しいバイナリの`--schema`を使って、旧ノードを動かしたまま事前に列と索引を追加できる。大きな表での索引作成は書き込みを待たせるため、下記のconcurrent作成を先に行う。 + +```bash +~/.local/share/amitoki/plugins/postgres/amitoki-plugin-postgres --schema > schema.sql +# 既存のpending表が大きい場合だけ先に実行する。トランザクションに入れない。 +psql "$AMITOKI_POSTGRES_URL" --set ON_ERROR_STOP=1 \ + --command 'CREATE INDEX CONCURRENTLY IF NOT EXISTS pending_frame ON stegrdb_relay.pending (channel, frame_id)' +psql "$AMITOKI_POSTGRES_URL" --set ON_ERROR_STOP=1 -f schema.sql +``` + +新規DBは`schema.sql`だけで初期化する。concurrent索引作成を中断した場合は`pg_index.indisvalid`を確認し、無効な索引を作り直してから再適用する。`IF NOT EXISTS`は無効な索引の修復をしない。 + +旧ノードの列は削除無効で初期化される。各ノードが新しい版で再接続すると設定を更新するため、一部だけの更新では掃除を開始しない。停止中の登録も含めて移行する。旧版へ戻す場合は、先に新しい版で`retention_ms=0`へ変更して再接続し、掃除を無効にしてから旧版へ戻す。列と索引は残してよい。 + +DB再起動によるノード占有接続の切断は、従来どおり恒久エラーとして停止する。本体を明示的に再起動して占有を取り直す。自動再接続とOSサンドボックスはこの変更に含まない。 diff --git a/examples/relay_bench.rs b/examples/relay_bench.rs new file mode 100644 index 0000000..23ba377 --- /dev/null +++ b/examples/relay_bench.rs @@ -0,0 +1,233 @@ +//! 専用DBで複数ノードの送信・受信・ACKを並行実行し、全フレームの内容を照合する。 +use amitoki_relay::{Frame, Relay, RelayContext, RelayPlugin}; +use amitoki_relay_postgres::PostgresPlugin; +use bytes::Bytes; +use serde_json::json; +use std::{collections::HashMap, error::Error, sync::Arc, time::Duration}; +use tokio::{sync::Barrier, task::JoinSet, time::Instant}; +use tokio_postgres::{Client, NoTls}; +use uuid::Uuid; + +type BenchError = Box; +// 既存のDBを誤って測定しないよう、試験専用の変数を必須にする。 +const CONNECTION_ENV: &str = "AMITOKI_BENCH_POSTGRES_URL"; +const RECEIVE_BATCH: usize = 128; +// 送信開始前の空ポーリングだけを抑える。配送待ちには全体の期限を適用する。 +const EMPTY_POLL_INTERVAL: Duration = Duration::from_millis(1); +const CASE_DEADLINE: Duration = Duration::from_secs(120); + +#[derive(Clone, Copy)] +struct Settings { + nodes: usize, + batch_size: usize, + frame_bytes: usize, + frames_per_node: usize, + warmup_frames: usize, + duration_seconds: u64, +} + +impl Settings { + fn from_arguments() -> Result { + let mut values: Vec = std::env::args().skip(1).map(|value| value.parse()).collect::>()?; + let duration_seconds = if values.len() == 6 { values.pop().unwrap() as u64 } else { 0 }; + if duration_seconds > 3600 { + return Err("耐久試験は最大3600秒です".into()); + } + let [nodes, batch_size, frame_bytes, frames_per_node, warmup_frames] = values.as_slice() else { + return Err("引数: ノード数 バッチ件数 フレーム長 ノードあたり送信件数 ウォームアップ件数 [耐久試験の秒数]".into()); + }; + if !(2..=8).contains(nodes) || !(1..=128).contains(batch_size) || !(14..=65535).contains(frame_bytes) || !(1..=100_000).contains(frames_per_node) || *warmup_frames > 10_000 + { + return Err("測定条件が範囲外です".into()); + } + Ok(Self { + nodes: *nodes, + batch_size: *batch_size, + frame_bytes: *frame_bytes, + frames_per_node: *frames_per_node, + warmup_frames: *warmup_frames, + duration_seconds, + }) + } +} + +fn frames(settings: &Settings, node: usize) -> Vec { + (0..settings.frames_per_node) + .map(|sequence| { + let payload: Vec = (0..settings.frame_bytes).map(|offset| (offset ^ sequence ^ (node << 4)) as u8).collect(); + Frame::new(Bytes::from(payload)).expect("検証済みのフレーム長") + }) + .collect() +} + +async fn publish(relay: Arc, batches: (Vec, usize), start: Arc) -> Result, BenchError> { + let (frames, batch_size) = batches; + let mut latencies = Vec::new(); + start.wait().await; + for batch in frames.chunks(batch_size) { + let started = Instant::now(); + relay.publish(batch).await?; + latencies.push(started.elapsed().as_secs_f64() * 1000.0); + } + Ok(latencies) +} + +async fn receive(relay: Arc, mut expected: HashMap, start: Arc) -> Result<(), BenchError> { + start.wait().await; + while !expected.is_empty() { + let deliveries = relay.receive(RECEIVE_BATCH).await?; + if deliveries.is_empty() { + tokio::time::sleep(EMPTY_POLL_INTERVAL).await; + continue; + } + let mut receipts = Vec::new(); + for delivery in deliveries { + let Some(bytes) = expected.remove(&delivery.frame.id) else { + return Err("不明またはACK済みのフレームを受信しました".into()); + }; + if bytes != delivery.frame.bytes { + return Err("受信内容が送信内容と一致しません".into()); + } + receipts.push(delivery.receipt); + } + relay.acknowledge(&receipts).await?; + } + Ok(()) +} + +async fn warm_up(relays: &[Arc], settings: &Settings, client: &Client) -> Result<(), BenchError> { + if settings.warmup_frames == 0 { + return Ok(()); + } + let warmup = Settings { + frames_per_node: settings.warmup_frames, + ..*settings + }; + let packets: Vec> = (0..settings.nodes).map(|node| frames(&warmup, node)).collect(); + for (relay, packets) in relays.iter().zip(&packets) { + for batch in packets.chunks(RECEIVE_BATCH) { + relay.publish(batch).await?; + } + } + // 空テーブルの統計に引きずられる初期状態と、統計更新後の測定を分ける。 + client.batch_execute("ANALYZE stegrdb_relay.frames; ANALYZE stegrdb_relay.pending").await?; + for (node, relay) in relays.iter().enumerate() { + receive(relay.clone(), expected_frames(&packets, node), Arc::new(Barrier::new(1))).await?; + } + Ok(()) +} + +fn expected_frames(packets: &[Vec], receiver: usize) -> HashMap { + packets.iter().enumerate().filter(|(sender, _)| *sender != receiver).flat_map(|(_, frames)| frames.iter().map(|frame| (frame.id, frame.bytes.clone()))).collect() +} + +async fn measure(settings: &Settings, channel: &str, client: &Client) -> Result { + let mut relays = Vec::new(); + let retention_ms: u64 = std::env::var("AMITOKI_BENCH_RETENTION_MS").unwrap_or_else(|_| "0".into()).parse()?; + for node in 0..settings.nodes { + relays.push( + PostgresPlugin + .connect( + RelayContext { + channel: channel.into(), + node_id: format!("node-{node}"), + }, + json!({"connection_env": CONNECTION_ENV, "replay_window_ms": 0, "retention_ms": retention_ms, "cleanup_interval_ms": 100}), + ) + .await?, + ); + } + warm_up(&relays, settings, client).await?; + if settings.duration_seconds == 0 { + return measure_round(settings, &relays).await; + } + let started = Instant::now(); + let mut rounds = Vec::new(); + let mut samples = Vec::new(); + // 大きな表の件数集計を毎バッチ走らせず、10秒ごとに状態を観測する。 + let sample_interval = Duration::from_secs(10); + let mut next_sample = Duration::ZERO; + while started.elapsed() < Duration::from_secs(settings.duration_seconds) { + rounds.push(tokio::time::timeout(CASE_DEADLINE, measure_round(settings, &relays)).await??); + if started.elapsed() >= next_sample { + let row = client.query_one("SELECT (SELECT count(*) FROM stegrdb_relay.frames WHERE channel=$1), (SELECT count(*) FROM stegrdb_relay.pending WHERE channel=$1), pg_total_relation_size('stegrdb_relay.frames')", &[&channel]).await?; + let sample = json!({"elapsed_seconds": started.elapsed().as_secs_f64(), "retained_frames": row.get::<_, i64>(0), "pending": row.get::<_, i64>(1), "frames_table_bytes": row.get::<_, i64>(2), "resident_kib": resident_kib()?}); + eprintln!("{sample}"); + samples.push(sample); + next_sample = started.elapsed() + sample_interval; + } + } + let published: u64 = rounds.iter().map(|round| round["published_frames"].as_u64().unwrap()).sum(); + Ok( + json!({"duration_seconds": started.elapsed().as_secs_f64(), "retention_ms": retention_ms, "published_frames": published, + "verified_deliveries": published * (settings.nodes as u64 - 1), "rounds": rounds, "samples": samples, "errors": 0}), + ) +} + +fn resident_kib() -> Result { + let status = std::fs::read_to_string("/proc/self/status")?; + let value = status.lines().find_map(|line| line.strip_prefix("VmRSS:")).ok_or("VmRSSがありません")?; + Ok(value.split_whitespace().next().ok_or("VmRSSの値がありません")?.parse()?) +} + +async fn measure_round(settings: &Settings, relays: &[Arc]) -> Result { + let packets: Vec> = (0..settings.nodes).map(|node| frames(settings, node)).collect(); + let start = Arc::new(Barrier::new(settings.nodes * 2 + 1)); + let mut senders = JoinSet::new(); + let mut receivers = JoinSet::new(); + for (node, relay) in relays.iter().enumerate() { + let expected = expected_frames(&packets, node); + receivers.spawn(receive(relay.clone(), expected, start.clone())); + senders.spawn(publish(relay.clone(), (packets[node].clone(), settings.batch_size), start.clone())); + } + let started = Instant::now(); + start.wait().await; + let mut publish_latency_ms = Vec::new(); + while let Some(completed) = senders.join_next().await { + publish_latency_ms.extend(completed??); + } + let publish_seconds = started.elapsed().as_secs_f64(); + while let Some(completed) = receivers.join_next().await { + completed??; + } + let elapsed_seconds = started.elapsed().as_secs_f64(); + for relay in relays { + if !relay.receive(RECEIVE_BATCH).await?.is_empty() { + return Err("照合後も余分な配送が残っています".into()); + } + } + publish_latency_ms.sort_by(f64::total_cmp); + let percentile = |percent: usize| publish_latency_ms[(publish_latency_ms.len() * percent).div_ceil(100).saturating_sub(1)]; + let frame_count = settings.nodes * settings.frames_per_node; + Ok(json!({ + "nodes": settings.nodes, "batch_size": settings.batch_size, "frame_bytes": settings.frame_bytes, + "frames_per_node": settings.frames_per_node, "published_frames": frame_count, + "warmup_frames_per_node": settings.warmup_frames, + "verified_deliveries": frame_count * (settings.nodes - 1), + "publish_seconds": publish_seconds, "elapsed_seconds": elapsed_seconds, + "published_frames_per_second": frame_count as f64 / publish_seconds, + "completed_frames_per_second": frame_count as f64 / elapsed_seconds, + "publish_batch_latency_ms": {"p50": percentile(50), "p95": percentile(95), "p99": percentile(99)}, + "errors": 0 + })) +} + +#[tokio::main] +async fn main() -> Result<(), BenchError> { + let settings = Settings::from_arguments()?; + let channel = format!("bench-{}", Uuid::new_v4()); + let (client, connection) = tokio_postgres::connect(&std::env::var(CONNECTION_ENV)?, NoTls).await?; + let connection_task = tokio::spawn(connection); + let measured = tokio::time::timeout(CASE_DEADLINE + Duration::from_secs(settings.duration_seconds), measure(&settings, &channel, &client)).await; + let retained: i64 = client.query_one("SELECT count(*) FROM stegrdb_relay.frames WHERE channel=$1", &[&channel]).await?.get(0); + let pending: i64 = client.query_one("SELECT count(*) FROM stegrdb_relay.pending WHERE channel=$1", &[&channel]).await?.get(0); + // 例外時も、この測定が作ったchannelだけを掃除する。 + client.execute("DELETE FROM stegrdb_relay.nodes WHERE channel=$1", &[&channel]).await?; + client.execute("DELETE FROM stegrdb_relay.frames WHERE channel=$1", &[&channel]).await?; + connection_task.abort(); + let mut report = measured??; + report["retained_frames_after_ack"] = json!(retained); + report["pending_after_ack"] = json!(pending); + println!("{report}"); + Ok(()) +} diff --git a/schema.sql b/schema.sql index b8588d0..d2d034d 100644 --- a/schema.sql +++ b/schema.sql @@ -7,6 +7,9 @@ CREATE TABLE IF NOT EXISTS stegrdb_relay.nodes ( node_id TEXT NOT NULL, PRIMARY KEY (channel, node_id) ); +-- 旧ノードの設定が不明な間は削除を無効にする。再接続時に各ノードが更新する。 +ALTER TABLE stegrdb_relay.nodes ADD COLUMN IF NOT EXISTS replay_window_ms BIGINT NOT NULL DEFAULT 2147483647; +ALTER TABLE stegrdb_relay.nodes ADD COLUMN IF NOT EXISTS retention_ms BIGINT NOT NULL DEFAULT 0; CREATE TABLE IF NOT EXISTS stegrdb_relay.frames ( channel TEXT NOT NULL, @@ -29,3 +32,20 @@ CREATE TABLE IF NOT EXISTS stegrdb_relay.pending ( FOREIGN KEY (channel, frame_id) REFERENCES stegrdb_relay.frames ON DELETE CASCADE ); CREATE INDEX IF NOT EXISTS pending_delivery_order ON stegrdb_relay.pending (channel, node_id, position); +-- node_idをまたぐ未ACK確認とframes削除時の外部キー検査に使う。 +CREATE INDEX IF NOT EXISTS pending_frame ON stegrdb_relay.pending (channel, frame_id); + +-- 高頻度のINSERT/DELETEで古い行・統計が蓄積しないよう、中継表だけを調整する。 +-- 末尾ページの縮小はACCESS EXCLUSIVEを取るため、オンラインでは領域の再利用を優先する。 +ALTER TABLE stegrdb_relay.pending SET ( + autovacuum_vacuum_scale_factor = 0.02, + autovacuum_analyze_scale_factor = 0.02, + autovacuum_vacuum_cost_limit = 1000, + vacuum_truncate = false +); +ALTER TABLE stegrdb_relay.frames SET ( + autovacuum_vacuum_scale_factor = 0.02, + autovacuum_analyze_scale_factor = 0.02, + autovacuum_vacuum_cost_limit = 1000, + vacuum_truncate = false +); diff --git a/scripts/bench-postgres.py b/scripts/bench-postgres.py new file mode 100644 index 0000000..e09b1c3 --- /dev/null +++ b/scripts/bench-postgres.py @@ -0,0 +1,153 @@ +#!/usr/bin/env python3 +"""一時コンテナで実配送を測り、任意の比較用バイナリとSQL実行計画を記録する。""" + +import argparse +from datetime import datetime +import hashlib +from itertools import product +import json +import os +from pathlib import Path +import platform +import subprocess +import time + +from postgres_query_plans import measure_query_plans + +ROOT = Path(__file__).resolve().parents[1] +READY_TIMEOUT_SECONDS = 60 +READY_POLL_SECONDS = 1 +READINESS_COMMAND_TIMEOUT_SECONDS = 10 +COMMAND_TIMEOUT_SECONDS = 180 +# 3ノードの比較を一定条件で行い、ホストの資源を使い切らない。 +DATABASE_CPUS = 4 +DATABASE_MEMORY_MIB = 2048 +# 小さいバッチでも秒単位で測り、大きいバッチでは配送件数を増やす。 +FRAMES_PER_NODE = {1: 1024, 32: 4096, 128: 4096} + + +def command(arguments, **options): + return subprocess.check_output(arguments, text=True, timeout=COMMAND_TIMEOUT_SECONDS, **options).strip() + + +def sql(container, statement): + return command(["docker", "exec", "-i", container, "psql", "-U", "postgres", "-At", "-v", "ON_ERROR_STOP=1"], input=statement) + + +def await_database(container): + deadline = time.monotonic() + READY_TIMEOUT_SECONDS + while time.monotonic() < deadline: + # 初期化中のUnix socketではなく、本稼働のTCP待受を確認する。 + ready = subprocess.run(["docker", "exec", container, "pg_isready", "-h", "127.0.0.1", "-U", "postgres"], + stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=READINESS_COMMAND_TIMEOUT_SECONDS) + if ready.returncode == 0: + return + time.sleep(READY_POLL_SECONDS) + raise TimeoutError("PostgreSQLの起動が完了しません") + + +def benchmark(container, settings, *, report, destination): + binaries = {"current": settings.binary.resolve()} + if settings.compare_binary: + binaries["candidate"] = settings.compare_binary.resolve() + report["binaries"] = {name: {"path": str(path), "sha256": hashlib.sha256(path.read_bytes()).hexdigest()} for name, path in binaries.items()} + if len({details["sha256"] for details in report["binaries"].values()}) != len(binaries): + raise ValueError("比較元と候補の実行ファイルが同一です。ビルド先を分けてください") + port = command(["docker", "port", container, "5432/tcp"]).rsplit(":", 1)[1] + environment = dict(os.environ, AMITOKI_BENCH_POSTGRES_URL=f"host=127.0.0.1 port={port} user=postgres dbname=postgres sslmode=disable") + for repetition, nodes, frame_bytes, batch in product(range(settings.repetitions), (2, 3), (64, 1400), FRAMES_PER_NODE): + frames = FRAMES_PER_NODE[batch] + order = list(binaries.items()) + if repetition % 2: + order.reverse() + for name, binary in order: + # 前の条件の行・統計を引き継がず、試験DB内だけを初期化する。 + sql(container, "TRUNCATE stegrdb_relay.pending, stegrdb_relay.frames, stegrdb_relay.nodes RESTART IDENTITY; " + "ANALYZE stegrdb_relay.frames; ANALYZE stegrdb_relay.pending; CHECKPOINT;") + measured = json.loads(command([str(binary), str(nodes), str(batch), str(frame_bytes), str(frames), str(settings.warmup_frames)], env=environment)) + assert measured["verified_deliveries"] == nodes * frames * (nodes - 1), measured + assert measured["pending_after_ack"] == 0 and measured["errors"] == 0, measured + report["cases"].append({"variant": name, "repetition": repetition, **measured}) + save_report(destination, report) + print(f"{name} repeat={repetition + 1} nodes={nodes} bytes={frame_bytes} batch={batch}: " + f"{measured['completed_frames_per_second']:.0f} frames/s", flush=True) + + +def save_report(destination, report): + destination.write_text(json.dumps(report, indent=2) + "\n") + + +def soak(container, settings): + port = command(["docker", "port", container, "5432/tcp"]).rsplit(":", 1)[1] + environment = dict(os.environ, AMITOKI_BENCH_POSTGRES_URL=f"host=127.0.0.1 port={port} user=postgres dbname=postgres sslmode=disable", + AMITOKI_BENCH_RETENTION_MS=str(settings.retention_ms)) + binary = settings.binary.resolve() + binary_sha256 = hashlib.sha256(binary.read_bytes()).hexdigest() + measured = subprocess.check_output([str(binary), "3", "128", "1400", "4096", str(settings.warmup_frames), str(settings.soak_seconds)], + text=True, env=environment, timeout=settings.soak_seconds + COMMAND_TIMEOUT_SECONDS) + measured = json.loads(measured) + assert measured["errors"] == 0 and measured["pending_after_ack"] == 0, measured + assert measured["verified_deliveries"] == 2 * measured["published_frames"], measured + return {"binary_sha256": binary_sha256, **measured} + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--binary", type=Path, default=ROOT / "target/release/examples/relay_bench") + parser.add_argument("--compare-binary", type=Path) + parser.add_argument("--repetitions", type=int, choices=range(1, 11), default=3) + parser.add_argument("--warmup-frames", type=int, default=1024) + parser.add_argument("--plans-only", action="store_true", help="受信と初回再生のSQLだけを比較する") + parser.add_argument("--baseline-queries", type=Path, help="比較元のsrc/queries.rs。指定時は現行SQLと結果/実行計画を比較する") + parser.add_argument("--soak-seconds", type=int, default=0, help="3ノード/1400B/batch128で接続を維持して耐久試験する（最大3600秒）") + parser.add_argument("--retention-ms", type=int, default=0, help="耐久試験で使う保持期間。0で削除無効") + parser.add_argument("--image", default="postgres:17-alpine") + parser.add_argument("--output", type=Path, required=True) + settings = parser.parse_args() + if not 0 <= settings.warmup_frames <= 10_000: + parser.error("--warmup-framesは0〜10000件で指定してください") + if not 0 <= settings.soak_seconds <= 3600 or not 0 <= settings.retention_ms <= 2_147_483_647: + parser.error("耐久試験の秒数または保持期間が範囲外です") + if settings.soak_seconds and (settings.compare_binary or settings.plans_only): + parser.error("耐久試験と比較/SQLのみの測定は別々に実行してください") + if settings.retention_ms and not settings.soak_seconds: + parser.error("--retention-msは耐久試験で指定してください") + settings.output.parent.mkdir(parents=True, exist_ok=True) + report = {"started_at": datetime.now().astimezone().isoformat(), "status": "running", "cases": [], + "host": {"machine": platform.machine(), "logical_cpus": os.cpu_count()}, + "database_limits": {"cpus": DATABASE_CPUS, "memory_bytes": DATABASE_MEMORY_MIB * 1024**2}, + "plugin_head": command(["git", "-C", str(ROOT), "rev-parse", "HEAD"])} + container = command(["docker", "run", "--detach", "--rm", "--publish", "127.0.0.1::5432", + "--cpus", str(DATABASE_CPUS), "--memory", f"{DATABASE_MEMORY_MIB}m", "--env", "POSTGRES_HOST_AUTH_METHOD=trust", + settings.image, "-c", "shared_preload_libraries=pg_stat_statements"]) + try: + await_database(container) + report["image"] = command(["docker", "inspect", "--format", "{{.Image}}", container]) + sql(container, (ROOT / "schema.sql").read_text() + "\nCREATE EXTENSION pg_stat_statements;") + report["database"] = json.loads(sql(container, "SELECT json_build_object('version',version(),'fsync',current_setting('fsync')," + "'synchronous_commit',current_setting('synchronous_commit'),'shared_buffers',current_setting('shared_buffers')," + "'autovacuum_naptime',current_setting('autovacuum_naptime'))")) + report["table_settings"] = json.loads(sql(container, "SELECT json_agg(json_build_object('table',relname,'options',reloptions)) " + "FROM pg_class WHERE relnamespace='stegrdb_relay'::regnamespace AND relkind='r'")) + if settings.soak_seconds: + report["soak"] = soak(container, settings) + elif not settings.plans_only: + benchmark(container, settings, report=report, destination=settings.output) + report["query_plans"] = measure_query_plans(lambda statement: sql(container, statement), ROOT, baseline=settings.baseline_queries) + report["sql_statistics"] = json.loads(sql(container, "SELECT coalesce(json_agg(summary),'[]'::json) FROM " + "(SELECT query,calls,total_exec_time,rows,shared_blks_hit,shared_blks_read,wal_bytes " + "FROM pg_stat_statements ORDER BY total_exec_time DESC LIMIT 15) AS summary")) + report["status"] = "passed" + except BaseException as error: + report["status"] = "interrupted" if isinstance(error, KeyboardInterrupt) else "failed" + report["error"] = str(error) + raise + finally: + report["finished_at"] = datetime.now().astimezone().isoformat() + save_report(settings.output, report) + # 匿名volumeも消し、この試験が作ったコンテナだけを片付ける。 + command(["docker", "rm", "--force", "--volumes", container]) + + +if __name__ == "__main__": + main() diff --git a/scripts/postgres_query_plans.py b/scripts/postgres_query_plans.py new file mode 100644 index 0000000..baeacfd --- /dev/null +++ b/scripts/postgres_query_plans.py @@ -0,0 +1,98 @@ +"""受信と初回再生のSQLを、初期統計・更新後の統計・蓄積済み履歴で比較する。""" + +import json +import re + +HISTORY_FRAMES = 250_000 +RECEIVE_FRAMES = 10_000 +PLAN_REPETITIONS = 3 +RECEIVE_LIMIT = 128 + + +def query(source, name, parameters): + statement = re.search(rf'const {name}: &str = "(.*?)";', source, re.DOTALL).group(1) + for position, value in enumerate(parameters, 1): + statement = statement.replace(f"${position}", value) + return statement + + +def compare_plans(sql, candidates): + plans = {name: [] for name in candidates} + for repetition in range(PLAN_REPETITIONS): + order = list(candidates.items()) + if repetition % 2: + order.reverse() + for name, statement in order: + plans[name].append(json.loads(sql(f"EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON) {statement}"))[0]) + return plans + + +def frame_rows_examined(plan): + rows = 0 + if plan.get("Relation Name") == "frames": + rows = (plan["Actual Rows"] + plan.get("Rows Removed by Filter", 0)) * plan["Actual Loops"] + return rows + sum(frame_rows_examined(child) for child in plan.get("Plans", [])) + + +def check_receive_work(plans): + # 時間の閾値ではなく、128件のために履歴を128回全走査する回帰を検出する。 + for execution in plans: + examined = frame_rows_examined(execution["Plan"]) + assert examined <= RECEIVE_FRAMES + RECEIVE_LIMIT, f"受信が履歴を繰り返し走査しました: {examined}行" + + +def measure_receive(sql, sources): + # 自動ANALYZEとの競合を避け、この診断中だけ試験テーブルを手動管理する。 + sql(f""" + ALTER TABLE stegrdb_relay.frames SET (autovacuum_enabled=false); + ALTER TABLE stegrdb_relay.pending SET (autovacuum_enabled=false); + TRUNCATE stegrdb_relay.pending,stegrdb_relay.frames,stegrdb_relay.nodes RESTART IDENTITY; + INSERT INTO stegrdb_relay.frames(channel,id,sender,payload) + SELECT 'previous-channel',md5(value::text)::uuid,'source',decode(repeat('ab',1400),'hex') + FROM generate_series(1,{RECEIVE_FRAMES}) AS value; + ANALYZE stegrdb_relay.frames; + TRUNCATE stegrdb_relay.pending,stegrdb_relay.frames,stegrdb_relay.nodes RESTART IDENTITY; + ANALYZE stegrdb_relay.frames; ANALYZE stegrdb_relay.pending; + INSERT INTO stegrdb_relay.nodes VALUES ('bench-receive','receiver'); + INSERT INTO stegrdb_relay.frames(channel,id,sender,payload) + SELECT 'bench-receive', md5(value::text)::uuid, 'source', decode(repeat('ab',64),'hex') + FROM generate_series(1,{RECEIVE_FRAMES}) AS value; + INSERT INTO stegrdb_relay.pending + SELECT channel,'receiver',id,position FROM stegrdb_relay.frames ORDER BY position DESC LIMIT {RECEIVE_LIMIT}; + """) + candidates = {name: query(source, "RECEIVE", ["'bench-receive'", "'receiver'", str(RECEIVE_LIMIT)]) for name, source in sources.items()} + # 受領情報の世代列が増えても、配送するUUID・本文・順序を比較する。 + def frames(statement): + return [line.split("|")[:2] for line in sql(statement).splitlines()] + + received = frames(candidates["current"]) + assert len(received) == RECEIVE_LIMIT, "受信SQLが期待した128件を返しません" + for statement in candidates.values(): + assert received == frames(statement), "受信SQLの比較元と結果が変わりました" + cold = compare_plans(sql, candidates) + check_receive_work(cold["current"]) + sql("ANALYZE stegrdb_relay.frames; ANALYZE stegrdb_relay.pending;") + analyzed = compare_plans(sql, candidates) + check_receive_work(analyzed["current"]) + sql("ALTER TABLE stegrdb_relay.frames RESET (autovacuum_enabled); ALTER TABLE stegrdb_relay.pending RESET (autovacuum_enabled);") + return {"frames": RECEIVE_FRAMES, "pending": RECEIVE_LIMIT, "equivalent_output": True, "bounded_work": True, "initial_statistics": cold, "updated_statistics": analyzed} + + +def measure_replay(sql, sources): + sql(f""" + INSERT INTO stegrdb_relay.nodes VALUES ('bench-history','late'); + INSERT INTO stegrdb_relay.frames(channel,id,sender,payload,created_at) + SELECT 'bench-history',md5(value::text)::uuid,'source',decode(repeat('ab',64),'hex'), + statement_timestamp()-interval '1 day' + FROM generate_series(1,{HISTORY_FRAMES}) AS value; + ANALYZE stegrdb_relay.frames; + """) + candidates = {name: query(source, "REPLAY", ["'bench-history'", "'late'", "4000"]) for name, source in sources.items()} + return {"history_frames": HISTORY_FRAMES, "plans": compare_plans(sql, candidates)} + + +def measure_query_plans(sql, root, *, baseline=None): + sources = {"current": (root / "src/queries.rs").read_text()} + if baseline: + sources["baseline"] = baseline.read_text() + return {"receive": measure_receive(sql, sources), "replay": measure_replay(sql, sources)} diff --git a/scripts/publish-release.py b/scripts/publish-release.py new file mode 100644 index 0000000..a357391 --- /dev/null +++ b/scripts/publish-release.py @@ -0,0 +1,105 @@ +#!/usr/bin/env python3 +"""検証済みのプラグイン・manifest・SQLを版タグから公開する。""" +import argparse +import hashlib +import json +from pathlib import Path +import subprocess +import tomllib + +ROOT = Path(__file__).resolve().parents[1] +# 現在のプラグインCIがネイティブでビルド・検証するCPU。 +TARGET = "x86_64-unknown-linux-gnu" + + +def output(*arguments): + return subprocess.check_output(arguments, cwd=ROOT, text=True).strip() + + +def digest(path): + with path.open("rb") as stream: + return hashlib.file_digest(stream, "sha256").hexdigest() + + +def verify_package(directory, *, release_version, schema): + manifest_path = directory / f"plugin-{TARGET}.json" + binary = directory / f"amitoki-plugin-postgres-{TARGET}" + package = json.loads(manifest_path.read_text()) + if package["target"] != TARGET or package["binary"] != "amitoki-plugin-postgres": + raise ValueError("配布物のCPUまたは実行ファイル名が違います") + manifest = package["manifest"] + if manifest["name"] != "postgres" or manifest["version"] != release_version: + raise ValueError("manifestと公開する版が一致しません") + if digest(binary) != package["sha256"]: + raise ValueError("実行ファイルのSHA256が一致しません") + # ActionsのArtifactから復元すると実行権限が失われる。 + binary.chmod(0o755) + if json.loads(output(str(binary), "--describe")) != manifest: + raise ValueError("実行ファイルとmanifestが一致しません") + packaged_schema = directory / "schema.sql" + if packaged_schema.read_bytes() != schema.read_bytes(): + raise ValueError("配布するSQLとタグのSQLが一致しません") + embedded_schema = subprocess.check_output([str(binary), "--schema"]) + if embedded_schema != schema.read_bytes(): + raise ValueError("実行ファイルに埋め込まれたSQLが一致しません") + assets = [binary, manifest_path, packaged_schema] + checksums = directory / "SHA256SUMS" + checksums.write_text("".join(f"{digest(path)} {path.name}\n" for path in assets)) + return [*assets, checksums] + + +def release_preflight(tag): + release_version = tomllib.loads((ROOT / "Cargo.toml").read_text())["package"]["version"] + if tag != f"v{release_version}": + raise ValueError("タグとCargo.tomlの版が一致しません") + commit = output("git", "rev-parse", "HEAD") + if output("git", "rev-parse", f"refs/tags/{tag}^{{commit}}") != commit: + raise ValueError("チェックアウトしたcommitとタグが一致しません") + subprocess.run(["git", "merge-base", "--is-ancestor", commit, "origin/main"], cwd=ROOT, check=True) + notes = ROOT / f"docs/releases/{tag}.md" + if not notes.is_file(): + raise ValueError("リリースノートがありません") + return release_version, notes + + +def verify_uploaded(release, *, assets, notes): + expected = {path.name: f"sha256:{digest(path)}" for path in assets} + uploaded = {asset["name"]: asset["digest"] for asset in release["assets"]} + if uploaded != expected or release["body"].strip() != notes.read_text().strip(): + raise ValueError("GitHub上の配布物またはノートが一致しません。公開済み版は上書きしません") + + +def publish(directory, *, tag, repository): + release_version, notes = release_preflight(tag) + assets = verify_package(directory, release_version=release_version, schema=ROOT / "schema.sql") + base = ["gh", "release"] + view = [*base, "view", tag, "--repo", repository, "--json", "isDraft,body,assets"] + existing = subprocess.run(view, capture_output=True, text=True, check=False) + if existing.returncode == 0: + release = json.loads(existing.stdout) + if not release["isDraft"]: + verify_uploaded(release, assets=assets, notes=notes) + print("同じ内容が公開済みです。変更しません") + return + subprocess.run([*base, "edit", tag, "--repo", repository, "--title", tag, + "--notes-file", str(notes)], check=True) + else: + subprocess.run([*base, "create", tag, "--repo", repository, "--verify-tag", "--draft", + "--title", tag, "--notes-file", str(notes)], check=True) + # 完了前のdraftにだけ再アップロードし、全ファイルの一致後に公開する。 + subprocess.run([*base, "upload", tag, "--repo", repository, "--clobber", *map(str, assets)], check=True) + verify_uploaded(json.loads(output(*view)), assets=assets, notes=notes) + subprocess.run([*base, "edit", tag, "--repo", repository, "--draft=false", "--latest"], check=True) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--directory", type=Path, required=True) + parser.add_argument("--tag", required=True) + parser.add_argument("--repository", required=True) + arguments = parser.parse_args() + publish(arguments.directory.resolve(), tag=arguments.tag, repository=arguments.repository) + + +if __name__ == "__main__": + main() diff --git a/scripts/test-postgres.sh b/scripts/test-postgres.sh index 79aa856..977b173 100644 --- a/scripts/test-postgres.sh +++ b/scripts/test-postgres.sh @@ -5,18 +5,44 @@ cd "$(dirname "$0")/.." # 初回のDB起動を最大60秒まで、readinessが返るまで待つ。 READY_ATTEMPTS=60 READY_INTERVAL_SECONDS=1 +POSTGRES_IMAGE=${POSTGRES_IMAGE:-postgres:17-alpine} # ランダムなホスト側ポートと一時コンテナで既存DBから隔離する。 container=$(docker run --detach --rm --publish 127.0.0.1::5432 \ - --env POSTGRES_HOST_AUTH_METHOD=trust postgres:17-alpine) -trap 'docker stop "$container" >/dev/null' EXIT + --env POSTGRES_HOST_AUTH_METHOD=trust "$POSTGRES_IMAGE") +trap 'docker rm --force --volumes "$container" >/dev/null' EXIT # 初期化用の一時サーバはUnix socketだけで待受するため、本稼働のTCPを確認する。 for ((attempt=0; attempt/dev/null 2>&1; then break; fi sleep "$READY_INTERVAL_SECONDS" done docker exec "$container" pg_isready --host 127.0.0.1 --username postgres >/dev/null +# 旧スキーマのデータと未ACKが更新で消えないことを先に確認する。 +docker exec -i "$container" psql --username postgres --set ON_ERROR_STOP=1 < tests/fixtures/schema-v0.1.2.sql +docker exec -i "$container" psql --username postgres --set ON_ERROR_STOP=1 <<'SQL' +INSERT INTO stegrdb_relay.nodes VALUES ('schema-upgrade', 'receiver'); +INSERT INTO stegrdb_relay.frames(channel,id,sender,payload) +VALUES ('schema-upgrade',md5('schema-upgrade')::uuid,'sender',decode(repeat('ab',64),'hex')); +INSERT INTO stegrdb_relay.pending +SELECT channel,'receiver',id,position FROM stegrdb_relay.frames WHERE channel='schema-upgrade'; +SQL docker exec -i "$container" psql --username postgres --set ON_ERROR_STOP=1 < schema.sql port=$(docker port "$container" 5432/tcp | cut -d: -f2) export AMITOKI_TEST_POSTGRES_URL="host=127.0.0.1 port=$port user=postgres dbname=postgres sslmode=disable" -cargo test --test delivery --locked -- --ignored --nocapture +# schema.sqlは既存データを維持したまま再適用できる必要がある。 +docker exec -i "$container" psql --username postgres --set ON_ERROR_STOP=1 < schema.sql +docker exec -i "$container" psql --username postgres --set ON_ERROR_STOP=1 <<'SQL' +DO $$ BEGIN + IF (SELECT count(*) FROM stegrdb_relay.pending WHERE channel='schema-upgrade') <> 1 + OR (SELECT encode(payload,'hex') FROM stegrdb_relay.frames WHERE channel='schema-upgrade') IS DISTINCT FROM repeat('ab',64) + OR NOT EXISTS (SELECT 1 FROM stegrdb_relay.nodes WHERE channel='schema-upgrade' AND retention_ms=0 AND replay_window_ms=2147483647) + THEN RAISE EXCEPTION 'Schema update changed existing delivery or retention'; END IF; +END $$; +DELETE FROM stegrdb_relay.nodes WHERE channel='schema-upgrade'; +DELETE FROM stegrdb_relay.frames WHERE channel='schema-upgrade'; +SQL +cargo test --tests --locked -- --ignored --nocapture +# DBの既定が強い分離レベルでも、登録/送信はロック後の最新状態を参照する。 +docker exec "$container" psql --username postgres --set ON_ERROR_STOP=1 \ + --command "ALTER DATABASE postgres SET default_transaction_isolation TO 'repeatable read'" +cargo test --test concurrency --locked -- --ignored --nocapture diff --git a/src/lib.rs b/src/lib.rs index f895d90..d75cb37 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,3 +1,4 @@ +mod maintenance; pub mod manifest; mod options; mod queries; @@ -11,6 +12,7 @@ use options::Options; use serde_json::Value; use session::NodeSession; use std::{error::Error, io, sync::Arc}; +use tokio_postgres::IsolationLevel; use tokio_postgres_rustls::MakeRustlsConnect; use uuid::Uuid; @@ -20,6 +22,7 @@ struct PostgresRelay { context: RelayContext, pool: Pool, session: NodeSession, + maintenance: Option, } fn database_error(error: tokio_postgres::Error) -> RelayError { @@ -33,6 +36,14 @@ fn database_error(error: tokio_postgres::Error) -> RelayError { } } +fn pool_error(error: PoolError) -> RelayError { + match error { + PoolError::Backend(error) => database_error(error), + PoolError::Timeout(_) => RelayError::retryable("PostgreSQL接続プールの待機がタイムアウトしました"), + _ => RelayError::permanent("PostgreSQL接続プールから接続を取得できません"), + } +} + fn is_connection_interrupted(error: &tokio_postgres::Error) -> bool { if error.is_closed() { return true; @@ -80,8 +91,14 @@ impl RelayPlugin for PostgresPlugin { let tls = tls_connector()?; let session = NodeSession::connect(&config, tls.clone(), &context).await?; let pool = Pool::builder(Manager::new(config, tls)).max_size(options.max_connections).build().map_err(|_| RelayError::permanent("PostgreSQL接続プールを作成できません"))?; - let relay = PostgresRelay { context, pool, session }; - relay.register(options.replay_window_ms as i32).await?; + let mut relay = PostgresRelay { + context, + pool, + session, + maintenance: None, + }; + relay.register(&options).await?; + relay.maintenance = maintenance::Maintenance::start(relay.pool.clone(), relay.context.channel.clone(), &options); Ok(Arc::new(relay)) } } @@ -89,20 +106,24 @@ impl RelayPlugin for PostgresPlugin { impl PostgresRelay { async fn acquire_client(&self) -> Result { self.session.ensure_connected()?; - self.pool.get().await.map_err(|error| match error { - PoolError::Backend(error) => database_error(error), - PoolError::Timeout(_) => RelayError::retryable("PostgreSQL接続プールの待機がタイムアウトしました"), - _ => RelayError::permanent("PostgreSQL接続プールから接続を取得できません"), - }) + self.pool.get().await.map_err(pool_error) } - async fn register(&self, replay_window_ms: i32) -> Result<(), RelayError> { + async fn register(&self, options: &Options) -> Result<(), RelayError> { let mut client = self.acquire_client().await?; - let transaction = client.transaction().await.map_err(database_error)?; + // ロック待ち後のSQLで最新の登録/送信を見る。DB側の既定分離レベルに依存しない。 + let transaction = client.build_transaction().isolation_level(IsolationLevel::ReadCommitted).start().await.map_err(database_error)?; // 登録と送信が交差して、どちらからもpendingが作られない競合を防ぐ。 transaction.query_one(queries::LOCK_CHANNEL, &[&self.context.channel]).await.map_err(database_error)?; - let inserted = transaction.execute(queries::REGISTER, &[&self.context.channel, &self.context.node_id]).await.map_err(database_error)?; - if inserted != 0 { + let replay_window_ms = options.replay_window_ms as i32; + let replay_window_db = options.replay_window_ms as i64; + let retention_ms = options.retention_ms as i64; + let parameters: &[&(dyn tokio_postgres::types::ToSql + Sync)] = &[&self.context.channel, &self.context.node_id, &replay_window_db, &retention_ms]; + let inserted = transaction.execute(queries::REGISTER, parameters).await.map_err(database_error)?; + if inserted == 0 { + transaction.execute(queries::UPDATE_NODE_OPTIONS, parameters).await.map_err(database_error)?; + } + if inserted != 0 && replay_window_ms > 0 { transaction.execute(queries::REPLAY, &[&self.context.channel, &self.context.node_id, &replay_window_ms]).await.map_err(database_error)?; } transaction.commit().await.map_err(database_error) @@ -121,8 +142,10 @@ impl Relay for PostgresRelay { let ids: Vec<_> = frames.iter().map(|frame| frame.id).collect(); let payloads: Vec<&[u8]> = frames.iter().map(|frame| frame.bytes.as_ref()).collect(); let mut client = self.acquire_client().await?; - let transaction = client.transaction().await.map_err(database_error)?; - let lock = transaction.prepare_cached(queries::LOCK_CHANNEL).await.map_err(database_error)?; + let transaction = client.build_transaction().isolation_level(IsolationLevel::ReadCommitted).start().await.map_err(database_error)?; + // 送信同士は並行実行し、登録とは排他にする。ロックとINSERTは別SQLにして + // 待機中に完了した登録を次のREAD COMMITTEDスナップショットへ反映する。 + let lock = transaction.prepare_cached(queries::LOCK_CHANNEL_SHARED).await.map_err(database_error)?; transaction.query_one(&lock, &[&self.context.channel]).await.map_err(database_error)?; let statement = transaction.prepare_cached(queries::PUBLISH).await.map_err(database_error)?; transaction.execute(&statement, &[&self.context.channel, &self.context.node_id, &ids, &payloads]).await.map_err(database_error)?; @@ -138,9 +161,11 @@ impl Relay for PostgresRelay { .map(|row| { let id: Uuid = row.try_get("id").map_err(database_error)?; let payload: Vec = row.try_get("payload").map_err(database_error)?; + let position: i64 = row.try_get("position").map_err(database_error)?; Ok(Delivery { frame: Frame { id, bytes: Bytes::from(payload) }, - receipt: Receipt(id.to_string()), + // 期限切れ削除後の同一UUIDを、遅延した古いACKから区別する。 + receipt: Receipt(format!("{id}:{position}")), }) }) .collect() @@ -150,11 +175,22 @@ impl Relay for PostgresRelay { if receipts.is_empty() { return Ok(()); } - let ids: Result, _> = receipts.iter().map(|receipt| Uuid::parse_str(&receipt.0)).collect(); - let ids = ids.map_err(|_| RelayError::permanent("postgresの受領情報が不正です"))?; + let mut ids = Vec::with_capacity(receipts.len()); + let mut positions = Vec::with_capacity(receipts.len()); + for receipt in receipts { + let (id, position) = parse_receipt(receipt).ok_or_else(|| RelayError::permanent("postgresの受領情報が不正です"))?; + ids.push(id); + positions.push(position); + } let client = self.acquire_client().await?; let statement = client.prepare_cached(queries::ACKNOWLEDGE).await.map_err(database_error)?; - client.execute(&statement, &[&self.context.channel, &self.context.node_id, &ids]).await.map_err(database_error)?; + client.execute(&statement, &[&self.context.channel, &self.context.node_id, &ids, &positions]).await.map_err(database_error)?; Ok(()) } } + +fn parse_receipt(receipt: &Receipt) -> Option<(Uuid, i64)> { + let (id, position) = receipt.0.split_once(':')?; + let position: i64 = position.parse().ok()?; + (position > 0).then_some((Uuid::parse_str(id).ok()?, position)) +} diff --git a/src/maintenance.rs b/src/maintenance.rs new file mode 100644 index 0000000..f625af4 --- /dev/null +++ b/src/maintenance.rs @@ -0,0 +1,128 @@ +use crate::{database_error, options::Options, pool_error}; +use amitoki_relay::RelayError; +use deadpool_postgres::Pool; +use std::time::Duration; +use tokio::task::JoinHandle; +use tokio_postgres::IsolationLevel; + +// 掃除が統計不足や大量の未ACKに遭遇しても、接続と登録を占有し続けない。 +const CLEANUP_STATEMENT_TIMEOUT: &str = "SET LOCAL statement_timeout = '250ms'"; +const RETRY_INTERVAL: Duration = Duration::from_secs(60); +// 一時的なVACUUM競合などの後は、削除量を減らして早めに追従する。 +const TIMEOUT_RETRY_INTERVAL: Duration = Duration::from_secs(1); +const FAST_BATCHES_BEFORE_GROWTH: usize = 16; +const FAST_BATCH_DURATION: Duration = Duration::from_millis(50); +const TRY_CHANNEL_LOCK: &str = "SELECT pg_try_advisory_xact_lock_shared(hashtextextended($1, 0))"; +// 1ノードでも無効を選んでいれば削除しない。停止中ノードの設定も保護する。 +const CHANNEL_RETENTION: &str = " +SELECT CASE WHEN bool_and(retention_ms > 0) + THEN greatest(max(retention_ms), max(replay_window_ms)) ELSE 0 END +FROM stegrdb_relay.nodes WHERE channel = $1"; +// 相関NOT EXISTSのOFFSET 0は、古い統計でpending全体とのanti joinへ +// 展開されることを防ぐ。候補ごとにpending_frameの索引を調べる。 +const PRUNE: &str = " +WITH expired AS MATERIALIZED ( + SELECT frames.channel, frames.id FROM stegrdb_relay.frames AS frames + WHERE frames.channel = $1 + AND frames.created_at < statement_timestamp() - $2::bigint * INTERVAL '1 millisecond' + AND NOT EXISTS ( + SELECT 1 FROM stegrdb_relay.pending + WHERE pending.channel = frames.channel AND pending.frame_id = frames.id + OFFSET 0 + ) + ORDER BY frames.created_at LIMIT $3 + FOR UPDATE OF frames SKIP LOCKED +) +DELETE FROM stegrdb_relay.frames AS frames USING expired +WHERE frames.channel = expired.channel AND frames.id = expired.id"; + +pub(crate) struct Maintenance(JoinHandle<()>); + +#[derive(Debug)] +enum CleanupError { + TimedOut, + Database(RelayError), +} + +impl From for CleanupError { + fn from(error: tokio_postgres::Error) -> Self { + if error.code() == Some(&tokio_postgres::error::SqlState::QUERY_CANCELED) { + Self::TimedOut + } else { + Self::Database(database_error(error)) + } + } +} + +impl Maintenance { + pub fn start(pool: Pool, channel: String, options: &Options) -> Option { + if options.retention_ms == 0 { + return None; + } + let interval = Duration::from_millis(options.cleanup_interval_ms); + let maximum_batch_size = options.cleanup_batch_size as i64; + Some(Self(tokio::spawn(async move { + let mut next_interval = interval; + let mut batch_size = maximum_batch_size; + let mut fast_batches = 0; + loop { + tokio::time::sleep(next_interval).await; + let started = tokio::time::Instant::now(); + match prune(&pool, &channel, batch_size).await { + Ok(removed) => { + next_interval = interval; + if removed == batch_size as u64 && started.elapsed() < FAST_BATCH_DURATION { + fast_batches += 1; + } else { + fast_batches = 0; + } + if fast_batches == FAST_BATCHES_BEFORE_GROWTH { + batch_size = (batch_size + (batch_size / 4).max(1)).min(maximum_batch_size); + fast_batches = 0; + } + }, + Err(CleanupError::TimedOut) => { + batch_size = (batch_size / 2).max(1); + fast_batches = 0; + next_interval = interval.max(TIMEOUT_RETRY_INTERVAL); + eprintln!("PostgreSQLの掃除が期限を超えたため、次回の削除上限を{batch_size}件に減らします"); + }, + Err(CleanupError::Database(error)) => { + // SDKのstdoutはプロトコル専用。DB詳細や接続文字列をログに出さない。 + eprintln!("PostgreSQLの期限切れフレーム掃除を延期します: {error}"); + next_interval = RETRY_INTERVAL; + }, + } + } + }))) + } +} + +impl Drop for Maintenance { + fn drop(&mut self) { + self.0.abort(); + } +} + +async fn prune(pool: &Pool, channel: &str, batch_size: i64) -> Result { + let mut client = pool.get().await.map_err(|error| CleanupError::Database(pool_error(error)))?; + let transaction = client.build_transaction().isolation_level(IsolationLevel::ReadCommitted).start().await?; + transaction.batch_execute(CLEANUP_STATEMENT_TIMEOUT).await?; + let locked: bool = transaction.query_one(TRY_CHANNEL_LOCK, &[&channel]).await?.get(0); + if !locked { + transaction.rollback().await?; + return Ok(0); + } + // 登録とは排他、送信とは共有。ロック後に合意済み設定の最新スナップショットを読む。 + let retention_ms: i64 = transaction.query_one(CHANNEL_RETENTION, &[&channel]).await?.get(0); + let removed = if retention_ms > 0 { + transaction.execute(PRUNE, &[&channel, &retention_ms, &batch_size]).await? + } else { + 0 + }; + transaction.commit().await?; + Ok(removed) +} + +#[cfg(test)] +mod tests; diff --git a/src/maintenance/tests.rs b/src/maintenance/tests.rs new file mode 100644 index 0000000..eb83fad --- /dev/null +++ b/src/maintenance/tests.rs @@ -0,0 +1,106 @@ +use super::prune; +#[path = "../../tests/common/database.rs"] +mod database; +use database::{database, CONNECTION_ENV}; +use deadpool_postgres::{Manager, Pool}; +use tokio_postgres::{Client, NoTls}; +use uuid::Uuid; + +fn pool() -> Pool { + let config = std::env::var(CONNECTION_ENV).unwrap().parse().unwrap(); + Pool::builder(Manager::new(config, NoTls)).max_size(1).build().unwrap() +} + +async fn expired_channel(client: &Client, count: i32) -> String { + let channel = Uuid::new_v4().to_string(); + client + .execute( + "INSERT INTO stegrdb_relay.nodes(channel,node_id,replay_window_ms,retention_ms) VALUES ($1,'node',0,60000)", + &[&channel], + ) + .await + .unwrap(); + client.execute("INSERT INTO stegrdb_relay.frames(channel,id,sender,payload,created_at) SELECT $1,md5(value::text)::uuid,'source',decode(repeat('ab',64),'hex'),statement_timestamp()-interval '5 minutes' FROM generate_series(1,$2::int) AS value", &[&channel, &count]).await.unwrap(); + channel +} + +#[tokio::test] +#[ignore = "専用DBでscripts/test-postgres.shを実行する"] +async fn cleanup_requires_every_node_to_enable_retention_and_stays_within_its_channel() { + let (client, connection) = database().await; + let channel = expired_channel(&client, 1).await; + let other = expired_channel(&client, 1).await; + client.execute("INSERT INTO stegrdb_relay.nodes(channel,node_id) VALUES ($1,'legacy')", &[&channel]).await.unwrap(); + let pool = pool(); + assert_eq!(prune(&pool, &channel, 10).await.unwrap(), 0); + client.execute("UPDATE stegrdb_relay.nodes SET retention_ms=60000,replay_window_ms=0 WHERE channel=$1", &[&channel]).await.unwrap(); + assert_eq!(prune(&pool, &channel, 10).await.unwrap(), 1); + let retained: i64 = client.query_one("SELECT count(*) FROM stegrdb_relay.frames WHERE channel=$1", &[&other]).await.unwrap().get(0); + assert_eq!(retained, 1); + connection.abort(); +} + +#[tokio::test] +#[ignore = "専用DBでscripts/test-postgres.shを実行する"] +async fn cleanup_preserves_unacknowledged_frames_and_the_longest_replay_or_retention_window() { + let (client, connection) = database().await; + let channel = expired_channel(&client, 4).await; + client + .execute( + "INSERT INTO stegrdb_relay.nodes(channel,node_id,replay_window_ms,retention_ms) VALUES ($1,'long',180000,120000)", + &[&channel], + ) + .await + .unwrap(); + client + .execute( + "INSERT INTO stegrdb_relay.pending SELECT channel,'node',id,position FROM stegrdb_relay.frames WHERE channel=$1 AND id=md5('1')::uuid", + &[&channel], + ) + .await + .unwrap(); + client + .execute( + "UPDATE stegrdb_relay.frames SET created_at=statement_timestamp()-interval '2 minutes' WHERE channel=$1 AND id=md5('2')::uuid", + &[&channel], + ) + .await + .unwrap(); + client + .execute( + "UPDATE stegrdb_relay.frames SET created_at=statement_timestamp() WHERE channel=$1 AND id=md5('3')::uuid", + &[&channel], + ) + .await + .unwrap(); + let pool = pool(); + assert_eq!(prune(&pool, &channel, 10).await.unwrap(), 1); + let remaining: Vec = client.query("SELECT id FROM stegrdb_relay.frames WHERE channel=$1", &[&channel]).await.unwrap().into_iter().map(|row| row.get(0)).collect(); + assert_eq!(remaining.len(), 3); + let pending: i64 = client.query_one("SELECT count(*) FROM stegrdb_relay.pending WHERE channel=$1", &[&channel]).await.unwrap().get(0); + assert_eq!(pending, 1); + client.execute("DELETE FROM stegrdb_relay.pending WHERE channel=$1", &[&channel]).await.unwrap(); + assert_eq!(prune(&pool, &channel, 10).await.unwrap(), 1); + connection.abort(); +} + +#[tokio::test] +#[ignore = "専用DBでscripts/test-postgres.shを実行する"] +async fn cleanup_limits_each_batch_skips_locked_rows_and_yields_to_registration() { + let (mut client, connection) = database().await; + let channel = expired_channel(&client, 5).await; + let pool = pool(); + let transaction = client.transaction().await.unwrap(); + transaction.query_one("SELECT id FROM stegrdb_relay.frames WHERE channel=$1 AND id=md5('1')::uuid FOR UPDATE", &[&channel]).await.unwrap(); + assert_eq!(prune(&pool, &channel, 2).await.unwrap(), 2); + assert_eq!(prune(&pool, &channel, 2).await.unwrap(), 2); + assert_eq!(prune(&pool, &channel, 2).await.unwrap(), 0); + transaction.rollback().await.unwrap(); + let transaction = client.transaction().await.unwrap(); + transaction.query_one(crate::queries::LOCK_CHANNEL, &[&channel]).await.unwrap(); + assert_eq!(prune(&pool, &channel, 2).await.unwrap(), 0); + transaction.rollback().await.unwrap(); + assert_eq!(prune(&pool, &channel, 2).await.unwrap(), 1); + assert_eq!(prune(&pool, &channel, 2).await.unwrap(), 0); + connection.abort(); +} diff --git a/src/manifest.rs b/src/manifest.rs index 6d7aee1..7c48e01 100644 --- a/src/manifest.rs +++ b/src/manifest.rs @@ -1,6 +1,8 @@ +use crate::options::{Options, MAX_CLEANUP_BATCH_SIZE, MAX_CLEANUP_INTERVAL_MS, MAX_CONNECTIONS, MAX_WINDOW_MS, MIN_CLEANUP_INTERVAL_MS}; use amitoki_plugin_sdk::{PluginManifest, PROTOCOL_VERSION}; pub fn manifest() -> PluginManifest { + let defaults = Options::default(); PluginManifest { name: "postgres".into(), version: env!("CARGO_PKG_VERSION").into(), @@ -9,9 +11,12 @@ pub fn manifest() -> PluginManifest { config_schema: serde_json::json!({ "type":"object", "additionalProperties":false, "properties": { - "connection_env":{"type":"string","minLength":1,"default":"AMITOKI_POSTGRES_URL","description":"接続文字列を読む環境変数名"}, - "max_connections":{"type":"integer","minimum":1,"maximum":64,"default":4,"description":"最大接続数"}, - "replay_window_ms":{"type":"integer","minimum":0,"maximum":2147483647,"default":4000,"description":"初回登録時に過去のフレームを受信する期間（ミリ秒）"} + "connection_env":{"type":"string","minLength":1,"default":defaults.connection_env,"description":"接続文字列を読む環境変数名"}, + "max_connections":{"type":"integer","minimum":1,"maximum":MAX_CONNECTIONS,"default":defaults.max_connections,"description":"最大接続数"}, + "replay_window_ms":{"type":"integer","minimum":0,"maximum":MAX_WINDOW_MS,"default":defaults.replay_window_ms,"description":"初回登録時に過去のフレームを受信する期間（ミリ秒）"}, + "retention_ms":{"type":"integer","minimum":0,"maximum":MAX_WINDOW_MS,"default":defaults.retention_ms,"description":"ACK済みフレームの最低保持期間。0で自動削除を無効化。channel内の全ノードの合意が必要（ミリ秒）"}, + "cleanup_interval_ms":{"type":"integer","minimum":MIN_CLEANUP_INTERVAL_MS,"maximum":MAX_CLEANUP_INTERVAL_MS,"default":defaults.cleanup_interval_ms,"description":"保持期限を過ぎたフレームの掃除間隔（ミリ秒）"}, + "cleanup_batch_size":{"type":"integer","minimum":1,"maximum":MAX_CLEANUP_BATCH_SIZE,"default":defaults.cleanup_batch_size,"description":"1回の掃除で削除する最大フレーム数"} } }), } diff --git a/src/options.rs b/src/options.rs index d33455f..4cc3812 100644 --- a/src/options.rs +++ b/src/options.rs @@ -6,10 +6,18 @@ use tokio_postgres::{config::SslMode, Config}; // 送受信を並行処理でき、ノード数に比例して接続を増やしすぎない規定値。 const DEFAULT_CONNECTIONS: usize = 4; // ノードごとの誤設定がDB全体の接続枠を使い切らないようにする。 -const MAX_CONNECTIONS: usize = 64; +pub(crate) const MAX_CONNECTIONS: usize = 64; // 初回登録だけは、起動直前のパケットも取り込む。 const DEFAULT_REPLAY_WINDOW_MS: u64 = 4000; const CONNECT_TIMEOUT: Duration = Duration::from_secs(10); +// 掃除のトランザクションを短くし、送信を長時間止めないための既定値。 +const DEFAULT_CLEANUP_INTERVAL_MS: u64 = 1000; +const DEFAULT_CLEANUP_BATCH_SIZE: usize = 4096; +// PostgreSQLの整数ミリ秒パラメータと、掃除の負荷調整に許す範囲。 +pub(crate) const MAX_WINDOW_MS: u64 = i32::MAX as u64; +pub(crate) const MIN_CLEANUP_INTERVAL_MS: u64 = 100; +pub(crate) const MAX_CLEANUP_INTERVAL_MS: u64 = 3_600_000; +pub(crate) const MAX_CLEANUP_BATCH_SIZE: usize = 65_536; #[derive(Deserialize)] #[serde(default, deny_unknown_fields)] @@ -17,6 +25,9 @@ pub(crate) struct Options { pub connection_env: String, pub max_connections: usize, pub replay_window_ms: u64, + pub retention_ms: u64, + pub cleanup_interval_ms: u64, + pub cleanup_batch_size: usize, } impl Default for Options { @@ -25,15 +36,16 @@ impl Default for Options { connection_env: "AMITOKI_POSTGRES_URL".into(), max_connections: DEFAULT_CONNECTIONS, replay_window_ms: DEFAULT_REPLAY_WINDOW_MS, + retention_ms: 0, + cleanup_interval_ms: DEFAULT_CLEANUP_INTERVAL_MS, + cleanup_batch_size: DEFAULT_CLEANUP_BATCH_SIZE, } } } impl Options { pub fn connection_config(&self) -> Result { - if self.max_connections == 0 || self.max_connections > MAX_CONNECTIONS || self.replay_window_ms > i32::MAX as u64 { - return Err(RelayError::permanent("postgresの接続数または再生期間が範囲外です")); - } + self.validate()?; let connection_string = std::env::var(&self.connection_env).map_err(|_| RelayError::permanent(format!("接続情報の環境変数が未設定です: {}", self.connection_env)))?; // 接続文字列を含むパースエラーは認証情報が混ざるため外へ出さない。 let mut config: Config = connection_string.parse().map_err(|_| RelayError::permanent("PostgreSQL接続情報の書式が不正です"))?; @@ -43,4 +55,46 @@ impl Options { config.connect_timeout(CONNECT_TIMEOUT); Ok(config) } + + fn validate(&self) -> Result<(), RelayError> { + if self.connection_env.is_empty() || self.max_connections == 0 || self.max_connections > MAX_CONNECTIONS || self.replay_window_ms > MAX_WINDOW_MS { + return Err(RelayError::permanent("postgresの接続数または再生期間が範囲外です")); + } + // replayと同じ整数ミリ秒範囲にし、DB側のinterval計算を正確に保つ。 + if self.retention_ms > MAX_WINDOW_MS || (self.retention_ms != 0 && self.retention_ms < self.replay_window_ms) { + return Err(RelayError::permanent( + "postgresの保持期間は0（無効）または再生期間以上の2147483647ミリ秒以下で指定してください", + )); + } + if !(MIN_CLEANUP_INTERVAL_MS..=MAX_CLEANUP_INTERVAL_MS).contains(&self.cleanup_interval_ms) || !(1..=MAX_CLEANUP_BATCH_SIZE).contains(&self.cleanup_batch_size) { + return Err(RelayError::permanent("postgresの掃除間隔は100〜3600000ミリ秒、1回の削除上限は1〜65536件で指定してください")); + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::Options; + + #[test] + fn retention_is_disabled_by_default_and_cannot_shorten_replay() { + assert_eq!(Options::default().retention_ms, 0); + assert!(Options::default().validate().is_ok()); + let options = Options { + retention_ms: 1, + ..Options::default() + }; + assert!(options.validate().is_err()); + let options = Options { + retention_ms: i32::MAX as u64 + 1, + ..Options::default() + }; + assert!(options.validate().is_err()); + let options = Options { + cleanup_batch_size: 0, + ..Options::default() + }; + assert!(options.validate().is_err()); + } } diff --git a/src/queries.rs b/src/queries.rs index ad855a0..a0cd059 100644 --- a/src/queries.rs +++ b/src/queries.rs @@ -1,4 +1,5 @@ pub(crate) const LOCK_CHANNEL: &str = "SELECT pg_advisory_xact_lock(hashtextextended($1, 0))"; +pub(crate) const LOCK_CHANNEL_SHARED: &str = "SELECT pg_advisory_xact_lock_shared(hashtextextended($1, 0))"; pub(crate) const PUBLISH: &str = " WITH inserted AS ( @@ -12,24 +13,40 @@ SELECT inserted.channel, nodes.node_id, inserted.id, inserted.position FROM inserted JOIN stegrdb_relay.nodes AS nodes ON nodes.channel = inserted.channel WHERE nodes.node_id <> $2"; +// LIMITを本文JOINより前に適用し、OFFSET 0でLATERALの展開を防ぐ。 +// channelをInitPlanから渡し、別channelの古い頻度統計で本文が1行と +// 誤推定されて期間索引を全走査する計画を避ける。UUIDとの主キー参照を優先する。 pub(crate) const RECEIVE: &str = " -SELECT frames.id, frames.payload -FROM stegrdb_relay.pending AS pending -JOIN stegrdb_relay.frames AS frames ON frames.channel = pending.channel AND frames.id = pending.frame_id -WHERE pending.channel = $1 AND pending.node_id = $2 -ORDER BY pending.position LIMIT $3"; +SELECT frame.id, frame.payload, delivery.position +FROM ( + SELECT frame_id, position FROM stegrdb_relay.pending + WHERE channel = $1 AND node_id = $2 + ORDER BY position LIMIT $3 +) AS delivery +CROSS JOIN LATERAL ( + SELECT id, payload FROM stegrdb_relay.frames + WHERE channel = (SELECT $1::text) AND id = delivery.frame_id + OFFSET 0 +) AS frame +ORDER BY delivery.position"; pub(crate) const ACKNOWLEDGE: &str = " -DELETE FROM stegrdb_relay.pending -WHERE channel = $1 AND node_id = $2 AND frame_id = ANY($3::uuid[])"; +DELETE FROM stegrdb_relay.pending AS pending +USING unnest($3::uuid[], $4::bigint[]) AS acknowledged(id, position) +WHERE pending.channel = $1 AND pending.node_id = $2 +AND pending.frame_id = acknowledged.id AND pending.position = acknowledged.position"; pub(crate) const REGISTER: &str = " -INSERT INTO stegrdb_relay.nodes (channel, node_id) VALUES ($1, $2) +INSERT INTO stegrdb_relay.nodes (channel, node_id, replay_window_ms, retention_ms) VALUES ($1, $2, $3, $4) ON CONFLICT (channel, node_id) DO NOTHING"; +pub(crate) const UPDATE_NODE_OPTIONS: &str = " +UPDATE stegrdb_relay.nodes SET replay_window_ms = $3, retention_ms = $4 +WHERE channel = $1 AND node_id = $2"; + pub(crate) const REPLAY: &str = " INSERT INTO stegrdb_relay.pending (channel, node_id, frame_id, position) SELECT channel, $2, id, position FROM stegrdb_relay.frames WHERE channel = $1 AND sender <> $2 -AND created_at >= clock_timestamp() - $3::int * INTERVAL '1 millisecond' +AND created_at >= statement_timestamp() - $3::int * INTERVAL '1 millisecond' ON CONFLICT DO NOTHING"; diff --git a/tests/common/database.rs b/tests/common/database.rs new file mode 100644 index 0000000..cd4c7e6 --- /dev/null +++ b/tests/common/database.rs @@ -0,0 +1,12 @@ +use tokio_postgres::{Client, NoTls}; + +pub const CONNECTION_ENV: &str = "AMITOKI_TEST_POSTGRES_URL"; + +pub async fn database() -> (Client, tokio::task::JoinHandle<()>) { + let connection_string = std::env::var(CONNECTION_ENV).expect("専用のPostgreSQLテスト接続を設定してください"); + let (client, connection) = tokio_postgres::connect(&connection_string, NoTls).await.expect("テスト用DBへの接続"); + let task = tokio::spawn(async move { + connection.await.expect("テスト用DBの接続維持"); + }); + (client, task) +} diff --git a/tests/common/mod.rs b/tests/common/mod.rs new file mode 100644 index 0000000..3a15b40 --- /dev/null +++ b/tests/common/mod.rs @@ -0,0 +1,35 @@ +use amitoki_relay::{Frame, Relay, RelayContext, RelayPlugin}; +use amitoki_relay_postgres::PostgresPlugin; +use bytes::Bytes; +use serde_json::json; +use std::sync::Arc; +use tokio_postgres::Client; +mod database; +pub use database::database; +use database::CONNECTION_ENV; + +pub async fn connect(channel: &str, node: &str) -> Result, amitoki_relay::RelayError> { + connect_with_options(channel, node, json!({})).await +} + +pub async fn connect_with_options(channel: &str, node: &str, mut options: serde_json::Value) -> Result, amitoki_relay::RelayError> { + options["connection_env"] = json!(CONNECTION_ENV); + PostgresPlugin + .connect( + RelayContext { + channel: channel.into(), + node_id: node.into(), + }, + options, + ) + .await +} + +pub fn frame(value: u8) -> Frame { + Frame::new(Bytes::from(vec![value; 1024])).unwrap() +} + +pub async fn remove_channel(client: &Client, channel: &str) { + client.execute("DELETE FROM stegrdb_relay.nodes WHERE channel = $1", &[&channel]).await.unwrap(); + client.execute("DELETE FROM stegrdb_relay.frames WHERE channel = $1", &[&channel]).await.unwrap(); +} diff --git a/tests/concurrency.rs b/tests/concurrency.rs new file mode 100644 index 0000000..716b617 --- /dev/null +++ b/tests/concurrency.rs @@ -0,0 +1,208 @@ +mod common; + +use amitoki_relay::{Frame, Relay}; +use common::{connect, connect_with_options, database, frame, remove_channel}; +use serde_json::json; +use std::{collections::HashMap, sync::Arc, time::Duration}; +use tokio::task::JoinSet; +use tokio_postgres::Client; +use uuid::Uuid; + +// 成功条件はDBが返す待機関係で判断し、期限は回帰時のハング検出だけに使う。 +const TEST_DEADLINE: Duration = Duration::from_secs(15); +const LOCK_POLL_INTERVAL: Duration = Duration::from_millis(5); + +async fn wait_for_blocked_backend(client: &Client, blocker: i32) -> i32 { + tokio::time::timeout(TEST_DEADLINE, async { + loop { + if let Some(row) = client.query_opt("SELECT pid FROM pg_stat_activity WHERE $1 = ANY(pg_blocking_pids(pid)) LIMIT 1", &[&blocker]).await.unwrap() { + return row.get(0); + } + tokio::time::sleep(LOCK_POLL_INTERVAL).await; + } + }) + .await + .expect("DBでロック待ちが発生すること") +} + +#[tokio::test] +#[ignore = "専用DBでscripts/test-postgres.shを実行する"] +async fn publishers_can_commit_while_another_publisher_holds_the_channel_lock() { + let channel = Uuid::new_v4().to_string(); + let a = connect(&channel, "a").await.unwrap(); + let b = connect(&channel, "b").await.unwrap(); + let receiver = connect(&channel, "receiver").await.unwrap(); + let (mut client, connection) = database().await; + let transaction = client.transaction().await.unwrap(); + transaction.query_one("SELECT pg_advisory_xact_lock_shared(hashtextextended($1, 0))", &[&channel]).await.unwrap(); + let packets = [frame(1), frame(2)]; + tokio::time::timeout(TEST_DEADLINE, async { + let (a_sent, b_sent) = tokio::join!(a.publish(&packets[..1]), b.publish(&packets[1..])); + a_sent.unwrap(); + b_sent.unwrap(); + }) + .await + .expect("共有ロック中も別の送信が完了すること"); + transaction.rollback().await.unwrap(); + assert_eq!(receiver.receive(10).await.unwrap().len(), 2); + remove_channel(&client, &channel).await; + connection.abort(); +} + +#[tokio::test] +#[ignore = "専用DBでscripts/test-postgres.shを実行する"] +async fn a_publisher_waiting_for_registration_delivers_to_the_new_node_without_replay() { + let channel = Uuid::new_v4().to_string(); + let sender = connect(&channel, "sender").await.unwrap(); + let (mut blocker, blocker_connection) = database().await; + let (observer, observer_connection) = database().await; + let blocker_pid: i32 = blocker.query_one("SELECT pg_backend_pid()", &[]).await.unwrap().get(0); + let transaction = blocker.transaction().await.unwrap(); + // REGISTERを一意制約で待たせ、登録側がchannelの排他ロックを持つ状態を作る。 + transaction.execute("INSERT INTO stegrdb_relay.nodes VALUES ($1, 'receiver')", &[&channel]).await.unwrap(); + let joining_channel = channel.clone(); + let joining = tokio::spawn(async move { connect_with_options(&joining_channel, "receiver", json!({"replay_window_ms": 0})).await }); + let registration_pid = wait_for_blocked_backend(&observer, blocker_pid).await; + let packet = frame(3); + let sent = packet.clone(); + let publishing = tokio::spawn(async move { sender.publish(&[sent]).await }); + wait_for_blocked_backend(&observer, registration_pid).await; + transaction.rollback().await.unwrap(); + let receiver = tokio::time::timeout(TEST_DEADLINE, joining).await.unwrap().unwrap().unwrap(); + tokio::time::timeout(TEST_DEADLINE, publishing).await.unwrap().unwrap().unwrap(); + let deliveries = receiver.receive(10).await.unwrap(); + assert_eq!(deliveries.len(), 1); + assert_eq!(deliveries[0].frame, packet); + remove_channel(&observer, &channel).await; + blocker_connection.abort(); + observer_connection.abort(); +} + +#[tokio::test] +#[ignore = "専用DBでscripts/test-postgres.shを実行する"] +async fn registration_waits_for_a_publisher_then_replays_its_committed_frame() { + let channel = Uuid::new_v4().to_string(); + let _sender = connect(&channel, "sender").await.unwrap(); + let (mut blocker, blocker_connection) = database().await; + let (observer, observer_connection) = database().await; + let blocker_pid: i32 = blocker.query_one("SELECT pg_backend_pid()", &[]).await.unwrap().get(0); + let transaction = blocker.transaction().await.unwrap(); + transaction.query_one("SELECT pg_advisory_xact_lock_shared(hashtextextended($1, 0))", &[&channel]).await.unwrap(); + let packet = frame(4); + transaction + .execute( + "INSERT INTO stegrdb_relay.frames(channel,id,sender,payload) VALUES ($1,$2,'sender',$3)", + &[&channel, &packet.id, &packet.bytes.as_ref()], + ) + .await + .unwrap(); + let joining_channel = channel.clone(); + let joining = tokio::spawn(async move { connect_with_options(&joining_channel, "receiver", json!({"replay_window_ms": 60_000})).await }); + wait_for_blocked_backend(&observer, blocker_pid).await; + transaction.commit().await.unwrap(); + let receiver = tokio::time::timeout(TEST_DEADLINE, joining).await.unwrap().unwrap().unwrap(); + let deliveries = receiver.receive(10).await.unwrap(); + assert_eq!(deliveries.len(), 1); + assert_eq!(deliveries[0].frame, packet); + remove_channel(&observer, &channel).await; + blocker_connection.abort(); + observer_connection.abort(); +} + +#[tokio::test] +#[ignore = "専用DBでscripts/test-postgres.shを実行する"] +async fn cancelling_a_blocked_insert_rolls_back_the_batch_and_releases_the_connection() { + let channel = Uuid::new_v4().to_string(); + let sender = connect_with_options(&channel, "sender", json!({"max_connections": 1})).await.unwrap(); + let receiver = connect(&channel, "receiver").await.unwrap(); + let (mut blocker, blocker_connection) = database().await; + let (observer, observer_connection) = database().await; + let blocker_pid: i32 = blocker.query_one("SELECT pg_backend_pid()", &[]).await.unwrap().get(0); + let transaction = blocker.transaction().await.unwrap(); + let packets = [frame(5), frame(6)]; + transaction + .execute( + "INSERT INTO stegrdb_relay.frames(channel,id,sender,payload) VALUES ($1,$2,'sender',$3)", + &[&channel, &packets[0].id, &packets[0].bytes.as_ref()], + ) + .await + .unwrap(); + let sending = sender.clone(); + let publishing = tokio::spawn(async move { sending.publish(&packets).await }); + wait_for_blocked_backend(&observer, blocker_pid).await; + publishing.abort(); + assert!(publishing.await.unwrap_err().is_cancelled()); + transaction.rollback().await.unwrap(); + let final_packet = frame(7); + tokio::time::timeout(TEST_DEADLINE, sender.publish(std::slice::from_ref(&final_packet))).await.unwrap().unwrap(); + let deliveries = receiver.receive(10).await.unwrap(); + assert_eq!(deliveries.len(), 1, "キャンセルしたバッチが部分的に保存されないこと"); + assert_eq!(deliveries[0].frame, final_packet); + remove_channel(&observer, &channel).await; + blocker_connection.abort(); + observer_connection.abort(); +} + +#[tokio::test] +#[ignore = "専用DBでscripts/test-postgres.shを実行する"] +async fn concurrent_retries_of_the_same_batch_deliver_each_uuid_once_in_input_order() { + let channel = Uuid::new_v4().to_string(); + let sender = connect(&channel, "sender").await.unwrap(); + let receiver = connect(&channel, "receiver").await.unwrap(); + let packets: Vec<_> = (0..128).map(frame).collect(); + let mut sending = JoinSet::new(); + for _ in 0..8 { + let sender = sender.clone(); + let packets = packets.clone(); + sending.spawn(async move { sender.publish(&packets).await }); + } + tokio::time::timeout(TEST_DEADLINE, async { + while let Some(sent) = sending.join_next().await { + sent.unwrap().unwrap(); + } + }) + .await + .unwrap(); + let deliveries = receiver.receive(1024).await.unwrap(); + assert_eq!(deliveries.iter().map(|delivery| &delivery.frame).collect::>(), packets.iter().collect::>()); + receiver.acknowledge(&deliveries.into_iter().map(|delivery| delivery.receipt).collect::>()).await.unwrap(); + sender.publish(&packets).await.unwrap(); + assert!(receiver.receive(1024).await.unwrap().is_empty()); + let (client, connection) = database().await; + remove_channel(&client, &channel).await; + connection.abort(); +} + +async fn publish_until_committed(relay: Arc, packets: Vec) { + loop { + match relay.publish(&packets).await { + Ok(()) => return, + // 重なるUUIDを逆順でINSERTするとDBが片方を中断し得る。 + // 呼び出し側の再送でバッチ全体が冪等に復旧することを確かめる。 + Err(error) => assert!(error.is_retryable(), "再送可能であること: {error}"), + } + } +} + +#[tokio::test] +#[ignore = "専用DBでscripts/test-postgres.shを実行する"] +async fn overlapping_batches_from_different_senders_remain_idempotent_after_database_retries() { + let channel = Uuid::new_v4().to_string(); + let a = connect(&channel, "a").await.unwrap(); + let b = connect(&channel, "b").await.unwrap(); + let receiver = connect(&channel, "receiver").await.unwrap(); + let packets: Vec<_> = (0..128).map(frame).collect(); + let reversed: Vec<_> = packets.iter().rev().cloned().collect(); + tokio::time::timeout(TEST_DEADLINE, async { + tokio::join!(publish_until_committed(a, packets.clone()), publish_until_committed(b, reversed)); + }) + .await + .expect("DBが中断した送信も再試行で完了すること"); + let delivered = receiver.receive(1024).await.unwrap(); + assert_eq!(delivered.len(), packets.len()); + let contents: HashMap<_, _> = delivered.into_iter().map(|delivery| (delivery.frame.id, delivery.frame.bytes)).collect(); + assert_eq!(contents, packets.into_iter().map(|packet| (packet.id, packet.bytes)).collect()); + let (client, connection) = database().await; + remove_channel(&client, &channel).await; + connection.abort(); +} diff --git a/tests/delivery.rs b/tests/delivery.rs index 53be76b..ffa0fee 100644 --- a/tests/delivery.rs +++ b/tests/delivery.rs @@ -1,44 +1,17 @@ -use amitoki_relay::{Frame, Relay, RelayContext, RelayPlugin}; -use amitoki_relay_postgres::PostgresPlugin; -use bytes::Bytes; -use serde_json::json; +mod common; + +use amitoki_relay::Relay; +use common::{connect, database, frame, remove_channel}; use std::{ collections::HashSet, sync::Arc, time::{Duration, Instant}, }; -use tokio_postgres::{Client, NoTls}; use uuid::Uuid; -const CONNECTION_ENV: &str = "AMITOKI_TEST_POSTGRES_URL"; const BATCH_SIZE: usize = 128; const FRAME_COUNT: usize = 2048; -async fn database() -> (Client, tokio::task::JoinHandle<()>) { - let connection_string = std::env::var(CONNECTION_ENV).expect("専用のPostgreSQLテスト接続を設定してください"); - let (client, connection) = tokio_postgres::connect(&connection_string, NoTls).await.expect("テスト用DBへの接続"); - let task = tokio::spawn(async move { - connection.await.expect("テスト用DBの接続維持"); - }); - (client, task) -} - -async fn connect(channel: &str, node: &str) -> Result, amitoki_relay::RelayError> { - PostgresPlugin - .connect( - RelayContext { - channel: channel.into(), - node_id: node.into(), - }, - json!({"connection_env": CONNECTION_ENV}), - ) - .await -} - -fn frame(value: u8) -> Frame { - Frame::new(Bytes::from(vec![value; 1024])).unwrap() -} - async fn drain(relay: &Arc) -> HashSet { let mut received = HashSet::new(); loop { @@ -156,8 +129,3 @@ async fn concurrent_node_registration_and_publish_do_not_lose_or_duplicate_deliv remove_channel(&client, &channel).await; connection.abort(); } - -async fn remove_channel(client: &Client, channel: &str) { - client.execute("DELETE FROM stegrdb_relay.nodes WHERE channel = $1", &[&channel]).await.unwrap(); - client.execute("DELETE FROM stegrdb_relay.frames WHERE channel = $1", &[&channel]).await.unwrap(); -} diff --git a/tests/fixtures/schema-v0.1.2.sql b/tests/fixtures/schema-v0.1.2.sql new file mode 100644 index 0000000..b8588d0 --- /dev/null +++ b/tests/fixtures/schema-v0.1.2.sql @@ -0,0 +1,31 @@ +-- 旧packets/node_list/firewall_settingsには触れない、中継専用のスキーマ。 +-- 既存の未ACKキューを引き継ぐためスキーマ名を維持する。 +CREATE SCHEMA IF NOT EXISTS stegrdb_relay; + +CREATE TABLE IF NOT EXISTS stegrdb_relay.nodes ( + channel TEXT NOT NULL, + node_id TEXT NOT NULL, + PRIMARY KEY (channel, node_id) +); + +CREATE TABLE IF NOT EXISTS stegrdb_relay.frames ( + channel TEXT NOT NULL, + id UUID NOT NULL, + sender TEXT NOT NULL, + position BIGINT GENERATED ALWAYS AS IDENTITY, + payload BYTEA NOT NULL CHECK (octet_length(payload) BETWEEN 14 AND 65535), + created_at TIMESTAMPTZ NOT NULL DEFAULT clock_timestamp(), + PRIMARY KEY (channel, id) +); +CREATE INDEX IF NOT EXISTS frames_channel_created_at ON stegrdb_relay.frames (channel, created_at); + +CREATE TABLE IF NOT EXISTS stegrdb_relay.pending ( + channel TEXT NOT NULL, + node_id TEXT NOT NULL, + frame_id UUID NOT NULL, + position BIGINT NOT NULL, + PRIMARY KEY (channel, node_id, frame_id), + FOREIGN KEY (channel, node_id) REFERENCES stegrdb_relay.nodes ON DELETE CASCADE, + FOREIGN KEY (channel, frame_id) REFERENCES stegrdb_relay.frames ON DELETE CASCADE +); +CREATE INDEX IF NOT EXISTS pending_delivery_order ON stegrdb_relay.pending (channel, node_id, position); diff --git a/tests/process.rs b/tests/process.rs new file mode 100644 index 0000000..6c5e3b2 --- /dev/null +++ b/tests/process.rs @@ -0,0 +1,42 @@ +mod common; + +use amitoki_plugin_sdk::ProcessRelay; +use amitoki_relay::{Relay, RelayContext}; +use amitoki_relay_postgres::manifest::manifest; +use common::{connect, database, frame, remove_channel}; +use serde_json::json; +use std::path::Path; +use uuid::Uuid; + +#[tokio::test] +#[ignore = "専用DBでscripts/test-postgres.shを実行する"] +async fn the_executable_delivers_payloads_and_acknowledges_opaque_receipts_over_ipc() { + let channel = Uuid::new_v4().to_string(); + let sender = connect(&channel, "sender").await.unwrap(); + let receiver = ProcessRelay::connect( + Path::new(env!("CARGO_BIN_EXE_amitoki-plugin-postgres")), + &manifest(), + ( + RelayContext { + channel: channel.clone(), + node_id: "receiver".into(), + }, + json!({"connection_env": "AMITOKI_TEST_POSTGRES_URL"}), + ), + ) + .await + .unwrap(); + let packets: Vec<_> = (0..128).map(frame).collect(); + sender.publish(&packets).await.unwrap(); + let deliveries = receiver.receive(128).await.unwrap(); + assert_eq!(deliveries.iter().map(|delivery| &delivery.frame).collect::>(), packets.iter().collect::>()); + let receipts: Vec<_> = deliveries.into_iter().map(|delivery| delivery.receipt).collect(); + receiver.acknowledge(&receipts).await.unwrap(); + receiver.acknowledge(&receipts).await.unwrap(); + sender.publish(&packets).await.unwrap(); + assert!(receiver.receive(128).await.unwrap().is_empty()); + drop(receiver); + let (client, connection) = database().await; + remove_channel(&client, &channel).await; + connection.abort(); +} diff --git a/tests/release/test_publish.py b/tests/release/test_publish.py new file mode 100644 index 0000000..33b74ad --- /dev/null +++ b/tests/release/test_publish.py @@ -0,0 +1,93 @@ +import importlib.util +import json +from pathlib import Path +import tempfile +import unittest +from unittest.mock import patch + +SCRIPT = Path(__file__).resolve().parents[2] / "scripts/publish-release.py" +SPEC = importlib.util.spec_from_file_location("publish_release", SCRIPT) +publisher = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(publisher) + + +class ReleaseTests(unittest.TestCase): + def setUp(self): + self.temporary = tempfile.TemporaryDirectory() + self.addCleanup(self.temporary.cleanup) + self.directory = Path(self.temporary.name) + self.binary = self.directory / f"amitoki-plugin-postgres-{publisher.TARGET}" + self.binary.write_bytes(b"verified executable") + self.manifest = {"name": "postgres", "version": "0.2.0", "protocol_version": 1} + self.package = {"manifest": self.manifest, "target": publisher.TARGET, + "binary": "amitoki-plugin-postgres", "sha256": publisher.digest(self.binary)} + self.manifest_path = self.directory / f"plugin-{publisher.TARGET}.json" + self.manifest_path.write_text(json.dumps(self.package)) + self.schema = self.directory / "source.sql" + self.schema.write_bytes(b"CREATE SCHEMA example;\n") + (self.directory / "schema.sql").write_bytes(self.schema.read_bytes()) + + def verify(self): + return publisher.verify_package(self.directory, release_version="0.2.0", schema=self.schema) + + def test_matching_package_produces_checksums_for_all_public_files(self): + with patch.object(publisher, "output", return_value=json.dumps(self.manifest)), \ + patch.object(publisher.subprocess, "check_output", return_value=self.schema.read_bytes()): + assets = self.verify() + self.assertEqual(len(assets), 4) + self.assertEqual((self.directory / "SHA256SUMS").read_text().splitlines(), + [f"{publisher.digest(path)} {path.name}" for path in assets[:-1]]) + + def test_corrupted_binary_is_rejected_before_execution(self): + self.binary.write_bytes(b"corrupted executable") + with patch.object(publisher, "output") as execute: + with self.assertRaisesRegex(ValueError, "SHA256"): + self.verify() + execute.assert_not_called() + + def test_wrong_version_or_architecture_is_rejected(self): + for field, value in (("target", "aarch64-unknown-linux-gnu"), ("version", "0.1.2")): + with self.subTest(field=field): + package = json.loads(json.dumps(self.package)) + if field == "version": + package["manifest"][field] = value + else: + package[field] = value + self.manifest_path.write_text(json.dumps(package)) + with self.assertRaises(ValueError): + self.verify() + + def test_stale_embedded_schema_is_rejected(self): + with patch.object(publisher, "output", return_value=json.dumps(self.manifest)), \ + patch.object(publisher.subprocess, "check_output", return_value=b"old SQL"): + with self.assertRaisesRegex(ValueError, "埋め込まれたSQL"): + self.verify() + + def test_missing_uploaded_asset_or_changed_notes_prevent_publication(self): + notes = self.directory / "notes.md" + notes.write_text("release notes\n") + release = {"body": notes.read_text(), "assets": [ + {"name": self.binary.name, "digest": f"sha256:{publisher.digest(self.binary)}"}]} + publisher.verify_uploaded(release, assets=[self.binary], notes=notes) + with self.assertRaises(ValueError): + publisher.verify_uploaded(release, assets=[self.binary, self.schema], notes=notes) + release["body"] = "different release notes" + with self.assertRaises(ValueError): + publisher.verify_uploaded(release, assets=[self.binary], notes=notes) + + def test_different_published_release_is_never_overwritten(self): + notes = self.directory / "notes.md" + notes.write_text("release notes\n") + existing = publisher.subprocess.CompletedProcess([], 0, json.dumps( + {"isDraft": False, "body": notes.read_text(), "assets": []})) + with patch.object(publisher, "release_preflight", return_value=("0.2.0", notes)), \ + patch.object(publisher, "verify_package", return_value=[self.binary]), \ + patch.object(publisher.subprocess, "run", return_value=existing) as command: + with self.assertRaises(ValueError): + publisher.publish(self.directory, tag="v0.2.0", repository="amitoki/amitoki-plugin-postgres") + self.assertEqual(command.call_count, 1) + self.assertEqual(command.call_args.args[0][:3], ["gh", "release", "view"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/replay.rs b/tests/replay.rs new file mode 100644 index 0000000..cbe026a --- /dev/null +++ b/tests/replay.rs @@ -0,0 +1,63 @@ +mod common; + +use common::{connect, connect_with_options, database, frame, remove_channel}; +use serde_json::json; +use uuid::Uuid; + +#[tokio::test] +#[ignore = "専用DBでscripts/test-postgres.shを実行する"] +async fn first_registration_replays_only_recent_frames_and_zero_disables_history() { + let channel = Uuid::new_v4().to_string(); + let sender = connect(&channel, "sender").await.unwrap(); + let (client, connection) = database().await; + let packets = [frame(1), frame(2), frame(3)]; + sender.publish(&packets).await.unwrap(); + client + .execute( + "UPDATE stegrdb_relay.frames SET created_at = statement_timestamp() - interval '1 day' WHERE channel=$1 AND id=$2", + &[&channel, &packets[0].id], + ) + .await + .unwrap(); + // DB時計のずれなどで未来の日付があっても、0なら履歴を取り込まない。 + client + .execute( + "UPDATE stegrdb_relay.frames SET created_at = statement_timestamp() + interval '1 day' WHERE channel=$1 AND id=$2", + &[&channel, &packets[2].id], + ) + .await + .unwrap(); + let recent = connect_with_options(&channel, "recent", json!({"replay_window_ms": 60_000})).await.unwrap(); + let disabled = connect_with_options(&channel, "disabled", json!({"replay_window_ms": 0})).await.unwrap(); + let deliveries = recent.receive(10).await.unwrap(); + assert_eq!( + deliveries.iter().map(|delivery| &delivery.frame).collect::>(), + packets[1..].iter().collect::>() + ); + assert!(disabled.receive(10).await.unwrap().is_empty()); + let fresh = frame(4); + sender.publish(std::slice::from_ref(&fresh)).await.unwrap(); + assert_eq!(disabled.receive(1).await.unwrap()[0].frame, fresh); + remove_channel(&client, &channel).await; + connection.abort(); +} + +#[tokio::test] +#[ignore = "専用DBでscripts/test-postgres.shを実行する"] +async fn receive_respects_zero_and_page_limits_with_a_large_history() { + let channel = Uuid::new_v4().to_string(); + let sender = connect(&channel, "sender").await.unwrap(); + let receiver = connect_with_options(&channel, "receiver", json!({"replay_window_ms": 0})).await.unwrap(); + let (client, connection) = database().await; + client.execute("INSERT INTO stegrdb_relay.frames(channel,id,sender,payload) SELECT $1,md5(value::text)::uuid,'sender',decode(repeat('ab',64),'hex') FROM generate_series(1,10000) AS value", &[&channel]).await.unwrap(); + let packets: Vec<_> = (0..130).map(frame).collect(); + sender.publish(&packets).await.unwrap(); + assert!(receiver.receive(0).await.unwrap().is_empty()); + let first = receiver.receive(128).await.unwrap(); + assert_eq!(first.iter().map(|delivery| &delivery.frame).collect::>(), packets[..128].iter().collect::>()); + receiver.acknowledge(&first.into_iter().map(|delivery| delivery.receipt).collect::>()).await.unwrap(); + let last = receiver.receive(128).await.unwrap(); + assert_eq!(last.iter().map(|delivery| &delivery.frame).collect::>(), packets[128..].iter().collect::>()); + remove_channel(&client, &channel).await; + connection.abort(); +} diff --git a/tests/retention.rs b/tests/retention.rs new file mode 100644 index 0000000..c5f7b65 --- /dev/null +++ b/tests/retention.rs @@ -0,0 +1,93 @@ +mod common; + +use common::{connect, connect_with_options, database, frame, remove_channel}; +use serde_json::json; +use std::time::Duration; +use tokio_postgres::Client; +use uuid::Uuid; + +// 実ワーカの完了を行の消滅で確認し、固定sleepで成功を推測しない。 +const CLEANUP_DEADLINE: Duration = Duration::from_secs(15); +const CLEANUP_POLL: Duration = Duration::from_millis(10); + +async fn wait_for_removal(client: &Client, channel: &str, id: Uuid) { + tokio::time::timeout(CLEANUP_DEADLINE, async { + while client.query_one("SELECT EXISTS(SELECT 1 FROM stegrdb_relay.frames WHERE channel=$1 AND id=$2)", &[&channel, &id]).await.unwrap().get::<_, bool>(0) { + tokio::time::sleep(CLEANUP_POLL).await; + } + }) + .await + .expect("設定した期限切れフレームが掃除されること"); +} + +#[tokio::test] +#[ignore = "専用DBでscripts/test-postgres.shを実行する"] +async fn background_cleanup_preserves_pending_and_allows_a_new_delivery_after_the_deduplication_window() { + let channel = Uuid::new_v4().to_string(); + let settings = json!({"retention_ms": 60_000, "cleanup_interval_ms": 100, "cleanup_batch_size": 1}); + let sender = connect_with_options(&channel, "sender", settings.clone()).await.unwrap(); + let receiver = connect_with_options(&channel, "receiver", settings).await.unwrap(); + let (client, connection) = database().await; + let packet = frame(1); + sender.publish(std::slice::from_ref(&packet)).await.unwrap(); + let sentinel = frame(2); + client + .execute( + "INSERT INTO stegrdb_relay.frames(channel,id,sender,payload,created_at) VALUES ($1,$2,'sender',$3,statement_timestamp()-interval '5 minutes')", + &[&channel, &sentinel.id, &sentinel.bytes.as_ref()], + ) + .await + .unwrap(); + client + .execute( + "UPDATE stegrdb_relay.frames SET created_at=statement_timestamp()-interval '5 minutes' WHERE channel=$1", + &[&channel], + ) + .await + .unwrap(); + wait_for_removal(&client, &channel, sentinel.id).await; + let deliveries = receiver.receive(10).await.unwrap(); + assert_eq!(deliveries.len(), 1); + assert_eq!(deliveries[0].frame, packet); + let old_receipt = deliveries[0].receipt.clone(); + receiver.acknowledge(std::slice::from_ref(&old_receipt)).await.unwrap(); + wait_for_removal(&client, &channel, packet.id).await; + sender.publish(std::slice::from_ref(&packet)).await.unwrap(); + receiver.acknowledge(std::slice::from_ref(&old_receipt)).await.unwrap(); + let renewed = receiver.receive(10).await.unwrap(); + assert_eq!(renewed[0].frame, packet); + assert_ne!(renewed[0].receipt, old_receipt); + receiver.acknowledge(&[renewed[0].receipt.clone()]).await.unwrap(); + assert!(receiver.receive(10).await.unwrap().is_empty()); + remove_channel(&client, &channel).await; + connection.abort(); +} + +#[tokio::test] +#[ignore = "専用DBでscripts/test-postgres.shを実行する"] +async fn a_restarted_node_updates_its_retention_choice_without_replaying_acknowledged_frames() { + let channel = Uuid::new_v4().to_string(); + let sender = connect(&channel, "sender").await.unwrap(); + let receiver = connect(&channel, "receiver").await.unwrap(); + let packet = frame(3); + sender.publish(&[packet]).await.unwrap(); + let deliveries = receiver.receive(10).await.unwrap(); + receiver.acknowledge(&[deliveries[0].receipt.clone()]).await.unwrap(); + drop(receiver); + let receiver = tokio::time::timeout(CLEANUP_DEADLINE, async { + loop { + if let Ok(receiver) = connect_with_options(&channel, "receiver", json!({"retention_ms": 60_000})).await { + break receiver; + } + tokio::time::sleep(CLEANUP_POLL).await; + } + }) + .await + .unwrap(); + assert!(receiver.receive(10).await.unwrap().is_empty()); + let (client, connection) = database().await; + let retention: i64 = client.query_one("SELECT retention_ms FROM stegrdb_relay.nodes WHERE channel=$1 AND node_id='receiver'", &[&channel]).await.unwrap().get(0); + assert_eq!(retention, 60_000); + remove_channel(&client, &channel).await; + connection.abort(); +}