diff --git a/.dockerignore b/.dockerignore index 8011b19..83e544f 100644 --- a/.dockerignore +++ b/.dockerignore @@ -1,6 +1,9 @@ -# The image is built from the manifest and the sources alone. +# The images are built from the manifests and the sources alone: the +# workspace's crates are part of the manifest, their build output is not. * !Cargo.toml !Cargo.lock !clippy.toml !src +!crates +crates/*/target diff --git a/.github/workflows/ci-pr.yml b/.github/workflows/ci-pr.yml index c570de2..2adb0ac 100644 --- a/.github/workflows/ci-pr.yml +++ b/.github/workflows/ci-pr.yml @@ -46,14 +46,17 @@ jobs: - name: Install rust cache uses: Swatinem/rust-cache@v2.7.7 - - name: Clippy, warnings denied - run: cargo clippy --all-targets --locked -- -D warnings + - name: Clippy, warnings denied, every crate of the workspace + run: cargo clippy --workspace --all-targets --locked -- -D warnings test: name: Tests (accuracy) runs-on: ubuntu-latest env: STREAMGRES_PG_DSN: postgresql://postgres:postgres@localhost:5432/postgres + # The Diesel wrapper's and the demo's tests read and write their own + # tables in the same database. + DATABASE_URL: postgresql://postgres:postgres@localhost:5432/postgres steps: - name: Checkout repository uses: actions/checkout@v4 @@ -77,8 +80,8 @@ jobs: -c wal_level=logical -c max_replication_slots=32 -c max_wal_senders=32 until docker exec pg pg_isready -U postgres; do sleep 1; done - - name: Unit, scenario, randomized and live PostgreSQL tests - run: cargo test --release --locked + - name: Unit, scenario, randomized and live PostgreSQL tests, every crate + run: cargo test --release --locked --workspace perf: name: Benchmark against the base branch (performance) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 5384412..360f94b 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -31,22 +31,24 @@ Build and test: ```bash cargo build -cargo test # unit and scenario tests -cargo clippy --all-targets -- -D warnings +cargo test --workspace # unit and scenario tests, every crate +cargo clippy --workspace --all-targets -- -D warnings cargo run --bin streamgres # Streamgres demo: every check must end in PASS cargo run --release --bin bench # benchmarks cargo run --release --bin load # the server over the wire: latency, cores, memory (needs STREAMGRES_PG_DSN) ``` The tests against a real PostgreSQL run when `STREAMGRES_PG_DSN` is set and -report themselves skipped otherwise. To run them the way CI does: +report themselves skipped otherwise; the Diesel wrapper's and the demo's +tests read `DATABASE_URL`. To run them the way CI does: ```bash docker run -d --name pg -e POSTGRES_PASSWORD=postgres -p 5432:5432 postgres:17 \ -c wal_level=logical -c max_replication_slots=32 -c max_wal_senders=32 STREAMGRES_PG_DSN=postgresql://postgres:postgres@localhost:5432/postgres \ - cargo test --release +DATABASE_URL=postgresql://postgres:postgres@localhost:5432/postgres \ + cargo test --release --workspace ``` To run the server locally, follow the [Quick start](README.md#quick-start) in @@ -58,8 +60,8 @@ Every pull request runs: | Job | What it runs | | --- | --- | -| Lints | `cargo clippy --all-targets --locked -- -D warnings` | -| Tests | `cargo test --release --locked` against PostgreSQL 17 | +| Lints | `cargo clippy --workspace --all-targets --locked -- -D warnings` | +| Tests | `cargo test --release --locked --workspace` against PostgreSQL 17 | | Benchmark | `bench` on your branch and the base branch, five alternating rounds, compared by `scripts/bench-compare.py` | | Image | builds `docker/server/Dockerfile` and drives it end to end with `scripts/smoke.mjs` | diff --git a/Cargo.lock b/Cargo.lock index d2b9643..5fce639 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -23,6 +23,15 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "aae1277d39aeec15cb388266ecc24b11c80469deae6067e17a1a7aa9e5c1f234" +[[package]] +name = "aho-corasick" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c982642fa9e8606056828ee9a8505737230110bb1099153c79efe865c59d12ba" +dependencies = [ + "memchr", +] + [[package]] name = "aligned-vec" version = "0.6.4" @@ -47,6 +56,56 @@ dependencies = [ "libc", ] +[[package]] +name = "anstream" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "824a212faf96e9acacdbd09febd34438f8f711fb84e09a8916013cd7815ca28d" +dependencies = [ + "anstyle", + "anstyle-parse", + "anstyle-query", + "anstyle-wincon", + "colorchoice", + "is_terminal_polyfill", + "utf8parse", +] + +[[package]] +name = "anstyle" +version = "1.0.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" + +[[package]] +name = "anstyle-parse" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52ce7f38b242319f7cabaa6813055467063ecdc9d355bbb4ce0c68908cd8130e" +dependencies = [ + "utf8parse", +] + +[[package]] +name = "anstyle-query" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "anstyle-wincon" +version = "3.0.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" +dependencies = [ + "anstyle", + "once_cell_polyfill", + "windows-sys 0.61.2", +] + [[package]] name = "anyhow" version = "1.0.104" @@ -117,10 +176,10 @@ dependencies = [ "serde_json", "serde_path_to_error", "serde_urlencoded", - "sha1", + "sha1 0.10.7", "sync_wrapper", "tokio", - "tokio-tungstenite", + "tokio-tungstenite 0.29.0", "tower", "tower-layer", "tower-service", @@ -173,6 +232,18 @@ version = "0.23.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ac07cdecf99051d9a5238b80f35af32cdeba5b336e55d957b318b50137e18da5" +[[package]] +name = "bb8" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "457d7ed3f888dfd2c7af56d4975cade43c622f74bdcddfed6d4352f57acc6310" +dependencies = [ + "futures-util", + "parking_lot", + "portable-atomic", + "tokio", +] + [[package]] name = "bitflags" version = "2.13.1" @@ -255,16 +326,65 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1aa79e62e7697b8e29b513a68abacf485adcd1fe8284a4316c5ae868e6633327" dependencies = [ "iana-time-zone", + "js-sys", "num-traits", + "serde", + "wasm-bindgen", "windows-link", ] +[[package]] +name = "clap" +version = "4.6.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aa8876b300ab35ba921adea3dfd70157a46249b33f95c9084ae5709785478946" +dependencies = [ + "clap_builder", + "clap_derive", +] + +[[package]] +name = "clap_builder" +version = "4.6.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0797fb7aeb1406c84efac526901f7ec3ead2124f946b494e72879d4b54704d" +dependencies = [ + "anstream", + "anstyle", + "clap_lex", + "strsim", +] + +[[package]] +name = "clap_derive" +version = "4.6.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9c751b79415d4e559e3d1fcf128e09e720eb673a06d26cf6f392d37d75b66e0" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn 3.0.4", +] + +[[package]] +name = "clap_lex" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c133bc6a41be0d194c306b5506d15e6feeea7b1d6604bd3f8310dfb2ca96486" + [[package]] name = "cmov" version = "0.5.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0c9ea0ac24bc397ab3c98583a3c9ba74fa56b09a4449bbe172b9b1ddb016027a" +[[package]] +name = "colorchoice" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570" + [[package]] name = "combine" version = "4.6.8" @@ -275,6 +395,46 @@ dependencies = [ "memchr", ] +[[package]] +name = "console-api" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8599749b6667e2f0c910c1d0dff6901163ff698a52d5a39720f61b5be4b20d3" +dependencies = [ + "futures-core", + "prost", + "prost-types", + "tonic", + "tonic-prost", + "tracing-core", +] + +[[package]] +name = "console-subscriber" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb4915b7d8dd960457a1b6c380114c2944f728e7c65294ab247ae6b6f1f37592" +dependencies = [ + "console-api", + "crossbeam-channel", + "crossbeam-utils", + "futures-task", + "hdrhistogram", + "humantime", + "hyper-util", + "prost", + "prost-types", + "serde", + "serde_json", + "thread_local", + "tokio", + "tokio-stream", + "tonic", + "tracing", + "tracing-core", + "tracing-subscriber", +] + [[package]] name = "const-oid" version = "0.10.2" @@ -324,6 +484,30 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "crossbeam-channel" +version = "0.5.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "98b0cc327b5bc766e7fda9c9260cc0fa81b43a8e240440422dff70788e3f9ef1" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc74980687109a3b14c72fd458107bf0baa1da1a1a805e178d15501ba9b86d9d" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a31eee39dddec8330830986fcd7625edb5a24ec90ea038215273bbc3adb08ac6" + [[package]] name = "crypto-common" version = "0.1.7" @@ -352,6 +536,41 @@ dependencies = [ "cmov", ] +[[package]] +name = "darling" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc7f46116c46ff9ab3eb1597a45688b6715c6e628b5c133e288e709a29bcb4ee" +dependencies = [ + "darling_core", + "darling_macro", +] + +[[package]] +name = "darling_core" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d00b9596d185e565c2207a0b01f8bd1a135483d02d9b7b0a54b11da8d53412e" +dependencies = [ + "fnv", + "ident_case", + "proc-macro2", + "quote", + "strsim", + "syn 2.0.119", +] + +[[package]] +name = "darling_macro" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead" +dependencies = [ + "darling_core", + "quote", + "syn 2.0.119", +] + [[package]] name = "dary_heap" version = "0.3.9" @@ -373,6 +592,77 @@ dependencies = [ "uuid", ] +[[package]] +name = "diesel" +version = "2.2.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff3e1edb1f37b4953dd5176916347289ed43d7119cc2e6c7c3f7849ff44ea506" +dependencies = [ + "bitflags", + "byteorder", + "chrono", + "diesel_derives", + "itoa", + "pq-sys", + "r2d2", + "serde_json", +] + +[[package]] +name = "diesel-async" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7fcc26599f590c7e5b182a05061cfb445f216bb069df72eb31f38cffde8ca598" +dependencies = [ + "bb8", + "diesel", + "futures-core", + "futures-util", + "scoped-futures", + "tokio", + "tokio-postgres", +] + +[[package]] +name = "diesel-streamgres" +version = "0.1.0" +dependencies = [ + "base64 0.23.1", + "chrono", + "diesel", + "diesel-async", + "futures-util", + "moka", + "rustls", + "rustls-native-certs", + "scoped-futures", + "serde_json", + "tokio", + "tokio-tungstenite 0.30.0", +] + +[[package]] +name = "diesel_derives" +version = "2.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b96984c469425cb577bf6f17121ecb3e4fe1e81de5d8f780dd372802858d756" +dependencies = [ + "diesel_table_macro_syntax", + "dsl_auto_type", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "diesel_table_macro_syntax" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "209c735641a413bc68c4923a9d6ad4bcb3ca306b794edaa7eb0b3228a99ffb25" +dependencies = [ + "syn 2.0.119", +] + [[package]] name = "digest" version = "0.10.7" @@ -407,6 +697,20 @@ dependencies = [ "syn 3.0.4", ] +[[package]] +name = "dsl_auto_type" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "139ae9aca7527f85f26dd76483eb38533fd84bd571065da1739656ef71c5ff5b" +dependencies = [ + "darling", + "either", + "heck", + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "either" version = "1.18.0" @@ -496,6 +800,12 @@ dependencies = [ "zlib-rs", ] +[[package]] +name = "fnv" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" + [[package]] name = "foldhash" version = "0.2.0" @@ -547,6 +857,17 @@ version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "53c0fa8157de1303bfffdaa1cc2a673bfffb60102f76b0ef4441659124373fed" +[[package]] +name = "futures-macro" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9fb9654ba8355388abeb8dcb4fc62f511300867002afc858860463bdd9fe0c44" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.4", +] + [[package]] name = "futures-sink" version = "0.3.34" @@ -567,6 +888,7 @@ checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc" dependencies = [ "futures-core", "futures-io", + "futures-macro", "futures-sink", "futures-task", "memchr", @@ -638,6 +960,25 @@ dependencies = [ "stable_deref_trait", ] +[[package]] +name = "h2" +version = "0.4.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d29020232d6aa3fb1daca64c1127cf662cf97f254ae16c18c05b8ab635fc118" +dependencies = [ + "atomic-waker", + "bytes", + "fnv", + "futures-core", + "futures-sink", + "http", + "indexmap", + "slab", + "tokio", + "tokio-util", + "tracing", +] + [[package]] name = "hashbrown" version = "0.16.1" @@ -649,6 +990,31 @@ dependencies = [ "foldhash", ] +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "hdrhistogram" +version = "7.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f49d1053f4708f0af3cf9fc5bffc7e68a914a3c45becb231c80068c9c3f78bea" +dependencies = [ + "base64 0.22.1", + "byteorder", + "flate2", + "nom", + "num-traits", +] + +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + [[package]] name = "hmac" version = "0.12.1" @@ -712,6 +1078,12 @@ version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" +[[package]] +name = "humantime" +version = "2.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15cdd26707701c53297e2fa6afb323d55fbc1d0810c3aec078ae3ef0424c3c15" + [[package]] name = "hybrid-array" version = "0.4.15" @@ -731,6 +1103,7 @@ dependencies = [ "bytes", "futures-channel", "futures-core", + "h2", "http", "http-body", "httparse", @@ -758,6 +1131,19 @@ dependencies = [ "webpki-roots", ] +[[package]] +name = "hyper-timeout" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b90d566bffbce6a75bd8b09a05aa8c2cb1fabb6cb348f8840c9e4c90a0d83b0" +dependencies = [ + "hyper", + "hyper-util", + "pin-project-lite", + "tokio", + "tower-service", +] + [[package]] name = "hyper-util" version = "0.1.20" @@ -888,6 +1274,12 @@ dependencies = [ "zerovec", ] +[[package]] +name = "ident_case" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39" + [[package]] name = "idna" version = "1.1.0" @@ -909,12 +1301,28 @@ dependencies = [ "icu_properties", ] +[[package]] +name = "indexmap" +version = "2.14.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc4e190f5d26ca7051642629da2c52fc03bde85a03197c99408dcd291734c855" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", +] + [[package]] name = "ipnet" version = "2.12.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "791930b43c0d5973160d90a8f3894509f2b273430f5c5c73b668636d0287c5c0" +[[package]] +name = "is_terminal_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695" + [[package]] name = "itertools" version = "0.14.0" @@ -1038,7 +1446,7 @@ version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ff7a10e427698aef6eef269482776debfef63384d30f13aad39a1a95e0e098fd" dependencies = [ - "hashbrown", + "hashbrown 0.16.1", "no_std_io2", "rle-decode-fast", ] @@ -1109,6 +1517,15 @@ dependencies = [ "tracing", ] +[[package]] +name = "matchers" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1525a2a28c7f4fa0fc98bb91ae755d1e2d1505079e05539e35bc876b5d65ae9" +dependencies = [ + "regex-automata", +] + [[package]] name = "matchit" version = "0.8.4" @@ -1176,6 +1593,23 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "moka" +version = "0.12.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4293f18e7567a1caf3c584855554377025c65e0aa445344d04171f5ad63d19b9" +dependencies = [ + "crossbeam-channel", + "crossbeam-epoch", + "crossbeam-utils", + "equivalent", + "parking_lot", + "portable-atomic", + "smallvec", + "tagptr", + "uuid", +] + [[package]] name = "nix" version = "0.31.3" @@ -1197,6 +1631,15 @@ dependencies = [ "memchr", ] +[[package]] +name = "nom" +version = "8.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df9761775871bdef83bee530e60050f7e54b1105350d6884eb0fb4f46c2f9405" +dependencies = [ + "memchr", +] + [[package]] name = "num" version = "0.4.3" @@ -1313,6 +1756,12 @@ version = "1.21.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" +[[package]] +name = "once_cell_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" + [[package]] name = "openssl-probe" version = "0.2.1" @@ -1395,19 +1844,39 @@ dependencies = [ name = "phf" version = "0.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c1562dc717473dbaa4c1f85a36410e03c047b2e7df7f45ee938fbef64ae7fadf" +checksum = "c1562dc717473dbaa4c1f85a36410e03c047b2e7df7f45ee938fbef64ae7fadf" +dependencies = [ + "phf_shared", + "serde", +] + +[[package]] +name = "phf_shared" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e57fef6bc5981e38c2ce2d63bfa546861309f875b8a75f092d1d54ae2d64f266" +dependencies = [ + "siphasher", +] + +[[package]] +name = "pin-project" +version = "1.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2466b2336ed02bcdca6b294417127b90ec92038d1d5c4fbeac971a922e0e0924" dependencies = [ - "phf_shared", - "serde", + "pin-project-internal", ] [[package]] -name = "phf_shared" -version = "0.13.1" +name = "pin-project-internal" +version = "1.1.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e57fef6bc5981e38c2ce2d63bfa546861309f875b8a75f092d1d54ae2d64f266" +checksum = "c96395f0a926bc13b1c17622aaddda1ecb55d49c8f1bf9777e4d877800a43f8b" dependencies = [ - "siphasher", + "proc-macro2", + "quote", + "syn 2.0.119", ] [[package]] @@ -1416,6 +1885,18 @@ version = "0.2.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" +[[package]] +name = "pkg-config" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f6b464fbc74e149a392436b17d523f769e057cb6877f6a5c4618bc6f11800548" + +[[package]] +name = "portable-atomic" +version = "1.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05c8b63e8d9609db387f0324918f81d68fe27748f084ef092fb35954d0539a85" + [[package]] name = "postgres-protocol" version = "0.6.12" @@ -1478,6 +1959,17 @@ dependencies = [ "zerocopy", ] +[[package]] +name = "pq-sys" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9152654059e4fb5b66a4052acf3ee43a7a7132332e44cdb72a81ee93bd038950" +dependencies = [ + "libc", + "pkg-config", + "vcpkg", +] + [[package]] name = "proc-macro2" version = "1.0.107" @@ -1510,6 +2002,15 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "prost-types" +version = "0.14.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f94967dc7688f3054c7fac87473ffae4cc4c3904800e2d9f5b857246d8963b0a" +dependencies = [ + "prost", +] + [[package]] name = "pyroscope" version = "2.1.1" @@ -1617,6 +2118,17 @@ version = "6.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" +[[package]] +name = "r2d2" +version = "0.8.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51de85fb3fb6524929c8a2eb85e6b6d363de4e8c48f9e2c2eac4944abc181c93" +dependencies = [ + "log", + "parking_lot", + "scheduled-thread-pool", +] + [[package]] name = "rand" version = "0.9.5" @@ -1681,6 +2193,23 @@ dependencies = [ "bitflags", ] +[[package]] +name = "regex-automata" +version = "0.4.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + [[package]] name = "reqwest" version = "0.12.28" @@ -1924,6 +2453,24 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "scheduled-thread-pool" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3cbc66816425a074528352f5789333ecff06ca41b36b0b0efdfbb29edc391a19" +dependencies = [ + "parking_lot", +] + +[[package]] +name = "scoped-futures" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b24aae2d0636530f359e9d5ef0c04669d11c5e756699b27a6a6d845d8329091" +dependencies = [ + "pin-project-lite", +] + [[package]] name = "scopeguard" version = "1.2.0" @@ -1991,9 +2538,9 @@ dependencies = [ [[package]] name = "serde_json" -version = "1.0.151" +version = "1.0.154" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c841b55ecdae098c80dcae9cf767f6f8a0c2cdb3416bbef72181df4d0fe73f14" +checksum = "e7e9cc8b1b85264074fbcc02a88680c4096b1e47df8f739dceb03bf482f04bd6" dependencies = [ "itoa", "memchr", @@ -2036,6 +2583,17 @@ dependencies = [ "digest 0.10.7", ] +[[package]] +name = "sha1" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aacc4cc499359472b4abe1bf11d0b12e688af9a805fa5e3016f9a386dc2d0214" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.1", + "digest 0.11.3", +] + [[package]] name = "sha2" version = "0.10.9" @@ -2058,6 +2616,15 @@ dependencies = [ "digest 0.11.3", ] +[[package]] +name = "sharded-slab" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6" +dependencies = [ + "lazy_static", +] + [[package]] name = "shlex" version = "2.0.1" @@ -2139,6 +2706,58 @@ version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" +[[package]] +name = "streamgres" +version = "0.1.0" +dependencies = [ + "arc-swap", + "axum", + "base64 0.22.1", + "bytes", + "chrono", + "console-subscriber", + "fallible-iterator 0.2.0", + "futures-util", + "jemalloc_pprof", + "libc", + "log", + "memchr", + "once_cell", + "percent-encoding", + "pgoutput", + "pgwire-replication", + "postgres-protocol", + "pyroscope", + "reqwest 0.12.28", + "rustc-hash", + "rustls", + "serde", + "serde_json", + "socket2", + "tikv-jemallocator", + "tokio", + "tokio-postgres", + "tokio-tungstenite 0.29.0", + "uuid", +] + +[[package]] +name = "streamgres-demo" +version = "0.1.0" +dependencies = [ + "axum", + "chrono", + "clap", + "diesel", + "diesel-async", + "diesel-streamgres", + "futures-util", + "rand 0.9.5", + "serde", + "serde_json", + "tokio", +] + [[package]] name = "stringprep" version = "0.1.5" @@ -2150,6 +2769,12 @@ dependencies = [ "unicode-properties", ] +[[package]] +name = "strsim" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" + [[package]] name = "subtle" version = "2.6.1" @@ -2220,6 +2845,12 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "tagptr" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b2093cf4c8eb1e67749a6762251bc9cd836b6fc171623bd0a9d324d37af2417" + [[package]] name = "tempfile" version = "3.27.0" @@ -2273,6 +2904,15 @@ dependencies = [ "syn 3.0.4", ] +[[package]] +name = "thread_local" +version = "1.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ad99c4c6d32803332c548b1af0540b357b3f5fc0be8f6c6bfe8b2e6ae784070" +dependencies = [ + "cfg-if", +] + [[package]] name = "tikv-jemalloc-ctl" version = "0.7.0" @@ -2331,9 +2971,9 @@ checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" [[package]] name = "tokio" -version = "1.53.1" +version = "1.53.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "202caea871b69668250d242070849eb495be178ed697a3e98aebce5bc81a0bed" +checksum = "ce3335fa71841cda333a58d7615b03901380ecf09d59b3296d21f8bbac0dde4e" dependencies = [ "bytes", "libc", @@ -2343,6 +2983,7 @@ dependencies = [ "signal-hook-registry", "socket2", "tokio-macros", + "tracing", "windows-sys 0.61.2", ] @@ -2393,6 +3034,17 @@ dependencies = [ "tokio", ] +[[package]] +name = "tokio-stream" +version = "0.1.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a3d06f0b082ba57c26b79407372e57cf2a1e28124f78e9479fe80322cf53420b" +dependencies = [ + "futures-core", + "pin-project-lite", + "tokio", +] + [[package]] name = "tokio-tungstenite" version = "0.29.0" @@ -2402,7 +3054,23 @@ dependencies = [ "futures-util", "log", "tokio", - "tungstenite", + "tungstenite 0.29.0", +] + +[[package]] +name = "tokio-tungstenite" +version = "0.30.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "17a073bfed563fa236697a068031408a93cd9522e08abf9933ead3e73411bd71" +dependencies = [ + "futures-util", + "log", + "rustls", + "rustls-native-certs", + "rustls-pki-types", + "tokio", + "tokio-rustls", + "tungstenite 0.30.0", ] [[package]] @@ -2419,6 +3087,46 @@ dependencies = [ "tokio", ] +[[package]] +name = "tonic" +version = "0.14.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac2a5518c70fa84342385732db33fb3f44bc4cc748936eb5833d2df34d6445ef" +dependencies = [ + "async-trait", + "axum", + "base64 0.22.1", + "bytes", + "h2", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-timeout", + "hyper-util", + "percent-encoding", + "pin-project", + "socket2", + "sync_wrapper", + "tokio", + "tokio-stream", + "tower", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "tonic-prost" +version = "0.14.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "50849f68853be452acf590cde0b146665b8d507b3b8af17261df47e02c209ea0" +dependencies = [ + "bytes", + "prost", + "tonic", +] + [[package]] name = "tower" version = "0.5.3" @@ -2427,9 +3135,12 @@ checksum = "ebe5ef63511595f1344e2d5cfa636d973292adc0eec1f0ad45fae9f0851ab1d4" dependencies = [ "futures-core", "futures-util", + "indexmap", "pin-project-lite", + "slab", "sync_wrapper", "tokio", + "tokio-util", "tower-layer", "tower-service", "tracing", @@ -2495,6 +3206,22 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" dependencies = [ "once_cell", + "valuable", +] + +[[package]] +name = "tracing-subscriber" +version = "0.3.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb7f578e5945fb242538965c2d0b04418d38ec25c79d160cd279bf0731c8d319" +dependencies = [ + "matchers", + "once_cell", + "regex-automata", + "sharded-slab", + "thread_local", + "tracing", + "tracing-core", ] [[package]] @@ -2515,7 +3242,25 @@ dependencies = [ "httparse", "log", "rand 0.9.5", - "sha1", + "sha1 0.10.7", + "thiserror 2.0.20", +] + +[[package]] +name = "tungstenite" +version = "0.30.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e48ac77174b19c110a50ab2128b24215ac9cb40e0e12e093fb602d175c569d22" +dependencies = [ + "bytes", + "data-encoding", + "http", + "httparse", + "log", + "rand 0.10.2", + "rustls", + "rustls-pki-types", + "sha1 0.11.0", "thiserror 2.0.20", ] @@ -2582,6 +3327,12 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" +[[package]] +name = "utf8parse" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" + [[package]] name = "uuid" version = "1.26.1" @@ -2593,6 +3344,18 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "valuable" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" + +[[package]] +name = "vcpkg" +version = "0.2.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "accd4ea62f7bb7a82fe23066fb0957d48ef677f6eeb8215f372f52e48bb32426" + [[package]] name = "version_check" version = "0.9.5" @@ -2941,40 +3704,6 @@ version = "0.6.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3ad82d2a33cdc9674dc7465672f271e096168fcdbe0f799d9e6db8c5892679dc" -[[package]] -name = "streamgres" -version = "0.1.0" -dependencies = [ - "arc-swap", - "axum", - "base64 0.22.1", - "bytes", - "chrono", - "fallible-iterator 0.2.0", - "futures-util", - "jemalloc_pprof", - "libc", - "log", - "memchr", - "once_cell", - "percent-encoding", - "pgoutput", - "pgwire-replication", - "postgres-protocol", - "pyroscope", - "reqwest 0.12.28", - "rustc-hash", - "rustls", - "serde", - "serde_json", - "socket2", - "tikv-jemallocator", - "tokio", - "tokio-postgres", - "tokio-tungstenite", - "uuid", -] - [[package]] name = "yoke" version = "0.8.3" diff --git a/Cargo.toml b/Cargo.toml index fe758ee..38fff5b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,3 +1,6 @@ +[workspace] +members = ["crates/*"] + [package] name = "streamgres" version = "0.1.0" @@ -5,7 +8,15 @@ edition = "2024" rust-version = "1.88" license = "Apache-2.0" +[features] +# `tokio-console` support for local debugging: build with +# `RUSTFLAGS="--cfg tokio_unstable" cargo build --features console`, run +# the server, then `tokio-console` attaches on 127.0.0.1:6669. +console = ["dep:console-subscriber"] + [dependencies] +console-subscriber = { version = "0.5", optional = true } + # The hasher of the engine's integer-keyed maps (subscription, row, read # and tree ids): one multiply per key instead of SipHash (see # `model::ids`). Keys that are user data keep the default hasher. @@ -98,3 +109,8 @@ log = "0.4" # the same. [profile.release] debug = "limited" + +[lints.rust] +# `tokio_unstable` is set with `RUSTFLAGS="--cfg tokio_unstable"` for +# `tokio-console` builds (see the `console` feature). +unexpected_cfgs = { level = "warn", check-cfg = ['cfg(tokio_unstable)'] } diff --git a/README.md b/README.md index 1a1eafb..346eba7 100644 --- a/README.md +++ b/README.md @@ -117,6 +117,30 @@ The server listens on `0.0.0.0:4848`. Point the client's `server` option at `.env.example`. `GET /health` returns `200` once the server is ready to serve queries. +### Try it in five minutes + +`demo/docker-compose.yml` brings up Postgres, the server and a small +ticket-desk backend built on the Diesel wrapper, with a background writer +keeping the data moving; `demo/README.md` walks through it with curl, and + shows the cache's counters live. The backend is +`crates/streamgres-demo`, which is also the wrapper's test bed. + +### Server-to-server clients + +A backend that caches what it reads does not speak the Zero protocol. It +connects to `ws://:4848//s2s/v1` (with the bearer token in +`STREAMGRES_S2S_TOKEN`, through TLS in front of the server off a private +network) and subscribes the SQL its ORM rendered, `$n` binds and all; +the server answers `accepted` with the hydrated rows (per table for a +join; `bytea` as hex text, `OFFSET` left for the client to apply) or +`refused`, then sends `delta` with the rows each committed transaction +put or deleted, so the backend never reads Postgres for a query it +holds. The protocol is documented in +`src/client/s2s.rs`. The Diesel wrapper that uses it lives in this +repository as [`crates/diesel-streamgres`](crates/diesel-streamgres): +a drop-in `PgConnection` replacement whose reads come from an in-memory +cache the engine keeps fresh. + Configuration uses `STREAMGRES_*` environment variables, each documented in [`.env.example`](.env.example). Existing `XYNE_SYNC_*` variables are still accepted as fallbacks; a `STREAMGRES_*` value takes precedence when both are diff --git a/crates/diesel-streamgres/.gitignore b/crates/diesel-streamgres/.gitignore new file mode 100644 index 0000000..2f7896d --- /dev/null +++ b/crates/diesel-streamgres/.gitignore @@ -0,0 +1 @@ +target/ diff --git a/crates/diesel-streamgres/Cargo.toml b/crates/diesel-streamgres/Cargo.toml new file mode 100644 index 0000000..eb1d57a --- /dev/null +++ b/crates/diesel-streamgres/Cargo.toml @@ -0,0 +1,32 @@ +[package] +name = "diesel-streamgres" +version = "0.1.0" +edition = "2024" +rust-version = "1.88" +license = "Apache-2.0" +description = "A Diesel connection wrapper serving reads from a cache the Streamgres engine keeps fresh" + +[features] +# `AsyncSyncedConnection`, the same wrapper over diesel-async's +# `AsyncPgConnection`. +async = ["dep:diesel-async", "diesel-async/postgres", "diesel-async/pool"] + +[dependencies] +base64 = "0.23.1" +diesel = { version = "2.2.10", features = ["postgres", "r2d2", "i-implement-a-third-party-backend-and-opt-into-breaking-changes"] } +diesel-async = { version = "0.6", features = ["postgres"], optional = true } +futures-util = "0.3.34" +moka = { version = "0.12", features = ["sync"] } +serde_json = "1.0.154" +tokio = { version = "1.53.2", features = ["rt-multi-thread", "macros", "time", "net", "sync"] } +# `wss://` through rustls with the ring provider and the system's roots, +# configured per socket rather than installed for the whole process. +tokio-tungstenite = { version = "0.30.0", features = ["rustls-tls-native-roots"] } +rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] } +rustls-native-certs = "0.8" + +[dev-dependencies] +chrono = "0.4.45" +scoped-futures = "0.1" +diesel = { version = "2.2.10", features = ["chrono", "postgres", "r2d2", "serde_json"] } +serde_json = "1.0.154" diff --git a/crates/diesel-streamgres/README.md b/crates/diesel-streamgres/README.md new file mode 100644 index 0000000..96d16e1 --- /dev/null +++ b/crates/diesel-streamgres/README.md @@ -0,0 +1,70 @@ +# diesel_streamgres + +A Diesel connection wrapper for Postgres that an application swaps in for +`PgConnection`. Reads are served from a local cache the Streamgres engine +keeps current; writes, transactions and queries routed to Postgres go to +Postgres as before. + +A read on a connection routed to Streamgres, outside a transaction, is +keyed by its SQL and bind bytes and subscribed with the engine over its +S2S endpoint. When every projected column has a type the client can +encode exactly, the engine serves the rows and keeps them current with +each commit: no Postgres read, and a commit lands in the cache in about +0.3 ms locally. Otherwise the rows are read from Postgres once and read +again when the engine reports a change. Queries the engine refuses go to +Postgres. An engine restart resubscribes and refills, and asks again +about what was refused. See `docs/design.md`. + +## Using it + +```rust +// Pool size and subscription cap from STREAMGRES_S2S_CONNECTIONS and +// STREAMGRES_S2S_MAX_SUBSCRIPTIONS (defaults 2 and 10000); the bearer +// token from STREAMGRES_S2S_TOKEN when the engine asks for one. Use +// wss:// (or a TLS-terminating proxy) unless the engine is on this +// machine or a private network: the token travels in the upgrade. +let engine = EngineClient::connect("wss://engine:4848/zero/s2s/v1"); +for warning in engine.warnings() { + log::warn!("{warning}"); +} +let cache = Arc::new(Cache::new(256 * 1024 * 1024, Arc::new(engine.clone()))); +engine.attach(&cache); +// One cache per database (one engine streams one database), a route per pool. +diesel_streamgres::install(Settings::single( + cache, + |url| if url.contains("replica") { Route::Streamgres } else { Route::Postgres }, +)).ok(); +// Several databases: Settings::by_url(|url| Some(Attachment { cache: ..., route: ... })) +// then use SyncedConnection where PgConnection was, pools included. +// One query can override its pool's route: +// configs::table.find(key).via(Route::Postgres).first(conn) +// conn.with_route(Route::Streamgres, |c| configs::table.find(key).first(c)) +``` + +## diesel-async + +With the `async` feature, `AsyncSyncedConnection` wraps diesel-async's +`AsyncPgConnection` the same way and pools through bb8 or deadpool: + +```rust +let conn = AsyncSyncedConnection::establish(&url).await?; // picks up the installed settings +configs::table.find(key).first::(&mut conn).await?; // cached +configs::table.find(key).via(Route::Postgres).first::(&mut conn).await?; +``` + +## Live tests + +With an engine streaming the test database (see `docs/design.md` for a +local setup): + +```sh +STREAMGRES_S2S_URL=ws://127.0.0.1:4848/zero/s2s/v1 \ +DATABASE_URL=postgres://sg@127.0.0.1:5434/sg \ +cargo test --test engine_live +``` + +```sh +# needs a Postgres; default is the local socket database diesel_streamgres_test +createdb diesel_streamgres_test +cargo test +``` diff --git a/crates/diesel-streamgres/docs/design.md b/crates/diesel-streamgres/docs/design.md new file mode 100644 index 0000000..c61470c --- /dev/null +++ b/crates/diesel-streamgres/docs/design.md @@ -0,0 +1,550 @@ +# diesel_streamgres: design + +`diesel_streamgres` is a Diesel connection that serves an application's +reads from a local cache the Streamgres engine keeps fresh, and sends +everything else to Postgres. The application swaps its connection type +and changes nothing else. + +```text + application code (Diesel DSL, sql_query, transactions, pools) + | + v + SyncedConnection or AsyncSyncedConnection + | load / execute / batch_execute + | + +-- writes, transactions, non-SELECT, flagged Postgres --> Postgres + | + +-- cacheable read + | + v + Cache (moka, keyed by SQL + bind bytes) <---- delta / touched + | miss ^ + v | + EngineClient (pool of WebSockets) --- subscribe ---> Streamgres engine + | accepted (rows) / delta / touched / refused | + v | logical + Materialized store per query (rows mode) | replication + v + Postgres +``` + +## Modules + +| module | role | +|---|---| +| `lib.rs` | `SyncedConnection`, `Settings`, `install`, the interception of Diesel's three connection methods | +| `asynchronous.rs` | `AsyncSyncedConnection`, the same wrapper over diesel-async (feature `async`) | +| `route.rs` | `Route`, `RouteDsl`, `Via`, `Served`: where a read goes and where it came from | +| `key.rs` | `QueryKey`: what identifies a read, and a refusal | +| `cache.rs` | `Cache`, `Limits`, `Stats`, the `Live` trait, reservations and generations | +| `subscriptions.rs` | the `Subscriptions` trait between the cache and the engine side | +| `engine.rs` | `EngineClient`: the pool of sockets to the engine, `Config` | +| `materialized.rs` | a query's result kept from the engine's rows; joins composed here | +| `types.rs` | exact column types of a table, from `pg_attribute` | +| `encode.rs` | engine JSON cells to Postgres binary, per column type | +| `row.rs` | `OwnedRow`: a row that outlives its statement and decodes like a live one | +| `cursor.rs` | one cursor type for rows from the cache and rows from Postgres | + +## The interception surface + +Diesel funnels every operation into three methods on the connection type: + +| application call | connection method | +|------------------------------------------|---------------------------------------| +| `load`, `first`, `get_result(s)`, `count`| `LoadConnection::load` | +| `execute`, writes without `RETURNING` | `Connection::execute_returning_count` | +| `sql_query` with several statements | `SimpleConnection::batch_execute` | + +`SyncedConnection` implements all three and delegates to the inner +connection `C`, normally `PgConnection`. Transactions use Diesel's ANSI +manager over the inner connection's state, so there is one transaction +depth. Writes always go to Postgres. Pools (`r2d2`, and `bb8` or +`deadpool` through diesel-async) and `#[derive(Queryable)]` work +unchanged, because the wrapper is just another connection type. + +Diesel seals `Connection` behind the cargo feature +`i-implement-a-third-party-backend-and-opt-into-breaking-changes`; the +crate enables it. + +## A read + +At `load` the wrapper holds the typed Diesel AST. From it, with no +network, it renders the SQL with `$n` placeholders (`to_sql` into a +`PgQueryBuilder`), collects every bind as its Postgres type OID and +binary bytes (`collect_binds`), and reads the route flag if the query +carries one. Then: + +1. **Route.** The query's flag, else a `with_route` scope, else the + connection's route. `Route::Postgres` reads live. +2. **Reads that are never cached** go live: a statement that is not a + `SELECT` (Diesel sends `INSERT ... RETURNING` through `load`), a read + inside an open transaction, a read inside the read-your-writes window, + `FOR UPDATE` and its variants, a `WITH` whose expressions modify data. +3. **Behind.** While the engine says it is further behind the database + than `STREAMGRES_S2S_MAX_LAG_MS`, or has said nothing for that long, + reads go live and `lagging` counts them. A miss whose engine rows + arrive with a lag past it, as on a socket that just reconnected to + a lagging engine, keeps them for later reads and serves this one + live. +4. **Refused recently.** A query the engine refused, with the same + values, within `STREAMGRES_REFUSED_TTL_MS` goes live without asking + again. +5. **Hit.** The entry under the key is served; in rows mode the result + is built from the engine's rows if a change came since it last was. +6. **Miss.** The key is reserved and the query subscribed. Rows mode + builds the entry from the engine's rows; refill mode reads Postgres + and keeps what it returned; a refusal goes live. + +`SyncedConnection::last_served()` reports where the connection's last +read came from: `Served::Cache`, `Served::Engine` or `Served::Postgres`. + +## The cache key + +A key is the rendered SQL plus every bind as Postgres receives it, OID +and bytes. Two reads with equal keys get equal rows. (Diesel's own +statement-cache key is the `TypeId` of the query shape, which ignores +argument values.) + +A refusal is remembered under the same key, values and all: the engine +refuses on values too (a time at infinity, more rows than it follows for +one user), so a refusal for one value says nothing of another. + +## Two ways an entry is kept fresh + +Every subscribe asks the engine for rows. Which mode a query ends up in +is decided by its projected columns: + +| | rows mode | refill mode | +|---|---|---| +| when | every projected column has a type the client encodes exactly | a projected column cannot be reproduced exactly from the engine's values | +| miss | built from the engine's `accepted` rows; Postgres is not read | Postgres is read once; rows kept byte for byte | +| change | the engine's `delta` updates the store in place; the next read builds the result | the engine's `touched` marks the entry stale; the next read refills | +| reconnect | rebuilt from the engine's snapshot | refilled on next read | + +`EngineClient::fallbacks()` lists the queries in refill mode and why; +`EngineClient::refusals()` lists refused queries and why. + +### Rows mode + +The engine's `accepted` carries the result per table of the query, root +first: the table, its primary-key columns, its join edge to its parent, +and every column of every row as JSON. The body also carries the +projected columns and their output names, a `COUNT(*)` flag, the root's +order, the limit and the offset. Each later transaction that changes the +result sends a `delta`: per table, the rows that entered or changed and +the keys of rows that left. + +**Types.** The engine has coarse types (string, integer, float, time); +Diesel decodes Postgres's binary format, where an `int4` is four bytes +and an `int8` eight. So the client looks up each table's exact column +types from `pg_attribute`, once per table, through the connection the +read came in on, and encodes each cell for its type. A cached row then +decodes through the same `FromSql` as a live one. A snapshot that names +a column the known types lack has the table's types looked up again. + +| encoded exactly (rows mode) | not encodable (refill mode) | +|---|---| +| `bool`, `int2`, `int4`, `int8`, `oid`, `float4`, `float8` | `numeric` (the engine holds a float) | +| `text`, `varchar`, `bpchar`, `name`, `char` | `interval` | +| `uuid`, `json`, `jsonb`, `bytea` | ranges, composites | +| `date`, `timestamp`, `timestamptz` at any precision | any other type | +| user enums | arrays of an unencodable element type | +| one-dimensional arrays of the above | | + +The check runs once per query over every projected column, so an +unencodable column puts the whole query in refill mode up front. + +Values on the wire: times are microseconds since the Unix epoch, +Postgres's own precision; `bytea` is Postgres's hex text; an enum is its +label; a float's `NaN` and infinities are the strings `"NaN"`, +`"Infinity"` and `"-Infinity"`. + +**The materialized store** keeps the rows by primary key per table. The +cache entry holds the store itself (`cache::Live`): a delta moves rows +in and out of it on the socket's thread, and the next read builds the +result from it, on the reading thread, and keeps it until the next +change. A burst of deltas costs one build, and none if nobody reads in +between. The build: + +1. The root table's rows in query order. `ASC` puts NULLs last, `DESC` + first, as Postgres does by default. Strings compare bytewise; the + engine only accepts an `ORDER BY` on a column Postgres orders the + same way (see below). Rows that tie, and child rows under one parent, + come out in primary-key order by value (9 before 10). +2. A hash join down each edge, the joined tables depth first. An inner + edge drops a parent with no match; a left edge extends it with NULLs. +3. `OFFSET` rows skipped, then `LIMIT` applied over the joined rows. The + engine keeps `limit + offset` root rows so there is enough to skip. +4. Projection to the selected columns and names, each cell encoded. A + `COUNT(*)` counts every joined row (the engine keeps every row for + it), and its `LIMIT` and `OFFSET` cut the one `int8` row it is, as + Postgres applies them after counting. + +A value that cannot be encoded after all (a delta brings one) falls the +query back to refill mode at the read that meets it. + +### Refill mode + +On a miss the client subscribes first, then reads Postgres through the +inner connection and keeps each row's OID and raw bytes (`OwnedRow`). A +cached row is byte-identical to the one Postgres sent, so any column +type works. On `touched` the entry is marked stale and the next read +refills. The cost is one Postgres read per change per affected query. + +## What the engine follows + +The engine parses the SQL as Diesel renders it: quoted and qualified +identifiers, a column list, `*` or `COUNT(*)`, `$n` placeholders. + +| accepted | refused, with a reason | +|---|---| +| `=`, `<>`/`!=`, `<`, `<=`, `>`, `>=`, `[NOT] IN`, `IS [NOT] NULL`, `AND`, `OR`, parentheses | `DISTINCT`, `GROUP BY`, aggregates other than `COUNT(*)` | +| `= ANY($n)` with an array bind (`eq_any`), read as `IN`; `!= ALL($n)` (`ne_all`), read as `NOT IN`, an empty array being no condition | `LIKE`, `ILIKE`, array operators such as `@>` | +| `ORDER BY` on the root table by a column Postgres orders as the engine does: numbers, booleans, times, `uuid`, `bytea`, `name`, `time`, and strings under the `C` collation; `LIMIT`, `OFFSET` | an `ORDER BY` on an enum, on a string under any other collation, on JSON; table aliases, unknown tables or columns | +| `INNER` and `LEFT [OUTER]` joins on a one-column equality each side; an `INNER JOIN` inside a `LEFT JOIN`'s parentheses | a table joined twice; `RIGHT`, `FULL`, `CROSS`; compound join keys; an `INNER JOIN` onto a table under a `LEFT JOIN`, written after it | +| `ON` conditions and `WHERE` conjuncts that name one table | an `OR` across tables; a `WHERE` on any table under a `LEFT JOIN` | +| joined tables ordered by their primary keys ascending, depth first, after root columns that cover the root's whole key | any other `ORDER BY` on a joined table; an unqualified column in a join | +| binds of integers, floats, numerics, booleans, text-like types, UUIDs, JSON, timestamps, dates, enums (one created after the engine started is learnt on its first bind) | any other bind type; a time or date at infinity; a result over the engine's row limit | +| filters that expand to at most 1,024 alternatives, with at most 10,000 `IN` and `NOT IN` values in all (`STREAMGRES_S2S_MAX_DISJUNCTS`, `STREAMGRES_S2S_MAX_LIST_VALUES` on the engine) | a larger query: an `AND` of `OR`s multiplies its alternatives, and the engine expands them on the thread every subscriber's changes go through | + +In a join, Diesel's `FROM ("a" INNER JOIN ("b" ...) ON (...))` becomes a +tree: each `JOIN` is an edge from the table on its left (the parent) to +the item on its right (the child). Each top-level `WHERE` conjunct and +each extra `ON` condition becomes the filter of the one table it names. +`ORDER BY` and `LIMIT` belong to the root. + +## The protocol + +WebSocket at `{base}/s2s/v1`; JSON text frames, each `[tag, body]`. + +Client to engine: + +| frame | meaning | +|---|---| +| `["subscribe", {"id", "sql", "binds": [{"oid", "bytes": base64 or null}], "rows": true}]` | follow this query; `id` is the client's own, per socket | +| `["unsubscribe", {"id"}]` | the cache dropped the key | +| `["ping"]` | keepalive; the client sends one while the socket is quiet | + +Engine to client: + +| frame | meaning | +|---|---| +| `["hello", {"protocol": 1, "max_subscriptions": n, "database": {"system", "name"}}]` | first frame; `max_subscriptions` is the per-connection limit, `null` for none; `database` is what the engine streams (the cluster's `system_identifier` and the database's name); an engine that cannot read it closes the socket instead of greeting, and an older engine sends none | +| `["accepted", {"id", "nodes", "columns", "names", "count", "order", "limit", "offset"}]` | the engine follows the query; the body is the hydrated result, `nodes` per table with `table`, `key`, `parent`, `parent_column`, `child_column`, `inner`, `rows` | +| `["delta", {"id", "changes": [{"table", "put": [row], "del": [key]}]}]` | one transaction's change to a rows-mode result; one entry per key, its last state, never in both lists, so apply order does not matter | +| `["touched", {"ids"}]` | one transaction changed these refill-mode subscriptions | +| `["refused", {"id", "reason"}]` | the engine cannot follow the query; after an `accepted`, it stopped following it | +| `["pong"]` | reply to `ping` | +| `["lag", {"ms"}]` | right after `hello`, then every 500 ms: how far at most the engine's applied changes trail the database, the age of the newest heartbeat it applied; `null` before the first heartbeat is sent | +| `["error", {"reason"}]` | a frame that could not be read and names no id; a `subscribe` whose id reads is answered `refused` instead | + +An id is never reused on a socket: frames queued for an unsubscribed id +may still arrive after the `unsubscribe`. + +On the engine, one task owns every S2S subscription. It consumes the +same engine events as the Zero client groups, with its own index of +subscriptions, and maps the engine's ids back to the socket and the +client's id. A row a transaction changed is serialized once and spliced +into every `delta` that carries it. Each socket has a request queue of +1,024 frames and an outgoing queue of 4,096 frames and 64 MiB; a client +that stops reading is closed when its outgoing queue fills rather than +buffered for, and its socket is cut off 5 seconds later if its writer +still cannot flush. + +## Correctness rules + +**Subscribe before read.** The engine accepts first and the client reads +second. A change after the read is reported; a change between the two +is reported redundantly, which costs one refill and never a stale row. + +**Generations.** `Cache::reserve` takes the key's generation before a +fill; `Cache::fill` stores the rows only if nothing bumped it since. A +touch, a delta and every path that ends a subscription bump it. The map +holds generations only for keys with a reservation in flight. + +The engine side follows the same rule. A snapshot for a resubscribe +reserves the key before building its store; `Cache::live` bumps the +generation, so a slower Postgres fill with older rows loses, and checks +it again after the insert. A delta only changes the store the entry +already holds, and stores it again only when its weight drifted. Every path that ends a subscription (dropped read, failed read, +oversized result, eviction, expiry) bumps the generation as it drops the +key, under the worker's lock, so rows built for a key unsubscribed +meanwhile are dropped. The wrapper never unsubscribes directly; +`Cache::abandon` does it in that order. No lock is held across a moka +insert, because the eviction listener runs on the inserting thread; a +fill checks the generation before and after instead. + +**Stale, not gone.** A touch, a lost fill and a dropped socket leave the +key marked stale rather than removed: the mark weighs only the key and +keeps the entry's clocks, so the subscription of a query nobody reads +again still ends when its idle time runs out. + +**The `Building` store.** When `accepted` arrives, the reader thread +puts an empty `Building` store in place under the lock that ends the +subscriber's wait, so a delta right behind it queues there. +`materialize` drains the queue and turns that store `Ready` in place. + +**Late answers.** A subscribe the client gave up on is withdrawn. An +`accepted` for an id nobody holds is answered with `unsubscribe`. + +## The engine client + +`EngineClient` holds a pool of sockets and a cap on subscriptions, from +the environment (`Config::from_env`) or explicitly +(`EngineClient::with_config`): + +| variable | meaning | default | +|-------------------------------------|-------------------------------------------|---------| +| `STREAMGRES_S2S_CONNECTIONS` | sockets to the engine, each on its thread | 2 | +| `STREAMGRES_S2S_MAX_SUBSCRIPTIONS` | subscriptions held across all sockets | 10000 | +| `STREAMGRES_S2S_VERDICT_TIMEOUT_MS` | wait for `accepted` or `refused` | 5000 | +| `STREAMGRES_S2S_SILENCE_MS` | a socket silent this long is taken for dead | 10000 | +| `STREAMGRES_S2S_MAX_LAG_MS` | the engine's lag past which reads go live; 0 trusts the cache | 2000 | +| `STREAMGRES_S2S_TOKEN` | bearer token sent when a socket opens | none | + +The token goes out as `Authorization: Bearer ` on each socket's +upgrade, for an engine started with the same `STREAMGRES_S2S_TOKEN`. An +engine that wants a token and is not given the right one answers 401. +The worker then retries like any failed connect, so every read goes to +Postgres until the two agree. The token never appears in `Config`'s +`Debug` output. Over `ws://` the token crosses the network in plain +text; `wss://` connects with TLS (rustls, the system's root +certificates), and `EngineClient::warnings` names a token sent over +`ws://` to another machine, for the application to log. + +**Connect failures.** A worker keeps why its latest connect failed: a +401 (the token), a 503 (the engine not serving yet, or out of +connections), another status, a network error, or a URL or token that +cannot go in a request. While the socket is down, each refusal it gives +says so (`the engine is not connected: ...`), and +`EngineClient::connect_errors` lists them. A successful connect clears +it. + +**The engine's database.** The engine's `hello` names the database it +streams. The first miss through a cache reads the same of the database +its reads go to (`cache::DATABASE_SQL`, readable by any role), once per +cache, before it subscribes. A probe that fails is not recorded: that +read goes live without asking the engine, and the next miss probes +again. A socket whose engine +streams another database is not asked for rows, which would otherwise +be served as this database's: every subscribe through it is refused, +naming both, and the reads go to Postgres. A key resubscribed after a +reconnect is checked the same way when its `accepted` arrives. An engine +whose `hello` names no database, as engines from before the check do, +is taken on trust; this engine closes a socket it cannot name its +database on rather than greet it. + +**The engine's lag.** A large transaction, a slow apply or a stalled +replication stream delay the cache without dropping the socket. The +engine writes a heartbeat into the WAL every 500 ms +(`pg_logical_emit_message`, not transactional, under the prefix +`/heartbeat`; `STREAMGRES_HEARTBEAT_MS`, 0 off) and notes when it +sent each. When its replication stream brings a heartbeat back and every +change before it is applied and queued to the sockets, that heartbeat is +applied, and so is every change committed before it was sent. Its lag is the age of the newest +applied heartbeat, a bound on how stale it can be: under one interval +while it keeps up, growing while the stream is stuck. Each socket is +told the lag right after `hello`, so before any answer, and every +500 ms. The wrapper adds the time since the last +`lag` frame, so an engine that stops sending them is taken as falling +behind too, and goes live while the largest lag over the connected +sockets is past `STREAMGRES_S2S_MAX_LAG_MS`. An engine that reports no +lag (`null`, or one that never sends the frame) is trusted, as before. + +**The engine's limit.** The engine's `hello` names how many +subscriptions it takes on one connection (`max_subscriptions`). The pool +passes over a socket at that limit. With every socket full, a subscribe +is refused here, not by the engine, so the read goes live at once +instead of after a round trip. A subscribe sent before the `hello` +lands may still be refused by the engine; both refusals are transient. + +Each socket is a worker with its own thread, runtime, ids, waiting +subscribes and rows-mode stores. The pool maps each key to its worker. A +miss goes to the worker already holding the key, else to the connected +worker with the fewest subscriptions. Past the cap the engine is not +asked and the read goes live; a key freed by eviction or refusal gives +its slot back. + +A subscribe blocks the calling thread on a channel until the answer +arrives or the verdict timeout passes, so the wrapper needs no async +context. + +| timer | value | on expiry | +|---|---|---| +| verdict timeout | 5 s, configurable | the read goes live; the id is withdrawn | +| send timeout | 10 s | the socket is given up on, so an engine that stops reading cannot block the worker | +| ping | a third of the silence window | a quiet socket is pinged | +| silence | 10 s, configurable | nothing heard: the socket is taken for dead (an engine killed without a close, a network gone without a reset), and its keys go stale; checked every thirtieth of the window, a second at most, so the cache serves a frozen engine's rows for the window and no longer | +| reconnect delay | 1 s | a lost socket is retried every second | + +**Refusals.** A refusal from the engine is remembered for the query and +its values for `STREAMGRES_REFUSED_TTL_MS`. Some refusals are not remembered: +- those from the client's own side (not connected, over the cap, no + answer in time); +- the engine's per-connection limit (`the subscription limit of ...`), + which frees up as other keys go. All +remembered refusals are forgotten when a socket drops, since the engine +that comes back may know types the old one did not; so are the column +types learnt, since a migration may have changed one. + +**Reconnect.** A dropped socket refuses the subscribes waiting on it, +marks the cache entries of its keys stale, which go live until it is +back, and +on reconnect resubscribes every key it holds under the same ids. +Rows-mode keys rebuild from the snapshot with no Postgres read; +refill-mode keys refill on their next read; keys now refused are +forgotten. A key is pending until the engine accepts it again, and reads +of it go live meanwhile. One socket dropping affects only its keys. With +the engine unreachable, every subscribe is refused at once, so an engine +outage costs the application its cache and nothing else. + +## Cache limits + +`cache::Limits`, from the environment (`Limits::from_env`) or in code +(`Cache::with_limits`): + +| variable | meaning | default | +|------------------------------------|----------------------------------------------|---------| +| `STREAMGRES_CACHE_MAX_BYTES` | rows and keys held | 256 MiB | +| `STREAMGRES_CACHE_MAX_ENTRY_BYTES` | the largest result kept; 0 no limit | 16 MiB | +| `STREAMGRES_CACHE_IDLE_MS` | drop an entry unread for this long; 0 never | 30 min | +| `STREAMGRES_CACHE_TTL_MS` | drop an entry this old however read; 0 never | off | +| `STREAMGRES_REFUSED_TTL_MS` | skip the engine for a refused query this long; 0 always ask | 30 s | + +A rows-mode entry weighs its store twice, for the rows and for the +result built from them, and is weighed again when the store drifts by a +quarter (64 KiB at least). A refill-mode entry weighs its rows; a stale +mark its key. A result over the entry limit is served but not kept, and +its subscription ends. + +When the cache is full, what stays is decided by how often a key is +asked for, not by size or age (moka's frequency-based admission): a key +read again and again displaces an entry read once, and a run of keys +each read once cannot displace it. Reads are counted, misses included, +once the cache is past half its budget. + +Idle time makes the cache follow what the backend reads now; the engine +changing a query's rows is not a read. Time to live drops an entry on a +timer regardless of reads; it is off by default because the engine +keeps entries fresh. The byte budget and both clocks end the +subscription and bump the generation. + +`cache::Stats` counts: `hits`, `misses`, `refused`, `refused_recently`, +`touched`, `evicted`, `expired`, `stale_fills`, `oversized`, `replaced`, +`after_write`, `locking`, `lagging`. + +## Routing + +A read goes through the cache (`Route::Streamgres`) or live +(`Route::Postgres`). The most specific level wins: + +1. **The query.** `query.via(Route::Postgres)`, or `via_postgres()` / + `via_streamgres()` from `RouteDsl`, renders a leading comment such as + `/*streamgres:postgres*/`. The wrapper reads and strips it, so the + key and the SQL the engine sees are the plain query and flagged and + unflagged reads share an entry. `conn.with_route(route, |c| ...)` + flags every read in the closure. +2. **The pool.** Each connection's route and cache come from the URL it + was established with, through the process-wide `Settings` + (`install`, once per process). `Settings::single` shares one cache + and routes by URL; `Settings::by_url` gives each database its own + cache, since one engine streams one database. A URL with no cache + is a passthrough. `set_route` changes one connection. +3. **The default:** Streamgres. + +`Consistency::Eventual` and `Consistency::Strong` are aliases of the two +routes. + +**Read your writes.** Between a commit and the engine's delta, a +Postgres read sees the new row and a cache read the old one. So after a +write through a cache (`execute`, raw SQL with any statement but +`SELECT`, `BEGIN`, `ROLLBACK`, `SAVEPOINT` or `SET`, `INSERT ... +RETURNING`, a `COMMIT`), every read through that cache goes live for +`STREAMGRES_READ_YOUR_WRITES_MS` (default 1000, 0 off; +`set_read_your_writes` per connection), counted from when the write +ends, however long it took. The clock lives on the cache, so a write on +one pooled connection covers reads on the others. A Postgres read never +writes into the cache. + +## diesel-async + +`AsyncSyncedConnection` (feature `async`) wraps `AsyncPgConnection` and +implements `AsyncConnection` and `PoolableConnection`, so bb8 and +deadpool pools work unchanged. Routing, the cache, read-your-writes, the +settings and the engine client are shared with the blocking wrapper. + +`AsyncConnection::load` must return a future that borrows neither the +connection nor the query. So inside `load`, synchronously, the key is +rendered, the route decided, the cache checked, and the futures the read +might need are built: the inner connection's load (polled only if the +read goes to Postgres) and a lookup of every table's column types +(`types::ALL_TABLES_SQL`, polled only when the engine ships rows for a +table whose types are unknown). The subscribe itself runs on a blocking +thread. This is why `Subscriptions::subscribe` has two steps: it answers +`NeedsTypes` for unknown tables, and `materialize` finishes once the +caller has looked them up. + +An enum bind needs its OID for the key. The async wrapper cannot look it +up inside `load`, so the first read with such a bind goes live while the +OID is resolved and stored in the cache for every connection sharing it; +later reads are cached normally. An array of an enum (`eq_any`) cannot +be rendered until the OID is known and takes the same path. + +diesel-async allows several pending futures per connection and panics +if its transaction state is queried while one is pending. The wrapper +counts its own futures in flight and, while any is, sends reads live +without asking. A future dropped after it subscribed ends its +subscription, and a dropped `with_route` scope restores the outer route. + +## Engine behaviour the wrapper relies on + +- **Catalog.** The engine loads its catalog at start and follows schema + changes through the DDL event trigger. An added column is absorbed; + new tables keep the known enum types; a new enum type is learnt the + first time a subscribe binds it. A dropped + column or a `TRUNCATE` of a published table stops the engine; it + restarts and serves the schema as it is, and the client reconnects. +- **Times.** The engine keeps every time in UTC microseconds, and every + storage and replication session runs with `TimeZone=UTC`, so a time in + a filter means the same instant to Postgres and to the engine. + +## Tests + +| file | covers | +|---|---| +| `tests/passthrough.rs` | every Diesel read and write behaves the same through the wrapper | +| `tests/owned_row.rs` | a captured row decodes like the live row | +| `tests/cached_reads.rs` | cache rules against a fake engine side: generations, limits, clocks, refusals, routing, read-your-writes, admission | +| `tests/async_cached.rs` | the same through diesel-async | +| `tests/settings.rs` | cache and route chosen by URL | +| `tests/join_sql_shape.rs` | the exact join SQL Diesel renders, which the engine parser expects | +| `tests/silent_engine.rs` | an engine that never answers, and one that drops every socket | +| `tests/engine_live.rs`, `tests/async_live.rs` | the whole path against a running engine | +| `tests/shape_coverage.rs` | production-shaped lookups compared with Postgres before and after a change | +| `tests/engine_latency.rs` | commit-to-cache latency | +| `tests/s2s_live.rs` (engine crate) | the protocol frame by frame | +| `crates/streamgres-demo/tests/scenario.rs` | every lookup of the demo app through both wrappers under writes, an engine restart and migrations | + +`crates/streamgres-demo` is a ticket desk built on the wrapper: a +list of query shapes, each declaring rows, refill or refused; a writer +that makes every kind of change; a verify loop comparing both routes; +latency and migration commands. `demo/` runs it with Postgres and the +engine under docker compose. + +## Running it locally + +- Postgres 15 or newer with `wal_level = logical`. A Homebrew + `postgresql@16` on its own port works: `initdb`, set + `wal_level = logical` and a short `unix_socket_directories`, `pg_ctl + start`. +- In the database: the demo's publication and schema-change trigger + stack, both for app `demo`: `psql -f demo/postgres-init.sql` and + `psql -f demo/ddl-triggers.sql`. +- The engine: `STREAMGRES_APP_ID=demo STREAMGRES_PUBLICATION=demo_pub STREAMGRES_PG_DSN=... STREAMGRES_QUERY_URL=http://127.0.0.1:1/q STREAMGRES_MUTATE_URL=http://127.0.0.1:1/m cargo run --release --bin server`. + The query and mutate URLs are required by the config and unused on + the S2S path. +- Or everything at once: `docker compose -f demo/docker-compose.yml up --build`. +- tokio-console: build the server with + `RUSTFLAGS="--cfg tokio_unstable" cargo build --release --bin server --features console`, + run it, and `tokio-console` attaches on `127.0.0.1:6669`; the S2S + task is named `s2s`. diff --git a/crates/diesel-streamgres/src/asynchronous.rs b/crates/diesel-streamgres/src/asynchronous.rs new file mode 100644 index 0000000..18972ed --- /dev/null +++ b/crates/diesel-streamgres/src/asynchronous.rs @@ -0,0 +1,757 @@ +//! The same wrapper over diesel-async's `AsyncPgConnection`. +//! +//! `AsyncSyncedConnection` implements `AsyncConnection`, so an +//! application on diesel-async swaps its connection type and nothing +//! else changes: the DSL, `sql_query`, transactions and the bb8 or +//! deadpool pools go through unchanged. The routing, the cache, the +//! read-your-writes window and the engine side are the ones the blocking +//! wrapper uses; only how a read reaches them differs. +//! +//! # Why the futures are built up front +//! +//! `AsyncConnection::load` returns a future that may not borrow the +//! connection (the trait's lifetimes do not relate `'conn` to `'query`) +//! and may not carry the query (it need not be `Send`). diesel-async +//! handles this by rendering the query to SQL and bind bytes before the +//! future exists; its own `load` future is `'query` and owns nothing of +//! the connection. So this wrapper does the same: everything that needs +//! the query or the connection happens synchronously in `load`, and the +//! future it returns owns the pieces it may need: the inner connection's +//! own load future (polled only when the read goes to Postgres), a +//! lookup of every table's column types (polled only when the engine +//! ships rows for a table not yet known), and the cache. A subscribe +//! blocks on the engine, so it runs on a blocking thread. +//! +//! One consequence: diesel-async records a query as started when its +//! future is built, and as finished when the future completes, so a +//! miss the engine then serves from its own rows leaves one started +//! query in the instrumentation with no finish. A hit builds nothing. +//! +//! # Pipelining +//! +//! diesel-async lets several futures from one connection run at once. +//! Its transaction state is shared with every such future, and asking +//! for it while one is pending is a panic in diesel-async itself. The +//! wrapper counts its own futures in flight and, while any is, sends a +//! read to Postgres without asking: a read pipelined behind a write +//! must see that write, which only Postgres can promise. +//! +//! A bind of a user-defined type needs its OID, which the blocking +//! wrapper resolves on the spot through the connection. Here the first +//! read with such a bind goes to Postgres while the OID is resolved and +//! remembered in the cache, for every connection sharing it; from the +//! next read on the query is cached like any other. + +use std::borrow::Cow; +use std::collections::HashMap; +use std::sync::atomic::{AtomicBool, AtomicU8, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use diesel::QueryableByName; +use diesel::connection::Instrumentation; +use diesel::deserialize::FromSqlRow; +use diesel::pg::{ + FailedToLookupTypeError, Pg, PgMetadataCache, PgMetadataCacheKey, PgMetadataLookup, + PgTypeMetadata, +}; +use diesel::query_builder::{AsQuery, QueryFragment, QueryId}; +use diesel::result::Error; +use diesel::sql_query; +use diesel::sql_types::{Oid, Text, Untyped}; +use diesel::{ConnectionResult, QueryResult}; +use diesel_async::pooled_connection::PoolableConnection; +use diesel_async::{ + AnsiTransactionManager, AsyncConnection, AsyncPgConnection, CacheSize, SimpleAsyncConnection, + TransactionManager, +}; +use futures_util::future::{BoxFuture, FutureExt, TryFutureExt}; +use futures_util::stream::{self, BoxStream, StreamExt, TryStreamExt}; +use tokio::task::spawn_blocking; + +use crate::cache::{Cache, DATABASE_SQL, Database}; +use crate::cursor::SyncedRow; +use crate::key::QueryKey; +use crate::route::{Route, Served}; +use crate::row::{OwnedCursor, OwnedRow, capture_all}; +use crate::subscriptions::STALE_TYPES; +use crate::materialized::Snapshot; +use crate::subscriptions::{Outcome, Subscriptions, Verdict}; +use crate::types::{ALL_TABLES_SQL, TableAttribute, TableTypes, all_table_types}; +use crate::{Attachment, SETTINGS, may_write, read_your_writes_from_env, rendered}; + +/// diesel-async's Postgres row, which it does not export by name. +pub type PgRow = ::Row<'static, 'static>; + +/// The row an `AsyncSyncedConnection` yields: diesel-async's own, or one +/// the cache holds. +pub type AsyncSyncedRow = SyncedRow; + +/// A diesel-async Postgres connection whose reads can be served from a +/// synced local copy. See [`crate::SyncedConnection`] for what is +/// cached and when; the rules are the same. +pub struct AsyncSyncedConnection { + inner: AsyncPgConnection, + cache: Option>, + route: Route, + scoped: Option, + read_your_writes: Duration, + /// OIDs of user-defined types seen in binds, for the cache key: + /// the cache's, shared by every connection through it, or this + /// connection's own when it is a passthrough. + metadata: Arc>, + /// Whether this connection has loaded every table's types for the + /// engine side; false again when the engine named a table that was + /// not among them. + catalog_loaded: Arc, + /// Cloned into every future the wrapper hands out: while the count + /// is above one, a future is pending on this connection. + inflight: Arc<()>, + /// Where the last read was served from, written by the read's + /// future once it knows: 0 before any read, else [`Served`] + 1. + served: Arc, +} + +impl AsyncSyncedConnection { + /// Wraps an established connection as a passthrough. + pub fn new(inner: AsyncPgConnection) -> Self { + AsyncSyncedConnection { + inner, + cache: None, + route: Route::Postgres, + scoped: None, + read_your_writes: Duration::ZERO, + metadata: Arc::new(Mutex::new(PgMetadataCache::new())), + catalog_loaded: Arc::new(AtomicBool::new(false)), + inflight: Arc::new(()), + served: Arc::new(AtomicU8::new(0)), + } + } + + /// Where this connection's last read was served from: the cache, + /// the engine's rows, or Postgres. `None` before any read. A read + /// still in flight has not written it yet. + pub fn last_served(&self) -> Option { + match self.served.load(Ordering::Acquire) { + 1 => Some(Served::Cache), + 2 => Some(Served::Engine), + 3 => Some(Served::Postgres), + _ => None, + } + } + + /// Wraps a connection with a cache and the route its reads take + /// unless a query or a scope says otherwise. + pub fn with_cache(inner: AsyncPgConnection, cache: Arc, route: Route) -> Self { + AsyncSyncedConnection { + metadata: Arc::clone(cache.type_metadata()), + cache: Some(cache), + route, + read_your_writes: read_your_writes_from_env(), + ..Self::new(inner) + } + } + + /// Changes the route this connection's reads take by default. + pub fn set_route(&mut self, route: Route) { + self.route = route; + } + + /// The route this connection's reads take by default. + pub fn route(&self) -> Route { + self.route + } + + /// How long after a write reads through this connection's cache go + /// to Postgres; see [`crate::SyncedConnection::set_read_your_writes`]. + pub fn set_read_your_writes(&mut self, window: Duration) { + self.read_your_writes = window; + } + + /// How long after a write reads go to Postgres. + pub fn read_your_writes(&self) -> Duration { + self.read_your_writes + } + + /// The cache this connection reads through, if any. + pub fn cache(&self) -> Option<&Arc> { + self.cache.as_ref() + } + + /// The connection that talks to Postgres. + pub fn inner(&mut self) -> &mut AsyncPgConnection { + &mut self.inner + } + + /// Gives the Postgres connection back. + pub fn into_inner(self) -> AsyncPgConnection { + self.inner + } + + /// Runs `f` with every read inside it routed `route`; a flag on the + /// query itself still wins. The closure gets the connection back. + /// The outer route is restored however the future ends: completed, + /// dropped before completion, or unwound. + pub async fn with_route(&mut self, route: Route, f: F) -> R + where + F: for<'a> FnOnce(&'a mut Self) -> BoxFuture<'a, R>, + { + /// Puts the outer route back when dropped. + struct Restore<'a> { + conn: &'a mut AsyncSyncedConnection, + outer: Option, + } + impl Drop for Restore<'_> { + /// Restores the outer route. + fn drop(&mut self) { + self.conn.scoped = self.outer; + } + } + let outer = self.scoped.replace(route); + let restore = Restore { conn: self, outer }; + f(&mut *restore.conn).await + } + + /// Whether a future this wrapper handed out is still pending. + fn pipelined(&self) -> bool { + Arc::strong_count(&self.inflight) > 1 + } + + /// Marks the write now, and again when `future` completes: a read + /// issued after the write's future resolves is what the window is + /// for, and the clock should start as late as the write did. + fn written<'a, T: Send + 'a>( + &mut self, + future: BoxFuture<'a, QueryResult>, + ) -> BoxFuture<'a, QueryResult> { + let cache = self.cache.clone(); + if let Some(cache) = &cache { + cache.mark_write(); + } + let token = Arc::clone(&self.inflight); + async move { + let _pending = token; + let result = future.await; + if let Some(cache) = cache { + cache.mark_write(); + } + result + } + .boxed() + } + + /// A read future with the in-flight token attached. + fn tracked<'a, T: Send + 'a>( + &self, + future: BoxFuture<'a, QueryResult>, + ) -> BoxFuture<'a, QueryResult> { + let token = Arc::clone(&self.inflight); + async move { + let _pending = token; + future.await + } + .boxed() + } + + /// Whether a write went through the cache within the read-your-writes + /// window. + fn within_write_window(&self) -> bool { + self.cache + .as_ref() + .is_some_and(|cache| cache.wrote_within(self.read_your_writes)) + } + + /// Whether a transaction is open; an unreadable state counts as one. + fn in_transaction(&mut self) -> bool { + AnsiTransactionManager::transaction_manager_status_mut(self) + .transaction_depth() + .map(|depth| depth.is_some()) + .unwrap_or(true) + } + + /// The key of a query, and the user-defined types its binds name + /// whose OIDs are not known yet. A query that cannot be rendered + /// until those are known (an array of an enum, whose element OID + /// goes into the bind) has no key yet: `Ok((None, unknown))`. + fn key_of>( + &mut self, + query: &T, + ) -> QueryResult<(Option, Vec)> { + let mut lookup = Lookup { + cache: Arc::clone(&self.metadata), + unknown: Vec::new(), + }; + let key = QueryKey::of(query, &mut lookup); + match key { + Ok(key) => Ok((Some(key), lookup.unknown)), + Err(error) if !lookup.unknown.is_empty() => { + let _ = error; + Ok((None, lookup.unknown)) + } + Err(error) => Err(error), + } + } + + /// A Postgres read's rows as the wrapper's. + fn live_stream<'a>( + stream: BoxStream<'a, QueryResult>, + ) -> BoxStream<'a, QueryResult> { + stream.map_ok(SyncedRow::Live).boxed() + } + + /// Cached rows as the wrapper's. + fn cached_stream(rows: Arc<[OwnedRow]>) -> BoxStream<'static, QueryResult> { + stream::iter(OwnedCursor::new(rows).map(|row| row.map(SyncedRow::Cached))).boxed() + } + + /// The live read, with every unknown type resolved first and + /// remembered for the next read. + fn resolving<'query>( + &mut self, + unknown: Vec, + live: LoadFuture<'query>, + ) -> LoadFuture<'query> { + let resolvers: Vec> = unknown + .iter() + .map(|name| self.inner.load(type_oids_query(name))) + .collect(); + let metadata = Arc::clone(&self.metadata); + async move { + for (name, resolver) in unknown.into_iter().zip(resolvers) { + let rows: Vec = resolver.await?.try_collect().await?; + if let Some(row) = rows.first() { + let oids = >::build_from_row(row) + .map_err(Error::DeserializationError)?; + metadata + .lock() + .expect("metadata") + .store_type(name.key(), (oids.oid, oids.typarray)); + } + } + live.await + } + .boxed() + } +} + +/// A user-defined type a bind named, as Diesel names it. +#[derive(Debug, Clone, PartialEq, Eq)] +struct TypeName { + schema: Option, + name: String, +} + +impl TypeName { + /// The name as Diesel's metadata cache keys it. + fn key(&self) -> PgMetadataCacheKey<'static> { + PgMetadataCacheKey::new( + self.schema.clone().map(Cow::Owned), + Cow::Owned(self.name.clone()), + ) + } +} + +/// A type lookup over the OIDs already known, remembering what it did +/// not know rather than asking the database (it cannot, synchronously). +/// The cache is shared by every connection through one `Cache`, so it +/// is locked per lookup, never taken out: a connection that took it out +/// would leave the others an empty one and lose their writes. +struct Lookup { + cache: Arc>, + unknown: Vec, +} + +impl PgMetadataLookup for Lookup { + /// The type's OIDs when known; otherwise noted as unknown, and a + /// failed lookup for Diesel. + fn lookup_type(&mut self, type_name: &str, schema: Option<&str>) -> PgTypeMetadata { + let key = PgMetadataCacheKey::new(schema.map(Cow::Borrowed), Cow::Borrowed(type_name)); + if let Some(metadata) = self.cache.lock().expect("metadata").lookup_type(&key) { + return metadata; + } + let name = TypeName { + schema: schema.map(str::to_string), + name: type_name.to_string(), + }; + if !self.unknown.contains(&name) { + self.unknown.push(name); + } + PgTypeMetadata::from_result(Err(FailedToLookupTypeError::new(key.into_owned()))) + } + + /// The lookup itself. + fn as_any<'a>(&mut self) -> &mut (dyn std::any::Any + 'a) + where + Self: 'a, + { + self + } +} + +/// A type's OID and its array type's. +#[derive(QueryableByName)] +struct TypeOids { + #[diesel(sql_type = Oid)] + oid: u32, + #[diesel(sql_type = Oid)] + typarray: u32, +} + +/// Diesel's own resolution of a type name, as a query. +fn type_oids_query( + name: &TypeName, +) -> diesel::query_builder::BoxedSqlQuery<'static, Pg, diesel::query_builder::SqlQuery> { + match &name.schema { + Some(schema) => sql_query( + "SELECT t.oid, t.typarray FROM pg_type t JOIN pg_namespace n ON n.oid = t.typnamespace \ + WHERE t.typname = $1 AND n.nspname = $2", + ) + .into_boxed() + .bind::(name.name.clone()) + .bind::(schema.clone()), + None => { + sql_query("SELECT oid, typarray FROM pg_type WHERE oid = quote_ident($1)::regtype::oid") + .into_boxed() + .bind::(name.name.clone()) + } + } +} + +/// A Postgres read in flight. +type LoadFuture<'query> = BoxFuture<'query, QueryResult>>>; + +impl SimpleAsyncConnection for AsyncSyncedConnection { + /// Runs SQL text, marking a possible write now and again when it + /// completes. diesel-async's future borrows the connection, so it + /// cannot be pipelined and needs no in-flight token. + fn batch_execute( + &mut self, + query: &str, + ) -> impl std::future::Future> + Send { + let cache = if may_write(query) { + self.cache.clone() + } else { + None + }; + if let Some(cache) = &cache { + cache.mark_write(); + } + let future = self.inner.batch_execute(query); + async move { + let result = future.await; + if let Some(cache) = cache { + cache.mark_write(); + } + result + } + } +} + +impl AsyncConnection for AsyncSyncedConnection { + type ExecuteFuture<'conn, 'query> = BoxFuture<'query, QueryResult>; + type LoadFuture<'conn, 'query> = BoxFuture<'query, QueryResult>>; + type Stream<'conn, 'query> = BoxStream<'static, QueryResult>; + type Row<'conn, 'query> = AsyncSyncedRow; + type Backend = Pg; + type TransactionManager = AnsiTransactionManager; + + /// Connects, attached to the cache [`crate::install`] set up for + /// this URL, if any. + async fn establish(database_url: &str) -> ConnectionResult { + let inner = AsyncPgConnection::establish(database_url).await?; + Ok( + match SETTINGS + .get() + .and_then(|settings| settings.attachment_for(database_url)) + { + Some(Attachment { cache, route }) => Self::with_cache(inner, cache, route), + None => Self::new(inner), + }, + ) + } + + /// A read, from the cache when it can be. A query whose binds name + /// types not known yet is read live this time, the types resolved + /// for the next (one that cannot even be keyed until then too). + /// + /// The first miss through a cache reads which database the cache's + /// reads go to before it subscribes, for the engine side to check its + /// engine streams the same one. + /// + /// On a miss, the futures the read may need are built at once; which + /// are polled depends on the engine's answer. The read's future owns + /// the subscription from the moment the subscribe may have made it + /// until the fill; dropped before that, it ends it ([`Unfilled`]). + /// The subscribe runs on a blocking thread and cannot be cancelled, + /// so it checks the guard's flag once it returns: a subscription + /// made for a read that has gone is ended there, by the same path. + fn load<'conn, 'query, T>(&'conn mut self, source: T) -> Self::LoadFuture<'conn, 'query> + where + T: AsQuery + 'query, + T::Query: QueryFragment + QueryId + 'query, + { + let query = source.as_query(); + self.served.store(3, Ordering::Release); + let Some(cache) = self.cache.clone() else { + let live = self.inner.load(query).map_ok(Self::live_stream).boxed(); + return self.tracked(live); + }; + let (key, unknown) = match self.key_of(&query) { + Ok(found) => found, + Err(error) => return std::future::ready(Err(error)).boxed(), + }; + let Some(mut key) = key else { + let live = self.inner.load(query); + let live = self + .resolving(unknown, live) + .map_ok(Self::live_stream) + .boxed(); + return self.tracked(live); + }; + let route = key.take_route().or(self.scoped).unwrap_or(self.route); + if !key.is_select() { + let live = self.inner.load(query).map_ok(Self::live_stream).boxed(); + return self.written(live); + } + if route == Route::Postgres || self.pipelined() || self.in_transaction() { + let live = self.inner.load(query).map_ok(Self::live_stream).boxed(); + return self.tracked(live); + } + if key.is_locking() { + cache.stats.locking.fetch_add(1, Ordering::Relaxed); + let live = self.inner.load(query).map_ok(Self::live_stream).boxed(); + return self.tracked(live); + } + if self.within_write_window() { + cache.stats.after_write.fetch_add(1, Ordering::Relaxed); + let live = self.inner.load(query).map_ok(Self::live_stream).boxed(); + return self.tracked(live); + } + if cache.subscriptions().behind() { + cache.stats.lagging.fetch_add(1, Ordering::Relaxed); + let live = self.inner.load(query).map_ok(Self::live_stream).boxed(); + return self.tracked(live); + } + if !unknown.is_empty() { + let live = self.inner.load(query); + let live = self + .resolving(unknown, live) + .map_ok(Self::live_stream) + .boxed(); + return self.tracked(live); + } + if let Some(rows) = cache.get(&key) { + self.served.store(1, Ordering::Release); + return std::future::ready(Ok(Self::cached_stream(rows))).boxed(); + } + if cache.refused_recently(&key).is_some() { + let live = self.inner.load(query).map_ok(Self::live_stream).boxed(); + return self.tracked(live); + } + let probe: Option> = (!cache.knows_database()) + .then(|| self.inner.load(sql_query(DATABASE_SQL))); + let live = self.inner.load(query); + let catalog: Option> = if self.catalog_loaded.load(Ordering::Acquire) { + None + } else { + Some(self.inner.load(sql_query(ALL_TABLES_SQL))) + }; + let catalog_loaded = Arc::clone(&self.catalog_loaded); + let subscriptions = cache.shared_subscriptions(); + let token = Arc::clone(&self.inflight); + let served = Arc::clone(&self.served); + async move { + let _pending = token; + if let Some(probe) = probe { + match read_database(probe).await { + Some(database) => cache.learn_database(Some(database)), + None => return live.await.map(Self::live_stream), + } + } + let reservation = cache.reserve(&key); + let guard = Unfilled { + cache: Arc::clone(&cache), + key: key.clone(), + gone: Arc::new(AtomicBool::new(false)), + armed: true, + }; + let outcome = { + let subscriptions = Arc::clone(&subscriptions); + let cache = Arc::clone(&cache); + let key = key.clone(); + let gone = Arc::clone(&guard.gone); + spawn_blocking(move || { + let outcome = subscriptions.subscribe(&key); + if gone.load(Ordering::Acquire) { + cache.abandon(&key); + } + outcome + }) + .await + .map_err(|error| Error::QueryBuilderError(Box::new(error)))? + }; + let verdict = match outcome { + Outcome::Verdict(verdict) => verdict, + Outcome::NeedsTypes { snapshot, missing } => { + with_types(&*subscriptions, &key, snapshot, &missing, catalog, &catalog_loaded) + .await? + } + }; + match verdict { + Verdict::Refused(reason) => { + guard.disarm(); + cache.refused(&key, &reason); + live.await.map(Self::live_stream) + } + Verdict::Rows { rows, live: store } => { + guard.disarm(); + cache.fill_live(reservation, store, &rows); + if subscriptions.behind() { + cache.stats.lagging.fetch_add(1, Ordering::Relaxed); + return live.await.map(Self::live_stream); + } + served.store(2, Ordering::Release); + Ok(Self::cached_stream(rows)) + } + Verdict::Accepted => { + let pg_rows: Vec = live.await?.try_collect().await?; + let rows: Arc<[OwnedRow]> = + Arc::from(capture_all(pg_rows.into_iter().map(Ok))?); + guard.disarm(); + cache.fill(reservation, Arc::clone(&rows)); + Ok(Self::cached_stream(rows)) + } + } + } + .boxed() + } + + /// Runs a statement, marking a write when it may be one. + fn execute_returning_count<'conn, 'query, T>( + &'conn mut self, + source: T, + ) -> Self::ExecuteFuture<'conn, 'query> + where + T: QueryFragment + QueryId + 'query, + { + let writes = may_write(&rendered(&source)); + let future = self.inner.execute_returning_count(source); + if writes { + self.written(future) + } else { + self.tracked(future) + } + } + + /// The inner connection's transaction state. + fn transaction_state(&mut self) -> &mut AnsiTransactionManager { + self.inner.transaction_state() + } + + /// The inner connection's instrumentation. + fn instrumentation(&mut self) -> &mut dyn Instrumentation { + self.inner.instrumentation() + } + + /// Sets the inner connection's instrumentation. + fn set_instrumentation(&mut self, instrumentation: impl Instrumentation) { + self.inner.set_instrumentation(instrumentation) + } + + /// Sizes the inner connection's statement cache. + fn set_prepared_statement_cache_size(&mut self, size: CacheSize) { + self.inner.set_prepared_statement_cache_size(size) + } +} + +/// The database `probe` ([`DATABASE_SQL`]) reads; `None` when it fails, +/// and the read goes live without asking the engine, the next miss +/// reading it again. +async fn read_database(probe: LoadFuture<'_>) -> Option { + let rows: Vec = probe.await.ok()?.try_collect().await.ok()?; + rows.first() + .and_then(|row| >::build_from_row(row).ok()) +} + +/// The engine's rows for a key whose tables' types the engine side does +/// not know yet: every table's types are read from the catalog, not only +/// the missing ones, since the engine side keeps them and the next query +/// on another table then needs no lookup at all. A table created since +/// this load is not among them, so the next miss loads the catalog +/// again. A connection that already loaded the catalog has no lookup to +/// make: the types it lacks are stale, and the read goes live. +async fn with_types( + subscriptions: &dyn Subscriptions, + key: &QueryKey, + snapshot: Snapshot, + missing: &[String], + catalog: Option>, + catalog_loaded: &AtomicBool, +) -> QueryResult { + let Some(catalog) = catalog else { + catalog_loaded.store(false, Ordering::Release); + subscriptions.unsubscribe(key); + return Ok(Verdict::Refused(STALE_TYPES.to_string())); + }; + let rows: Vec = catalog.await?.try_collect().await?; + let attributes = rows + .iter() + .map(|row| { + >::build_from_row(row) + .map_err(Error::DeserializationError) + }) + .collect::>>()?; + let all = all_table_types(attributes); + catalog_loaded.store( + missing.iter().all(|table| all.contains_key(table)), + Ordering::Release, + ); + let types: HashMap> = all + .into_iter() + .map(|(table, found)| (table, Arc::new(found))) + .collect(); + Ok(subscriptions.materialize(key, snapshot, types)) +} + +/// A subscription made for a read whose future may be dropped before +/// the fill (while the subscribe runs on its blocking thread, during +/// the `Accepted` read of Postgres, during the catalog lookup). Armed +/// until the fill; dropped armed, it ends the subscription through the +/// cache, which also fails any other reader's fill in flight for the +/// key, and raises `gone` for the blocking thread. Ending it twice is +/// harmless: the cache side is idempotent, and a subscription another +/// read made since is protected by that reader's own generation. +struct Unfilled { + cache: Arc, + key: QueryKey, + gone: Arc, + armed: bool, +} + +impl Unfilled { + /// The read filled: the subscription is the cache's from here. + fn disarm(mut self) { + self.armed = false; + } +} + +impl Drop for Unfilled { + /// Ends the subscription if the read never filled. + fn drop(&mut self) { + if self.armed { + self.gone.store(true, Ordering::Release); + self.cache.abandon(&self.key); + } + } +} + +/// Lets bb8 and deadpool pool the wrapper. diesel-async's own check +/// also asks the socket whether it closed, which it keeps private; a +/// pool that verifies connections on checkout (the default) catches +/// that on its ping. +impl PoolableConnection for AsyncSyncedConnection { + /// Whether the transaction manager says the connection is broken. + fn is_broken(&mut self) -> bool { + AnsiTransactionManager::is_broken_transaction_manager(self) + } +} diff --git a/crates/diesel-streamgres/src/cache.rs b/crates/diesel-streamgres/src/cache.rs new file mode 100644 index 0000000..29a06f4 --- /dev/null +++ b/crates/diesel-streamgres/src/cache.rs @@ -0,0 +1,974 @@ +//! The result cache: rows by query key, bounded by bytes, kept fresh by the +//! engine's touches. +//! +//! An entry is one of three things. The rows one Postgres read +//! returned, byte for byte, for a query the engine reports touches for. +//! The engine's own rows for a query it serves ([`Live`]): the result is +//! built from them when a read asks for it after a change, so a burst of +//! deltas costs one build, on the reading thread. Or a mark that the +//! query is stale: the engine reported it touched, and the next read +//! refills it from Postgres. The mark weighs only its key; it is kept so +//! that the clocks below still end the subscription of a query nobody +//! reads again. +//! +//! An entry is evicted when the cache is over its byte budget or the +//! entry has expired ([`Limits`]), either of which also ends the +//! subscription. Two clocks bound an entry's life. Idle time drops a +//! query nobody has read for a while, with its subscription, so a +//! backend's cache follows what it reads now; the engine changing a +//! query's rows is not a read. Time to live drops an entry on a timer +//! however often it is read: a safety net that makes the next read +//! resubscribe and refill from the engine, off by default since the +//! engine keeps the entry fresh. +//! +//! Between subscribing and filling there is a window: the read runs, and a +//! touch may arrive for the key before the rows are in. A generation per +//! key closes it. [`Cache::reserve`] takes the generation before the read, +//! [`Cache::fill`] stores the rows only if no touch bumped it since. A +//! generation exists only while a read holds a [`Reservation`] on its +//! key: a touch on a key nobody is filling has nothing to invalidate but +//! the entry, so the map holds the keys in flight, not every key ever +//! seen. +//! +//! The cache also keeps the one clock the read-your-writes window needs: +//! when any connection through it last wrote. A pool's connections +//! share a cache, so a write on one is followed by reads on another, and +//! those must see it too. + +use std::collections::HashMap; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex, OnceLock}; +use std::time::{Duration, Instant}; + +use diesel::pg::PgMetadataCache; +use moka::notification::RemovalCause; + +use crate::key::QueryKey; +use crate::row::OwnedRow; +use crate::subscriptions::Subscriptions; + +/// Rows the engine keeps current for one query, as an entry holds them. +pub trait Live: Send + Sync + std::fmt::Debug { + /// The result as of the latest change, built now if a change came + /// since it last was. `None` once the rows cannot be served (a value + /// arrived that cannot be encoded): the query then falls back to + /// refills, and the implementation marks it stale. + fn rows(&self) -> Option>; + + /// What the rows held weigh. + fn bytes(&self) -> usize; +} + +/// What an entry holds. +#[derive(Debug)] +enum Content { + /// The rows a Postgres read returned. + Rows(Arc<[OwnedRow]>), + /// The engine's rows, built into the result when read. + Live(Arc), + /// Touched: the next read refills. + Stale, +} + +/// One cached query. +#[derive(Debug)] +pub struct Entry { + content: Content, + bytes: usize, +} + +impl Entry { + /// A Postgres read's rows. + fn rows(rows: Arc<[OwnedRow]>, key: &QueryKey) -> Entry { + let bytes = key.byte_len() + rows_bytes(&rows); + Entry { + content: Content::Rows(rows), + bytes, + } + } + + /// The engine's rows: `built` is what the result weighs, when it has + /// just been built; otherwise it is taken to weigh what the rows do. + fn live(live: Arc, key: &QueryKey, built: Option) -> Entry { + let held = live.bytes(); + let bytes = key.byte_len() + held + built.unwrap_or(held); + Entry { + content: Content::Live(live), + bytes, + } + } + + /// The mark of a touched query. + fn stale(key: &QueryKey) -> Entry { + Entry { + content: Content::Stale, + bytes: key.byte_len(), + } + } + + /// What the entry weighs against the byte budget. + pub fn bytes(&self) -> usize { + self.bytes + } +} + +/// What a result's rows weigh. +fn rows_bytes(rows: &[OwnedRow]) -> usize { + rows.iter().map(OwnedRow::byte_len).sum() +} + +/// How many bytes of keys of refused queries are remembered at most. +const REFUSED_BYTES: u64 = 8 * 1024 * 1024; + +/// What bounds the cache: bytes, idle time, time to live, and how long a +/// refusal is remembered. From the environment through +/// [`Limits::from_env`]: +/// +/// | variable | meaning | default | +/// |------------------------------------|----------------------------------------------|-----------| +/// | `STREAMGRES_CACHE_MAX_BYTES` | rows and keys held, in bytes | 268435456 | +/// | `STREAMGRES_CACHE_MAX_ENTRY_BYTES` | the largest result kept; 0 no limit | 16777216 | +/// | `STREAMGRES_CACHE_IDLE_MS` | drop an entry unread for this long; 0 never | 1800000 | +/// | `STREAMGRES_CACHE_TTL_MS` | drop an entry this old however read; 0 never | 0 | +/// | `STREAMGRES_REFUSED_TTL_MS` | do not re-ask the engine about a refused query for this long; 0 always ask | 30000 | +/// +/// A value that is missing or not a number leaves the default. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Limits { + pub max_bytes: u64, + /// A result larger than this is served but not kept, and its + /// subscription ends: one query must not take the whole budget. + pub max_entry_bytes: u64, + pub time_to_idle: Option, + pub time_to_live: Option, + /// A query the engine refused (a shape outside its subset) would + /// otherwise cost a round trip to the engine on every read; for this + /// long it goes straight to Postgres instead. + pub refused_for: Option, +} + +/// The cache's byte budget when the environment says nothing: 256 MiB. +pub const DEFAULT_MAX_BYTES: u64 = 256 * 1024 * 1024; +/// The largest result kept when the environment says nothing: 16 MiB. +pub const DEFAULT_MAX_ENTRY_BYTES: u64 = 16 * 1024 * 1024; +/// How long an unread entry stays when the environment says nothing. +pub const DEFAULT_TIME_TO_IDLE: Duration = Duration::from_secs(30 * 60); +/// How long a refusal is remembered when the environment says nothing. +pub const DEFAULT_REFUSED_FOR: Duration = Duration::from_secs(30); + +impl Default for Limits { + /// The documented defaults: idle expiry on, time to live off. + fn default() -> Limits { + Limits { + max_bytes: DEFAULT_MAX_BYTES, + max_entry_bytes: DEFAULT_MAX_ENTRY_BYTES, + time_to_idle: Some(DEFAULT_TIME_TO_IDLE), + time_to_live: None, + refused_for: Some(DEFAULT_REFUSED_FOR), + } + } +} + +impl Limits { + /// The limits from the environment, each missing or unreadable one + /// at its default; a zero turns a clock or the entry limit off. + pub fn from_env() -> Limits { + let defaults = Limits::default(); + let millis = |name: &str, default: Option| match env_number(name) { + Some(0) => None, + Some(millis) => Some(Duration::from_millis(millis)), + None => default, + }; + Limits { + max_bytes: env_number("STREAMGRES_CACHE_MAX_BYTES").unwrap_or(defaults.max_bytes), + max_entry_bytes: match env_number("STREAMGRES_CACHE_MAX_ENTRY_BYTES") { + Some(0) => u64::MAX, + Some(bytes) => bytes, + None => defaults.max_entry_bytes, + }, + time_to_idle: millis("STREAMGRES_CACHE_IDLE_MS", defaults.time_to_idle), + time_to_live: millis("STREAMGRES_CACHE_TTL_MS", defaults.time_to_live), + refused_for: millis("STREAMGRES_REFUSED_TTL_MS", defaults.refused_for), + } + } +} + +/// A number from the environment. +fn env_number(name: &str) -> Option { + std::env::var(name).ok()?.trim().parse::().ok() +} + +/// Counters for dashboards and tests. +#[derive(Debug, Default)] +pub struct Stats { + pub hits: AtomicU64, + pub misses: AtomicU64, + /// Reads the engine (or the client's own limits) refused; each one + /// asked. + pub refused: AtomicU64, + /// Reads sent to Postgres without asking because the query was + /// refused recently. + pub refused_recently: AtomicU64, + pub touched: AtomicU64, + /// Entries dropped for the byte budget. + pub evicted: AtomicU64, + /// Entries dropped by the idle or time-to-live clock. + pub expired: AtomicU64, + /// Fills dropped because the key was touched during the read. + pub stale_fills: AtomicU64, + /// Results served but not kept because they were larger than + /// [`Limits::max_entry_bytes`]. + pub oversized: AtomicU64, + /// Entries the engine's own rows replaced (rows mode). + pub replaced: AtomicU64, + /// Reads sent to Postgres because their connection had just written. + pub after_write: AtomicU64, + /// Reads sent to Postgres because the read's lock (`FOR UPDATE`) + /// cannot be taken on a cache. + pub locking: AtomicU64, + /// Reads sent to Postgres because the engine was further behind the + /// database than the cache may be. + pub lagging: AtomicU64, +} + +impl Stats { + /// Counts one more. + fn bump(counter: &AtomicU64) { + counter.fetch_add(1, Ordering::Relaxed); + } +} + +/// A key's generation while reads are in flight for it. +#[derive(Default)] +struct Slot { + generation: u64, + inflight: u32, +} + +/// The generation of every key with a read in flight. +type Generations = Arc>>; + +/// A read's claim on a key between reserving it and filling it. Dropping +/// it, with or without a fill, lets the key's generation go once no +/// other read holds one. +pub struct Reservation { + generations: Generations, + key: QueryKey, + generation: u64, +} + +impl Reservation { + /// The key reserved. + pub fn key(&self) -> &QueryKey { + &self.key + } +} + +impl Drop for Reservation { + /// Lets the key's generation go once no other read holds one. + fn drop(&mut self) { + let mut generations = self.generations.lock().expect("generations"); + if let Some(slot) = generations.get_mut(&self.key) { + slot.inflight = slot.inflight.saturating_sub(1); + if slot.inflight == 0 { + generations.remove(&self.key); + } + } + } +} + +/// Query results by key, kept current by the engine side. +pub struct Cache { + entries: moka::sync::Cache>, + generations: Generations, + subscriptions: Arc, + /// Queries refused recently, with their values, each with the + /// reason, dropped on their own clock. By the whole key, values and + /// all: the engine refuses on values too (a time at infinity, more + /// rows than it follows for this user), and a refusal for one value + /// must not send every other value of the query to Postgres. + refused: Option>>, + limits: Limits, + /// When a connection through this cache last wrote, as nanoseconds + /// since `started` plus one; 0 is never. + last_write: AtomicU64, + started: Instant, + /// The OIDs of user-defined types, resolved once per database for + /// the async wrapper (the blocking one asks its connection). + type_metadata: Arc>, + /// The database this cache's reads go to, read once through the + /// first connection that misses; `None` inside when it could not be + /// read. The engine side compares it with the one its engine streams. + database: OnceLock>, + pub stats: Arc, +} + +/// A database as Postgres identifies it to any role: the cluster's +/// `system_identifier` and the database's name. Two clusters restored +/// from one backup share an identifier, so the name is kept too. +#[derive(Debug, Clone, PartialEq, Eq, diesel::QueryableByName)] +pub struct Database { + #[diesel(sql_type = diesel::sql_types::Text)] + pub system: String, + #[diesel(sql_type = diesel::sql_types::Text)] + pub name: String, +} + +/// Reads a [`Database`] through a connection; the engine announces the +/// same in its `hello`. +pub const DATABASE_SQL: &str = "SELECT system_identifier::text AS system, current_database()::text AS name FROM pg_control_system()"; + +impl Cache { + /// A cache holding at most `max_bytes` of rows and keys, with the + /// default clocks, subscribing through `subscriptions`. + pub fn new(max_bytes: u64, subscriptions: Arc) -> Cache { + Self::with_limits( + Limits { + max_bytes, + ..Limits::default() + }, + subscriptions, + ) + } + + /// A cache bounded by `limits`, subscribing through `subscriptions`. + pub fn with_limits(limits: Limits, subscriptions: Arc) -> Cache { + let stats = Arc::new(Stats::default()); + let generations: Generations = Arc::default(); + let listener_subscriptions = Arc::clone(&subscriptions); + let listener_stats = Arc::clone(&stats); + let listener_generations = Arc::clone(&generations); + let mut builder = moka::sync::Cache::builder() + .max_capacity(limits.max_bytes) + .weigher(|_: &QueryKey, entry: &Arc| { + u32::try_from(entry.bytes).unwrap_or(u32::MAX) + }) + .eviction_listener(move |key: Arc, _: Arc, cause| { + let ended = match cause { + RemovalCause::Size => &listener_stats.evicted, + RemovalCause::Expired => &listener_stats.expired, + RemovalCause::Explicit | RemovalCause::Replaced => return, + }; + Stats::bump(ended); + bump(&listener_generations, &key); + listener_subscriptions.unsubscribe(&key); + }); + if let Some(idle) = limits.time_to_idle { + builder = builder.time_to_idle(idle); + } + if let Some(live) = limits.time_to_live { + builder = builder.time_to_live(live); + } + let refused = limits.refused_for.map(|window| { + moka::sync::Cache::builder() + .max_capacity(REFUSED_BYTES) + .weigher(|key: &QueryKey, _: &Arc| { + u32::try_from(key.byte_len()).unwrap_or(u32::MAX) + }) + .time_to_live(window) + .build() + }); + Cache { + entries: builder.build(), + generations, + subscriptions, + refused, + limits, + last_write: AtomicU64::new(0), + started: Instant::now(), + type_metadata: Arc::new(Mutex::new(PgMetadataCache::new())), + database: OnceLock::new(), + stats, + } + } + + /// Counts a refusal of `key`, and when the engine side says the + /// reason will hold, remembers it so reads of the same query with the + /// same values skip the engine for [`Limits::refused_for`]. + pub fn refused(&self, key: &QueryKey, reason: &str) { + Stats::bump(&self.stats.refused); + if let Some(refused) = &self.refused + && self.subscriptions.refusal_is_final(reason) + { + refused.insert(key.clone(), Arc::from(reason)); + } + } + + /// Why `key` was refused recently, if it was; counts the skipped ask. + pub fn refused_recently(&self, key: &QueryKey) -> Option> { + let reason = self.refused.as_ref()?.get(key)?; + Stats::bump(&self.stats.refused_recently); + Some(reason) + } + + /// Forgets every remembered refusal: the engine connection dropped, + /// and the engine that answers next may be a restarted one with a + /// catalog that knows more (a type created after the old one + /// started, say), so each refused shape is asked about once more. + pub fn forget_refusals(&self) { + if let Some(refused) = &self.refused { + refused.invalidate_all(); + } + } + + /// What bounds the cache. + pub fn limits(&self) -> &Limits { + &self.limits + } + + /// The rows for a key, if the cache has them: a Postgres read's as + /// they were, the engine's built as of its latest change. A stale + /// entry has none. + pub fn get(&self, key: &QueryKey) -> Option> { + let rows = self.entries.get(key).and_then(|entry| match &entry.content { + Content::Rows(rows) => Some(Arc::clone(rows)), + Content::Live(live) => live.rows(), + Content::Stale => None, + }); + Stats::bump(match rows { + Some(_) => &self.stats.hits, + None => &self.stats.misses, + }); + rows + } + + /// The engine side the cache subscribes through. + pub fn subscriptions(&self) -> &dyn Subscriptions { + &*self.subscriptions + } + + /// The same, owned, for a subscribe on another thread. + pub fn shared_subscriptions(&self) -> Arc { + Arc::clone(&self.subscriptions) + } + + /// The key's generation, taken before the read that will fill it. + pub fn reserve(&self, key: &QueryKey) -> Reservation { + let generation = { + let mut generations = self.generations.lock().expect("generations"); + let slot = generations.entry(key.clone()).or_default(); + slot.inflight += 1; + slot.generation + }; + Reservation { + generations: Arc::clone(&self.generations), + key: key.clone(), + generation, + } + } + + /// Stores the rows a Postgres read returned, unless the key was + /// touched since the reservation was taken or the result is over the + /// entry limit. Returns whether they were stored. + pub fn fill(&self, reservation: Reservation, rows: Arc<[OwnedRow]>) -> bool { + let entry = Entry::rows(rows, &reservation.key); + self.fill_entry(reservation, entry) + } + + /// [`Cache::fill`] for a query the engine serves from its own rows: + /// `rows` is the result just built from them, which a later read + /// gets again until a change. + pub fn fill_live( + &self, + reservation: Reservation, + live: Arc, + rows: &[OwnedRow], + ) -> bool { + let entry = Entry::live(live, &reservation.key, Some(rows_bytes(rows))); + self.fill_entry(reservation, entry) + } + + /// The generation is checked before and after the insert: a touch + /// or a replace between the two would otherwise leave the older + /// rows in, since nothing holds a lock across the insert (moka runs + /// the eviction listener on the inserting thread, and the listener + /// takes the generations lock). Losing to something newer leaves the + /// key stale, which costs one refill, never a stale row, and keeps + /// its clocks running. + fn fill_entry(&self, reservation: Reservation, entry: Entry) -> bool { + if self.generation(&reservation.key) != reservation.generation { + Stats::bump(&self.stats.stale_fills); + return false; + } + if !self.insert(&reservation.key, entry) { + return false; + } + if self.generation(&reservation.key) != reservation.generation { + Stats::bump(&self.stats.stale_fills); + self.stale(&reservation.key); + return false; + } + true + } + + /// The engine side holds rows for the key afresh (a store built for + /// a resubscribe), or its rows weigh much more or less than when the + /// entry was weighed: store them, failing any fill in flight, whose + /// rows are older. The caller reserves the key before it builds the + /// store, so a subscription ended meanwhile (which bumps the + /// generation once more) leaves the key stale again: the engine + /// side's writes are held to the same rule as a Postgres read's. + pub fn live(&self, reservation: Reservation, live: Arc) { + Stats::bump(&self.stats.replaced); + let key = &reservation.key; + bump(&self.generations, key); + let entry = Entry::live(live, key, None); + if self.insert(key, entry) && self.generation(key) != reservation.generation + 1 { + Stats::bump(&self.stats.stale_fills); + self.stale(key); + } + } + + /// The database this cache's reads go to, once it has been read: + /// `None` until then, and when it could not be read. + pub fn database(&self) -> Option<&Database> { + self.database.get().and_then(Option::as_ref) + } + + /// Whether the database has been looked for, found or not. + pub fn knows_database(&self) -> bool { + self.database.get().is_some() + } + + /// Records the database a connection read; the first record stands. + /// `None` says there is none to check the engine's against, and + /// trusts it: the wrapper records only a database it read, and goes + /// live on a read that failed. + pub fn learn_database(&self, database: Option) { + let _ = self.database.set(database); + } + + /// The engine changed a query it serves from its own rows: counted; + /// the entry builds the new result on its next read. + pub fn changed(&self) { + Stats::bump(&self.stats.replaced); + } + + /// Fails any fill or replace in flight for `key`, without touching + /// the entry. For the engine side to call while it drops the key, + /// under its own lock, so that no rows built before the drop can + /// land after it. + pub fn bump_generation(&self, key: &QueryKey) { + bump(&self.generations, key); + } + + /// Stores `entry`, unless it is over the entry limit, in which case + /// the key is abandoned. Whether it was stored. + fn insert(&self, key: &QueryKey, entry: Entry) -> bool { + if entry.bytes as u64 > self.limits.max_entry_bytes { + Stats::bump(&self.stats.oversized); + self.abandon(key); + return false; + } + self.entries.insert(key.clone(), Arc::new(entry)); + true + } + + /// The engine reported a change to the query's rows: mark the key + /// stale, so the next read refills from Postgres, and fail any fill + /// in flight. + pub fn touched(&self, key: &QueryKey) { + Stats::bump(&self.stats.touched); + bump(&self.generations, key); + self.stale(key); + } + + /// Leaves `key` stale: no rows, its subscription kept, its clocks + /// running. + fn stale(&self, key: &QueryKey) { + self.entries.insert(key.clone(), Arc::new(Entry::stale(key))); + } + + /// Ends the key's subscription from the cache's side: a read gave up + /// on it (dropped, or its Postgres read failed) or its result was + /// too large to keep. The generation is bumped first, so a fill in + /// flight for the key by another read is dropped rather than kept + /// with nothing to update it: the sequence the eviction listener + /// uses, and the only way the wrapper may unsubscribe. + /// + /// The subscription goes first: the engine side bumps the generation + /// as it drops the key, so an engine delta that checked the key + /// before this point and inserts after it is dropped by its own + /// re-check. Then the bump again for a Postgres fill in flight, and + /// the entry. + pub fn abandon(&self, key: &QueryKey) { + self.subscriptions.unsubscribe(key); + bump(&self.generations, key); + self.entries.invalidate(key); + } + + /// A connection through this cache wrote. + pub fn mark_write(&self) { + let nanos = u64::try_from(self.started.elapsed().as_nanos()).unwrap_or(u64::MAX - 1); + self.last_write.fetch_max(nanos + 1, Ordering::Release); + } + + /// Whether a connection through this cache wrote within `window`. + pub fn wrote_within(&self, window: Duration) -> bool { + if window.is_zero() { + return false; + } + match self.last_write.load(Ordering::Acquire) { + 0 => false, + marked => { + let since = self + .started + .elapsed() + .saturating_sub(Duration::from_nanos(marked - 1)); + since < window + } + } + } + + /// The OIDs of user-defined types resolved so far for this database. + pub fn type_metadata(&self) -> &Arc> { + &self.type_metadata + } + + /// Runs the maintenance the clocks need, so expired entries are + /// dropped (and unsubscribed) now rather than on the next access. + /// moka does this on its own as the cache is used; this is for + /// tests and quiet processes. + pub fn run_pending_tasks(&self) { + self.entries.run_pending_tasks(); + } + + /// Leaves every entry stale and fails every fill in flight: every + /// key's next read refills, and every subscription keeps its clocks. + pub fn clear(&self) { + for slot in self.generations.lock().expect("generations").values_mut() { + slot.generation += 1; + } + let keys: Vec> = self.entries.iter().map(|(key, _)| key).collect(); + for key in keys { + self.stale(&key); + } + } + + /// How many entries are in, after pending maintenance. + pub fn len(&self) -> u64 { + self.entries.run_pending_tasks(); + self.entries.entry_count() + } + + /// Whether no entry is in, stale ones counted. + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// The bytes the entries weigh, after pending maintenance. + pub fn weighted_size(&self) -> u64 { + self.entries.run_pending_tasks(); + self.entries.weighted_size() + } + + /// How many keys have a read in flight, for tests. + pub fn reservations(&self) -> usize { + self.generations.lock().expect("generations").len() + } + + /// The key's generation now; 0 when no read is in flight. + fn generation(&self, key: &QueryKey) -> u64 { + self.generations + .lock() + .expect("generations") + .get(key) + .map_or(0, |slot| slot.generation) + } +} + +/// Fails every fill in flight for `key`. A key with none has no +/// generation to bump: its entry, if any, is what the caller drops. +fn bump(generations: &Generations, key: &QueryKey) { + if let Some(slot) = generations.lock().expect("generations").get_mut(key) { + slot.generation += 1; + } +} + +impl std::fmt::Debug for Cache { + /// The entry count and the counters. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Cache") + .field("entries", &self.entries.entry_count()) + .field("stats", &self.stats) + .finish() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::materialized::Snapshot; + use crate::subscriptions::{Outcome, Unsynced, Verdict}; + use crate::types::TableTypes; + + /// A key for `sql` with no binds. + fn key(sql: &str) -> QueryKey { + QueryKey { + sql: sql.to_string(), + binds: Vec::new(), + } + } + + /// An empty result. + fn no_rows() -> Arc<[OwnedRow]> { + Arc::from(Vec::new()) + } + + /// Accepts everything and records what is unsubscribed. + #[derive(Default)] + struct Recording { + unsubscribed: Mutex>, + } + + impl Subscriptions for Recording { + /// Accepted. + fn subscribe(&self, _: &QueryKey) -> Outcome { + Outcome::Verdict(Verdict::Accepted) + } + + /// Accepted, with no rows. + fn materialize( + &self, + _: &QueryKey, + _: Snapshot, + _: HashMap>, + ) -> Verdict { + Verdict::Accepted + } + + /// Records the key. + fn unsubscribe(&self, key: &QueryKey) { + self.unsubscribed.lock().unwrap().push(key.clone()); + } + } + + /// Two readers of one key; the second gives up (its future was + /// dropped, or its read failed) and ends the subscription. The first + /// must not then store an entry nothing will invalidate. + #[test] + fn abandoning_a_key_fails_the_other_readers_fill() { + let engine = Arc::new(Recording::default()); + let cache = Cache::new(1 << 20, engine.clone()); + let k = key("SELECT 1"); + let first = cache.reserve(&k); + let second = cache.reserve(&k); + drop(second); + cache.abandon(&k); + assert_eq!(engine.unsubscribed.lock().unwrap().len(), 1); + assert!( + !cache.fill(first, no_rows()), + "the first reader's fill is refused" + ); + assert!(cache.is_empty()); + assert_eq!(cache.reservations(), 0); + } + + /// A rows-mode caller holds the snapshot's rows; a delta replaces the + /// entry before the caller fills. The older rows must lose. + #[test] + fn engine_rows_beat_a_fill_in_flight() { + let cache = Cache::new(1 << 20, Arc::new(Unsynced)); + let k = key("SELECT 1"); + let reservation = cache.reserve(&k); + let columns: crate::row::Columns = Arc::from(vec!["a".to_string()]); + let newer: Arc<[OwnedRow]> = Arc::from(vec![OwnedRow::new(columns, vec![None])]); + let engine_side = cache.reserve(&k); + cache.live(engine_side, Arc::new(Fixed(Arc::clone(&newer)))); + assert!( + !cache.fill(reservation, no_rows()), + "the snapshot's rows are older than the delta's" + ); + assert_eq!( + cache.get(&k).map(|rows| rows.len()), + Some(1), + "the delta's rows stay" + ); + assert_eq!(cache.stats.stale_fills.load(Ordering::Relaxed), 1); + } + + /// The engine side reserved before building its rows; the key is + /// abandoned (which bumps as the subscription goes) in between, so + /// the rows are not served. + #[test] + fn a_replace_after_the_key_was_abandoned_does_not_stick() { + let engine = Arc::new(Recording::default()); + let cache = Cache::new(1 << 20, engine.clone()); + let k = key("SELECT 1"); + let engine_side = cache.reserve(&k); + cache.abandon(&k); + cache.live(engine_side, Arc::new(Fixed(no_rows()))); + assert!( + cache.get(&k).is_none(), + "rows built for a subscription that ended are not served" + ); + } + + /// The engine's rows over the entry limit end the subscription, and + /// with it another reader's smaller fill. + #[test] + fn an_oversized_result_ends_the_subscription_and_the_other_fill() { + let engine = Arc::new(Recording::default()); + let cache = Cache::with_limits( + Limits { + max_entry_bytes: 8, + ..Limits::default() + }, + engine.clone(), + ); + let k = key("SELECT a long enough statement"); + let other = cache.reserve(&k); + let engine_side = cache.reserve(&k); + cache.live(engine_side, Arc::new(Fixed(no_rows()))); + assert_eq!(engine.unsubscribed.lock().unwrap().len(), 1); + assert!( + !cache.fill(other, no_rows()), + "the other reader's smaller fill cannot keep a key with no subscription" + ); + } + + /// The defaults are the documented ones. + #[test] + fn limits_default_to_idle_on_and_ttl_off() { + let limits = Limits::default(); + assert_eq!(limits.max_bytes, DEFAULT_MAX_BYTES); + assert_eq!(limits.max_entry_bytes, DEFAULT_MAX_ENTRY_BYTES); + assert_eq!(limits.time_to_idle, Some(DEFAULT_TIME_TO_IDLE)); + assert_eq!(limits.time_to_live, None); + assert_eq!(env_number("STREAMGRES_CACHE_TEST_UNSET_VARIABLE"), None); + } + + /// A key's generation lives only while a read of it is in flight: + /// touches, evictions and expiries of keys nobody is reading leave + /// nothing behind, and a reservation dropped without a fill goes too. + #[test] + fn generations_live_only_while_a_read_is_in_flight() { + let cache = Cache::new(1 << 20, Arc::new(Unsynced)); + let k = key("SELECT 1"); + assert_eq!(cache.reservations(), 0); + let reservation = cache.reserve(&k); + assert_eq!(cache.reservations(), 1); + cache.touched(&k); + assert!( + !cache.fill(reservation, no_rows()), + "touched during the read" + ); + assert_eq!( + cache.reservations(), + 0, + "the reservation went with the fill" + ); + for i in 0..1000 { + cache.touched(&key(&format!("SELECT {i}"))); + } + assert_eq!(cache.reservations(), 0); + let dropped = cache.reserve(&k); + drop(dropped); + assert_eq!(cache.reservations(), 0); + } + + /// A read that started before a touch cannot fill after a read that + /// started after it did. + #[test] + fn an_older_reservation_cannot_fill_after_a_newer_one_did() { + let cache = Cache::new(1 << 20, Arc::new(Unsynced)); + let k = key("SELECT 1"); + let older = cache.reserve(&k); + cache.touched(&k); + let newer = cache.reserve(&k); + assert!( + cache.fill(newer, no_rows()), + "the newer read saw the change" + ); + assert!(!cache.fill(older, no_rows()), "the older read did not"); + assert_eq!(cache.reservations(), 0); + } + + /// A result over the entry limit is not kept, from Postgres or from + /// the engine. + #[test] + fn a_result_over_the_entry_limit_is_not_kept() { + let cache = Cache::with_limits( + Limits { + max_entry_bytes: 8, + ..Limits::default() + }, + Arc::new(Unsynced), + ); + let k = key("SELECT a long enough statement"); + let reservation = cache.reserve(&k); + assert!(!cache.fill(reservation, no_rows())); + assert_eq!(cache.stats.oversized.load(Ordering::Relaxed), 1); + assert!(cache.is_empty()); + let engine_side = cache.reserve(&k); + cache.live(engine_side, Arc::new(Fixed(no_rows()))); + assert!( + cache.is_empty(), + "the engine's rows are bounded the same way" + ); + } + + /// The write clock is the cache's, and a zero window is off. + #[test] + fn the_write_clock_is_shared() { + let cache = Cache::new(1 << 20, Arc::new(Unsynced)); + assert!(!cache.wrote_within(Duration::from_secs(1))); + cache.mark_write(); + assert!(cache.wrote_within(Duration::from_secs(1))); + assert!(!cache.wrote_within(Duration::ZERO), "a zero window is off"); + } + + /// A refusal is remembered for the query with its values: the + /// engine refuses on values too, and another value asks again. + #[test] + fn a_refusal_is_remembered_for_its_values() { + let cache = Cache::new(1 << 20, Arc::new(Unsynced)); + let mut a = key("SELECT $1"); + a.binds.push(crate::key::Bind { + type_oid: 25, + bytes: Some(b"a".to_vec()), + }); + let mut b = a.clone(); + b.binds[0].bytes = Some(b"b".to_vec()); + cache.refused(&a, "unsupported"); + assert!(cache.refused_recently(&a).is_some()); + assert!(cache.refused_recently(&b).is_none()); + } + + /// A touched key stays stale rather than gone, so the idle clock + /// still ends the subscription of a query nobody reads again. + #[test] + fn a_touched_key_unread_still_ends_its_subscription() { + let engine = Arc::new(Recording::default()); + let cache = Cache::with_limits( + Limits { + time_to_idle: Some(Duration::from_millis(50)), + ..Limits::default() + }, + engine.clone(), + ); + let k = key("SELECT 1"); + assert!(cache.fill(cache.reserve(&k), no_rows())); + cache.touched(&k); + assert!(cache.get(&k).is_none(), "stale: the next read refills"); + std::thread::sleep(Duration::from_millis(120)); + cache.run_pending_tasks(); + assert_eq!(*engine.unsubscribed.lock().unwrap(), vec![k]); + } + + /// The engine's rows for one query, fixed. + #[derive(Debug)] + struct Fixed(Arc<[OwnedRow]>); + + impl Live for Fixed { + /// The fixed rows. + fn rows(&self) -> Option> { + Some(Arc::clone(&self.0)) + } + + /// Weighs nothing. + fn bytes(&self) -> usize { + 0 + } + } +} diff --git a/crates/diesel-streamgres/src/cursor.rs b/crates/diesel-streamgres/src/cursor.rs new file mode 100644 index 0000000..7456ecc --- /dev/null +++ b/crates/diesel-streamgres/src/cursor.rs @@ -0,0 +1,142 @@ +//! One cursor type for both sources of rows. +//! +//! Diesel's `LoadConnection` names a single cursor and row type per +//! connection, so a read served from the cache and one served live by +//! Postgres come back through the same enum, the way a +//! `#[derive(MultiConnection)]` connection carries its backends' rows. + +use diesel::QueryResult; +use diesel::pg::{Pg, PgValue}; +use diesel::row::{Field, PartialRow, Row, RowIndex, RowSealed}; + +use crate::row::{OwnedCursor, OwnedField, OwnedRow}; + +/// Rows on their way to the application. +pub enum SyncedCursor { + Live(I), + Cached(OwnedCursor), +} + +impl Iterator for SyncedCursor +where + I: Iterator>, +{ + type Item = QueryResult>; + + /// The next row, live or cached. + fn next(&mut self) -> Option { + match self { + SyncedCursor::Live(inner) => inner.next().map(|row| row.map(SyncedRow::Live)), + SyncedCursor::Cached(inner) => inner.next().map(|row| row.map(SyncedRow::Cached)), + } + } + + /// The inner cursor's hint. + fn size_hint(&self) -> (usize, Option) { + match self { + SyncedCursor::Live(inner) => inner.size_hint(), + SyncedCursor::Cached(inner) => inner.size_hint(), + } + } +} + +impl ExactSizeIterator for SyncedCursor where I: ExactSizeIterator> {} + +/// One row, live or cached. +pub enum SyncedRow { + Live(R), + Cached(OwnedRow), +} + +impl RowSealed for SyncedRow {} + +impl RowIndex for SyncedRow +where + R: RowIndex, +{ + /// The position, as the row behind it says. + fn idx(&self, idx: usize) -> Option { + match self { + SyncedRow::Live(row) => row.idx(idx), + SyncedRow::Cached(row) => row.idx(idx), + } + } +} + +impl<'a, R> RowIndex<&'a str> for SyncedRow +where + R: RowIndex<&'a str>, +{ + /// The position of a named column, as the row behind it says. + fn idx(&self, idx: &'a str) -> Option { + match self { + SyncedRow::Live(row) => row.idx(idx), + SyncedRow::Cached(row) => row.idx(idx), + } + } +} + +impl<'a, R> Row<'a, Pg> for SyncedRow +where + R: Row<'a, Pg>, +{ + type Field<'f> + = SyncedField<'f, R::Field<'f>> + where + 'a: 'f, + Self: 'f; + type InnerPartialRow = Self; + + /// How many columns the row has. + fn field_count(&self) -> usize { + match self { + SyncedRow::Live(row) => row.field_count(), + SyncedRow::Cached(row) => row.field_count(), + } + } + + /// One column, by position or name. + fn get<'b, I>(&'b self, idx: I) -> Option> + where + 'a: 'b, + Self: RowIndex, + { + let index = self.idx(idx)?; + match self { + SyncedRow::Live(row) => row.get(index).map(SyncedField::Live), + SyncedRow::Cached(row) => row.get(index).map(SyncedField::Cached), + } + } + + /// The columns in `range`, for a nested `Queryable`. + fn partial_row(&self, range: std::ops::Range) -> PartialRow<'_, Self::InnerPartialRow> { + PartialRow::new(self, range) + } +} + +/// One column of a [`SyncedRow`]. +pub enum SyncedField<'a, F> { + Live(F), + Cached(OwnedField<'a>), +} + +impl<'a, F> Field<'a, Pg> for SyncedField<'a, F> +where + F: Field<'a, Pg>, +{ + /// The column's name. + fn field_name(&self) -> Option<&str> { + match self { + SyncedField::Live(field) => field.field_name(), + SyncedField::Cached(field) => field.field_name(), + } + } + + /// The column's value; `None` for `NULL`. + fn value(&self) -> Option> { + match self { + SyncedField::Live(field) => field.value(), + SyncedField::Cached(field) => field.value(), + } + } +} diff --git a/crates/diesel-streamgres/src/encode.rs b/crates/diesel-streamgres/src/encode.rs new file mode 100644 index 0000000..1fbcc68 --- /dev/null +++ b/crates/diesel-streamgres/src/encode.rs @@ -0,0 +1,468 @@ +//! JSON cells from the engine, encoded as the Postgres binary Diesel +//! decodes. +//! +//! The engine ships a row as JSON: strings, integers, floats, booleans, +//! JSON documents, times as microseconds since the Unix epoch, arrays. +//! Each cell is encoded for the column's exact Postgres type, so a cached +//! row decodes through the same `FromSql` as a live one. A type whose +//! value cannot be reproduced exactly from what the engine has is +//! refused ([`Unsupported`]), and the query keeps refilling from Postgres +//! instead: a `numeric` (the engine holds a float), ranges, intervals, +//! composites. A time comes back exact at any declared precision, since +//! the engine keeps Postgres's own microseconds; a float's `NaN` and +//! infinities come as the names Postgres gives them. + +use std::num::NonZeroU32; + +use serde_json::Value as Json; + +use crate::row::OwnedValue; +use crate::types::{ColumnType, TypeKind}; + +/// Why a column cannot be encoded from the engine's cell. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Unsupported(pub String); + +mod oid { + pub const BOOL: u32 = 16; + pub const BYTEA: u32 = 17; + pub const CHAR: u32 = 18; + pub const NAME: u32 = 19; + pub const INT8: u32 = 20; + pub const INT2: u32 = 21; + pub const INT4: u32 = 23; + pub const TEXT: u32 = 25; + pub const OID: u32 = 26; + pub const JSON: u32 = 114; + pub const FLOAT4: u32 = 700; + pub const FLOAT8: u32 = 701; + pub const BPCHAR: u32 = 1042; + pub const VARCHAR: u32 = 1043; + pub const DATE: u32 = 1082; + pub const TIMESTAMP: u32 = 1114; + pub const TIMESTAMPTZ: u32 = 1184; + pub const UUID: u32 = 2950; + pub const JSONB: u32 = 3802; +} + +const PG_EPOCH_MICROS: i64 = 946_684_800_000_000; + +/// Whether every value of this column can be encoded exactly; checked +/// once per query so an unsupported column refuses the whole query up +/// front rather than at the first odd row. +pub fn supports(column: &ColumnType) -> Result<(), Unsupported> { + match &column.kind { + TypeKind::Base => supports_scalar(column.oid, column.typmod, &column.name), + TypeKind::Enum => Ok(()), + TypeKind::Array { + element, + element_kind, + } => match element_kind.as_ref() { + TypeKind::Base => supports_scalar(*element, -1, &column.name), + TypeKind::Enum => Ok(()), + _ => Err(Unsupported(format!( + "column `{}` is an array of a type that cannot be encoded", + column.name + ))), + }, + TypeKind::Other => Err(Unsupported(format!( + "column `{}` has a type that cannot be encoded (oid {})", + column.name, column.oid + ))), + } +} + +/// Whether a base type can be encoded; a `numeric` cannot, since the +/// engine holds it as a float and its exact digits are gone. +fn supports_scalar(type_oid: u32, typmod: i32, name: &str) -> Result<(), Unsupported> { + match type_oid { + oid::BOOL + | oid::INT2 + | oid::INT4 + | oid::INT8 + | oid::OID + | oid::FLOAT4 + | oid::FLOAT8 + | oid::TEXT + | oid::VARCHAR + | oid::BPCHAR + | oid::NAME + | oid::CHAR + | oid::UUID + | oid::JSON + | oid::JSONB + | oid::BYTEA + | oid::DATE + | oid::TIMESTAMP + | oid::TIMESTAMPTZ => { + let _ = typmod; + Ok(()) + } + 1700 => Err(Unsupported(format!( + "column `{name}` is a numeric, which the engine holds as a float" + ))), + other => Err(Unsupported(format!( + "column `{name}` has type oid {other}, which cannot be encoded" + ))), + } +} + +/// One cell for one column; `None` for SQL `NULL`. +pub fn encode(cell: &Json, column: &ColumnType) -> Result, Unsupported> { + if cell.is_null() { + return Ok(None); + } + let type_oid = NonZeroU32::new(column.oid) + .ok_or_else(|| Unsupported(format!("column `{}` has type oid 0", column.name)))?; + let bytes = match &column.kind { + TypeKind::Base => scalar(cell, column.oid, &column.name)?, + TypeKind::Enum => text(cell, &column.name)?.into_bytes(), + TypeKind::Array { + element, + element_kind, + } => array(cell, *element, element_kind, &column.name)?, + TypeKind::Other => return Err(Unsupported(format!("column `{}`", column.name))), + }; + Ok(Some(OwnedValue::new(type_oid, bytes))) +} + +/// One base-type cell in Postgres's binary format. The engine carries a +/// `bytea` as Postgres's hex text, `\x` and two digits per byte; the +/// binary format is the bytes themselves. A time travels in +/// microseconds since the Unix epoch, and a date as such a time; +/// `infinity` and `-infinity` travel as the extreme `i64`s and become +/// Postgres's own binary infinities, so a cached row decodes exactly as +/// the live one does. +fn scalar(cell: &Json, type_oid: u32, name: &str) -> Result, Unsupported> { + Ok(match type_oid { + oid::BOOL => vec![u8::from( + cell.as_bool().ok_or_else(|| mismatch(name, "a boolean"))?, + )], + oid::INT2 => i16::try_from(integer(cell, name)?) + .map_err(|_| mismatch(name, "a 16-bit integer"))? + .to_be_bytes() + .to_vec(), + oid::INT4 => i32::try_from(integer(cell, name)?) + .map_err(|_| mismatch(name, "a 32-bit integer"))? + .to_be_bytes() + .to_vec(), + oid::OID => u32::try_from(integer(cell, name)?) + .map_err(|_| mismatch(name, "an oid"))? + .to_be_bytes() + .to_vec(), + oid::INT8 => integer(cell, name)?.to_be_bytes().to_vec(), + oid::FLOAT4 => (float(cell, name)? as f32).to_be_bytes().to_vec(), + oid::FLOAT8 => float(cell, name)?.to_be_bytes().to_vec(), + oid::TEXT | oid::VARCHAR | oid::BPCHAR | oid::NAME | oid::CHAR => { + text(cell, name)?.into_bytes() + } + oid::UUID => { + let text = text(cell, name)?; + let hex: String = text.chars().filter(|c| *c != '-').collect(); + let bytes = hex_bytes(&hex).ok_or_else(|| mismatch(name, "a UUID"))?; + if bytes.len() != 16 { + return Err(mismatch(name, "a UUID")); + } + bytes + } + oid::BYTEA => { + let text = text(cell, name)?; + let hex = text + .strip_prefix("\\x") + .ok_or_else(|| mismatch(name, "hex bytea"))?; + hex_bytes(hex).ok_or_else(|| mismatch(name, "hex bytea"))? + } + oid::JSON => serde_json::to_vec(cell).map_err(|_| mismatch(name, "JSON"))?, + oid::JSONB => { + let mut bytes = vec![1u8]; + bytes.extend(serde_json::to_vec(cell).map_err(|_| mismatch(name, "JSON"))?); + bytes + } + oid::TIMESTAMP | oid::TIMESTAMPTZ => match integer(cell, name)? { + i64::MAX => i64::MAX.to_be_bytes().to_vec(), + i64::MIN => i64::MIN.to_be_bytes().to_vec(), + micros => (micros - PG_EPOCH_MICROS).to_be_bytes().to_vec(), + }, + oid::DATE => match integer(cell, name)? { + i64::MAX => i32::MAX.to_be_bytes().to_vec(), + i64::MIN => i32::MIN.to_be_bytes().to_vec(), + micros => { + let days = (micros - PG_EPOCH_MICROS).div_euclid(86_400_000_000); + i32::try_from(days) + .map_err(|_| mismatch(name, "a date"))? + .to_be_bytes() + .to_vec() + } + }, + other => return Err(Unsupported(format!("column `{name}` has type oid {other}"))), + }) +} + +/// A one-dimensional array: the header (dimensions, a has-nulls flag, +/// the element type), one dimension (length, lower bound 1), then each +/// element as its length and bytes, `-1` for NULL. +fn array( + cell: &Json, + element: u32, + element_kind: &TypeKind, + name: &str, +) -> Result, Unsupported> { + let items = cell.as_array().ok_or_else(|| mismatch(name, "an array"))?; + let mut out = Vec::new(); + let has_nulls = items.iter().any(Json::is_null); + out.extend_from_slice(&1i32.to_be_bytes()); + out.extend_from_slice(&i32::from(has_nulls).to_be_bytes()); + out.extend_from_slice(&element.to_be_bytes()); + out.extend_from_slice(&(items.len() as i32).to_be_bytes()); + out.extend_from_slice(&1i32.to_be_bytes()); + for item in items { + if item.is_null() { + out.extend_from_slice(&(-1i32).to_be_bytes()); + continue; + } + let bytes = match element_kind { + TypeKind::Enum => text(item, name)?.into_bytes(), + _ => scalar(item, element, name)?, + }; + out.extend_from_slice(&(bytes.len() as i32).to_be_bytes()); + out.extend_from_slice(&bytes); + } + Ok(out) +} + +/// An integer cell; a float with no fraction counts as one. +fn integer(cell: &Json, name: &str) -> Result { + cell.as_i64() + .or_else(|| cell.as_f64().filter(|f| f.fract() == 0.0).map(|f| f as i64)) + .ok_or_else(|| mismatch(name, "an integer")) +} + +/// A float cell: a number, or the name of a value JSON has no number +/// for (`NaN`, `Infinity`, `-Infinity`), as the engine sends them. +fn float(cell: &Json, name: &str) -> Result { + match cell { + Json::String(text) => match text.as_str() { + "NaN" => Ok(f64::NAN), + "Infinity" => Ok(f64::INFINITY), + "-Infinity" => Ok(f64::NEG_INFINITY), + _ => Err(mismatch(name, "a number")), + }, + other => other.as_f64().ok_or_else(|| mismatch(name, "a number")), + } +} + +/// Decodes hex digits, two per byte. `None` for an odd length or +/// anything that is not an ASCII hex digit; never panics, whatever the +/// text holds. +fn hex_bytes(hex: &str) -> Option> { + let digits = hex.as_bytes(); + if !digits.len().is_multiple_of(2) { + return None; + } + digits + .chunks_exact(2) + .map(|pair| { + let high = (pair[0] as char).to_digit(16)?; + let low = (pair[1] as char).to_digit(16)?; + Some((high * 16 + low) as u8) + }) + .collect() +} + +/// A string cell. +fn text(cell: &Json, name: &str) -> Result { + match cell { + Json::String(text) => Ok(text.clone()), + other => Err(mismatch(name, &format!("a string, found {other}"))), + } +} + +/// A cell that is not what its column's type wants. +fn mismatch(name: &str, want: &str) -> Unsupported { + Unsupported(format!("column `{name}`: expected {want}")) +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + /// A base-type column named `c`. + fn base(oid: u32) -> ColumnType { + ColumnType { + name: "c".into(), + oid, + typmod: -1, + kind: TypeKind::Base, + } + } + + /// An encoded value's bytes. + fn bytes(value: Option) -> Vec { + value.expect("a value").bytes().to_vec() + } + + /// Each scalar takes its Postgres width; a value that does not fit is + /// refused. + #[test] + fn scalars_take_their_postgres_widths() { + assert_eq!( + bytes(encode(&json!(7), &base(oid::INT4)).unwrap()), + 7i32.to_be_bytes() + ); + assert_eq!( + bytes(encode(&json!(7), &base(oid::INT8)).unwrap()), + 7i64.to_be_bytes() + ); + assert_eq!( + bytes(encode(&json!(-2), &base(oid::INT2)).unwrap()), + (-2i16).to_be_bytes() + ); + assert!(encode(&json!(70000), &base(oid::INT2)).is_err()); + assert_eq!( + bytes(encode(&json!(2.5), &base(oid::FLOAT8)).unwrap()), + 2.5f64.to_be_bytes() + ); + assert_eq!( + bytes(encode(&json!("Infinity"), &base(oid::FLOAT8)).unwrap()), + f64::INFINITY.to_be_bytes() + ); + let nan = bytes(encode(&json!("NaN"), &base(oid::FLOAT8)).unwrap()); + assert!(f64::from_be_bytes(nan.try_into().unwrap()).is_nan()); + assert!(encode(&json!("lots"), &base(oid::FLOAT8)).is_err()); + assert_eq!(bytes(encode(&json!(true), &base(oid::BOOL)).unwrap()), [1]); + assert_eq!( + bytes(encode(&json!("hi"), &base(oid::TEXT)).unwrap()), + b"hi" + ); + assert_eq!(encode(&Json::Null, &base(oid::TEXT)).unwrap(), None); + assert_eq!( + bytes( + encode( + &json!("12345678-9abc-def0-1122-334455667788"), + &base(oid::UUID) + ) + .unwrap() + ), + [ + 0x12, 0x34, 0x56, 0x78, 0x9a, 0xbc, 0xde, 0xf0, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, + 0x77, 0x88 + ] + ); + } + + /// `jsonb` carries its version byte; times and dates count from + /// Postgres's epoch, and their infinities are Postgres's. + #[test] + fn json_and_times() { + let jsonb = bytes(encode(&json!({"a": [1, null]}), &base(oid::JSONB)).unwrap()); + assert_eq!(jsonb[0], 1); + assert_eq!( + serde_json::from_slice::(&jsonb[1..]).unwrap(), + json!({"a": [1, null]}) + ); + let micros = PG_EPOCH_MICROS + 1_500_000; + assert_eq!( + bytes(encode(&json!(micros), &base(oid::TIMESTAMPTZ)).unwrap()), + 1_500_000i64.to_be_bytes() + ); + assert_eq!( + bytes( + encode( + &json!(PG_EPOCH_MICROS + 2 * 86_400_000_000), + &base(oid::DATE) + ) + .unwrap() + ), + 2i32.to_be_bytes() + ); + for (micros, timestamp, date) in [ + (i64::MAX, i64::MAX, i32::MAX), + (i64::MIN, i64::MIN, i32::MIN), + ] { + assert_eq!( + bytes(encode(&json!(micros), &base(oid::TIMESTAMPTZ)).unwrap()), + timestamp.to_be_bytes(), + "an infinite time is Postgres's infinity" + ); + assert_eq!(bytes(encode(&json!(micros), &base(oid::DATE)).unwrap()), date.to_be_bytes()); + } + } + + /// Malformed hex is refused, never a panic: a multi-byte character + /// where a hex digit should be would make slicing by byte index land + /// inside it. + #[test] + fn malformed_hex_is_refused_not_panicked() { + let uuid = base(oid::UUID); + assert!(encode(&json!("é1234567-89ab-cdef-0123-456789abcdef"), &uuid).is_err()); + assert!(encode(&json!("ééééééééééééééééééééééééééééééé"), &uuid).is_err()); + assert!( + encode(&json!("12345678-89ab-cdef-0123-456789abcde"), &uuid).is_err(), + "too short" + ); + let bytea = base(oid::BYTEA); + assert!(encode(&json!("\\xé0"), &bytea).is_err()); + assert!(encode(&json!("\\xzz"), &bytea).is_err()); + assert!(encode(&json!("\\x0"), &bytea).is_err(), "odd length"); + assert_eq!(hex_bytes("00ff"), Some(vec![0, 255])); + } + + /// A `bytea` is decoded from its `\x` hex text. + #[test] + fn bytea_is_decoded_from_hex_text() { + let column = ColumnType { + name: "blob".into(), + oid: oid::BYTEA, + typmod: -1, + kind: TypeKind::Base, + }; + assert!(supports(&column).is_ok()); + let value = encode(&json!("\\x00ff10"), &column).unwrap().unwrap(); + assert_eq!(value.bytes(), &[0x00, 0xff, 0x10]); + assert!(encode(&json!("00ff"), &column).is_err(), "no prefix"); + assert!(encode(&json!("\\x0"), &column).is_err(), "odd digits"); + } + + /// What cannot be encoded is refused before any row is: an interval, + /// a numeric, a range. An array's header is one dimension, a + /// has-nulls flag, the element type, its length and lower bound 1. + #[test] + fn what_is_refused_up_front() { + assert!( + supports(&base(oid::TIMESTAMP)).is_ok(), + "microseconds, Postgres's own precision" + ); + let mut ms = base(oid::TIMESTAMP); + ms.typmod = 3; + assert!(supports(&ms).is_ok()); + assert!(supports(&base(1186)).is_err(), "interval"); + assert!(supports(&base(1700)).is_err(), "numeric"); + assert!( + supports(&base(oid::BYTEA)).is_ok(), + "bytea comes as hex text" + ); + assert!(supports(&base(3904)).is_err(), "an int4 range"); + let text_array = ColumnType { + name: "tags".into(), + oid: 1009, + typmod: -1, + kind: TypeKind::Array { + element: oid::TEXT, + element_kind: Box::new(TypeKind::Base), + }, + }; + assert!(supports(&text_array).is_ok()); + let encoded = bytes(encode(&json!(["a", null, "c,d"]), &text_array).unwrap()); + assert_eq!( + &encoded[..20], + &[0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 25, 0, 0, 0, 3, 0, 0, 0, 1] + ); + assert_eq!( + &encoded[20..], + &[ + 0, 0, 0, 1, b'a', 0xff, 0xff, 0xff, 0xff, 0, 0, 0, 3, b'c', b',', b'd' + ] + ); + } +} diff --git a/crates/diesel-streamgres/src/engine.rs b/crates/diesel-streamgres/src/engine.rs new file mode 100644 index 0000000..5a341d9 --- /dev/null +++ b/crates/diesel-streamgres/src/engine.rs @@ -0,0 +1,2230 @@ +//! The engine side of the cache: a client of the Streamgres S2S endpoint. +//! +//! A process holds a small pool of WebSockets to the engine, each on its +//! own thread with its own runtime, and spreads its subscriptions across +//! them. A read that misses asks [`EngineClient::subscribe`], which picks +//! the connection with the fewest subscriptions, sends the query's SQL +//! and bind bytes and waits for the engine's `accepted` or `refused`; the +//! engine's `touched` frames are turned into [`Cache::touched`] calls. +//! +//! A backend subscribes to whatever its code reads, with no screen to +//! bound the set, so the pool is capped: [`Config::max_subscriptions`] +//! refuses the subscribe past the limit and the read is served live, +//! which is the same thing that happens when the engine is down. Both +//! the pool size and the cap come from the environment +//! ([`Config::from_env`]). +//! +//! Every subscription asks for the rows. When the engine's `accepted` +//! carries them and every projected column has a type the client encodes +//! exactly, the rows are kept as a [`Materialized`] store that the cache +//! entry holds ([`crate::cache::Live`]): the engine's `delta` frames +//! update the store in place, and the next read builds the result from +//! it. Postgres is never read for that query, and the socket's thread +//! only moves rows in and out, never builds a result. When a column +//! cannot be encoded (a `numeric`, an interval, a range) the +//! subscription falls back to touches: the cache reads Postgres and the +//! engine's frames mark the entry stale. +//! +//! The engine is pinged while the socket is quiet; a socket that stays +//! silent past [`Config::silence`] is taken for dead (an engine killed without a +//! close, a network gone without a reset), so its keys go stale rather +//! than being served from rows nobody updates any more. +//! +//! The token travels in the upgrade request. Over `ws://` that is plain +//! text: use `wss://` (TLS, with the system's roots) or a TLS-terminating +//! proxy for anything but a loopback or a private network. +//! +//! The wrapper never needs an async context: a subscribe blocks the +//! calling thread on a channel until the reply lands or the wait runs +//! out. When the engine is unreachable every subscribe is refused at +//! once, so an engine outage costs the application nothing but its cache. +//! +//! # Reconnecting +//! +//! Dropping one socket drops the subscriptions the engine held for it. +//! The worker refuses every subscribe waiting on it, drops the cache +//! entries of its keys (a change may land while it is down, and until it +//! is back those keys are served live), and on reconnect resubscribes +//! every key it still holds under the same ids. A rows-mode key refills +//! from the engine's snapshot with no Postgres read; a refill-mode key +//! refills on its next read. Keys the engine then refuses are forgotten. + +use std::collections::HashMap; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::mpsc::{SyncSender, sync_channel}; +use std::sync::{Arc, Mutex, OnceLock, Weak}; +use std::thread; +use std::time::{Duration, Instant}; + +use base64::Engine as _; +use base64::engine::general_purpose::STANDARD as BASE64; +use futures_util::{SinkExt, StreamExt}; +use serde_json::{Value as Json, json}; +use tokio::sync::mpsc; +use tokio::time::timeout; +use tokio_tungstenite::Connector; +use tokio_tungstenite::tungstenite::Message; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; +use tokio_tungstenite::tungstenite::handshake::client::Request; +use tokio_tungstenite::tungstenite::http::HeaderValue; +use tokio_tungstenite::tungstenite::http::header::AUTHORIZATION; + +use crate::cache::{Cache, Database, Live}; +use crate::key::QueryKey; +use crate::row::OwnedRow; +use crate::materialized::{Materialized, Snapshot}; +use crate::subscriptions::{Outcome, Subscriptions, Verdict}; +use crate::types::TableTypes; + +/// How long a subscribe waits for the engine's verdict before the read is +/// served live instead. +pub const DEFAULT_VERDICT_TIMEOUT: Duration = Duration::from_secs(5); + +/// Sockets to the engine per process when the environment says nothing. +pub const DEFAULT_CONNECTIONS: usize = 2; + +/// Subscriptions a process may hold across its sockets when the +/// environment says nothing. +pub const DEFAULT_MAX_SUBSCRIPTIONS: usize = 10_000; + +/// How long to wait before trying the engine again after a lost socket. +const RECONNECT_DELAY: Duration = Duration::from_secs(1); + +/// How long one frame may take to go out before the socket is given up +/// on: an engine that holds the socket open but stops reading would +/// otherwise block the worker, and its thread, for good. +const SEND_TIMEOUT: Duration = Duration::from_secs(10); + +/// How long a socket may stay silent, pings unanswered, before it is +/// taken for dead and its keys go stale, when the environment says +/// nothing. A quiet socket is pinged three times in that time. +pub const DEFAULT_SILENCE: Duration = Duration::from_secs(10); + +/// How far behind the database the engine may be, by its own heartbeats, +/// before reads stop trusting the cache, when the environment says +/// nothing. +pub const DEFAULT_MAX_LAG: Duration = Duration::from_secs(2); + +/// How far a store's weight may drift from what its cache entry was +/// weighed at, as a share of it (a quarter) and at least this many +/// bytes, before the entry is weighed again. +const REWEIGH_BYTES: usize = 64 * 1024; + +/// The pool's shape. [`Config::from_env`] reads it from the environment: +/// +/// | variable | meaning | default | +/// |-------------------------------------|------------------------------------|---------| +/// | `STREAMGRES_S2S_CONNECTIONS` | sockets to the engine | 2 | +/// | `STREAMGRES_S2S_MAX_SUBSCRIPTIONS` | subscriptions held, all sockets | 10000 | +/// | `STREAMGRES_S2S_VERDICT_TIMEOUT_MS` | wait for `accepted` or `refused` | 5000 | +/// | `STREAMGRES_S2S_SILENCE_MS` | a silent socket is taken for dead | 10000 | +/// | `STREAMGRES_S2S_MAX_LAG_MS` | the engine's lag the cache allows; 0 trusts it | 2000 | +/// | `STREAMGRES_S2S_TOKEN` | bearer token the engine asks for | none | +/// +/// A number that is missing, not a number, or zero leaves the default; +/// an empty token is no token. +#[derive(Clone)] +pub struct Config { + pub connections: usize, + pub max_subscriptions: usize, + pub verdict_timeout: Duration, + /// How long a socket may stay silent before it is taken for dead. + pub silence: Duration, + /// How far behind the database the engine may be before reads go to + /// Postgres; zero trusts the cache whatever the engine says. + pub max_lag: Duration, + /// Sent as `Authorization: Bearer ` when each socket opens, + /// for an engine started with `STREAMGRES_S2S_TOKEN`. + pub token: Option, +} + +impl std::fmt::Debug for Config { + /// The configuration with its token masked. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Config") + .field("connections", &self.connections) + .field("max_subscriptions", &self.max_subscriptions) + .field("verdict_timeout", &self.verdict_timeout) + .field("silence", &self.silence) + .field("max_lag", &self.max_lag) + .field("token", &self.token.as_ref().map(|_| "..")) + .finish() + } +} + +impl Default for Config { + /// The documented defaults, with no token. + fn default() -> Config { + Config { + connections: DEFAULT_CONNECTIONS, + max_subscriptions: DEFAULT_MAX_SUBSCRIPTIONS, + verdict_timeout: DEFAULT_VERDICT_TIMEOUT, + silence: DEFAULT_SILENCE, + max_lag: DEFAULT_MAX_LAG, + token: None, + } + } +} + +impl Config { + /// The pool's shape from the environment (see [`Config`]). + pub fn from_env() -> Config { + let defaults = Config::default(); + Config { + connections: env_number("STREAMGRES_S2S_CONNECTIONS").unwrap_or(defaults.connections), + max_subscriptions: env_number("STREAMGRES_S2S_MAX_SUBSCRIPTIONS") + .unwrap_or(defaults.max_subscriptions), + verdict_timeout: env_number("STREAMGRES_S2S_VERDICT_TIMEOUT_MS") + .map(|millis| Duration::from_millis(millis as u64)) + .unwrap_or(defaults.verdict_timeout), + silence: env_number("STREAMGRES_S2S_SILENCE_MS") + .map(|millis| Duration::from_millis(millis as u64)) + .unwrap_or(defaults.silence), + max_lag: match std::env::var("STREAMGRES_S2S_MAX_LAG_MS").as_deref().map(str::trim) { + Ok("0") => Duration::ZERO, + _ => env_number("STREAMGRES_S2S_MAX_LAG_MS") + .map(|millis| Duration::from_millis(millis as u64)) + .unwrap_or(defaults.max_lag), + }, + token: std::env::var("STREAMGRES_S2S_TOKEN") + .ok() + .filter(|token| !token.is_empty()), + } + } +} + +/// A positive number from the environment. +fn env_number(name: &str) -> Option { + std::env::var(name) + .ok()? + .trim() + .parse::() + .ok() + .filter(|value| *value > 0) +} + +/// What a connection task is asked to send. +enum Outgoing { + Subscribe { + id: u64, + sql: String, + binds: Vec, + }, + Unsubscribe { + id: u64, + }, + /// The pool is gone: close the socket and end the thread. + Stop, +} + +/// What the reader hands a waiting subscribe. +enum Reply { + Verdict(Verdict), + /// Accepted with rows: the caller, which has a connection, builds the + /// store. + Snapshot(Snapshot), +} + +/// What every worker shares: the cache, the types learnt, the fallback +/// log and the one map saying which worker holds which key. +struct Shared { + /// Which worker holds each key. Also the count against the cap. + by_key: Mutex>, + max_subscriptions: usize, + /// The bearer token each socket opens with, if the engine asks for one. + token: Option, + /// How long a socket may stay silent before it is taken for dead. + silence: Duration, + /// How far behind the engine may be before reads go to Postgres. + max_lag: Duration, + /// Column types by table, learnt on the first rows-mode subscribe of + /// a table and reused for every delta and resubscribe after; looked + /// up again when a snapshot names a column they lack, which is how a + /// column a migration added becomes known. + types: Mutex>>, + /// Why the latest subscriptions fell back from rows to refills, most + /// recent last, at most a few. + fallbacks: Mutex>, + /// Why the latest subscriptions were refused, by the engine or by + /// this client, most recent last, at most a few. + refusals: Mutex>, + /// The cache this client feeds, once one is attached. + cache: Mutex>>, + /// Set when the pool is dropped; the worker threads then end. + stopped: AtomicBool, +} + +impl Shared { + /// The attached cache, while it lives. + fn cache(&self) -> Option> { + self.cache.lock().expect("cache").as_ref()?.upgrade() + } + + /// Notes why a subscription fell back from rows to refills. + fn fell_back(&self, reason: String) { + Self::remember(&self.fallbacks, reason); + } + + /// Notes a refusal and returns it as the verdict. + fn refused(&self, reason: String) -> Verdict { + Self::remember(&self.refusals, reason.clone()); + Verdict::Refused(reason) + } + + /// The reasons this client refuses on its own, all of which pass. + const TRANSIENT: [&'static str; 9] = [ + crate::subscriptions::STALE_TYPES, + "the engine has not said which database it streams", + "the engine is not connected", + "the engine client has stopped", + "the engine did not answer in time", + "the engine connection dropped", + "the subscription limit of ", + "the subscription is pending", + "the subscription ended before its rows were built", + ]; + + /// Appends `reason` to `log`, keeping the last eight. + fn remember(log: &Mutex>, reason: String) { + let mut log = log.lock().expect("log"); + if log.len() >= 8 { + log.remove(0); + } + log.push(reason); + } + + /// Drops `key` from the routing map if `worker` is the one holding + /// it. + fn release(&self, key: &QueryKey, worker: usize) { + let mut by_key = self.by_key.lock().expect("by_key"); + if by_key.get(key) == Some(&worker) { + by_key.remove(key); + } + } +} + +/// What one worker holds, under its lock. +struct WorkerState { + /// The id the next subscription is sent under; never reused. + next_id: u64, + /// Each held key's subscription id. + by_key: HashMap, + /// Each subscription id's key. + by_id: HashMap, + /// Subscriptions the engine has not accepted yet: a subscribe + /// waiting for its verdict, with where to deliver it, or `None` for + /// a key resubscribed after a reconnect, which nobody waits for. A + /// read of a pending key goes live: until the engine accepts, a + /// commit can land that no frame reports. + pending: HashMap>>, + /// Subscriptions served from the engine's rows, each behind its own + /// lock so a delta to one is applied without holding up the rest. + materialized: HashMap>, +} + +/// A rows-mode subscription's store, from the moment the engine's +/// snapshot lands. +enum Store { + /// The snapshot is with the caller, which is looking the column + /// types up; the deltas that land meanwhile wait here and are + /// applied once the store is built, so none is lost. + Building(Vec), + Ready(Materialized), +} + +/// A rows-mode subscription's store as the worker and the cache entry +/// share it: the worker applies the engine's deltas, a read builds the +/// result. +struct LiveStore { + id: u64, + key: QueryKey, + store: Mutex, + worker: Weak, + /// What the store weighed when its cache entry was last weighed. + weighed: AtomicUsize, +} + +impl LiveStore { + /// An empty store for subscription `id`, waiting for its snapshot. + fn building(id: u64, key: QueryKey, worker: Weak) -> Arc { + Arc::new(LiveStore { + id, + key, + store: Mutex::new(Store::Building(Vec::new())), + worker, + weighed: AtomicUsize::new(0), + }) + } + + /// Whether the store, now `bytes`, has drifted from what its entry + /// was weighed at by more than [`REWEIGH_BYTES`] allows; if so it + /// counts as weighed at `bytes` from here. + fn drifted(&self, bytes: usize) -> bool { + let weighed = self.weighed.load(Ordering::Relaxed); + if bytes.abs_diff(weighed) > (weighed / 4).max(REWEIGH_BYTES) { + self.weighed.store(bytes, Ordering::Relaxed); + true + } else { + false + } + } +} + +impl std::fmt::Debug for LiveStore { + /// The store by its subscription id; its rows are not printed. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("LiveStore").field("id", &self.id).finish() + } +} + +impl Live for LiveStore { + /// Built under the store's lock, which the socket's thread waits on + /// for at most this one build. A value that cannot be encoded falls + /// the query back to refills, once the lock is let go. + fn rows(&self) -> Option> { + let built = match &mut *self.store.lock().expect("store") { + Store::Ready(materialized) => materialized.rows(), + Store::Building(_) => return None, + }; + match built { + Ok(rows) => Some(rows), + Err(error) => { + if let Some(worker) = self.worker.upgrade() { + worker.fall_back(self.id, &self.key, error.0); + } + None + } + } + } + + /// What the store weighs now; nothing while it is building. + fn bytes(&self) -> usize { + match &*self.store.lock().expect("store") { + Store::Ready(materialized) => materialized.byte_len(), + Store::Building(_) => 0, + } + } +} + +/// What an engine's `hello` said of the database it streams. +#[derive(Debug, Clone, PartialEq)] +enum Greeting { + /// No `hello` on this socket yet. + Waiting, + /// A `hello` that names no database: an engine from before it did, + /// or one that could not read it. Its rows are taken on trust. + Unnamed, + /// The database the engine streams. + Named(Database), +} + +/// One socket to the engine and the subscriptions it carries. +struct Worker { + index: usize, + /// The worker itself, for the stores it hands the cache. + this: Weak, + url: String, + verdict_timeout: Duration, + outgoing: mpsc::UnboundedSender, + connected: AtomicBool, + /// How many keys this worker holds, kept beside the state so the + /// pool can pick the least loaded worker without locking each one. + held: AtomicUsize, + /// The most subscriptions the engine takes on one connection, from + /// its `hello`; 0 when it named no limit. + room: AtomicUsize, + /// Why the latest attempt to connect failed, until one succeeds. + connect_error: Mutex>, + /// What the engine on the socket said it streams, from its `hello`. + greeting: Mutex, + /// The engine's lag as its latest `lag` frame said, and when that + /// came; `None` until one has, and while the engine measures none. + lag: Mutex>, + state: Mutex, + shared: Arc, +} + +/// The workers, and what they share, for the life of the client. +struct Pool { + workers: Vec>, + shared: Arc, + config: Config, +} + +impl Drop for Pool { + /// Stops every worker's thread. + fn drop(&mut self) { + self.shared.stopped.store(true, Ordering::Release); + for worker in &self.workers { + let _ = worker.outgoing.send(Outgoing::Stop); + } + } +} + +/// A client of the engine's S2S endpoint, shared by every connection of a +/// process. +#[derive(Clone)] +pub struct EngineClient { + pool: Arc, +} + +impl EngineClient { + /// Starts the pool for `url` (`ws://host:port//s2s/v1`) with + /// [`Config::from_env`]. Returns at once; the sockets connect in the + /// background and subscribes are refused until one has. + pub fn connect(url: impl Into) -> EngineClient { + Self::with_config(url, Config::from_env()) + } + + /// [`Config::from_env`] with another verdict timeout. + pub fn connect_with_timeout(url: impl Into, verdict_timeout: Duration) -> EngineClient { + Self::with_config( + url, + Config { + verdict_timeout, + ..Config::from_env() + }, + ) + } + + /// A client with `config` (zero connections taken as one). Returns at + /// once, as [`EngineClient::connect`] does. + pub fn with_config(url: impl Into, config: Config) -> EngineClient { + let url = url.into(); + let config = Config { + connections: config.connections.max(1), + ..config + }; + let shared = Arc::new(Shared { + by_key: Mutex::new(HashMap::new()), + max_subscriptions: config.max_subscriptions, + token: config.token.clone(), + silence: config.silence, + max_lag: config.max_lag, + types: Mutex::new(HashMap::new()), + fallbacks: Mutex::new(Vec::new()), + refusals: Mutex::new(Vec::new()), + cache: Mutex::new(None), + stopped: AtomicBool::new(false), + }); + let workers = (0..config.connections) + .map(|index| { + Worker::start( + index, + url.clone(), + config.verdict_timeout, + Arc::clone(&shared), + ) + }) + .collect(); + EngineClient { + pool: Arc::new(Pool { + workers, + shared, + config, + }), + } + } + + /// The cache whose entries this client's `touched` frames drop. Set + /// once the cache exists; it holds the client, so the client holds it + /// weakly. + pub fn attach(&self, cache: &Arc) { + *self.pool.shared.cache.lock().expect("cache") = Some(Arc::downgrade(cache)); + } + + /// The configuration the client runs with. + pub fn config(&self) -> &Config { + &self.pool.config + } + + /// Whether any socket to the engine is up. + pub fn is_connected(&self) -> bool { + self.connections() > 0 + } + + /// How many of the pool's sockets are up. + pub fn connections(&self) -> usize { + self.pool + .workers + .iter() + .filter(|worker| worker.connected.load(Ordering::Acquire)) + .count() + } + + /// How many subscriptions the engine holds for this client, over + /// every socket. + pub fn subscribed(&self) -> usize { + self.pool.shared.by_key.lock().expect("by_key").len() + } + + /// How many subscriptions each socket carries, in pool order. + pub fn per_worker(&self) -> Vec { + self.pool + .workers + .iter() + .map(|worker| worker.held.load(Ordering::Relaxed)) + .collect() + } + + /// What the rows-mode stores weigh, roughly, over every socket. + pub fn store_bytes(&self) -> usize { + self.pool + .workers + .iter() + .map(|worker| { + let stores: Vec> = worker + .state + .lock() + .expect("state") + .materialized + .values() + .cloned() + .collect(); + stores.iter().map(|store| store.bytes()).sum::() + }) + .sum() + } + + /// How many subscriptions are served from the engine's rows rather + /// than refilled from Postgres. + pub fn materialized(&self) -> usize { + self.pool + .workers + .iter() + .map(|worker| worker.state.lock().expect("state").materialized.len()) + .sum() + } + + /// What the application should know about how the client is set + /// up, for its log: a bearer token that would cross the network in + /// plain text, over `ws://` to another host. + pub fn warnings(&self) -> Vec { + let url = self.pool.workers.first().map(|worker| worker.url.as_str()); + match (url, &self.pool.shared.token) { + (Some(url), Some(_)) if token_in_clear(url) => vec![format!( + "STREAMGRES_S2S_TOKEN goes to `{url}` in plain text; use wss:// or a TLS-terminating proxy" + )], + _ => Vec::new(), + } + } + + /// Why recent subscriptions fell back from rows to refills, oldest + /// first; the last few only. + pub fn fallbacks(&self) -> Vec { + self.pool + .shared + .fallbacks + .lock() + .expect("fallbacks") + .clone() + } + + /// Why recent subscriptions were refused, by the engine or by this + /// client (not connected, over the cap, no answer in time), oldest + /// first; the last few only. + pub fn refusals(&self) -> Vec { + self.pool.shared.refusals.lock().expect("refusals").clone() + } + + /// The connected worker carrying the fewest subscriptions among those + /// the engine has room on; else why there is none. + fn least_loaded(&self) -> Result<&Arc, String> { + let workers = &self.pool.workers; + let connected = || { + workers + .iter() + .filter(|worker| worker.connected.load(Ordering::Acquire)) + }; + if let Some(worker) = connected() + .filter(|worker| worker.has_room()) + .min_by_key(|worker| worker.held.load(Ordering::Relaxed)) + { + return Ok(worker); + } + match connected().next() { + Some(full) => Err(format!( + "the subscription limit of {} per connection is reached on every socket", + full.room.load(Ordering::Relaxed) + )), + None => Err(workers + .iter() + .find(|worker| { + worker + .connect_error + .lock() + .expect("connect_error") + .is_some() + }) + .or(workers.first()) + .map_or_else( + || "the engine is not connected".to_string(), + |worker| worker.not_connected(), + )), + } + } + + /// How far behind the database the engine is, by its heartbeats: the + /// most any connected socket says, `None` when none says. + pub fn lag(&self) -> Option { + self.pool.workers.iter().filter_map(|worker| worker.lag()).max() + } + + /// Why each socket that is down failed to connect, for those that + /// have tried and failed. + pub fn connect_errors(&self) -> Vec { + self.pool + .workers + .iter() + .filter(|worker| !worker.connected.load(Ordering::Acquire)) + .filter_map(|worker| worker.connect_error.lock().expect("connect_error").clone()) + .collect() + } +} + +impl Subscriptions for EngineClient { + /// Routes the key to the worker already holding it, else the least + /// loaded one, unless the cap is reached. The key is claimed in the + /// routing map before the worker is asked, so concurrent subscribes + /// count against the cap and land on one worker. + fn subscribe(&self, key: &QueryKey) -> Outcome { + let shared = &self.pool.shared; + let worker = { + let mut by_key = shared.by_key.lock().expect("by_key"); + match by_key.get(key) { + Some(index) => Arc::clone(&self.pool.workers[*index]), + None => { + if by_key.len() >= shared.max_subscriptions { + return Outcome::Verdict(shared.refused(format!( + "the subscription limit of {} is reached", + shared.max_subscriptions + ))); + } + let worker = match self.least_loaded() { + Ok(worker) => worker, + Err(reason) => return Outcome::Verdict(shared.refused(reason)), + }; + by_key.insert(key.clone(), worker.index); + Arc::clone(worker) + } + } + }; + worker.subscribe(key) + } + + /// Learns the looked-up `types` and builds the key's store on the + /// worker holding it. + fn materialize( + &self, + key: &QueryKey, + snapshot: Snapshot, + types: HashMap>, + ) -> Verdict { + self.pool.shared.types.lock().expect("types").extend(types); + let worker = self + .pool + .shared + .by_key + .lock() + .expect("by_key") + .get(key) + .map(|index| Arc::clone(&self.pool.workers[*index])); + match worker { + Some(worker) => worker.materialize_key(key, snapshot), + None => { + Verdict::Refused("the subscription ended before its rows were built".to_string()) + } + } + } + + /// Every refusal but this client's own transient ones. + fn refusal_is_final(&self, reason: &str) -> bool { + !Shared::TRANSIENT + .iter() + .any(|transient| reason.starts_with(transient)) + } + + /// Whether the engine is further behind the database than + /// [`Config::max_lag`] allows, on any socket that says how far. + fn behind(&self) -> bool { + let max = self.pool.shared.max_lag; + !max.is_zero() && self.lag().is_some_and(|lag| lag > max) + } + + /// Releases the key from the routing map and its worker. + fn unsubscribe(&self, key: &QueryKey) { + let worker = self + .pool + .shared + .by_key + .lock() + .expect("by_key") + .remove(key) + .map(|index| Arc::clone(&self.pool.workers[index])); + if let Some(worker) = worker { + worker.unsubscribe(key); + } + } +} + +impl Worker { + /// A worker and the receiving end of its outgoing frames, not yet + /// running. + fn new( + index: usize, + url: String, + verdict_timeout: Duration, + shared: Arc, + ) -> (Arc, mpsc::UnboundedReceiver) { + let (outgoing, outgoing_rx) = mpsc::unbounded_channel(); + let worker = Arc::new_cyclic(|this| Worker { + index, + this: this.clone(), + url, + verdict_timeout, + outgoing, + connected: AtomicBool::new(false), + held: AtomicUsize::new(0), + room: AtomicUsize::new(0), + connect_error: Mutex::new(None), + greeting: Mutex::new(Greeting::Waiting), + lag: Mutex::new(None), + state: Mutex::new(WorkerState { + next_id: 1, + by_key: HashMap::new(), + by_id: HashMap::new(), + pending: HashMap::new(), + materialized: HashMap::new(), + }), + shared, + }); + (worker, outgoing_rx) + } + + /// A worker running its socket on a thread of its own. + fn start( + index: usize, + url: String, + verdict_timeout: Duration, + shared: Arc, + ) -> Arc { + let (worker, outgoing_rx) = Worker::new(index, url, verdict_timeout, shared); + let task = Arc::clone(&worker); + thread::Builder::new() + .name(format!("streamgres-engine-{index}")) + .spawn(move || { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("engine client runtime"); + runtime.block_on(connection_loop(task, outgoing_rx)); + }) + .expect("engine client thread"); + worker + } + + /// Why this socket's engine must not be asked for `key`'s rows: it has + /// not greeted yet, or it streams another database than the cache's + /// reads go to, whose rows would be served as this one's. + fn wrong_database(&self) -> Option { + let greeting = self.greeting.lock().expect("greeting").clone(); + let ours = self.shared.cache(); + let ours = ours.as_ref().and_then(|cache| cache.database()); + match (greeting, ours) { + (Greeting::Waiting, _) => Some("the engine has not said which database it streams".to_string()), + (Greeting::Named(theirs), Some(ours)) if theirs != *ours => Some(format!( + "the engine streams database `{}` of system {}, but reads here go to `{}` of system {}", + theirs.name, theirs.system, ours.name, ours.system + )), + _ => None, + } + } + + /// Subscribes `key` on this worker's socket and waits for the + /// engine's verdict. + /// + /// While the socket is down, a key this worker holds stays routed to + /// it, to be resubscribed when the socket is back; a new one is + /// released for another worker or another time. A held key whose + /// verdict another read is waiting for is refused for now: accepting + /// would fill an entry that a refusal or a timeout then leaves + /// without a subscription, so this read goes live and the next one + /// finds the answer. A key already subscribed (a refill after a + /// touch, or a resubscribe the reader could not materialise without + /// a connection) gets the rows if its store has them. + /// + /// The reader thread puts a `Building` store in place before it + /// hands a snapshot over, so every delta from that moment on queues + /// there while the snapshot's types are looked up. + fn subscribe(&self, key: &QueryKey) -> Outcome { + let (reply, verdict) = sync_channel(1); + let id = { + let mut state = self.state.lock().expect("state"); + let held = state.by_key.get(key).copied(); + if !self.connected.load(Ordering::Acquire) { + drop(state); + if held.is_none() { + self.shared.release(key, self.index); + } + return Outcome::Verdict(self.shared.refused(self.not_connected())); + } + if let Some(reason) = self.wrong_database() { + drop(state); + if held.is_none() { + self.shared.release(key, self.index); + } + return Outcome::Verdict(self.shared.refused(reason)); + } + if let Some(id) = held { + if state.pending.contains_key(&id) { + return Outcome::Verdict( + self.shared + .refused("the subscription is pending".to_string()), + ); + } + let store = state.materialized.get(&id).cloned(); + drop(state); + return Outcome::Verdict(held_rows(store)); + } + let id = state.next_id; + state.next_id += 1; + state.by_key.insert(key.clone(), id); + state.by_id.insert(id, key.clone()); + state.pending.insert(id, Some(reply)); + self.note_held(&state); + id + }; + if self + .outgoing + .send(Outgoing::Subscribe { + id, + sql: key.sql.clone(), + binds: key.binds.iter().map(bind_json).collect(), + }) + .is_err() + { + self.forget(id); + return Outcome::Verdict( + self.shared + .refused("the engine client has stopped".to_string()), + ); + } + match verdict.recv_timeout(self.verdict_timeout) { + Ok(Reply::Verdict(verdict)) => Outcome::Verdict(verdict), + Ok(Reply::Snapshot(snapshot)) => { + let missing: Vec = { + let types = self.shared.types.lock().expect("types"); + snapshot + .tables() + .filter(|table| { + types.get(*table).is_none_or(|known| { + snapshot + .columns_of(table) + .any(|column| known.column(column).is_none()) + }) + }) + .map(str::to_string) + .collect() + }; + if missing.is_empty() { + Outcome::Verdict(self.materialize(id, snapshot)) + } else { + Outcome::NeedsTypes { snapshot, missing } + } + } + Err(_) => { + self.forget(id); + let _ = self.outgoing.send(Outgoing::Unsubscribe { id }); + Outcome::Verdict( + self.shared + .refused("the engine did not answer in time".to_string()), + ) + } + } + } + + /// The second step of a subscribe, once the caller has looked the + /// missing types up and put them in the shared map. + fn materialize_key(&self, key: &QueryKey, snapshot: Snapshot) -> Verdict { + let id = self.state.lock().expect("state").by_key.get(key).copied(); + match id { + Some(id) => self.materialize(id, snapshot), + None => { + Verdict::Refused("the subscription ended before its rows were built".to_string()) + } + } + } + + /// Ends the key's subscription on this worker and with the engine. + fn unsubscribe(&self, key: &QueryKey) { + let id = { + let mut state = self.state.lock().expect("state"); + let Some(id) = state.by_key.get(key).copied() else { + return; + }; + self.drop_id(&mut state, id); + id + }; + let _ = self.outgoing.send(Outgoing::Unsubscribe { id }); + } + + /// Removes a subscription from the worker's maps, under the state + /// lock the caller holds, and fails any rows in flight for its key + /// in the cache: an engine delta or snapshot that found the key held + /// a moment ago must not land after this. + fn drop_id(&self, state: &mut WorkerState, id: u64) -> Option { + let key = state.by_id.remove(&id); + if let Some(key) = &key { + state.by_key.remove(key); + if let Some(cache) = self.shared.cache() { + cache.bump_generation(key); + } + } + state.pending.remove(&id); + state.materialized.remove(&id); + self.note_held(state); + key + } + + /// Builds the store for an accepted subscription from its snapshot + /// and the types already known, applies the deltas that landed + /// while the types were looked up, and installs it. Falls back to + /// touches when a table's types are not known or the rows cannot be + /// encoded. A subscription that ended meanwhile gets nothing. + /// + /// The store is transitioned in place, under the state lock and its + /// own: the one `Arc` lives for the subscription, so a delta handler + /// that took a clone of it before this point finds it `Ready` and + /// applies to the live store, and no delta lands between taking the + /// queue and installing the rows. + fn materialize(&self, id: u64, snapshot: Snapshot) -> Verdict { + let shared = &self.shared; + let fallback = |reason: String| { + self.state.lock().expect("state").materialized.remove(&id); + shared.fell_back(reason); + Verdict::Accepted + }; + let types: HashMap> = { + let known = shared.types.lock().expect("types"); + let mut types = HashMap::new(); + for table in snapshot.tables() { + match known.get(table) { + Some(found) => { + types.insert(table.to_string(), Arc::clone(found)); + } + None => { + drop(known); + return fallback(format!("types of `{table}` unavailable")); + } + } + } + types + }; + let mut materialized = match Materialized::new(snapshot, &types) { + Ok(materialized) => materialized, + Err(error) => return fallback(error.0), + }; + let mut state = self.state.lock().expect("state"); + let Some(key) = state.by_id.get(&id).cloned() else { + state.materialized.remove(&id); + return Verdict::Refused( + "the subscription ended before its rows were built".to_string(), + ); + }; + let live = Arc::clone( + state + .materialized + .entry(id) + .or_insert_with(|| LiveStore::building(id, key, self.this.clone())), + ); + let mut store = live.store.lock().expect("store"); + let queued = match &mut *store { + Store::Building(queued) => std::mem::take(queued), + Store::Ready(_) => Vec::new(), + }; + for body in queued { + apply_delta(&mut materialized, body); + } + match materialized.rows() { + Ok(rows) => { + live.weighed + .store(materialized.byte_len(), Ordering::Relaxed); + *store = Store::Ready(materialized); + drop(store); + Verdict::Rows { rows, live } + } + Err(error) => { + drop(store); + state.materialized.remove(&id); + drop(state); + shared.fell_back(error.0); + Verdict::Accepted + } + } + } + + /// Drops a subscription from this worker and from the routing map. + fn forget(&self, id: u64) { + let key = { + let mut state = self.state.lock().expect("state"); + self.drop_id(&mut state, id) + }; + if let Some(key) = key { + self.shared.release(&key, self.index); + } + } + + /// A rows-mode subscription's rows can no longer be served (a value + /// arrived that cannot be encoded): its store goes, and its key is + /// stale, so the next read refills from Postgres and the engine's + /// frames touch it from then on. + fn fall_back(&self, id: u64, key: &QueryKey, reason: String) { + self.state.lock().expect("state").materialized.remove(&id); + self.shared.fell_back(reason); + if let Some(cache) = self.shared.cache() { + cache.touched(key); + } + } + + /// Publishes how many keys this worker holds, for the pool's choice. + fn note_held(&self, state: &WorkerState) { + self.held.store(state.by_key.len(), Ordering::Relaxed); + } + + /// Whether the engine takes another subscription on this socket: it + /// counts the ones still registering too, as `held` does. + fn has_room(&self) -> bool { + let room = self.room.load(Ordering::Relaxed); + room == 0 || self.held.load(Ordering::Relaxed) < room + } + + /// The refusal while this socket is down, with why it last failed to + /// connect when it has. + fn not_connected(&self) -> String { + match &*self.connect_error.lock().expect("connect_error") { + Some(error) => format!("the engine is not connected: {error}"), + None => "the engine is not connected".to_string(), + } + } + + /// The socket is up: the worker counts as connected, and returns what + /// it still holds, to resubscribe under the same ids. Their cache + /// entries went when the socket did, and each is pending again until + /// the engine accepts it, so nothing stale is served meanwhile. + fn up(&self) -> Vec<(u64, QueryKey)> { + let mut state = self.state.lock().expect("state"); + self.connected.store(true, Ordering::Release); + let held: Vec<(u64, QueryKey)> = state + .by_id + .iter() + .map(|(id, key)| (*id, key.clone())) + .collect(); + for (id, _) in &held { + state.pending.insert(*id, None); + } + held + } + + /// How far behind the engine on this socket is: what it last said, + /// and the time since it said it, so an engine that stops saying + /// anything (frozen, or its frames held up) counts as falling further + /// behind; `None` when it has said nothing measurable. + fn lag(&self) -> Option { + if !self.connected.load(Ordering::Acquire) { + return None; + } + let (lag, heard) = (*self.lag.lock().expect("lag"))?; + Some(lag + heard.elapsed()) + } + + /// Records why the latest connect failed, or that one succeeded. + fn connect_failed(&self, error: Option) { + *self.connect_error.lock().expect("connect_error") = error; + } +} + +/// The verdict for a key already subscribed: its store's rows when it +/// has them, else accepted, for the caller to read Postgres. +fn held_rows(store: Option>) -> Verdict { + match store { + Some(store) => match store.rows() { + Some(rows) => Verdict::Rows { rows, live: store }, + None => Verdict::Accepted, + }, + None => Verdict::Accepted, + } +} + +/// Applies one `delta` body to a store: every table's departures, then +/// its arrivals, the rows moved in rather than copied. The engine +/// coalesces a frame per key, so no key is in both lists and the order +/// between them does not matter. +fn apply_delta(materialized: &mut Materialized, mut body: Json) { + let Some(Json::Array(changes)) = body.get_mut("changes").map(Json::take) else { + return; + }; + for mut change in changes { + let Some(table) = change.get("table").and_then(Json::as_str).map(str::to_owned) else { + continue; + }; + if let Some(Json::Array(keys)) = change.get_mut("del").map(Json::take) { + for key in &keys { + materialized.del(&table, key); + } + } + if let Some(Json::Array(rows)) = change.get_mut("put").map(Json::take) { + for row in rows { + materialized.put(&table, row); + } + } + } +} + +/// The bind as the subscribe frame carries it. +fn bind_json(bind: &crate::key::Bind) -> Json { + json!({ + "oid": bind.type_oid, + "bytes": bind.bytes.as_ref().map(|bytes| BASE64.encode(bytes)), + }) +} + +/// The upgrade request for `url`, with the bearer token when there is +/// one; else why neither can go in a request. The worker then keeps +/// retrying like any other failed connect, every read is served live, +/// and the reason is in each refusal. +fn upgrade_request(url: &str, token: Option<&str>) -> Result { + let mut request = url + .into_client_request() + .map_err(|error| format!("`{url}` is not an engine URL: {error}"))?; + if let Some(token) = token { + let value = HeaderValue::from_str(&format!("Bearer {token}")).map_err(|_| { + "STREAMGRES_S2S_TOKEN holds a character a header cannot carry".to_string() + })?; + request.headers_mut().insert(AUTHORIZATION, value); + } + Ok(request) +} + +/// The TLS a `wss://` socket opens with: rustls with the ring provider +/// and the system's root certificates, built once and given to each +/// socket rather than installed for the whole process, which belongs to +/// the application. `None` when the roots cannot be loaded; a `wss://` +/// connect then fails, and says so in the refusals. A `ws://` socket +/// ignores it. +fn tls() -> Option { + static CONFIG: OnceLock>> = OnceLock::new(); + CONFIG + .get_or_init(|| { + let mut roots = rustls::RootCertStore::empty(); + for certificate in rustls_native_certs::load_native_certs().certs { + let _ = roots.add(certificate); + } + let provider = Arc::new(rustls::crypto::ring::default_provider()); + let config = rustls::ClientConfig::builder_with_provider(provider) + .with_safe_default_protocol_versions() + .ok()? + .with_root_certificates(roots) + .with_no_client_auth(); + Some(Arc::new(config)) + }) + .clone() + .map(Connector::Rustls) +} + +/// Whether a bearer token sent to `url` would cross the network in plain +/// text: a `ws://` URL whose host is not this machine. +fn token_in_clear(url: &str) -> bool { + let Some(rest) = url.strip_prefix("ws://") else { + return false; + }; + let authority = rest.split(['/', '?', '#']).next().unwrap_or(""); + let host = match authority.rsplit_once(':') { + Some((host, port)) if port.bytes().all(|byte| byte.is_ascii_digit()) => host, + _ => authority, + }; + !matches!(host, "localhost" | "127.0.0.1" | "[::1]") +} + +/// Why a connect failed, as a refusal names it. +fn connect_failure(error: &tokio_tungstenite::tungstenite::Error) -> String { + use tokio_tungstenite::tungstenite::Error; + match error { + Error::Http(response) if response.status() == 401 => { + "the engine refused the token (check STREAMGRES_S2S_TOKEN)".to_string() + } + Error::Http(response) if response.status() == 503 => { + "the engine is not serving yet, or has no free connection".to_string() + } + Error::Http(response) => format!("the engine answered {}", response.status()), + other => other.to_string(), + } +} + +/// Connects, serves until the socket breaks either way, reconnects, +/// until the pool stops: one worker's socket. A quiet socket is pinged +/// every third of the silence window, and how long it has been silent +/// is looked at every thirtieth (a second at most), so it is taken for +/// dead within that of the window, not at the next ping. +async fn connection_loop(worker: Arc, mut outgoing: mpsc::UnboundedReceiver) { + let mut first = true; + loop { + if !first { + tokio::time::sleep(RECONNECT_DELAY).await; + } + first = false; + if worker.shared.stopped.load(Ordering::Acquire) { + return; + } + let request = match upgrade_request(&worker.url, worker.shared.token.as_deref()) { + Ok(request) => request, + Err(error) => { + worker.connect_failed(Some(error)); + continue; + } + }; + let socket = match tokio_tungstenite::connect_async_tls_with_config( + request, + None, + false, + tls(), + ) + .await + { + Ok((socket, _)) => { + worker.connect_failed(None); + socket + } + Err(error) => { + worker.connect_failed(Some(connect_failure(&error))); + continue; + } + }; + let (mut sink, mut stream) = socket.split(); + let mut heard = Instant::now(); + let silence = worker.shared.silence; + let every = (silence / 3).max(Duration::from_millis(1)); + let mut ping = tokio::time::interval_at(tokio::time::Instant::now() + every, every); + ping.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + let glance = (silence / 30).clamp(Duration::from_millis(1), Duration::from_secs(1)); + let mut check = tokio::time::interval_at(tokio::time::Instant::now() + glance, glance); + check.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + + let resubscribe = worker.up(); + let mut up = true; + for (id, key) in resubscribe { + let frame = subscribe_frame( + id, + &key.sql, + &key.binds.iter().map(bind_json).collect::>(), + ); + if !send_frame(&mut sink, frame).await { + up = false; + break; + } + } + while up { + up = tokio::select! { + frame = stream.next() => { + heard = Instant::now(); + match frame { + Some(Ok(Message::Text(text))) => { + inbound(&worker, &text); + true + } + Some(Ok(Message::Ping(_))) | Some(Ok(Message::Pong(_))) | Some(Ok(Message::Binary(_))) | Some(Ok(Message::Frame(_))) => true, + Some(Ok(Message::Close(_))) | Some(Err(_)) | None => false, + } + }, + _ = ping.tick() => { + heard.elapsed() < silence + && send_frame(&mut sink, json!(["ping"]).to_string()).await + }, + _ = check.tick() => heard.elapsed() < silence, + message = outgoing.recv() => match message { + Some(Outgoing::Subscribe { id, sql, binds }) => { + send_frame(&mut sink, subscribe_frame(id, &sql, &binds)).await + } + Some(Outgoing::Unsubscribe { id }) => { + send_frame(&mut sink, json!(["unsubscribe", {"id": id}]).to_string()).await + } + Some(Outgoing::Stop) | None => { + let _ = sink.close().await; + return; + } + }, + }; + } + disconnected(&worker); + } +} + +/// Sends one frame, giving the socket up if it does not take it in time. +async fn send_frame(sink: &mut S, frame: String) -> bool +where + S: futures_util::Sink + Unpin, +{ + matches!( + timeout(SEND_TIMEOUT, sink.send(Message::Text(frame.into()))).await, + Ok(Ok(())) + ) +} + +/// The socket is down. Every wait ends refused, and what it waited for +/// is forgotten: the engine never accepted it, so there is nothing to +/// resubscribe, and a frame for it still queued would otherwise +/// register it twice. The stores go, the keys the engine had accepted +/// stay held, pending again, and leave the cache until the engine is +/// back. The other workers' keys are untouched. The cache's memory of +/// refusals goes too, and the column types learnt: the engine coming +/// back may be a restarted one, after a migration that changed a +/// column's type to one of the same width. +fn disconnected(worker: &Worker) { + let (waiting, released, keys) = { + let mut state = worker.state.lock().expect("state"); + worker.connected.store(false, Ordering::Release); + *worker.greeting.lock().expect("greeting") = Greeting::Waiting; + *worker.lag.lock().expect("lag") = None; + state.materialized.clear(); + let mut waiting: Vec> = Vec::new(); + let mut released: Vec = Vec::new(); + for (id, reply) in state.pending.drain().collect::>() { + if let Some(reply) = reply { + waiting.push(reply); + if let Some(key) = worker.drop_id(&mut state, id) { + released.push(key); + } + } + } + let held: Vec = state.by_id.keys().copied().collect(); + for id in held { + state.pending.insert(id, None); + } + let keys: Vec = state.by_key.keys().cloned().collect(); + (waiting, released, keys) + }; + for key in &released { + worker.shared.release(key, worker.index); + } + for reply in &waiting { + let _ = reply.try_send(Reply::Verdict( + worker + .shared + .refused("the engine connection dropped".to_string()), + )); + } + if let Some(cache) = worker.shared.cache() { + for key in &keys { + cache.touched(key); + } + cache.forget_refusals(); + } + worker.shared.types.lock().expect("types").clear(); +} + +/// A subscribe frame asking for the rows. +fn subscribe_frame(id: u64, sql: &str, binds: &[Json]) -> String { + json!(["subscribe", {"id": id, "sql": sql, "binds": binds, "rows": true}]).to_string() +} + +/// One frame from the engine, on one worker's socket. +fn inbound(worker: &Worker, text: &str) { + let Ok(Json::Array(mut array)) = serde_json::from_str::(text) else { + return; + }; + let body = if array.len() > 1 { + array.swap_remove(1) + } else { + Json::Null + }; + match array.first().and_then(Json::as_str).unwrap_or("") { + "accepted" => on_accepted(worker, body), + "hello" => on_hello(worker, &body), + "refused" => on_refused(worker, &body), + "delta" => on_delta(worker, body), + "touched" => on_touched(worker, &body), + "lag" => on_lag(worker, &body), + _ => {} + } +} + +/// `accepted`: the wait ends, and for a rows-mode subscription the +/// `Building` store is in place, under the same lock, before anyone can +/// see the snapshot, so a delta right behind this frame queues there +/// rather than being lost. +/// +/// When nobody holds the id (the wait timed out, and the frame for it +/// was still queued) the engine must not keep following it. A +/// resubscribe after a reconnect has no caller waiting: in rows mode +/// the store is built from the types already known and the cache +/// refilled without a read, the key reserved before the rows are built +/// as a Postgres read reserves it; in touch mode it is accepted, and the +/// next read refills it. +fn on_accepted(worker: &Worker, mut body: Json) { + let shared = &worker.shared; + let Some(id) = body.get("id").and_then(Json::as_u64) else { + return; + }; + let carried_rows = body.get("nodes").is_some(); + let snapshot = Snapshot::from_json(&mut body); + if snapshot.is_none() && carried_rows { + shared.fell_back("the accepted frame was not understood".to_string()); + } + let (waiter, key) = { + let mut state = worker.state.lock().expect("state"); + let waiter = state.pending.remove(&id); + let key = state.by_id.get(&id).cloned(); + if let Some(key) = &key + && snapshot.is_some() + { + state + .materialized + .entry(id) + .or_insert_with(|| LiveStore::building(id, key.clone(), worker.this.clone())); + } + (waiter, key) + }; + let Some(key) = key else { + let _ = worker.outgoing.send(Outgoing::Unsubscribe { id }); + return; + }; + if !matches!(waiter, Some(Some(_))) + && let Some(reason) = worker.wrong_database() + { + resubscribed_elsewhere(worker, id, &key, reason); + return; + } + match (waiter, snapshot) { + (Some(Some(reply)), Some(snapshot)) => { + let _ = reply.try_send(Reply::Snapshot(snapshot)); + } + (Some(Some(reply)), None) => { + let _ = reply.try_send(Reply::Verdict(Verdict::Accepted)); + } + (Some(None) | None, Some(snapshot)) => { + let reservation = shared.cache().map(|cache| cache.reserve(&key)); + if let Verdict::Rows { live, .. } = worker.materialize(id, snapshot) + && let (Some(cache), Some(reservation)) = (shared.cache(), reservation) + { + cache.live(reservation, live); + } + } + (Some(None) | None, None) => {} + } +} + +/// A key resubscribed after a reconnect, accepted by an engine that +/// streams another database than the cache's reads go to: its rows are +/// not this database's, so the key is let go as a refusal would, and its +/// entry goes stale for the next read to take from Postgres. +fn resubscribed_elsewhere(worker: &Worker, id: u64, key: &QueryKey, reason: String) { + let shared = &worker.shared; + { + let mut state = worker.state.lock().expect("state"); + worker.drop_id(&mut state, id); + } + shared.release(key, worker.index); + let _ = shared.refused(reason); + let _ = worker.outgoing.send(Outgoing::Unsubscribe { id }); + if let Some(cache) = shared.cache() { + cache.touched(key); + } +} + +/// `hello`: how many subscriptions the engine takes on this socket, and +/// which database it streams. +fn on_hello(worker: &Worker, body: &Json) { + let room = body + .get("max_subscriptions") + .and_then(Json::as_u64) + .map_or(0, |room| usize::try_from(room).unwrap_or(usize::MAX)); + worker.room.store(room, Ordering::Relaxed); + let field = |name: &str| { + body.get("database") + .and_then(|database| database.get(name)) + .and_then(Json::as_str) + .map(str::to_owned) + }; + *worker.greeting.lock().expect("greeting") = match (field("system"), field("name")) { + (Some(system), Some(name)) => Greeting::Named(Database { system, name }), + _ => Greeting::Unnamed, + }; +} + +/// `refused`: the waiting subscribe gets the reason. A refusal after +/// acceptance (the engine stopped following the query), or of a +/// resubscribe, has nobody waiting, and the cache must not serve the +/// key any more. +fn on_refused(worker: &Worker, body: &Json) { + let shared = &worker.shared; + let Some(id) = body.get("id").and_then(Json::as_u64) else { + return; + }; + let reason = body + .get("reason") + .and_then(Json::as_str) + .unwrap_or("refused by the engine") + .to_string(); + let (reply, key) = { + let mut state = worker.state.lock().expect("state"); + let reply = state.pending.remove(&id); + let key = worker.drop_id(&mut state, id); + (reply, key) + }; + if let Some(key) = &key { + shared.release(key, worker.index); + } + match reply { + Some(Some(reply)) => { + let _ = reply.try_send(Reply::Verdict(shared.refused(reason))); + } + Some(None) | None => { + if let (Some(key), Some(cache)) = (key, shared.cache()) { + cache.touched(&key); + } + } + } +} + +/// `delta`: applied to the subscription's store, or queued while the +/// store is building; a subscription without a store is touched. When +/// the store's weight drifted, its cache entry is weighed again: the key +/// is reserved first, and the store checked to still be the one held, +/// so a subscription ended in between leaves the key alone. +fn on_delta(worker: &Worker, body: Json) { + let shared = &worker.shared; + let Some(id) = body.get("id").and_then(Json::as_u64) else { + return; + }; + let (key, live) = { + let state = worker.state.lock().expect("state"); + let Some(key) = state.by_id.get(&id).cloned() else { + return; + }; + (key, state.materialized.get(&id).cloned()) + }; + let Some(live) = live else { + if let Some(cache) = shared.cache() { + cache.touched(&key); + } + return; + }; + let reweigh = { + let mut store = live.store.lock().expect("store"); + match &mut *store { + Store::Building(queued) => { + queued.push(body); + return; + } + Store::Ready(materialized) => { + apply_delta(materialized, body); + live.drifted(materialized.byte_len()) + } + } + }; + let Some(cache) = shared.cache() else { + return; + }; + cache.changed(); + if reweigh { + let reservation = cache.reserve(&key); + let held = worker + .state + .lock() + .expect("state") + .materialized + .get(&id) + .is_some_and(|held| Arc::ptr_eq(held, &live)); + if held { + cache.live(reservation, live); + } + } +} + +/// `lag`: how far behind the database the engine is, by its heartbeats; +/// `null` while it measures none. +fn on_lag(worker: &Worker, body: &Json) { + let lag = body + .get("ms") + .and_then(Json::as_u64) + .map(|ms| (Duration::from_millis(ms), Instant::now())); + *worker.lag.lock().expect("lag") = lag; +} + +/// `touched`: every named subscription's key goes stale. +fn on_touched(worker: &Worker, body: &Json) { + let Some(ids) = body.get("ids").and_then(Json::as_array) else { + return; + }; + let keys: Vec = { + let state = worker.state.lock().expect("state"); + ids.iter() + .filter_map(Json::as_u64) + .filter_map(|id| state.by_id.get(&id).cloned()) + .collect() + }; + if let Some(cache) = worker.shared.cache() { + for key in &keys { + cache.touched(key); + } + } +} + +#[cfg(test)] +mod tests { + use std::time::Instant; + + use super::*; + use crate::row::OwnedRow; + use crate::types::{ColumnType, TypeKind}; + use diesel::deserialize::FromSqlRow; + use diesel::pg::Pg; + use diesel::sql_types::{Integer, Text}; + + /// A worker with no socket: `connected` is set by hand and frames are + /// fed to `inbound` directly, so the orderings the reader thread and + /// the subscribing thread can take are forced rather than hoped for. + fn worker() -> (Arc, mpsc::UnboundedReceiver, Arc) { + worker_waiting(Duration::from_secs(5)) + } + + /// As [`worker`], giving a verdict `wait` to arrive. + fn worker_waiting( + wait: Duration, + ) -> (Arc, mpsc::UnboundedReceiver, Arc) { + let shared = Arc::new(Shared { + by_key: Mutex::new(HashMap::new()), + max_subscriptions: 100, + token: None, + silence: DEFAULT_SILENCE, + max_lag: DEFAULT_MAX_LAG, + types: Mutex::new(HashMap::new()), + fallbacks: Mutex::new(Vec::new()), + refusals: Mutex::new(Vec::new()), + cache: Mutex::new(None), + stopped: AtomicBool::new(false), + }); + let (worker, outgoing) = Worker::new( + 0, + "ws://127.0.0.1:1/never".to_string(), + wait, + Arc::clone(&shared), + ); + greeted(&worker); + (worker, outgoing, shared) + } + + /// The socket is up and the engine has greeted, naming no database. + fn greeted(worker: &Worker) { + worker.connected.store(true, Ordering::Release); + *worker.greeting.lock().unwrap() = Greeting::Unnamed; + } + + /// The types of table `t`: an `int4` id and a `text` name. + fn t_types() -> Arc { + let base = |name: &str, oid: u32| ColumnType { + name: name.into(), + oid, + typmod: -1, + kind: TypeKind::Base, + }; + Arc::new(TableTypes { + columns: vec![base("id", 23), base("name", 25)], + }) + } + + /// The key every test subscribes. + fn key() -> QueryKey { + QueryKey { + sql: "SELECT id, name FROM t".to_string(), + binds: Vec::new(), + } + } + + /// An `accepted` for `id` carrying row 1 of `t`. + fn accepted_with_rows(id: u64) -> String { + json!(["accepted", { + "id": id, + "nodes": [{"table": "t", "key": ["id"], "parent": null, "rows": [{"id": 1, "name": "a"}]}], + "columns": [["t", "id"], ["t", "name"]], + "names": ["id", "name"], + "count": false, + "order": [["id", "asc"]], + "limit": null, + "offset": 0 + }]) + .to_string() + } + + /// A `delta` for `id` putting row 2 of `t`. + fn delta(id: u64) -> String { + json!(["delta", {"id": id, "changes": [{"table": "t", "put": [{"id": 2, "name": "b"}], "del": []}]}]).to_string() + } + + /// The rows as `(id, name)` pairs. + fn decoded(rows: &Arc<[OwnedRow]>) -> Vec<(i32, String)> { + rows.iter() + .map(|row| { + <(i32, String) as FromSqlRow<(Integer, Text), Pg>>::build_from_row(row).unwrap() + }) + .collect() + } + + /// Starts a subscribe on another thread and returns once the worker + /// has it pending. + fn subscribing(worker: &Arc) -> std::thread::JoinHandle { + let subscriber = Arc::clone(worker); + let handle = std::thread::spawn(move || subscriber.subscribe(&key())); + let started = Instant::now(); + while worker.state.lock().unwrap().pending.is_empty() { + assert!( + started.elapsed() < Duration::from_secs(5), + "the subscribe never became pending" + ); + std::thread::sleep(Duration::from_millis(1)); + } + handle + } + + /// The reader thread handles `accepted` and a delta back to back; + /// whether the subscribing thread has woken up in between must not + /// matter: the store has the delta's row. + #[test] + fn a_delta_right_behind_accepted_is_in_the_store() { + let (worker, _outgoing, shared) = worker(); + shared + .types + .lock() + .unwrap() + .insert("t".to_string(), t_types()); + let handle = subscribing(&worker); + inbound(&worker, &accepted_with_rows(1)); + inbound(&worker, &delta(1)); + let Outcome::Verdict(Verdict::Rows { .. }) = handle.join().unwrap() else { + panic!("rows mode"); + }; + let Outcome::Verdict(Verdict::Rows { rows, .. }) = worker.subscribe(&key()) else { + panic!("held in rows mode"); + }; + assert_eq!( + decoded(&rows), + vec![(1, "a".into()), (2, "b".into())], + "the store has the delta's row" + ); + } + + /// Types learnt before a migration lack the column it added: a + /// snapshot naming that column asks for the table's types again, + /// and the fresh ones serve the rows. + #[test] + fn a_column_the_known_types_lack_asks_for_the_types_again() { + let (worker, _outgoing, shared) = worker(); + let stale = Arc::new(TableTypes { + columns: vec![ColumnType { + name: "id".into(), + oid: 23, + typmod: -1, + kind: TypeKind::Base, + }], + }); + shared.types.lock().unwrap().insert("t".to_string(), stale); + let handle = subscribing(&worker); + inbound(&worker, &accepted_with_rows(1)); + let Outcome::NeedsTypes { snapshot, missing } = handle.join().unwrap() else { + panic!("stale types are looked up again"); + }; + assert_eq!(missing, vec!["t".to_string()]); + shared + .types + .lock() + .unwrap() + .insert("t".to_string(), t_types()); + let Verdict::Rows { .. } = worker.materialize_key(&key(), snapshot) else { + panic!("rows mode with the fresh types"); + }; + } + + /// A delta that lands while the caller is off looking the types up + /// queues and is applied once the store is built; a delta after that + /// reaches the same store. + #[test] + fn deltas_during_the_type_lookup_queue_and_apply_once() { + let (worker, _outgoing, shared) = worker(); + let handle = subscribing(&worker); + inbound(&worker, &accepted_with_rows(1)); + let Outcome::NeedsTypes { snapshot, missing } = handle.join().unwrap() else { + panic!("the types of t are not known yet"); + }; + assert_eq!(missing, vec!["t".to_string()]); + inbound(&worker, &delta(1)); + shared + .types + .lock() + .unwrap() + .insert("t".to_string(), t_types()); + let Verdict::Rows { rows, .. } = worker.materialize_key(&key(), snapshot) else { + panic!( + "rows mode once the types are in: {:?}", + shared.fallbacks.lock().unwrap() + ); + }; + assert_eq!(decoded(&rows), vec![(1, "a".into()), (2, "b".into())]); + inbound(&worker, &json!(["delta", {"id": 1, "changes": [{"table": "t", "put": [], "del": [{"id": 1}]}]}]).to_string()); + let Outcome::Verdict(Verdict::Rows { rows, .. }) = worker.subscribe(&key()) else { + panic!("held in rows mode"); + }; + assert_eq!(decoded(&rows), vec![(2, "b".into())]); + } + + /// A subscribe whose socket drops is refused and forgotten. Its frame + /// was still queued and went out on the next socket, so the engine's + /// acceptance names an id nobody holds: the engine is told to drop + /// it. + #[test] + fn a_wait_cut_by_a_socket_drop_is_forgotten_and_a_late_accept_is_unsubscribed() { + let (worker, mut outgoing, _shared) = worker(); + let handle = subscribing(&worker); + disconnected(&worker); + let Outcome::Verdict(Verdict::Refused(reason)) = handle.join().unwrap() else { + panic!("refused"); + }; + assert_eq!(reason, "the engine connection dropped"); + assert!( + worker.state.lock().unwrap().by_id.is_empty(), + "nothing to resubscribe" + ); + assert_eq!(worker.held.load(Ordering::Relaxed), 0); + greeted(&worker); + inbound(&worker, &accepted_with_rows(1)); + let mut unsubscribed = false; + while let Ok(message) = outgoing.try_recv() { + if matches!(message, Outgoing::Unsubscribe { id: 1 }) { + unsubscribed = true; + } + } + assert!(unsubscribed, "the engine is told to drop it"); + assert!( + worker.state.lock().unwrap().materialized.is_empty(), + "no store for an id nobody holds" + ); + } + + /// When the socket drops and comes back, an accepted key is held, + /// resubscribed, and not to be trusted until the engine says so: a + /// read in the window goes live. + #[test] + fn a_resubscribed_key_is_pending_until_the_engine_accepts_it() { + let (worker, _outgoing, _shared) = worker(); + let handle = subscribing(&worker); + inbound(&worker, &json!(["accepted", {"id": 1}]).to_string()); + assert!(matches!( + handle.join().unwrap(), + Outcome::Verdict(Verdict::Accepted) + )); + disconnected(&worker); + assert!( + worker.state.lock().unwrap().by_id.contains_key(&1), + "an accepted key stays held" + ); + greeted(&worker); + let Outcome::Verdict(Verdict::Refused(reason)) = worker.subscribe(&key()) else { + panic!("a read in the window goes live"); + }; + assert_eq!(reason, "the subscription is pending"); + inbound(&worker, &json!(["accepted", {"id": 1}]).to_string()); + assert!(matches!( + worker.subscribe(&key()), + Outcome::Verdict(Verdict::Accepted) + )); + } + + /// An engine that takes the subscribe and never answers: the wait + /// ends at the verdict timeout, refused for a reason that passes, + /// and the id is withdrawn, so a late `accepted` registers nothing. + #[test] + fn a_subscribe_nobody_answers_is_given_up_in_time() { + let (worker, mut outgoing, shared) = worker_waiting(Duration::from_millis(50)); + let started = Instant::now(); + let Outcome::Verdict(Verdict::Refused(reason)) = worker.subscribe(&key()) else { + panic!("refused"); + }; + assert!(started.elapsed() >= Duration::from_millis(50)); + assert_eq!(reason, "the engine did not answer in time"); + assert!(Shared::TRANSIENT.contains(&reason.as_str())); + assert!(matches!( + outgoing.try_recv(), + Ok(Outgoing::Subscribe { id: 1, .. }) + )); + assert!(matches!( + outgoing.try_recv(), + Ok(Outgoing::Unsubscribe { id: 1 }) + )); + let state = worker.state.lock().unwrap(); + assert!(state.pending.is_empty(), "nothing waits"); + assert!( + state.by_key.is_empty() && state.by_id.is_empty(), + "nothing is held" + ); + assert!( + shared.by_key.lock().unwrap().is_empty(), + "the key is not counted against the cap" + ); + } + + /// The engine that comes back after a dropped socket may know a type + /// it refused before, so its refusals are forgotten with the socket. + #[test] + fn a_dropped_socket_forgets_the_engines_refusals() { + let (worker, _outgoing, shared) = worker(); + let cache = Arc::new(Cache::new( + 1 << 20, + Arc::new(crate::subscriptions::Unsynced), + )); + *shared.cache.lock().unwrap() = Some(Arc::downgrade(&cache)); + cache.refused(&key(), "`$2`: a bind of type oid 16406 is not supported"); + assert!( + cache.refused_recently(&key()).is_some(), + "remembered while the engine is the same" + ); + disconnected(&worker); + assert!( + cache.refused_recently(&key()).is_none(), + "forgotten with the socket" + ); + } + + /// An engine that streams another database than the cache's reads go + /// to is not asked for rows, which would be served as this one's; + /// the right one is, and one that names no database is taken on + /// trust. A key resubscribed after a reconnect to the wrong one is + /// let go and its entry left stale. + #[test] + fn an_engine_on_another_database_is_not_asked() { + let ours = Database { + system: "7699115118699778086".into(), + name: "tickets".into(), + }; + let hello = |system: &str, name: &str| json!({"database": {"system": system, "name": name}}); + let (worker, mut outgoing, shared) = worker_waiting(Duration::from_millis(50)); + let cache = Arc::new(Cache::new(1 << 20, Arc::new(crate::subscriptions::Unsynced))); + *shared.cache.lock().unwrap() = Some(Arc::downgrade(&cache)); + cache.learn_database(Some(ours.clone())); + + on_hello(&worker, &hello("7699114362382213158", "tickets")); + let Outcome::Verdict(Verdict::Refused(reason)) = worker.subscribe(&key()) else { + panic!("refused"); + }; + assert!(reason.contains("streams database `tickets` of system 7699114362382213158"), "{reason}"); + assert!(outgoing.try_recv().is_err(), "the engine was not asked"); + + on_hello(&worker, &json!({})); + assert!(worker.wrong_database().is_none(), "an engine naming no database is trusted"); + on_hello(&worker, &hello(&ours.system, &ours.name)); + assert!(worker.wrong_database().is_none()); + let handle = subscribing(&worker); + inbound(&worker, &json!(["accepted", {"id": 1}]).to_string()); + assert!(matches!(handle.join().unwrap(), Outcome::Verdict(Verdict::Accepted))); + assert!(matches!(outgoing.try_recv(), Ok(Outgoing::Subscribe { id: 1, .. }))); + + disconnected(&worker); + assert!(worker.state.lock().unwrap().by_id.contains_key(&1), "held to resubscribe"); + greeted(&worker); + on_hello(&worker, &hello(&ours.system, "staging")); + inbound(&worker, &json!(["accepted", {"id": 1}]).to_string()); + assert!(matches!(outgoing.try_recv(), Ok(Outgoing::Unsubscribe { id: 1 }))); + assert!(worker.state.lock().unwrap().by_id.is_empty(), "the key is let go"); + } + + /// The engine's lag is what its latest frame said plus the time since, + /// so an engine that stops reporting counts as falling further behind; + /// none is known until a frame says, nor while it says `null`, nor + /// once the socket drops. + #[test] + fn the_lag_grows_while_the_engine_says_nothing() { + let (worker, _outgoing, _shared) = worker(); + assert_eq!(worker.lag(), None); + inbound(&worker, &json!(["lag", {"ms": null}]).to_string()); + assert_eq!(worker.lag(), None); + inbound(&worker, &json!(["lag", {"ms": 1500}]).to_string()); + let lag = worker.lag().expect("said"); + assert!(lag >= Duration::from_millis(1500) && lag < Duration::from_millis(1600), "{lag:?}"); + *worker.lag.lock().unwrap() = Some((Duration::from_millis(1500), Instant::now() - Duration::from_secs(1))); + assert!(worker.lag().unwrap() >= Duration::from_millis(2500), "a second of silence adds a second"); + disconnected(&worker); + assert_eq!(worker.lag(), None, "a dropped socket says nothing"); + } + + /// Reads stop trusting the cache once the engine is further behind + /// than the limit, and trust it again when it catches up; a zero + /// limit always trusts it. + #[test] + fn a_lagging_engine_sends_reads_to_postgres() { + let client = EngineClient::with_config("ws://127.0.0.1:1/never/s2s/v1", Config::default()); + let worker = &client.pool.workers[0]; + worker.connected.store(true, Ordering::Release); + assert!(!client.behind(), "no frame, no lag"); + inbound(worker, &json!(["lag", {"ms": 50}]).to_string()); + assert!(!client.behind()); + inbound(worker, &json!(["lag", {"ms": 2500}]).to_string()); + assert!(client.behind()); + assert!(client.lag().unwrap() >= Duration::from_millis(2500)); + inbound(worker, &json!(["lag", {"ms": 0}]).to_string()); + assert!(!client.behind(), "caught up"); + let trusting = EngineClient::with_config( + "ws://127.0.0.1:1/never/s2s/v1", + Config { + max_lag: Duration::ZERO, + ..Config::default() + }, + ); + trusting.pool.workers[0].connected.store(true, Ordering::Release); + inbound(&trusting.pool.workers[0], &json!(["lag", {"ms": 60000}]).to_string()); + assert!(!trusting.behind()); + } + + /// This client's own refusal reasons all pass; the engine's are final. + #[test] + fn every_client_side_reason_is_transient() { + let client = EngineClient::with_config("ws://127.0.0.1:1/never/s2s/v1", Config::default()); + for reason in Shared::TRANSIENT { + assert!( + !client.refusal_is_final(&format!("{reason} and more")), + "{reason}" + ); + } + assert!(client.refusal_is_final("the query joins more tables than the engine follows")); + } + + /// The defaults are the documented ones, and an unset variable + /// leaves its default. + #[test] + fn config_defaults_and_env() { + let defaults = Config::default(); + assert_eq!(defaults.connections, DEFAULT_CONNECTIONS); + assert_eq!(defaults.max_subscriptions, DEFAULT_MAX_SUBSCRIPTIONS); + assert_eq!(defaults.verdict_timeout, DEFAULT_VERDICT_TIMEOUT); + assert_eq!(env_number("STREAMGRES_S2S_TEST_UNSET_VARIABLE"), None); + } + + /// A token is in plain text over `ws://` to anything but this + /// machine, and never over `wss://`. + #[test] + fn a_token_over_plain_ws_to_another_host_is_named() { + assert!(token_in_clear("ws://engine:4848/zero/s2s/v1")); + assert!(token_in_clear("ws://10.0.0.7/zero/s2s/v1")); + assert!(!token_in_clear("ws://127.0.0.1:4848/zero/s2s/v1")); + assert!(!token_in_clear("ws://localhost/zero/s2s/v1")); + assert!(!token_in_clear("ws://[::1]:4848/zero/s2s/v1")); + assert!(!token_in_clear("wss://engine:4848/zero/s2s/v1")); + assert!(tls().is_some(), "the TLS config builds"); + } + + /// The token goes out as a bearer, and only when there is one; it + /// never shows in the configuration's `Debug`. + #[test] + fn the_token_goes_out_as_a_bearer() { + let url = "ws://127.0.0.1:4848/zero/s2s/v1"; + let plain = upgrade_request(url, None).expect("a request"); + assert!(plain.headers().get(AUTHORIZATION).is_none()); + let with = upgrade_request(url, Some("s3cret")).expect("a request"); + assert_eq!(with.headers()[AUTHORIZATION], "Bearer s3cret"); + let bad = upgrade_request(url, Some("bad\nline")).expect_err("no header"); + assert!(bad.contains("STREAMGRES_S2S_TOKEN"), "{bad}"); + assert!(upgrade_request("not a url", None).is_err()); + let config = Config { + token: Some("s3cret".into()), + ..Config::default() + }; + assert!(!format!("{config:?}").contains("s3cret")); + } + + /// A pool over `workers`, as `with_config` would build it. + fn pool_of(workers: Vec>, shared: Arc) -> EngineClient { + EngineClient { + pool: Arc::new(Pool { + workers, + shared, + config: Config::default(), + }), + } + } + + /// The engine's `hello` names its per-connection limit: a full socket + /// is passed over, and with every socket full the subscribe is + /// refused here, transiently, without asking the engine. + #[test] + fn a_socket_at_the_engine_limit_is_passed_over() { + let (worker, mut outgoing, shared) = worker(); + assert!(worker.has_room(), "no limit until the engine names one"); + inbound( + &worker, + r#"["hello", {"protocol": 1, "max_subscriptions": 1}]"#, + ); + assert_eq!(worker.room.load(Ordering::Relaxed), 1); + let client = pool_of(vec![Arc::clone(&worker)], Arc::clone(&shared)); + assert!(client.least_loaded().is_ok()); + + worker.held.store(1, Ordering::Relaxed); + let reason = client.least_loaded().err().expect("full"); + assert!( + reason.starts_with("the subscription limit of 1 per connection"), + "{reason}" + ); + assert!(!client.refusal_is_final(&reason)); + match client.subscribe(&key()) { + Outcome::Verdict(Verdict::Refused(refused)) => assert_eq!(refused, reason), + _ => panic!("expected a refusal"), + } + assert!(outgoing.try_recv().is_err(), "the engine was not asked"); + + inbound( + &worker, + r#"["hello", {"protocol": 1, "max_subscriptions": null}]"#, + ); + assert!(worker.has_room(), "a later hello with no limit lifts it"); + } + + /// While a socket is down, its refusals say why it last failed to + /// connect, and the client lists it; a success clears it. + #[test] + fn a_failed_connect_is_named_in_the_refusals() { + let (worker, _outgoing, shared) = worker(); + worker.connected.store(false, Ordering::Release); + let client = pool_of(vec![Arc::clone(&worker)], Arc::clone(&shared)); + assert_eq!( + client.least_loaded().err().as_deref(), + Some("the engine is not connected") + ); + + let response = tokio_tungstenite::tungstenite::http::Response::builder() + .status(401) + .body(None) + .unwrap(); + let error = tokio_tungstenite::tungstenite::Error::Http(Box::new(response)); + worker.connect_failed(Some(connect_failure(&error))); + let reason = client.least_loaded().err().expect("down"); + assert_eq!( + reason, + "the engine is not connected: the engine refused the token (check STREAMGRES_S2S_TOKEN)" + ); + assert!(!client.refusal_is_final(&reason)); + assert_eq!(client.connect_errors().len(), 1); + + worker.connect_failed(None); + assert!(client.connect_errors().is_empty()); + } + + /// A pool asked for no sockets opens one. + #[test] + fn zero_connections_become_one() { + let client = EngineClient::with_config( + "ws://127.0.0.1:1/never/s2s/v1", + Config { + connections: 0, + ..Config::default() + }, + ); + assert_eq!(client.config().connections, 1); + assert_eq!(client.per_worker(), vec![0]); + assert!(!client.is_connected()); + } +} diff --git a/crates/diesel-streamgres/src/key.rs b/crates/diesel-streamgres/src/key.rs new file mode 100644 index 0000000..90f5c2e --- /dev/null +++ b/crates/diesel-streamgres/src/key.rs @@ -0,0 +1,170 @@ +//! What identifies a read: its SQL and its bind values. +//! +//! Diesel's own statement cache keys a query by the `TypeId` of its shape, +//! which is the same for `WHERE key = 'a'` and `WHERE key = 'b'`. A result +//! cache needs the values too, so the key is the rendered SQL (with `$n` +//! placeholders) and every bind as Postgres received it: its type OID and +//! its bytes. Two reads with equal keys get equal rows. + +use diesel::QueryResult; +use diesel::pg::{Pg, PgMetadataLookup, PgQueryBuilder}; +use diesel::query_builder::{QueryBuilder, QueryFragment, bind_collector::RawBytesBindCollector}; + +use crate::route::Route; + +/// One bind parameter as the wire carries it. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct Bind { + pub type_oid: u32, + /// `None` is SQL `NULL`. + pub bytes: Option>, +} + +/// The identity of one read. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct QueryKey { + pub sql: String, + pub binds: Vec, +} + +/// Whether `text` starts with `word` as a whole word, case-insensitively. +fn starts_with_word(text: &str, word: &str) -> bool { + text.get(..word.len()) + .is_some_and(|head| head.eq_ignore_ascii_case(word)) + && text[word.len()..] + .chars() + .next() + .is_none_or(|next| !next.is_alphanumeric() && next != '_') +} + +/// Whether `word` occurs in `sql` as a whole word, case-insensitively. +fn has_word(sql: &str, word: &str) -> bool { + sql.split(|c: char| !c.is_alphanumeric() && c != '_') + .any(|token| token.eq_ignore_ascii_case(word)) +} + +impl QueryKey { + /// Renders a query the way the connection is about to send it. The + /// lookup resolves custom types' OIDs, so it is the connection itself. + pub fn of(query: &T, lookup: &mut (dyn PgMetadataLookup + 'static)) -> QueryResult + where + T: QueryFragment, + { + let mut builder = PgQueryBuilder::default(); + query.to_sql(&mut builder, &Pg)?; + let sql = builder.finish(); + + let mut collector = RawBytesBindCollector::::new(); + query.collect_binds(&mut collector, lookup, &Pg)?; + let binds = collector + .metadata + .into_iter() + .zip(collector.binds) + .map(|(metadata, bytes)| Bind { + type_oid: metadata.oid().unwrap_or(0), + bytes, + }) + .collect(); + Ok(QueryKey { sql, binds }) + } + + /// Splits a leading route marker off the SQL, if the query was flagged + /// with [`crate::RouteDsl::via`]. The key, and what the engine sees, + /// is the SQL without it. + pub fn take_route(&mut self) -> Option { + let (route, rest) = Route::strip(&self.sql)?; + self.sql = rest.to_string(); + Some(route) + } + + /// Whether the statement reads rather than writes: a `SELECT`, possibly + /// parenthesised, or a `WITH` whose common table expressions do not + /// modify data. Diesel sends `INSERT ... RETURNING` through the same + /// `load` path as a `SELECT`, and those must never be cached. + pub fn is_select(&self) -> bool { + let head = self + .sql + .trim_start_matches(|c: char| c.is_whitespace() || c == '('); + if starts_with_word(head, "SELECT") { + return true; + } + starts_with_word(head, "WITH") + && !["INSERT", "UPDATE", "DELETE", "MERGE"] + .iter() + .any(|word| has_word(&self.sql, word)) + } + + /// Whether the read takes row locks (`FOR UPDATE`, `FOR SHARE` and + /// their variants). The lock is the point of such a read, and a cache + /// cannot take one, so it always goes to Postgres. + pub fn is_locking(&self) -> bool { + let mut words = self + .sql + .split_whitespace() + .map(|word| word.trim_end_matches(';')); + while let Some(word) = words.next() { + if word.eq_ignore_ascii_case("FOR") + && let Some(next) = words.next() + && (next.eq_ignore_ascii_case("UPDATE") + || next.eq_ignore_ascii_case("SHARE") + || next.eq_ignore_ascii_case("NO") + || next.eq_ignore_ascii_case("KEY")) + { + return true; + } + } + false + } + + /// The bytes the binds hold, for a size-bounded cache. + pub fn byte_len(&self) -> usize { + self.sql.len() + + self + .binds + .iter() + .map(|bind| bind.bytes.as_ref().map_or(0, Vec::len)) + .sum::() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// A key for `sql` with no binds. + fn key(sql: &str) -> QueryKey { + QueryKey { + sql: sql.to_string(), + binds: Vec::new(), + } + } + + /// A `SELECT`, and a `WITH` that only reads, are reads; a `WITH` + /// that writes, and a word that only starts like one, are not. + #[test] + fn selects_and_reading_ctes_are_reads() { + assert!(key("SELECT 1").is_select()); + assert!(key(" (select 1)").is_select()); + assert!(key("WITH recent AS (SELECT * FROM t) SELECT * FROM recent").is_select()); + assert!( + key("with u as (select updated_at from t) select * from u").is_select(), + "a column name is not a statement" + ); + assert!(!key("WITH moved AS (DELETE FROM t RETURNING *) SELECT * FROM moved").is_select()); + assert!(!key("with x as (update t set a = 1 returning *) select 1").is_select()); + assert!(!key("INSERT INTO t VALUES (1) RETURNING id").is_select()); + assert!(!key("SELECTED").is_select()); + assert!(!key("WITHIN").is_select()); + } + + /// Every `FOR` locking clause marks a locking read. + #[test] + fn locking_reads_are_recognised() { + assert!(key("SELECT * FROM t WHERE id = $1 FOR UPDATE").is_locking()); + assert!(key("SELECT * FROM t FOR NO KEY UPDATE").is_locking()); + assert!(key("SELECT * FROM t for share skip locked").is_locking()); + assert!(key("SELECT * FROM t FOR KEY SHARE;").is_locking()); + assert!(!key("SELECT * FROM t ORDER BY id").is_locking()); + assert!(!key("SELECT * FROM forecasts").is_locking()); + } +} diff --git a/crates/diesel-streamgres/src/lib.rs b/crates/diesel-streamgres/src/lib.rs new file mode 100644 index 0000000..8e150d3 --- /dev/null +++ b/crates/diesel-streamgres/src/lib.rs @@ -0,0 +1,622 @@ +//! A Diesel connection that sits between an application and Postgres and +//! serves reads from a local copy the Streamgres engine keeps fresh. +//! +//! `SyncedConnection` wraps any Diesel Postgres connection `C` and +//! implements the same traits, so an application swaps its connection type +//! and nothing else changes: the query DSL, `sql_query`, transactions, pools +//! and `#[derive(Queryable)]` all go through unchanged. +//! +//! Every read Diesel performs reaches [`LoadConnection::load`]; every write +//! reaches [`Connection::execute_returning_count`]; raw multi-statement SQL +//! reaches [`SimpleConnection::batch_execute`]. Those three methods are the +//! whole interception surface, and the wrapper owns all of them. Writes +//! always go to Postgres. A read is answered from the [`cache::Cache`] +//! when all of these hold: +//! +//! - the read's route is [`Route::Streamgres`]: flagged so on the query +//! ([`RouteDsl::via`]), or in a [`SyncedConnection::with_route`] scope, +//! or by the connection's own route, in that order of precedence; +//! - no transaction is open on it (inside one, the application must see +//! its own writes), and the connection has not written for the last +//! [`SyncedConnection::read_your_writes`] window (after one, it must +//! see what it wrote before the engine has shipped it); +//! - the statement is a `SELECT` (Diesel sends `INSERT ... RETURNING` +//! through the same path); +//! - the engine accepted a subscription for the query +//! ([`subscriptions::Subscriptions`]). +//! +//! On a miss the query is subscribed first. When the engine can serve +//! the rows itself (every projected column has a type the client encodes +//! exactly, see [`encode`]), the result comes from the engine, is kept as +//! Diesel rows ([`materialized`]), and every later change replaces it in +//! place; Postgres is never read for that query. Otherwise the rows are +//! read from Postgres once and kept byte for byte ([`row::OwnedRow`]), +//! and when the engine reports the query touched they are dropped and the +//! next read refills them. Every other read goes to Postgres as before. +//! +//! # Why the opt-in feature +//! +//! Diesel seals `Connection` behind the cargo feature +//! `i-implement-a-third-party-backend-and-opt-into-breaking-changes`. +//! Consumers that already wrap their connection enable it, so this crate +//! does too rather than reaching into Diesel's hidden derive re-exports. + +#[cfg(feature = "async")] +pub mod asynchronous; +pub mod cache; +pub mod cursor; +pub mod encode; +pub mod engine; +pub mod key; +pub mod materialized; +pub mod route; +pub mod row; +pub mod subscriptions; +pub mod types; + +use std::sync::atomic::Ordering; +use std::sync::{Arc, OnceLock}; +use std::time::Duration; + +use diesel::connection::{ + AnsiTransactionManager, ConnectionSealed, DefaultLoadingMode, Instrumentation, LoadConnection, + SimpleConnection, TransactionManager, +}; +use diesel::expression::QueryMetadata; +use diesel::pg::{GetPgMetadataCache, Pg, PgMetadataCache}; +use diesel::query_builder::{Query, QueryFragment, QueryId}; +use diesel::r2d2::R2D2Connection; +use diesel::{Connection, ConnectionResult, QueryResult}; + +#[cfg(feature = "async")] +pub use crate::asynchronous::AsyncSyncedConnection; +use crate::cache::{Cache, DATABASE_SQL, Database}; +use crate::cursor::{SyncedCursor, SyncedRow}; +use crate::key::QueryKey; +pub use crate::route::{Route, RouteDsl, Served, Via}; +use crate::row::{OwnedCursor, OwnedRow, capture_all}; +use crate::subscriptions::Verdict; + +/// [`Route`] under another name: `Consistency::Eventual` is +/// [`Route::Streamgres`], `Consistency::Strong` is [`Route::Postgres`]. +pub type Consistency = Route; + +/// What a connection established with some URL gets: the cache its +/// reads go through and the route they take by default. +#[derive(Clone)] +pub struct Attachment { + pub cache: Arc, + pub route: Route, +} + +/// What a process shares between its connections, by the URL they are +/// established with. +/// +/// One engine streams one database, and a cache is kept fresh by one +/// engine, so a cache is per database: pools against the same database +/// (a master and its replicas) share one, pools against different +/// databases each get their own, attached to the engine streaming that +/// database. The route is +/// per pool: a master pool can go to Postgres while a replica pool goes +/// through Streamgres. +pub struct Settings { + for_url: Box, +} + +/// What [`Settings`] asks of a URL. +type AttachmentFor = dyn Fn(&str) -> Option + Send + Sync; + +impl Settings { + /// One database: every connection shares `cache`, routed by its URL. + pub fn single( + cache: Arc, + route_for: impl Fn(&str) -> Route + Send + Sync + 'static, + ) -> Settings { + Settings::by_url(move |url| { + Some(Attachment { + cache: Arc::clone(&cache), + route: route_for(url), + }) + }) + } + + /// A cache and a route per URL. A URL that yields `None` gets a + /// passthrough connection: every read goes to Postgres. + pub fn by_url( + for_url: impl Fn(&str) -> Option + Send + Sync + 'static, + ) -> Settings { + Settings { + for_url: Box::new(for_url), + } + } + + /// What a connection established with `url` gets. + pub fn attachment_for(&self, url: &str) -> Option { + (self.for_url)(url) + } +} + +/// The settings [`install`] put in place, for every connection +/// established after. +static SETTINGS: OnceLock = OnceLock::new(); + +/// How long after a write a connection's reads go to Postgres when the +/// environment says nothing. +pub const DEFAULT_READ_YOUR_WRITES: Duration = Duration::from_millis(1000); + +/// `STREAMGRES_READ_YOUR_WRITES_MS`: how long after a write a +/// connection's reads go to Postgres, so a flow that writes and then +/// reads sees its own write before the engine has shipped it. 0 turns +/// the window off. Read once per process. +pub fn read_your_writes_from_env() -> Duration { + static WINDOW: OnceLock = OnceLock::new(); + *WINDOW.get_or_init(|| { + std::env::var("STREAMGRES_READ_YOUR_WRITES_MS") + .ok() + .and_then(|value| value.trim().parse::().ok()) + .map(Duration::from_millis) + .unwrap_or(DEFAULT_READ_YOUR_WRITES) + }) +} + +/// Installs the process-wide settings. Connections established after this +/// (including by pools, which call [`Connection::establish`]) use them; +/// before it, every connection is a passthrough. Returns the settings back +/// if some were already installed. +pub fn install(settings: Settings) -> Result<(), Settings> { + SETTINGS.set(settings) +} + +/// A Postgres connection whose reads can be served from a synced local copy. +/// +/// `C` is the connection that talks to Postgres, normally +/// [`diesel::PgConnection`]. Everything the wrapper does not understand is +/// delegated to it untouched. +pub struct SyncedConnection { + inner: C, + cache: Option>, + route: Route, + /// A [`SyncedConnection::with_route`] scope in progress. + scoped: Option, + read_your_writes: Duration, + /// Where the last read was served from; `None` before the first. + served: Option, +} + +impl SyncedConnection { + /// Wraps an already established connection as a passthrough. + pub fn new(inner: C) -> Self { + SyncedConnection { + inner, + cache: None, + route: Route::Postgres, + scoped: None, + read_your_writes: Duration::ZERO, + served: None, + } + } + + /// Wraps a connection with a cache and the route its reads take + /// unless a query or a scope says otherwise. + pub fn with_cache(inner: C, cache: Arc, route: Route) -> Self { + SyncedConnection { + inner, + cache: Some(cache), + route, + scoped: None, + read_your_writes: read_your_writes_from_env(), + served: None, + } + } + + /// Where this connection's last read was served from: the cache, + /// the engine's rows, or Postgres. `None` before any read. + pub fn last_served(&self) -> Option { + self.served + } + + /// Changes the route this connection's reads take by default. + pub fn set_route(&mut self, route: Route) { + self.route = route; + } + + /// The route this connection's reads take by default. + pub fn route(&self) -> Route { + self.route + } + + /// [`SyncedConnection::set_route`] under another name. + pub fn set_consistency(&mut self, consistency: Consistency) { + self.set_route(consistency); + } + + /// [`SyncedConnection::route`] under another name. + pub fn consistency(&self) -> Consistency { + self.route() + } + + /// How long after a write reads through this connection's cache go + /// to Postgres. A write is followed, in the same flow, by reads that + /// expect to see it, often on another connection of the same pool; + /// the cache sees the write only once the engine has shipped the + /// commit, so for this window every connection sharing the cache + /// reads the database. The clock is the cache's; the window is this + /// connection's. From `STREAMGRES_READ_YOUR_WRITES_MS` by default. + pub fn set_read_your_writes(&mut self, window: Duration) { + self.read_your_writes = window; + } + + /// How long after a write reads go to Postgres. + pub fn read_your_writes(&self) -> Duration { + self.read_your_writes + } + + /// Marks a write on the cache's clock. A write is marked as it + /// starts and again as it ends, so the read-your-writes window runs + /// from its end however long it took: a read after it on any + /// connection of the cache goes to Postgres for the whole window. + fn wrote(&mut self) { + if let Some(cache) = &self.cache { + cache.mark_write(); + } + } + + /// Whether a write went through the cache within the read-your-writes + /// window. + fn within_write_window(&self) -> bool { + self.cache + .as_ref() + .is_some_and(|cache| cache.wrote_within(self.read_your_writes)) + } + + /// Runs `f` with every read inside it routed `route`, whatever the + /// connection's own route; a flag on the query itself still wins. + /// For code that holds the connection anyway, such as a pool's + /// `interact` closure. The outer route is restored however `f` + /// ends, a panic included. + pub fn with_route(&mut self, route: Route, f: impl FnOnce(&mut Self) -> R) -> R { + /// Puts the outer route back when dropped. + struct Restore<'a, C> { + conn: &'a mut SyncedConnection, + outer: Option, + } + impl Drop for Restore<'_, C> { + /// Restores the outer route. + fn drop(&mut self) { + self.conn.scoped = self.outer; + } + } + let outer = self.scoped.replace(route); + let restore = Restore { conn: self, outer }; + f(&mut *restore.conn) + } + + /// The cache this connection reads through, if any. + pub fn cache(&self) -> Option<&Arc> { + self.cache.as_ref() + } + + /// The connection that talks to Postgres. + pub fn inner(&mut self) -> &mut C { + &mut self.inner + } + + /// Gives the Postgres connection back. + pub fn into_inner(self) -> C { + self.inner + } +} + +impl SimpleConnection for SyncedConnection +where + C: SimpleConnection, +{ + /// Raw SQL may write; so does the ANSI manager's `COMMIT`, which is + /// when a transaction's writes become what the next read expects. + /// Opening, naming a point in, or abandoning a transaction writes + /// nothing. + fn batch_execute(&mut self, query: &str) -> QueryResult<()> { + let writes = may_write(query); + if writes { + self.wrote(); + } + let done = self.inner.batch_execute(query); + if writes { + self.wrote(); + } + done + } +} + +impl ConnectionSealed for SyncedConnection {} + +impl Connection for SyncedConnection +where + C: Connection, +{ + type Backend = Pg; + /// The ANSI manager is generic over the connection it drives: it + /// issues `BEGIN`/`COMMIT`/`ROLLBACK` through `batch_execute` and + /// keeps its depth in whatever `transaction_state` returns, so the + /// wrapper reuses it and points it at the inner connection's state. + /// One transaction state, not two that could disagree. + type TransactionManager = AnsiTransactionManager; + + /// Connects, attached to the cache [`install`] set up for this URL, + /// if any. + fn establish(database_url: &str) -> ConnectionResult { + let inner = C::establish(database_url)?; + Ok( + match SETTINGS + .get() + .and_then(|settings| settings.attachment_for(database_url)) + { + Some(Attachment { cache, route }) => { + SyncedConnection::with_cache(inner, cache, route) + } + None => SyncedConnection::new(inner), + }, + ) + } + + /// Begins Diesel's test transaction on the inner connection. + fn begin_test_transaction(&mut self) -> QueryResult<()> { + self.inner.begin_test_transaction() + } + + /// A pool verifies a connection on checkout with `SELECT 1` through + /// this path; that is not a write. + fn execute_returning_count(&mut self, source: &T) -> QueryResult + where + T: QueryFragment + QueryId, + { + let writes = may_write(&rendered(source)); + if writes { + self.wrote(); + } + let done = self.inner.execute_returning_count(source); + if writes { + self.wrote(); + } + done + } + + /// The inner connection's transaction state. + fn transaction_state(&mut self) -> &mut AnsiTransactionManager { + self.inner.transaction_state() + } + + /// The inner connection's instrumentation. + fn instrumentation(&mut self) -> &mut dyn Instrumentation { + self.inner.instrumentation() + } + + /// Sets the inner connection's instrumentation. + fn set_instrumentation(&mut self, instrumentation: impl Instrumentation) { + self.inner.set_instrumentation(instrumentation) + } +} + +impl LoadConnection for SyncedConnection +where + C: LoadConnection + + LoadConnection + + GetPgMetadataCache + + Connection + + 'static, +{ + type Cursor<'conn, 'query> + = SyncedCursor<>::Cursor<'conn, 'query>> + where + Self: 'conn; + type Row<'conn, 'query> + = SyncedRow<>::Row<'conn, 'query>> + where + Self: 'conn; + + /// A read, from the cache when it can be. + /// + /// The key is rendered before the route is known: a flag on the + /// query travels in its SQL, and is stripped from the key. A locking + /// read goes to Postgres, since the lock is the point: `FOR UPDATE` + /// outside a transaction still blocks a concurrent writer for the + /// statement. On a miss the generation is taken before subscribing, + /// so a touch, or an expiry that ends the subscription, between then + /// and the fill leaves nothing cached rather than something no one + /// updates. Once subscribed, the engine reports every change from + /// there on, so what Postgres returns then is current until the + /// first touch; a read that fails leaves no entry, and the + /// subscription ends with it rather than being held for nothing. The + /// first miss through a cache reads which database the cache's reads + /// go to, which the engine side checks its engine streams. While the + /// engine side is too far behind the database, reads go to Postgres. + fn load<'conn, 'query, T>( + &'conn mut self, + source: T, + ) -> QueryResult> + where + T: Query + QueryFragment + QueryId + 'query, + Pg: QueryMetadata, + { + self.served = Some(Served::Postgres); + let Some(cache) = self.cache.clone() else { + return LoadConnection::::load(&mut self.inner, source).map(SyncedCursor::Live); + }; + let mut key = QueryKey::of(&source, &mut self.inner)?; + let route = key.take_route().or(self.scoped).unwrap_or(self.route); + if !key.is_select() { + let cache = self.cache.clone(); + self.wrote(); + let done = LoadConnection::::load(&mut self.inner, source).map(SyncedCursor::Live); + if let Some(cache) = cache { + cache.mark_write(); + } + return done; + } + if route == Route::Postgres || in_transaction(&mut self.inner) { + return LoadConnection::::load(&mut self.inner, source).map(SyncedCursor::Live); + } + if key.is_locking() { + cache.stats.locking.fetch_add(1, Ordering::Relaxed); + return LoadConnection::::load(&mut self.inner, source).map(SyncedCursor::Live); + } + if self.within_write_window() { + cache.stats.after_write.fetch_add(1, Ordering::Relaxed); + return LoadConnection::::load(&mut self.inner, source).map(SyncedCursor::Live); + } + if cache.subscriptions().behind() { + cache.stats.lagging.fetch_add(1, Ordering::Relaxed); + return LoadConnection::::load(&mut self.inner, source).map(SyncedCursor::Live); + } + if let Some(rows) = cache.get(&key) { + self.served = Some(Served::Cache); + return Ok(SyncedCursor::Cached(OwnedCursor::new(rows))); + } + if cache.refused_recently(&key).is_some() { + return LoadConnection::::load(&mut self.inner, source).map(SyncedCursor::Live); + } + if !cache.knows_database() { + match database_of(&mut self.inner) { + Some(database) => cache.learn_database(Some(database)), + None => { + return LoadConnection::::load(&mut self.inner, source) + .map(SyncedCursor::Live); + } + } + } + let reservation = cache.reserve(&key); + match cache.subscriptions().subscribe_with(&key, &mut self.inner) { + Verdict::Refused(reason) => { + cache.refused(&key, &reason); + return LoadConnection::::load(&mut self.inner, source).map(SyncedCursor::Live); + } + Verdict::Rows { rows, live } => { + cache.fill_live(reservation, live, &rows); + if cache.subscriptions().behind() { + cache.stats.lagging.fetch_add(1, Ordering::Relaxed); + return LoadConnection::::load(&mut self.inner, source) + .map(SyncedCursor::Live); + } + self.served = Some(Served::Engine); + return Ok(SyncedCursor::Cached(OwnedCursor::new(rows))); + } + Verdict::Accepted => {} + } + let rows = LoadConnection::::load(&mut self.inner, source) + .and_then(capture_all); + let rows: Arc<[OwnedRow]> = match rows { + Ok(rows) => Arc::from(rows), + Err(error) => { + cache.abandon(&key); + return Err(error); + } + }; + cache.fill(reservation, Arc::clone(&rows)); + Ok(SyncedCursor::Cached(OwnedCursor::new(rows))) + } +} + +/// The database `conn` reads, as the engine names the one it streams; +/// `None` when it cannot be read, and the read goes live without asking +/// the engine, the next miss reading it again. +fn database_of(conn: &mut C) -> Option +where + C: LoadConnection + Connection, +{ + use diesel::RunQueryDsl; + diesel::sql_query(DATABASE_SQL) + .load::(conn) + .ok()? + .into_iter() + .next() +} + +/// The SQL a statement renders to, for [`may_write`]; a statement that +/// cannot be rendered is taken for a write. +pub(crate) fn rendered>(query: &T) -> String { + let mut builder = diesel::pg::PgQueryBuilder::default(); + match query.to_sql(&mut builder, &Pg) { + Ok(()) => diesel::query_builder::QueryBuilder::finish(builder), + Err(_) => String::new(), + } +} + +/// Whether raw SQL could change what a read sees: any statement of it +/// but a plain `SELECT` (a pool's `SELECT 1` on every checkout would +/// otherwise keep the read-your-writes window open for good) and the +/// statements that only move a transaction along without committing it. +/// A batch is split on every `;`, one inside a string literal included: +/// a piece that then starts oddly counts as a write, which only sends a +/// few more reads to Postgres. +pub(crate) fn may_write(sql: &str) -> bool { + sql.split(';') + .map(str::trim) + .filter(|statement| !statement.is_empty()) + .any(|statement| !reads_only(statement)) +} + +/// Whether one statement is a plain read or only moves a transaction +/// along. +fn reads_only(statement: &str) -> bool { + let head = statement.trim_start_matches('('); + let word = head + .split(|c: char| !c.is_alphanumeric() && c != '_') + .next() + .unwrap_or(""); + [ + "SELECT", + "BEGIN", + "START", + "SAVEPOINT", + "RELEASE", + "ROLLBACK", + "SET", + "RESET", + "SHOW", + ] + .iter() + .any(|statement| word.eq_ignore_ascii_case(statement)) +} + +/// Whether a transaction is open on the connection. Inside one, reads must +/// see the connection's own uncommitted writes, which no cache has. A +/// broken transaction state counts as open: the inner connection reports +/// it on the read. +fn in_transaction(conn: &mut C) -> bool +where + C: Connection, +{ + AnsiTransactionManager::transaction_manager_status_mut(conn) + .transaction_depth() + .map(|depth| depth.is_some()) + .unwrap_or(true) +} + +impl R2D2Connection for SyncedConnection +where + C: R2D2Connection + Connection, +{ + /// Pings the inner connection. + fn ping(&mut self) -> QueryResult<()> { + self.inner.ping() + } + + /// Whether the inner connection is broken. + fn is_broken(&mut self) -> bool { + self.inner.is_broken() + } +} + +/// Lets the wrapper resolve Postgres type OIDs the way the inner connection +/// does, which makes it a `PgMetadataLookup`. Diesel needs that whenever a +/// query binds a custom type, and applications that collect binds themselves +/// hand the connection in for the same reason. +impl GetPgMetadataCache for SyncedConnection +where + C: GetPgMetadataCache, +{ + /// The inner connection's type metadata cache. + fn get_metadata_cache(&mut self) -> &mut PgMetadataCache { + self.inner.get_metadata_cache() + } +} diff --git a/crates/diesel-streamgres/src/materialized.rs b/crates/diesel-streamgres/src/materialized.rs new file mode 100644 index 0000000..a7fc846 --- /dev/null +++ b/crates/diesel-streamgres/src/materialized.rs @@ -0,0 +1,1203 @@ +//! A query's result kept current from the engine's rows, joins composed +//! here. +//! +//! The engine ships a subscription's rows per table: every column of +//! each row, and afterwards the rows that entered or changed and the keys +//! of those that left. [`Materialized`] keeps them by primary key per +//! table and, on demand, produces the result as Diesel rows: the join +//! composed from the root table down each edge (an inner edge drops a +//! parent row with no match, a left edge extends it with NULLs), ordered +//! as the query orders its root table, cut to its limit, projected to its +//! columns, each cell encoded for its Postgres type +//! ([`crate::encode`]). +//! +//! Ordering follows Postgres's defaults: `ASC` puts NULLs last, `DESC` +//! first. Strings compare bytewise, which is Postgres's `C` collation; a +//! column under a linguistic collation may order differently here, so a +//! query ordering by such a column should not be cached in rows mode. +//! Rows that tie on the ordering, and the child rows under one parent, +//! come out in primary-key order by value (9 before 10), which is what +//! `ORDER BY child.id` gives in Postgres; the engine accepts exactly +//! that ordering on a joined table, since the composition yields it. + +use std::cmp::Ordering; +use std::collections::HashMap; +use std::sync::Arc; + +use serde_json::Value as Json; + +use crate::encode::{Unsupported, encode, supports}; +use crate::row::{Columns, OwnedRow, OwnedValue}; +use crate::types::TableTypes; + +/// One table of the query as the engine's `accepted` describes it. +#[derive(Debug, Clone, PartialEq)] +pub struct NodeSnapshot { + pub table: String, + pub key: Vec, + /// The table this one joins to; `None` for the root. + pub parent: Option, + pub parent_column: Option, + pub child_column: Option, + /// `false` for a LEFT join. + pub inner: bool, + pub rows: Vec, +} + +/// What the engine's `accepted` carries in rows mode. +#[derive(Debug, Clone, PartialEq)] +pub struct Snapshot { + /// Root first. + pub nodes: Vec, + /// The projected `(table, column)` pairs, `None` for `SELECT *`. + pub columns: Option>, + /// What each projected column comes out as (its alias or its name); + /// for a count, the one output name. `None` falls back to the + /// columns' own names. + pub names: Option>, + pub count: bool, + /// Root-table column and whether ascending. + pub order: Vec<(String, bool)>, + pub limit: Option, + /// Rows to skip before the limit; the engine hydrates `offset + + /// limit` rows and the client cuts. + pub offset: usize, +} + +impl Snapshot { + /// Reads an `accepted` body, its rows moved out of it; `None` when it + /// carries no rows. + pub fn from_json(body: &mut Json) -> Option { + let strings = |value: Option<&Json>| -> Option> { + value? + .as_array()? + .iter() + .map(|item| item.as_str().map(str::to_owned)) + .collect() + }; + let nodes = body + .get_mut("nodes")? + .as_array_mut()? + .iter_mut() + .map(|node| { + let rows = match node.get_mut("rows").map(Json::take) { + Some(Json::Array(rows)) => rows, + _ => return None, + }; + Some(NodeSnapshot { + table: node.get("table")?.as_str()?.to_owned(), + key: strings(node.get("key"))?, + parent: node.get("parent").and_then(Json::as_str).map(str::to_owned), + parent_column: node + .get("parent_column") + .and_then(Json::as_str) + .map(str::to_owned), + child_column: node + .get("child_column") + .and_then(Json::as_str) + .map(str::to_owned), + inner: node.get("inner").and_then(Json::as_bool).unwrap_or(true), + rows, + }) + }) + .collect::>>()?; + if nodes.is_empty() { + return None; + } + let order = body + .get("order") + .and_then(Json::as_array) + .map(|clauses| { + clauses + .iter() + .filter_map(|clause| { + let pair = clause.as_array()?; + Some(( + pair.first()?.as_str()?.to_owned(), + pair.get(1)?.as_str()? == "asc", + )) + }) + .collect() + }) + .unwrap_or_default(); + let columns = match body.get("columns") { + None | Some(Json::Null) => None, + Some(columns) => Some( + columns + .as_array()? + .iter() + .map(|pair| { + let pair = pair.as_array()?; + Some(( + pair.first()?.as_str()?.to_owned(), + pair.get(1)?.as_str()?.to_owned(), + )) + }) + .collect::>>()?, + ), + }; + Some(Snapshot { + nodes, + columns, + names: match body.get("names") { + None | Some(Json::Null) => None, + other => Some(strings(other)?), + }, + count: body.get("count").and_then(Json::as_bool).unwrap_or(false), + order, + offset: body + .get("offset") + .and_then(Json::as_u64) + .map(|offset| offset as usize) + .unwrap_or(0), + limit: body + .get("limit") + .and_then(Json::as_u64) + .map(|limit| limit as usize), + }) + } + + /// The tables whose column types the store needs. + pub fn tables(&self) -> impl Iterator { + self.nodes.iter().map(|node| node.table.as_str()) + } + + /// Every column of `table` the snapshot names: its key, the columns + /// of its joins, the root's order, and the projection. The column + /// types known for the table must cover them, or they are stale. + pub fn columns_of<'a>(&'a self, table: &'a str) -> impl Iterator + 'a { + let keys = self + .nodes + .iter() + .filter(move |node| node.table == table) + .flat_map(|node| { + node.key + .iter() + .map(String::as_str) + .chain(node.child_column.as_deref()) + }); + let joins = self + .nodes + .iter() + .filter(move |node| node.parent.as_deref() == Some(table)) + .filter_map(|node| node.parent_column.as_deref()); + let order = self + .order + .iter() + .filter(move |_| self.nodes.first().is_some_and(|root| root.table == table)) + .map(|(name, _)| name.as_str()); + let projected = self + .columns + .iter() + .flatten() + .filter(move |(of, _)| of == table) + .map(|(_, name)| name.as_str()); + keys.chain(joins).chain(order).chain(projected) + } +} + +/// What the result is made of. +enum Projection { + /// One row, the count of the joined rows. + Count, + /// `(node, column index in the node's types)` per output column. + Columns(Vec<(usize, usize)>), +} + +/// One table of the query, with the rows the engine shipped for it. +struct Node { + table: String, + types: Arc, + key: Vec, + /// Column index in the parent's types. + parent_column: Option, + /// Column index in this node's types. + child_column: Option, + inner: bool, + children: Vec, + rows: HashMap, +} + +/// One query's rows and how to present them. The result is built from +/// the rows when it is asked for and kept until the next change, so a +/// burst of deltas costs one build, by whichever read comes next. +pub struct Materialized { + nodes: Vec, + projection: Projection, + order: Vec<(usize, bool)>, + limit: Option, + offset: usize, + columns: Columns, + /// What the held rows weigh ([`Materialized::byte_len`]), kept as + /// they come and go. + bytes: usize, + /// The result as last built, until a change. + built: Option>, +} + +impl Materialized { + /// Checks the query can be served from engine rows (every projected + /// column encodable, every named column and table known) and loads + /// the snapshot. `types` must cover every table of the snapshot. + /// Without a projection (`SELECT *`) the result is every column of + /// every table, in `FROM` order. + pub fn new( + snapshot: Snapshot, + types: &HashMap>, + ) -> Result { + let mut nodes: Vec = Vec::with_capacity(snapshot.nodes.len()); + let node_index = |nodes: &[Node], table: &str| -> Result { + nodes + .iter() + .position(|node| node.table == table) + .ok_or_else(|| Unsupported(format!("`{table}` is not a table of the query"))) + }; + for node in &snapshot.nodes { + let node_types = types.get(&node.table).cloned().ok_or_else(|| { + Unsupported(format!("the types of `{}` are not known", node.table)) + })?; + let column = |name: &str| -> Result { + node_types + .column(name) + .map(|(index, _)| index) + .ok_or_else(|| { + Unsupported(format!( + "column `{name}` is not in the types of `{}`", + node.table + )) + }) + }; + let key = node + .key + .iter() + .map(|name| column(name)) + .collect::, _>>()?; + if key.is_empty() { + return Err(Unsupported(format!("`{}` has no primary key", node.table))); + } + let (parent, parent_column, child_column) = + match (&node.parent, &node.parent_column, &node.child_column) { + (Some(parent), Some(parent_column), Some(child_column)) => { + let parent = node_index(&nodes, parent)?; + let parent_column = nodes[parent] + .types + .column(parent_column) + .map(|(index, _)| index) + .ok_or_else(|| { + Unsupported(format!( + "join column `{parent_column}` is not in the types of `{}`", + nodes[parent].table + )) + })?; + ( + Some(parent), + Some(parent_column), + Some(column(child_column)?), + ) + } + (None, _, _) if nodes.is_empty() => (None, None, None), + _ => { + return Err(Unsupported(format!( + "`{}` has no join to a parent table", + node.table + ))); + } + }; + let index = nodes.len(); + if let Some(parent) = parent { + nodes[parent].children.push(index); + } + nodes.push(Node { + table: node.table.clone(), + types: node_types, + key, + parent_column, + child_column, + inner: node.inner, + children: Vec::new(), + rows: HashMap::new(), + }); + } + + let root = &nodes[0]; + let order = snapshot + .order + .iter() + .map(|(name, ascending)| { + root.types + .column(name) + .map(|(index, _)| (index, *ascending)) + .ok_or_else(|| { + Unsupported(format!("order column `{name}` is not in the root's types")) + }) + }) + .collect::, _>>()?; + + let (projection, columns): (Projection, Columns) = if snapshot.count { + let name = snapshot + .names + .as_ref() + .and_then(|names| names.first().cloned()) + .unwrap_or_else(|| "count".to_string()); + (Projection::Count, Arc::from(vec![name])) + } else { + let pairs: Vec<(usize, usize)> = match &snapshot.columns { + Some(columns) => columns + .iter() + .map(|(table, name)| { + let node = node_index(&nodes, table)?; + let column = nodes[node] + .types + .column(name) + .map(|(index, _)| index) + .ok_or_else(|| { + Unsupported(format!( + "column `{name}` is not in the types of `{table}`" + )) + })?; + Ok((node, column)) + }) + .collect::>()?, + None => nodes + .iter() + .enumerate() + .flat_map(|(n, node)| (0..node.types.columns.len()).map(move |c| (n, c))) + .collect(), + }; + for &(node, column) in &pairs { + supports(&nodes[node].types.columns[column])?; + } + let names = match &snapshot.names { + Some(names) if names.len() == pairs.len() => names.clone(), + _ => pairs + .iter() + .map(|&(node, column)| nodes[node].types.columns[column].name.clone()) + .collect::>(), + }; + (Projection::Columns(pairs), Arc::from(names)) + }; + + let mut materialized = Materialized { + nodes, + projection, + order, + limit: snapshot.limit, + offset: snapshot.offset, + columns, + bytes: 0, + built: None, + }; + for node in snapshot.nodes { + for row in node.rows { + materialized.put(&node.table, row); + } + } + Ok(materialized) + } + + /// A row of `table` entered or changed. + pub fn put(&mut self, table: &str, row: Json) { + if let Some(index) = self.nodes.iter().position(|node| node.table == table) { + let key = self.key_of(index, &row); + let (key_bytes, row_bytes) = (key.len(), json_bytes(&row)); + match self.nodes[index].rows.insert(key, row) { + Some(old) => { + self.bytes = (self.bytes + row_bytes).saturating_sub(json_bytes(&old)); + } + None => self.bytes += key_bytes + row_bytes, + } + self.built = None; + } + } + + /// A row of `table` left; `key` is the JSON object of its key columns. + pub fn del(&mut self, table: &str, key: &Json) { + if let Some(index) = self.nodes.iter().position(|node| node.table == table) { + let key = self.key_of(index, key); + if let Some(old) = self.nodes[index].rows.remove(&key) { + self.bytes = self.bytes.saturating_sub(key.len() + json_bytes(&old)); + } + self.built = None; + } + } + + /// What the store weighs, roughly: every row's JSON as it is held. + pub fn byte_len(&self) -> usize { + self.bytes + } + + + /// How many rows are held, over every table. + pub fn len(&self) -> usize { + self.nodes.iter().map(|node| node.rows.len()).sum() + } + + /// Whether no row is held. + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// The row's primary key as one string: each key column's JSON text, + /// each followed by U+001F. Two keys cannot run together into one, + /// because JSON text never holds a raw U+001F: inside a string + /// `serde_json` writes it as `\u001f`, and outside strings it is not + /// a JSON character at all. + fn key_of(&self, node: usize, object: &Json) -> String { + let node = &self.nodes[node]; + let mut key = String::new(); + for &index in &node.key { + let name = &node.types.columns[index].name; + key.push_str(&object.get(name).cloned().unwrap_or(Json::Null).to_string()); + key.push('\u{1f}'); + } + key + } + + /// The result as Diesel rows, built if a change came since it last + /// was. A count's `LIMIT` and `OFFSET` cut the one row it is, as + /// Postgres applies them after counting. + pub fn rows(&mut self) -> Result, Unsupported> { + if let Some(built) = &self.built { + return Ok(Arc::clone(built)); + } + let built = self.build()?; + self.built = Some(Arc::clone(&built)); + Ok(built) + } + + /// The result built from the rows held now. + fn build(&self) -> Result, Unsupported> { + let combos = self.assemble(); + let limit = self.limit.unwrap_or(usize::MAX); + let pairs = match &self.projection { + Projection::Count => { + let count = OwnedValue::new( + std::num::NonZeroU32::new(20).expect("int8"), + (combos.len() as i64).to_be_bytes().to_vec(), + ); + let row = OwnedRow::new(Arc::clone(&self.columns), vec![Some(count)]); + return Ok(std::iter::once(row) + .skip(self.offset) + .take(limit) + .collect::>() + .into()); + } + Projection::Columns(pairs) => pairs, + }; + combos + .into_iter() + .skip(self.offset) + .take(limit) + .map(|combo| { + let values = pairs + .iter() + .map(|&(node, column)| { + let column = &self.nodes[node].types.columns[column]; + match combo[node] { + Some(row) => { + encode(row.get(&column.name).unwrap_or(&Json::Null), column) + } + None => Ok(None), + } + }) + .collect::, _>>()?; + Ok(OwnedRow::new(Arc::clone(&self.columns), values)) + }) + .collect::, _>>() + .map(Arc::from) + } + + /// The joined rows: one entry per result row, holding each table's + /// row or `None` where a left edge found no match. Root rows in the + /// query's order; everything else in key order. + /// + /// Each child node's rows are indexed by the value of its join + /// column, in primary-key order by value: `ORDER BY child.id` in + /// Postgres puts 9 before 10, and so must this. + fn assemble(&self) -> Vec>> { + let mut by_value: Vec>> = Vec::with_capacity(self.nodes.len()); + for node in &self.nodes { + let mut index: HashMap> = HashMap::new(); + if let Some(child_column) = node.child_column { + let name = &node.types.columns[child_column].name; + let mut keyed: Vec<(&String, &Json)> = node.rows.iter().collect(); + keyed.sort_by(|(_, a), (_, b)| self.key_order(node, a, b)); + for (_, row) in keyed { + match row.get(name) { + Some(value) if !value.is_null() => { + index.entry(value.to_string()).or_default().push(row); + } + _ => {} + } + } + } + by_value.push(index); + } + + let root = &self.nodes[0]; + let mut roots: Vec<(&String, &Json)> = root.rows.iter().collect(); + roots.sort_by(|(_, a), (_, b)| { + for &(index, ascending) in &self.order { + let name = &root.types.columns[index].name; + let ordering = compare(a.get(name), b.get(name), ascending); + if ordering != Ordering::Equal { + return ordering; + } + } + self.key_order(root, a, b) + }); + + let mut out = Vec::new(); + for (_, row) in roots { + self.combos(0, row, &by_value, &mut out); + } + out + } + + /// Two rows of `node` by their primary key, column by column, as + /// Postgres orders them: numbers numerically, strings bytewise. + fn key_order(&self, node: &Node, a: &Json, b: &Json) -> Ordering { + for &index in &node.key { + let name = &node.types.columns[index].name; + let ordering = compare(a.get(name), b.get(name), true); + if ordering != Ordering::Equal { + return ordering; + } + } + Ordering::Equal + } + + /// Every result row the subtree under `node` makes with `row` as that + /// node's row, appended to `out`. An inner edge with no match makes + /// none; a left edge with no match leaves its subtree NULL. + fn combos<'a>( + &'a self, + node: usize, + row: &'a Json, + by_value: &[HashMap>], + out: &mut Vec>>, + ) { + let mut partial: Vec>> = vec![vec![None; self.nodes.len()]]; + partial[0][node] = Some(row); + for &child in &self.nodes[node].children { + let child_node = &self.nodes[child]; + let parent_column = child_node + .parent_column + .expect("a child has a parent column"); + let value = row.get(&self.nodes[node].types.columns[parent_column].name); + let matches: &[&Json] = match value { + Some(value) if !value.is_null() => by_value[child] + .get(&value.to_string()) + .map(Vec::as_slice) + .unwrap_or(&[]), + _ => &[], + }; + let mut child_combos: Vec>> = Vec::new(); + for &child_row in matches { + self.combos(child, child_row, by_value, &mut child_combos); + } + if child_combos.is_empty() { + if child_node.inner { + return; + } + continue; + } + let mut next = Vec::with_capacity(partial.len() * child_combos.len()); + for left in &partial { + for right in &child_combos { + let mut merged = left.clone(); + for (slot, value) in merged.iter_mut().zip(right) { + if value.is_some() { + *slot = *value; + } + } + next.push(merged); + } + } + partial = next; + } + out.extend(partial); + } +} + +/// A rough size of a JSON value in memory: what it holds plus what each +/// node costs. +fn json_bytes(value: &Json) -> usize { + const NODE: usize = 32; + match value { + Json::Null | Json::Bool(_) | Json::Number(_) => NODE, + Json::String(s) => NODE + s.len(), + Json::Array(items) => NODE + items.iter().map(json_bytes).sum::(), + Json::Object(fields) => { + NODE + fields + .iter() + .map(|(k, v)| k.len() + NODE + json_bytes(v)) + .sum::() + } + } +} + +/// Postgres's order for one clause: NULLs last ascending, first +/// descending (the reversal for descending applies to them as it does +/// to values); numbers numerically, strings bytewise, booleans false +/// first. +fn compare(a: Option<&Json>, b: Option<&Json>, ascending: bool) -> Ordering { + let a = a.filter(|v| !v.is_null()); + let b = b.filter(|v| !v.is_null()); + match (a, b) { + (None, None) => Ordering::Equal, + (None, Some(_)) => { + if ascending { + Ordering::Greater + } else { + Ordering::Less + } + } + (Some(_), None) => { + if ascending { + Ordering::Less + } else { + Ordering::Greater + } + } + (Some(a), Some(b)) => { + let natural = match (a, b) { + (Json::Number(x), Json::Number(y)) => match (x.as_i64(), y.as_i64()) { + (Some(x), Some(y)) => x.cmp(&y), + _ => x + .as_f64() + .partial_cmp(&y.as_f64()) + .unwrap_or(Ordering::Equal), + }, + (Json::String(x), Json::String(y)) => x.as_bytes().cmp(y.as_bytes()), + (Json::Bool(x), Json::Bool(y)) => x.cmp(y), + (x, y) => x.to_string().cmp(&y.to_string()), + }; + if ascending { + natural + } else { + natural.reverse() + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::types::{ColumnType, TypeKind}; + use diesel::deserialize::FromSqlRow; + use diesel::pg::Pg; + use diesel::sql_types::{BigInt, Integer, Nullable, Text}; + use serde_json::json; + + /// A base-type column. + fn base(name: &str, oid: u32) -> ColumnType { + ColumnType { + name: name.into(), + oid, + typmod: -1, + kind: TypeKind::Base, + } + } + + /// The types of `t`, `books` and `reviews`. + fn types() -> HashMap> { + HashMap::from([ + ( + "t".to_string(), + Arc::new(TableTypes { + columns: vec![base("id", 23), base("name", 25), base("rank", 23)], + }), + ), + ( + "books".to_string(), + Arc::new(TableTypes { + columns: vec![base("id", 23), base("author_id", 23), base("title", 25)], + }), + ), + ( + "reviews".to_string(), + Arc::new(TableTypes { + columns: vec![base("id", 23), base("book_id", 23), base("stars", 23)], + }), + ), + ]) + } + + /// The root node, table `t`, holding `rows`. + fn root(rows: Vec) -> NodeSnapshot { + NodeSnapshot { + table: "t".into(), + key: vec!["id".into()], + parent: None, + parent_column: None, + child_column: None, + inner: true, + rows, + } + } + + /// A node joined under `parent` on `parent_column = child_column`. + fn child( + table: &str, + parent: &str, + parent_column: &str, + child_column: &str, + inner: bool, + rows: Vec, + ) -> NodeSnapshot { + NodeSnapshot { + table: table.into(), + key: vec!["id".into()], + parent: Some(parent.into()), + parent_column: Some(parent_column.into()), + child_column: Some(child_column.into()), + inner, + rows, + } + } + + /// A snapshot of `t` alone, projecting `id` and `name`. + fn snapshot(rows: Vec, order: Vec<(&str, bool)>, limit: Option) -> Snapshot { + Snapshot { + nodes: vec![root(rows)], + columns: Some(vec![("t".into(), "id".into()), ("t".into(), "name".into())]), + names: None, + count: false, + order: order.into_iter().map(|(c, a)| (c.to_string(), a)).collect(), + limit, + offset: 0, + } + } + + /// The rows as `(id, name)` pairs. + fn decoded(rows: &Arc<[OwnedRow]>) -> Vec<(i32, String)> { + rows.iter() + .map(|row| { + <(i32, String) as FromSqlRow<(Integer, Text), Pg>>::build_from_row(row).unwrap() + }) + .collect() + } + + /// `key_of` relies on this: a raw U+001F never appears in JSON + /// text, so it can end each key column without a collision. + #[test] + fn json_text_never_holds_the_key_separator() { + assert_eq!(Json::from("a\u{1f}b").to_string(), r#""a\u001fb""#); + } + + /// An offset skips rows before the limit takes its rows. + #[test] + fn an_offset_skips_rows_before_the_limit() { + let mut snap = snapshot( + vec![ + json!({"id": 1, "key": "a", "value": "1"}), + json!({"id": 2, "key": "b", "value": "2"}), + json!({"id": 3, "key": "c", "value": "3"}), + json!({"id": 4, "key": "d", "value": "4"}), + ], + vec![("id", true)], + Some(2), + ); + snap.offset = 1; + let mut store = Materialized::new(snap, &types()).unwrap(); + let rows = store.rows().unwrap(); + let ids: Vec = rows + .iter() + .map(|row| i32::from_be_bytes(row.value(0).unwrap().bytes().try_into().unwrap())) + .collect(); + assert_eq!(ids, vec![2, 3], "skipped one, took two"); + } + + /// Rows come out in the query's order, NULLs last ascending, cut to + /// the limit and projected. + #[test] + fn rows_come_out_ordered_cut_and_projected() { + let snap = snapshot( + vec![ + json!({"id": 1, "name": "b", "rank": 2}), + json!({"id": 2, "name": "a", "rank": null}), + json!({"id": 3, "name": "c", "rank": 1}), + ], + vec![("rank", true), ("id", true)], + Some(2), + ); + let mut m = Materialized::new(snap, &types()).unwrap(); + assert_eq!( + decoded(&m.rows().unwrap()), + vec![(3, "c".into()), (1, "b".into())], + "NULL rank last, cut to 2" + ); + } + + /// Under `DESC` a NULL comes first, as Postgres puts it. + #[test] + fn descending_puts_nulls_first_as_postgres_does() { + let snap = snapshot( + vec![ + json!({"id": 1, "name": "b", "rank": 2}), + json!({"id": 2, "name": "a", "rank": null}), + json!({"id": 3, "name": "c", "rank": 1}), + ], + vec![("rank", false), ("id", true)], + None, + ); + let mut m = Materialized::new(snap, &types()).unwrap(); + assert_eq!( + decoded(&m.rows().unwrap()), + vec![(2, "a".into()), (1, "b".into()), (3, "c".into())], + "NULL first under DESC, then 2, 1" + ); + } + + /// A second, longer row weighs more. + #[test] + fn the_store_weighs_its_rows() { + let snap = snapshot(vec![json!({"id": 1, "name": "b", "rank": 2})], vec![], None); + let mut m = Materialized::new(snap, &types()).unwrap(); + let one = m.byte_len(); + assert!(one > 0); + m.put( + "t", + json!({"id": 2, "name": "a much longer name than before", "rank": 1}), + ); + assert!( + m.byte_len() > one * 2 - 16, + "a second, longer row weighs more" + ); + } + + /// A joined table's rows come out in key order by value (9, 10, 100), + /// not by text. + #[test] + fn child_rows_come_out_in_numeric_key_order() { + let mut snap = snapshot( + vec![json!({"id": 1, "name": "a", "rank": 1})], + vec![("id", true)], + None, + ); + snap.nodes.push(NodeSnapshot { + table: "books".to_string(), + key: vec!["id".to_string()], + parent: Some("t".to_string()), + parent_column: Some("id".to_string()), + child_column: Some("author_id".to_string()), + inner: true, + rows: vec![ + json!({"id": 10, "author_id": 1, "title": "ten"}), + json!({"id": 9, "author_id": 1, "title": "nine"}), + json!({"id": 100, "author_id": 1, "title": "hundred"}), + ], + }); + snap.columns = Some(vec![ + ("t".to_string(), "id".to_string()), + ("books".to_string(), "title".to_string()), + ]); + snap.names = None; + let mut m = Materialized::new(snap, &types()).unwrap(); + assert_eq!( + decoded(&m.rows().unwrap()), + vec![(1, "nine".into()), (1, "ten".into()), (1, "hundred".into())], + "9, 10, 100: by value, not by text" + ); + } + + /// Puts and deletes change the result. + #[test] + fn deltas_change_the_result() { + let snap = snapshot( + vec![ + json!({"id": 1, "name": "b", "rank": 2}), + json!({"id": 2, "name": "a", "rank": 1}), + ], + vec![("name", false)], + None, + ); + let mut m = Materialized::new(snap, &types()).unwrap(); + assert_eq!( + decoded(&m.rows().unwrap()), + vec![(1, "b".into()), (2, "a".into())] + ); + m.put("t", json!({"id": 1, "name": "z", "rank": 2})); + m.put("t", json!({"id": 3, "name": "m", "rank": 0})); + m.del("t", &json!({"id": 2})); + assert_eq!( + decoded(&m.rows().unwrap()), + vec![(1, "z".into()), (3, "m".into())] + ); + assert_eq!(m.len(), 2); + } + + /// A count is one `int8` row, which counts every row whatever the + /// limit. + #[test] + fn a_count_is_one_int8() { + let mut snap = snapshot(vec![json!({"id": 1}), json!({"id": 2})], vec![], Some(1)); + snap.count = true; + snap.columns = None; + let mut m = Materialized::new(snap, &types()).unwrap(); + let rows = m.rows().unwrap(); + let count = >::build_from_row(&rows[0]).unwrap(); + assert_eq!(count, 2, "a count ignores the limit"); + } + + /// A count's LIMIT and OFFSET cut the one row it is, as Postgres + /// applies them after counting: an OFFSET past it, or a LIMIT of + /// zero, leaves no row. + #[test] + fn a_count_is_cut_by_its_limit_and_offset() { + let counted = |limit: Option, offset: usize| { + let mut snap = snapshot(vec![json!({"id": 1}), json!({"id": 2})], vec![], limit); + snap.count = true; + snap.columns = None; + snap.offset = offset; + Materialized::new(snap, &types()).unwrap().rows().unwrap().len() + }; + assert_eq!(counted(Some(1), 0), 1); + assert_eq!(counted(None, 1), 0); + assert_eq!(counted(Some(0), 0), 0); + } + + /// The result is built once per change: read again unchanged it is + /// the same rows; a put or a delete builds it afresh; the weight + /// follows the rows held. + #[test] + fn the_result_is_built_once_per_change() { + let mut m = Materialized::new( + snapshot(vec![json!({"id": 1, "name": "a"})], vec![("id", true)], None), + &types(), + ) + .unwrap(); + let first = m.rows().unwrap(); + assert!(Arc::ptr_eq(&first, &m.rows().unwrap()), "kept while unchanged"); + let one_row = m.byte_len(); + m.put("t", json!({"id": 2, "name": "b"})); + let second = m.rows().unwrap(); + assert!(!Arc::ptr_eq(&first, &second)); + assert_eq!(decoded(&second), vec![(1, "a".into()), (2, "b".into())]); + assert!(m.byte_len() > one_row); + m.put("t", json!({"id": 2, "name": "bee"})); + m.del("t", &json!({"id": 2})); + assert_eq!(decoded(&m.rows().unwrap()), vec![(1, "a".into())]); + assert_eq!(m.byte_len(), one_row, "a row that came and went weighs nothing"); + } + + /// A projected column the client cannot encode refuses rows mode; an + /// unprojected one does not matter. + #[test] + fn an_unencodable_projected_column_refuses_rows_mode() { + let mut t = types(); + let mut with_created = (*t["t"]).clone(); + with_created.columns.push(base("created", 1186)); + t.insert("t".into(), Arc::new(with_created)); + let mut snap = snapshot(vec![], vec![], None); + snap.columns = Some(vec![ + ("t".into(), "id".into()), + ("t".into(), "created".into()), + ]); + assert!(Materialized::new(snap.clone(), &t).is_err()); + snap.columns = Some(vec![("t".into(), "id".into())]); + assert!( + Materialized::new(snap, &t).is_ok(), + "unprojected columns do not matter" + ); + } + + /// Output columns are named by the statement's aliases. + #[test] + fn output_names_follow_the_aliases() { + use diesel::deserialize::QueryableByName; + /// A row read by its output names. + #[derive(QueryableByName, Debug, PartialEq)] + struct Named { + #[diesel(sql_type = Text)] + n: String, + #[diesel(sql_type = Integer)] + i: i32, + } + let snap = Snapshot { + columns: Some(vec![("t".into(), "name".into()), ("t".into(), "id".into())]), + names: Some(vec!["n".into(), "i".into()]), + ..snapshot(vec![json!({"id": 4, "name": "x", "rank": 1})], vec![], None) + }; + let mut m = Materialized::new(snap, &types()).unwrap(); + let rows = m.rows().unwrap(); + assert_eq!( + Named::build(&rows[0]).unwrap(), + Named { + n: "x".into(), + i: 4 + } + ); + assert_eq!( + rows[0] + .columns() + .iter() + .map(String::as_str) + .collect::>(), + ["n", "i"] + ); + } + + /// A NULL decodes into a nullable column. + #[test] + fn a_nullable_projection_decodes() { + let snap = Snapshot { + columns: Some(vec![("t".into(), "rank".into())]), + ..snapshot( + vec![json!({"id": 1, "name": "x", "rank": null})], + vec![], + None, + ) + }; + let mut m = Materialized::new(snap, &types()).unwrap(); + let rows = m.rows().unwrap(); + let rank = + as FromSqlRow, Pg>>::build_from_row(&rows[0]).unwrap(); + assert_eq!(rank, None); + } + + type Joined = (String, Option); + type JoinedSql = (Text, Nullable); + + /// The rows as `(name, title)` pairs. + fn joined(rows: &Arc<[OwnedRow]>) -> Vec { + rows.iter() + .map(|row| >::build_from_row(row).unwrap()) + .collect() + } + + /// Authors `t` joined to their `books`, inner or left, in id order. + fn authors_books(inner: bool, limit: Option) -> Snapshot { + Snapshot { + nodes: vec![ + root(vec![ + json!({"id": 1, "name": "ann", "rank": 1}), + json!({"id": 2, "name": "bob", "rank": 2}), + json!({"id": 3, "name": "cy", "rank": 3}), + ]), + child( + "books", + "t", + "id", + "author_id", + inner, + vec![ + json!({"id": 10, "author_id": 1, "title": "a1"}), + json!({"id": 11, "author_id": 1, "title": "a2"}), + json!({"id": 12, "author_id": 3, "title": "c1"}), + ], + ), + ], + columns: Some(vec![ + ("t".into(), "name".into()), + ("books".into(), "title".into()), + ]), + names: None, + count: false, + order: vec![("id".into(), true)], + limit, + offset: 0, + } + } + + /// An inner join drops a parent with no match. + #[test] + fn an_inner_join_drops_parents_without_a_match() { + let mut m = Materialized::new(authors_books(true, None), &types()).unwrap(); + assert_eq!( + joined(&m.rows().unwrap()), + vec![ + ("ann".into(), Some("a1".into())), + ("ann".into(), Some("a2".into())), + ("cy".into(), Some("c1".into())), + ] + ); + } + + /// A left join keeps a parent with no match, NULL-extended, and the + /// limit counts joined rows. + #[test] + fn a_left_join_null_extends_and_the_limit_cuts_joined_rows() { + let mut m = Materialized::new(authors_books(false, Some(3)), &types()).unwrap(); + assert_eq!( + joined(&m.rows().unwrap()), + vec![ + ("ann".into(), Some("a1".into())), + ("ann".into(), Some("a2".into())), + ("bob".into(), None), + ], + "bob has no books and the limit counts joined rows" + ); + } + + /// Changes on either side of a join recompose it; a parent that + /// leaves takes its rows with it. + #[test] + fn deltas_on_either_table_recompose_the_join() { + let mut m = Materialized::new(authors_books(true, None), &types()).unwrap(); + m.put("books", json!({"id": 13, "author_id": 2, "title": "b1"})); + m.del("books", &json!({"id": 11})); + m.put("t", json!({"id": 3, "name": "cyrus", "rank": 3})); + assert_eq!( + joined(&m.rows().unwrap()), + vec![ + ("ann".into(), Some("a1".into())), + ("bob".into(), Some("b1".into())), + ("cyrus".into(), Some("c1".into())), + ] + ); + m.del("t", &json!({"id": 1})); + assert_eq!( + joined(&m.rows().unwrap()).len(), + 2, + "the parent left, its books are orphans" + ); + } + + /// A join under a join, and a count of the rows it makes. + #[test] + fn a_nested_join_and_a_count_of_joined_rows() { + let mut snap = authors_books(true, None); + snap.nodes.push(child( + "reviews", + "books", + "id", + "book_id", + true, + vec![ + json!({"id": 100, "book_id": 10, "stars": 5}), + json!({"id": 101, "book_id": 10, "stars": 3}), + json!({"id": 102, "book_id": 12, "stars": 4}), + ], + )); + snap.columns = Some(vec![ + ("t".into(), "name".into()), + ("reviews".into(), "stars".into()), + ]); + let mut m = Materialized::new(snap.clone(), &types()).unwrap(); + let rows: Vec<(String, i32)> = m + .rows() + .unwrap() + .iter() + .map(|row| { + <(String, i32) as FromSqlRow<(Text, Integer), Pg>>::build_from_row(row).unwrap() + }) + .collect(); + assert_eq!( + rows, + vec![("ann".into(), 5), ("ann".into(), 3), ("cy".into(), 4)], + "a2 has no reviews, so it is gone" + ); + + snap.count = true; + snap.columns = None; + let mut m = Materialized::new(snap, &types()).unwrap(); + let rows = m.rows().unwrap(); + assert_eq!( + >::build_from_row(&rows[0]).unwrap(), + 3 + ); + } +} diff --git a/crates/diesel-streamgres/src/route.rs b/crates/diesel-streamgres/src/route.rs new file mode 100644 index 0000000..50ac813 --- /dev/null +++ b/crates/diesel-streamgres/src/route.rs @@ -0,0 +1,195 @@ +//! Where a read goes: through the cache the engine keeps fresh, or +//! straight to Postgres. +//! +//! The route is decided at three levels, the most specific winning: +//! +//! 1. **The query.** [`RouteDsl::via`] wraps a Diesel query in a [`Via`] +//! marker that renders a leading SQL comment, `/*streamgres:postgres*/` +//! or `/*streamgres:streamgres*/`. The wrapper reads the comment at its +//! interception point, strips it, and routes on it; Postgres ignores +//! comments, so a flagged query sent live is the same query. +//! [`crate::SyncedConnection::with_route`] is the same flag for every +//! read inside a closure, for code that holds the connection anyway. +//! 2. **The connection pool.** Every connection carries a route, chosen +//! by the database URL it was established with (see +//! [`crate::Settings`]), so a master pool can go to Postgres and a +//! replica pool through Streamgres. +//! 3. **The process default**, [`Route::default`], which is Streamgres. +//! +//! Three things go to Postgres whatever the flags say: a statement that +//! is not a `SELECT`, a read inside an open transaction, and a query the +//! engine refuses. + +use diesel::QueryResult; +use diesel::pg::Pg; +use diesel::query_builder::{AstPass, Query, QueryFragment, QueryId}; +use diesel::query_dsl::RunQueryDsl; +use diesel::query_dsl::methods::{LimitDsl, OffsetDsl}; + +/// Where a connection's last read was served from, as +/// [`crate::SyncedConnection::last_served`] reports it. Per connection, +/// so it is exact under concurrency where the cache's shared counters +/// are not. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum Served { + /// The cache had the rows: no round trip at all. + Cache, + /// The engine shipped the rows when the query was subscribed: no + /// Postgres read. + Engine, + /// Postgres ran the query (a miss the engine only accepted, a + /// refusal, a write window, a transaction, or the Postgres route). + Postgres, +} + +/// Where a read goes. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)] +pub enum Route { + /// Through the cache: the result may lag a commit by the engine's + /// propagation delay. + #[default] + Streamgres, + /// Straight to Postgres, every time. + Postgres, +} + +impl Route { + /// Another name for [`Route::Streamgres`], under [`crate::Consistency`]. + #[allow(non_upper_case_globals)] + pub const Eventual: Route = Route::Streamgres; + /// Another name for [`Route::Postgres`], under [`crate::Consistency`]. + #[allow(non_upper_case_globals)] + pub const Strong: Route = Route::Postgres; + + const MARKER_STREAMGRES: &'static str = "/*streamgres:streamgres*/ "; + const MARKER_POSTGRES: &'static str = "/*streamgres:postgres*/ "; + + /// The comment a flagged query starts with. + pub(crate) fn marker(self) -> &'static str { + match self { + Route::Streamgres => Self::MARKER_STREAMGRES, + Route::Postgres => Self::MARKER_POSTGRES, + } + } + + /// Splits a leading marker off rendered SQL. + pub(crate) fn strip(sql: &str) -> Option<(Route, &str)> { + if let Some(rest) = sql.strip_prefix(Self::MARKER_POSTGRES) { + return Some((Route::Postgres, rest)); + } + if let Some(rest) = sql.strip_prefix(Self::MARKER_STREAMGRES) { + return Some((Route::Streamgres, rest)); + } + None + } +} + +/// A query flagged with the route it takes. Built by [`RouteDsl::via`]; +/// runs like the query it wraps. +#[derive(Debug, Clone, Copy)] +pub struct Via { + query: Q, + route: Route, +} + +impl Via { + /// The route the query is flagged with. + pub fn route(&self) -> Route { + self.route + } + + /// The query without its flag. + pub fn into_inner(self) -> Q { + self.query + } +} + +/// Flags a query with its route: `query.via(Route::Postgres).first(conn)`. +pub trait RouteDsl: Sized { + /// This read takes `route`. + fn via(self, route: Route) -> Via { + Via { query: self, route } + } + + /// This read goes straight to Postgres. + fn via_postgres(self) -> Via { + self.via(Route::Postgres) + } + + /// This read may come from the cache, whatever its connection says. + fn via_streamgres(self) -> Via { + self.via(Route::Streamgres) + } +} + +impl RouteDsl for T {} + +impl> QueryFragment for Via { + /// The query's SQL behind the route's marker comment. + fn walk_ast<'b>(&'b self, mut out: AstPass<'_, 'b, Pg>) -> QueryResult<()> { + out.push_sql(self.route.marker()); + self.query.walk_ast(out.reborrow()) + } +} + +impl QueryId for Via { + type QueryId = Via; + /// The same query flagged both ways renders two SQL texts, so the + /// prepared statement is keyed by the text, not the type. + const HAS_STATIC_QUERY_ID: bool = false; +} + +impl Query for Via { + type SqlType = Q::SqlType; +} + +impl RunQueryDsl for Via {} + +impl LimitDsl for Via { + type Output = Via; + + /// The query limited, still flagged. + fn limit(self, limit: i64) -> Self::Output { + Via { + query: self.query.limit(limit), + route: self.route, + } + } +} + +impl OffsetDsl for Via { + type Output = Via; + + /// The query offset, still flagged. + fn offset(self, offset: i64) -> Self::Output { + Via { + query: self.query.offset(offset), + route: self.route, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// A marker is stripped back to its route; other comments are not + /// markers. + #[test] + fn markers_round_trip() { + for route in [Route::Streamgres, Route::Postgres] { + let sql = format!("{}SELECT 1", route.marker()); + assert_eq!(Route::strip(&sql), Some((route, "SELECT 1"))); + } + assert_eq!(Route::strip("SELECT 1"), None); + assert_eq!(Route::strip("/* a comment */ SELECT 1"), None); + } + + /// The old route names are aliases of the new ones. + #[test] + fn old_names_still_mean_the_same() { + assert_eq!(Route::Eventual, Route::Streamgres); + assert_eq!(Route::Strong, Route::Postgres); + assert_eq!(Route::default(), Route::Streamgres); + } +} diff --git a/crates/diesel-streamgres/src/row.rs b/crates/diesel-streamgres/src/row.rs new file mode 100644 index 0000000..13b37ae --- /dev/null +++ b/crates/diesel-streamgres/src/row.rs @@ -0,0 +1,258 @@ +//! Rows the cache can hold. +//! +//! A `PgRow` borrows from the libpq result it came from, so it can neither +//! outlive the statement nor cross threads. [`OwnedRow`] keeps what a +//! `PgRow` yields, each value's type OID and raw bytes in the format +//! Postgres sent them in, and implements Diesel's [`Row`] over that, so +//! `#[derive(Queryable)]`, `#[derive(QueryableByName)]` and every +//! `FromSql` decode it exactly as they decode a live row. A row read from +//! Postgres is kept as the bytes Postgres served, nothing re-encoded; a +//! row built from the engine's rows carries the bytes [`crate::encode`] +//! writes for its column's type, which decode the same. +//! +//! The column names are the same for every row of a result set, so they +//! are captured once per set and shared ([`Columns`]). The OID travels +//! with each value: Diesel only ever asks for it through a value, and a +//! NULL has none. + +use std::num::NonZeroU32; +use std::sync::Arc; + +use diesel::QueryResult; +use diesel::pg::{Pg, PgValue}; +use diesel::result::Error; +use diesel::row::{Field, PartialRow, Row, RowIndex, RowSealed}; + +/// The column names of one result set, shared by all its rows. +pub type Columns = Arc<[String]>; + +/// One non-NULL value as Postgres sent it, or as encoded to the same +/// format. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct OwnedValue { + type_oid: NonZeroU32, + bytes: Box<[u8]>, +} + +impl OwnedValue { + /// A value from bytes already in Postgres's binary format for `type_oid`. + pub fn new(type_oid: NonZeroU32, bytes: Vec) -> OwnedValue { + OwnedValue { + type_oid, + bytes: bytes.into_boxed_slice(), + } + } + + /// The value's type. + pub fn type_oid(&self) -> NonZeroU32 { + self.type_oid + } + + /// The value in Postgres's binary format. + pub fn bytes(&self) -> &[u8] { + &self.bytes + } +} + +/// A row that owns its bytes. Cloning shares them: a cached row is handed +/// to every read that hits it. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct OwnedRow { + columns: Columns, + /// One entry per column; `None` is SQL `NULL`. + values: Arc<[Option]>, +} + +impl OwnedRow { + /// A row from values already in Postgres's binary format, one per + /// column of `columns`. + pub fn new(columns: Columns, values: Vec>) -> OwnedRow { + debug_assert_eq!(columns.len(), values.len()); + OwnedRow { + columns, + values: Arc::from(values), + } + } + + /// Copies a live row. The column list is taken from `columns` when one + /// is given, so the rows of one result set share it; otherwise it is + /// read off the row. + /// + /// `None` when a column has no name, which Diesel's row types never + /// produce for a Postgres result. + pub fn capture<'a, R>(row: &R, columns: Option<&Columns>) -> Option + where + R: Row<'a, Pg>, + { + let count = row.field_count(); + let mut values = Vec::with_capacity(count); + let mut fresh = match columns { + Some(_) => None, + None => Some(Vec::with_capacity(count)), + }; + for index in 0..count { + let field = row.get(index)?; + values.push(field.value().map(|value| OwnedValue { + type_oid: value.get_oid(), + bytes: Box::from(value.as_bytes()), + })); + if let Some(fresh) = fresh.as_mut() { + fresh.push(field.field_name()?.to_string()); + } + } + let columns = match (columns, fresh) { + (Some(shared), _) => Arc::clone(shared), + (None, Some(fresh)) => Arc::from(fresh), + (None, None) => unreachable!("one of the two is always set"), + }; + Some(OwnedRow { + columns, + values: Arc::from(values), + }) + } + + /// The columns this row has, shared with the rest of its result set. + pub fn columns(&self) -> &Columns { + &self.columns + } + + /// One column's value, `None` for NULL or out of range. + pub fn value(&self, index: usize) -> Option<&OwnedValue> { + self.values.get(index)?.as_ref() + } + + /// What the row costs, for a size-bounded cache: its values' bytes + /// plus what each value and the row itself cost to hold. A rough + /// figure; the budget is a ceiling, not an accounting. + pub fn byte_len(&self) -> usize { + self.values + .iter() + .map(|value| VALUE_OVERHEAD + value.as_ref().map_or(0, |value| value.bytes.len())) + .sum::() + + ROW_OVERHEAD + } +} + +/// A boxed value's pointer, length, type and the `Option` around it. +const VALUE_OVERHEAD: usize = 32; +/// The row's two `Arc`s and the vector of values. +const ROW_OVERHEAD: usize = 64; + +impl RowSealed for OwnedRow {} + +impl RowIndex for OwnedRow { + /// The position, if the row has that many columns. + fn idx(&self, idx: usize) -> Option { + (idx < self.values.len()).then_some(idx) + } +} + +impl<'a> RowIndex<&'a str> for OwnedRow { + /// The position of the column named `name`. + fn idx(&self, name: &'a str) -> Option { + self.columns.iter().position(|column| column == name) + } +} + +impl<'a> Row<'a, Pg> for OwnedRow { + type Field<'f> + = OwnedField<'f> + where + 'a: 'f, + Self: 'f; + type InnerPartialRow = Self; + + /// How many columns the row has. + fn field_count(&self) -> usize { + self.values.len() + } + + /// One column, by position or name. + fn get<'b, I>(&'b self, idx: I) -> Option> + where + 'a: 'b, + Self: RowIndex, + { + let index = self.idx(idx)?; + Some(OwnedField { row: self, index }) + } + + /// The columns in `range`, for a nested `Queryable`. + fn partial_row(&self, range: std::ops::Range) -> PartialRow<'_, Self::InnerPartialRow> { + PartialRow::new(self, range) + } +} + +/// One column of an [`OwnedRow`]. +pub struct OwnedField<'a> { + row: &'a OwnedRow, + index: usize, +} + +impl<'a> Field<'a, Pg> for OwnedField<'a> { + /// The column's name. + fn field_name(&self) -> Option<&str> { + Some(&self.row.columns[self.index]) + } + + /// The column's value; `None` for `NULL`. + fn value(&self) -> Option> { + let value = self.row.value(self.index)?; + Some(PgValue::new(&value.bytes, &value.type_oid)) + } +} + +/// Copies every row a live cursor yields. Stops at the first error and +/// returns it, so a half-read result is never kept. +pub fn capture_all<'a, I, R>(cursor: I) -> QueryResult> +where + I: Iterator>, + R: Row<'a, Pg>, +{ + let mut rows = Vec::new(); + let mut columns: Option = None; + for row in cursor { + let row = row?; + let owned = OwnedRow::capture(&row, columns.as_ref()).ok_or( + Error::DeserializationError("a result row without column names cannot be kept".into()), + )?; + if columns.is_none() { + columns = Some(Arc::clone(owned.columns())); + } + rows.push(owned); + } + Ok(rows) +} + +/// Yields owned rows the way a live cursor yields live ones, so Diesel's +/// `load` machinery can decode a cached result. +pub struct OwnedCursor { + rows: Arc<[OwnedRow]>, + next: usize, +} + +impl OwnedCursor { + /// A cursor over `rows`, from the first. + pub fn new(rows: Arc<[OwnedRow]>) -> Self { + OwnedCursor { rows, next: 0 } + } +} + +impl Iterator for OwnedCursor { + type Item = QueryResult; + + /// The next row, shared with the cache. + fn next(&mut self) -> Option { + let row = self.rows.get(self.next)?.clone(); + self.next += 1; + Some(Ok(row)) + } + + /// Exactly the rows left. + fn size_hint(&self) -> (usize, Option) { + let left = self.rows.len() - self.next; + (left, Some(left)) + } +} + +impl ExactSizeIterator for OwnedCursor {} diff --git a/crates/diesel-streamgres/src/subscriptions.rs b/crates/diesel-streamgres/src/subscriptions.rs new file mode 100644 index 0000000..8b0c05d --- /dev/null +++ b/crates/diesel-streamgres/src/subscriptions.rs @@ -0,0 +1,150 @@ +//! The engine side of a cached read, as the cache sees it. +//! +//! The cache does not talk to the engine. It asks a [`Subscriptions`] to +//! register a query before it fills an entry for it, and tells it when the +//! entry is gone. Whoever owns the engine connection implements the trait +//! and, when the engine reports a query touched, calls +//! [`Cache::touched`](crate::cache::Cache::touched). +//! +//! The order matters for correctness: the subscription is registered +//! before Postgres is read, so a change that lands after the read is +//! reported, and one that lands between the two is reported redundantly, +//! which costs one refill and never a stale row. +//! +//! Subscribing is two steps because the second may need the database. +//! [`Subscriptions::subscribe`] sends the query and waits for the +//! engine. When the engine ships rows for tables whose column types are +//! not yet known, the answer is [`Outcome::NeedsTypes`], the caller +//! looks the types up on whatever connection it has (blocking or async) +//! and hands them to [`Subscriptions::materialize`], which builds the +//! store and gives the verdict. [`Subscriptions::subscribe_with`] does +//! both on a blocking connection. + +use std::collections::HashMap; +use std::sync::Arc; + +use crate::cache::Live; +use crate::key::QueryKey; +use crate::materialized::Snapshot; +use crate::row::OwnedRow; +use crate::types::{TableTypes, TypeLookup}; + +/// The engine's answer to a subscription. +#[derive(Debug, Clone)] +pub enum Verdict { + /// The engine will report every change to this query's rows; the + /// cache reads Postgres for them and refills on each report. + Accepted, + /// The engine serves the rows itself: here is the result, built + /// from `live`, which the cache keeps and builds the result from + /// again on a read after each change. No Postgres read. + Rows { + rows: Arc<[OwnedRow]>, + live: Arc, + }, + /// The engine cannot follow this query; the read goes to Postgres and + /// is not cached. The reason is for logs and stats. + Refused(String), +} + +/// The refusal for a subscribe whose table's known column types lack a +/// column the engine names, or whose table was not in the catalog when +/// it was loaded: that read goes live, the catalog is loaded again on +/// the next miss, and that subscribe is served from rows. Transient, so +/// the cache does not remember it. +pub const STALE_TYPES: &str = "the column types are being reloaded"; + +/// The engine's answer to [`Subscriptions::subscribe`]. +pub enum Outcome { + Verdict(Verdict), + /// The engine ships rows, and `missing` are the tables whose column + /// types the engine side does not know yet. Look them up and call + /// [`Subscriptions::materialize`]. + NeedsTypes { + snapshot: Snapshot, + missing: Vec, + }, +} + +/// What the cache needs from the engine. +pub trait Subscriptions: Send + Sync { + /// Registers a query. Returns once the engine has it, since the cache + /// reads Postgres only after that. Blocks the calling thread for up + /// to the verdict timeout. + fn subscribe(&self, key: &QueryKey) -> Outcome; + + /// Finishes an [`Outcome::NeedsTypes`]: `types` are the column types + /// of the missing tables (a table left out falls the query back to + /// refills). Pure computation, no I/O. + fn materialize( + &self, + key: &QueryKey, + snapshot: Snapshot, + types: HashMap>, + ) -> Verdict; + + /// The cache dropped the entry for good (evicted for size or age, or + /// abandoned by a read), so the engine need not follow the query any + /// more. + fn unsubscribe(&self, key: &QueryKey); + + /// Whether a refusal with this reason will hold for a while (the + /// engine cannot follow the query's shape), so the cache may stop + /// asking for a time, or is transient (the engine side was down or + /// full, or did not answer) and the next read should ask again. + fn refusal_is_final(&self, _reason: &str) -> bool { + true + } + + /// Whether the engine side is too far behind the database for the + /// cache to be trusted: reads then go to Postgres until it catches up. + fn behind(&self) -> bool { + false + } + + /// Both steps on a blocking connection: `lookup` is the connection + /// the read is on. A lookup that fails stops the lookups; the engine + /// side then reports the fallback for the types it lacks. + fn subscribe_with(&self, key: &QueryKey, lookup: &mut dyn TypeLookup) -> Verdict { + match self.subscribe(key) { + Outcome::Verdict(verdict) => verdict, + Outcome::NeedsTypes { snapshot, missing } => { + let mut types = HashMap::new(); + for table in missing { + match lookup.table_types(&table) { + Ok(found) => { + types.insert(table, Arc::new(found)); + } + Err(_) => break, + } + } + self.materialize(key, snapshot, types) + } + } + } +} + +/// No engine: refuses every query, so the connection is a plain +/// passthrough. +#[derive(Debug, Default, Clone, Copy)] +pub struct Unsynced; + +impl Subscriptions for Unsynced { + /// Refused: there is no engine. + fn subscribe(&self, _: &QueryKey) -> Outcome { + Outcome::Verdict(Verdict::Refused("no engine".to_string())) + } + + /// Refused: there is no engine. + fn materialize( + &self, + _: &QueryKey, + _: Snapshot, + _: HashMap>, + ) -> Verdict { + Verdict::Refused("no engine".to_string()) + } + + /// Nothing to end. + fn unsubscribe(&self, _: &QueryKey) {} +} diff --git a/crates/diesel-streamgres/src/types.rs b/crates/diesel-streamgres/src/types.rs new file mode 100644 index 0000000..5ebcfa6 --- /dev/null +++ b/crates/diesel-streamgres/src/types.rs @@ -0,0 +1,198 @@ +//! What Postgres says a table's columns are, for encoding rows the engine +//! ships into the bytes Diesel decodes. +//! +//! The engine knows its own coarse types (string, int, float, time); +//! Diesel decodes Postgres's binary format, where an `int4` is four bytes +//! and an `int8` eight. So the exact column types come from Postgres +//! itself, once per table, through the connection the read came in on. + +use diesel::connection::{DefaultLoadingMode, LoadConnection}; +use diesel::pg::Pg; +use diesel::prelude::*; +use diesel::sql_query; +use diesel::sql_types::{Integer, Nullable, Oid, Text}; +use std::collections::HashMap; + +/// One column of a table as Postgres declares it. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ColumnType { + pub name: String, + pub oid: u32, + /// The type modifier: a `timestamp(3)`'s precision, a `varchar(64)`'s + /// length plus four, `-1` when none. + pub typmod: i32, + pub kind: TypeKind, +} + +/// How a type's values are laid out. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum TypeKind { + /// A built-in scalar, by OID. + Base, + /// A user enum: its labels travel as text. + Enum, + /// A one-dimensional array of `element`, which is `Base` or `Enum`. + Array { + element: u32, + element_kind: Box, + }, + /// Composite, range, domain, pseudo: not encoded. + Other, +} + +/// A table's columns in declaration order, which is also the order of +/// `SELECT *`. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct TableTypes { + pub columns: Vec, +} + +impl TableTypes { + /// The column named `name`, with its position. + pub fn column(&self, name: &str) -> Option<(usize, &ColumnType)> { + self.columns + .iter() + .enumerate() + .find(|(_, column)| column.name == name) + } +} + +/// Where the engine client asks for a table's types: the connection the +/// read is on. +pub trait TypeLookup { + /// The column types of `table`, in declaration order. + fn table_types(&mut self, table: &str) -> QueryResult; +} + +/// One column as `pg_attribute` and `pg_type` describe it. +#[derive(QueryableByName)] +struct Attribute { + #[diesel(sql_type = Text)] + attname: String, + #[diesel(sql_type = Oid)] + atttypid: u32, + #[diesel(sql_type = Integer)] + atttypmod: i32, + #[diesel(sql_type = Text)] + typtype: String, + #[diesel(sql_type = Oid)] + typelem: u32, + #[diesel(sql_type = Text)] + typcategory: String, + #[diesel(sql_type = Nullable)] + elemtyptype: Option, +} + +impl Attribute { + /// The column's type. An array is a type of the array category: a + /// fixed-width type like `name` or `point` also names an element + /// type, and is not one. + fn into_column(self) -> ColumnType { + let kind = match self.typtype.as_str() { + "b" if self.typcategory == "A" && self.elemtyptype.is_some() => { + let element_kind = match self.elemtyptype.as_deref() { + Some("b") => TypeKind::Base, + Some("e") => TypeKind::Enum, + _ => TypeKind::Other, + }; + TypeKind::Array { + element: self.typelem, + element_kind: Box::new(element_kind), + } + } + "b" => TypeKind::Base, + "e" => TypeKind::Enum, + _ => TypeKind::Other, + }; + ColumnType { + name: self.attname, + oid: self.atttypid, + typmod: self.atttypmod, + kind, + } + } +} + +const ATTRIBUTE_COLUMNS: &str = "a.attname, a.atttypid, a.atttypmod, t.typtype::text AS typtype, \ + t.typelem, t.typcategory::text AS typcategory, et.typtype::text AS elemtyptype"; + +impl TypeLookup for C +where + C: LoadConnection + Connection, +{ + /// Reads the table's columns from the catalog on this connection. + fn table_types(&mut self, table: &str) -> QueryResult { + let attributes = sql_query(format!( + "SELECT {ATTRIBUTE_COLUMNS} \ + FROM pg_attribute a \ + JOIN pg_type t ON t.oid = a.atttypid \ + LEFT JOIN pg_type et ON et.oid = t.typelem AND t.typelem <> 0 \ + WHERE a.attrelid = quote_ident($1)::regclass AND a.attnum > 0 AND NOT a.attisdropped \ + ORDER BY a.attnum" + )) + .bind::(table) + .load::(self)?; + Ok(TableTypes { + columns: attributes.into_iter().map(Attribute::into_column).collect(), + }) + } +} + +/// One column of one table, from [`ALL_TABLES_SQL`]: how a connection +/// that cannot run a query per table (an async one, whose futures are +/// built before the engine names the tables) learns every table's types +/// at once. +#[derive(QueryableByName)] +pub struct TableAttribute { + #[diesel(sql_type = Text)] + relname: String, + #[diesel(embed)] + attribute: Attribute, +} + +/// Every column of every table in the search path, tables in no +/// particular order, columns in declaration order. A table name present +/// in several schemas resolves to the first in the search path, as a +/// bare name does. +pub const ALL_TABLES_SQL: &str = "SELECT c.relname, a.attname, a.atttypid, a.atttypmod, \ + t.typtype::text AS typtype, t.typelem, t.typcategory::text AS typcategory, \ + et.typtype::text AS elemtyptype \ + FROM pg_class c \ + JOIN pg_namespace n ON n.oid = c.relnamespace \ + JOIN pg_attribute a ON a.attrelid = c.oid \ + JOIN pg_type t ON t.oid = a.atttypid \ + LEFT JOIN pg_type et ON et.oid = t.typelem AND t.typelem <> 0 \ + WHERE c.relkind IN ('r', 'p', 'v', 'm') \ + AND n.nspname = ANY(current_schemas(false)) \ + AND a.attnum > 0 AND NOT a.attisdropped \ + ORDER BY array_position(current_schemas(false), n.nspname), c.relname, a.attnum"; + +/// Groups [`ALL_TABLES_SQL`]'s rows by table. Rows of one table are +/// contiguous; a later run under a name already taken is the same name +/// in a later schema, which a bare name does not mean, so it is skipped +/// whole. +pub fn all_table_types(rows: Vec) -> std::collections::HashMap { + let mut tables: HashMap = HashMap::new(); + let mut current: Option<(String, bool)> = None; + for row in rows { + let skip = match ¤t { + Some((name, skip)) if *name == row.relname => *skip, + _ => { + let skip = tables.contains_key(&row.relname); + current = Some((row.relname.clone(), skip)); + skip + } + }; + if skip { + continue; + } + tables + .entry(row.relname) + .or_insert_with(|| TableTypes { + columns: Vec::new(), + }) + .columns + .push(row.attribute.into_column()); + } + tables +} diff --git a/crates/diesel-streamgres/tests/async_cached.rs b/crates/diesel-streamgres/tests/async_cached.rs new file mode 100644 index 0000000..596a94f --- /dev/null +++ b/crates/diesel-streamgres/tests/async_cached.rs @@ -0,0 +1,685 @@ +//! The async wrapper against a fake engine side: the same rules as the +//! blocking one, reached through diesel-async. + +#![cfg(feature = "async")] + +use std::collections::HashMap; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use diesel::prelude::*; +use diesel_async::{AsyncConnection, AsyncPgConnection, RunQueryDsl, SimpleAsyncConnection}; +use diesel_streamgres::cache::Cache; +use diesel_streamgres::key::QueryKey; +use diesel_streamgres::materialized::Snapshot; +use diesel_streamgres::subscriptions::{Outcome, Subscriptions, Verdict}; +use diesel_streamgres::types::TableTypes; +use diesel_streamgres::{AsyncSyncedConnection, Route, RouteDsl, Served}; +use scoped_futures::ScopedFutureExt; + +/// diesel-async's `RunQueryDsl::load` shadows `AtomicU64::load`. +fn stat(counter: &AtomicU64) -> u64 { + AtomicU64::load(counter, Ordering::Relaxed) +} + +diesel::table! { + sg_configs (id) { + id -> Int4, + key -> Text, + value -> Text, + } +} + +/// A user-defined type: Diesel must look its OIDs up, and an array of it +/// cannot even be rendered until it has. +#[derive(diesel::sql_types::SqlType, diesel::query_builder::QueryId)] +#[diesel(postgres_type(name = "sg_mood"))] +pub struct SgMood; + +/// A value of the `sg_mood` enum. +#[derive(Debug, Clone, Copy, PartialEq, diesel::AsExpression, diesel::FromSqlRow)] +#[diesel(sql_type = SgMood)] +enum Mood { + Happy, + Sad, +} + +impl diesel::serialize::ToSql for Mood { + /// The label. + fn to_sql<'b>( + &'b self, + out: &mut diesel::serialize::Output<'b, '_, diesel::pg::Pg>, + ) -> diesel::serialize::Result { + use std::io::Write; + out.write_all(match self { + Mood::Happy => b"happy", + Mood::Sad => b"sad", + })?; + Ok(diesel::serialize::IsNull::No) + } +} + +impl diesel::deserialize::FromSql for Mood { + /// The variant the label names. + fn from_sql(bytes: diesel::pg::PgValue<'_>) -> diesel::deserialize::Result { + match bytes.as_bytes() { + b"happy" => Ok(Mood::Happy), + b"sad" => Ok(Mood::Sad), + other => Err(format!("not a mood: {other:?}").into()), + } + } +} + +diesel::table! { + use diesel::sql_types::*; + use super::SgMood; + sg_moods (id) { + id -> Int4, + key -> Text, + mood -> SgMood, + } +} + +/// The test database's URL. +fn url() -> String { + std::env::var("DATABASE_URL") + .unwrap_or_else(|_| "postgres:///diesel_streamgres_test?host=/tmp".to_string()) +} + +/// An engine side that accepts every read of a test table, refuses the +/// rest, and records what it was asked. +#[derive(Default)] +struct FakeEngine { + subscribed: Mutex>, + unsubscribed: Mutex>, +} + +impl Subscriptions for FakeEngine { + /// Accepted for a test table, refused otherwise. + fn subscribe(&self, key: &QueryKey) -> Outcome { + if !key.sql.contains("sg_") { + return Outcome::Verdict(Verdict::Refused("not a config read".into())); + } + self.subscribed.lock().unwrap().push(key.clone()); + Outcome::Verdict(Verdict::Accepted) + } + + /// Accepted, with no rows. + fn materialize( + &self, + _: &QueryKey, + _: Snapshot, + _: HashMap>, + ) -> Verdict { + Verdict::Accepted + } + + /// Records the key. + fn unsubscribe(&self, key: &QueryKey) { + self.unsubscribed.lock().unwrap().push(key.clone()); + } +} + +/// One test's wrapped connection and cache, its fake engine, a plain +/// connection for changes, and the prefix of its keys. +struct World { + conn: AsyncSyncedConnection, + cache: Arc, + engine: Arc, + other: AsyncPgConnection, + prefix: String, +} + +impl World { + /// A world for test `name`, the schema created and its keys cleared. + /// The tests run in parallel; an advisory lock serialises the DDL. + async fn new(name: &str) -> World { + let mut other = AsyncPgConnection::establish(&url()).await.expect("connect"); + other + .batch_execute( + "BEGIN; + SELECT pg_advisory_xact_lock(424242); + CREATE TABLE IF NOT EXISTS sg_configs ( + id serial PRIMARY KEY, key text NOT NULL UNIQUE, value text NOT NULL); + DO $$ BEGIN CREATE TYPE sg_mood AS ENUM ('happy', 'sad'); + EXCEPTION WHEN duplicate_object OR unique_violation THEN NULL; END $$; + CREATE TABLE IF NOT EXISTS sg_moods ( + id serial PRIMARY KEY, key text NOT NULL UNIQUE, mood sg_mood NOT NULL); + COMMIT", + ) + .await + .expect("schema"); + let engine = Arc::new(FakeEngine::default()); + let cache = Arc::new(Cache::new(64 * 1024 * 1024, engine.clone())); + let inner = AsyncPgConnection::establish(&url()).await.expect("connect"); + let conn = AsyncSyncedConnection::with_cache(inner, cache.clone(), Route::Streamgres); + let prefix = format!("{name}-{}-", std::process::id()); + diesel::delete(sg_configs::table.filter(sg_configs::key.like(format!("{prefix}%")))) + .execute(&mut other) + .await + .expect("clean"); + World { + conn, + cache, + engine, + other, + prefix, + } + } + + /// A key of this test's. + fn key(&self, name: &str) -> String { + format!("{}{name}", self.prefix) + } + + /// Sets a key on the plain connection, as another service would. + async fn put(&mut self, name: &str, value: &str) { + let key = self.key(name); + diesel::insert_into(sg_configs::table) + .values((sg_configs::key.eq(&key), sg_configs::value.eq(value))) + .on_conflict(sg_configs::key) + .do_update() + .set(sg_configs::value.eq(value)) + .execute(&mut self.other) + .await + .expect("put"); + } + + /// A key's value through the wrapper. + async fn read(&mut self, name: &str) -> Option { + let key = self.key(name); + sg_configs::table + .filter(sg_configs::key.eq(key)) + .select(sg_configs::value) + .first(&mut self.conn) + .await + .optional() + .expect("read") + } + + /// A key's value through the wrapper, flagged with `route`. + async fn read_via(&mut self, name: &str, route: Route) -> Option { + let key = self.key(name); + sg_configs::table + .filter(sg_configs::key.eq(key)) + .select(sg_configs::value) + .via(route) + .first(&mut self.conn) + .await + .optional() + .expect("read via") + } + + /// The cache key of a read of `name`. It binds built-in types only, + /// so no type lookup is needed. + fn key_of(&mut self, name: &str) -> QueryKey { + let key = self.key(name); + let query = sg_configs::table + .filter(sg_configs::key.eq(key)) + .select(sg_configs::value) + .limit(1); + /// A lookup that is never asked. + struct NoLookup; + impl diesel::pg::PgMetadataLookup for NoLookup { + /// Never called. + fn lookup_type(&mut self, _: &str, _: Option<&str>) -> diesel::pg::PgTypeMetadata { + unreachable!() + } + } + let mut lookup = NoLookup; + QueryKey::of(&query, &mut lookup).expect("key") + } + + /// Deletes this test's keys. + async fn clean(&mut self) { + let _ = diesel::delete( + sg_configs::table.filter(sg_configs::key.like(format!("{}%", self.prefix))), + ) + .execute(&mut self.other) + .await; + } + + /// The cache's hits so far. + fn hits(&self) -> u64 { + stat(&self.cache.stats.hits) + } +} + +/// A read the engine accepted is subscribed before it is read, and a repeat +/// is served from the cache. +/// +/// Without a touch the cache keeps the old row. +#[tokio::test] +async fn a_repeated_read_is_served_from_the_cache() { + let mut w = World::new("async-hit").await; + w.put("mode", "a").await; + assert_eq!(w.read("mode").await.as_deref(), Some("a")); + assert_eq!( + w.engine.subscribed.lock().unwrap().len(), + 1, + "subscribed before the read" + ); + assert_eq!(w.cache.len(), 1); + + let hits = w.hits(); + assert_eq!(w.read("mode").await.as_deref(), Some("a")); + assert_eq!(w.hits(), hits + 1, "served from the cache"); + + w.put("mode", "b").await; + assert_eq!(w.read("mode").await.as_deref(), Some("a")); + w.clean().await; +} + +/// The connection says where each read was served from: Postgres for the +/// accepted miss, the cache for the repeat, Postgres again under the flag. +#[tokio::test] +async fn the_connection_says_where_each_read_was_served_from() { + let mut w = World::new("async-served").await; + assert_eq!(w.conn.last_served(), None, "nothing read yet"); + w.put("k", "1").await; + + assert_eq!(w.read("k").await.as_deref(), Some("1")); + assert_eq!( + w.conn.last_served(), + Some(Served::Postgres), + "the miss the engine accepted" + ); + assert_eq!(w.read("k").await.as_deref(), Some("1")); + assert_eq!(w.conn.last_served(), Some(Served::Cache)); + assert_eq!(w.read_via("k", Route::Postgres).await.as_deref(), Some("1")); + assert_eq!( + w.conn.last_served(), + Some(Served::Postgres), + "the flag sent it live" + ); + w.clean().await; +} + +/// A touch makes the next read refill from Postgres, subscribing again. +#[tokio::test] +async fn a_touch_makes_the_next_read_refill() { + let mut w = World::new("async-touch").await; + w.put("mode", "a").await; + assert_eq!(w.read("mode").await.as_deref(), Some("a")); + w.put("mode", "b").await; + let key = w.key_of("mode"); + w.cache.touched(&key); + assert_eq!( + w.read("mode").await.as_deref(), + Some("b"), + "refilled from Postgres" + ); + assert_eq!(w.cache.len(), 1); + assert_eq!( + w.engine.subscribed.lock().unwrap().len(), + 2, + "the refill subscribed again, the fake forgets" + ); + w.clean().await; +} + +/// The query's flag and a route scope each send a read to Postgres without +/// disturbing the entry; the scope ends with its closure. +#[tokio::test] +async fn the_query_flag_and_the_scope_route_a_read() { + let mut w = World::new("async-route").await; + w.put("mode", "a").await; + assert_eq!(w.read("mode").await.as_deref(), Some("a")); + w.put("mode", "b").await; + assert_eq!( + w.read_via("mode", Route::Postgres).await.as_deref(), + Some("b"), + "the flag bypassed the cache" + ); + assert_eq!( + w.read("mode").await.as_deref(), + Some("a"), + "the entry is untouched" + ); + + let key = w.key("mode"); + let scoped = w + .conn + .with_route(Route::Postgres, |conn| { + async move { + sg_configs::table + .filter(sg_configs::key.eq(key)) + .select(sg_configs::value) + .first::(conn) + .await + .optional() + .unwrap() + } + .boxed() + }) + .await; + assert_eq!(scoped.as_deref(), Some("b")); + assert_eq!( + w.read("mode").await.as_deref(), + Some("a"), + "the scope ended" + ); + w.clean().await; +} + +use futures_util::FutureExt; + +/// Inside a transaction a read goes to Postgres and sees the connection's +/// own write, without subscribing; outside it the cache holds the old row +/// until a touch. +#[tokio::test] +async fn inside_a_transaction_reads_see_the_connections_own_writes() { + let mut w = World::new("async-txn").await; + w.conn.set_read_your_writes(Duration::ZERO); + w.put("count", "1").await; + assert_eq!(w.read("count").await.as_deref(), Some("1")); + let count_key = w.key("count"); + + let seen: String = w + .conn + .transaction::<_, diesel::result::Error, _>(|conn| { + let count_key = count_key.clone(); + async move { + diesel::update(sg_configs::table.filter(sg_configs::key.eq(&count_key))) + .set(sg_configs::value.eq("2")) + .execute(conn) + .await?; + sg_configs::table + .filter(sg_configs::key.eq(&count_key)) + .select(sg_configs::value) + .first(conn) + .await + } + .scope_boxed() + }) + .await + .expect("transaction"); + assert_eq!( + seen, "2", + "the read inside the transaction went to Postgres" + ); + assert_eq!( + w.engine.subscribed.lock().unwrap().len(), + 1, + "the transactional read did not subscribe" + ); + assert_eq!( + w.read("count").await.as_deref(), + Some("1"), + "outside it the cache holds the old row until a touch" + ); + w.clean().await; +} + +/// After its own write the connection reads Postgres for the read-your- +/// writes window, and the cache again once it closes. +#[tokio::test] +async fn after_a_write_the_connection_reads_postgres_for_a_window() { + let mut w = World::new("async-ryw").await; + w.conn.set_read_your_writes(Duration::from_millis(150)); + w.put("mode", "a").await; + assert_eq!(w.read("mode").await.as_deref(), Some("a")); + let key = w.key("mode"); + diesel::update(sg_configs::table.filter(sg_configs::key.eq(&key))) + .set(sg_configs::value.eq("b")) + .execute(&mut w.conn) + .await + .unwrap(); + assert_eq!( + w.read("mode").await.as_deref(), + Some("b"), + "its own write is seen at once" + ); + assert_eq!(stat(&w.cache.stats.after_write), 1); + tokio::time::sleep(Duration::from_millis(200)).await; + assert_eq!( + w.read("mode").await.as_deref(), + Some("a"), + "the window closed; the entry was not refreshed" + ); + w.clean().await; +} + +/// A query the engine refuses is read live and not kept. +#[tokio::test] +async fn a_query_the_engine_refuses_is_read_live() { + let mut w = World::new("async-refused").await; + let count: i64 = diesel::sql_query("SELECT count(*) AS n FROM pg_class") + .get_result::(&mut w.conn) + .await + .map(|c| c.n) + .unwrap(); + assert!(count > 0); + assert_eq!( + w.cache.len(), + 0, + "not a config read: the fake refused it, the read was live" + ); + assert_eq!(stat(&w.cache.stats.refused), 1); + w.clean().await; +} + +/// A count, read by name. +#[derive(QueryableByName)] +struct Count { + #[diesel(sql_type = diesel::sql_types::BigInt)] + n: i64, +} + +/// A read binding an array of an enum is live once, keyed and subscribed +/// the second time, cached the third. +/// +/// The array's element OID is unknown: the query cannot be keyed, so it +/// goes live while the type is resolved. +/// +/// The OIDs live in the cache: a sibling connection knows them at once. +#[tokio::test] +async fn an_enum_array_bind_is_live_once_then_cached() { + let mut w = World::new("async-enum-array").await; + let key = w.key("m"); + diesel::sql_query( + "INSERT INTO sg_moods (key, mood) VALUES ($1, 'happy') ON CONFLICT (key) DO NOTHING", + ) + .bind::(&key) + .execute(&mut w.other) + .await + .unwrap(); + /// The mood of `key`, if it is one of the two. + async fn read(conn: &mut AsyncSyncedConnection, key: String) -> Option { + sg_moods::table + .filter(sg_moods::key.eq(key)) + .filter(sg_moods::mood.eq_any(vec![Mood::Happy, Mood::Sad])) + .select(sg_moods::mood) + .first::(conn) + .await + .optional() + .unwrap() + } + assert_eq!(read(&mut w.conn, key.clone()).await, Some(Mood::Happy)); + assert_eq!( + w.engine.subscribed.lock().unwrap().len(), + 0, + "not subscribed the first time" + ); + assert_eq!(read(&mut w.conn, key.clone()).await, Some(Mood::Happy)); + assert_eq!( + w.engine.subscribed.lock().unwrap().len(), + 1, + "keyed and subscribed the second time" + ); + assert_eq!(read(&mut w.conn, key.clone()).await, Some(Mood::Happy)); + assert_eq!(w.hits(), 1, "cached the third time"); + + let inner = AsyncPgConnection::establish(&url()).await.expect("connect"); + let mut sibling = + AsyncSyncedConnection::with_cache(inner, Arc::clone(&w.cache), Route::Streamgres); + assert_eq!(read(&mut sibling, key.clone()).await, Some(Mood::Happy)); + assert_eq!( + w.hits(), + 2, + "the sibling hit the same entry without a live read" + ); + let _ = diesel::sql_query("DELETE FROM sg_moods WHERE key = $1") + .bind::(&key) + .execute(&mut w.other) + .await; + w.clean().await; +} + +/// Two reads pipelined on one connection do not ask for its transaction +/// state. +/// +/// Two futures from one connection, both pending at once: diesel-async +/// shares its transaction state with each, and asking for it while the +/// first is pending is a panic. +#[tokio::test] +async fn pipelined_reads_do_not_touch_the_transaction_state() { + let mut w = World::new("async-pipeline").await; + w.put("a", "1").await; + w.put("b", "2").await; + let query = |key: String| { + sg_configs::table + .filter(sg_configs::key.eq(key)) + .select(sg_configs::value) + .limit(1) + }; + let (key_a, key_b) = (w.key("a"), w.key("b")); + let first = AsyncConnection::load(&mut w.conn, query(key_a)); + let second = AsyncConnection::load(&mut w.conn, query(key_b)); + let (first, second) = futures_util::join!(first, second); + let first: Vec = first + .unwrap() + .map(|row| { + row.and_then(|row| { + >::build_from_row( + &row, + ) + .map_err(diesel::result::Error::DeserializationError) + }) + }) + .try_collect() + .await + .unwrap(); + let second: Vec = second + .unwrap() + .map(|row| { + row.and_then(|row| { + >::build_from_row( + &row, + ) + .map_err(diesel::result::Error::DeserializationError) + }) + }) + .try_collect() + .await + .unwrap(); + assert_eq!( + (first, second), + (vec!["1".to_string()], vec!["2".to_string()]) + ); + w.clean().await; +} + +/// A route scope dropped before it completes still restores the outer +/// route. +#[tokio::test] +async fn a_route_scope_ends_when_its_future_is_dropped() { + let mut w = World::new("async-scope-drop").await; + w.put("mode", "a").await; + assert_eq!(w.read("mode").await.as_deref(), Some("a")); + w.put("mode", "b").await; + let dropped = tokio::time::timeout( + Duration::from_millis(10), + w.conn + .with_route(Route::Postgres, |_| std::future::pending::<()>().boxed()), + ) + .await; + assert!(dropped.is_err(), "the scope's future was dropped mid-way"); + assert_eq!( + w.read("mode").await.as_deref(), + Some("a"), + "the scope was restored: the cache answers" + ); + w.clean().await; +} + +/// A read dropped after its subscribe may have landed ends the subscription +/// it made, and releases its reservation. +/// +/// Exactly one poll, then the drop: the subscribe may already have landed +/// on its blocking thread by then, and the Postgres read has at best been +/// sent. +/// +/// The drop ends the subscription at once (a no-op on the engine side if +/// the subscribe has not landed yet); the blocking thread's subscribe lands +/// after, and finds the read gone. +/// +/// The guard's drop and the blocking thread's check can both end it; the +/// second is a no-op on the engine side, so either count is right. The +/// cache is told its database first, which a first miss would read +/// before it subscribes. +#[tokio::test] +async fn a_read_dropped_after_subscribing_ends_its_subscription() { + let mut w = World::new("async-drop-read").await; + w.cache.learn_database(None); + w.put("mode", "a").await; + let query = sg_configs::table + .filter(sg_configs::key.eq(w.key("mode"))) + .select(sg_configs::value) + .limit(1); + let mut read = Box::pin(AsyncConnection::load(&mut w.conn, query)); + assert!( + futures_util::poll!(read.as_mut()).is_pending(), + "a miss cannot complete in one poll" + ); + drop(read); + for _ in 0..100 { + let subscribed = w.engine.subscribed.lock().unwrap().len(); + let unsubscribed = w.engine.unsubscribed.lock().unwrap().len(); + if subscribed == 1 && unsubscribed >= 1 { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + let subscribed = w.engine.subscribed.lock().unwrap().len(); + let unsubscribed = w.engine.unsubscribed.lock().unwrap().len(); + assert_eq!(subscribed, 1, "the subscribe landed on its blocking thread"); + assert!( + (1..=2).contains(&unsubscribed), + "the subscription the dropped read made was ended: {unsubscribed}" + ); + tokio::time::sleep(Duration::from_millis(20)).await; + assert_eq!(w.cache.reservations(), 0, "and its reservation released"); + assert_eq!(w.cache.len(), 0); + w.clean().await; +} + +/// A write on a sibling connection opens the read-your-writes window of the +/// cache they share. +#[tokio::test] +async fn a_write_on_a_sibling_connection_is_seen_through_the_shared_cache() { + let mut w = World::new("async-ryw-pool").await; + w.conn.set_read_your_writes(Duration::from_secs(5)); + w.put("mode", "a").await; + assert_eq!(w.read("mode").await.as_deref(), Some("a")); + let inner = AsyncPgConnection::establish(&url()).await.expect("connect"); + let mut sibling = + AsyncSyncedConnection::with_cache(inner, Arc::clone(&w.cache), Route::Streamgres); + let key = w.key("mode"); + diesel::update(sg_configs::table.filter(sg_configs::key.eq(&key))) + .set(sg_configs::value.eq("b")) + .execute(&mut sibling) + .await + .unwrap(); + assert_eq!( + w.read("mode").await.as_deref(), + Some("b"), + "the other connection reads Postgres for the window" + ); + assert_eq!(stat(&w.cache.stats.after_write), 1); + w.clean().await; +} + +use diesel::deserialize::FromSqlRow; +use futures_util::{StreamExt, TryStreamExt}; diff --git a/crates/diesel-streamgres/tests/async_live.rs b/crates/diesel-streamgres/tests/async_live.rs new file mode 100644 index 0000000..9e5fdc0 --- /dev/null +++ b/crates/diesel-streamgres/tests/async_live.rs @@ -0,0 +1,145 @@ +//! The async wrapper against a running engine: rows mode end to end, +//! including the one-shot load of every table's types. +//! +//! Needs `STREAMGRES_S2S_URL` and `DATABASE_URL`; skips itself otherwise. + +#![cfg(feature = "async")] + +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::{Duration, Instant}; + +use diesel::prelude::*; +use diesel_async::{AsyncConnection, AsyncPgConnection, RunQueryDsl, SimpleAsyncConnection}; +use diesel_streamgres::cache::Cache; +use diesel_streamgres::engine::EngineClient; +use diesel_streamgres::{AsyncSyncedConnection, Route}; + +/// diesel-async's `RunQueryDsl::load` shadows `AtomicU64::load`. +fn stat(counter: &AtomicU64) -> u64 { + AtomicU64::load(counter, Ordering::Relaxed) +} + +diesel::table! { + sg_configs (id) { + id -> Int4, + key -> Text, + value -> Text, + } +} + +/// Sets `key` to `value` on a plain connection, as another service +/// would. +async fn put(other: &mut AsyncPgConnection, key: &str, value: &str) { + diesel::insert_into(sg_configs::table) + .values((sg_configs::key.eq(key), sg_configs::value.eq(value))) + .on_conflict(sg_configs::key) + .do_update() + .set(sg_configs::value.eq(value)) + .execute(other) + .await + .expect("put"); +} + +/// `key`'s value, through the wrapper. +async fn read(conn: &mut AsyncSyncedConnection, key: &str) -> Option { + sg_configs::table + .filter(sg_configs::key.eq(key)) + .select(sg_configs::value) + .first::(conn) + .await + .optional() + .expect("read") +} + +/// Rows mode through the async wrapper against a live engine: the first +/// read is served from the engine's rows and the second from the cache. +/// +/// A change arrives as a delta and replaces the entry; no read. +/// +/// A second query on the same table: the types are known, no catalog load, +/// still rows mode. +#[tokio::test] +async fn rows_mode_through_the_async_wrapper() { + let Ok(url) = std::env::var("STREAMGRES_S2S_URL") else { + eprintln!("STREAMGRES_S2S_URL not set; skipping"); + return; + }; + let database = std::env::var("DATABASE_URL").expect("DATABASE_URL"); + let engine = EngineClient::connect(url); + let cache = Arc::new(Cache::new(64 * 1024 * 1024, Arc::new(engine.clone()))); + engine.attach(&cache); + let deadline = Instant::now() + Duration::from_secs(5); + while !engine.is_connected() { + assert!(Instant::now() < deadline, "the engine did not connect"); + tokio::time::sleep(Duration::from_millis(50)).await; + } + let mut other = AsyncPgConnection::establish(&database) + .await + .expect("connect"); + other + .batch_execute( + "CREATE TABLE IF NOT EXISTS sg_configs ( + id serial PRIMARY KEY, key text NOT NULL UNIQUE, value text NOT NULL)", + ) + .await + .expect("schema"); + let prefix = format!("async-live-{}-", std::process::id()); + let key = format!("{prefix}r"); + put(&mut other, &key, "1").await; + + let inner = AsyncPgConnection::establish(&database) + .await + .expect("connect"); + let mut conn = AsyncSyncedConnection::with_cache(inner, cache.clone(), Route::Streamgres); + assert_eq!(read(&mut conn, &key).await.as_deref(), Some("1")); + assert_eq!( + engine.materialized(), + 1, + "served from the engine's rows: {:?}", + engine.fallbacks() + ); + assert_eq!(engine.subscribed(), 1); + let hits = stat(&cache.stats.hits); + assert_eq!(read(&mut conn, &key).await.as_deref(), Some("1")); + assert_eq!(stat(&cache.stats.hits), hits + 1); + + put(&mut other, &key, "2").await; + let started = Instant::now(); + loop { + if read(&mut conn, &key).await.as_deref() == Some("2") { + break; + } + assert!( + started.elapsed() < Duration::from_secs(10), + "the change did not arrive" + ); + tokio::time::sleep(Duration::from_millis(20)).await; + } + eprintln!( + "ASYNC rows mode commit to replaced entry: {:?}", + started.elapsed() + ); + assert!(stat(&cache.stats.replaced) >= 1); + assert_eq!(stat(&cache.stats.touched), 0); + + let key2 = format!("{prefix}r2"); + diesel::insert_into(sg_configs::table) + .values((sg_configs::key.eq(&key2), sg_configs::value.eq("x"))) + .execute(&mut other) + .await + .unwrap(); + let got: Option = sg_configs::table + .filter(sg_configs::key.eq(&key2)) + .select(sg_configs::value) + .first(&mut conn) + .await + .optional() + .unwrap(); + assert_eq!(got.as_deref(), Some("x")); + assert_eq!(engine.materialized(), 2, "{:?}", engine.fallbacks()); + + let _ = diesel::delete(sg_configs::table.filter(sg_configs::key.like(format!("{prefix}%")))) + .execute(&mut other) + .await; +} diff --git a/crates/diesel-streamgres/tests/cached_reads.rs b/crates/diesel-streamgres/tests/cached_reads.rs new file mode 100644 index 0000000..4bbae60 --- /dev/null +++ b/crates/diesel-streamgres/tests/cached_reads.rs @@ -0,0 +1,1120 @@ +//! Reads through a cache that a fake engine side keeps honest. + +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; + +use diesel::connection::{Instrumentation, InstrumentationEvent, SimpleConnection}; +use diesel::prelude::*; +use diesel::sql_query; +use diesel::sql_types::{Integer, Text}; +use diesel_streamgres::cache::{Cache, Limits}; +use diesel_streamgres::key::QueryKey; +use diesel_streamgres::subscriptions::{Outcome, Subscriptions, Verdict}; +use diesel_streamgres::{Consistency, Route, RouteDsl, Served, SyncedConnection}; + +diesel::table! { + sg_configs (id) { + id -> Int4, + key -> Text, + value -> Text, + } +} + +/// A config row. +#[derive(Queryable, QueryableByName, Selectable, Debug, PartialEq, Clone)] +#[diesel(table_name = sg_configs)] +struct Config { + id: i32, + key: String, + value: String, +} + +/// The test database's URL. +fn url() -> String { + std::env::var("DATABASE_URL") + .unwrap_or_else(|_| "postgres:///diesel_streamgres_test?host=/tmp".to_string()) +} + +/// The engine side as a test sees it: accepts every `SELECT` on +/// `sg_configs`, refuses the rest, and remembers every call. +#[derive(Default)] +struct FakeEngine { + subscribed: Mutex>, + unsubscribed: Mutex>, + refuse_all: std::sync::atomic::AtomicBool, + /// Says the engine is too far behind for the cache to be trusted. + behind: std::sync::atomic::AtomicBool, + /// The next subscribe ships no rows, the engine's view, and leaves + /// the engine side behind, the way a socket that just reconnected to + /// a lagging engine hears its lag before the answer. + rows_from_behind: std::sync::atomic::AtomicBool, +} + +/// An engine-side store that holds no rows. +#[derive(Debug)] +struct NoRows; + +impl diesel_streamgres::cache::Live for NoRows { + /// Always empty. + fn rows(&self) -> Option> { + Some(Arc::from(Vec::new())) + } + + /// Weighs nothing. + fn bytes(&self) -> usize { + 0 + } +} + +impl Subscriptions for FakeEngine { + /// Accepted for `sg_configs` unless told to refuse all, refused + /// otherwise. + fn subscribe(&self, key: &QueryKey) -> Outcome { + if self.refuse_all.load(Ordering::SeqCst) || !key.sql.contains("sg_configs") { + return Outcome::Verdict(Verdict::Refused("not a config read".into())); + } + self.subscribed.lock().unwrap().push(key.clone()); + if self.rows_from_behind.swap(false, Ordering::SeqCst) { + self.behind.store(true, Ordering::SeqCst); + return Outcome::Verdict(Verdict::Rows { + rows: Arc::from(Vec::new()), + live: Arc::new(NoRows), + }); + } + Outcome::Verdict(Verdict::Accepted) + } + + /// Accepted, with no rows. + fn materialize( + &self, + _: &QueryKey, + _: diesel_streamgres::materialized::Snapshot, + _: std::collections::HashMap>, + ) -> Verdict { + Verdict::Accepted + } + + /// Records the key. + fn unsubscribe(&self, key: &QueryKey) { + self.unsubscribed.lock().unwrap().push(key.clone()); + } + + /// As set by the test. + fn behind(&self) -> bool { + self.behind.load(Ordering::SeqCst) + } +} + +/// Counts the queries the connection starts. +struct CountQueries(Arc); + +impl Instrumentation for CountQueries { + /// Counts a query start. + fn on_connection_event(&mut self, event: InstrumentationEvent<'_>) { + if let InstrumentationEvent::StartQuery { .. } = event { + self.0.fetch_add(1, Ordering::SeqCst); + } + } +} + +/// One test's wrapped connection and cache, its fake engine, a plain +/// connection for changes, and the prefix of its keys. +struct World { + conn: SyncedConnection, + cache: Arc, + engine: Arc, + /// Queries the inner connection actually sent to Postgres. + sent: Arc, + /// A second, plain connection for writes "from elsewhere". + other: PgConnection, + prefix: String, +} + +impl World { + /// Each test gets its own key prefix so the tests can run in parallel + /// against committed rows, and cleans up after itself. + fn new(name: &str) -> World { + Self::with_budget(name, 64 * 1024 * 1024) + } + + /// A world whose cache holds at most `max_bytes`. + fn with_budget(name: &str, max_bytes: u64) -> World { + Self::with_limits( + name, + Limits { + max_bytes, + ..Limits::default() + }, + ) + } + + /// A world whose cache has `limits`, the table created once. + fn with_limits(name: &str, limits: Limits) -> World { + static SCHEMA_READY: std::sync::Once = std::sync::Once::new(); + SCHEMA_READY.call_once(|| { + let mut conn = PgConnection::establish(&url()).expect("connect"); + conn.batch_execute( + "CREATE TABLE IF NOT EXISTS sg_configs ( + id serial PRIMARY KEY, key text NOT NULL UNIQUE, value text NOT NULL)", + ) + .expect("schema"); + }); + let engine = Arc::new(FakeEngine::default()); + let cache = Arc::new(Cache::with_limits(limits, engine.clone())); + let mut inner = PgConnection::establish(&url()).expect("connect"); + let sent = Arc::new(AtomicUsize::new(0)); + inner.set_instrumentation(CountQueries(sent.clone())); + let conn = SyncedConnection::with_cache(inner, cache.clone(), Consistency::Eventual); + let mut other = PgConnection::establish(&url()).expect("connect"); + let prefix = format!("{name}-{}-", std::process::id()); + diesel::delete(sg_configs::table.filter(sg_configs::key.like(format!("{prefix}%")))) + .execute(&mut other) + .expect("clean"); + World { + conn, + cache, + engine, + sent, + other, + prefix, + } + } + + /// A key of this test's. + fn key(&self, name: &str) -> String { + format!("{}{name}", self.prefix) + } + + /// Writes a config from "elsewhere": another connection, no cache. + fn put(&mut self, name: &str, value: &str) { + let key = self.key(name); + diesel::insert_into(sg_configs::table) + .values((sg_configs::key.eq(&key), sg_configs::value.eq(value))) + .on_conflict(sg_configs::key) + .do_update() + .set(sg_configs::value.eq(value)) + .execute(&mut self.other) + .expect("put"); + } + + /// A key's value through the wrapper. + fn read(&mut self, name: &str) -> Option { + let key = self.key(name); + sg_configs::table + .filter(sg_configs::key.eq(key)) + .select(sg_configs::value) + .first(&mut self.conn) + .optional() + .expect("read") + } + + /// The key the cache uses for `read(name)`: `first` adds `LIMIT 1` as + /// a bind, so the key carries two. + fn key_of(&mut self, name: &str) -> QueryKey { + let key = self.key(name); + let query = sg_configs::table + .filter(sg_configs::key.eq(key)) + .select(sg_configs::value) + .limit(1); + QueryKey::of(&query, self.conn.inner()).expect("key") + } + + /// How many queries reached Postgres so far. + fn sent(&self) -> usize { + self.sent.load(Ordering::SeqCst) + } + + /// `read(name)` with the query itself flagged `route`. + fn read_via(&mut self, name: &str, route: Route) -> Option { + let key = self.key(name); + sg_configs::table + .filter(sg_configs::key.eq(key)) + .select(sg_configs::value) + .via(route) + .first(&mut self.conn) + .optional() + .expect("read via") + } +} + +impl Drop for World { + /// Deletes this test's keys. + fn drop(&mut self) { + let _ = diesel::delete( + sg_configs::table.filter(sg_configs::key.like(format!("{}%", self.prefix))), + ) + .execute(&mut self.other); + } +} + +/// A repeated read is served from the cache, without reaching Postgres. +#[test] +fn a_repeated_read_is_served_from_the_cache() { + let mut w = World::new("repeat"); + w.put("timeout", "30"); + + assert_eq!(w.read("timeout").as_deref(), Some("30")); + let after_first = w.sent(); + assert_eq!(w.read("timeout").as_deref(), Some("30")); + assert_eq!(w.read("timeout").as_deref(), Some("30")); + + assert_eq!( + w.sent(), + after_first, + "the second and third read sent nothing to Postgres" + ); + assert_eq!(w.cache.stats.hits.load(Ordering::Relaxed), 2); + assert_eq!(w.cache.stats.misses.load(Ordering::Relaxed), 1); + assert_eq!( + w.engine.subscribed.lock().unwrap().len(), + 1, + "subscribed once" + ); +} + +/// The connection says where each read was served from. +#[test] +fn the_connection_says_where_each_read_was_served_from() { + let mut w = World::new("served"); + assert_eq!(w.conn.last_served(), None, "nothing read yet"); + w.put("k", "1"); + + assert_eq!(w.read("k").as_deref(), Some("1")); + assert_eq!( + w.conn.last_served(), + Some(Served::Postgres), + "the miss the engine accepted" + ); + assert_eq!(w.read("k").as_deref(), Some("1")); + assert_eq!(w.conn.last_served(), Some(Served::Cache)); + + let key = w.key("k"); + let value: Option = sg_configs::table + .filter(sg_configs::key.eq(&key)) + .select(sg_configs::value) + .via(Route::Postgres) + .first(&mut w.conn) + .optional() + .expect("read via"); + assert_eq!(value.as_deref(), Some("1")); + assert_eq!( + w.conn.last_served(), + Some(Served::Postgres), + "the flag sent it live" + ); + + w.engine.refuse_all.store(true, Ordering::SeqCst); + assert_eq!(w.read("other").as_deref(), None); + assert_eq!( + w.conn.last_served(), + Some(Served::Postgres), + "refused, so live" + ); +} + +/// The subscription is registered before Postgres is read, so no commit in +/// between goes unreported. The cache is told its database first, which +/// its first miss would otherwise read, so the one read counted is the +/// query's. +#[test] +fn the_subscription_is_registered_before_postgres_is_read() { + let mut w = World::new("order"); + w.cache.learn_database(None); + w.put("a", "1"); + let before = w.sent(); + w.read("a"); + let key = w.key_of("a"); + let subscribed = w.engine.subscribed.lock().unwrap(); + assert_eq!(subscribed.as_slice(), std::slice::from_ref(&key)); + assert_eq!(w.sent(), before + 1, "exactly one read reached Postgres"); + assert!(key.is_select()); + assert_eq!(key.binds.len(), 2, "the key's value and the LIMIT"); + assert_eq!(key.binds[0].type_oid, 25, "the first bind is a text"); + assert_eq!(key.binds[1].type_oid, 20, "the limit is an int8"); +} + +/// The first miss through a cache reads which database the cache's +/// reads go to, once, beside the query; later misses do not. +#[test] +fn the_first_miss_reads_the_database_once() { + let mut w = World::new("database"); + w.put("a", "1"); + w.put("b", "2"); + assert!(!w.cache.knows_database()); + let before = w.sent(); + w.read("a"); + assert_eq!(w.sent(), before + 2, "the database and the read"); + let database = w.cache.database().expect("read").clone(); + assert!(!database.system.is_empty() && !database.name.is_empty(), "{database:?}"); + w.read("b"); + assert_eq!(w.sent(), before + 3, "only the read"); +} + +/// While the engine side says it is too far behind, reads go to Postgres +/// and see a change the engine has not reported; once it catches up the +/// cache answers again, with what it held. +#[test] +fn a_lagging_engine_sends_reads_to_postgres() { + let mut w = World::new("lagging"); + w.put("mode", "a"); + assert_eq!(w.read("mode").as_deref(), Some("a")); + w.put("mode", "b"); + assert_eq!(w.read("mode").as_deref(), Some("a"), "cached, unreported"); + w.engine.behind.store(true, Ordering::SeqCst); + assert_eq!(w.read("mode").as_deref(), Some("b"), "the engine is behind: Postgres"); + assert_eq!(w.conn.last_served(), Some(Served::Postgres)); + assert_eq!(w.cache.stats.lagging.load(Ordering::Relaxed), 1); + w.engine.behind.store(false, Ordering::SeqCst); + assert_eq!(w.read("mode").as_deref(), Some("a"), "caught up: the cache again"); + assert_eq!(w.conn.last_served(), Some(Served::Cache)); +} + +/// An engine side found behind only once its answer arrives does not +/// have that answer served: the read goes to Postgres. +#[test] +fn rows_from_a_lagging_engine_are_not_served() { + let mut w = World::new("rows_lagging"); + w.put("mode", "a"); + w.engine.rows_from_behind.store(true, Ordering::SeqCst); + assert_eq!(w.read("mode").as_deref(), Some("a"), "Postgres, not the empty rows"); + assert_eq!(w.conn.last_served(), Some(Served::Postgres)); + assert_eq!(w.cache.stats.lagging.load(Ordering::Relaxed), 1); +} + +/// The same query with different binds makes different entries. +#[test] +fn different_binds_are_different_entries() { + let mut w = World::new("binds"); + w.put("a", "1"); + w.put("b", "2"); + assert_eq!(w.read("a").as_deref(), Some("1")); + assert_eq!(w.read("b").as_deref(), Some("2")); + assert_eq!(w.read("a").as_deref(), Some("1")); + assert_eq!(w.cache.len(), 2); + assert_ne!(w.key_of("a"), w.key_of("b")); +} + +/// A change from elsewhere is not seen by the cache until the engine +/// reports the query touched; the next read then refills from Postgres. +#[test] +fn a_touch_makes_the_next_read_refill_from_postgres() { + let mut w = World::new("touch"); + w.put("limit", "10"); + assert_eq!(w.read("limit").as_deref(), Some("10")); + + w.put("limit", "20"); + assert_eq!( + w.read("limit").as_deref(), + Some("10"), + "stale until the engine says so" + ); + + let key = w.key_of("limit"); + w.cache.touched(&key); + let before = w.sent(); + assert_eq!(w.read("limit").as_deref(), Some("20")); + assert_eq!(w.sent(), before + 1, "refilled with one read"); + assert_eq!(w.read("limit").as_deref(), Some("20")); + assert_eq!(w.sent(), before + 1, "and cached again"); + assert!( + w.engine.unsubscribed.lock().unwrap().is_empty(), + "a touch keeps the subscription" + ); +} + +/// A touch that lands while a miss is being read is not lost: the fill is +/// dropped and the next read refills. +/// +/// What `load` does on a miss, with a touch landing between the reservation +/// and the fill. +#[test] +fn a_touch_during_the_read_is_not_lost() { + let mut w = World::new("race"); + w.put("flag", "off"); + let key = w.key_of("flag"); + + let reservation = w.cache.reserve(&key); + w.put("flag", "on"); + w.cache.touched(&key); + let stale_rows: Arc<[diesel_streamgres::row::OwnedRow]> = Arc::from(Vec::new()); + assert!( + !w.cache.fill(reservation, stale_rows), + "the fill is refused" + ); + assert_eq!(w.cache.stats.stale_fills.load(Ordering::Relaxed), 1); + assert!(w.cache.get(&key).is_none(), "the key is stale, nothing is served"); + + assert_eq!( + w.read("flag").as_deref(), + Some("on"), + "the next read sees the change" + ); +} + +/// A connection routed to Postgres never reads the cache. +#[test] +fn a_strong_connection_never_reads_the_cache() { + let mut w = World::new("strong"); + w.put("mode", "a"); + assert_eq!(w.read("mode").as_deref(), Some("a")); + w.put("mode", "b"); + + w.conn.set_consistency(Consistency::Strong); + assert_eq!(w.read("mode").as_deref(), Some("b")); + w.conn.set_consistency(Consistency::Eventual); + assert_eq!( + w.read("mode").as_deref(), + Some("a"), + "the eventual entry is untouched" + ); +} + +/// Inside a transaction reads go to Postgres and see the connection's own +/// writes. +/// +/// Right after the COMMIT the connection is in its read-your-writes window +/// and reads Postgres; the entry itself is untouched until the engine says +/// so. +#[test] +fn inside_a_transaction_reads_see_the_connections_own_writes() { + let mut w = World::new("txn"); + w.put("count", "1"); + assert_eq!(w.read("count").as_deref(), Some("1")); + let count_key = w.key("count"); + + let seen: String = w + .conn + .transaction(|c| { + diesel::update(sg_configs::table.filter(sg_configs::key.eq(&count_key))) + .set(sg_configs::value.eq("2")) + .execute(c)?; + sg_configs::table + .filter(sg_configs::key.eq(&count_key)) + .select(sg_configs::value) + .first(c) + }) + .expect("transaction"); + assert_eq!( + seen, "2", + "the read inside the transaction went to Postgres" + ); + assert_eq!( + w.engine.subscribed.lock().unwrap().len(), + 1, + "the transactional read did not subscribe" + ); + assert_eq!(w.read("count").as_deref(), Some("2")); + let cache_key = w.key_of("count"); + assert_eq!( + w.cache.get(&cache_key).map(|rows| rows.len()), + Some(1), + "the entry is still the old row" + ); + w.conn.set_read_your_writes(std::time::Duration::ZERO); + assert_eq!( + w.read("count").as_deref(), + Some("1"), + "outside the window the cache still holds the old row until a touch" + ); +} + +/// A write that returns rows is never cached. +#[test] +fn a_write_returning_rows_is_never_cached() { + let mut w = World::new("returning"); + let key = w.key("new"); + let inserted: Config = diesel::insert_into(sg_configs::table) + .values((sg_configs::key.eq(&key), sg_configs::value.eq("x"))) + .get_result(&mut w.conn) + .expect("insert returning"); + assert_eq!(inserted.value, "x"); + let updated: Config = diesel::update(sg_configs::table.filter(sg_configs::key.eq(&key))) + .set(sg_configs::value.eq("y")) + .get_result(&mut w.conn) + .expect("update returning"); + assert_eq!(updated.value, "y"); + assert!(w.cache.is_empty()); + assert!(w.engine.subscribed.lock().unwrap().is_empty()); + assert_eq!( + w.cache.stats.misses.load(Ordering::Relaxed), + 0, + "writes are not even looked up" + ); +} + +/// A query the engine refuses is read live every time. +/// +/// The engine was asked once; the refusal is remembered and the second read +/// went to Postgres without asking. +/// +/// Raw SQL against another table is refused by this engine too. +#[test] +fn a_refused_query_is_read_live_every_time() { + let mut w = World::new("refused"); + w.put("k", "v"); + w.engine.refuse_all.store(true, Ordering::SeqCst); + assert_eq!(w.read("k").as_deref(), Some("v")); + w.put("k", "w"); + assert_eq!( + w.read("k").as_deref(), + Some("w"), + "live, so the change is seen" + ); + assert_eq!(w.cache.stats.refused.load(Ordering::Relaxed), 1); + assert_eq!(w.cache.stats.refused_recently.load(Ordering::Relaxed), 1); + assert_eq!(w.engine.subscribed.lock().unwrap().len(), 0); + assert!(w.cache.is_empty()); + + let rows: Vec = sql_query("SELECT id, key, value FROM sg_configs WHERE key = $1") + .bind::(w.key("k")) + .load(&mut w.conn) + .expect("sql_query"); + assert_eq!(rows.len(), 1); + let n: i32 = diesel::select(diesel::dsl::sql::("41 + 1")) + .get_result(&mut w.conn) + .expect("select 42"); + assert_eq!(n, 42); +} + +/// Entries over the byte budget are evicted, and their subscriptions end. +#[test] +fn over_budget_entries_are_evicted_and_unsubscribed() { + let mut w = World::with_budget("evict", 600); + for i in 0..20 { + w.put(&format!("k{i}"), &"x".repeat(100)); + } + for i in 0..20 { + assert_eq!(w.read(&format!("k{i}")).unwrap().len(), 100); + } + let held = w.cache.len(); + assert!(held < 20, "the budget holds fewer than all 20, held {held}"); + assert!(w.cache.weighted_size() <= 600); + let unsubscribed = w.engine.unsubscribed.lock().unwrap().len(); + assert_eq!( + unsubscribed as u64, + 20 - held, + "every evicted entry ended its subscription" + ); + assert_eq!( + w.cache.stats.evicted.load(Ordering::Relaxed), + unsubscribed as u64 + ); +} + +/// Cached rows decode into whole structs, by name, and as counts. +#[test] +fn cached_rows_decode_whole_structs_by_name_and_count() { + let mut w = World::new("shapes"); + w.put("p", "1"); + w.put("q", "2"); + let like = format!("{}%", w.prefix); + + let whole = || { + sg_configs::table + .filter(sg_configs::key.like(&like)) + .order(sg_configs::key) + .select(Config::as_select()) + }; + let first: Vec = whole().load(&mut w.conn).expect("load"); + let second: Vec = whole().load(&mut w.conn).expect("load again"); + assert_eq!(first, second); + assert_eq!( + first.iter().map(|c| c.value.as_str()).collect::>(), + ["1", "2"] + ); + + let count = || { + sg_configs::table + .filter(sg_configs::key.like(&like)) + .count() + }; + let n: i64 = count().get_result(&mut w.conn).expect("count"); + let m: i64 = count().get_result(&mut w.conn).expect("count again"); + assert_eq!((n, m), (2, 2)); + + /// Two columns of a config, read by name. + #[derive(QueryableByName, Debug, PartialEq)] + struct Pair { + #[diesel(sql_type = Text)] + key: String, + #[diesel(sql_type = Integer)] + id: i32, + } + let by_name = || { + sql_query("SELECT key, id FROM sg_configs WHERE key LIKE $1 ORDER BY key") + .bind::(&like) + }; + let a: Vec = by_name().load(&mut w.conn).expect("by name"); + let b: Vec = by_name().load(&mut w.conn).expect("by name again"); + assert_eq!(a, b); + assert_eq!(a[0].key, w.key("p")); + assert_eq!(w.cache.stats.hits.load(Ordering::Relaxed), 3); +} + +/// Clearing the cache makes every next read refill. +#[test] +fn clearing_the_cache_refills_everything() { + let mut w = World::new("clear"); + w.put("x", "1"); + w.read("x"); + w.put("x", "2"); + w.cache.clear(); + assert_eq!(w.read("x").as_deref(), Some("2")); +} + +/// A query flagged Postgres bypasses the cache on a connection routed +/// through Streamgres. +/// +/// The flag wins over the connection: Postgres is read, the cached entry is +/// neither used nor replaced, and nothing is subscribed. +#[test] +fn a_query_flagged_postgres_bypasses_the_cache_on_a_streamgres_connection() { + let mut w = World::new("via-pg"); + w.put("mode", "a"); + assert_eq!(w.read("mode").as_deref(), Some("a")); + w.put("mode", "b"); + + let sent = w.sent(); + let subscribed = w.engine.subscribed.lock().unwrap().len(); + assert_eq!(w.read_via("mode", Route::Postgres).as_deref(), Some("b")); + assert_eq!(w.sent(), sent + 1); + assert_eq!(w.engine.subscribed.lock().unwrap().len(), subscribed); + assert_eq!( + w.read("mode").as_deref(), + Some("a"), + "the entry is untouched" + ); +} + +/// A query flagged Streamgres is cached on a connection routed to Postgres. +#[test] +fn a_query_flagged_streamgres_is_cached_on_a_postgres_connection() { + let mut w = World::new("via-sg"); + w.conn.set_route(Route::Postgres); + w.put("mode", "a"); + assert_eq!(w.read_via("mode", Route::Streamgres).as_deref(), Some("a")); + assert_eq!(w.cache.len(), 1, "subscribed and filled"); + w.put("mode", "b"); + assert_eq!( + w.read_via("mode", Route::Streamgres).as_deref(), + Some("a"), + "served from the cache" + ); + assert_eq!( + w.read("mode").as_deref(), + Some("b"), + "the connection's own route is Postgres" + ); +} + +/// The route marker is part of neither the cache key nor the SQL the engine +/// sees. +#[test] +fn the_marker_is_not_part_of_the_key_or_what_the_engine_sees() { + let mut w = World::new("via-key"); + w.put("mode", "a"); + assert_eq!(w.read_via("mode", Route::Streamgres).as_deref(), Some("a")); + let subscribed = w.engine.subscribed.lock().unwrap().clone(); + assert_eq!(subscribed.len(), 1); + assert!(!subscribed[0].sql.contains("/*"), "{}", subscribed[0].sql); + assert_eq!( + subscribed[0], + w.key_of("mode"), + "flagged and unflagged reads share the entry" + ); + let sent = w.sent(); + assert_eq!(w.read("mode").as_deref(), Some("a")); + assert_eq!( + w.sent(), + sent, + "the unflagged read hit the flagged read's entry" + ); +} + +/// A route scope covers every read inside it, and a query's own flag still +/// wins. +#[test] +fn a_route_scope_covers_every_read_inside_it_and_the_query_flag_still_wins() { + let mut w = World::new("scope"); + w.put("mode", "a"); + assert_eq!(w.read("mode").as_deref(), Some("a")); + w.put("mode", "b"); + + let World { conn, prefix, .. } = &mut w; + let key = format!("{prefix}mode"); + let (scoped, flagged) = conn.with_route(Route::Postgres, |conn| { + let scoped: Option = sg_configs::table + .filter(sg_configs::key.eq(&key)) + .select(sg_configs::value) + .first(conn) + .optional() + .unwrap(); + let flagged: Option = sg_configs::table + .filter(sg_configs::key.eq(&key)) + .select(sg_configs::value) + .via_streamgres() + .first(conn) + .optional() + .unwrap(); + (scoped, flagged) + }); + assert_eq!(scoped.as_deref(), Some("b"), "the scope routed to Postgres"); + assert_eq!( + flagged.as_deref(), + Some("a"), + "the query flag beat the scope" + ); + assert_eq!(w.read("mode").as_deref(), Some("a"), "the scope ended"); +} + +/// An entry nobody reads expires on the idle clock, and its subscription +/// ends. +/// +/// Reads keep it alive; once it has expired, the next read subscribes +/// again. +#[test] +fn an_entry_nobody_reads_expires_and_is_unsubscribed() { + let mut w = World::with_limits( + "idle", + Limits { + time_to_idle: Some(std::time::Duration::from_millis(80)), + ..Limits::default() + }, + ); + w.put("mode", "a"); + assert_eq!(w.read("mode").as_deref(), Some("a")); + for _ in 0..4 { + std::thread::sleep(std::time::Duration::from_millis(40)); + assert_eq!(w.read("mode").as_deref(), Some("a")); + } + assert_eq!(w.cache.len(), 1); + assert!(w.engine.unsubscribed.lock().unwrap().is_empty()); + + std::thread::sleep(std::time::Duration::from_millis(160)); + assert_eq!(w.cache.len(), 0, "idle past the clock"); + assert_eq!(w.engine.unsubscribed.lock().unwrap().len(), 1); + assert_eq!(w.cache.stats.expired.load(Ordering::Relaxed), 1); + + w.put("mode", "b"); + assert_eq!(w.read("mode").as_deref(), Some("b")); + assert_eq!(w.engine.subscribed.lock().unwrap().len(), 2); +} + +/// The time-to-live clock drops an entry however often it is read. +#[test] +fn time_to_live_drops_an_entry_however_often_it_is_read() { + let mut w = World::with_limits( + "ttl", + Limits { + time_to_live: Some(std::time::Duration::from_millis(100)), + ..Limits::default() + }, + ); + w.put("mode", "a"); + assert_eq!(w.read("mode").as_deref(), Some("a")); + let started = std::time::Instant::now(); + while started.elapsed() < std::time::Duration::from_millis(80) { + assert_eq!(w.read("mode").as_deref(), Some("a")); + std::thread::sleep(std::time::Duration::from_millis(10)); + } + std::thread::sleep(std::time::Duration::from_millis(60)); + w.cache.run_pending_tasks(); + assert_eq!( + w.engine.unsubscribed.lock().unwrap().len(), + 1, + "ended on the timer" + ); + w.put("mode", "b"); + assert_eq!( + w.read("mode").as_deref(), + Some("b"), + "resubscribed and refilled" + ); + assert_eq!(w.engine.subscribed.lock().unwrap().len(), 2); +} + +/// After its own write the connection reads Postgres for the read-your- +/// writes window. +/// +/// The write goes through the synced connection itself. +/// +/// The window closes; the entry was neither used nor refreshed by the +/// bypass, so without an engine touch the old rows are back. +/// +/// A transaction's writes count from its COMMIT. +#[test] +fn after_a_write_the_connection_reads_postgres_for_a_window() { + let mut w = World::new("ryw"); + w.conn + .set_read_your_writes(std::time::Duration::from_millis(150)); + w.put("mode", "a"); + assert_eq!(w.read("mode").as_deref(), Some("a")); + + let key = w.key("mode"); + diesel::update(sg_configs::table.filter(sg_configs::key.eq(&key))) + .set(sg_configs::value.eq("b")) + .execute(&mut w.conn) + .unwrap(); + let sent = w.sent(); + assert_eq!( + w.read("mode").as_deref(), + Some("b"), + "its own write is seen at once" + ); + assert_eq!(w.sent(), sent + 1, "read from Postgres"); + assert_eq!(w.cache.stats.after_write.load(Ordering::Relaxed), 1); + + std::thread::sleep(std::time::Duration::from_millis(200)); + let sent = w.sent(); + assert_eq!(w.read("mode").as_deref(), Some("a")); + assert_eq!(w.sent(), sent, "served from the cache again"); + + w.conn + .transaction::<_, diesel::result::Error, _>(|conn| { + diesel::update(sg_configs::table.filter(sg_configs::key.eq(&key))) + .set(sg_configs::value.eq("c")) + .execute(conn)?; + Ok(()) + }) + .unwrap(); + assert_eq!(w.read("mode").as_deref(), Some("c")); +} + +/// A locking read always goes to Postgres. +#[test] +fn a_locking_read_always_goes_to_postgres() { + let mut w = World::new("locking"); + w.put("row", "a"); + let key = w.key("row"); + for _ in 0..2 { + let seen: Option = sg_configs::table + .filter(sg_configs::key.eq(&key)) + .select(sg_configs::value) + .for_update() + .first(&mut w.conn) + .optional() + .unwrap(); + assert_eq!(seen.as_deref(), Some("a")); + } + assert_eq!(w.cache.stats.locking.load(Ordering::Relaxed), 2); + assert!(w.cache.is_empty(), "nothing cached"); + assert!( + w.engine.subscribed.lock().unwrap().is_empty(), + "nothing subscribed" + ); +} + +/// A route scope restores the outer route when its closure panics. +#[test] +fn a_route_scope_ends_when_the_closure_panics() { + let mut w = World::new("scope-panic"); + w.put("mode", "a"); + assert_eq!(w.read("mode").as_deref(), Some("a")); + w.put("mode", "b"); + let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + w.conn + .with_route(Route::Postgres, |_| panic!("inside the scope")); + })); + assert!(outcome.is_err()); + assert_eq!( + w.read("mode").as_deref(), + Some("a"), + "the scope was restored, the cache answers" + ); +} + +/// Opening and abandoning a transaction is not a write: it opens no read- +/// your-writes window. +#[test] +fn opening_and_abandoning_a_transaction_is_not_a_write() { + let mut w = World::new("txn-control"); + w.conn + .set_read_your_writes(std::time::Duration::from_secs(5)); + w.put("mode", "a"); + assert_eq!(w.read("mode").as_deref(), Some("a")); + w.put("mode", "b"); + w.conn.batch_execute("BEGIN").unwrap(); + w.conn.batch_execute("ROLLBACK").unwrap(); + w.conn.batch_execute("SET statement_timeout = 0").unwrap(); + w.conn.batch_execute("SELECT 1").unwrap(); + diesel::select(diesel::dsl::sql::("1")) + .execute(&mut w.conn) + .unwrap(); + assert_eq!( + w.read("mode").as_deref(), + Some("a"), + "still from the cache: a pool's checkout ping is not a write" + ); + assert_eq!(w.cache.stats.after_write.load(Ordering::Relaxed), 0); + let key = w.key("other"); + w.conn + .batch_execute(&format!( + "INSERT INTO sg_configs (key, value) VALUES ('{key}', 'x') ON CONFLICT (key) DO NOTHING" + )) + .unwrap(); + assert_eq!( + w.read("mode").as_deref(), + Some("b"), + "any other raw SQL counts as a write" + ); + assert_eq!(w.cache.stats.after_write.load(Ordering::Relaxed), 1); +} + +/// A write on one connection is seen by reads on another connection of the +/// same cache. +/// +/// The second connection stands for another of the pool's: the same +/// cache, its own socket. +#[test] +fn a_write_on_one_connection_is_seen_by_reads_on_another_of_the_same_cache() { + let mut w = World::new("ryw-pool"); + w.conn + .set_read_your_writes(std::time::Duration::from_secs(5)); + w.put("mode", "a"); + assert_eq!(w.read("mode").as_deref(), Some("a")); + + let mut sibling = SyncedConnection::with_cache( + PgConnection::establish(&url()).expect("connect"), + Arc::clone(&w.cache), + Consistency::Eventual, + ); + let key = w.key("mode"); + diesel::update(sg_configs::table.filter(sg_configs::key.eq(&key))) + .set(sg_configs::value.eq("b")) + .execute(&mut sibling) + .unwrap(); + assert_eq!( + w.read("mode").as_deref(), + Some("b"), + "the first connection reads Postgres for the window" + ); + assert_eq!(w.cache.stats.after_write.load(Ordering::Relaxed), 1); +} + +/// The window runs from a write's end, not its start: a write that takes +/// longer than the window still sends the next read to Postgres. A +/// batch that starts with a `SELECT` writes all the same. +#[test] +fn a_slow_write_opens_the_window_as_it_ends() { + let mut w = World::new("ryw-slow"); + w.conn + .set_read_your_writes(std::time::Duration::from_millis(300)); + w.put("mode", "a"); + assert_eq!(w.read("mode").as_deref(), Some("a")); + let key = w.key("mode"); + w.conn + .batch_execute(&format!( + "SELECT pg_sleep(0.5); UPDATE sg_configs SET value = 'b' WHERE key = '{key}'" + )) + .unwrap(); + assert_eq!( + w.read("mode").as_deref(), + Some("b"), + "read from Postgres: the write ended inside the window" + ); +} + +/// A read that fails after the engine accepted it ends its subscription. +#[test] +fn a_read_that_fails_after_the_engine_accepted_ends_its_subscription() { + let mut w = World::new("accepted-error"); + let failing = sql_query("SELECT no_such_column FROM sg_configs WHERE key = $1") + .bind::(w.key("x")); + let outcome: Result, _> = failing.load(&mut w.conn); + assert!(outcome.is_err()); + assert_eq!( + w.engine.subscribed.lock().unwrap().len(), + 1, + "the engine accepted it" + ); + assert_eq!( + w.engine.unsubscribed.lock().unwrap().len(), + 1, + "and was told to drop it" + ); + assert!(w.cache.is_empty()); +} + +/// A refusal is remembered for the query with the values it was for: +/// the engine refuses on values too (a time at infinity, more rows than +/// it follows), so another value is asked about afresh. +#[test] +fn a_refusal_is_remembered_for_its_values_not_the_shape() { + let mut w = World::new("refused-values"); + w.put("a", "1"); + w.put("b", "2"); + w.engine.refuse_all.store(true, Ordering::SeqCst); + assert_eq!(w.read("a").as_deref(), Some("1")); + w.engine.refuse_all.store(false, Ordering::SeqCst); + assert_eq!(w.read("a").as_deref(), Some("1")); + assert_eq!(w.read("b").as_deref(), Some("2")); + assert_eq!( + w.cache.stats.refused.load(Ordering::Relaxed), + 1, + "asked once about `a`" + ); + assert_eq!( + w.cache.stats.refused_recently.load(Ordering::Relaxed), + 1, + "`a` again skipped the ask" + ); + assert_eq!( + w.engine.subscribed.lock().unwrap().len(), + 1, + "`b` was asked about, and accepted" + ); +} + +/// What a full cache keeps is decided by how often a key is asked for, +/// not by its size or age: a key asked for again and again gets in, +/// displacing an entry read once, and a run of keys each asked for +/// once cannot displace it. (moka counts reads, misses included, once +/// the cache is past half full; the pending work is run by hand so the +/// counts are applied before each decision.) +#[test] +fn a_hot_entry_outlives_cold_ones_under_the_budget() { + let mut w = World::with_budget("hot", 2000); + for i in 0..8 { + w.put(&format!("p{i}"), &"x".repeat(100)); + assert_eq!(w.read(&format!("p{i}")).unwrap().len(), 100); + w.cache.run_pending_tasks(); + } + let full = w.cache.len(); + assert!( + (2..8).contains(&full), + "the budget holds some, not all: {full}" + ); + assert!(w.cache.weighted_size() > 1000, "past half full"); + + w.put("hot", &"h".repeat(100)); + let mut admitted_after = None; + for read in 1..=8 { + assert_eq!(w.read("hot").unwrap().len(), 100); + w.cache.run_pending_tasks(); + if w.conn.last_served() == Some(Served::Cache) { + admitted_after = Some(read); + break; + } + } + let admitted_after = admitted_after.expect("a key asked for repeatedly gets in"); + assert!( + admitted_after > 1, + "not on its first miss, the cache being full" + ); + + for i in 0..20 { + w.put(&format!("c{i}"), &"x".repeat(100)); + assert_eq!(w.read(&format!("c{i}")).unwrap().len(), 100); + w.cache.run_pending_tasks(); + } + assert!(w.cache.len() <= full + 1, "the budget still holds"); + let hits = w.cache.stats.hits.load(Ordering::Relaxed); + assert_eq!(w.read("hot").unwrap().len(), 100); + assert_eq!( + w.conn.last_served(), + Some(Served::Cache), + "the hot entry outlived the cold ones; held {}", + w.cache.len() + ); + assert_eq!(w.cache.stats.hits.load(Ordering::Relaxed), hits + 1); +} diff --git a/crates/diesel-streamgres/tests/engine_latency.rs b/crates/diesel-streamgres/tests/engine_latency.rs new file mode 100644 index 0000000..d0d9e0b --- /dev/null +++ b/crates/diesel-streamgres/tests/engine_latency.rs @@ -0,0 +1,141 @@ +//! Commit-to-cache latency against a running engine, on any table. +//! +//! Flips one text column of one existing row back and forth from a +//! separate connection and measures how long a cached read of that row +//! takes to show each new value: that is the engine's propagation (WAL +//! decode, routing, the touch over the socket) plus one Postgres refill. +//! +//! Driven by environment so it runs against a real schema: +//! - `STREAMGRES_S2S_URL`, `DATABASE_URL` +//! - `STREAMGRES_LATENCY_TABLE`, `STREAMGRES_LATENCY_KEY_COLUMN`, +//! `STREAMGRES_LATENCY_VALUE_COLUMN`, `STREAMGRES_LATENCY_ROW_KEY` +//! (an existing row; its value column is overwritten and restored) +//! - `STREAMGRES_LATENCY_SAMPLES` (default 25) +//! +//! Skips itself when any is missing. + +use std::sync::Arc; +use std::sync::atomic::Ordering; +use std::time::{Duration, Instant}; + +use diesel::prelude::*; +use diesel::sql_query; +use diesel::sql_types::Text; +use diesel_streamgres::cache::Cache; +use diesel_streamgres::engine::EngineClient; +use diesel_streamgres::{Consistency, SyncedConnection}; + +/// The value column, read by name. +#[derive(QueryableByName)] +struct Val { + #[diesel(sql_type = Text)] + v: String, +} + +/// A variable from the environment, unless empty. +fn env(name: &str) -> Option { + std::env::var(name).ok().filter(|value| !value.is_empty()) +} + +/// How long a commit takes to reach a refilled read, against a live engine +/// and a table the environment names; skipped unless it is named. +/// The row is put back afterwards. +#[test] +fn commit_to_refilled_read_latency() { + let (Some(url), Some(database), Some(table), Some(key_column), Some(value_column), Some(row_key)) = ( + env("STREAMGRES_S2S_URL"), + env("DATABASE_URL"), + env("STREAMGRES_LATENCY_TABLE"), + env("STREAMGRES_LATENCY_KEY_COLUMN"), + env("STREAMGRES_LATENCY_VALUE_COLUMN"), + env("STREAMGRES_LATENCY_ROW_KEY"), + ) else { + eprintln!("latency env not set; skipping"); + return; + }; + let samples: usize = env("STREAMGRES_LATENCY_SAMPLES") + .and_then(|value| value.parse().ok()) + .unwrap_or(25); + + let engine = EngineClient::connect(url); + let cache = Arc::new(Cache::new(64 * 1024 * 1024, Arc::new(engine.clone()))); + engine.attach(&cache); + let deadline = Instant::now() + Duration::from_secs(10); + while !engine.is_connected() { + assert!(Instant::now() < deadline, "the engine did not connect"); + std::thread::sleep(Duration::from_millis(50)); + } + let mut writer = PgConnection::establish(&database).expect("connect"); + let mut conn = SyncedConnection::with_cache( + PgConnection::establish(&database).expect("connect"), + cache.clone(), + Consistency::Eventual, + ); + + let select = format!( + "SELECT \"{value_column}\" AS v FROM \"{table}\" WHERE \"{key_column}\" = $1" + ); + let update = format!( + "UPDATE \"{table}\" SET \"{value_column}\" = $1 WHERE \"{key_column}\" = $2" + ); + let read = |conn: &mut SyncedConnection| -> String { + sql_query(&select) + .bind::(&row_key) + .load::(conn) + .expect("read") + .pop() + .expect("the row exists") + .v + }; + let original = read(&mut conn); + assert_eq!(engine.subscribed(), 1, "the engine accepted the read"); + assert_eq!(read(&mut conn), original); + assert_eq!(cache.stats.hits.load(Ordering::Relaxed), 1, "the second read hit"); + + let mut took = Vec::with_capacity(samples); + for sample in 0..samples { + let next = format!("sg-latency-{sample}"); + let started = Instant::now(); + sql_query(&update) + .bind::(&next) + .bind::(&row_key) + .execute(&mut writer) + .expect("update"); + let committed = Instant::now(); + loop { + if read(&mut conn) == next { + break; + } + assert!( + started.elapsed() < Duration::from_secs(10), + "sample {sample} never arrived" + ); + std::thread::sleep(Duration::from_micros(200)); + } + took.push((committed.elapsed(), started.elapsed())); + } + sql_query(&update) + .bind::(&original) + .bind::(&row_key) + .execute(&mut writer) + .expect("restore"); + + let mut after_commit: Vec = took.iter().map(|(a, _)| *a).collect(); + let mut with_write: Vec = took.iter().map(|(_, b)| *b).collect(); + after_commit.sort(); + with_write.sort(); + let at = |sorted: &[Duration], q: f64| sorted[((sorted.len() - 1) as f64 * q).round() as usize]; + eprintln!( + "LATENCY samples={samples} table={table} commit->refilled read: p50 {:?} p90 {:?} max {:?}; including the write: p50 {:?} p90 {:?} max {:?}", + at(&after_commit, 0.5), + at(&after_commit, 0.9), + after_commit[after_commit.len() - 1], + at(&with_write, 0.5), + at(&with_write, 0.9), + with_write[with_write.len() - 1], + ); + let changes = cache.stats.touched.load(Ordering::Relaxed) + cache.stats.replaced.load(Ordering::Relaxed); + eprintln!("LATENCY mode: {}", if engine.materialized() == 1 { "rows (no Postgres read per change)" } else { "refill" }); + assert!(changes as usize >= samples); + assert_eq!(engine.subscribed(), 1); +} diff --git a/crates/diesel-streamgres/tests/engine_live.rs b/crates/diesel-streamgres/tests/engine_live.rs new file mode 100644 index 0000000..ddf9c40 --- /dev/null +++ b/crates/diesel-streamgres/tests/engine_live.rs @@ -0,0 +1,1219 @@ +//! The whole path against a running engine: a read is subscribed, served +//! from the cache, and refilled after a commit the engine reports. +//! +//! Needs `STREAMGRES_S2S_URL` (`ws://host:port//s2s/v1`) and +//! `DATABASE_URL` pointing at the database that engine streams. Skips +//! itself otherwise. + +use std::sync::Arc; +use std::sync::atomic::Ordering; +use std::time::{Duration, Instant}; + +use diesel::connection::SimpleConnection; +use diesel::prelude::*; +use diesel::sql_types::Text; +use diesel_streamgres::cache::{Cache, DATABASE_SQL, Database}; +use diesel_streamgres::engine::EngineClient; +use diesel_streamgres::subscriptions::Subscriptions; +use diesel_streamgres::{Consistency, Served, SyncedConnection}; + +diesel::table! { + sg_configs (id) { + id -> Int4, + key -> Text, + value -> Text, + } +} + +/// A config row. +#[derive(Queryable, Selectable, Debug, PartialEq, Clone)] +#[diesel(table_name = sg_configs)] +struct Config { + id: i32, + key: String, + value: String, +} + +/// One test against a live engine: the wrapped connection and its +/// cache, the engine client, a plain connection for changes, and the +/// prefix of its keys. +struct Live { + conn: SyncedConnection, + cache: Arc, + engine: EngineClient, + other: PgConnection, + prefix: String, +} + +impl Live { + /// Connects to the engine `STREAMGRES_S2S_URL` names; `None`, and the + /// test skipped, when it names none. The cache is told which database + /// it reads up front, so the statements a test counts are its own. + fn start(name: &str) -> Option { + let Ok(url) = std::env::var("STREAMGRES_S2S_URL") else { + eprintln!("STREAMGRES_S2S_URL not set; skipping"); + return None; + }; + let database = std::env::var("DATABASE_URL").expect("DATABASE_URL"); + let engine = EngineClient::connect(url); + let cache = Arc::new(Cache::new(64 * 1024 * 1024, Arc::new(engine.clone()))); + engine.attach(&cache); + let deadline = Instant::now() + Duration::from_secs(5); + while !engine.is_connected() { + assert!(Instant::now() < deadline, "the engine did not connect"); + std::thread::sleep(Duration::from_millis(50)); + } + let mut other = PgConnection::establish(&database).expect("connect"); + other + .batch_execute( + "CREATE TABLE IF NOT EXISTS sg_configs ( + id serial PRIMARY KEY, key text NOT NULL UNIQUE, value text NOT NULL)", + ) + .expect("schema"); + let identity = diesel::sql_query(DATABASE_SQL) + .load::(&mut other) + .expect("the database's identity"); + cache.learn_database(identity.into_iter().next()); + let prefix = format!("{name}-{}-", std::process::id()); + let inner = PgConnection::establish(&database).expect("connect"); + let conn = SyncedConnection::with_cache(inner, cache.clone(), Consistency::Eventual); + Some(Live { + conn, + cache, + engine, + other, + prefix, + }) + } + + /// A key of this test's. + fn key(&self, name: &str) -> String { + format!("{}{name}", self.prefix) + } + + /// Sets a key on the plain connection, as another service would. + fn put(&mut self, name: &str, value: &str) { + let key = self.key(name); + diesel::insert_into(sg_configs::table) + .values((sg_configs::key.eq(&key), sg_configs::value.eq(value))) + .on_conflict(sg_configs::key) + .do_update() + .set(sg_configs::value.eq(value)) + .execute(&mut self.other) + .expect("put"); + } + + /// A key's value through the wrapper. + fn read(&mut self, name: &str) -> Option { + let key = self.key(name); + sg_configs::table + .filter(sg_configs::key.eq(key)) + .select(sg_configs::value) + .first(&mut self.conn) + .optional() + .expect("read") + } + + /// Reads until the value is seen or the wait runs out, returning how + /// long it took. + fn read_until(&mut self, name: &str, want: &str, wait: Duration) -> Option { + let started = Instant::now(); + while started.elapsed() < wait { + if self.read(name).as_deref() == Some(want) { + return Some(started.elapsed()); + } + std::thread::sleep(Duration::from_millis(20)); + } + None + } +} + +/// Changes the engine reported: touches (refill mode) plus replacements +/// (rows mode). +fn changes(cache: &Cache) -> u64 { + cache.stats.touched.load(Ordering::Relaxed) + cache.stats.replaced.load(Ordering::Relaxed) +} + +impl Drop for Live { + /// Deletes this test's keys. + fn drop(&mut self) { + let _ = diesel::delete( + sg_configs::table.filter(sg_configs::key.like(format!("{}%", self.prefix))), + ) + .execute(&mut self.other); + } +} + +/// A commit the engine reports makes the next read current: an update, +/// a delete, and a row appearing under the key the query asks for. +#[test] +fn a_commit_the_engine_reports_refills_the_cache() { + let Some(mut live) = Live::start("refill") else { + return; + }; + live.put("rate", "1"); + + assert_eq!( + live.read("rate").as_deref(), + Some("1"), + "first read, from Postgres" + ); + assert_eq!(live.engine.subscribed(), 1, "the engine accepted the query"); + assert_eq!( + live.read("rate").as_deref(), + Some("1"), + "second read, from the cache" + ); + assert_eq!(live.cache.stats.hits.load(Ordering::Relaxed), 1); + + live.put("rate", "2"); + let took = live + .read_until("rate", "2", Duration::from_secs(10)) + .expect("the change reached the cache through the engine"); + eprintln!("commit to refilled read: {took:?}"); + assert!(changes(&live.cache) >= 1); + assert_eq!( + live.engine.subscribed(), + 1, + "a change keeps the subscription" + ); + + diesel::delete(sg_configs::table.filter(sg_configs::key.eq(live.key("rate")))) + .execute(&mut live.other) + .expect("delete"); + let started = Instant::now(); + while live.read("rate").is_some() { + assert!( + started.elapsed() < Duration::from_secs(10), + "the delete never reached the cache" + ); + std::thread::sleep(Duration::from_millis(20)); + } + + live.put("rate", "3"); + live.read_until("rate", "3", Duration::from_secs(10)) + .expect("the insert reached the cache"); +} + +/// A commit to a row the query does not read leaves the entry alone. +#[test] +fn an_unrelated_commit_does_not_touch_the_entry() { + let Some(mut live) = Live::start("unrelated") else { + return; + }; + live.put("a", "1"); + live.put("b", "1"); + assert_eq!(live.read("a").as_deref(), Some("1")); + let touched_before = live.cache.stats.touched.load(Ordering::Relaxed); + + live.put("b", "2"); + std::thread::sleep(Duration::from_millis(500)); + assert_eq!(live.read("a").as_deref(), Some("1")); + assert_eq!( + live.cache.stats.touched.load(Ordering::Relaxed), + touched_before, + "a change to another row's key is not a touch of this query" + ); + assert_eq!(live.cache.stats.hits.load(Ordering::Relaxed), 1); +} + +/// A query outside the engine's subset is refused and served live. +/// +/// Raw SQL in the subset is accepted. +#[test] +fn a_query_outside_the_engines_subset_is_served_live() { + let Some(mut live) = Live::start("live") else { + return; + }; + live.put("x", "1"); + let like = format!("{}%", live.prefix); + let read = |live: &mut Live| -> Vec { + sg_configs::table + .filter(sg_configs::key.like(&like)) + .select(Config::as_select()) + .load(&mut live.conn) + .expect("like") + }; + assert_eq!(read(&mut live).len(), 1); + assert_eq!( + live.cache.stats.refused.load(Ordering::Relaxed), + 1, + "LIKE is refused" + ); + live.put("y", "1"); + assert_eq!( + read(&mut live).len(), + 2, + "live, so the new row is seen at once" + ); + assert_eq!(live.engine.subscribed(), 0); + + /// The value column, read by name. + #[derive(QueryableByName, Debug, PartialEq)] + struct Val { + #[diesel(sql_type = Text)] + value: String, + } + let raw = |live: &mut Live| -> Vec { + diesel::sql_query("SELECT value FROM sg_configs WHERE key = $1") + .bind::(live.key("x")) + .load(&mut live.conn) + .expect("raw") + }; + assert_eq!(raw(&mut live)[0].value, "1"); + assert_eq!(live.engine.subscribed(), 1); + live.put("x", "2"); + let started = Instant::now(); + while raw(&mut live)[0].value != "2" { + assert!(started.elapsed() < Duration::from_secs(10)); + std::thread::sleep(Duration::from_millis(20)); + } +} + +/// Evicting an entry ends its subscription with the engine. +/// +/// A tiny budget: the second entry pushes the first out. +#[test] +fn evicting_an_entry_unsubscribes_it_from_the_engine() { + let Some(mut live) = Live::start("evict") else { + return; + }; + let engine = live.engine.clone(); + let small = Arc::new(Cache::new(400, Arc::new(engine.clone()))); + engine.attach(&small); + let inner = live.conn.inner(); + let database = std::env::var("DATABASE_URL").unwrap(); + let _ = inner; + let mut conn = SyncedConnection::with_cache( + PgConnection::establish(&database).unwrap(), + small.clone(), + Consistency::Eventual, + ); + for i in 0..6 { + live.put(&format!("e{i}"), &"v".repeat(120)); + } + for i in 0..6 { + let key = live.key(&format!("e{i}")); + let _: Option = sg_configs::table + .filter(sg_configs::key.eq(key)) + .select(sg_configs::value) + .first(&mut conn) + .optional() + .unwrap(); + } + let held = small.len() as usize; + assert!(held < 6, "held {held}"); + let deadline = Instant::now() + Duration::from_secs(2); + while engine.subscribed() != held { + assert!( + Instant::now() < deadline, + "subscribed {} held {held}", + engine.subscribed() + ); + std::thread::sleep(Duration::from_millis(20)); + } +} + +/// Needs `STREAMGRES_RESTART_CMD`, a shell command that restarts the +/// engine and returns once it serves again. +/// +/// While the engine was away, the socket dropped and the cache was cleared; +/// a change made meanwhile must be seen once reads resume. +/// +/// The change made while the socket was down reaches the cache: in rows +/// mode the entry is refilled from the engine's snapshot and then its +/// deltas, so the read is eventual by the engine's propagation. +/// +/// And the resubscription is live: a further change is reported. +#[test] +fn an_engine_restart_is_survived_without_a_stale_read() { + let Some(mut live) = Live::start("restart") else { + return; + }; + let Ok(restart) = std::env::var("STREAMGRES_RESTART_CMD") else { + eprintln!("STREAMGRES_RESTART_CMD not set; skipping"); + return; + }; + live.put("v", "1"); + assert_eq!(live.read("v").as_deref(), Some("1")); + assert_eq!(live.engine.subscribed(), 1); + + let status = std::process::Command::new("sh") + .arg("-c") + .arg(&restart) + .status() + .expect("restart command"); + assert!(status.success(), "the restart command failed"); + + live.put("v", "2"); + let deadline = Instant::now() + Duration::from_secs(10); + while !live.engine.is_connected() { + assert!(Instant::now() < deadline, "the client did not reconnect"); + std::thread::sleep(Duration::from_millis(50)); + } + let took = live + .read_until("v", "2", Duration::from_secs(10)) + .expect("the change made during the restart was seen"); + eprintln!("after restart, change seen within: {took:?}"); + assert_eq!( + live.engine.subscribed(), + 1, + "resubscribed under the same key" + ); + assert_eq!(live.cache.len(), 1, "the entry is back"); + let changes_before = changes(&live.cache); + let hits_before = live.cache.stats.hits.load(Ordering::Relaxed); + assert_eq!(live.read("v").as_deref(), Some("2")); + assert_eq!( + live.cache.stats.hits.load(Ordering::Relaxed), + hits_before + 1, + "served from the cache" + ); + + live.put("v", "3"); + let took = live + .read_until("v", "3", Duration::from_secs(10)) + .expect("the change reached the cache through the resubscribed engine"); + eprintln!("after restart, commit to refilled read: {took:?}"); + assert!( + changes(&live.cache) > changes_before, + "the refill came from the engine, not from a missing entry" + ); +} + +/// Counts the queries the connection starts. +struct CountQueries(Arc); + +impl diesel::connection::Instrumentation for CountQueries { + /// Counts a query start. + fn on_connection_event(&mut self, event: diesel::connection::InstrumentationEvent<'_>) { + if let diesel::connection::InstrumentationEvent::StartQuery { .. } = event { + self.0.fetch_add(1, Ordering::SeqCst); + } + } +} + +/// In rows mode Postgres is read only for the table's column types, once. +/// +/// A second query on the same table reuses the types: no lookup. +#[test] +fn rows_mode_reads_postgres_only_for_the_table_types() { + let Some(mut live) = Live::start("rows") else { + return; + }; + let sent = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + live.conn + .inner() + .set_instrumentation(CountQueries(sent.clone())); + live.put("r", "1"); + + assert_eq!(live.read("r").as_deref(), Some("1")); + assert_eq!( + live.engine.materialized(), + 1, + "served from the engine's rows" + ); + assert_eq!(live.conn.last_served(), Some(Served::Engine)); + assert_eq!( + sent.load(Ordering::SeqCst), + 1, + "the only statement Postgres saw was the column-type lookup, not the rows" + ); + assert_eq!(live.read("r").as_deref(), Some("1")); + assert_eq!(live.conn.last_served(), Some(Served::Cache)); + + live.put("r", "2"); + let took = live + .read_until("r", "2", Duration::from_secs(10)) + .expect("the change arrived as rows"); + eprintln!("ROWS MODE commit to replaced entry: {took:?}"); + assert_eq!( + sent.load(Ordering::SeqCst), + 1, + "no Postgres read for the change" + ); + assert!( + live.cache.stats.replaced.load(Ordering::Relaxed) >= 1, + "the change replaced the entry" + ); + assert_eq!(live.cache.stats.touched.load(Ordering::Relaxed), 0); + + live.put("r2", "x"); + assert_eq!(live.read("r2").as_deref(), Some("x")); + assert_eq!(sent.load(Ordering::SeqCst), 1); + assert_eq!(live.engine.materialized(), 2); +} + +/// Ordered, limited and counted results served from the engine's rows match +/// Postgres. +/// +/// A change that reorders the window. +/// +/// A count. +#[test] +fn ordered_limited_and_counted_results_match_postgres() { + let Some(mut live) = Live::start("order") else { + return; + }; + for (name, value) in [("k1", "m"), ("k2", "z"), ("k3", "a"), ("k4", "q")] { + live.put(name, value); + } + let keys: Vec = ["k1", "k2", "k3", "k4"] + .iter() + .map(|k| live.key(k)) + .collect(); + let database = std::env::var("DATABASE_URL").unwrap(); + let mut plain = PgConnection::establish(&database).unwrap(); + + let top = |conn: &mut SyncedConnection, keys: &[String]| -> Vec { + sg_configs::table + .filter(sg_configs::key.eq_any(keys)) + .order(sg_configs::value.desc()) + .limit(2) + .select(Config::as_select()) + .load(conn) + .expect("top") + }; + let expected: Vec = sg_configs::table + .filter(sg_configs::key.eq_any(&keys)) + .order(sg_configs::value.desc()) + .limit(2) + .select(Config::as_select()) + .load(&mut plain) + .unwrap(); + assert_eq!(top(&mut live.conn, &keys), expected); + assert_eq!(live.engine.materialized(), 1); + assert_eq!( + expected + .iter() + .map(|c| c.value.as_str()) + .collect::>(), + ["z", "q"] + ); + + live.put("k3", "zz"); + let started = Instant::now(); + loop { + let now = top(&mut live.conn, &keys); + if now.iter().map(|c| c.value.as_str()).collect::>() == ["zz", "z"] { + break; + } + assert!( + started.elapsed() < Duration::from_secs(10), + "window never reordered: {now:?}" + ); + std::thread::sleep(Duration::from_millis(10)); + } + + let count = |conn: &mut SyncedConnection| -> i64 { + sg_configs::table + .filter(sg_configs::key.eq_any(&keys)) + .count() + .get_result(conn) + .expect("count") + }; + assert_eq!(count(&mut live.conn), 4); + assert_eq!(live.engine.materialized(), 2); + diesel::delete(sg_configs::table.filter(sg_configs::key.eq(live.key("k1")))) + .execute(&mut live.other) + .unwrap(); + let started = Instant::now(); + while count(&mut live.conn) != 3 { + assert!( + started.elapsed() < Duration::from_secs(10), + "count never dropped" + ); + std::thread::sleep(Duration::from_millis(10)); + } +} + +diesel::table! { + sg_items (id) { + id -> Int4, + name -> Text, + qty -> Int4, + note -> Nullable, + tags -> Array, + payload -> Jsonb, + created -> Timestamptz, + lead -> Nullable, + } +} + +/// Microsecond timestamps are served exactly; an interval column falls the +/// query back to refills. +/// +/// The engine learns a new column from the DDL trigger; give it a moment. +/// +/// This row and any left by an earlier run that died before its own +/// cleanup: the passthrough and owned-row suites read the whole table. +/// +/// `created` is timestamptz(6): served from engine rows, to the +/// microsecond. +/// +/// An interval has no encoding here: the read goes to Postgres and refills +/// on touches. +#[test] +fn microsecond_timestamps_are_exact_and_an_interval_falls_back() { + let Some(mut live) = Live::start("fallback") else { + return; + }; + live.other + .batch_execute( + "CREATE TABLE IF NOT EXISTS sg_items ( + id serial PRIMARY KEY, name text NOT NULL, qty int NOT NULL, note text, + tags text[] NOT NULL DEFAULT '{}', payload jsonb NOT NULL DEFAULT '{}', + created timestamptz NOT NULL DEFAULT now()); + ALTER TABLE sg_items ADD COLUMN IF NOT EXISTS lead interval", + ) + .unwrap(); + std::thread::sleep(Duration::from_millis(500)); + let name = live.key("item"); + diesel::delete(sg_items::table.filter(sg_items::name.like("fallback-%"))) + .execute(&mut live.other) + .unwrap(); + diesel::insert_into(sg_items::table) + .values(( + sg_items::name.eq(&name), + sg_items::qty.eq(1), + sg_items::tags.eq(vec!["a", "b"]), + sg_items::payload.eq(serde_json::json!({"k": [1, 2]})), + sg_items::lead.eq(diesel::data_types::PgInterval::from_microseconds( + 3_600_000_000, + )), + )) + .execute(&mut live.other) + .unwrap(); + let sent = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + live.conn + .inner() + .set_instrumentation(CountQueries(sent.clone())); + + let with_created = + |conn: &mut SyncedConnection| -> (i32, chrono::DateTime) { + sg_items::table + .filter(sg_items::name.eq(&name)) + .select((sg_items::qty, sg_items::created)) + .first(conn) + .expect("with created") + }; + let first = with_created(&mut live.conn); + let from_postgres: (i32, chrono::DateTime) = sg_items::table + .filter(sg_items::name.eq(&name)) + .select((sg_items::qty, sg_items::created)) + .first(&mut live.other) + .unwrap(); + assert_eq!(first, from_postgres, "exact to the microsecond"); + assert_eq!( + live.engine.materialized(), + 1, + "rows mode: {:?}", + live.engine.fallbacks() + ); + assert_eq!(sent.load(Ordering::SeqCst), 1, "the type lookup only"); + assert_eq!(with_created(&mut live.conn), first, "and cached"); + assert_eq!(sent.load(Ordering::SeqCst), 1); + + let with_lead = |conn: &mut SyncedConnection| -> (i32, Option) { + sg_items::table + .filter(sg_items::name.eq(&name)) + .select((sg_items::qty, sg_items::lead)) + .first(conn) + .expect("with lead") + }; + let lead = with_lead(&mut live.conn); + assert_eq!(lead.1.map(|i| i.microseconds), Some(3_600_000_000)); + assert_eq!(live.engine.materialized(), 1, "fell back to refills"); + assert!( + live.engine + .fallbacks() + .last() + .is_some_and(|why| why.contains("oid 1186")), + "{:?}", + live.engine.fallbacks() + ); + assert_eq!(live.engine.subscribed(), 2); + assert_eq!(sent.load(Ordering::SeqCst), 2, "the row read"); + assert_eq!(with_lead(&mut live.conn), lead, "and cached"); + assert!( + sent.load(Ordering::SeqCst) <= 3, + "at most one redundant refill" + ); + let settled = sent.load(Ordering::SeqCst); + + diesel::update(sg_items::table.filter(sg_items::name.eq(&name))) + .set((sg_items::qty.eq(2), sg_items::tags.eq(vec!["c"]))) + .execute(&mut live.other) + .unwrap(); + let started = Instant::now(); + loop { + let a = with_created(&mut live.conn).0; + let b = with_lead(&mut live.conn).0; + if a == 2 && b == 2 { + break; + } + assert!( + started.elapsed() < Duration::from_secs(10), + "the update never arrived: {a} {b}" + ); + std::thread::sleep(Duration::from_millis(10)); + } + assert_eq!( + sent.load(Ordering::SeqCst), + settled + 1, + "one refill read for the fallback query, none for rows mode" + ); + diesel::delete(sg_items::table.filter(sg_items::name.eq(&name))) + .execute(&mut live.other) + .unwrap(); +} + +diesel::table! { + sg_authors (id) { + id -> Int4, + name -> Text, + } +} + +diesel::table! { + sg_books (id) { + id -> Int4, + author_id -> Int4, + title -> Text, + } +} + +diesel::joinable!(sg_books -> sg_authors (author_id)); +diesel::allow_tables_to_appear_in_same_query!(sg_authors, sg_books); + +/// Joins are composed from the engine's rows and stay current as either +/// side changes. +/// +/// Clean, then seed: ann has two books, bob none, cy one. +/// +/// authors.inner_join(books), filtered by the authors in play, ordered. +/// +/// A book for bob appears; one of ann's goes; cy is renamed. +/// +/// The cached query still filters on the old names, so the renamed author +/// leaves it, as it leaves the same query on Postgres. +/// +/// A left join shows authors without books, and a count counts joined rows. +#[test] +fn joins_are_composed_from_the_engines_rows() { + let Some(mut live) = Live::start("join") else { + return; + }; + live.other + .batch_execute( + "CREATE TABLE IF NOT EXISTS sg_authors (id serial PRIMARY KEY, name text NOT NULL); + CREATE TABLE IF NOT EXISTS sg_books (id serial PRIMARY KEY, author_id int NOT NULL REFERENCES sg_authors(id), title text NOT NULL)", + ) + .unwrap(); + let database = std::env::var("DATABASE_URL").unwrap(); + let mut plain = PgConnection::establish(&database).unwrap(); + let p = live.prefix.clone(); + let names: Vec = ["ann", "bob", "cy"] + .iter() + .map(|n| format!("{p}{n}")) + .collect(); + diesel::delete(sg_books::table.filter(sg_books::title.like(format!("{p}%")))) + .execute(&mut plain) + .unwrap(); + diesel::delete(sg_authors::table.filter(sg_authors::name.like(format!("{p}%")))) + .execute(&mut plain) + .unwrap(); + let ids: Vec = diesel::insert_into(sg_authors::table) + .values( + names + .iter() + .map(|n| sg_authors::name.eq(n)) + .collect::>(), + ) + .returning(sg_authors::id) + .get_results(&mut plain) + .unwrap(); + diesel::insert_into(sg_books::table) + .values(vec![ + ( + sg_books::author_id.eq(ids[0]), + sg_books::title.eq(format!("{p}a1")), + ), + ( + sg_books::author_id.eq(ids[0]), + sg_books::title.eq(format!("{p}a2")), + ), + ( + sg_books::author_id.eq(ids[2]), + sg_books::title.eq(format!("{p}c1")), + ), + ]) + .execute(&mut plain) + .unwrap(); + let sent = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + live.conn + .inner() + .set_instrumentation(CountQueries(sent.clone())); + + type Pair = (String, String); + let inner = |conn: &mut PgConnection, names: &[String]| -> Vec { + sg_authors::table + .inner_join(sg_books::table) + .filter(sg_authors::name.eq_any(names)) + .order(sg_authors::id) + .select((sg_authors::name, sg_books::title)) + .load(conn) + .unwrap() + }; + let through = |conn: &mut SyncedConnection, names: &[String]| -> Vec { + sg_authors::table + .inner_join(sg_books::table) + .filter(sg_authors::name.eq_any(names)) + .order(sg_authors::id) + .select((sg_authors::name, sg_books::title)) + .load(conn) + .unwrap() + }; + let expected = inner(&mut plain, &names); + assert_eq!(expected.len(), 3); + assert_eq!(through(&mut live.conn, &names), expected); + assert_eq!( + live.engine.materialized(), + 1, + "{:?}", + live.engine.fallbacks() + ); + assert_eq!( + sent.load(Ordering::SeqCst), + 2, + "two type lookups, no row read" + ); + assert_eq!(through(&mut live.conn, &names), expected); + assert_eq!(sent.load(Ordering::SeqCst), 2); + + diesel::insert_into(sg_books::table) + .values(( + sg_books::author_id.eq(ids[1]), + sg_books::title.eq(format!("{p}b1")), + )) + .execute(&mut plain) + .unwrap(); + diesel::delete(sg_books::table.filter(sg_books::title.eq(format!("{p}a2")))) + .execute(&mut plain) + .unwrap(); + diesel::update(sg_authors::table.filter(sg_authors::id.eq(ids[2]))) + .set(sg_authors::name.eq(format!("{p}cyrus"))) + .execute(&mut plain) + .unwrap(); + let expected = inner(&mut plain, &names); + assert_eq!(expected.len(), 2); + let started = Instant::now(); + loop { + let now = through(&mut live.conn, &names); + if now == expected { + break; + } + assert!( + started.elapsed() < Duration::from_secs(10), + "join never caught up: {now:?} vs {expected:?}" + ); + std::thread::sleep(Duration::from_millis(10)); + } + assert_eq!(sent.load(Ordering::SeqCst), 2, "still no Postgres read"); + + let mut names_now = names.clone(); + names_now[2] = format!("{p}cyrus"); + type Left = (String, Option); + let left: Vec = sg_authors::table + .left_join(sg_books::table) + .filter(sg_authors::name.eq_any(&names_now)) + .order(sg_authors::id) + .select((sg_authors::name, sg_books::title.nullable())) + .load(&mut live.conn) + .unwrap(); + let left_expected: Vec = sg_authors::table + .left_join(sg_books::table) + .filter(sg_authors::name.eq_any(&names_now)) + .order(sg_authors::id) + .select((sg_authors::name, sg_books::title.nullable())) + .load(&mut plain) + .unwrap(); + assert_eq!(left, left_expected); + let count: i64 = sg_authors::table + .inner_join(sg_books::table) + .filter(sg_authors::name.eq_any(&names_now)) + .count() + .get_result(&mut live.conn) + .unwrap(); + assert_eq!(count, 3); + assert_eq!( + live.engine.materialized(), + 3, + "{:?}", + live.engine.fallbacks() + ); + + diesel::delete(sg_books::table.filter(sg_books::title.like(format!("{p}%")))) + .execute(&mut plain) + .unwrap(); + diesel::delete(sg_authors::table.filter(sg_authors::name.like(format!("{p}%")))) + .execute(&mut plain) + .unwrap(); +} + +/// A pool of sockets spreads the subscriptions across them, and a key +/// stays with the socket that first took it. +/// +/// A change on a key held by any socket reaches the cache. +/// +/// Dropping one subscription frees its slot on its socket only. +#[test] +fn a_pool_spreads_subscriptions_across_its_sockets() { + let Some(mut live) = Live::start("pool") else { + return; + }; + let url = std::env::var("STREAMGRES_S2S_URL").unwrap(); + let database = std::env::var("DATABASE_URL").unwrap(); + let engine = EngineClient::with_config( + url, + diesel_streamgres::engine::Config { + connections: 3, + ..diesel_streamgres::engine::Config::from_env() + }, + ); + let cache = Arc::new(Cache::new(64 * 1024 * 1024, Arc::new(engine.clone()))); + engine.attach(&cache); + let deadline = Instant::now() + Duration::from_secs(5); + while engine.connections() < 3 { + assert!( + Instant::now() < deadline, + "only {} sockets connected", + engine.connections() + ); + std::thread::sleep(Duration::from_millis(50)); + } + let mut conn = SyncedConnection::with_cache( + PgConnection::establish(&database).unwrap(), + cache.clone(), + Consistency::Eventual, + ); + for i in 0..6 { + live.put(&format!("p{i}"), &format!("v{i}")); + } + let read = |conn: &mut SyncedConnection, key: String| -> Option { + sg_configs::table + .filter(sg_configs::key.eq(key)) + .select(sg_configs::value) + .first(conn) + .optional() + .unwrap() + }; + for i in 0..6 { + assert_eq!( + read(&mut conn, live.key(&format!("p{i}"))).as_deref(), + Some(format!("v{i}").as_str()) + ); + } + assert_eq!(engine.subscribed(), 6); + assert_eq!(engine.per_worker(), vec![2, 2, 2], "least-loaded routing"); + assert_eq!(engine.materialized(), 6, "{:?}", engine.fallbacks()); + + let before = changes(&cache); + live.put("p4", "v4b"); + let started = Instant::now(); + loop { + if read(&mut conn, live.key("p4")).as_deref() == Some("v4b") { + break; + } + assert!( + started.elapsed() < Duration::from_secs(10), + "the change on a pooled key was not seen" + ); + std::thread::sleep(Duration::from_millis(20)); + } + assert!(changes(&cache) > before); + assert_eq!( + engine.per_worker(), + vec![2, 2, 2], + "a change keeps the key where it was" + ); + assert_eq!(engine.subscribed(), 6); + + let key = query_key(&mut conn, live.key("p4")); + engine.unsubscribe(&key); + cache.touched(&key); + assert_eq!(engine.subscribed(), 5); + assert_eq!(engine.per_worker().iter().sum::(), 5); + assert_eq!(read(&mut conn, live.key("p4")).as_deref(), Some("v4b")); + assert_eq!(engine.subscribed(), 6, "the next miss subscribes again"); +} + +/// The cache key of the test's lookup: what `first` renders, bind values +/// included. +fn query_key( + conn: &mut SyncedConnection, + key: String, +) -> diesel_streamgres::key::QueryKey { + let query = sg_configs::table + .filter(sg_configs::key.eq(key)) + .select(sg_configs::value) + .limit(1); + diesel_streamgres::key::QueryKey::of(&query, conn.inner()).unwrap() +} + +/// Past the subscription cap the engine is not asked and the read is +/// served live, like a refused query; a held key still comes from the +/// cache. +/// +/// The third key is live every time, so a change is seen at once; the held +/// keys are cached. +/// +/// Freeing a slot lets the next miss subscribe. +#[test] +fn the_subscription_cap_serves_the_rest_live() { + let Some(mut live) = Live::start("cap") else { + return; + }; + let url = std::env::var("STREAMGRES_S2S_URL").unwrap(); + let database = std::env::var("DATABASE_URL").unwrap(); + let engine = EngineClient::with_config( + url, + diesel_streamgres::engine::Config { + connections: 1, + max_subscriptions: 2, + ..diesel_streamgres::engine::Config::from_env() + }, + ); + let cache = Arc::new(Cache::new(64 * 1024 * 1024, Arc::new(engine.clone()))); + engine.attach(&cache); + let deadline = Instant::now() + Duration::from_secs(5); + while !engine.is_connected() { + assert!(Instant::now() < deadline, "the engine did not connect"); + std::thread::sleep(Duration::from_millis(50)); + } + let mut conn = SyncedConnection::with_cache( + PgConnection::establish(&database).unwrap(), + cache.clone(), + Consistency::Eventual, + ); + for i in 0..3 { + live.put(&format!("c{i}"), &format!("v{i}")); + } + let read = |conn: &mut SyncedConnection, key: String| -> Option { + sg_configs::table + .filter(sg_configs::key.eq(key)) + .select(sg_configs::value) + .first(conn) + .optional() + .unwrap() + }; + for i in 0..3 { + assert_eq!( + read(&mut conn, live.key(&format!("c{i}"))).as_deref(), + Some(format!("v{i}").as_str()) + ); + } + assert_eq!(engine.subscribed(), 2, "the cap holds"); + assert_eq!(cache.len(), 2); + assert_eq!( + cache.stats.refused.load(Ordering::Relaxed), + 1, + "the third was refused by the client" + ); + + live.put("c2", "v2b"); + assert_eq!(read(&mut conn, live.key("c2")).as_deref(), Some("v2b")); + let hits = cache.stats.hits.load(Ordering::Relaxed); + assert_eq!(read(&mut conn, live.key("c0")).as_deref(), Some("v0")); + assert_eq!(cache.stats.hits.load(Ordering::Relaxed), hits + 1); + assert_eq!(engine.subscribed(), 2); + + let key = query_key(&mut conn, live.key("c0")); + engine.unsubscribe(&key); + cache.touched(&key); + assert_eq!(engine.subscribed(), 1); + assert_eq!(read(&mut conn, live.key("c2")).as_deref(), Some("v2b")); + assert_eq!( + engine.subscribed(), + 2, + "the freed slot went to the next miss" + ); +} + +diesel::table! { + sg_blobs (id) { + id -> Int4, + name -> Text, + data -> Nullable, + n -> Int4, + } +} + +/// `bytea` travels as hex text and comes back as bytes; `OFFSET` is +/// applied by the client over the rows the engine keeps. +/// +/// A change inside the page arrives as a delta; the NULL row and the offset +/// still apply. +#[test] +fn bytea_and_offset_are_served_from_engine_rows() { + let Some(mut live) = Live::start("blobs") else { + return; + }; + live.other + .batch_execute( + "CREATE TABLE IF NOT EXISTS sg_blobs ( + id serial PRIMARY KEY, name text NOT NULL, data bytea, n int NOT NULL)", + ) + .unwrap(); + let prefix = live.key("b"); + diesel::delete(sg_blobs::table.filter(sg_blobs::name.like(format!("{prefix}%")))) + .execute(&mut live.other) + .unwrap(); + let rows: Vec<_> = (0..5) + .map(|i| { + ( + sg_blobs::name.eq(format!("{prefix}{i}")), + sg_blobs::data.eq(if i == 3 { + None + } else { + Some(vec![i as u8, 0xff, 0x10]) + }), + sg_blobs::n.eq(i), + ) + }) + .collect(); + diesel::insert_into(sg_blobs::table) + .values(&rows) + .execute(&mut live.other) + .unwrap(); + + let names: Vec = (0..5).map(|i| format!("{prefix}{i}")).collect(); + let page = |conn: &mut SyncedConnection| -> Vec<(i32, Option>)> { + sg_blobs::table + .filter(sg_blobs::name.eq_any(&names)) + .select((sg_blobs::n, sg_blobs::data)) + .order(sg_blobs::n.asc()) + .limit(2) + .offset(1) + .load(conn) + .unwrap() + }; + let got = page(&mut live.conn); + assert_eq!( + got, + vec![ + (1, Some(vec![1, 0xff, 0x10])), + (2, Some(vec![2, 0xff, 0x10])) + ], + "skipped one, took two, bytes intact" + ); + assert_eq!( + live.engine.materialized(), + 1, + "rows mode: {:?}", + live.engine.fallbacks() + ); + assert_eq!(live.engine.subscribed(), 1); + + diesel::update(sg_blobs::table.filter(sg_blobs::name.eq(format!("{prefix}2")))) + .set(sg_blobs::data.eq(Some(vec![0xab]))) + .execute(&mut live.other) + .unwrap(); + let started = Instant::now(); + loop { + let got = page(&mut live.conn); + if got[1].1.as_deref() == Some(&[0xab][..]) { + break; + } + assert!( + started.elapsed() < Duration::from_secs(10), + "the delta did not arrive: {got:?}" + ); + std::thread::sleep(Duration::from_millis(20)); + } + let third: Vec<(i32, Option>)> = sg_blobs::table + .filter(sg_blobs::name.eq_any(&names)) + .select((sg_blobs::n, sg_blobs::data)) + .order(sg_blobs::n.asc()) + .limit(1) + .offset(3) + .load(&mut live.conn) + .unwrap(); + assert_eq!(third, vec![(3, None)], "a NULL bytea stays NULL"); + assert_eq!( + live.engine.materialized(), + 2, + "{:?}", + live.engine.fallbacks() + ); + let _ = diesel::delete(sg_blobs::table.filter(sg_blobs::name.like(format!("{prefix}%")))) + .execute(&mut live.other); +} + +/// A row inserted and deleted in one transaction is not kept. +/// +/// One transaction: a matching row comes and goes, and another row changes +/// so the delta's arrival can be seen. The engine folds the ghost into an +/// add and a delete; one frame must not leave it behind. +#[test] +fn a_row_inserted_and_deleted_in_one_transaction_is_not_kept() { + let Some(mut live) = Live::start("ghost") else { + return; + }; + live.put("real", "1"); + let names = vec![live.key("real"), live.key("ghost")]; + let page = |conn: &mut SyncedConnection| -> Vec<(String, String)> { + sg_configs::table + .filter(sg_configs::key.eq_any(&names)) + .select((sg_configs::key, sg_configs::value)) + .order(sg_configs::key.asc()) + .load(conn) + .unwrap() + }; + assert_eq!( + page(&mut live.conn), + vec![(live.key("real"), "1".to_string())] + ); + assert_eq!( + live.engine.materialized(), + 1, + "rows mode: {:?}", + live.engine.fallbacks() + ); + + let (real, ghost) = (live.key("real"), live.key("ghost")); + live.other + .transaction::<_, diesel::result::Error, _>(|conn| { + diesel::insert_into(sg_configs::table) + .values((sg_configs::key.eq(&ghost), sg_configs::value.eq("boo"))) + .execute(conn)?; + diesel::delete(sg_configs::table.filter(sg_configs::key.eq(&ghost))).execute(conn)?; + diesel::update(sg_configs::table.filter(sg_configs::key.eq(&real))) + .set(sg_configs::value.eq("2")) + .execute(conn)?; + Ok(()) + }) + .unwrap(); + let started = Instant::now(); + loop { + let got = page(&mut live.conn); + if got.iter().any(|(_, value)| value == "2") { + assert_eq!( + got, + vec![(real.clone(), "2".to_string())], + "the ghost row must not be in the result" + ); + break; + } + assert!( + started.elapsed() < Duration::from_secs(10), + "the delta did not arrive: {got:?}" + ); + std::thread::sleep(Duration::from_millis(20)); + } + assert_eq!( + live.cache.stats.touched.load(Ordering::Relaxed), + 0, + "served from engine rows throughout" + ); +} diff --git a/crates/diesel-streamgres/tests/join_sql_shape.rs b/crates/diesel-streamgres/tests/join_sql_shape.rs new file mode 100644 index 0000000..ced09a0 --- /dev/null +++ b/crates/diesel-streamgres/tests/join_sql_shape.rs @@ -0,0 +1,42 @@ +//! What Diesel renders for joins: the engine's parser is written against +//! these exact shapes, so a change in Diesel's rendering shows up here. +use diesel::pg::Pg; +use diesel::prelude::*; + +diesel::table! { authors (id) { id -> Int4, name -> Text } } +diesel::table! { books (id) { id -> Int4, author_id -> Int4, title -> Text } } +diesel::table! { reviews (id) { id -> Int4, book_id -> Int4, stars -> Int4 } } +diesel::joinable!(books -> authors (author_id)); +diesel::joinable!(reviews -> books (book_id)); +diesel::allow_tables_to_appear_in_same_query!(authors, books, reviews); + +/// Diesel renders an inner join, a left join, a nested join and an +/// explicit `ON` in the shapes the engine's parser reads. +#[test] +fn diesel_renders_joins_as_the_engine_expects() { + let inner = authors::table + .inner_join(books::table) + .filter(authors::name.eq("x")) + .select((authors::id, books::title)) + .order(authors::id.desc()) + .limit(5); + assert_eq!( + diesel::debug_query::(&inner).to_string(), + r#"SELECT "authors"."id", "books"."title" FROM ("authors" INNER JOIN "books" ON ("books"."author_id" = "authors"."id")) WHERE ("authors"."name" = $1) ORDER BY "authors"."id" DESC LIMIT $2 -- binds: ["x", 5]"# + ); + let left = authors::table.left_join(books::table).select((authors::name, books::title.nullable())); + assert_eq!( + diesel::debug_query::(&left).to_string(), + r#"SELECT "authors"."name", "books"."title" FROM ("authors" LEFT OUTER JOIN "books" ON ("books"."author_id" = "authors"."id")) -- binds: []"# + ); + let three = authors::table.inner_join(books::table.inner_join(reviews::table)).select((authors::name, reviews::stars)); + assert_eq!( + diesel::debug_query::(&three).to_string(), + r#"SELECT "authors"."name", "reviews"."stars" FROM ("authors" INNER JOIN ("books" INNER JOIN "reviews" ON ("reviews"."book_id" = "books"."id")) ON ("books"."author_id" = "authors"."id")) -- binds: []"# + ); + let explicit = authors::table.inner_join(books::table.on(books::author_id.eq(authors::id).and(books::title.eq("t")))).select(authors::id); + assert_eq!( + diesel::debug_query::(&explicit).to_string(), + r#"SELECT "authors"."id" FROM ("authors" INNER JOIN "books" ON (("books"."author_id" = "authors"."id") AND ("books"."title" = $1))) -- binds: ["t"]"# + ); +} diff --git a/crates/diesel-streamgres/tests/owned_row.rs b/crates/diesel-streamgres/tests/owned_row.rs new file mode 100644 index 0000000..566a305 --- /dev/null +++ b/crates/diesel-streamgres/tests/owned_row.rs @@ -0,0 +1,189 @@ +//! A captured row decodes exactly like the live row it was copied from. + +use diesel::connection::{DefaultLoadingMode, LoadConnection, SimpleConnection}; +use diesel::pg::Pg; +use diesel::row::Row; +use diesel::deserialize::{FromSqlRow, QueryableByName}; +use diesel::prelude::*; +use diesel::sql_query; +use diesel::sql_types::{Integer, Nullable, Text}; +use diesel_streamgres::SyncedConnection; +use diesel_streamgres::row::{OwnedCursor, OwnedRow, capture_all}; + +diesel::table! { + sg_items (id) { + id -> Int4, + name -> Text, + qty -> Int4, + note -> Nullable, + tags -> Array, + payload -> Jsonb, + created -> Timestamptz, + } +} + +/// A row of every column type the capture must keep. +#[derive(Queryable, Selectable, Insertable, Debug, PartialEq, Clone)] +#[diesel(table_name = sg_items)] +struct Item { + id: i32, + name: String, + qty: i32, + note: Option, + tags: Vec, + payload: serde_json::Value, + created: chrono::DateTime, +} + +/// The test database's URL. +fn url() -> String { + std::env::var("DATABASE_URL") + .unwrap_or_else(|_| "postgres:///diesel_streamgres_test?host=/tmp".to_string()) +} + +/// A wrapped connection inside a test transaction, the table created +/// once. +fn conn() -> SyncedConnection { + static SCHEMA_READY: std::sync::Once = std::sync::Once::new(); + SCHEMA_READY.call_once(|| { + let mut conn = PgConnection::establish(&url()).expect("connect"); + conn.batch_execute( + "CREATE TABLE IF NOT EXISTS sg_items ( + id serial PRIMARY KEY, name text NOT NULL, qty int NOT NULL, note text, + tags text[] NOT NULL DEFAULT '{}', payload jsonb NOT NULL DEFAULT '{}', + created timestamptz NOT NULL DEFAULT now())", + ) + .expect("schema"); + }); + let mut conn = SyncedConnection::::establish(&url()).expect("connect"); + conn.begin_test_transaction().expect("test transaction"); + conn +} + +/// Three rows covering NULLs, empty and quoted array items, JSON and +/// microsecond times; returns them as read live. +fn seed(conn: &mut SyncedConnection) -> Vec { + conn.batch_execute( + "INSERT INTO sg_items (name, qty, note, tags, payload, created) VALUES + ('bolt', 3, 'm6', '{metal,small}', '{\"len_mm\": 20, \"zinc\": true}', '2026-10-13 11:22:33.123456+05:30'), + ('washer', 0, NULL, '{}', '[1, 2.5, \"x\", null]', '1999-12-31 23:59:59Z'), + ('nut', -7, '', '{\"a b\",\"c,d\",\"\"}', '{}', '2026-01-01 00:00:00Z')", + ) + .expect("seed"); + sg_items::table + .order(sg_items::id) + .select(Item::as_select()) + .load(conn) + .expect("live rows") +} + +/// Loads the whole table as owned rows through Diesel's own cursor. +fn captured(conn: &mut SyncedConnection) -> Vec { + let cursor = as LoadConnection>::load( + conn, + sg_items::table.order(sg_items::id), + ) + .expect("load"); + capture_all(cursor).expect("capture") +} + +/// A captured row decodes exactly as the live row did. +#[test] +fn a_captured_row_decodes_like_the_live_one() { + let mut conn = conn(); + let live = seed(&mut conn); + let rows = captured(&mut conn); + assert_eq!(rows.len(), 3); + + let decoded: Vec = rows + .iter() + .map(|row| >::build_from_row(row).expect("decode")) + .collect(); + assert_eq!(decoded, live); + assert_eq!(decoded[1].note, None); + assert_eq!(decoded[2].tags, vec!["a b", "c,d", ""]); + assert_eq!(decoded[0].created.timestamp_subsec_micros(), 123456); +} + +/// The rows of one result share one list of column names. +#[test] +fn rows_of_one_result_share_their_column_names() { + let mut conn = conn(); + seed(&mut conn); + let rows = captured(&mut conn); + assert!(std::sync::Arc::ptr_eq(rows[0].columns(), rows[2].columns())); + assert_eq!( + rows[0].columns().iter().map(String::as_str).collect::>(), + ["id", "name", "qty", "note", "tags", "payload", "created"] + ); + assert!(rows[1].value(3).is_none(), "the washer's note is NULL"); + assert!(rows[0].byte_len() > 0); +} + +/// Three columns read by name, in another order than selected. +#[derive(QueryableByName, Debug, PartialEq)] +struct NameNote { + #[diesel(sql_type = Text)] + name: String, + #[diesel(sql_type = Nullable)] + note: Option, + #[diesel(sql_type = Integer)] + qty: i32, +} + +/// A captured row decodes by column name. +#[test] +fn a_captured_row_decodes_by_column_name() { + let mut conn = conn(); + seed(&mut conn); + let cursor = as LoadConnection>::load( + &mut conn, + sql_query("SELECT qty, note, name FROM sg_items ORDER BY id"), + ) + .expect("load"); + let rows = capture_all(cursor).expect("capture"); + let decoded: Vec = rows + .iter() + .map(|row| NameNote::build(row).expect("by name")) + .collect(); + assert_eq!( + decoded, + vec![ + NameNote { name: "bolt".into(), note: Some("m6".into()), qty: 3 }, + NameNote { name: "washer".into(), note: None, qty: 0 }, + NameNote { name: "nut".into(), note: Some("".into()), qty: -7 }, + ] + ); +} + +/// An owned cursor yields its rows in order, and knows how many. +#[test] +fn an_owned_cursor_yields_the_rows_in_order() { + let mut conn = conn(); + seed(&mut conn); + let rows = captured(&mut conn); + let cursor = OwnedCursor::new(rows.clone().into()); + assert_eq!(cursor.len(), 3); + let back: Vec = cursor.map(|row| row.expect("row")).collect(); + assert_eq!(back, rows); +} + +/// A nested tuple reads its part of the row. +#[test] +fn nested_tuples_read_partial_rows() { + let mut conn = conn(); + seed(&mut conn); + let rows = captured(&mut conn); + type Nested = (i32, (String, i32), Option); + type NestedSql = ( + diesel::sql_types::Int4, + (Text, Integer), + Nullable, + ); + let first: Nested = { + let row = rows[0].partial_row(0..4); + >::build_from_row(&row).expect("nested") + }; + assert_eq!(first.1, ("bolt".to_string(), 3)); + assert_eq!(first.2, Some("m6".to_string())); +} diff --git a/crates/diesel-streamgres/tests/passthrough.rs b/crates/diesel-streamgres/tests/passthrough.rs new file mode 100644 index 0000000..128ffc8 --- /dev/null +++ b/crates/diesel-streamgres/tests/passthrough.rs @@ -0,0 +1,348 @@ +//! The wrapper is transparent: every way Diesel reads and writes through a +//! `PgConnection` behaves the same through `SyncedConnection`. +//! +//! Needs a Postgres at `DATABASE_URL` (default: the local socket database +//! `diesel_streamgres_test`). Every test runs inside a test transaction, so +//! nothing it writes outlives it. + +use diesel::connection::{Instrumentation, InstrumentationEvent, SimpleConnection}; +use diesel::pg::{PgMetadataLookup, PgRowByRowLoadingMode}; +use diesel::prelude::*; +use diesel::r2d2::{ConnectionManager, Pool}; +use diesel::sql_query; +use diesel::sql_types::{Integer, Text}; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; +use diesel_streamgres::SyncedConnection; + +diesel::table! { + sg_items (id) { + id -> Int4, + name -> Text, + qty -> Int4, + note -> Nullable, + tags -> Array, + payload -> Jsonb, + created -> Timestamptz, + } +} + +/// A row as read. +#[derive(Queryable, Selectable, Debug, PartialEq, Clone)] +#[diesel(table_name = sg_items)] +struct Item { + id: i32, + name: String, + qty: i32, + note: Option, + tags: Vec, + payload: serde_json::Value, + created: chrono::DateTime, +} + +/// A row to insert. +#[derive(Insertable)] +#[diesel(table_name = sg_items)] +struct NewItem<'a> { + name: &'a str, + qty: i32, + note: Option<&'a str>, + tags: Vec<&'a str>, + payload: serde_json::Value, +} + +const SCHEMA: &str = " + CREATE TABLE IF NOT EXISTS sg_items ( + id serial PRIMARY KEY, + name text NOT NULL, + qty int NOT NULL, + note text, + tags text[] NOT NULL DEFAULT '{}', + payload jsonb NOT NULL DEFAULT '{}', + created timestamptz NOT NULL DEFAULT now() + ); +"; + +/// The test database's URL. +fn url() -> String { + std::env::var("DATABASE_URL") + .unwrap_or_else(|_| "postgres:///diesel_streamgres_test?host=/tmp".to_string()) +} + +/// Creates the table once per process: the tests run in parallel, and +/// concurrent `CREATE TABLE IF NOT EXISTS` on a missing table races in +/// Postgres. +fn ensure_schema() { + static SCHEMA_READY: std::sync::Once = std::sync::Once::new(); + SCHEMA_READY.call_once(|| { + let mut conn = PgConnection::establish(&url()).expect("connect"); + conn.batch_execute(SCHEMA).expect("schema"); + }); +} + +/// A wrapped connection inside a test transaction with the table present. +fn conn() -> SyncedConnection { + ensure_schema(); + let mut conn = SyncedConnection::::establish(&url()).expect("connect"); + conn.begin_test_transaction().expect("test transaction"); + conn +} + +/// Inserts three rows and returns them as read. +fn seed(conn: &mut SyncedConnection) -> Vec { + let rows = vec![ + NewItem { + name: "bolt", + qty: 3, + note: Some("m6"), + tags: vec!["metal", "small"], + payload: serde_json::json!({"len_mm": 20, "zinc": true}), + }, + NewItem { + name: "washer", + qty: 0, + note: None, + tags: vec![], + payload: serde_json::json!([1, 2.5, "x", null]), + }, + NewItem { + name: "nut", + qty: 7, + note: Some("m6"), + tags: vec!["metal"], + payload: serde_json::json!({}), + }, + ]; + diesel::insert_into(sg_items::table) + .values(&rows) + .get_results::(conn) + .expect("seed") +} + +/// The wrapper loads rows exactly as the inner connection does. +#[test] +fn loads_rows_exactly_like_the_inner_connection() { + let mut conn = conn(); + let inserted = seed(&mut conn); + + let through_wrapper = sg_items::table + .filter(sg_items::note.eq("m6")) + .order(sg_items::id) + .select(Item::as_select()) + .load(&mut conn) + .expect("wrapped load"); + let through_inner = sg_items::table + .filter(sg_items::note.eq("m6")) + .order(sg_items::id) + .select(Item::as_select()) + .load(conn.inner()) + .expect("inner load"); + + assert_eq!(through_wrapper, through_inner); + assert_eq!(through_wrapper.len(), 2); + assert_eq!(through_wrapper[0], inserted[0]); + assert_eq!(through_wrapper[1], inserted[2]); + assert_eq!(through_wrapper[0].tags, vec!["metal", "small"]); + assert_eq!(through_wrapper[0].payload["len_mm"], 20); +} + +/// `first` and `count` go through. +#[test] +fn first_and_count_go_through() { + let mut conn = conn(); + seed(&mut conn); + + let count: i64 = sg_items::table.count().get_result(&mut conn).expect("count"); + assert_eq!(count, 3); + + let washer: Item = sg_items::table + .filter(sg_items::name.eq("washer")) + .select(Item::as_select()) + .first(&mut conn) + .expect("first"); + assert_eq!(washer.note, None); + assert_eq!(washer.tags, Vec::::new()); + + let missing = sg_items::table + .filter(sg_items::name.eq("gear")) + .select(Item::as_select()) + .first(&mut conn) + .optional() + .expect("optional"); + assert_eq!(missing, None); +} + +/// Writes return their counts and their `RETURNING` rows. +#[test] +fn writes_return_counts_and_returning_rows() { + let mut conn = conn(); + seed(&mut conn); + + let updated = diesel::update(sg_items::table.filter(sg_items::note.eq("m6"))) + .set(sg_items::qty.eq(sg_items::qty + 1)) + .execute(&mut conn) + .expect("update"); + assert_eq!(updated, 2); + + let renamed: Item = diesel::update(sg_items::table.filter(sg_items::name.eq("washer"))) + .set(sg_items::name.eq("spring washer")) + .get_result(&mut conn) + .expect("update returning"); + assert_eq!(renamed.name, "spring washer"); + + let deleted = diesel::delete(sg_items::table.filter(sg_items::qty.gt(5))) + .execute(&mut conn) + .expect("delete"); + assert_eq!(deleted, 1, "only the nut had qty > 5 after the bump"); + + let left: Vec = sg_items::table + .order(sg_items::id) + .select(sg_items::name) + .load(&mut conn) + .expect("load"); + assert_eq!(left, vec!["bolt", "spring washer"]); +} + +/// Transactions roll back and nest through the wrapper: a failed +/// transaction undoes its delete, and a nested savepoint that succeeds +/// inside one that fails is undone with it. +#[test] +fn transactions_roll_back_and_nest_through_the_wrapper() { + let mut conn = conn(); + seed(&mut conn); + + let failed: Result<(), diesel::result::Error> = conn.transaction(|c| { + diesel::delete(sg_items::table).execute(c)?; + Err(diesel::result::Error::RollbackTransaction) + }); + assert!(matches!( + failed, + Err(diesel::result::Error::RollbackTransaction) + )); + let count: i64 = sg_items::table.count().get_result(&mut conn).expect("count"); + assert_eq!(count, 3, "the inner transaction's delete was rolled back"); + + let outer: Result<(), diesel::result::Error> = conn.transaction(|c| { + c.transaction(|c| { + diesel::insert_into(sg_items::table) + .values(NewItem { + name: "gear", + qty: 1, + note: None, + tags: vec![], + payload: serde_json::json!(null), + }) + .execute(c)?; + Ok::<_, diesel::result::Error>(()) + })?; + let n: i64 = sg_items::table.count().get_result(c)?; + assert_eq!(n, 4); + Err(diesel::result::Error::RollbackTransaction) + }); + assert!(outer.is_err()); + let count: i64 = sg_items::table.count().get_result(&mut conn).expect("count"); + assert_eq!(count, 3, "the outer rollback also undid the nested insert"); +} + +/// Two columns read by name. +#[derive(QueryableByName, Debug, PartialEq)] +struct NameQty { + #[diesel(sql_type = Text)] + name: String, + #[diesel(sql_type = Integer)] + qty: i32, +} + +/// Raw SQL read by column name goes through. +#[test] +fn raw_sql_queries_by_name_go_through() { + let mut conn = conn(); + seed(&mut conn); + + let rows = sql_query("SELECT name, qty FROM sg_items WHERE qty >= $1 ORDER BY qty") + .bind::(3) + .load::(&mut conn) + .expect("sql_query"); + assert_eq!( + rows, + vec![ + NameQty { name: "bolt".into(), qty: 3 }, + NameQty { name: "nut".into(), qty: 7 }, + ] + ); +} + +/// Postgres's row-by-row loading mode goes through. +#[test] +fn row_by_row_loading_mode_goes_through() { + let mut conn = conn(); + seed(&mut conn); + + let names: Vec = sg_items::table + .order(sg_items::id) + .select(sg_items::name) + .load_iter::(&mut conn) + .expect("load_iter") + .collect::>() + .expect("rows"); + assert_eq!(names, vec!["bolt", "washer", "nut"]); +} + +/// An r2d2 pool hands out wrapped connections, and takes them back. +#[test] +fn a_pool_hands_out_wrapped_connections() { + let pool: Pool>> = Pool::builder() + .max_size(2) + .build(ConnectionManager::new(url())) + .expect("pool"); + + let mut a = pool.get().expect("conn a"); + let mut b = pool.get().expect("conn b"); + let one: i32 = diesel::select(diesel::dsl::sql::("1")) + .get_result(&mut a) + .expect("select 1"); + let two: i32 = diesel::select(diesel::dsl::sql::("2")) + .get_result(&mut b) + .expect("select 2"); + assert_eq!((one, two), (1, 2)); + assert!(pool.get_timeout(std::time::Duration::from_millis(50)).is_err()); + drop(a); + assert!(pool.get_timeout(std::time::Duration::from_millis(500)).is_ok()); +} + +/// Counts the queries the connection starts. +struct CountQueries(Arc); + +impl Instrumentation for CountQueries { + /// Counts a query start. + fn on_connection_event(&mut self, event: InstrumentationEvent<'_>) { + if let InstrumentationEvent::StartQuery { .. } = event { + self.0.fetch_add(1, Ordering::SeqCst); + } + } +} + +/// Instrumentation set on the wrapper sees the inner connection's +/// queries. +#[test] +fn instrumentation_reaches_the_inner_connection() { + let mut conn = conn(); + let started = Arc::new(AtomicUsize::new(0)); + conn.set_instrumentation(CountQueries(started.clone())); + + seed(&mut conn); + let _: i64 = sg_items::table.count().get_result(&mut conn).expect("count"); + + assert_eq!(started.load(Ordering::SeqCst), 2, "one insert, one count"); +} + +/// The wrapper resolves type metadata as the inner connection does. +#[test] +fn the_wrapper_resolves_type_metadata_like_the_inner_connection() { + let mut conn = conn(); + let through_wrapper = PgMetadataLookup::lookup_type(&mut conn, "int4", None); + let through_inner = PgMetadataLookup::lookup_type(conn.inner(), "int4", None); + assert_eq!(through_wrapper.oid().ok(), through_inner.oid().ok()); + assert_eq!(through_wrapper.oid().ok(), Some(23)); + assert_eq!(through_wrapper.array_oid().ok(), Some(1007)); +} diff --git a/crates/diesel-streamgres/tests/settings.rs b/crates/diesel-streamgres/tests/settings.rs new file mode 100644 index 0000000..b09ca9b --- /dev/null +++ b/crates/diesel-streamgres/tests/settings.rs @@ -0,0 +1,88 @@ +//! The process-wide settings pick a cache and a route by the URL a +//! connection is established with. Its own binary: `install` is once per +//! process. + +use std::sync::Arc; + +use diesel::prelude::*; +use diesel_streamgres::cache::Cache; +use diesel_streamgres::key::QueryKey; +use diesel_streamgres::materialized::Snapshot; +use diesel_streamgres::subscriptions::{Outcome, Subscriptions, Verdict}; +use diesel_streamgres::types::TableTypes; +use diesel_streamgres::{Attachment, Route, Settings, SyncedConnection, install}; + +/// An engine side that refuses everything. +struct Unsynced; + +impl Subscriptions for Unsynced { + /// Refused. + fn subscribe(&self, _: &QueryKey) -> Outcome { + Outcome::Verdict(Verdict::Refused("test".into())) + } + /// Refused. + fn materialize(&self, _: &QueryKey, _: Snapshot, _: std::collections::HashMap>) -> Verdict { + Verdict::Refused("test".into()) + } + /// Nothing to end. + fn unsubscribe(&self, _: &QueryKey) {} +} + +/// The test database's URL. +fn url() -> String { + std::env::var("DATABASE_URL") + .unwrap_or_else(|_| "postgres:///diesel_streamgres_test?host=/tmp".to_string()) +} + +/// Settings attach a cache per database and a route per pool, by the URL a +/// connection is established with. +/// +/// The application name stands in for the database here; a real +/// configuration keys on the database the URL names. +#[test] +fn a_cache_per_database_and_a_route_per_pool() { + let main = Arc::new(Cache::new(1024, Arc::new(Unsynced))); + let accounts = Arc::new(Cache::new(1024, Arc::new(Unsynced))); + let (main_for, accounts_for) = (Arc::clone(&main), Arc::clone(&accounts)); + install(Settings::by_url(move |url| { + if url.contains("application_name=accounts") { + Some(Attachment { + cache: Arc::clone(&accounts_for), + route: Route::Streamgres, + }) + } else if url.contains("application_name=main_replica") { + Some(Attachment { + cache: Arc::clone(&main_for), + route: Route::Streamgres, + }) + } else if url.contains("application_name=main") { + Some(Attachment { + cache: Arc::clone(&main_for), + route: Route::Postgres, + }) + } else { + None + } + })) + .ok(); + + let sep = if url().contains('?') { "&" } else { "?" }; + let establish = |name: &str| { + SyncedConnection::::establish(&format!("{}{sep}application_name={name}", url())) + .expect("connect") + }; + + let master = establish("main"); + let replica = establish("main_replica"); + let accounts_conn = establish("accounts"); + let other = establish("elsewhere"); + + assert!(Arc::ptr_eq(master.cache().unwrap(), &main)); + assert!(Arc::ptr_eq(replica.cache().unwrap(), &main), "one cache per database"); + assert!(Arc::ptr_eq(accounts_conn.cache().unwrap(), &accounts), "another database, another cache"); + assert_eq!(master.route(), Route::Postgres, "the master pool goes to Postgres"); + assert_eq!(replica.route(), Route::Streamgres, "the replica pool through Streamgres"); + assert_eq!(accounts_conn.route(), Route::Streamgres); + assert!(other.cache().is_none(), "a URL with no attachment is a passthrough"); + assert_eq!(other.route(), Route::Postgres); +} diff --git a/crates/diesel-streamgres/tests/shape_coverage.rs b/crates/diesel-streamgres/tests/shape_coverage.rs new file mode 100644 index 0000000..5516511 --- /dev/null +++ b/crates/diesel-streamgres/tests/shape_coverage.rs @@ -0,0 +1,530 @@ +//! Production-shaped lookups through the synced connection, one shape +//! after another, side by side with Postgres: whether the wrapper's +//! answer equals Postgres's before and after a change, that the engine +//! serves each shape from its rows, and how long a change takes to reach +//! the cached answer. +//! +//! Every case runs the same query on a passthrough connection and on the +//! synced one, compares the two, changes the rows behind it with raw +//! SQL, waits for the synced answer to move, and compares again. The +//! shapes are the ones an application's model layer produces: a primary +//! key, a two-column key, `AND` and `OR` predicates, an enum bind, a +//! list, `ORDER BY ... LIMIT`, `LIMIT OFFSET`, a `bytea`, a timestamp +//! bind and a count. +//! +//! Needs `STREAMGRES_S2S_URL` and `DATABASE_URL` as `engine_live.rs` +//! does; skips itself otherwise. The enum type must exist before the +//! engine starts, since the engine reads its type map at start. + +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use chrono::{NaiveDate, NaiveDateTime}; +use diesel::connection::SimpleConnection; +use diesel::prelude::*; +use diesel_streamgres::cache::{Cache, Limits}; +use diesel_streamgres::engine::EngineClient; +use diesel_streamgres::{Route, Served, SyncedConnection}; + +/// A user-defined type, so the enum bind goes through Diesel's OID +/// lookup like an application's would. +#[derive(diesel::sql_types::SqlType, diesel::query_builder::QueryId)] +#[diesel(postgres_type(name = "sg_mood"))] +pub struct SgMood; + +/// The enum's values on the Rust side. +#[derive(Debug, Clone, Copy, PartialEq, diesel::AsExpression, diesel::FromSqlRow)] +#[diesel(sql_type = SgMood)] +enum Mood { + Happy, + Sad, +} + +impl diesel::serialize::ToSql for Mood { + /// The label Postgres knows the value by. + fn to_sql<'b>( + &'b self, + out: &mut diesel::serialize::Output<'b, '_, diesel::pg::Pg>, + ) -> diesel::serialize::Result { + use std::io::Write; + out.write_all(match self { + Mood::Happy => b"happy", + Mood::Sad => b"sad", + })?; + Ok(diesel::serialize::IsNull::No) + } +} + +impl diesel::deserialize::FromSql for Mood { + /// The value from its label. + fn from_sql(bytes: diesel::pg::PgValue<'_>) -> diesel::deserialize::Result { + match bytes.as_bytes() { + b"happy" => Ok(Mood::Happy), + b"sad" => Ok(Mood::Sad), + other => Err(format!("not a mood: {other:?}").into()), + } + } +} + +diesel::table! { + use diesel::sql_types::*; + use super::SgMood; + + sg_shapes (id) { + id -> Text, + merchant -> Text, + mood -> SgMood, + amount -> Int8, + created -> Timestamp, + blob -> Nullable, + note -> Nullable, + active -> Bool, + } +} + +/// A whole row, the way a model struct reads one. +#[derive(Queryable, Selectable, Debug, Clone, PartialEq)] +#[diesel(table_name = sg_shapes)] +struct Shape { + id: String, + merchant: String, + mood: Mood, + amount: i64, + created: NaiveDateTime, + blob: Option>, + note: Option, + active: bool, +} + +/// The blocking wrapper every shape is read through. +type Conn = SyncedConnection; + +/// The names this run's rows go by, unique per process so parallel runs +/// on one database do not meet. +struct Ids { + prefix: String, +} + +impl Ids { + /// The id of seeded row `n`. + fn id(&self, n: u8) -> String { + format!("{}{n}", self.prefix) + } + + /// The merchant four of the five seeded rows belong to. + fn merchant(&self) -> String { + format!("{}m", self.prefix) + } + + /// The SQL with `{p}` standing for this run's prefix. + fn sql(&self, template: &str) -> String { + template.replace("{p}", &self.prefix) + } +} + +/// One query shape, how to run it, and a change that alters its answer. +struct Case { + name: &'static str, + run: fn(&mut Conn, &Ids) -> String, + change: &'static str, +} + +/// The rows as text, so answers compare whatever the shape returns. +fn text(rows: T) -> String { + format!("{rows:?}") +} + +/// The shapes in the order they run; each change leaves the rows in a +/// state the later cases still read correctly. +fn cases() -> Vec { + vec![ + Case { + name: "primary key", + run: by_primary_key, + change: "UPDATE sg_shapes SET note = coalesce(note, '') || 'x' WHERE id = '{p}1'", + }, + Case { + name: "two-column key", + run: by_two_columns, + change: "UPDATE sg_shapes SET note = coalesce(note, '') || 'x' WHERE id = '{p}2'", + }, + Case { + name: "AND of three predicates", + run: and_of_three, + change: "UPDATE sg_shapes SET active = NOT active WHERE id = '{p}3'", + }, + Case { + name: "OR predicate", + run: or_predicate, + change: "UPDATE sg_shapes SET amount = 100 WHERE id = '{p}2'", + }, + Case { + name: "enum bind", + run: enum_bind, + change: "UPDATE sg_shapes SET mood = 'sad' WHERE id = '{p}1'", + }, + Case { + name: "list of keys", + run: list_of_keys, + change: "UPDATE sg_shapes SET note = coalesce(note, '') || 'x' WHERE id = '{p}5'", + }, + Case { + name: "ORDER BY timestamp DESC LIMIT", + run: latest_two, + change: "UPDATE sg_shapes SET created = '2026-01-01 01:00:00' WHERE id = '{p}4'", + }, + Case { + name: "LIMIT OFFSET", + run: second_page, + change: "UPDATE sg_shapes SET note = coalesce(note, '') || 'x' WHERE id = '{p}3'", + }, + Case { + name: "bytea", + run: blobs, + change: "UPDATE sg_shapes SET blob = '\\xab'::bytea WHERE id = '{p}2'", + }, + Case { + name: "timestamp bind", + run: since, + change: "UPDATE sg_shapes SET created = '2026-01-01 00:09:00' WHERE id = '{p}1'", + }, + Case { + name: "COUNT", + run: count, + change: "UPDATE sg_shapes SET merchant = '{p}m' WHERE id = '{p}5'", + }, + ] +} + +/// `WHERE id = $1`. +fn by_primary_key(conn: &mut Conn, ids: &Ids) -> String { + text( + sg_shapes::table + .filter(sg_shapes::id.eq(ids.id(1))) + .select(Shape::as_select()) + .first::(conn) + .optional() + .unwrap(), + ) +} + +/// `WHERE merchant = $1 AND id = $2`, a composite lookup. +fn by_two_columns(conn: &mut Conn, ids: &Ids) -> String { + text( + sg_shapes::table + .filter( + sg_shapes::merchant + .eq(ids.merchant()) + .and(sg_shapes::id.eq(ids.id(2))), + ) + .select(Shape::as_select()) + .first::(conn) + .optional() + .unwrap(), + ) +} + +/// Three predicates on a text, a bool and an integer column. +fn and_of_three(conn: &mut Conn, ids: &Ids) -> String { + text( + sg_shapes::table + .filter(sg_shapes::merchant.eq(ids.merchant())) + .filter(sg_shapes::active.eq(true)) + .filter(sg_shapes::amount.gt(20)) + .order(sg_shapes::id) + .select(Shape::as_select()) + .load::(conn) + .unwrap(), + ) +} + +/// `WHERE merchant = $1 AND (mood = $2 OR amount >= $3)`. +fn or_predicate(conn: &mut Conn, ids: &Ids) -> String { + text( + sg_shapes::table + .filter( + sg_shapes::merchant.eq(ids.merchant()).and( + sg_shapes::mood + .eq(Mood::Happy) + .or(sg_shapes::amount.ge(100)), + ), + ) + .order(sg_shapes::id) + .select(Shape::as_select()) + .load::(conn) + .unwrap(), + ) +} + +/// A single enum bind, which Diesel sends by its looked-up OID. +fn enum_bind(conn: &mut Conn, ids: &Ids) -> String { + text( + sg_shapes::table + .filter(sg_shapes::merchant.eq(ids.merchant())) + .filter(sg_shapes::mood.eq(Mood::Sad)) + .order(sg_shapes::id) + .select(Shape::as_select()) + .load::(conn) + .unwrap(), + ) +} + +/// `WHERE id = ANY($1)`. +fn list_of_keys(conn: &mut Conn, ids: &Ids) -> String { + text( + sg_shapes::table + .filter(sg_shapes::id.eq_any([ids.id(1), ids.id(3), ids.id(5)])) + .order(sg_shapes::id) + .select(Shape::as_select()) + .load::(conn) + .unwrap(), + ) +} + +/// The two newest rows of the merchant. +fn latest_two(conn: &mut Conn, ids: &Ids) -> String { + text( + sg_shapes::table + .filter(sg_shapes::merchant.eq(ids.merchant())) + .order(sg_shapes::created.desc()) + .limit(2) + .select(Shape::as_select()) + .load::(conn) + .unwrap(), + ) +} + +/// The second page of two, by id. +fn second_page(conn: &mut Conn, ids: &Ids) -> String { + text( + sg_shapes::table + .filter(sg_shapes::merchant.eq(ids.merchant())) + .order(sg_shapes::id) + .limit(2) + .offset(1) + .select(Shape::as_select()) + .load::(conn) + .unwrap(), + ) +} + +/// A `bytea` column, present and `NULL`. +fn blobs(conn: &mut Conn, ids: &Ids) -> String { + text( + sg_shapes::table + .filter(sg_shapes::id.eq_any([ids.id(1), ids.id(2), ids.id(3)])) + .order(sg_shapes::id) + .select((sg_shapes::id, sg_shapes::blob)) + .load::<(String, Option>)>(conn) + .unwrap(), + ) +} + +/// A `timestamp(6)` bound and compared at microsecond precision. +fn since(conn: &mut Conn, ids: &Ids) -> String { + let cutoff = NaiveDate::from_ymd_opt(2026, 1, 1) + .unwrap() + .and_hms_micro_opt(0, 2, 30, 500_000) + .unwrap(); + text( + sg_shapes::table + .filter(sg_shapes::merchant.eq(ids.merchant())) + .filter(sg_shapes::created.ge(cutoff)) + .order(sg_shapes::id) + .select(Shape::as_select()) + .load::(conn) + .unwrap(), + ) +} + +/// `SELECT count(*)` under a predicate. +fn count(conn: &mut Conn, ids: &Ids) -> String { + text( + sg_shapes::table + .filter(sg_shapes::merchant.eq(ids.merchant())) + .count() + .get_result::(conn) + .unwrap(), + ) +} + +/// The two connections under test, a third for seeding and changes, +/// and the engine client behind the cache. +struct Live { + synced: Conn, + plain: Conn, + other: PgConnection, + engine: EngineClient, + ids: Ids, +} + +impl Live { + /// Connects to the engine and the database, makes sure the table and + /// the enum exist, waits until the engine serves the table, and seeds + /// five rows under this run's prefix. `None` when the environment + /// names no engine. The cache keeps no refusal memory here: the + /// engine hears of a new table through the DDL trigger a moment after + /// it is created, and a subscribe refused in that moment would + /// otherwise be remembered for the run. + fn start() -> Option { + let Ok(url) = std::env::var("STREAMGRES_S2S_URL") else { + eprintln!("STREAMGRES_S2S_URL not set; skipping"); + return None; + }; + let database = std::env::var("DATABASE_URL").expect("DATABASE_URL"); + let engine = EngineClient::connect(url); + let limits = Limits { + refused_for: None, + ..Limits::default() + }; + let cache = Arc::new(Cache::with_limits(limits, Arc::new(engine.clone()))); + engine.attach(&cache); + let deadline = Instant::now() + Duration::from_secs(5); + while !engine.is_connected() { + assert!(Instant::now() < deadline, "the engine did not connect"); + std::thread::sleep(Duration::from_millis(50)); + } + let mut other = PgConnection::establish(&database).expect("connect"); + other + .batch_execute( + "BEGIN; + SELECT pg_advisory_xact_lock(424242); + DO $$ BEGIN CREATE TYPE sg_mood AS ENUM ('happy', 'sad'); + EXCEPTION WHEN duplicate_object OR unique_violation THEN NULL; END $$; + CREATE TABLE IF NOT EXISTS sg_shapes ( + id text PRIMARY KEY, + merchant text NOT NULL, + mood sg_mood NOT NULL, + amount int8 NOT NULL, + created timestamp(6) NOT NULL, + blob bytea, + note text, + active bool NOT NULL); + COMMIT", + ) + .expect("schema"); + let ids = Ids { + prefix: format!("shapes-{}-", std::process::id()), + }; + let mut live = Live { + synced: SyncedConnection::with_cache( + PgConnection::establish(&database).expect("connect"), + cache, + Route::Streamgres, + ), + plain: SyncedConnection::new(PgConnection::establish(&database).expect("connect")), + other, + engine, + ids, + }; + live.seed(); + live.wait_until_known(); + Some(live) + } + + /// Asks the engine for a count on the table until it answers from + /// its rows, which is when the DDL trigger's message has landed. + fn wait_until_known(&mut self) { + let started = Instant::now(); + loop { + let before = self.engine.materialized(); + let probe = sg_shapes::table + .filter(sg_shapes::id.eq(self.ids.sql("{p}probe"))) + .count() + .get_result::(&mut self.synced) + .expect("probe"); + assert_eq!(probe, 0); + if self.engine.materialized() > before { + return; + } + assert!( + started.elapsed() < Duration::from_secs(10), + "the engine never served the table: {:?}", + self.engine.refusals() + ); + std::thread::sleep(Duration::from_millis(50)); + } + } + + /// Five rows: four of one merchant and one of another, with both + /// moods, a `NULL` and a present `bytea`, and distinct microsecond + /// timestamps a minute apart. + fn seed(&mut self) { + self.cleanup(); + self.change( + "INSERT INTO sg_shapes (id, merchant, mood, amount, created, blob, note, active) VALUES + ('{p}1', '{p}m', 'happy', 10, '2026-01-01 00:01:00.123456', '\\x01ff10'::bytea, 'one', true), + ('{p}2', '{p}m', 'sad', 50, '2026-01-01 00:02:00.123456', '\\x02ff10'::bytea, NULL, true), + ('{p}3', '{p}m', 'happy', 100, '2026-01-01 00:03:00.123456', NULL, 'three', true), + ('{p}4', '{p}m', 'sad', 150, '2026-01-01 00:04:00.123456', '\\x04'::bytea, 'four', false), + ('{p}5', '{p}other', 'happy', 7, '2026-01-01 00:05:00.123456', NULL, 'five', false)", + ); + } + + /// Runs a change on the third connection, outside the wrapper. + fn change(&mut self, template: &str) { + let sql = self.ids.sql(template); + self.other.batch_execute(&sql).expect("change"); + } + + /// Removes this run's rows. + fn cleanup(&mut self) { + self.change("DELETE FROM sg_shapes WHERE id LIKE '{p}%'"); + } +} + +/// Every shape agrees with Postgres before and after its change, is +/// served from the engine's rows, and sees the change within seconds. +#[test] +fn every_shape_matches_postgres_and_is_served_from_engine_rows() { + let Some(mut live) = Live::start() else { + return; + }; + let materialized = live.engine.materialized(); + let subscribed = live.engine.subscribed(); + for (n, case) in cases().into_iter().enumerate() { + let name = case.name; + let expected = (case.run)(&mut live.plain, &live.ids); + let first = (case.run)(&mut live.synced, &live.ids); + assert_eq!(first, expected, "{name}: first answer"); + assert_eq!( + live.engine.materialized(), + materialized + n + 1, + "{name}: rows mode; fallbacks {:?}, refusals {:?}", + live.engine.fallbacks(), + live.engine.refusals() + ); + let again = (case.run)(&mut live.synced, &live.ids); + assert_eq!(again, expected, "{name}: repeat"); + assert_eq!( + live.synced.last_served(), + Some(Served::Cache), + "{name}: repeat is a hit" + ); + + live.change(case.change); + let started = Instant::now(); + let moved = loop { + let now = (case.run)(&mut live.synced, &live.ids); + if now != first { + break now; + } + assert!( + started.elapsed() < Duration::from_secs(10), + "{name}: the change never reached the cached answer" + ); + std::thread::sleep(Duration::from_millis(5)); + }; + let took = started.elapsed(); + let expected = (case.run)(&mut live.plain, &live.ids); + assert_eq!(moved, expected, "{name}: answer after the change"); + assert_eq!( + live.synced.last_served(), + Some(Served::Cache), + "{name}: still cached" + ); + eprintln!("SHAPE {name}: rows mode, change seen after {took:?}"); + } + assert_eq!(live.engine.subscribed(), subscribed + cases().len()); + live.cleanup(); +} diff --git a/crates/diesel-streamgres/tests/silent_engine.rs b/crates/diesel-streamgres/tests/silent_engine.rs new file mode 100644 index 0000000..c2cdff9 --- /dev/null +++ b/crates/diesel-streamgres/tests/silent_engine.rs @@ -0,0 +1,339 @@ +//! The engine client against sockets that misbehave: one that takes +//! every subscribe and never answers, and one that drops each connection +//! the moment it is made. The reads need `DATABASE_URL`. + +use std::sync::atomic::Ordering; +use std::sync::{Arc, Mutex}; +use std::thread; +use std::time::{Duration, Instant}; + +use diesel::connection::SimpleConnection; +use diesel::prelude::*; +use diesel_streamgres::cache::{Cache, Limits}; +use diesel_streamgres::engine::{Config, EngineClient}; +use diesel_streamgres::key::QueryKey; +use diesel_streamgres::{Consistency, Served, SyncedConnection}; +use futures_util::{SinkExt, StreamExt}; +use tokio::net::TcpListener; +use tokio_tungstenite::tungstenite::Message; + +diesel::table! { + sg_configs (id) { + id -> Int4, + key -> Text, + value -> Text, + } +} + +/// What the fake does with a connection once the handshake is done. +#[derive(Clone, Copy)] +enum Answer { + /// Greets, reads every frame, answers none. + Never, + /// Closes at once. + Close, + /// Greets this long after the handshake, then reads every frame and + /// answers none: the last frame heard falls between two pings. + GreetLate(Duration), +} + +/// What a fake server records: when each handshake completed, and every +/// frame received. +type Seen = (Arc>>, Arc>>); + +/// A WebSocket server on its own thread that behaves one way. +struct Fake { + url: String, + accepts: Arc>>, + frames: Arc>>, +} + +impl Fake { + /// Binds a free port and serves until the test ends. + fn start(answer: Answer) -> Fake { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("runtime"); + let listener = runtime + .block_on(TcpListener::bind("127.0.0.1:0")) + .expect("bind"); + let port = listener.local_addr().expect("addr").port(); + let accepts = Arc::new(Mutex::new(Vec::new())); + let frames = Arc::new(Mutex::new(Vec::new())); + let seen = (Arc::clone(&accepts), Arc::clone(&frames)); + thread::spawn(move || runtime.block_on(serve(listener, answer, seen))); + Fake { + url: format!("ws://127.0.0.1:{port}/zero/s2s/v1"), + accepts, + frames, + } + } + + /// When each connection completed its handshake. + fn accepts(&self) -> Vec { + self.accepts.lock().unwrap().clone() + } + + /// Every frame received so far, in order. + fn frames(&self) -> Vec { + self.frames.lock().unwrap().clone() + } + + /// Waits until `count` frames have arrived. + fn frames_after(&self, count: usize, wait: Duration) -> Vec { + let deadline = Instant::now() + wait; + while self.frames.lock().unwrap().len() < count { + assert!(Instant::now() < deadline, "only {:?}", self.frames()); + thread::sleep(Duration::from_millis(10)); + } + self.frames() + } +} + +/// Accepts connections forever, each handled by `answer`. +async fn serve(listener: TcpListener, answer: Answer, seen: Seen) { + loop { + let Ok((stream, _)) = listener.accept().await else { + return; + }; + let Ok(mut socket) = tokio_tungstenite::accept_async(stream).await else { + continue; + }; + seen.0.lock().unwrap().push(Instant::now()); + if let Answer::Close = answer { + continue; + } + let frames = Arc::clone(&seen.1); + tokio::spawn(async move { + if let Answer::GreetLate(delay) = answer { + tokio::time::sleep(delay).await; + } + let hello = serde_json::json!(["hello", {"protocol": 1}]).to_string(); + if socket.send(Message::Text(hello.into())).await.is_err() { + return; + } + while let Some(Ok(Message::Text(text))) = socket.next().await { + if let Ok(frame) = serde_json::from_str(text.as_str()) { + frames.lock().unwrap().push(frame); + } + } + }); + } +} + +/// A client on the fake, connected, with a cache that keeps refusals +/// for the default 30 seconds. +fn client(fake: &Fake, verdict_timeout: Duration) -> (EngineClient, Arc) { + let engine = EngineClient::with_config( + fake.url.clone(), + Config { + connections: 1, + verdict_timeout, + ..Config::default() + }, + ); + let cache = Arc::new(Cache::with_limits( + Limits::default(), + Arc::new(engine.clone()), + )); + engine.attach(&cache); + let deadline = Instant::now() + Duration::from_secs(5); + while !engine.is_connected() { + assert!(Instant::now() < deadline, "the client did not connect"); + thread::sleep(Duration::from_millis(20)); + } + (engine, cache) +} + +/// The test database's URL. +fn url() -> String { + std::env::var("DATABASE_URL") + .unwrap_or_else(|_| "postgres:///diesel_streamgres_test?host=/tmp".to_string()) +} + +/// A row to read, written outside the wrapper, and the wrapper to read +/// it through. +fn reader(cache: Arc, key: &str) -> (SyncedConnection, String) { + let mut other = PgConnection::establish(&url()).expect("connect"); + other + .batch_execute( + "CREATE TABLE IF NOT EXISTS sg_configs ( + id serial PRIMARY KEY, key text NOT NULL UNIQUE, value text NOT NULL)", + ) + .expect("schema"); + let key = format!("silent-{}-{key}", std::process::id()); + diesel::insert_into(sg_configs::table) + .values((sg_configs::key.eq(&key), sg_configs::value.eq("v"))) + .on_conflict(sg_configs::key) + .do_update() + .set(sg_configs::value.eq("v")) + .execute(&mut other) + .expect("put"); + let inner = PgConnection::establish(&url()).expect("connect"); + ( + SyncedConnection::with_cache(inner, cache, Consistency::Eventual), + key, + ) +} + +/// The frame's tag and the id in its body. +fn tagged(frame: &serde_json::Value) -> (&str, u64) { + ( + frame[0].as_str().unwrap_or(""), + frame[1]["id"].as_u64().unwrap_or(0), + ) +} + +/// An engine that takes the subscribe and never answers: the read +/// waits the verdict timeout and no longer, goes live, and withdraws its +/// id; the timeout is not remembered as a refusal, so the next read asks +/// again. +#[test] +fn a_hung_engine_is_given_up_on_in_time_and_asked_again() { + let fake = Fake::start(Answer::Never); + let (engine, cache) = client(&fake, Duration::from_millis(300)); + let (mut conn, key) = reader(Arc::clone(&cache), "hung"); + let query = || { + sg_configs::table + .filter(sg_configs::key.eq(&key)) + .select(sg_configs::value) + .limit(1) + }; + let query_key = QueryKey::of(&query(), conn.inner()).expect("key"); + + let started = Instant::now(); + let value: Option = query().first(&mut conn).optional().expect("read"); + let took = started.elapsed(); + assert_eq!(value.as_deref(), Some("v")); + assert_eq!(conn.last_served(), Some(Served::Postgres), "served live"); + assert!( + took >= Duration::from_millis(300), + "waited the timeout: {took:?}" + ); + assert!(took < Duration::from_secs(3), "not much longer: {took:?}"); + assert_eq!(cache.stats.refused.load(Ordering::Relaxed), 1); + assert!( + engine + .refusals() + .iter() + .any(|reason| reason.contains("did not answer in time")), + "{:?}", + engine.refusals() + ); + let frames = fake.frames_after(2, Duration::from_secs(2)); + assert_eq!(tagged(&frames[0]), ("subscribe", 1)); + assert_eq!( + tagged(&frames[1]), + ("unsubscribe", 1), + "the wait withdrew its id" + ); + assert_eq!(engine.subscribed(), 0); + assert!( + cache.refused_recently(&query_key).is_none(), + "a timeout is not remembered as a refusal" + ); + + let value: Option = query().first(&mut conn).optional().expect("read again"); + assert_eq!(value.as_deref(), Some("v")); + assert_eq!(conn.last_served(), Some(Served::Postgres)); + assert_eq!( + cache.stats.refused.load(Ordering::Relaxed), + 2, + "asked again" + ); + assert_eq!(cache.stats.refused_recently.load(Ordering::Relaxed), 0); + let frames = fake.frames_after(4, Duration::from_secs(2)); + assert_eq!(tagged(&frames[2]), ("subscribe", 2)); + assert_eq!(tagged(&frames[3]), ("unsubscribe", 2)); +} + +/// An engine that greets and then goes silent (killed without a close, +/// a network gone without a reset) is taken for dead once the socket has +/// been quiet for the silence window, and dialled again. +#[test] +fn a_silent_engine_is_taken_for_dead() { + let fake = Fake::start(Answer::Never); + let engine = EngineClient::with_config( + fake.url.clone(), + Config { + connections: 1, + silence: Duration::from_millis(600), + ..Config::default() + }, + ); + let deadline = Instant::now() + Duration::from_secs(5); + while fake.accepts().len() < 2 && Instant::now() < deadline { + thread::sleep(Duration::from_millis(20)); + } + assert!( + fake.accepts().len() >= 2, + "{} connections: the silent one was not given up", + fake.accepts().len() + ); + assert!( + fake.frames().iter().any(|frame| frame[0] == "ping"), + "the quiet socket was pinged" + ); + drop(engine); +} + +/// A silent socket is given up when the silence window runs out, not +/// at the next ping up to a third of a window later. The engine greets +/// 2.9 s in and then says nothing; with a 6 s window it is dead at 8.9 +/// s, and the 1 s reconnect delay brings the second connection at +/// about 9.9 s. Looking only at each ping (2, 4, ... s) it was 11 s. +#[test] +fn a_silent_socket_is_given_up_on_time() { + let fake = Fake::start(Answer::GreetLate(Duration::from_millis(2900))); + let engine = EngineClient::with_config( + fake.url.clone(), + Config { + connections: 1, + silence: Duration::from_secs(6), + ..Config::default() + }, + ); + let started = Instant::now(); + while fake.accepts().len() < 2 { + assert!(started.elapsed() < Duration::from_secs(20), "never redialled"); + thread::sleep(Duration::from_millis(20)); + } + let accepts = fake.accepts(); + let gap = accepts[1] - accepts[0]; + assert!( + gap >= Duration::from_millis(9800) && gap < Duration::from_millis(10_500), + "redialled {gap:?} after the first connection; silent from 2.9 s, the window is 6 s" + ); + drop(engine); +} + +/// An engine that closes every socket is dialled again about once a +/// second, never faster. +#[test] +fn a_dropped_socket_is_tried_again_every_second() { + let fake = Fake::start(Answer::Close); + let engine = EngineClient::with_config( + fake.url.clone(), + Config { + connections: 1, + ..Config::default() + }, + ); + thread::sleep(Duration::from_millis(3600)); + let accepts = fake.accepts(); + assert!( + (3..=5).contains(&accepts.len()), + "{} connections in 3.6 seconds", + accepts.len() + ); + for pair in accepts.windows(2) { + let gap = pair[1] - pair[0]; + assert!( + gap >= Duration::from_millis(900) && gap <= Duration::from_millis(2000), + "a reconnect came after {gap:?}" + ); + } + assert!(!engine.is_connected(), "never up for long"); + drop(engine); +} diff --git a/crates/streamgres-demo/Cargo.toml b/crates/streamgres-demo/Cargo.toml new file mode 100644 index 0000000..56a17b1 --- /dev/null +++ b/crates/streamgres-demo/Cargo.toml @@ -0,0 +1,25 @@ +[package] +name = "streamgres-demo" +version = "0.1.0" +edition = "2024" +rust-version = "1.88" +license = "Apache-2.0" +description = "A ticket-desk backend on the Streamgres Diesel wrapper: the engine's demo, and the wrapper's test bed" +publish = false + +[[bin]] +name = "demo" +path = "src/main.rs" + +[dependencies] +diesel-streamgres = { path = "../diesel-streamgres", features = ["async"] } +diesel = { version = "2.2.10", features = ["postgres", "r2d2", "chrono"] } +diesel-async = { version = "0.6", features = ["postgres", "bb8"] } +axum = "0.8" +chrono = { version = "0.4", features = ["serde"] } +clap = { version = "4", features = ["derive"] } +serde = { version = "1", features = ["derive"] } +serde_json = "1" +tokio = { version = "1", features = ["full"] } +rand = "0.9" +futures-util = "0.3" diff --git a/crates/streamgres-demo/README.md b/crates/streamgres-demo/README.md new file mode 100644 index 0000000..b4334fb --- /dev/null +++ b/crates/streamgres-demo/README.md @@ -0,0 +1,134 @@ +# streamgres-demo + +A ticket desk on the Streamgres Diesel wrapper. It is the engine's demo +(`demo/README.md` at the repository root walks through it) and the +wrapper's test bed: a fixed list of query shapes, a writer that makes +every kind of change, a verify loop that compares every shape across +the routes while the writer runs, and a scenario test that holds every +shape to what it expects of the engine. + +## The query list + +`src/queries.rs` holds every read the desk does, as `Lookup`: one named +shape each, runnable through the blocking wrapper (`run_sync`) and the +async one (`run_async`), with what it expects of the engine +(`expected()`): served from the engine's rows, refilled from Postgres +on each change, or refused and served live. + +| lookup | exercises | expects | +|-------------------|-------------------------------------------------------------|---------| +| `ticket` | primary key, enum columns, timestamps | rows | +| `queue` | enum array bind, `ORDER BY` timestamp `DESC`, `LIMIT`, `OFFSET` | rows | +| `queue_count` | `COUNT(*)` with an enum array bind | rows | +| `comments` | inner join composed from engine rows | rows | +| `attachment` | `bytea` | rows | +| `agents` | `text[]` column | rows | +| `sla` | `interval` column | refill | +| `search` | `ILIKE` | live | +| `recent` | `ORDER BY updated_at DESC LIMIT` | rows | +| `org_ticket` | two-column key | rows | +| `unassigned` | `IS NULL` | rows | +| `escalated` | `OR` of two enum binds | rows | +| `open_except` | `<>` and `NOT IN` | rows | +| `created_between` | timestamp range, two binds | rows | +| `many` | `IN` list of a hundred ids | rows | +| `queue_end` | page past the end: empty, subscribed, stays empty | rows | +| `nothing` | never any rows: empty, subscribed, stays empty | rows | +| `all_tickets` | largest result; not kept under a small per-entry cap | rows | +| `comment_authors` | left join with an empty side | rows | +| `skilled` | array `@>` | live | +| `statuses` | `DISTINCT` | live | +| `workload` | `GROUP BY` | live | +| `last_activity` | `MAX` | live | +| `of_kind` | an enum type added by `demo migrate new-type`; not in the demo list | rows | + +A shape that starts to matter to the wrapper goes in this list, so it +is exercised the day it changes. + +## The writer + +`src/writer.rs` makes one change per tick, drawn from a mix: a status +moved on, a comment added, a ticket inserted into the first +organization's queue, one of the writer's own tickets deleted with its +comments cascading, one moved to another organization, an agent renamed +under the joins, an assignee flipped between an agent and `NULL`, +several tickets touched in one transaction, a change rolled back, and a +comment with a body too large for its row written and then updated +without touching the body. `DEMO_WRITER_MIX` names the operations to +draw from; the default is all of them. The seeded rows keep their ids, +so the demo's URLs stay valid. + +## The verify loop + +`src/verify.rs` reads every lookup through the cache and through +Postgres and compares. A pair that differs is read again until it +agrees or a settle time runs out; one that never agrees is a mismatch, +kept with both answers. `demo verify` runs it for a while under the +writer and exits 1 on any mismatch; `serve` runs it in the background +when `DEMO_VERIFY_MS` is set; `/verify` runs one pass; `/stats` +carries the counters. + +## The scenario test + +`tests/scenario.rs` runs against a live engine and asserts what the +wrapper promises: + +- every lookup matches Postgres on the first read, is served the way + it expects (rows, refills or live), is a hit by the third read, shares + its entry between the two wrappers, and sees a change made elsewhere + within a bound, after which it matches Postgres again; the empty + results stay empty; +- a write on one connection is seen by the next read on another of the + same cache; +- a flag on the query beats the connection's route both ways, a scope + routes everything inside it, and the read-your-writes window ends; +- a result over the per-entry cap is served and correct but not kept; +- a read dropped mid-flight leaves no subscription behind; +- sixteen readers comparing every lookup across the routes while the + writer runs its whole mix never see the cache disagree with Postgres + for long; +- the same under an engine restart in the middle (needs + `STREAMGRES_RESTART_CMD`); +- the three kinds of schema change: a widened table is absorbed and + the cache keeps serving, a new enum type is absorbed and learnt on + its first bind, and a dropped column stops the + engine, which restarted serves the schema as it is (needs + `STREAMGRES_RESTART_CMD`). + +The load tests tell a wrong answer in this process's cache from one in +the engine, by reading the lookup again through a fresh subscription, +and fail on either, named. The restart test is the one that found the +engine reading a timestamp range in the server's own time zone. + +The tests reseed the dataset, so run them against a database nobody is +watching, one at a time: + +```bash +STREAMGRES_S2S_URL=ws://127.0.0.1:4848/zero/s2s/v1 \ +DATABASE_URL=postgres://user@host:5432/db \ +STREAMGRES_RESTART_CMD='scripts/restart-engine.sh' \ +cargo test -p streamgres-demo --test scenario -- --test-threads=1 --nocapture +``` + +`STREAMGRES_RESTART_CMD` is a shell command that restarts the engine +and returns once it serves again; without it the two tests that need +one skip themselves. + +## The binary + +```bash +demo seed # schema and dataset; safe to repeat +demo serve # the HTTP API and the stats page +demo bench --lookup queue # p50/p95/p99 per route, in process +demo bench --list +demo verify --for 60 # every lookup on both routes, under the writer; exits 1 on a mismatch +demo lag --rate 20 --for 10 # commit to cached answer, under 20 writes/s +demo migrate widen|new-type|drop # a schema change of each kind +``` + +`DATABASE_URL` and `STREAMGRES_S2S_URL` are required. Optional: +`DEMO_ADDR` (`0.0.0.0:8080`), `DEMO_WRITER_MS` (`1000`; `0` turns +the writer off), `DEMO_WRITER_MIX` (all operations), `DEMO_VERIFY_MS` +(`0`: off), `DEMO_SCALE` (`1`; multiplies the seed's tickets, comments +and attachments), and the wrapper's own `STREAMGRES_CACHE_*`, +`STREAMGRES_S2S_*` and `STREAMGRES_READ_YOUR_WRITES_MS`. diff --git a/crates/streamgres-demo/src/api.rs b/crates/streamgres-demo/src/api.rs new file mode 100644 index 0000000..0d98170 --- /dev/null +++ b/crates/streamgres-demo/src/api.rs @@ -0,0 +1,540 @@ +//! The HTTP face: one endpoint per lookup, a write, and the stats. +//! +//! Every read answers with the rows, which route it took, how long it +//! took, and whether it came from the cache. The last is read off the +//! cache's hit counter before and after the query, which is exact when +//! one request runs at a time and a good guess under load; the counters +//! on `/stats` are the exact figures. +//! +//! Query parameters on every read: +//! +//! - `route=postgres` or `route=streamgres` flags the query itself, +//! which beats the pool's route. +//! - `pool=reporting` uses the Postgres-routed pool instead of the main +//! one. +//! +//! `/verify` reads every lookup through both routes and reports the +//! pairs that never agreed; `/stats` carries the verify loop's and the +//! writer's counters next to the cache's and the engine's. + +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::{Duration, Instant}; + +use chrono::{DateTime, TimeZone, Utc}; + +use axum::extract::{Path, Query, State}; +use axum::http::StatusCode; +use axum::response::{Html, IntoResponse}; +use axum::routing::{get, post}; +use axum::{Json, Router}; +use diesel::prelude::*; +use diesel_async::RunQueryDsl; +use diesel_streamgres::Served; +use serde::Deserialize; +use serde_json::{Value, json}; + +use crate::models::{TicketKind, TicketStatus}; +use crate::pools::Stack; +use crate::queries::Lookup; +use crate::schema::tickets; +use crate::verify; + +/// What every handler shares. +pub type App = Arc; + +/// Every route of the demo's HTTP API. +pub fn router(stack: App) -> Router { + Router::new() + .route("/", get(index)) + .route("/stats", get(stats)) + .route("/lookups", get(lookups)) + .route("/verify", get(verify_once)) + .route("/tickets", get(many)) + .route("/tickets/{id}", get(ticket)) + .route("/tickets/{id}/comments", get(comments)) + .route("/tickets/{id}/authors", get(comment_authors)) + .route("/tickets/{id}/status", post(set_status)) + .route("/attachments/{id}", get(attachment)) + .route("/orgs/{org}/queue", get(queue)) + .route("/orgs/{org}/queue/count", get(queue_count)) + .route("/orgs/{org}/queue/end", get(queue_end)) + .route("/orgs/{org}/agents", get(agents)) + .route("/orgs/{org}/sla", get(sla)) + .route("/orgs/{org}/search", get(search)) + .route("/orgs/{org}/recent", get(recent)) + .route("/orgs/{org}/tickets", get(all_tickets)) + .route("/orgs/{org}/tickets/{id}", get(org_ticket)) + .route("/orgs/{org}/unassigned", get(unassigned)) + .route("/orgs/{org}/escalated", get(escalated)) + .route("/orgs/{org}/open", get(open_except)) + .route("/orgs/{org}/created", get(created_between)) + .route("/orgs/{org}/nothing", get(nothing)) + .route("/orgs/{org}/skilled", get(skilled)) + .route("/orgs/{org}/statuses", get(statuses)) + .route("/orgs/{org}/workload", get(workload)) + .route("/orgs/{org}/activity", get(last_activity)) + .route("/orgs/{org}/kind/{kind}", get(of_kind)) + .with_state(stack) +} + +/// The stats page. +async fn index() -> Html<&'static str> { + Html(include_str!("index.html")) +} + +/// The counters: the cache's, the engine client's, the writer's and +/// the verify loop's. +async fn stats(State(app): State) -> Json { + Json(stats_json(&app)) +} + +/// One pass of the verify loop, with every pair that never agreed. +async fn verify_once(State(app): State) -> Json { + let outcome = verify::once(&app, &Lookup::demo(), Duration::from_secs(1)).await; + Json(json!({ + "compared": outcome.compared, + "lagged": outcome.lagged, + "mismatches": outcome.mismatches, + "totals": app.verify.json(), + })) +} + +/// The counters as JSON, for `/stats` and the stats page. +pub fn stats_json(app: &Stack) -> Value { + let s = &app.cache.stats; + let n = |counter: &AtomicU64| AtomicU64::load(counter, Ordering::Relaxed); + json!({ + "cache": { + "hits": n(&s.hits), + "misses": n(&s.misses), + "touched": n(&s.touched), + "replaced": n(&s.replaced), + "refused": n(&s.refused), + "refused_recently": n(&s.refused_recently), + "after_write": n(&s.after_write), + "lagging": n(&s.lagging), + "locking": n(&s.locking), + "stale_fills": n(&s.stale_fills), + "evicted": n(&s.evicted), + "expired": n(&s.expired), + "oversized": n(&s.oversized), + "entries": app.cache.len(), + "bytes": app.cache.weighted_size(), + }, + "engine": { + "connected": app.engine.is_connected(), + "connections": app.engine.connections(), + "subscribed": app.engine.subscribed(), + "materialized": app.engine.materialized(), + "lag_ms": app.engine.lag().map(|lag| lag.as_millis() as u64), + "store_bytes": app.engine.store_bytes(), + "per_worker": app.engine.per_worker(), + "fallbacks": app.engine.fallbacks(), + "refusals": app.engine.refusals(), + }, + "writer": app.writer.json(), + "verify": app.verify.json(), + }) +} + +/// The lookups and what each exercises. +async fn lookups() -> Json { + Json(Value::Array( + Lookup::demo() + .into_iter() + .map(|lookup| json!({"name": lookup.name(), "exercises": lookup.exercises(), "args": lookup})) + .collect(), + )) +} + +/// A read's query parameters: the route and pool to take, and each +/// lookup's own arguments. +#[derive(Deserialize, Default)] +struct ReadParams { + route: Option, + pool: Option, + status: Option, + page: Option, + q: Option, + limit: Option, + except: Option, + from: Option, + to: Option, + ids: Option, + skill: Option, +} + +impl ReadParams { + /// The statuses asked for, open and pending unless named. + fn statuses(&self) -> Vec { + match &self.status { + Some(list) => list.split(',').filter_map(TicketStatus::parse).collect(), + None => vec![TicketStatus::Open, TicketStatus::Pending], + } + } + + /// A comma-separated list of ids, or `default`. + fn id_list(list: Option<&str>, default: Vec) -> Vec { + match list { + Some(list) => list + .split(',') + .filter_map(|id| id.trim().parse().ok()) + .collect(), + None => default, + } + } + + /// An RFC 3339 instant, or `default`. + fn instant(text: Option<&str>, default: DateTime) -> DateTime { + text.and_then(|text| text.parse::>().ok()) + .unwrap_or(default) + } +} + +/// A handler's answer: the read's report, or an error status. +type Reply = Result, (StatusCode, String)>; + +/// Runs a lookup on the pool and route the parameters ask for, and +/// reports what happened. A route for one request is set on the pooled +/// connection for the read and put back before the connection returns +/// to the pool. Where the rows came from is read per connection, so it +/// is right under concurrent requests, where a delta of the cache's +/// shared counters would not be. +async fn read(app: &Stack, params: &ReadParams, lookup: Lookup) -> Reply { + let pool = match params.pool.as_deref() { + Some("reporting") => &app.reporting, + _ => &app.main, + }; + let mut conn = pool + .get() + .await + .map_err(|error| (StatusCode::SERVICE_UNAVAILABLE, error.to_string()))?; + let route = match params.route.as_deref() { + Some("postgres") => Some(diesel_streamgres::Route::Postgres), + Some("streamgres") => Some(diesel_streamgres::Route::Streamgres), + _ => None, + }; + let started = Instant::now(); + let pool_route = conn.route(); + if let Some(route) = route { + conn.set_route(route); + } + let rows = lookup.run_async(&mut conn).await; + conn.set_route(pool_route); + let rows = rows.map_err(|error| (StatusCode::INTERNAL_SERVER_ERROR, error.to_string()))?; + let elapsed = started.elapsed(); + let source = match conn.last_served() { + Some(Served::Cache) => "cache", + Some(Served::Engine) => "engine", + Some(Served::Postgres) | None => "postgres", + }; + Ok(Json(json!({ + "lookup": lookup.name(), + "route": route.map(|route| format!("{route:?}").to_lowercase()).unwrap_or_else(|| match params.pool.as_deref() { + Some("reporting") => "postgres (reporting pool)".to_string(), + _ => "streamgres (main pool)".to_string(), + }), + "source": source, + "micros": elapsed.as_micros(), + "rows": rows, + }))) +} + +/// `GET /tickets/{id}`: the [`Lookup::Ticket`] lookup. +async fn ticket( + State(app): State, + Path(id): Path, + Query(params): Query, +) -> Reply { + read(&app, ¶ms, Lookup::Ticket { id }).await +} + +/// `GET /tickets/{id}/comments`: the [`Lookup::Comments`] lookup. +async fn comments( + State(app): State, + Path(id): Path, + Query(params): Query, +) -> Reply { + read(&app, ¶ms, Lookup::Comments { ticket: id }).await +} + +/// `GET /attachments/{id}`: the [`Lookup::Attachment`] lookup. +async fn attachment( + State(app): State, + Path(id): Path, + Query(params): Query, +) -> Reply { + read(&app, ¶ms, Lookup::Attachment { id }).await +} + +/// `GET /orgs/{org}/queue`: the [`Lookup::Queue`] lookup. +async fn queue( + State(app): State, + Path(org): Path, + Query(params): Query, +) -> Reply { + let lookup = Lookup::Queue { + org, + statuses: params.statuses(), + page: params.page.unwrap_or(0), + }; + read(&app, ¶ms, lookup).await +} + +/// `GET /orgs/{org}/queue/count`: the [`Lookup::QueueCount`] lookup. +async fn queue_count( + State(app): State, + Path(org): Path, + Query(params): Query, +) -> Reply { + let lookup = Lookup::QueueCount { + org, + statuses: params.statuses(), + }; + read(&app, ¶ms, lookup).await +} + +/// `GET /orgs/{org}/agents`: the [`Lookup::Agents`] lookup. +async fn agents( + State(app): State, + Path(org): Path, + Query(params): Query, +) -> Reply { + read(&app, ¶ms, Lookup::Agents { org }).await +} + +/// `GET /orgs/{org}/sla`: the [`Lookup::Sla`] lookup. +async fn sla( + State(app): State, + Path(org): Path, + Query(params): Query, +) -> Reply { + read(&app, ¶ms, Lookup::Sla { org }).await +} + +/// `GET /orgs/{org}/search`: the [`Lookup::Search`] lookup. +async fn search( + State(app): State, + Path(org): Path, + Query(params): Query, +) -> Reply { + let pattern = format!( + "%{}%", + params.q.clone().unwrap_or_else(|| "login".to_string()) + ); + read(&app, ¶ms, Lookup::Search { org, pattern }).await +} + +/// `GET /orgs/{org}/recent`: the [`Lookup::Recent`] lookup. +async fn recent( + State(app): State, + Path(org): Path, + Query(params): Query, +) -> Reply { + let lookup = Lookup::Recent { + org, + limit: params.limit.unwrap_or(5), + }; + read(&app, ¶ms, lookup).await +} + +/// `GET /orgs/{org}/tickets/{id}`: the [`Lookup::OrgTicket`] lookup. +async fn org_ticket( + State(app): State, + Path((org, id)): Path<(i64, i64)>, + Query(params): Query, +) -> Reply { + read(&app, ¶ms, Lookup::OrgTicket { org, id }).await +} + +/// `GET /orgs/{org}/unassigned`: the [`Lookup::Unassigned`] lookup. +async fn unassigned( + State(app): State, + Path(org): Path, + Query(params): Query, +) -> Reply { + read(&app, ¶ms, Lookup::Unassigned { org }).await +} + +/// `GET /orgs/{org}/escalated`: the [`Lookup::Escalated`] lookup. +async fn escalated( + State(app): State, + Path(org): Path, + Query(params): Query, +) -> Reply { + read(&app, ¶ms, Lookup::Escalated { org }).await +} + +/// `GET /orgs/{org}/open`: the [`Lookup::OpenExcept`] lookup. +async fn open_except( + State(app): State, + Path(org): Path, + Query(params): Query, +) -> Reply { + let except = ReadParams::id_list(params.except.as_deref(), vec![1, 2, 3]); + read(&app, ¶ms, Lookup::OpenExcept { org, except }).await +} + +/// `GET /orgs/{org}/created`: the [`Lookup::CreatedBetween`] lookup. +async fn created_between( + State(app): State, + Path(org): Path, + Query(params): Query, +) -> Reply { + let lookup = Lookup::CreatedBetween { + org, + from: ReadParams::instant( + params.from.as_deref(), + Utc.with_ymd_and_hms(2026, 10, 1, 9, 0, 0).unwrap(), + ), + to: ReadParams::instant( + params.to.as_deref(), + Utc.with_ymd_and_hms(2026, 10, 1, 12, 0, 0).unwrap(), + ), + }; + read(&app, ¶ms, lookup).await +} + +/// `GET /tickets`: the [`Lookup::Many`] lookup, a hundred ids unless +/// `ids` names others. +async fn many(State(app): State, Query(params): Query) -> Reply { + let ids = ReadParams::id_list(params.ids.as_deref(), (1..=100).collect()); + read(&app, ¶ms, Lookup::Many { ids }).await +} + +/// `GET /orgs/{org}/queue/end`: the [`Lookup::QueueEnd`] lookup. +async fn queue_end( + State(app): State, + Path(org): Path, + Query(params): Query, +) -> Reply { + read(&app, ¶ms, Lookup::QueueEnd { org }).await +} + +/// `GET /orgs/{org}/nothing`: the [`Lookup::Nothing`] lookup. +async fn nothing( + State(app): State, + Path(org): Path, + Query(params): Query, +) -> Reply { + read(&app, ¶ms, Lookup::Nothing { org }).await +} + +/// `GET /orgs/{org}/tickets`: the [`Lookup::AllTickets`] lookup. +async fn all_tickets( + State(app): State, + Path(org): Path, + Query(params): Query, +) -> Reply { + read(&app, ¶ms, Lookup::AllTickets { org }).await +} + +/// `GET /tickets/{id}/authors`: the [`Lookup::CommentAuthors`] lookup. +async fn comment_authors( + State(app): State, + Path(id): Path, + Query(params): Query, +) -> Reply { + read(&app, ¶ms, Lookup::CommentAuthors { ticket: id }).await +} + +/// `GET /orgs/{org}/skilled`: the [`Lookup::Skilled`] lookup. +async fn skilled( + State(app): State, + Path(org): Path, + Query(params): Query, +) -> Reply { + let skill = params + .skill + .clone() + .unwrap_or_else(|| "billing".to_string()); + read(&app, ¶ms, Lookup::Skilled { org, skill }).await +} + +/// `GET /orgs/{org}/statuses`: the [`Lookup::Statuses`] lookup. +async fn statuses( + State(app): State, + Path(org): Path, + Query(params): Query, +) -> Reply { + read(&app, ¶ms, Lookup::Statuses { org }).await +} + +/// `GET /orgs/{org}/workload`: the [`Lookup::Workload`] lookup. +async fn workload( + State(app): State, + Path(org): Path, + Query(params): Query, +) -> Reply { + read(&app, ¶ms, Lookup::Workload { org }).await +} + +/// `GET /orgs/{org}/activity`: the [`Lookup::LastActivity`] lookup. +async fn last_activity( + State(app): State, + Path(org): Path, + Query(params): Query, +) -> Reply { + read(&app, ¶ms, Lookup::LastActivity { org }).await +} + +/// `GET /orgs/{org}/kind/{kind}`: the [`Lookup::OfKind`] lookup. +async fn of_kind( + State(app): State, + Path((org, kind)): Path<(i64, String)>, + Query(params): Query, +) -> Reply { + let kind = TicketKind::parse(&kind) + .ok_or_else(|| (StatusCode::BAD_REQUEST, format!("no kind `{kind}`")))?; + read(&app, ¶ms, Lookup::OfKind { org, kind }).await +} + +/// The body of a status change. +#[derive(Deserialize)] +struct StatusChange { + status: TicketStatus, +} + +/// A write through the main pool: the cache marks it, and for the +/// read-your-writes window every read through the cache goes to +/// Postgres, so the caller sees its own change at once. +async fn set_status( + State(app): State, + Path(id): Path, + Json(change): Json, +) -> Reply { + let mut conn = app + .main + .get() + .await + .map_err(|error| (StatusCode::SERVICE_UNAVAILABLE, error.to_string()))?; + let changed = diesel::update(tickets::table.filter(tickets::id.eq(id))) + .set(( + tickets::status.eq(change.status), + tickets::updated_at.eq(diesel::dsl::now), + )) + .execute(&mut conn) + .await + .map_err(|error| (StatusCode::INTERNAL_SERVER_ERROR, error.to_string()))?; + if changed == 0 { + return Err((StatusCode::NOT_FOUND, format!("no ticket {id}"))); + } + Ok(Json( + json!({"id": id, "status": change.status, "read_your_writes_ms": diesel_streamgres::read_your_writes_from_env().as_millis()}), + )) +} + +/// For handlers that take a map of parameters rather than the struct. +#[allow(dead_code)] +fn param<'a>(params: &'a HashMap, name: &str) -> Option<&'a str> { + params.get(name).map(String::as_str) +} + +impl IntoResponse for Lookup { + /// The lookup as JSON. + fn into_response(self) -> axum::response::Response { + Json(self).into_response() + } +} diff --git a/crates/streamgres-demo/src/bench.rs b/crates/streamgres-demo/src/bench.rs new file mode 100644 index 0000000..e3df081 --- /dev/null +++ b/crates/streamgres-demo/src/bench.rs @@ -0,0 +1,92 @@ +//! Latency per route, in process: no HTTP in the way, just the wrapper. +//! +//! Every task takes a pooled connection, pins the route on it, and runs +//! the lookup `requests / concurrency` times, timing each. The first +//! run per route is the warm-up (the subscribe and the fill) and is +//! not counted. + +use std::time::{Duration, Instant}; + +use diesel_streamgres::Route; + +use crate::pools::Stack; +use crate::queries::Lookup; + +/// One route's latencies over a run. +#[derive(Debug, Clone)] +pub struct Report { + pub route: Route, + pub requests: usize, + pub p50: Duration, + pub p95: Duration, + pub p99: Duration, + pub max: Duration, + pub elapsed: Duration, +} + +impl Report { + /// Requests per second over the run. + pub fn per_second(&self) -> f64 { + self.requests as f64 / self.elapsed.as_secs_f64().max(1e-9) + } +} + +/// Times `requests` reads of `lookup` on `route`, `concurrency` at a +/// time, after one warm-up read (the subscribe and the fill) so the +/// measurement is of the steady state on both routes. +pub async fn run( + stack: &Stack, + lookup: &Lookup, + route: Route, + concurrency: usize, + requests: usize, +) -> Result { + let per_task = requests.div_ceil(concurrency.max(1)); + { + let mut conn = stack.main.get().await.map_err(|error| error.to_string())?; + let pool_route = conn.route(); + conn.set_route(route); + let warmed = lookup.run_async(&mut conn).await; + conn.set_route(pool_route); + warmed.map_err(|error| error.to_string())?; + } + let started = Instant::now(); + let mut tasks = Vec::with_capacity(concurrency); + for _ in 0..concurrency.max(1) { + let pool = stack.main.clone(); + let lookup = lookup.clone(); + tasks.push(tokio::spawn(async move { + let mut samples = Vec::with_capacity(per_task); + let mut conn = pool.get().await.map_err(|error| error.to_string())?; + let pool_route = conn.route(); + conn.set_route(route); + let mut outcome = Ok(()); + for _ in 0..per_task { + let at = Instant::now(); + if let Err(error) = lookup.run_async(&mut conn).await { + outcome = Err(error.to_string()); + break; + } + samples.push(at.elapsed()); + } + conn.set_route(pool_route); + outcome.map(|()| samples) + })); + } + let mut samples = Vec::with_capacity(requests); + for task in tasks { + samples.extend(task.await.map_err(|error| error.to_string())??); + } + let elapsed = started.elapsed(); + samples.sort(); + let at = |q: f64| samples[((samples.len() - 1) as f64 * q) as usize]; + Ok(Report { + route, + requests: samples.len(), + p50: at(0.50), + p95: at(0.95), + p99: at(0.99), + max: *samples.last().unwrap_or(&Duration::ZERO), + elapsed, + }) +} diff --git a/crates/streamgres-demo/src/index.html b/crates/streamgres-demo/src/index.html new file mode 100644 index 0000000..8cc5c13 --- /dev/null +++ b/crates/streamgres-demo/src/index.html @@ -0,0 +1,91 @@ + + + + + +Streamgres ticket desk + + + +

Streamgres ticket desk

+

Reads go through the Diesel wrapper; the engine keeps the cache fresh. A background writer makes a change every second; /verify compares every read with Postgres.

+ +
+ + + + +
lookupexerciseslastsourceµs
+ +

Engine:

+ + + + diff --git a/crates/streamgres-demo/src/lag.rs b/crates/streamgres-demo/src/lag.rs new file mode 100644 index 0000000..68db330 --- /dev/null +++ b/crates/streamgres-demo/src/lag.rs @@ -0,0 +1,131 @@ +//! Propagation lag: how long a commit takes to reach the cached answer, +//! under a write rate. +//! +//! A sentinel is written to one ticket through the plain pool, and the +//! cached lookup of that ticket is polled until it carries it; the +//! elapsed time is one sample. Meanwhile a load task cycles the status +//! of random tickets of the second organization at `rate` per second, +//! so the engine has other changes to carry. This is the number people +//! ask about ("how stale is the cache?"), measured instead of estimated. + +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::time::{Duration, Instant}; + +use diesel::sql_query; +use diesel::sql_types::{BigInt, Text}; +use diesel_async::RunQueryDsl; +use rand::Rng; + +use crate::pools::Stack; +use crate::queries::Lookup; +use crate::seed; + +/// The lag samples, and the load that ran under them. +#[derive(Debug, Clone)] +pub struct Report { + pub samples: usize, + pub p50: Duration, + pub p95: Duration, + pub p99: Duration, + pub max: Duration, + pub writes: u64, + pub elapsed: Duration, +} + +impl Report { + /// The load actually applied, in writes per second. + pub fn writes_per_second(&self) -> f64 { + self.writes as f64 / self.elapsed.as_secs_f64().max(1e-9) + } +} + +/// The ticket the sentinel goes to. +const SENTINEL: i64 = 1; + +/// Measures for `duration` under `rate` background writes per second. +pub async fn run(stack: &Stack, rate: u32, duration: Duration) -> Result { + let stop = Arc::new(AtomicBool::new(false)); + let writes = Arc::new(AtomicU64::new(0)); + let load = (rate > 0).then(|| { + let pool = stack.plain.clone(); + let stop = Arc::clone(&stop); + let writes = Arc::clone(&writes); + let per_org = seed::tickets_per_org(seed::scale()); + let every = Duration::from_secs_f64(1.0 / f64::from(rate)); + tokio::spawn(async move { + let mut next = tokio::time::Instant::now(); + while !AtomicBool::load(&stop, Ordering::Relaxed) { + next += every; + tokio::time::sleep_until(next).await; + let Ok(mut conn) = pool.get().await else { + continue; + }; + let id = per_org + rand::rng().random_range(1..=per_org); + let done = sql_query( + "UPDATE tickets SET status = (CASE status + WHEN 'open' THEN 'pending' WHEN 'pending' THEN 'closed' ELSE 'open' + END)::ticket_status, updated_at = now() WHERE id = $1", + ) + .bind::(id) + .execute(&mut conn) + .await; + if done.is_ok() { + writes.fetch_add(1, Ordering::Relaxed); + } + } + }) + }); + + let lookup = Lookup::Ticket { id: SENTINEL }; + let mut reader = stack.main.get().await.map_err(|error| error.to_string())?; + let mut writer = stack.plain.get().await.map_err(|error| error.to_string())?; + lookup + .run_async(&mut reader) + .await + .map_err(|error| error.to_string())?; + let started = Instant::now(); + let mut samples = Vec::new(); + let mut n = 0u64; + while started.elapsed() < duration { + n += 1; + let title = format!("lag-{}-{n}", std::process::id()); + sql_query("UPDATE tickets SET title = $1, updated_at = now() WHERE id = $2") + .bind::(&title) + .bind::(SENTINEL) + .execute(&mut writer) + .await + .map_err(|error| error.to_string())?; + let committed = Instant::now(); + loop { + let rows = lookup + .run_async(&mut reader) + .await + .map_err(|error| error.to_string())?; + if rows[0]["title"].as_str() == Some(title.as_str()) { + samples.push(committed.elapsed()); + break; + } + if committed.elapsed() > Duration::from_secs(10) { + return Err(format!("the sentinel {title} never reached the cache")); + } + tokio::time::sleep(Duration::from_millis(1)).await; + } + tokio::time::sleep(Duration::from_millis(200)).await; + } + stop.store(true, Ordering::Relaxed); + if let Some(load) = load { + let _ = load.await; + } + samples.sort(); + let at = |q: f64| samples[((samples.len() - 1) as f64 * q) as usize]; + Ok(Report { + samples: samples.len(), + p50: at(0.50), + p95: at(0.95), + p99: at(0.99), + max: *samples.last().unwrap_or(&Duration::ZERO), + writes: AtomicU64::load(&writes, Ordering::Relaxed), + elapsed: started.elapsed(), + }) +} diff --git a/crates/streamgres-demo/src/lib.rs b/crates/streamgres-demo/src/lib.rs new file mode 100644 index 0000000..1a12590 --- /dev/null +++ b/crates/streamgres-demo/src/lib.rs @@ -0,0 +1,33 @@ +//! A ticket desk on the Streamgres Diesel wrapper. +//! +//! Two things at once. For someone who found Streamgres: a backend +//! whose reads come from the engine's cache, runnable with one compose +//! file, with an HTTP API to curl and a stats page to watch. For the +//! people working on the wrapper: a fixed list of query shapes +//! ([`queries::Lookup`]) that the scenario test drives through both +//! wrappers and both routes, and the bench times. +//! +//! The modules, in the order to read them: +//! +//! - [`schema`] and [`models`]: five tables and the rows they yield. +//! - [`queries`]: the reads, each a named shape, runnable sync or async. +//! - [`seed`]: the deterministic dataset. +//! - [`pools`]: the connections and how each is routed. +//! - [`writer`]: the background changes that keep the demo moving. +//! - [`verify`]: the demo checking itself, every lookup on both routes. +//! - [`migrate`]: the schema changes, one of each kind. +//! - [`api`]: the HTTP face, including `/stats` and `/verify`. +//! - [`bench`]: latency per route, in process. +//! - [`lag`]: how long a commit takes to reach the cache, under load. + +pub mod api; +pub mod bench; +pub mod lag; +pub mod migrate; +pub mod models; +pub mod pools; +pub mod queries; +pub mod schema; +pub mod seed; +pub mod verify; +pub mod writer; diff --git a/crates/streamgres-demo/src/main.rs b/crates/streamgres-demo/src/main.rs new file mode 100644 index 0000000..bf79fa4 --- /dev/null +++ b/crates/streamgres-demo/src/main.rs @@ -0,0 +1,319 @@ +//! `demo serve | seed | bench | verify | lag | migrate`. +//! +//! Environment: +//! +//! | variable | meaning | default | +//! |-----------------------|------------------------------------------------|----------------| +//! | `DATABASE_URL` | the Postgres the engine follows | required | +//! | `STREAMGRES_S2S_URL` | `ws://host:port//s2s/v1` | required | +//! | `DEMO_ADDR` | where `serve` listens | `0.0.0.0:8080` | +//! | `DEMO_WRITER_MS` | the background writer's tick; 0 turns it off | `1000` | +//! | `DEMO_WRITER_MIX` | the writer's operations, comma separated | all of them | +//! | `DEMO_VERIFY_MS` | `serve` also verifies every lookup this often; 0 turns it off | `0` | +//! | `DEMO_SCALE` | multiplies the seed's tickets, comments and attachments | `1` | +//! +//! Plus the wrapper's own `STREAMGRES_CACHE_*`, `STREAMGRES_S2S_*` and +//! `STREAMGRES_READ_YOUR_WRITES_MS`. + +use std::sync::Arc; +use std::time::Duration; + +use clap::{Parser, Subcommand}; +use diesel_streamgres::Route; +use streamgres_demo::migrate::Step; +use streamgres_demo::pools::Stack; +use streamgres_demo::queries::Lookup; +use streamgres_demo::{api, bench, lag, seed, verify, writer}; + +/// The demo's command line. +#[derive(Parser)] +#[command( + name = "demo", + about = "A ticket desk on the Streamgres Diesel wrapper" +)] +struct Cli { + #[command(subcommand)] + command: Command, +} + +/// What the demo does. +#[derive(Subcommand)] +enum Command { + /// Serves the HTTP API and the stats page. + Serve, + /// Creates the schema and loads the dataset; safe to run again. + Seed, + /// Times a lookup on each route, in process. + Bench { + /// A lookup by name (`demo bench --list`). + #[arg(long, default_value = "queue")] + lookup: String, + /// `streamgres`, `postgres` or `both`. + #[arg(long, default_value = "both")] + route: String, + #[arg(long, default_value_t = 16)] + concurrency: usize, + #[arg(long, default_value_t = 5000)] + requests: usize, + /// Lists the lookups and exits. + #[arg(long)] + list: bool, + }, + /// Reads every lookup through the cache and through Postgres, + /// repeatedly, and reports every pair that never agreed. Exits 1 + /// when any did not. + Verify { + /// Milliseconds between passes. + #[arg(long, default_value_t = 250)] + every: u64, + /// Seconds to run for. + #[arg(long = "for", default_value_t = 30)] + seconds: u64, + /// Milliseconds a differing pair may take to agree before it is + /// a mismatch. + #[arg(long, default_value_t = 2000)] + settle: u64, + /// Run the background writer meanwhile, at this tick in + /// milliseconds; 0 for none. + #[arg(long, default_value_t = 50)] + writer_ms: u64, + }, + /// Measures how long a commit takes to reach the cached answer, + /// under a write rate. + Lag { + /// Background writes per second. + #[arg(long, default_value_t = 20)] + rate: u32, + /// Seconds to measure for. + #[arg(long = "for", default_value_t = 10)] + seconds: u64, + }, + /// Runs a schema change: `widen`, `new-type` or `drop`. + Migrate { step: String }, +} + +/// A required variable from the environment. +fn env(name: &str) -> Result { + std::env::var(name).map_err(|_| format!("{name} is not set")) +} + +/// Runs the command; exits 1 with the reason when it fails. +#[tokio::main] +async fn main() { + if let Err(error) = run().await { + eprintln!("demo: {error}"); + std::process::exit(1); + } +} + +/// A connected stack, or why not. +async fn connect() -> Result { + let stack = Stack::connect(&env("DATABASE_URL")?, &env("STREAMGRES_S2S_URL")?).await?; + if !stack.wait_for_engine(Duration::from_secs(10)).await { + return Err("the engine is not reachable".to_string()); + } + Ok(stack) +} + +/// A number from the environment, or the default. +fn env_number(name: &str, default: u64) -> u64 { + std::env::var(name) + .ok() + .and_then(|value| value.parse::().ok()) + .unwrap_or(default) +} + +/// Runs the command the line names. +async fn run() -> Result<(), String> { + match Cli::parse().command { + Command::Seed => { + let database_url = env("DATABASE_URL")?; + let mut conn = Stack::plain_sync(&database_url).map_err(|error| error.to_string())?; + seed::run(&mut conn).map_err(|error| error.to_string())?; + println!("seeded: {}", seed::summary(seed::scale())); + Ok(()) + } + Command::Serve => { + let stack = + Arc::new(Stack::connect(&env("DATABASE_URL")?, &env("STREAMGRES_S2S_URL")?).await?); + for warning in stack.engine.warnings() { + eprintln!("demo: {warning}"); + } + if !stack.wait_for_engine(Duration::from_secs(10)).await { + eprintln!( + "demo: the engine is not reachable yet; reads go to Postgres until it is" + ); + } + let every = env_number("DEMO_WRITER_MS", 1000); + if every > 0 { + writer::spawn( + stack.plain.clone(), + Duration::from_millis(every), + writer::mix_from_env(), + Arc::clone(&stack.writer), + ); + } + let verify_every = env_number("DEMO_VERIFY_MS", 0); + if verify_every > 0 { + let stack = Arc::clone(&stack); + tokio::spawn(async move { + let lookups = Lookup::demo(); + loop { + tokio::time::sleep(Duration::from_millis(verify_every)).await; + let outcome = verify::once(&stack, &lookups, Duration::from_secs(2)).await; + for mismatch in &outcome.mismatches { + eprintln!("verify: {} never agreed", mismatch.lookup); + } + } + }); + } + let addr = std::env::var("DEMO_ADDR").unwrap_or_else(|_| "0.0.0.0:8080".to_string()); + let listener = tokio::net::TcpListener::bind(&addr) + .await + .map_err(|error| format!("bind {addr}: {error}"))?; + println!("demo: serving on http://{addr} (stats at /stats, a check at /verify)"); + axum::serve(listener, api::router(stack)) + .await + .map_err(|error| error.to_string()) + } + Command::Bench { + lookup, + route, + concurrency, + requests, + list, + } => { + if list { + for lookup in Lookup::demo() { + println!("{:<16} {}", lookup.name(), lookup.exercises()); + } + return Ok(()); + } + let lookup = Lookup::named(&lookup) + .ok_or_else(|| format!("no lookup `{lookup}`; see --list"))?; + let stack = connect().await?; + let routes: Vec = match route.as_str() { + "streamgres" => vec![Route::Streamgres], + "postgres" => vec![Route::Postgres], + _ => vec![Route::Postgres, Route::Streamgres], + }; + println!( + "{} ({}), {concurrency} tasks, {requests} requests per route", + lookup.name(), + lookup.exercises() + ); + println!( + "{:<12} {:>9} {:>9} {:>9} {:>9} {:>10}", + "route", "p50", "p95", "p99", "max", "req/s" + ); + for route in routes { + let report = bench::run(&stack, &lookup, route, concurrency, requests).await?; + println!( + "{:<12} {:>9} {:>9} {:>9} {:>9} {:>10.0}", + format!("{route:?}").to_lowercase(), + micros(report.p50), + micros(report.p95), + micros(report.p99), + micros(report.max), + report.per_second() + ); + } + let stats = api::stats_json(&stack); + println!("cache {}", stats["cache"]); + Ok(()) + } + Command::Verify { + every, + seconds, + settle, + writer_ms, + } => { + let stack = connect().await?; + if writer_ms > 0 { + writer::spawn( + stack.plain.clone(), + Duration::from_millis(writer_ms), + writer::mix_from_env(), + Arc::clone(&stack.writer), + ); + } + let lookups = Lookup::demo(); + let started = std::time::Instant::now(); + let mut mismatches = 0u64; + let mut last_report = std::time::Instant::now(); + while started.elapsed() < Duration::from_secs(seconds) { + let outcome = verify::once(&stack, &lookups, Duration::from_millis(settle)).await; + for mismatch in &outcome.mismatches { + mismatches += 1; + println!( + "MISMATCH {}: cache {} | postgres {}", + mismatch.lookup, + truncate(&mismatch.cached.to_string()), + truncate(&mismatch.postgres.to_string()) + ); + } + if last_report.elapsed() >= Duration::from_secs(5) { + last_report = std::time::Instant::now(); + println!( + "{:>4}s verify {} writer {}", + started.elapsed().as_secs(), + stack.verify.json(), + stack.writer.json() + ); + } + tokio::time::sleep(Duration::from_millis(every)).await; + } + println!("verify {}", stack.verify.json()); + println!("writer {}", stack.writer.json()); + if mismatches > 0 { + return Err(format!("{mismatches} pair(s) never agreed")); + } + println!("every pair agreed"); + Ok(()) + } + Command::Lag { rate, seconds } => { + let stack = connect().await?; + let report = lag::run(&stack, rate, Duration::from_secs(seconds)).await?; + println!( + "commit to cached answer over {} samples, {:.0} writes/s of load:", + report.samples, + report.writes_per_second() + ); + println!( + "p50 {} p95 {} p99 {} max {}", + micros(report.p50), + micros(report.p95), + micros(report.p99), + micros(report.max) + ); + Ok(()) + } + Command::Migrate { step } => { + let step = Step::parse(&step).ok_or_else(|| { + format!("no migration `{step}`; the steps are widen, new-type, drop") + })?; + let mut conn = + Stack::plain_sync(&env("DATABASE_URL")?).map_err(|error| error.to_string())?; + step.run(&mut conn).map_err(|error| error.to_string())?; + println!("{step:?} done: {}", step.expectation()); + Ok(()) + } + } +} + +/// A duration in microseconds, for a report line. +fn micros(duration: Duration) -> String { + format!("{}µs", duration.as_micros()) +} + +/// The first part of a long JSON answer, for a report line. +fn truncate(text: &str) -> String { + if text.len() <= 200 { + return text.to_string(); + } + let mut end = 200; + while !text.is_char_boundary(end) { + end -= 1; + } + format!("{}…", &text[..end]) +} diff --git a/crates/streamgres-demo/src/migrate.rs b/crates/streamgres-demo/src/migrate.rs new file mode 100644 index 0000000..398dc88 --- /dev/null +++ b/crates/streamgres-demo/src/migrate.rs @@ -0,0 +1,75 @@ +//! Schema changes, one of each kind the engine treats differently. +//! +//! The engine hears of a migration through the schema-change trigger, +//! in the WAL, at the position of the change. A column added with a +//! constant default it absorbs: held rows are widened and nothing +//! restarts (`widen`). A new enum type and a column of it it absorbs +//! too, and learns the type the first time a read binds it +//! (`new-type`). A column dropped it cannot follow: the engine stops, +//! and restarted serves the schema as it is (`drop`, which also undoes +//! the other two). + +use diesel::QueryResult; +use diesel::connection::SimpleConnection; + +/// One migration. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Step { + /// `tags text NOT NULL DEFAULT ''` on tickets: absorbed. + Widen, + /// The enum `ticket_kind` and `kind ticket_kind` on tickets, filled + /// in: absorbed, and the type learnt on its first bind. + NewType, + /// Drops both columns and the type: the engine stops. + Drop, +} + +impl Step { + /// A step by its command-line name. + pub fn parse(name: &str) -> Option { + match name { + "widen" => Some(Step::Widen), + "new-type" => Some(Step::NewType), + "drop" => Some(Step::Drop), + _ => None, + } + } + + /// The SQL the step runs. + pub fn sql(self) -> &'static str { + match self { + Step::Widen => "ALTER TABLE tickets ADD COLUMN IF NOT EXISTS tags text NOT NULL DEFAULT ''", + Step::NewType => " + DO $$ BEGIN + CREATE TYPE ticket_kind AS ENUM ('bug', 'task', 'question'); + EXCEPTION WHEN duplicate_object OR unique_violation THEN NULL; END $$; + ALTER TABLE tickets ADD COLUMN IF NOT EXISTS kind ticket_kind; + UPDATE tickets SET kind = (ARRAY['bug', 'task', 'question'])[1 + id % 3]::ticket_kind + WHERE kind IS NULL", + Step::Drop => " + ALTER TABLE tickets DROP COLUMN IF EXISTS tags; + ALTER TABLE tickets DROP COLUMN IF EXISTS kind; + DROP TYPE IF EXISTS ticket_kind", + } + } + + /// What to expect of the engine after the step. + pub fn expectation(self) -> &'static str { + match self { + Step::Widen => { + "the engine absorbs the column; cached reads keep working and keep being served from rows" + } + Step::NewType => { + "the engine absorbs the column, and learns the type the first time `of_kind` binds it; from then on it is served from rows" + } + Step::Drop => { + "the engine cannot follow a dropped column: it stops, and restarted serves the schema as it is; reads go to Postgres meanwhile" + } + } + } + + /// Runs the step on a blocking connection. + pub fn run(self, conn: &mut C) -> QueryResult<()> { + conn.batch_execute(self.sql()) + } +} diff --git a/crates/streamgres-demo/src/models.rs b/crates/streamgres-demo/src/models.rs new file mode 100644 index 0000000..9d3edd0 --- /dev/null +++ b/crates/streamgres-demo/src/models.rs @@ -0,0 +1,215 @@ +//! The rows the queries return, and the two Postgres enums. + +use std::io::Write; + +use chrono::{DateTime, Utc}; +use diesel::data_types::PgInterval; +use diesel::deserialize::{self, FromSql, FromSqlRow}; +use diesel::expression::AsExpression; +use diesel::pg::{Pg, PgValue}; +use diesel::prelude::*; +use diesel::serialize::{self, IsNull, Output, ToSql}; +use serde::{Deserialize, Serialize}; + +use crate::schema::{agents, attachments, comments, organizations, tickets}; + +/// The `ticket_status` enum type. +#[derive(diesel::sql_types::SqlType, diesel::query_builder::QueryId)] +#[diesel(postgres_type(name = "ticket_status"))] +pub struct TicketStatusType; + +/// The `ticket_priority` enum type. +#[derive(diesel::sql_types::SqlType, diesel::query_builder::QueryId)] +#[diesel(postgres_type(name = "ticket_priority"))] +pub struct TicketPriorityType; + +/// The enum type a migration adds after the engine started +/// (`demo migrate new-type`): the engine learns it the first time a +/// read binds it. +#[derive(diesel::sql_types::SqlType, diesel::query_builder::QueryId)] +#[diesel(postgres_type(name = "ticket_kind"))] +pub struct TicketKindType; + +/// A Postgres enum as a Rust enum: its labels on the wire, lower case +/// in JSON. +macro_rules! pg_enum { + ($name:ident, $sql:ty, [$($variant:ident => $label:literal),+ $(,)?]) => { + #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, AsExpression, FromSqlRow, Serialize, Deserialize)] + #[diesel(sql_type = $sql)] + #[serde(rename_all = "lowercase")] + pub enum $name { + $($variant),+ + } + + impl $name { + pub const ALL: &'static [$name] = &[$($name::$variant),+]; + + /// The label Postgres knows the variant by. + pub fn label(self) -> &'static str { + match self { + $($name::$variant => $label),+ + } + } + + /// The variant a label names. + pub fn parse(label: &str) -> Option<$name> { + match label { + $($label => Some($name::$variant),)+ + _ => None, + } + } + } + + impl ToSql<$sql, Pg> for $name { + /// The label, as Postgres's binary format sends an enum. + fn to_sql<'b>(&'b self, out: &mut Output<'b, '_, Pg>) -> serialize::Result { + out.write_all(self.label().as_bytes())?; + Ok(IsNull::No) + } + } + + impl FromSql<$sql, Pg> for $name { + /// The variant the label names. + fn from_sql(value: PgValue<'_>) -> deserialize::Result { + let text = std::str::from_utf8(value.as_bytes())?; + Self::parse(text).ok_or_else(|| format!("not a {}: {text}", stringify!($name)).into()) + } + } + }; +} + +pg_enum!(TicketStatus, TicketStatusType, [Open => "open", Pending => "pending", Closed => "closed"]); +pg_enum!(TicketPriority, TicketPriorityType, [Low => "low", Normal => "normal", High => "high"]); +pg_enum!(TicketKind, TicketKindType, [Bug => "bug", Task => "task", Question => "question"]); + +impl TicketStatus { + /// The next status the background writer moves a ticket to. + pub fn next(self) -> TicketStatus { + match self { + TicketStatus::Open => TicketStatus::Pending, + TicketStatus::Pending => TicketStatus::Closed, + TicketStatus::Closed => TicketStatus::Open, + } + } +} + +/// An organization. +#[derive(Queryable, Selectable, Serialize, Debug, Clone, PartialEq)] +#[diesel(table_name = organizations)] +pub struct Organization { + pub id: i64, + pub slug: String, + pub name: String, +} + +/// An agent of an organization. +#[derive(Queryable, Selectable, Serialize, Debug, Clone, PartialEq)] +#[diesel(table_name = agents)] +pub struct Agent { + pub id: i64, + pub org_id: i64, + pub name: String, + pub skills: Vec, +} + +/// A ticket. +#[derive(Queryable, Selectable, Serialize, Debug, Clone, PartialEq)] +#[diesel(table_name = tickets)] +pub struct Ticket { + pub id: i64, + pub org_id: i64, + pub agent_id: Option, + pub title: String, + pub status: TicketStatus, + pub priority: TicketPriority, + pub created_at: DateTime, + pub updated_at: DateTime, +} + +/// A comment on a ticket. +#[derive(Queryable, Selectable, Serialize, Debug, Clone, PartialEq)] +#[diesel(table_name = comments)] +pub struct Comment { + pub id: i64, + pub ticket_id: i64, + pub agent_id: Option, + pub body: String, + pub created_at: DateTime, +} + +/// A file attached to a ticket. +#[derive(Queryable, Selectable, Serialize, Debug, Clone, PartialEq)] +#[diesel(table_name = attachments)] +pub struct Attachment { + pub id: i64, + pub ticket_id: i64, + pub name: String, + #[serde(serialize_with = "hex")] + pub data: Vec, +} + +/// A ticket with its `sla`, an `interval` the wrapper cannot encode: +/// the query reading it is refilled from Postgres on each change. +#[derive(Queryable, Selectable, Serialize, Debug, Clone, PartialEq)] +#[diesel(table_name = crate::schema::tickets_with_kind)] +pub struct TicketOfKind { + pub id: i64, + pub org_id: i64, + pub kind: Option, +} + +/// A comment with the name of the agent who wrote it, `None` where +/// nobody did: what the left join returns. +#[derive(Queryable, Serialize, Debug, Clone, PartialEq)] +pub struct CommentAuthor { + pub comment_id: i64, + pub body: String, + pub agent: Option, +} + +/// How many tickets each agent holds, `None` for the unassigned: the +/// `GROUP BY` the engine refuses. +#[derive(QueryableByName, Serialize, Debug, Clone, PartialEq)] +pub struct Workload { + #[diesel(sql_type = diesel::sql_types::Nullable)] + pub agent_id: Option, + #[diesel(sql_type = diesel::sql_types::BigInt)] + pub tickets: i64, +} + +/// A ticket's title and its comment, as the join returns them. +#[derive(Queryable, Serialize, Debug, Clone, PartialEq)] +pub struct TicketComment { + pub ticket_id: i64, + pub title: String, + pub comment_id: i64, + pub body: String, + pub created_at: DateTime, +} + +/// A ticket with its service-level agreement, an `interval`: the one +/// column the engine cannot ship rows for. +#[derive(Debug, Clone, PartialEq, Serialize)] +pub struct TicketSla { + pub id: i64, + pub title: String, + /// The interval in seconds, ignoring its month part. + pub sla_seconds: Option, +} + +impl TicketSla { + /// The row as read, its interval in seconds. + pub fn from_row((id, title, sla): (i64, String, Option)) -> TicketSla { + TicketSla { + id, + title, + sla_seconds: sla.map(|sla| sla.microseconds / 1_000_000 + i64::from(sla.days) * 86_400), + } + } +} + +/// Bytes as hex text in JSON. +fn hex(bytes: &[u8], serializer: S) -> Result { + let text: String = bytes.iter().map(|byte| format!("{byte:02x}")).collect(); + serializer.serialize_str(&text) +} diff --git a/crates/streamgres-demo/src/pools.rs b/crates/streamgres-demo/src/pools.rs new file mode 100644 index 0000000..0b037c8 --- /dev/null +++ b/crates/streamgres-demo/src/pools.rs @@ -0,0 +1,137 @@ +//! The connections the desk uses, and how each is routed. +//! +//! Three pools against one database: +//! +//! - `main`: the application's pool, through the cache, routed +//! Streamgres. Every read in the HTTP API goes through it. +//! - `reporting`: the same wrapper, routed Postgres by its URL, the way +//! a pool against a primary would be. Shows the per-pool flag. +//! - `plain`: diesel-async's own connection, no wrapper. The background +//! writer uses it, so its changes reach the cache the way another +//! service's would: through the engine, not the read-your-writes +//! window. +//! +//! And a small blocking pool of the sync wrapper, sharing the same +//! cache, for the paths that are not async. + +use std::sync::Arc; +use std::time::Duration; + +use diesel::PgConnection; +use diesel::r2d2::{ConnectionManager, Pool as SyncPool}; +use diesel_async::AsyncPgConnection; +use diesel_async::pooled_connection::AsyncDieselConnectionManager; +use diesel_async::pooled_connection::bb8::Pool; +use diesel_streamgres::cache::{Cache, Limits}; +use diesel_streamgres::engine::EngineClient; +use diesel_streamgres::{Attachment, Route, Settings, install}; + +use crate::queries::{AsyncConn, SyncConn}; +use crate::{verify, writer}; + +/// Everything a process holds: the cache, the engine client, the pools, +/// and the counters of the writer and the verify loop. +pub struct Stack { + pub cache: Arc, + pub engine: EngineClient, + pub main: Pool, + pub reporting: Pool, + pub plain: Pool, + pub sync: SyncPool>, + pub writer: Arc, + pub verify: Arc, +} + +/// The URL a pool is established with decides its route; the settings +/// read the application name off it. +pub fn pool_url(database_url: &str, application_name: &str) -> String { + let separator = if database_url.contains('?') { '&' } else { '?' }; + format!("{database_url}{separator}application_name={application_name}") +} + +/// Installs the process-wide settings once: one cache for the database, +/// a Postgres route for the reporting pool, Streamgres for the rest. +fn install_settings(cache: &Arc) { + let cache = Arc::clone(cache); + let _ = install(Settings::by_url(move |url| { + Some(Attachment { + cache: Arc::clone(&cache), + route: if url.contains("application_name=reporting") { + Route::Postgres + } else { + Route::Streamgres + }, + }) + })); +} + +impl Stack { + /// Connects to the engine and opens the pools. The engine's sockets + /// come up in the background; reads before that go live. + pub async fn connect(database_url: &str, s2s_url: &str) -> Result { + let engine = EngineClient::connect(s2s_url); + let cache = Arc::new(Cache::with_limits( + Limits::from_env(), + Arc::new(engine.clone()), + )); + engine.attach(&cache); + install_settings(&cache); + + let pool = |name: &str| { + let manager = + AsyncDieselConnectionManager::::new(pool_url(database_url, name)); + Pool::builder() + .max_size(16) + .connection_timeout(Duration::from_secs(10)) + .build(manager) + }; + let main = pool("main") + .await + .map_err(|error| format!("main pool: {error}"))?; + let reporting = pool("reporting") + .await + .map_err(|error| format!("reporting pool: {error}"))?; + let plain = Pool::builder() + .max_size(4) + .build(AsyncDieselConnectionManager::::new( + pool_url(database_url, "plain"), + )) + .await + .map_err(|error| format!("plain pool: {error}"))?; + let sync = SyncPool::builder() + .max_size(4) + .build(ConnectionManager::::new(pool_url( + database_url, + "sync", + ))) + .map_err(|error| format!("sync pool: {error}"))?; + Ok(Stack { + cache, + engine, + main, + reporting, + plain, + sync, + writer: Arc::default(), + verify: Arc::default(), + }) + } + + /// Waits for at least one engine socket, up to `wait`. + pub async fn wait_for_engine(&self, wait: Duration) -> bool { + let deadline = tokio::time::Instant::now() + wait; + while !self.engine.is_connected() { + if tokio::time::Instant::now() >= deadline { + return false; + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + true + } + + /// A plain, unwrapped blocking connection: for the seed. + pub fn plain_sync(database_url: &str) -> diesel::ConnectionResult { + use diesel::Connection; + PgConnection::establish(database_url) + } +} diff --git a/crates/streamgres-demo/src/queries.rs b/crates/streamgres-demo/src/queries.rs new file mode 100644 index 0000000..05b42b7 --- /dev/null +++ b/crates/streamgres-demo/src/queries.rs @@ -0,0 +1,741 @@ +//! The reads the ticket desk does, as one list: each is a named query +//! with a known shape, runnable through either wrapper. +//! +//! This list is the point of the crate. The HTTP API serves it, the +//! scenario test drives every entry through both wrappers and both +//! routes, the verify loop compares every entry across the routes, and +//! the bench times it. A shape that matters to the wrapper (an enum +//! bound as an array, an offset page, a join, a `bytea`, a column the +//! engine cannot ship, an operator it refuses) is here so it is +//! exercised the day it changes, not the day a user finds it. +//! +//! Each shape says what it expects of the engine ([`Lookup::expected`]): +//! served from the engine's rows, refilled from Postgres on each change +//! because a column cannot be shipped, or refused and served live. The +//! scenario test holds every shape to it, so a shape that drops from +//! rows mode to refills fails a test instead of passing more slowly. + +use chrono::{DateTime, TimeZone, Utc}; +use diesel::dsl::{AsSelect, InnerJoinQuerySource, LeftJoinQuerySource}; +use diesel::pg::Pg; +use diesel::prelude::*; +use diesel::query_builder::{BoxedSelectStatement, FromClause}; +use diesel::sql_types::{BigInt, Int8, Interval, Nullable, Text, Timestamptz}; +use diesel::{PgConnection, QueryResult}; +use diesel_streamgres::{AsyncSyncedConnection, SyncedConnection}; +use futures_util::future::{BoxFuture, FutureExt}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +use crate::models::{ + Agent, Attachment, CommentAuthor, Ticket, TicketComment, TicketKind, TicketOfKind, + TicketPriority, TicketSla, TicketStatus, TicketStatusType, Workload, +}; +use crate::schema::{agents, attachments, comments, tickets, tickets_with_kind}; + +/// The blocking wrapper the lookups run on. +pub type SyncConn = SyncedConnection; +/// The async wrapper the lookups run on. +pub type AsyncConn = AsyncSyncedConnection; + +/// Tickets per page of a queue. +pub const PAGE: i64 = 10; + +/// What a shape expects of the engine. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum Expected { + /// Subscribed, served from the engine's rows, every change applied + /// in place. + Rows, + /// Subscribed, but a projected column cannot be shipped, so each + /// change refills the entry from Postgres. + Refill, + /// Outside the engine's subset: refused, served by Postgres every + /// time. + Live, +} + +/// One read of the ticket desk. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(tag = "lookup", rename_all = "snake_case")] +pub enum Lookup { + /// One ticket by primary key. + Ticket { id: i64 }, + /// An organization's tickets in some statuses, newest first, one + /// page: an enum array bind, an order by timestamp, `LIMIT` and + /// `OFFSET`. + Queue { + org: i64, + statuses: Vec, + page: i64, + }, + /// How many tickets a queue has: a `COUNT(*)`. + QueueCount { + org: i64, + statuses: Vec, + }, + /// A ticket's comments with its title: an inner join, composed from + /// the engine's rows per table. + Comments { ticket: i64 }, + /// One attachment: a `bytea`, carried as hex by the engine. + Attachment { id: i64 }, + /// An organization's agents: a text array column. + Agents { org: i64 }, + /// Tickets with a service level: an `interval` column, which the + /// engine cannot ship rows for, so the query is refilled from + /// Postgres on every change instead. + Sla { org: i64 }, + /// Tickets whose title matches a pattern: `ILIKE`, outside the + /// engine's subset, so the read is refused and goes to Postgres + /// every time. + Search { org: i64, pattern: String }, + /// The organization's most recently updated tickets. + Recent { org: i64, limit: i64 }, + /// One ticket by its organization and id: a two-column key. + OrgTicket { org: i64, id: i64 }, + /// Tickets nobody holds: `IS NULL`. + Unassigned { org: i64 }, + /// Tickets that are high priority or pending: an `OR` of two enum + /// binds. + Escalated { org: i64 }, + /// Tickets not closed, a few excluded by id: `<>` and `NOT IN`. + OpenExcept { org: i64, except: Vec }, + /// Tickets created inside a window: two timestamp binds. + CreatedBetween { + org: i64, + from: DateTime, + to: DateTime, + }, + /// Many tickets by id: an `IN` list a hundred long. + Many { ids: Vec }, + /// A queue page far past the end: an empty result, subscribed, that + /// a change elsewhere leaves empty. + QueueEnd { org: i64 }, + /// A query that never has rows: subscribed, stays empty. + Nothing { org: i64 }, + /// Every ticket of an organization: the largest result; under a + /// small per-entry cap it is served but not kept. + AllTickets { org: i64 }, + /// A ticket's comments with the name of whoever wrote each: a left + /// join, `NULL` on the side with no agent. + CommentAuthors { ticket: i64 }, + /// Agents with a skill: the array operator `@>`, refused. + Skilled { org: i64, skill: String }, + /// The statuses in use: `DISTINCT`, refused. + Statuses { org: i64 }, + /// Tickets per agent: `GROUP BY`, refused. + Workload { org: i64 }, + /// When the organization was last active: `MAX`, an aggregate + /// other than `COUNT(*)`, refused. + LastActivity { org: i64 }, + /// Tickets of a kind, through the column and enum type `demo + /// migrate new-type` adds: before the migration the read fails + /// (no such column); after it the engine learns the type on the + /// first bind and serves it from rows. Not in the demo list. + OfKind { org: i64, kind: TicketKind }, +} + +impl Lookup { + /// The name the API, the bench and the stats page use. + pub fn name(&self) -> &'static str { + match self { + Lookup::Ticket { .. } => "ticket", + Lookup::Queue { .. } => "queue", + Lookup::QueueCount { .. } => "queue_count", + Lookup::Comments { .. } => "comments", + Lookup::Attachment { .. } => "attachment", + Lookup::Agents { .. } => "agents", + Lookup::Sla { .. } => "sla", + Lookup::Search { .. } => "search", + Lookup::Recent { .. } => "recent", + Lookup::OrgTicket { .. } => "org_ticket", + Lookup::Unassigned { .. } => "unassigned", + Lookup::Escalated { .. } => "escalated", + Lookup::OpenExcept { .. } => "open_except", + Lookup::CreatedBetween { .. } => "created_between", + Lookup::Many { .. } => "many", + Lookup::QueueEnd { .. } => "queue_end", + Lookup::Nothing { .. } => "nothing", + Lookup::AllTickets { .. } => "all_tickets", + Lookup::CommentAuthors { .. } => "comment_authors", + Lookup::Skilled { .. } => "skilled", + Lookup::Statuses { .. } => "statuses", + Lookup::Workload { .. } => "workload", + Lookup::LastActivity { .. } => "last_activity", + Lookup::OfKind { .. } => "of_kind", + } + } + + /// What the shape exercises in the wrapper, for the stats page and + /// the scenario report. + pub fn exercises(&self) -> &'static str { + match self { + Lookup::Ticket { .. } => "primary key, enum columns, timestamps", + Lookup::Queue { .. } => "enum array bind, ORDER BY timestamp DESC, LIMIT, OFFSET", + Lookup::QueueCount { .. } => "COUNT(*) with an enum array bind", + Lookup::Comments { .. } => "inner join composed from engine rows", + Lookup::Attachment { .. } => "bytea", + Lookup::Agents { .. } => "text[] column", + Lookup::Sla { .. } => "interval column: refill mode, not rows mode", + Lookup::Search { .. } => "ILIKE: refused by the engine, served live", + Lookup::Recent { .. } => "ORDER BY updated_at DESC LIMIT", + Lookup::OrgTicket { .. } => "two-column key", + Lookup::Unassigned { .. } => "IS NULL", + Lookup::Escalated { .. } => "OR of two enum binds", + Lookup::OpenExcept { .. } => "<> and NOT IN", + Lookup::CreatedBetween { .. } => "timestamp range: two binds", + Lookup::Many { .. } => "IN list of a hundred ids", + Lookup::QueueEnd { .. } => "page past the end: empty, subscribed, stays empty", + Lookup::Nothing { .. } => "never any rows: empty, subscribed, stays empty", + Lookup::AllTickets { .. } => "largest result: not kept under a small per-entry cap", + Lookup::CommentAuthors { .. } => "left join with an empty side", + Lookup::Skilled { .. } => "array @>: refused, served live", + Lookup::Statuses { .. } => "DISTINCT: refused, served live", + Lookup::Workload { .. } => "GROUP BY: refused, served live", + Lookup::LastActivity { .. } => "MAX: refused, served live", + Lookup::OfKind { .. } => "enum type added after the engine started", + } + } + + /// What the shape expects of the engine. + pub fn expected(&self) -> Expected { + match self { + Lookup::Sla { .. } => Expected::Refill, + Lookup::Search { .. } + | Lookup::Skilled { .. } + | Lookup::Statuses { .. } + | Lookup::Workload { .. } + | Lookup::LastActivity { .. } => Expected::Live, + _ => Expected::Rows, + } + } + + /// Whether the engine can follow the query at all. + pub fn is_cacheable(&self) -> bool { + self.expected() != Expected::Live + } + + /// Whether a change made for the shape leaves its result as it is: + /// the empty results, which must stay empty. + pub fn is_static(&self) -> bool { + matches!(self, Lookup::QueueEnd { .. } | Lookup::Nothing { .. }) + } + + /// The demo's lookups with their default arguments, in a fixed + /// order: what the scenario test, the verify loop and the bench + /// iterate. The shapes an operator refuses come last, so the table + /// reads from the common case down. + pub fn demo() -> Vec { + vec![ + Lookup::Ticket { id: 1 }, + Lookup::Queue { + org: 1, + statuses: vec![TicketStatus::Open, TicketStatus::Pending], + page: 1, + }, + Lookup::QueueCount { + org: 1, + statuses: vec![TicketStatus::Open], + }, + Lookup::Comments { ticket: 1 }, + Lookup::Attachment { id: 1 }, + Lookup::Agents { org: 1 }, + Lookup::Sla { org: 1 }, + Lookup::Search { + org: 1, + pattern: "%login%".to_string(), + }, + Lookup::Recent { org: 1, limit: 5 }, + Lookup::OrgTicket { org: 1, id: 5 }, + Lookup::Unassigned { org: 1 }, + Lookup::Escalated { org: 1 }, + Lookup::OpenExcept { + org: 1, + except: vec![1, 2, 3], + }, + Lookup::CreatedBetween { + org: 1, + from: Utc.with_ymd_and_hms(2026, 10, 1, 9, 0, 0).unwrap(), + to: Utc.with_ymd_and_hms(2026, 10, 1, 12, 0, 0).unwrap(), + }, + Lookup::Many { + ids: (1..=100).collect(), + }, + Lookup::QueueEnd { org: 1 }, + Lookup::Nothing { org: 1 }, + Lookup::AllTickets { org: 1 }, + Lookup::CommentAuthors { ticket: 1 }, + Lookup::Skilled { + org: 1, + skill: "billing".to_string(), + }, + Lookup::Statuses { org: 1 }, + Lookup::Workload { org: 1 }, + Lookup::LastActivity { org: 1 }, + ] + } + + /// A lookup by name: the demo list, plus the one that needs the + /// migration. + pub fn named(name: &str) -> Option { + if name == "of_kind" { + return Some(Lookup::OfKind { + org: 1, + kind: TicketKind::Bug, + }); + } + Self::demo() + .into_iter() + .find(|lookup| lookup.name() == name) + } + + /// Runs the lookup on the blocking wrapper. + pub fn run_sync(&self, conn: &mut SyncConn) -> QueryResult { + match self { + Lookup::Ticket { id } => json(sync_load(ticket(*id), conn)?), + Lookup::Queue { + org, + statuses, + page, + } => json(sync_load(queue(*org, statuses.clone(), *page), conn)?), + Lookup::QueueCount { org, statuses } => json(sync_load::<_, i64>( + queue_count(*org, statuses.clone()), + conn, + )?), + Lookup::Comments { ticket } => json(comments_of(sync_load::<_, CommentRow>( + ticket_comments(*ticket), + conn, + )?)), + Lookup::Attachment { id } => json(sync_load(attachment(*id), conn)?), + Lookup::Agents { org } => json(sync_load(org_agents(*org), conn)?), + Lookup::Sla { org } => json( + sync_load(sla(*org), conn)? + .into_iter() + .map(TicketSla::from_row) + .collect::>(), + ), + Lookup::Search { org, pattern } => { + json(sync_load(search(*org, pattern.clone()), conn)?) + } + Lookup::Recent { org, limit } => json(sync_load(recent(*org, *limit), conn)?), + Lookup::OrgTicket { org, id } => json(sync_load(org_ticket(*org, *id), conn)?), + Lookup::Unassigned { org } => json(sync_load(unassigned(*org), conn)?), + Lookup::Escalated { org } => json(sync_load(escalated(*org), conn)?), + Lookup::OpenExcept { org, except } => { + json(sync_load(open_except(*org, except.clone()), conn)?) + } + Lookup::CreatedBetween { org, from, to } => { + json(sync_load(created_between(*org, *from, *to), conn)?) + } + Lookup::Many { ids } => json(sync_load(many(ids.clone()), conn)?), + Lookup::QueueEnd { org } => json(sync_load(queue_end(*org), conn)?), + Lookup::Nothing { org } => json(sync_load(nothing(*org), conn)?), + Lookup::AllTickets { org } => json(sync_load(all_tickets(*org), conn)?), + Lookup::CommentAuthors { ticket } => json(authors_of(sync_load::<_, AuthorRow>( + comment_authors(*ticket), + conn, + )?)), + Lookup::Skilled { org, skill } => json(sync_load(skilled(*org, skill.clone()), conn)?), + Lookup::Statuses { org } => json(sync_load::<_, TicketStatus>(statuses(*org), conn)?), + Lookup::Workload { org } => json(sync_load::<_, Workload>(workload(*org), conn)?), + Lookup::LastActivity { org } => json(sync_load::<_, Option>>( + last_activity(*org), + conn, + )?), + Lookup::OfKind { org, kind } => json(sync_load(of_kind(*org, *kind), conn)?), + } + } + + /// Runs the lookup on the async wrapper. + pub async fn run_async(&self, conn: &mut AsyncConn) -> QueryResult { + match self { + Lookup::Ticket { id } => json(async_load(ticket(*id), conn).await?), + Lookup::Queue { + org, + statuses, + page, + } => json(async_load(queue(*org, statuses.clone(), *page), conn).await?), + Lookup::QueueCount { org, statuses } => { + json(async_load::<_, i64>(queue_count(*org, statuses.clone()), conn).await?) + } + Lookup::Comments { ticket } => json(comments_of( + async_load::<_, CommentRow>(ticket_comments(*ticket), conn).await?, + )), + Lookup::Attachment { id } => json(async_load(attachment(*id), conn).await?), + Lookup::Agents { org } => json(async_load(org_agents(*org), conn).await?), + Lookup::Sla { org } => json( + async_load(sla(*org), conn) + .await? + .into_iter() + .map(TicketSla::from_row) + .collect::>(), + ), + Lookup::Search { org, pattern } => { + json(async_load(search(*org, pattern.clone()), conn).await?) + } + Lookup::Recent { org, limit } => json(async_load(recent(*org, *limit), conn).await?), + Lookup::OrgTicket { org, id } => json(async_load(org_ticket(*org, *id), conn).await?), + Lookup::Unassigned { org } => json(async_load(unassigned(*org), conn).await?), + Lookup::Escalated { org } => json(async_load(escalated(*org), conn).await?), + Lookup::OpenExcept { org, except } => { + json(async_load(open_except(*org, except.clone()), conn).await?) + } + Lookup::CreatedBetween { org, from, to } => { + json(async_load(created_between(*org, *from, *to), conn).await?) + } + Lookup::Many { ids } => json(async_load(many(ids.clone()), conn).await?), + Lookup::QueueEnd { org } => json(async_load(queue_end(*org), conn).await?), + Lookup::Nothing { org } => json(async_load(nothing(*org), conn).await?), + Lookup::AllTickets { org } => json(async_load(all_tickets(*org), conn).await?), + Lookup::CommentAuthors { ticket } => json(authors_of( + async_load::<_, AuthorRow>(comment_authors(*ticket), conn).await?, + )), + Lookup::Skilled { org, skill } => { + json(async_load(skilled(*org, skill.clone()), conn).await?) + } + Lookup::Statuses { org } => { + json(async_load::<_, TicketStatus>(statuses(*org), conn).await?) + } + Lookup::Workload { org } => { + json(async_load::<_, Workload>(workload(*org), conn).await?) + } + Lookup::LastActivity { org } => { + json(async_load::<_, Option>>(last_activity(*org), conn).await?) + } + Lookup::OfKind { org, kind } => json(async_load(of_kind(*org, *kind), conn).await?), + } + } +} + +/// The rows as JSON, the form the API serves and the tests compare. +fn json(rows: T) -> QueryResult { + serde_json::to_value(rows) + .map_err(|error| diesel::result::Error::SerializationError(Box::new(error))) +} + +/// Loads a query on the blocking wrapper. +fn sync_load(query: Q, conn: &mut SyncConn) -> QueryResult> +where + Q: diesel::query_dsl::methods::LoadQuery<'static, SyncConn, U>, +{ + diesel::RunQueryDsl::load(query, conn) +} + +/// Boxed at a plain function rather than awaited in an `async fn`: the +/// compiler cannot prove the load future `Send` through the latter's +/// desugaring (the lifetime in diesel-async's `LoadQuery` bound comes +/// out higher-ranked), and can here. +fn async_load<'a, Q, U>(query: Q, conn: &'a mut AsyncConn) -> BoxFuture<'a, QueryResult>> +where + Q: diesel_async::methods::LoadQuery<'static, AsyncConn, U> + Send + 'static, + U: Send + 'static, +{ + diesel_async::RunQueryDsl::load(query, conn).boxed() +} + +/// Tickets, as the API serves them. +type TicketQuery = tickets::BoxedQuery<'static, Pg, AsSelect>; +/// A count of tickets. +type CountQuery = tickets::BoxedQuery<'static, Pg, BigInt>; +/// Tickets with their `interval` SLA. +type SlaQuery = tickets::BoxedQuery<'static, Pg, (Int8, Text, Nullable)>; +/// Agents. +type AgentsQuery = agents::BoxedQuery<'static, Pg, AsSelect>; +/// An attachment, `bytea` and all. +type AttachmentQuery = attachments::BoxedQuery<'static, Pg, AsSelect>; +/// Ticket statuses. +type StatusesQuery = tickets::BoxedQuery<'static, Pg, TicketStatusType>; +/// The latest update time. +type LastActivityQuery = tickets::BoxedQuery<'static, Pg, Nullable>; +/// Tickets with the kind a migration adds. +type OfKindQuery = tickets_with_kind::BoxedQuery<'static, Pg, AsSelect>; +/// Tickets joined to their comments. +type CommentsQuery = BoxedSelectStatement< + 'static, + (Int8, Text, Int8, Text, Timestamptz), + FromClause>, + Pg, +>; +/// Comments left-joined to their authors. +type AuthorsQuery = BoxedSelectStatement< + 'static, + (Int8, Text, Nullable), + FromClause>, + Pg, +>; +/// Raw SQL with its binds. +type WorkloadQuery = + diesel::query_builder::BoxedSqlQuery<'static, Pg, diesel::query_builder::SqlQuery>; + +/// One ticket by id. +fn ticket(id: i64) -> TicketQuery { + tickets::table + .filter(tickets::id.eq(id)) + .select(Ticket::as_select()) + .into_boxed() +} + +/// A page of an organization's queue, newest first: an enum bound as +/// an array, and an offset. +fn queue(org: i64, statuses: Vec, page: i64) -> TicketQuery { + tickets::table + .filter(tickets::org_id.eq(org)) + .filter(tickets::status.eq_any(statuses)) + .order((tickets::created_at.desc(), tickets::id.asc())) + .limit(PAGE) + .offset(page * PAGE) + .select(Ticket::as_select()) + .into_boxed() +} + +/// How many tickets the queue holds: a `COUNT(*)`. +fn queue_count(org: i64, statuses: Vec) -> CountQuery { + tickets::table + .filter(tickets::org_id.eq(org)) + .filter(tickets::status.eq_any(statuses)) + .count() + .into_boxed() +} + +/// A ticket's comments: an inner join, ordered by both keys. +fn ticket_comments(ticket: i64) -> CommentsQuery { + tickets::table + .inner_join(comments::table) + .filter(tickets::id.eq(ticket)) + .order((tickets::id.asc(), comments::id.asc())) + .select(( + tickets::id, + tickets::title, + comments::id, + comments::body, + comments::created_at, + )) + .into_boxed() +} + +/// One attachment: a `bytea`. +fn attachment(id: i64) -> AttachmentQuery { + attachments::table + .filter(attachments::id.eq(id)) + .select(Attachment::as_select()) + .into_boxed() +} + +/// An organization's agents: a `text[]` column. +fn org_agents(org: i64) -> AgentsQuery { + agents::table + .filter(agents::org_id.eq(org)) + .order(agents::id.asc()) + .select(Agent::as_select()) + .into_boxed() +} + +/// Tickets with an SLA: an `interval` the wrapper cannot encode. +fn sla(org: i64) -> SlaQuery { + tickets::table + .filter(tickets::org_id.eq(org)) + .filter(tickets::sla.is_not_null()) + .order(tickets::id.asc()) + .select((tickets::id, tickets::title, tickets::sla)) + .into_boxed() +} + +/// Tickets whose title matches: `ILIKE`, which the engine refuses. +fn search(org: i64, pattern: String) -> TicketQuery { + tickets::table + .filter(tickets::org_id.eq(org)) + .filter(tickets::title.ilike(pattern)) + .order(tickets::id.asc()) + .select(Ticket::as_select()) + .into_boxed() +} + +/// The most recently updated tickets: a limit. +fn recent(org: i64, limit: i64) -> TicketQuery { + tickets::table + .filter(tickets::org_id.eq(org)) + .order((tickets::updated_at.desc(), tickets::id.asc())) + .limit(limit) + .select(Ticket::as_select()) + .into_boxed() +} + +/// One ticket by organization and id: a two-column condition. +fn org_ticket(org: i64, id: i64) -> TicketQuery { + tickets::table + .filter(tickets::org_id.eq(org)) + .filter(tickets::id.eq(id)) + .select(Ticket::as_select()) + .into_boxed() +} + +/// Tickets with no agent: `IS NULL`. +fn unassigned(org: i64) -> TicketQuery { + tickets::table + .filter(tickets::org_id.eq(org)) + .filter(tickets::agent_id.is_null()) + .order(tickets::id.asc()) + .select(Ticket::as_select()) + .into_boxed() +} + +/// High-priority or pending tickets: an `OR`. +fn escalated(org: i64) -> TicketQuery { + tickets::table + .filter(tickets::org_id.eq(org)) + .filter( + tickets::priority + .eq(TicketPriority::High) + .or(tickets::status.eq(TicketStatus::Pending)), + ) + .order(tickets::id.asc()) + .select(Ticket::as_select()) + .into_boxed() +} + +/// Open tickets but the named ones: `!= ALL`. +fn open_except(org: i64, except: Vec) -> TicketQuery { + tickets::table + .filter(tickets::org_id.eq(org)) + .filter(tickets::status.ne(TicketStatus::Closed)) + .filter(tickets::id.ne_all(except)) + .order(tickets::id.asc()) + .select(Ticket::as_select()) + .into_boxed() +} + +/// Tickets created in a range: a timestamp range. +fn created_between(org: i64, from: DateTime, to: DateTime) -> TicketQuery { + tickets::table + .filter(tickets::org_id.eq(org)) + .filter(tickets::created_at.ge(from)) + .filter(tickets::created_at.lt(to)) + .order(tickets::id.asc()) + .select(Ticket::as_select()) + .into_boxed() +} + +/// Tickets by id: an `IN` list of up to a hundred. +fn many(ids: Vec) -> TicketQuery { + tickets::table + .filter(tickets::id.eq_any(ids)) + .order(tickets::id.asc()) + .select(Ticket::as_select()) + .into_boxed() +} + +/// A page past the end of the queue: an empty result. +fn queue_end(org: i64) -> TicketQuery { + queue(org, TicketStatus::ALL.to_vec(), 1000) +} + +/// A condition no ticket meets: an empty result. +fn nothing(org: i64) -> TicketQuery { + tickets::table + .filter(tickets::org_id.eq(org)) + .filter(tickets::id.lt(0)) + .order(tickets::id.asc()) + .select(Ticket::as_select()) + .into_boxed() +} + +/// Every ticket of an organization. +fn all_tickets(org: i64) -> TicketQuery { + tickets::table + .filter(tickets::org_id.eq(org)) + .order(tickets::id.asc()) + .select(Ticket::as_select()) + .into_boxed() +} + +/// A ticket's comments and their authors: a left join. +fn comment_authors(ticket: i64) -> AuthorsQuery { + comments::table + .left_join(agents::table) + .filter(comments::ticket_id.eq(ticket)) + .order(comments::id.asc()) + .select((comments::id, comments::body, agents::name.nullable())) + .into_boxed() +} + +/// Agents with a skill: the array operator `@>`, which the engine +/// refuses. +fn skilled(org: i64, skill: String) -> AgentsQuery { + agents::table + .filter(agents::org_id.eq(org)) + .filter(agents::skills.contains(vec![skill])) + .order(agents::id.asc()) + .select(Agent::as_select()) + .into_boxed() +} + +/// The statuses in use: `DISTINCT`, which the engine refuses. +fn statuses(org: i64) -> StatusesQuery { + tickets::table + .filter(tickets::org_id.eq(org)) + .select(tickets::status) + .distinct() + .order(tickets::status.asc()) + .into_boxed() +} + +/// Tickets per agent: `GROUP BY`, which the engine refuses. +fn workload(org: i64) -> WorkloadQuery { + diesel::sql_query( + "SELECT agent_id, COUNT(*) AS tickets FROM tickets WHERE org_id = $1 \ + GROUP BY agent_id ORDER BY agent_id", + ) + .into_boxed() + .bind::(org) +} + +/// The latest update: `MAX`, which the engine refuses. +fn last_activity(org: i64) -> LastActivityQuery { + tickets::table + .filter(tickets::org_id.eq(org)) + .select(diesel::dsl::max(tickets::updated_at)) + .into_boxed() +} + +/// Tickets of a kind: an enum created after the engine started. +fn of_kind(org: i64, kind: TicketKind) -> OfKindQuery { + tickets_with_kind::table + .filter(tickets_with_kind::org_id.eq(org)) + .filter(tickets_with_kind::kind.eq(kind)) + .order(tickets_with_kind::id.asc()) + .select(TicketOfKind::as_select()) + .into_boxed() +} + +/// A row of the comments join, as the tuple Diesel decodes. +pub type CommentRow = (i64, String, i64, String, DateTime); + +/// The join's rows as the struct the API serves. +pub fn comments_of(rows: Vec) -> Vec { + rows.into_iter() + .map( + |(ticket_id, title, comment_id, body, created_at)| TicketComment { + ticket_id, + title, + comment_id, + body, + created_at, + }, + ) + .collect() +} + +/// A row of the left join, as the tuple Diesel decodes. +pub type AuthorRow = (i64, String, Option); + +/// The left join's rows as the struct the API serves. +pub fn authors_of(rows: Vec) -> Vec { + rows.into_iter() + .map(|(comment_id, body, agent)| CommentAuthor { + comment_id, + body, + agent, + }) + .collect() +} diff --git a/crates/streamgres-demo/src/schema.rs b/crates/streamgres-demo/src/schema.rs new file mode 100644 index 0000000..7dbc9db --- /dev/null +++ b/crates/streamgres-demo/src/schema.rs @@ -0,0 +1,86 @@ +//! The ticket desk's tables, as Diesel sees them. +//! +//! Five tables, chosen for the shapes they let a query take: an enum +//! column to bind and to `eq_any`, a text array, a `bytea`, an +//! `interval` (which the engine cannot ship rows for, so the query +//! falls back to refills), timestamps to order by, and a parent-child +//! pair to join. `tickets_with_kind` is the tickets table again, seen +//! through the column `demo migrate new-type` adds. + +diesel::table! { + organizations (id) { + id -> Int8, + slug -> Text, + name -> Text, + } +} + +diesel::table! { + agents (id) { + id -> Int8, + org_id -> Int8, + name -> Text, + skills -> Array, + } +} + +diesel::table! { + use diesel::sql_types::*; + use crate::models::{TicketPriorityType, TicketStatusType}; + + tickets (id) { + id -> Int8, + org_id -> Int8, + agent_id -> Nullable, + title -> Text, + status -> TicketStatusType, + priority -> TicketPriorityType, + sla -> Nullable, + created_at -> Timestamptz, + updated_at -> Timestamptz, + } +} + +diesel::table! { + comments (id) { + id -> Int8, + ticket_id -> Int8, + agent_id -> Nullable, + body -> Text, + created_at -> Timestamptz, + } +} + +diesel::table! { + attachments (id) { + id -> Int8, + ticket_id -> Int8, + name -> Text, + data -> Bytea, + } +} + +diesel::table! { + use diesel::sql_types::*; + use crate::models::TicketKindType; + + #[sql_name = "tickets"] + tickets_with_kind (id) { + id -> Int8, + org_id -> Int8, + kind -> Nullable, + } +} + +diesel::joinable!(agents -> organizations (org_id)); +diesel::joinable!(comments -> agents (agent_id)); +diesel::joinable!(tickets -> organizations (org_id)); +diesel::joinable!(comments -> tickets (ticket_id)); +diesel::joinable!(attachments -> tickets (ticket_id)); +diesel::allow_tables_to_appear_in_same_query!( + organizations, + agents, + tickets, + comments, + attachments +); diff --git a/crates/streamgres-demo/src/seed.rs b/crates/streamgres-demo/src/seed.rs new file mode 100644 index 0000000..8c26d4e --- /dev/null +++ b/crates/streamgres-demo/src/seed.rs @@ -0,0 +1,163 @@ +//! The schema and a small, readable, deterministic dataset, at a scale. +//! +//! Three organizations, four agents each, sixty tickets each per unit of +//! scale with statuses and priorities spread evenly, comments on the +//! first third of each organization's tickets, an attachment on every +//! tenth, a service level on every third. Ids are fixed, so the demo's +//! URLs and the scenario test's expectations do not move between runs; +//! the first organization's tickets are `1..=60 * scale` at any scale. +//! Reseeding deletes and inserts rather than truncating, so an engine +//! already following the tables sees ordinary changes. +//! +//! `DEMO_SCALE` (default 1) multiplies the tickets, comments and +//! attachments, so pages run deep and the cache's byte limits matter. + +use diesel::QueryResult; +use diesel::connection::SimpleConnection; + +/// Tickets per organization at scale 1. +pub const TICKETS_PER_ORG: i64 = 60; + +/// Creates the types, tables and sequences if they do not exist. The +/// writer's rows take ids from the sequences, above the seed's. +pub const SCHEMA: &str = " +DO $$ BEGIN + CREATE TYPE ticket_status AS ENUM ('open', 'pending', 'closed'); +EXCEPTION WHEN duplicate_object OR unique_violation THEN NULL; END $$; +DO $$ BEGIN + CREATE TYPE ticket_priority AS ENUM ('low', 'normal', 'high'); +EXCEPTION WHEN duplicate_object OR unique_violation THEN NULL; END $$; +CREATE TABLE IF NOT EXISTS organizations ( + id bigint PRIMARY KEY, + slug text NOT NULL UNIQUE, + name text NOT NULL +); +CREATE TABLE IF NOT EXISTS agents ( + id bigint PRIMARY KEY, + org_id bigint NOT NULL REFERENCES organizations (id), + name text NOT NULL, + skills text[] NOT NULL DEFAULT '{}' +); +CREATE TABLE IF NOT EXISTS tickets ( + id bigint PRIMARY KEY, + org_id bigint NOT NULL REFERENCES organizations (id), + agent_id bigint REFERENCES agents (id), + title text NOT NULL, + status ticket_status NOT NULL DEFAULT 'open', + priority ticket_priority NOT NULL DEFAULT 'normal', + sla interval, + created_at timestamptz NOT NULL DEFAULT now(), + updated_at timestamptz NOT NULL DEFAULT now() +); +CREATE TABLE IF NOT EXISTS comments ( + id bigint PRIMARY KEY, + ticket_id bigint NOT NULL REFERENCES tickets (id) ON DELETE CASCADE, + agent_id bigint REFERENCES agents (id), + body text NOT NULL, + created_at timestamptz NOT NULL DEFAULT now() +); +CREATE TABLE IF NOT EXISTS attachments ( + id bigint PRIMARY KEY, + ticket_id bigint NOT NULL REFERENCES tickets (id) ON DELETE CASCADE, + name text NOT NULL, + data bytea NOT NULL +); +CREATE SEQUENCE IF NOT EXISTS comments_id_seq START 100000; +CREATE SEQUENCE IF NOT EXISTS tickets_id_seq START 1000000; +"; + +/// Replaces the data; `{tickets}`, `{commented}` and the fixed instants +/// make it the same on every run at a scale. +const DATA: &str = " +BEGIN; +DELETE FROM attachments; +DELETE FROM comments; +DELETE FROM tickets; +DELETE FROM agents; +DELETE FROM organizations; + +INSERT INTO organizations (id, slug, name) VALUES + (1, 'acme', 'Acme Corporation'), + (2, 'globex', 'Globex'), + (3, 'initech', 'Initech'); + +INSERT INTO agents (id, org_id, name, skills) +SELECT (org - 1) * 4 + n, + org, + (ARRAY['Ada', 'Grace', 'Linus', 'Margaret'])[n], + CASE n + WHEN 1 THEN ARRAY['billing', 'refunds'] + WHEN 2 THEN ARRAY['login', 'sso'] + WHEN 3 THEN ARRAY['api'] + ELSE ARRAY['billing', 'api', 'login'] + END +FROM generate_series(1, 3) AS org, generate_series(1, 4) AS n; + +INSERT INTO tickets (id, org_id, agent_id, title, status, priority, sla, created_at, updated_at) +SELECT (org - 1) * {tickets} + n, + org, + CASE WHEN n % 5 = 0 THEN NULL ELSE (org - 1) * 4 + 1 + (n % 4) END, + (ARRAY['Cannot login after password reset', 'Invoice total is wrong', 'API returns 500 on export', + 'SSO redirect loops', 'Refund not received', 'Login page blank on Safari'])[1 + n % 6] || ' #' || n, + (ARRAY['open', 'pending', 'closed'])[1 + n % 3]::ticket_status, + (ARRAY['low', 'normal', 'high'])[1 + (n / 3) % 3]::ticket_priority, + CASE WHEN n % 3 = 0 THEN (n % 2 + 1) * interval '4 hours' ELSE NULL END, + timestamptz '2026-10-01 09:00:00+00' + (n * interval '37 minutes'), + timestamptz '2026-10-01 09:00:00+00' + (n * interval '37 minutes') + interval '1 hour' +FROM generate_series(1, 3) AS org, generate_series(1, {tickets}) AS n; + +INSERT INTO comments (id, ticket_id, agent_id, body, created_at) +SELECT ((org - 1) * {tickets} + n) * 10 + c, + (org - 1) * {tickets} + n, + CASE WHEN c = 1 THEN (org - 1) * 4 + 1 + (n % 4) ELSE NULL END, + CASE WHEN c = 1 THEN 'Looking into it.' ELSE 'Any update?' END, + timestamptz '2026-10-01 10:00:00+00' + (n * interval '37 minutes') + (c * interval '15 minutes') +FROM generate_series(1, 3) AS org, generate_series(1, {commented}) AS n, generate_series(1, 2) AS c; + +INSERT INTO attachments (id, ticket_id, name, data) +SELECT ((org - 1) * {tickets} + n) / 10, + (org - 1) * {tickets} + n, + 'screenshot-' || n || '.png', + decode('89504e470d0a1a0a' || lpad(to_hex(n), 8, '0'), 'hex') +FROM generate_series(1, 3) AS org, generate_series(10, {tickets}, 10) AS n; +COMMIT; +"; + +/// The scale from `DEMO_SCALE`: 1 unless set to a larger number. +pub fn scale() -> i64 { + std::env::var("DEMO_SCALE") + .ok() + .and_then(|value| value.parse::().ok()) + .unwrap_or(1) + .max(1) +} + +/// Tickets per organization at the scale. +pub fn tickets_per_org(scale: i64) -> i64 { + TICKETS_PER_ORG * scale +} + +/// The data statement at a scale. +pub fn data(scale: i64) -> String { + let tickets = tickets_per_org(scale); + DATA.replace("{tickets}", &tickets.to_string()) + .replace("{commented}", &(tickets / 3).to_string()) +} + +/// What the seed loads at a scale, for the report. +pub fn summary(scale: i64) -> String { + let tickets = tickets_per_org(scale); + format!( + "3 organizations, 12 agents, {} tickets, {} comments, {} attachments (scale {scale})", + 3 * tickets, + 3 * (tickets / 3) * 2, + 3 * (tickets / 10) + ) +} + +/// Creates the schema and loads the data at `DEMO_SCALE` on a +/// blocking connection. +pub fn run(conn: &mut C) -> QueryResult<()> { + conn.batch_execute(SCHEMA)?; + conn.batch_execute(&data(scale())) +} diff --git a/crates/streamgres-demo/src/verify.rs b/crates/streamgres-demo/src/verify.rs new file mode 100644 index 0000000..fa5270c --- /dev/null +++ b/crates/streamgres-demo/src/verify.rs @@ -0,0 +1,138 @@ +//! The demo checking itself: every lookup read through the cache and +//! through Postgres, compared. +//! +//! A cached answer may lag a commit by the engine's propagation delay, +//! so a difference is not a mismatch until it has had time to settle: +//! the pair is re-read until the two agree or `settle` runs out. A pair +//! that agreed on a later read counts as `lagged`; one that never did is +//! a mismatch, kept with both answers. The verify loop (`demo +//! verify`) and `/verify` run this; `/stats` shows the counters. + +use std::sync::Mutex; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::{Duration, Instant}; + +use serde::Serialize; +use serde_json::Value; + +use crate::pools::Stack; +use crate::queries::Lookup; + +/// A pair that never agreed. +#[derive(Debug, Clone, Serialize)] +pub struct Mismatch { + pub lookup: String, + pub cached: Value, + pub postgres: Value, +} + +/// What one pass found. +#[derive(Debug, Default, Serialize)] +pub struct Outcome { + pub compared: u64, + pub lagged: u64, + pub mismatches: Vec, +} + +/// The counters over every pass so far, and the last mismatch. +#[derive(Default)] +pub struct Counts { + pub runs: AtomicU64, + pub compared: AtomicU64, + pub lagged: AtomicU64, + pub mismatches: AtomicU64, + pub errors: AtomicU64, + pub last: Mutex>, +} + +impl Counts { + /// The counters as JSON, with the last mismatch's lookup name. + pub fn json(&self) -> Value { + let last = self + .last + .lock() + .ok() + .and_then(|last| last.as_ref().map(|mismatch| mismatch.lookup.clone())); + serde_json::json!({ + "runs": AtomicU64::load(&self.runs, Ordering::Relaxed), + "compared": AtomicU64::load(&self.compared, Ordering::Relaxed), + "lagged": AtomicU64::load(&self.lagged, Ordering::Relaxed), + "mismatches": AtomicU64::load(&self.mismatches, Ordering::Relaxed), + "errors": AtomicU64::load(&self.errors, Ordering::Relaxed), + "last_mismatch": last, + }) + } + + /// Folds a pass into the counters. + fn record(&self, outcome: &Outcome) { + self.runs.fetch_add(1, Ordering::Relaxed); + self.compared.fetch_add(outcome.compared, Ordering::Relaxed); + self.lagged.fetch_add(outcome.lagged, Ordering::Relaxed); + self.mismatches + .fetch_add(outcome.mismatches.len() as u64, Ordering::Relaxed); + if let (Some(last), Ok(mut slot)) = (outcome.mismatches.last(), self.last.lock()) { + *slot = Some(last.clone()); + } + } +} + +/// Reads `lookup` through the cache and through Postgres until the two +/// agree or `settle` runs out: the answers, and whether a later read +/// was needed. +async fn compare( + stack: &Stack, + lookup: &Lookup, + settle: Duration, +) -> Result<(Value, Value, bool), String> { + let mut cached_conn = stack.main.get().await.map_err(|error| error.to_string())?; + let mut truth_conn = stack + .reporting + .get() + .await + .map_err(|error| error.to_string())?; + let started = Instant::now(); + let mut lagged = false; + loop { + let cached = lookup + .run_async(&mut cached_conn) + .await + .map_err(|error| format!("{}: cached read: {error}", lookup.name()))?; + let truth = lookup + .run_async(&mut truth_conn) + .await + .map_err(|error| format!("{}: postgres read: {error}", lookup.name()))?; + if cached == truth || started.elapsed() >= settle { + return Ok((cached, truth, lagged)); + } + lagged = true; + tokio::time::sleep(Duration::from_millis(25)).await; + } +} + +/// One pass over `lookups`, recorded in the stack's counters. +pub async fn once(stack: &Stack, lookups: &[Lookup], settle: Duration) -> Outcome { + let mut outcome = Outcome::default(); + for lookup in lookups { + match compare(stack, lookup, settle).await { + Ok((cached, postgres, lagged)) => { + outcome.compared += 1; + if lagged { + outcome.lagged += 1; + } + if cached != postgres { + outcome.mismatches.push(Mismatch { + lookup: lookup.name().to_string(), + cached, + postgres, + }); + } + } + Err(error) => { + stack.verify.errors.fetch_add(1, Ordering::Relaxed); + eprintln!("verify: {error}"); + } + } + } + stack.verify.record(&outcome); + outcome +} diff --git a/crates/streamgres-demo/src/writer.rs b/crates/streamgres-demo/src/writer.rs new file mode 100644 index 0000000..e7aba60 --- /dev/null +++ b/crates/streamgres-demo/src/writer.rs @@ -0,0 +1,308 @@ +//! A stream of changes, so the demo has something to keep fresh and the +//! verify loop something to catch. +//! +//! Every tick runs one operation drawn from the mix: a status moved on +//! (`status`), a comment added (`comment`), a ticket inserted into the +//! first organization's queue (`insert`), one of the writer's own +//! tickets deleted with its comments cascading (`delete`), one moved to +//! another organization so it leaves one subscription and enters +//! another (`move`), an agent renamed under the joins (`rename`), a +//! ticket's assignee flipped between an agent and `NULL` (`assign`), +//! several tickets touched in one transaction (`bulk`), a change rolled +//! back so nothing must reach the cache (`rollback`), and a comment +//! with a body too large for its row written then updated without +//! touching the body, so the engine sees a partial image (`toast`). +//! +//! Rows the writer creates take ids from the sequences, above the +//! seed's, and bodies marked `writer:`, so the seeded rows the demo's +//! URLs point at stay as they are apart from their status, assignee and +//! timestamps. The writer uses the plain pool, not the wrapper, so its +//! changes reach the cache the way another service's would: through the +//! engine, not the read-your-writes window. +//! +//! `DEMO_WRITER_MIX` names the operations to draw from, comma +//! separated; the default is all of them. + +use std::collections::VecDeque; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::Duration; + +use diesel::sql_query; +use diesel::sql_types::{BigInt, Text}; +use diesel_async::pooled_connection::bb8::Pool; +use diesel_async::{AsyncPgConnection, RunQueryDsl, SimpleAsyncConnection}; +use rand::Rng; +use serde_json::{Value, json}; +use tokio::task::JoinHandle; + +use crate::seed; + +/// The operations, in the order the counters report them. +pub const OPS: &[&str] = &[ + "status", "comment", "insert", "delete", "move", "rename", "assign", "bulk", "rollback", + "toast", +]; + +/// How many of the writer's own tickets it keeps around before an +/// insert becomes a delete. +const KEEP: usize = 50; + +/// How many times each operation ran, and how many failed. +#[derive(Default)] +pub struct Counts { + done: [AtomicU64; 10], + errors: AtomicU64, +} + +impl Counts { + /// The counters as JSON, by operation name. + pub fn json(&self) -> Value { + let mut map = serde_json::Map::new(); + for (index, name) in OPS.iter().enumerate() { + map.insert( + (*name).to_string(), + json!(AtomicU64::load(&self.done[index], Ordering::Relaxed)), + ); + } + map.insert( + "errors".to_string(), + json!(AtomicU64::load(&self.errors, Ordering::Relaxed)), + ); + Value::Object(map) + } + + /// The total of every operation that ran. + pub fn total(&self) -> u64 { + self.done + .iter() + .map(|count| AtomicU64::load(count, Ordering::Relaxed)) + .sum() + } + + /// How many operations failed. + pub fn errors(&self) -> u64 { + AtomicU64::load(&self.errors, Ordering::Relaxed) + } + + /// How many of `op` ran. + pub fn of(&self, op: &str) -> u64 { + OPS.iter().position(|name| *name == op).map_or(0, |index| { + AtomicU64::load(&self.done[index], Ordering::Relaxed) + }) + } +} + +/// The mix from `DEMO_WRITER_MIX`, every operation when unset; a +/// name that is not an operation is reported and skipped. +pub fn mix_from_env() -> Vec<&'static str> { + let Ok(list) = std::env::var("DEMO_WRITER_MIX") else { + return OPS.to_vec(); + }; + let mut mix = Vec::new(); + for name in list + .split(',') + .map(str::trim) + .filter(|name| !name.is_empty()) + { + match OPS.iter().find(|op| **op == name) { + Some(op) => mix.push(*op), + None => eprintln!("writer: no operation `{name}`; the operations are {OPS:?}"), + } + } + if mix.is_empty() { OPS.to_vec() } else { mix } +} + +/// A returned id. +#[derive(diesel::QueryableByName)] +struct Id { + #[diesel(sql_type = BigInt)] + id: i64, +} + +/// What the writer remembers between ticks: its own tickets, and the +/// comment it is using for the partial-image case. +struct State { + tick: u64, + tickets_per_org: i64, + created: VecDeque, + toast: Option, +} + +impl State { + /// A seeded ticket of the first organization. + fn seeded(&self) -> i64 { + rand::rng().random_range(1..=self.tickets_per_org) + } + + /// One of the writer's own tickets, if it has any. + fn own(&self) -> Option { + if self.created.is_empty() { + return None; + } + let index = rand::rng().random_range(0..self.created.len()); + self.created.get(index).copied() + } +} + +/// Starts the writer: one operation from `mix` every `every`, counted +/// in `counts`. +pub fn spawn( + plain: Pool, + every: Duration, + mix: Vec<&'static str>, + counts: std::sync::Arc, +) -> JoinHandle<()> { + tokio::spawn(async move { + let mut state = State { + tick: 0, + tickets_per_org: seed::tickets_per_org(seed::scale()), + created: VecDeque::new(), + toast: None, + }; + loop { + tokio::time::sleep(every).await; + state.tick += 1; + let op = mix[rand::rng().random_range(0..mix.len())]; + let Ok(mut conn) = plain.get().await else { + counts.errors.fetch_add(1, Ordering::Relaxed); + continue; + }; + match step(op, &mut conn, &mut state).await { + Ok(()) => { + let index = OPS.iter().position(|name| *name == op).unwrap_or(0); + counts.done[index].fetch_add(1, Ordering::Relaxed); + } + Err(error) => { + counts.errors.fetch_add(1, Ordering::Relaxed); + eprintln!("writer: {op}: {error}"); + } + } + } + }) +} + +/// Runs one operation. +async fn step( + op: &str, + conn: &mut AsyncPgConnection, + state: &mut State, +) -> Result<(), diesel::result::Error> { + match op { + "status" => { + sql_query( + "UPDATE tickets SET status = (CASE status + WHEN 'open' THEN 'pending' WHEN 'pending' THEN 'closed' ELSE 'open' + END)::ticket_status, updated_at = now() WHERE id = $1", + ) + .bind::(state.seeded()) + .execute(conn) + .await?; + } + "comment" => { + sql_query( + "INSERT INTO comments (id, ticket_id, body) + VALUES (nextval('comments_id_seq'), $1, 'writer: status changed')", + ) + .bind::(state.seeded()) + .execute(conn) + .await?; + } + "insert" => { + if state.created.len() >= KEEP { + return Box::pin(step("delete", conn, state)).await; + } + let row: Id = sql_query( + "INSERT INTO tickets (id, org_id, title, status, priority) + VALUES (nextval('tickets_id_seq'), 1, 'writer: ticket ' || $1, 'open', 'normal') + RETURNING id", + ) + .bind::(state.tick.to_string()) + .get_result(conn) + .await?; + state.created.push_back(row.id); + } + "delete" => { + let Some(id) = state.created.pop_front() else { + return Box::pin(step("insert", conn, state)).await; + }; + sql_query("DELETE FROM tickets WHERE id = $1") + .bind::(id) + .execute(conn) + .await?; + } + "move" => { + let Some(id) = state.own() else { + return Box::pin(step("insert", conn, state)).await; + }; + sql_query( + "UPDATE tickets SET org_id = CASE org_id WHEN 1 THEN 2 ELSE 1 END, updated_at = now() + WHERE id = $1", + ) + .bind::(id) + .execute(conn) + .await?; + } + "rename" => { + sql_query( + "UPDATE agents SET name = regexp_replace(name, ' \\(v\\d+\\)$', '') || ' (v' || $1 || ')' + WHERE id = $2", + ) + .bind::(state.tick.to_string()) + .bind::(1 + (state.tick % 4) as i64) + .execute(conn) + .await?; + } + "assign" => { + let id = state.seeded(); + sql_query( + "UPDATE tickets SET agent_id = CASE WHEN agent_id IS NULL THEN $2 ELSE NULL END, + updated_at = now() WHERE id = $1", + ) + .bind::(id) + .bind::(1 + id % 4) + .execute(conn) + .await?; + } + "bulk" => { + sql_query("UPDATE tickets SET updated_at = now() WHERE org_id = 1 AND id % 7 = $1") + .bind::((state.tick % 7) as i64) + .execute(conn) + .await?; + } + "rollback" => { + let id = state.seeded(); + conn.batch_execute(&format!( + "BEGIN; UPDATE tickets SET title = 'writer: rolled back' WHERE id = {id}; ROLLBACK" + )) + .await?; + } + "toast" => match state.toast { + None => { + let row: Id = sql_query( + "INSERT INTO comments (id, ticket_id, body) + VALUES (nextval('comments_id_seq'), $1, 'writer: toast ' || repeat('x', 4096)) + RETURNING id", + ) + .bind::(state.seeded()) + .get_result(conn) + .await?; + state.toast = Some(row.id); + } + Some(id) if state.tick.is_multiple_of(5) => { + sql_query("DELETE FROM comments WHERE id = $1") + .bind::(id) + .execute(conn) + .await?; + state.toast = None; + } + Some(id) => { + sql_query("UPDATE comments SET created_at = now() WHERE id = $1") + .bind::(id) + .execute(conn) + .await?; + } + }, + other => eprintln!("writer: no operation `{other}`"), + } + Ok(()) +} diff --git a/crates/streamgres-demo/tests/scenario.rs b/crates/streamgres-demo/tests/scenario.rs new file mode 100644 index 0000000..9a1511e --- /dev/null +++ b/crates/streamgres-demo/tests/scenario.rs @@ -0,0 +1,910 @@ +//! Every lookup, through both wrappers, against a running engine: the +//! invariants the wrapper promises, checked on the shapes the desk +//! actually reads. +//! +//! Needs `STREAMGRES_S2S_URL` and `DATABASE_URL` (the database that +//! engine follows); skips itself otherwise. The tests that restart the +//! engine also need `STREAMGRES_RESTART_CMD`, a shell command that +//! restarts it and returns once it serves again, and skip without it. +//! Every test reseeds the dataset, so run them against a database +//! nobody is watching, one at a time (`--test-threads=1`). + +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::{Duration, Instant}; + +use diesel::sql_query; +use diesel::{Connection, ExpressionMethods, PgConnection, QueryDsl, SelectableHelper}; +use diesel_async::pooled_connection::AsyncDieselConnectionManager; +use diesel_async::pooled_connection::bb8::Pool; +use diesel_async::{AsyncConnection, AsyncPgConnection, RunQueryDsl as _}; +use diesel_streamgres::cache::{Cache, Limits}; +use diesel_streamgres::engine::EngineClient; +use diesel_streamgres::{AsyncSyncedConnection, Route, RouteDsl, Served, SyncedConnection}; +use serde_json::Value; +use streamgres_demo::migrate::Step; +use streamgres_demo::models::{Ticket, TicketKind}; +use streamgres_demo::pools::Stack; +use streamgres_demo::queries::{Expected, Lookup, SyncConn}; +use streamgres_demo::schema::tickets; +use streamgres_demo::{seed, verify, writer}; + +/// The demo against a live Postgres and engine: the wrappers on one +/// cache, and a plain connection for the truth and the changes. +struct Live { + cache: Arc, + engine: EngineClient, + /// The async wrapper, routed Streamgres. + main: AsyncSyncedConnection, + /// The blocking wrapper on the same cache. + sync: SyncConn, + /// No wrapper: what Postgres says, and where the changes come from. + plain: AsyncPgConnection, + database_url: String, +} + +impl Live { + /// Reseeds, connects, and gives the engine a moment to hear of the + /// seed's DDL on a fresh database. `None` without an engine. + async fn start() -> Option { + Self::start_with(Limits { + refused_for: None, + ..Limits::default() + }) + .await + } + + /// As [`Live::start`], with the cache's limits chosen. No refusal + /// memory, so a refusal under test is re-asked rather than remembered. + async fn start_with(limits: Limits) -> Option { + let Ok(url) = std::env::var("STREAMGRES_S2S_URL") else { + eprintln!("STREAMGRES_S2S_URL not set; skipping"); + return None; + }; + let database_url = std::env::var("DATABASE_URL").expect("DATABASE_URL"); + let mut seeder = PgConnection::establish(&database_url).expect("connect"); + Step::Drop.run(&mut seeder).expect("undo the migrations"); + seed::run(&mut seeder).expect("seed"); + let engine = EngineClient::connect(url); + let cache = Arc::new(Cache::with_limits(limits, Arc::new(engine.clone()))); + engine.attach(&cache); + if !connected(&engine, Duration::from_secs(5)).await { + assert!( + restart(&engine).await, + "the engine did not connect (a migration left behind may have stopped it; set STREAMGRES_RESTART_CMD)" + ); + } + let main = AsyncSyncedConnection::with_cache( + AsyncPgConnection::establish(&database_url) + .await + .expect("connect"), + Arc::clone(&cache), + Route::Streamgres, + ); + let sync = SyncedConnection::with_cache( + PgConnection::establish(&database_url).expect("connect"), + Arc::clone(&cache), + Route::Streamgres, + ); + let plain = AsyncPgConnection::establish(&database_url) + .await + .expect("connect"); + tokio::time::sleep(Duration::from_millis(300)).await; + Some(Live { + cache, + engine, + main, + sync, + plain, + database_url, + }) + } + + /// One of the cache's counters now. + fn stat(&self, counter: &AtomicU64) -> u64 { + AtomicU64::load(counter, Ordering::Relaxed) + } + + /// What Postgres says, through the blocking wrapper routed Postgres. + fn truth(&mut self, lookup: &Lookup) -> Value { + self.sync + .with_route(Route::Postgres, |conn| lookup.run_sync(conn).unwrap()) + } + + /// Runs `sql` on the plain connection, as another service would. + async fn change(&mut self, sql: &str) { + sql_query(sql) + .execute(&mut self.plain) + .await + .expect("change"); + } + + /// Reads through the async wrapper until the result differs from + /// `before` or the wait runs out. + async fn until_changed( + &mut self, + lookup: &Lookup, + before: &Value, + wait: Duration, + ) -> Option<(Value, Duration)> { + let started = Instant::now(); + while started.elapsed() < wait { + let now = lookup.run_async(&mut self.main).await.unwrap(); + if &now != before { + return Some((now, started.elapsed())); + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + None + } + + /// Runs `STREAMGRES_RESTART_CMD` and waits for the client to be + /// back on the engine. `false` when the variable is not set. + async fn restart_engine(&self) -> bool { + restart(&self.engine).await + } +} + +/// Whether the client is on the engine within `wait`. +async fn connected(engine: &EngineClient, wait: Duration) -> bool { + let deadline = Instant::now() + wait; + while !engine.is_connected() { + if Instant::now() >= deadline { + return false; + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + true +} + +/// Runs `STREAMGRES_RESTART_CMD` and waits for the client to be back on +/// the engine. `false` when the variable is not set. +async fn restart(engine: &EngineClient) -> bool { + let Ok(restart) = std::env::var("STREAMGRES_RESTART_CMD") else { + eprintln!("STREAMGRES_RESTART_CMD not set; skipping"); + return false; + }; + let status = tokio::process::Command::new("sh") + .arg("-c") + .arg(&restart) + .status() + .await + .expect("restart command"); + assert!(status.success(), "the restart command failed"); + assert!( + connected(engine, Duration::from_secs(15)).await, + "the engine did not come back" + ); + true +} + +/// A change that each lookup's result reflects, or, for a static +/// lookup, a change nearby that must leave it as it is. +fn change_for(lookup: &Lookup) -> &'static str { + match lookup { + Lookup::Ticket { .. } => { + "UPDATE tickets SET priority = 'high', updated_at = now() WHERE id = 1" + } + Lookup::Queue { .. } => { + "UPDATE tickets SET title = title || '!' WHERE id = (SELECT id FROM tickets WHERE org_id = 1 AND status IN ('open', 'pending') ORDER BY created_at DESC, id ASC OFFSET 10 LIMIT 1)" + } + Lookup::QueueCount { .. } => "UPDATE tickets SET status = 'open' WHERE id = 2", + Lookup::Comments { .. } => { + "INSERT INTO comments (id, ticket_id, body) VALUES (99001, 1, 'scenario') ON CONFLICT (id) DO UPDATE SET body = comments.body || '!'" + } + Lookup::Attachment { .. } => { + "UPDATE attachments SET data = data || '\\x00'::bytea WHERE id = 1" + } + Lookup::Agents { .. } => { + "UPDATE agents SET skills = array_append(skills, 'scenario') WHERE id = 1" + } + Lookup::Sla { .. } => "UPDATE tickets SET sla = sla + interval '1 minute' WHERE id = 3", + Lookup::Search { .. } => "UPDATE tickets SET title = title || '?' WHERE id = 1", + Lookup::Recent { .. } => "UPDATE tickets SET updated_at = now() WHERE id = 7", + Lookup::OrgTicket { .. } => "UPDATE tickets SET title = title || '!' WHERE id = 5", + Lookup::Unassigned { .. } => "UPDATE tickets SET agent_id = NULL WHERE id = 1", + Lookup::Escalated { .. } => "UPDATE tickets SET priority = 'high' WHERE id = 2", + Lookup::OpenExcept { .. } => "UPDATE tickets SET status = 'closed' WHERE id = 4", + Lookup::CreatedBetween { .. } => { + "UPDATE tickets SET created_at = '2026-10-01 11:59:00+00' WHERE id = 6" + } + Lookup::Many { .. } => "UPDATE tickets SET title = title || '!' WHERE id = 50", + Lookup::QueueEnd { .. } | Lookup::Nothing { .. } => { + "UPDATE tickets SET title = title || '!' WHERE id = 8" + } + Lookup::AllTickets { .. } => "UPDATE tickets SET title = title || '!' WHERE id = 9", + Lookup::CommentAuthors { .. } => "UPDATE comments SET agent_id = 3 WHERE id = 12", + Lookup::Skilled { .. } => { + "UPDATE agents SET skills = array_append(skills, 'billing') WHERE id = 2" + } + Lookup::Statuses { .. } => { + "UPDATE tickets SET status = 'open' WHERE org_id = 1 AND status = 'closed'" + } + Lookup::Workload { .. } => "UPDATE tickets SET agent_id = 1 WHERE id = 10", + Lookup::LastActivity { .. } => "UPDATE tickets SET updated_at = now() WHERE id = 11", + Lookup::OfKind { .. } => "UPDATE tickets SET kind = 'bug' WHERE id = 2", + } +} + +/// Every demo lookup matches Postgres on the first read and after a +/// change, is served the way it expects (rows, refills or live), and +/// shares its entry between the two wrappers. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn every_lookup_is_served_correctly_through_both_wrappers() { + let Some(mut live) = Live::start().await else { + return; + }; + live.main.set_read_your_writes(Duration::ZERO); + live.sync.set_read_your_writes(Duration::ZERO); + println!( + "{:<16} {:<8} {:<7} {:<7} {:>8} exercises", + "lookup", "mode", "hit", "sync", "seen in" + ); + for lookup in Lookup::demo() { + let name = lookup.name(); + let hits_before = live.stat(&live.cache.stats.hits); + let refused_before = live.stat(&live.cache.stats.refused); + let materialized_before = live.engine.materialized(); + let truth = live.truth(&lookup); + let mut hit_on = None; + for read in 1..=3 { + let seen = lookup.run_async(&mut live.main).await.unwrap(); + assert_eq!(seen, truth, "{name}: read {read} matches Postgres"); + if live.stat(&live.cache.stats.hits) > hits_before { + hit_on = Some(read); + break; + } + } + let hits_so_far = live.stat(&live.cache.stats.hits); + let through_sync = lookup.run_sync(&mut live.sync).unwrap(); + assert_eq!(through_sync, truth, "{name}: the blocking wrapper agrees"); + let sync_hit = live.stat(&live.cache.stats.hits) > hits_so_far; + let mode = match lookup.expected() { + Expected::Live => { + assert!(hit_on.is_none(), "{name}: a refused shape never hits"); + assert!( + live.stat(&live.cache.stats.refused) > refused_before, + "{name}: the engine refused it" + ); + "live" + } + expected => { + assert!( + hit_on.is_some(), + "{name}: no hit in three reads; refusals {:?}", + live.engine.refusals() + ); + assert!(sync_hit, "{name}: the blocking wrapper shares the entry"); + let rows = live.engine.materialized() > materialized_before; + assert_eq!( + rows, + expected == Expected::Rows, + "{name}: expected {expected:?}; fallbacks {:?}", + live.engine.fallbacks() + ); + if rows { "rows" } else { "refill" } + } + }; + + live.change(change_for(&lookup)).await; + let seen = if lookup.is_static() { + tokio::time::sleep(Duration::from_millis(300)).await; + let after = lookup.run_async(&mut live.main).await.unwrap(); + assert_eq!(after, truth, "{name}: a change nearby leaves it as it is"); + assert_eq!(after, live.truth(&lookup), "{name}: and Postgres agrees"); + Duration::ZERO + } else if lookup.is_cacheable() { + let (after, seen_in) = live + .until_changed(&lookup, &truth, Duration::from_secs(5)) + .await + .unwrap_or_else(|| panic!("{name}: the change was not seen")); + assert_eq!( + after, + live.truth(&lookup), + "{name}: after the change, the cache matches Postgres" + ); + seen_in + } else { + let after = lookup.run_async(&mut live.main).await.unwrap(); + assert_eq!( + after, + live.truth(&lookup), + "{name}: live, so the change is seen at once" + ); + Duration::ZERO + }; + println!( + "{:<16} {:<8} {:<7} {:<7} {:>6}ms {}{}", + name, + mode, + hit_on + .map(|read| format!("read {read}")) + .unwrap_or_else(|| "-".to_string()), + if sync_hit { "shared" } else { "-" }, + seen.as_millis(), + lookup.exercises(), + if mode == "refill" { + format!( + " [{}]", + live.engine.fallbacks().last().cloned().unwrap_or_default() + ) + } else { + String::new() + } + ); + } +} + +/// A write through one wrapper is seen by the next read on another of +/// the same cache, inside the read-your-writes window. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn a_write_on_one_connection_is_seen_by_the_pools_other_connections() { + let Some(mut live) = Live::start().await else { + return; + }; + live.main.set_read_your_writes(Duration::from_secs(5)); + live.sync.set_read_your_writes(Duration::from_secs(5)); + let lookup = Lookup::Ticket { id: 20 }; + let before = lookup.run_async(&mut live.main).await.unwrap(); + let _ = lookup.run_async(&mut live.main).await.unwrap(); + diesel::update(tickets::table.filter(tickets::id.eq(20))) + .set(tickets::title.eq("read your writes")) + .execute(&mut live.main) + .await + .unwrap(); + let after = lookup.run_sync(&mut live.sync).unwrap(); + assert_ne!(after, before); + assert_eq!(after[0]["title"].as_str(), Some("read your writes")); + assert!(live.stat(&live.cache.stats.after_write) >= 1); +} + +/// A flag on the query beats the connection's route both ways, a scope +/// routes everything inside it, and the read-your-writes window ends. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn a_query_flag_beats_the_route_and_the_write_window_ends() { + let Some(mut live) = Live::start().await else { + return; + }; + live.main.set_read_your_writes(Duration::from_millis(300)); + let query = || { + tickets::table + .filter(tickets::id.eq(30)) + .select(Ticket::as_select()) + }; + let cached: Vec = query().load(&mut live.main).await.unwrap(); + let cached_again: Vec = query().load(&mut live.main).await.unwrap(); + assert_eq!(cached, cached_again); + assert_eq!(live.main.last_served(), Some(Served::Cache)); + let hits = live.stat(&live.cache.stats.hits); + + let flagged: Vec = query().via_postgres().load(&mut live.main).await.unwrap(); + assert_eq!(flagged, cached); + assert_eq!(live.main.last_served(), Some(Served::Postgres)); + assert_eq!( + live.stat(&live.cache.stats.hits), + hits, + "the flag skipped the cache" + ); + + live.main.set_route(Route::Postgres); + let _: Vec = query().load(&mut live.main).await.unwrap(); + assert_eq!(live.main.last_served(), Some(Served::Postgres)); + let _: Vec = query().via_streamgres().load(&mut live.main).await.unwrap(); + assert_eq!( + live.main.last_served(), + Some(Served::Cache), + "the flag reached the cache from a Postgres-routed connection" + ); + live.main.set_route(Route::Streamgres); + + let scoped = live + .main + .with_route(Route::Postgres, |conn| { + Box::pin(async move { + let rows: Vec = query().load(conn).await.unwrap(); + (rows, conn.last_served()) + }) + }) + .await; + assert_eq!( + scoped.1, + Some(Served::Postgres), + "the scope routed the read" + ); + assert_eq!( + live.main.route(), + Route::Streamgres, + "and the route came back" + ); + + diesel::update(tickets::table.filter(tickets::id.eq(30))) + .set(tickets::title.eq("inside the window")) + .execute(&mut live.main) + .await + .unwrap(); + let after_write = live.stat(&live.cache.stats.after_write); + let inside: Vec = query().load(&mut live.main).await.unwrap(); + assert_eq!(inside[0].title, "inside the window"); + assert_eq!(live.main.last_served(), Some(Served::Postgres)); + assert_eq!(live.stat(&live.cache.stats.after_write), after_write + 1); + tokio::time::sleep(Duration::from_millis(500)).await; + let outside: Vec = query().load(&mut live.main).await.unwrap(); + assert_eq!( + outside[0].title, "inside the window", + "the engine carried the write" + ); + assert_eq!( + live.main.last_served(), + Some(Served::Cache), + "the window ended, so the read is a hit again" + ); +} + +/// A result over the per-entry cap is served and correct, but not kept: +/// no hit, no subscription left behind. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn an_oversized_result_is_served_but_not_kept() { + let Some(mut live) = Live::start_with(Limits { + max_entry_bytes: 4096, + refused_for: None, + ..Limits::default() + }) + .await + else { + return; + }; + let lookup = Lookup::AllTickets { org: 1 }; + let truth = live.truth(&lookup); + let subscribed = live.engine.subscribed(); + for read in 1..=3 { + let seen = lookup.run_async(&mut live.main).await.unwrap(); + assert_eq!(seen, truth, "read {read}"); + assert_ne!( + live.main.last_served(), + Some(Served::Cache), + "read {read} is not a hit" + ); + } + assert!(live.stat(&live.cache.stats.oversized) >= 1); + assert_eq!(live.stat(&live.cache.stats.hits), 0); + let deadline = Instant::now() + Duration::from_secs(5); + while live.engine.subscribed() != subscribed && Instant::now() < deadline { + tokio::time::sleep(Duration::from_millis(20)).await; + } + assert_eq!( + live.engine.subscribed(), + subscribed, + "the subscription ended" + ); + let small = Lookup::Ticket { id: 1 }; + let _ = small.run_async(&mut live.main).await.unwrap(); + let _ = small.run_async(&mut live.main).await.unwrap(); + assert_eq!( + live.main.last_served(), + Some(Served::Cache), + "a small result is kept" + ); +} + +/// A read dropped before it completed leaves no subscription behind, +/// and the same read completed afterwards is cached as usual. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn a_read_dropped_mid_flight_leaves_no_subscription_behind() { + let Some(mut live) = Live::start().await else { + return; + }; + let lookup = Lookup::Queue { + org: 2, + statuses: vec![streamgres_demo::models::TicketStatus::Closed], + page: 0, + }; + let subscribed = live.engine.subscribed(); + let mut read = Box::pin(lookup.run_async(&mut live.main)); + assert!(futures_util::poll!(read.as_mut()).is_pending()); + drop(read); + let deadline = Instant::now() + Duration::from_secs(5); + while live.engine.subscribed() != subscribed && Instant::now() < deadline { + tokio::time::sleep(Duration::from_millis(20)).await; + } + assert_eq!( + live.engine.subscribed(), + subscribed, + "the subscription the dropped read made is gone" + ); + assert!(live.cache.is_empty()); + for _ in 0..3 { + let _ = lookup.run_async(&mut live.main).await.unwrap(); + } + assert_eq!( + live.engine.subscribed(), + subscribed + 1, + "refusals {:?}", + live.engine.refusals() + ); +} + +/// The process-wide stack the pool-based tests share: the wrapper's +/// settings install once per process. +async fn stack() -> Option> { + static STACK: tokio::sync::OnceCell> = tokio::sync::OnceCell::const_new(); + let (Ok(url), Ok(database_url)) = ( + std::env::var("STREAMGRES_S2S_URL"), + std::env::var("DATABASE_URL"), + ) else { + eprintln!("STREAMGRES_S2S_URL not set; skipping"); + return None; + }; + let stack = STACK + .get_or_init(|| async { + let stack = Stack::connect(&database_url, &url).await.expect("stack"); + assert!(stack.wait_for_engine(Duration::from_secs(5)).await); + Arc::new(stack) + }) + .await; + Some(Arc::clone(stack)) +} + +/// A plain pool for the writer, as the binary gives it. +async fn plain_pool(database_url: &str) -> Pool { + Pool::builder() + .max_size(4) + .build(AsyncDieselConnectionManager::::new( + database_url, + )) + .await + .expect("plain pool") +} + +/// The lookup's answer through a client that subscribes afresh, on its +/// own cache: what the engine says now, free of this process's stores. +/// Tells an engine-side answer from a client-side one in a mismatch. +async fn fresh_answer(name: &str) -> Value { + let lookup = Lookup::named(name).expect("a demo lookup"); + let url = std::env::var("STREAMGRES_S2S_URL").unwrap(); + let database_url = std::env::var("DATABASE_URL").unwrap(); + let engine = EngineClient::connect(url); + let cache = Arc::new(Cache::new(16 * 1024 * 1024, Arc::new(engine.clone()))); + engine.attach(&cache); + assert!(connected(&engine, Duration::from_secs(10)).await); + let mut conn = AsyncSyncedConnection::with_cache( + AsyncPgConnection::establish(&database_url).await.unwrap(), + cache, + Route::Streamgres, + ); + let mut last = Value::Null; + for _ in 0..3 { + last = lookup.run_async(&mut conn).await.unwrap(); + if conn.last_served() == Some(Served::Cache) { + break; + } + } + last +} + +/// Pairs that never agreed, told apart by where the wrong answer lives: +/// in this process's cache (a fresh subscription agrees with Postgres) +/// or in the engine (a fresh subscription disagrees the same way). +#[derive(Debug, Default)] +struct Disagreements { + client: u64, + engine: u64, +} + +impl Disagreements { + /// Fails on either kind, named, so the failure says whose answer + /// was wrong. + fn assert_none(&self) { + assert_eq!(self.client, 0, "pairs the cache got wrong: {self:?}"); + assert_eq!(self.engine, 0, "pairs the engine got wrong: {self:?}"); + } +} + +/// Runs `readers` tasks for `duration`, each comparing every lookup +/// across the routes; the pairs that never agreed. +async fn read_under_load(stack: &Arc, readers: usize, duration: Duration) -> Disagreements { + let mut tasks = Vec::new(); + for _ in 0..readers { + let stack = Arc::clone(stack); + tasks.push(tokio::spawn(async move { + let lookups = Lookup::demo(); + let started = Instant::now(); + let mut found = Disagreements::default(); + while started.elapsed() < duration { + let outcome = verify::once(&stack, &lookups, Duration::from_secs(3)).await; + for mismatch in &outcome.mismatches { + let fresh = fresh_answer(&mismatch.lookup).await; + let where_ = if fresh == mismatch.postgres { + found.client += 1; + "CLIENT" + } else { + found.engine += 1; + "ENGINE" + }; + eprintln!( + "MISMATCH ({where_}) {}: cache {} | postgres {}", + mismatch.lookup, mismatch.cached, mismatch.postgres + ); + } + } + found + })); + } + let mut found = Disagreements::default(); + for task in tasks { + let part = task.await.expect("reader"); + found.client += part.client; + found.engine += part.engine; + } + found +} + +/// Sixteen readers comparing every lookup across the routes while the +/// writer runs its whole mix: no pair ever disagrees for long. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn concurrent_readers_under_the_writer_never_see_a_wrong_row() { + let Some(_seeded) = Live::start().await else { + return; + }; + let stack = stack().await.unwrap(); + let database_url = std::env::var("DATABASE_URL").unwrap(); + let counts = Arc::new(writer::Counts::default()); + let writer = writer::spawn( + plain_pool(&database_url).await, + Duration::from_millis(20), + writer::OPS.to_vec(), + Arc::clone(&counts), + ); + let found = read_under_load(&stack, 16, Duration::from_secs(8)).await; + writer.abort(); + println!("writer {}", counts.json()); + println!("verify {}", stack.verify.json()); + found.assert_none(); + assert!(counts.total() >= 100, "the writer ran: {}", counts.json()); + assert_eq!(counts.errors(), 0, "{}", counts.json()); + for op in writer::OPS { + assert!( + counts.of(op) > 0, + "the mix reached `{op}`: {}", + counts.json() + ); + } +} + +/// The engine restarted in the middle of the same load: the readers go +/// to Postgres while it is away, resubscribe when it is back, and +/// neither the cache nor the engine disagrees with Postgres for long. +/// This is the test that found the engine reading a timestamp range in +/// the server's own zone: after a restart, with no replay buffer to +/// hide it, every row of the range shape was missing until its next +/// write. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn an_engine_restart_under_traffic_loses_no_correctness() { + let Some(live) = Live::start().await else { + return; + }; + if std::env::var("STREAMGRES_RESTART_CMD").is_err() { + eprintln!("STREAMGRES_RESTART_CMD not set; skipping"); + return; + } + let stack = stack().await.unwrap(); + let database_url = std::env::var("DATABASE_URL").unwrap(); + let counts = Arc::new(writer::Counts::default()); + let writer = writer::spawn( + plain_pool(&database_url).await, + Duration::from_millis(20), + writer::OPS.to_vec(), + Arc::clone(&counts), + ); + let readers = { + let stack = Arc::clone(&stack); + tokio::spawn(async move { read_under_load(&stack, 8, Duration::from_secs(10)).await }) + }; + tokio::time::sleep(Duration::from_secs(3)).await; + assert!(live.restart_engine().await); + let deadline = Instant::now() + Duration::from_secs(10); + while !stack.engine.is_connected() && Instant::now() < deadline { + tokio::time::sleep(Duration::from_millis(50)).await; + } + assert!(stack.engine.is_connected(), "the stack's client came back"); + let found = readers.await.expect("readers"); + writer.abort(); + println!("writer {}", counts.json()); + println!("verify {}", stack.verify.json()); + found.assert_none(); + assert!(stack.engine.subscribed() > 0, "subscriptions are back"); +} + +/// The three kinds of schema change: a widened table is absorbed and +/// the cache keeps serving; a new enum type is absorbed and learnt on +/// its first bind; a dropped column stops the engine, reads go to +/// Postgres meanwhile, and the restarted engine serves the schema as it +/// is. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn migrations_are_absorbed_learnt_or_stop_the_engine() { + let Some(mut live) = Live::start().await else { + return; + }; + if std::env::var("STREAMGRES_RESTART_CMD").is_err() { + eprintln!("STREAMGRES_RESTART_CMD not set; skipping"); + return; + } + live.main.set_read_your_writes(Duration::ZERO); + let mut migrator = PgConnection::establish(&live.database_url).unwrap(); + let ticket = Lookup::Ticket { id: 1 }; + let truth = live.truth(&ticket); + assert_eq!(ticket.run_async(&mut live.main).await.unwrap(), truth); + assert_eq!(ticket.run_async(&mut live.main).await.unwrap(), truth); + assert_eq!(live.main.last_served(), Some(Served::Cache)); + + Step::Widen.run(&mut migrator).unwrap(); + tokio::time::sleep(Duration::from_millis(500)).await; + assert!(live.engine.is_connected(), "a widened table is absorbed"); + assert_eq!(ticket.run_async(&mut live.main).await.unwrap(), truth); + assert_eq!( + live.main.last_served(), + Some(Served::Cache), + "the entry survived" + ); + live.change(change_for(&ticket)).await; + let (after, _) = live + .until_changed(&ticket, &truth, Duration::from_secs(5)) + .await + .expect("the change after the migration was seen"); + assert_eq!(after, live.truth(&ticket)); + + Step::NewType.run(&mut migrator).unwrap(); + tokio::time::sleep(Duration::from_millis(500)).await; + assert!(live.engine.is_connected(), "a new type is absorbed"); + let of_kind = Lookup::OfKind { + org: 1, + kind: TicketKind::Bug, + }; + let kinds = live.truth(&of_kind); + assert!(kinds.as_array().is_some_and(|rows| !rows.is_empty())); + for read in 1..=3 { + assert_eq!( + of_kind.run_async(&mut live.main).await.unwrap(), + kinds, + "read {read}" + ); + assert_ne!( + live.main.last_served(), + Some(Served::Cache), + "read {read}: refused" + ); + } + assert!( + live.engine + .refusals() + .iter() + .any(|reason| reason.contains("bind of type oid")), + "{:?}", + live.engine.refusals() + ); + + assert!(live.restart_engine().await); + let materialized = live.engine.materialized(); + let mut hit = false; + for _ in 1..=3 { + assert_eq!(of_kind.run_async(&mut live.main).await.unwrap(), kinds); + if live.main.last_served() == Some(Served::Cache) { + hit = true; + break; + } + } + assert!( + hit, + "after the restart the new type's bind is served: {:?}", + live.engine.refusals() + ); + assert_eq!( + live.engine.materialized(), + materialized + 1, + "in rows mode; fallbacks {:?}, refusals {:?}", + live.engine.fallbacks(), + live.engine.refusals() + ); + + Step::Drop.run(&mut migrator).unwrap(); + let deadline = Instant::now() + Duration::from_secs(10); + while live.engine.is_connected() && Instant::now() < deadline { + tokio::time::sleep(Duration::from_millis(50)).await; + } + assert!( + !live.engine.is_connected(), + "a dropped column stops the engine" + ); + let truth = live.truth(&ticket); + assert_eq!(ticket.run_async(&mut live.main).await.unwrap(), truth); + assert_eq!( + live.main.last_served(), + Some(Served::Postgres), + "served live meanwhile" + ); + + assert!(live.restart_engine().await); + for _ in 1..=3 { + assert_eq!(ticket.run_async(&mut live.main).await.unwrap(), truth); + } + assert_eq!(live.main.last_served(), Some(Served::Cache), "cached again"); +} + +/// A column added while the engine is down. The wrapper learnt the +/// table's column types before the outage; when the engine is back and +/// serves a query that names the new column, those types lack it, so +/// they are looked up again and the query is served from engine rows. +#[tokio::test] +async fn a_column_added_while_the_engine_was_down_is_served_when_it_returns() { + let Some(mut live) = Live::start().await else { + return; + }; + if std::env::var("STREAMGRES_RESTART_CMD").is_err() { + eprintln!("STREAMGRES_RESTART_CMD not set; skipping"); + return; + } + live.main.set_read_your_writes(Duration::ZERO); + let mut migrator = PgConnection::establish(&live.database_url).unwrap(); + let ticket = Lookup::Ticket { id: 1 }; + let truth = live.truth(&ticket); + assert_eq!(ticket.run_async(&mut live.main).await.unwrap(), truth); + assert_eq!(ticket.run_async(&mut live.main).await.unwrap(), truth); + assert_eq!(live.main.last_served(), Some(Served::Cache)); + + Step::Widen.run(&mut migrator).unwrap(); + Step::Drop.run(&mut migrator).unwrap(); + let deadline = Instant::now() + Duration::from_secs(10); + while live.engine.is_connected() && Instant::now() < deadline { + tokio::time::sleep(Duration::from_millis(50)).await; + } + assert!( + !live.engine.is_connected(), + "the dropped column stopped the engine" + ); + + Step::NewType.run(&mut migrator).unwrap(); + assert_eq!(ticket.run_async(&mut live.main).await.unwrap(), truth); + assert_eq!( + live.main.last_served(), + Some(Served::Postgres), + "served live while the engine is down" + ); + + assert!(live.restart_engine().await); + let of_kind = Lookup::OfKind { + org: 1, + kind: TicketKind::Bug, + }; + let kinds = live.truth(&of_kind); + assert!(kinds.as_array().is_some_and(|rows| !rows.is_empty())); + let materialized = live.engine.materialized(); + let mut served = Vec::new(); + for _ in 1..=4 { + assert_eq!(of_kind.run_async(&mut live.main).await.unwrap(), kinds); + served.push(live.main.last_served()); + } + assert_eq!( + served, + [ + Some(Served::Postgres), + Some(Served::Postgres), + Some(Served::Engine), + Some(Served::Cache) + ], + "live while the new type's OID is resolved, live again while the table's types are fetched afresh, then engine rows, then the cache; refusals {:?}, fallbacks {:?}", + live.engine.refusals(), + live.engine.fallbacks() + ); + assert_eq!( + live.engine.materialized(), + materialized + 1, + "from engine rows, with the types looked up again" + ); + for _ in 1..=3 { + assert_eq!(ticket.run_async(&mut live.main).await.unwrap(), truth); + } + assert_eq!( + live.main.last_served(), + Some(Served::Cache), + "the old entry is back too" + ); +} diff --git a/demo/README.md b/demo/README.md new file mode 100644 index 0000000..e009dfc --- /dev/null +++ b/demo/README.md @@ -0,0 +1,172 @@ +# Streamgres in five minutes + +A ticket-desk backend whose reads are served from an in-memory cache the +Streamgres engine keeps fresh from Postgres's replication stream. One +compose file brings up Postgres, the engine and the backend; a +background writer changes a ticket every second so there is something to +keep fresh. + +```bash +docker compose -f demo/docker-compose.yml up --build +``` + +When the three services are up, open for the +stats page, or walk through it with curl. If a port is taken on your +machine, override it: `DEMO_PORT=8081 ENGINE_PORT=4850 +POSTGRES_PORT=5435 docker compose -f demo/docker-compose.yml up --build`. +Every port is published on this machine only, and the engine's S2S +endpoint asks for a token (`STREAMGRES_S2S_TOKEN`, `demo-token` unless +set): the stack holds a well-known database password. + +## 1. A read, twice + +```bash +curl -s localhost:8080/tickets/1 | jq '{source, micros, rows}' +curl -s localhost:8080/tickets/1 | jq '{source, micros}' +``` + +The first read subscribes the query with the engine, which ships the +rows (`"source": "engine"`, no Postgres read); the second is served +from the cache (`"source": "cache"`), with no round trip at all. + +## 2. A change from outside + +Change the ticket with psql, the way another service would: + +```bash +docker compose -f demo/docker-compose.yml exec postgres \ + psql -U streamgres -d tickets -c "UPDATE tickets SET priority = 'high' WHERE id = 1" +curl -s localhost:8080/tickets/1 | jq '{source, rows: [.rows[].priority]}' +``` + +The read still comes from the cache, and shows the new priority: the +engine saw the commit and replaced the cached rows before you could type +the second command. + +## 3. The shapes + +```bash +curl -s localhost:8080/lookups | jq -r '.[] | "\(.name): \(.exercises)"' +``` + +Each lookup is a query shape the wrapper handles differently: an enum +bound as an array, a page with an offset, a join composed from the +engine's rows per table, a left join, a `bytea`, a two-column key, `IS +NULL`, `OR`, `NOT IN`, a timestamp range, a hundred-id `IN` list, empty +results, an `interval` column the engine cannot ship (so that query is +refilled from Postgres on each change instead), and the shapes the +engine refuses and serves live: `ILIKE`, `DISTINCT`, `GROUP BY`, `MAX` +and the array operator `@>`. The stats page has a button per shape; +`crates/streamgres-demo/README.md` lists them with their paths. + +## 4. Routing + +A read can be sent to Postgres regardless of the cache, per query or +per pool: + +```bash +curl -s 'localhost:8080/tickets/1?route=postgres' | jq '{route, source}' +curl -s 'localhost:8080/tickets/1?pool=reporting' | jq '{route, source}' +``` + +The first flags the query itself. The second uses a pool whose +connections are routed to Postgres by the URL they were opened with, the +way a pool against a primary would be. + +## 5. Your own writes + +```bash +curl -s -X POST localhost:8080/tickets/1/status -H 'content-type: application/json' -d '{"status":"closed"}' +curl -s localhost:8080/tickets/1 | jq '{source, rows: [.rows[].status]}' +``` + +The write went through the backend's own pool, so for the next second +every read through the cache goes to Postgres (`"source": "postgres"`) +and sees the write at once, before the engine has shipped it. + +## 6. Counters + +```bash +curl -s localhost:8080/stats | jq +``` + +`cache.hits` and `cache.misses` are what they say; `cache.replaced` is +how many times the engine's deltas replaced cached rows; `touched` how +many times a change made the cache refill from Postgres (the `interval` +query); `after_write` the reads sent to Postgres inside the +read-your-writes window; `lagging` the reads sent to Postgres because the +engine was more than two seconds behind the database, as `engine.lag_ms` +measures it; `engine.materialized` the queries served from engine rows +rather than refills; `writer` how many changes of each kind +the background writer made; `verify` what the self-check found. + +## 7. The demo checks itself + +```bash +curl -s localhost:8080/verify | jq '{compared, lagged, mismatches}' +``` + +Every lookup is read through the cache and through Postgres and +compared. A pair that differs is read again for up to a second, since +the cache may lag a commit by the engine's propagation delay; one that +never agrees is a mismatch, returned with both answers. With +`DEMO_VERIFY_MS=500` the backend runs this pass every half second +in the background, under the writer, and `/stats` carries the totals. +The writer's mix covers inserts, deletes, moves between organizations, +renames under the joins, assignees flipped to `NULL` and back, +multi-row transactions, rollbacks and oversized comment bodies +(`DEMO_WRITER_MIX` narrows it), so the stack is a continuous +correctness test of the wrapper, not only a demo. + +## 8. Schema changes + +```bash +docker compose -f demo/docker-compose.yml exec demo /app/demo migrate widen +curl -s localhost:8080/tickets/1 | jq '{source}' +``` + +A column added with a constant default is absorbed: the read above is +still served from the cache. `migrate new-type` adds an enum type and a +column of it, which the engine absorbs too. On a connection, the first +two reads that bind the new type (`/orgs/1/kind/bug`) go to Postgres: +the first teaches the wrapper the type's OID, the second teaches the +engine the type and the wrapper the table's new column. From then on +the query is served from the cache. +`migrate drop` removes both: the engine cannot follow a +dropped column, so it stops, compose restarts it, and reads go to +Postgres meanwhile. A `TRUNCATE` of a published table stops it the same +way: the replication stream carries no rows for it to follow. + +## 9. How stale is the cache? + +```bash +docker compose -f demo/docker-compose.yml exec demo /app/demo lag --rate 20 --for 10 +``` + +prints how long a commit takes to reach the cached answer, as p50, p95, +p99 and max over ten seconds, under twenty background writes per +second. + +## Latency + +Inside the demo's container, so no HTTP or port forwarding is in the +measurement: + +```bash +docker compose -f demo/docker-compose.yml exec demo \ + /app/demo bench --lookup queue --concurrency 16 --requests 5000 +``` + +prints p50, p95, p99 and throughput for the Postgres route and the +Streamgres route side by side. + +## Running it against your own Postgres + +The backend is `crates/streamgres-demo`; `demo seed`, `demo +serve` and `demo bench` take `DATABASE_URL`, `STREAMGRES_S2S_URL` and, +when the engine asks for one, `STREAMGRES_S2S_TOKEN`; `demo verify` and +`demo migrate` too. `DEMO_SCALE=10` seeds ten times the tickets, so +pages run deep and the cache's byte limits matter. Postgres needs +`wal_level = logical`, the publication and the schema-change trigger +stack for the engine's app id (`demo/postgres-init.sql` and +`demo/ddl-triggers.sql` are both for app `demo`); see `docs/deploy.md`. diff --git a/demo/ddl-triggers.sql b/demo/ddl-triggers.sql new file mode 100644 index 0000000..48c3900 --- /dev/null +++ b/demo/ddl-triggers.sql @@ -0,0 +1,322 @@ +-- The schema-change trigger stack for the demo: app `demo`, shard 0, +-- publication `demo_pub`, as the engine renders it +-- (`streamgres::client::ddl_triggers::trigger_stack_sql("demo", 0, +-- &["demo_pub"])`). The engine hears every migration through it. +CREATE SCHEMA IF NOT EXISTS "demo_0"; + +CREATE OR REPLACE FUNCTION "demo_0".get_trigger_context() +RETURNS record AS $$ +DECLARE + result record; +BEGIN + SELECT COALESCE(current_query(), 'current_query() returned NULL') AS "query" into result; + RETURN result; +END +$$ LANGUAGE plpgsql; + +CREATE OR REPLACE FUNCTION "demo_0".notice_ignore(reason TEXT, tag TEXT, target record) +RETURNS void AS $$ +BEGIN + RAISE NOTICE 'demo_0 ignoring % % %', reason, tag, + COALESCE(row_to_json(target)::text, ''); +END +$$ LANGUAGE plpgsql; + +DROP FUNCTION IF EXISTS "demo_0".schema_specs(); +CREATE FUNCTION "demo_0".schema_specs() +RETURNS JSON +STABLE +AS $$ +WITH published_columns AS (SELECT + pc.oid::int8 AS "oid", + nspname AS "schema", + pc.relnamespace::int8 AS "schemaOID" , + pc.relname AS "name", + pc.relreplident AS "replicaIdentity", + attnum AS "pos", + attname AS "col", + pt.typname AS "type", + atttypid::int8 AS "typeOID", + pt.typtype, + elem_pt.typtype AS "elemTyptype", + coll.collisdeterministic AS "collationIsDeterministic", + NULLIF(atttypmod, -1) AS "maxLen", + attndims "arrayDims", + attnotnull AS "notNull", + pg_get_expr(pd.adbin, pd.adrelid) as "dflt", + NULLIF(ARRAY_POSITION(conkey, attnum), -1) AS "keyPos", + pb.rowfilter as "rowFilter", + pb.pubname as "publication" +FROM pg_attribute +JOIN pg_class pc ON pc.oid = attrelid +JOIN pg_namespace pns ON pns.oid = relnamespace +JOIN pg_type pt ON atttypid = pt.oid +LEFT JOIN pg_type elem_pt ON elem_pt.oid = pt.typelem +LEFT JOIN pg_collation coll ON coll.oid = pg_attribute.attcollation +JOIN pg_publication_tables as pb ON + pb.schemaname = nspname AND + pb.tablename = pc.relname AND + attname = ANY(pb.attnames) +LEFT JOIN pg_constraint pk ON pk.contype = 'p' AND pk.connamespace = relnamespace AND pk.conrelid = attrelid +LEFT JOIN pg_attrdef pd ON pd.adrelid = attrelid AND pd.adnum = attnum +WHERE pb.pubname IN ('demo_pub') AND + (current_setting('server_version_num')::int >= 160000 OR attgenerated = '') +ORDER BY nspname, pc.relname), + +tables AS (SELECT json_build_object( + 'oid', "oid", + 'schema', "schema", + 'schemaOID', "schemaOID", + 'name', "name", + 'replicaIdentity', "replicaIdentity", + 'columns', json_object_agg( + DISTINCT + col, + jsonb_build_object( + 'pos', "pos", + 'dataType', CASE WHEN "arrayDims" = 0 + THEN "type" + ELSE substring("type" from 2) || repeat('[]', "arrayDims") END, + 'pgTypeClass', "typtype", + 'elemPgTypeClass', "elemTyptype", + 'typeOID', "typeOID", + 'collationIsDeterministic', "collationIsDeterministic", + 'characterMaximumLength', CASE WHEN "typeOID" = 1043 OR "typeOID" = 1042 + THEN "maxLen" - 4 + ELSE "maxLen" END, + 'notNull', "notNull", + 'dflt', "dflt" + ) + ), + 'primaryKey', ARRAY( SELECT json_object_keys( + json_strip_nulls( + json_object_agg( + DISTINCT "col", "keyPos" ORDER BY "keyPos" + ) + ) + )), + 'publications', json_object_agg( + DISTINCT + "publication", + jsonb_build_object('rowFilter', "rowFilter") + ) +) AS "table" FROM published_columns + GROUP BY "schema", "schemaOID", "name", "oid", "replicaIdentity"), + + indexed_columns AS (SELECT + pg_indexes.schemaname as "schema", + pg_indexes.tablename as "tableName", + pg_indexes.indexname as "name", + index_column.name as "col", + CASE WHEN pg_index.indoption[index_column.pos-1] & 1 = 1 THEN 'DESC' ELSE 'ASC' END as "dir", + (pg_index.indisunique AND pg_index.indpred IS NULL) as "unique", + pg_index.indisprimary as "isPrimaryKey", + pg_index.indisreplident as "isReplicaIdentity", + pg_index.indimmediate as "isImmediate", + pg_get_expr(pg_index.indpred, pg_index.indrelid, false) as "predicateSQL" + FROM pg_indexes + JOIN pg_namespace ON pg_indexes.schemaname = pg_namespace.nspname + JOIN pg_class pc ON + pc.relname = pg_indexes.indexname + AND pc.relnamespace = pg_namespace.oid + JOIN pg_publication_tables as pb ON + pb.schemaname = pg_indexes.schemaname AND + pb.tablename = pg_indexes.tablename + JOIN pg_index ON pg_index.indexrelid = pc.oid + JOIN LATERAL ( + SELECT array_agg(attname) as attnames, array_agg(attgenerated != '') as generated FROM pg_attribute + WHERE attrelid = pg_index.indrelid + AND attnum = ANY( (pg_index.indkey::smallint[] )[:pg_index.indnkeyatts - 1] ) + ) as indexed ON true + JOIN LATERAL ( + SELECT pg_attribute.attname as name, col.index_pos as pos + FROM UNNEST( (pg_index.indkey::smallint[])[:pg_index.indnkeyatts - 1] ) + WITH ORDINALITY as col(table_pos, index_pos) + JOIN pg_attribute ON attrelid = pg_index.indrelid AND attnum = col.table_pos + ) AS index_column ON true + LEFT JOIN pg_constraint ON pg_constraint.conindid = pc.oid + WHERE pb.pubname IN ('demo_pub') + AND pg_index.indexprs IS NULL + AND (true OR pg_index.indpred IS NULL) + AND (pg_constraint.contype IS NULL OR pg_constraint.contype IN ('p', 'u')) + AND indexed.attnames <@ pb.attnames + AND (current_setting('server_version_num')::int >= 160000 OR false = ALL(indexed.generated)) + ORDER BY + pg_indexes.schemaname, + pg_indexes.tablename, + pg_indexes.indexname, + index_column.pos ASC), + + indexes AS (SELECT json_build_object( + 'schema', "schema", + 'tableName', "tableName", + 'name', "name", + 'unique', "unique", + 'isPrimaryKey', "isPrimaryKey", + 'isReplicaIdentity', "isReplicaIdentity", + 'isImmediate', "isImmediate", + 'predicateSQL', "predicateSQL", + 'columns', json_object_agg("col", "dir") + ) AS index FROM indexed_columns + GROUP BY "schema", "tableName", "name", "unique", + "isPrimaryKey", "isReplicaIdentity", "isImmediate", "predicateSQL") + + SELECT json_build_object( + 'tables', COALESCE((SELECT json_agg("table") FROM tables), '[]'::json), + 'indexes', COALESCE((SELECT json_agg("index") FROM indexes), '[]'::json) + ) as "publishedSchema" +$$ LANGUAGE sql; + +CREATE TABLE IF NOT EXISTS "demo_0"."publishedSchema" ( + current JSON, + exists BOOL PRIMARY KEY DEFAULT true CHECK (exists) +); + +INSERT INTO "demo_0"."publishedSchema" (current) VALUES ("demo_0".schema_specs()) + ON CONFLICT (exists) DO + UPDATE SET current = excluded.current; + +CREATE OR REPLACE FUNCTION "demo_0".update_schemas(event_type text, tag text, target record) +RETURNS void AS $$ +DECLARE + prev_schema_specs JSON; + schema_specs JSON; + message TEXT; +BEGIN + SELECT current FROM "demo_0"."publishedSchema" INTO prev_schema_specs; + SELECT "demo_0".schema_specs() INTO schema_specs; + + IF prev_schema_specs::text != schema_specs::text THEN + UPDATE "demo_0"."publishedSchema" SET current = schema_specs; + ELSIF event_type = 'ddlStart' THEN + prev_schema_specs = NULL; + ELSIF event_type = 'ddlUpdate' THEN + PERFORM "demo_0".notice_ignore('noop', tag, target); + RETURN; + END IF; + + SELECT json_build_object( + 'type', event_type, + 'version', 1, + 'previousSchema', prev_schema_specs, + 'schema', schema_specs, + 'event', json_build_object('tag', tag), + 'context', "demo_0".get_trigger_context() + ) INTO message; + + PERFORM pg_logical_emit_message(true, 'demo/0/ddl', message); + + RAISE NOTICE 'Emitted demo_0 % for % %', event_type, tag, + COALESCE(row_to_json(target)::text, ''); +END +$$ LANGUAGE plpgsql; + +CREATE OR REPLACE FUNCTION "demo_0".update_schemas() +RETURNS void AS $$ +BEGIN + PERFORM "demo_0".update_schemas('schemaSnapshot', 'MANUAL', NULL); +END +$$ LANGUAGE plpgsql; + +CREATE OR REPLACE FUNCTION "demo_0".emit_ddl_start() +RETURNS event_trigger AS $$ +DECLARE + schema_specs JSON; + message TEXT; +BEGIN + PERFORM pg_advisory_xact_lock(4353719051050729648); + PERFORM "demo_0".update_schemas('ddlStart', TG_TAG, NULL); +END +$$ LANGUAGE plpgsql; + +CREATE OR REPLACE FUNCTION "demo_0".emit_ddl_end() +RETURNS event_trigger AS $$ +DECLARE + publications TEXT[]; + target RECORD; + relevant RECORD; + schema_specs JSON; + message TEXT; + event TEXT; +BEGIN + publications := ARRAY['demo_pub']; + + SELECT objid, object_type, object_identity + FROM pg_event_trigger_ddl_commands() + LIMIT 1 INTO target; + + SELECT true INTO relevant; + + IF (target.object_type = 'table' AND TG_TAG != 'ALTER TABLE') + OR target.object_type = 'table column' THEN + SELECT ns.nspname AS "schema", c.relname AS "name" FROM pg_class AS c + JOIN pg_namespace AS ns ON c.relnamespace = ns.oid + JOIN pg_publication_tables AS pb ON pb.schemaname = ns.nspname AND pb.tablename = c.relname + WHERE c.oid = target.objid AND pb.pubname = ANY (publications) + INTO relevant; + + ELSIF target.object_type = 'index' THEN + SELECT ns.nspname AS "schema", c.relname AS "name" FROM pg_class AS c + JOIN pg_namespace AS ns ON c.relnamespace = ns.oid + JOIN pg_indexes as ind ON ind.schemaname = ns.nspname AND ind.indexname = c.relname + JOIN pg_publication_tables AS pb ON pb.schemaname = ns.nspname AND pb.tablename = ind.tablename + WHERE c.oid = target.objid AND pb.pubname = ANY (publications) + INTO relevant; + + ELSIF target.object_type = 'publication relation' THEN + SELECT pb.pubname FROM pg_publication_rel AS rel + JOIN pg_publication AS pb ON pb.oid = rel.prpubid + WHERE rel.oid = target.objid AND pb.pubname = ANY (publications) + INTO relevant; + + ELSIF target.object_type = 'publication namespace' THEN + SELECT pb.pubname FROM pg_publication_namespace AS ns + JOIN pg_publication AS pb ON pb.oid = ns.pnpubid + WHERE ns.oid = target.objid AND pb.pubname = ANY (publications) + INTO relevant; + + ELSIF target.object_type = 'schema' THEN + SELECT ns.nspname AS "schema", c.relname AS "name" FROM pg_class AS c + JOIN pg_namespace AS ns ON c.relnamespace = ns.oid + JOIN pg_publication_tables AS pb ON pb.schemaname = ns.nspname AND pb.tablename = c.relname + WHERE ns.oid = target.objid AND pb.pubname = ANY (publications) + INTO relevant; + + ELSIF target.object_type = 'publication' THEN + SELECT 1 WHERE target.object_identity = ANY (publications) + INTO relevant; + + ELSIF TG_TAG LIKE 'CREATE %' AND target.object_type IS NULL THEN + relevant := NULL; + END IF; + + IF relevant IS NULL THEN + PERFORM "demo_0".notice_ignore('irrelevant', TG_TAG, target); + RETURN; + END IF; + + IF TG_TAG = 'COMMENT' THEN + IF target.object_type != 'publication' THEN + PERFORM "demo_0".notice_ignore('irrelevant', TG_TAG, target); + RETURN; + END IF; + PERFORM "demo_0".update_schemas('schemaSnapshot', TG_TAG, target); + ELSE + PERFORM "demo_0".update_schemas('ddlUpdate', TG_TAG, target); + END IF; + +END +$$ LANGUAGE plpgsql; + +DROP EVENT TRIGGER IF EXISTS "demo_ddl_start_0"; +DROP EVENT TRIGGER IF EXISTS "demo_ddl_end_0"; + +CREATE EVENT TRIGGER "demo_ddl_start_0" + ON ddl_command_start + WHEN TAG IN ('CREATE TABLE', 'ALTER TABLE', 'CREATE INDEX', 'DROP TABLE', 'DROP INDEX', 'ALTER PUBLICATION', 'ALTER SCHEMA') + EXECUTE PROCEDURE "demo_0".emit_ddl_start(); + +CREATE EVENT TRIGGER "demo_ddl_end_0" + ON ddl_command_end + WHEN TAG IN ('CREATE TABLE', 'ALTER TABLE', 'CREATE INDEX', 'DROP TABLE', 'DROP INDEX', 'ALTER PUBLICATION', 'ALTER SCHEMA', 'COMMENT') + EXECUTE PROCEDURE "demo_0".emit_ddl_end(); diff --git a/demo/docker-compose.yml b/demo/docker-compose.yml new file mode 100644 index 0000000..a66c930 --- /dev/null +++ b/demo/docker-compose.yml @@ -0,0 +1,105 @@ +# The whole thing on one machine: Postgres, the engine, the ticket desk. +# +# docker compose -f demo/docker-compose.yml up --build +# +# Then http://localhost:8080 for the stats page, and demo/README.md for +# the curl walk-through. Every port is published on this machine only, +# and the engine's S2S endpoint asks for a token: the stack holds a +# well-known database password and is for a laptop, not a network. +services: + postgres: + image: postgres:16 + command: + - postgres + - -c + - wal_level=logical + - -c + - max_replication_slots=16 + - -c + - max_wal_senders=16 + # A slot left by an engine that crashed (rather than stopped) + # holds WAL until it is dropped; this caps what it can hold, as + # docs/deploy.md advises, so a crash cannot fill the disk. + - -c + - max_slot_wal_keep_size=1GB + environment: + POSTGRES_USER: streamgres + POSTGRES_PASSWORD: streamgres + POSTGRES_DB: tickets + volumes: + # The publication the engine streams, then the schema-change + # trigger stack it requires, both for app `demo`, shard 0. + - ./postgres-init.sql:/docker-entrypoint-initdb.d/10-publication.sql:ro + - ./ddl-triggers.sql:/docker-entrypoint-initdb.d/20-ddl-triggers.sql:ro + ports: + # Host ports are overridable, for a machine that already has one + # of these services on the default port. + - "127.0.0.1:${POSTGRES_PORT:-5432}:5432" + healthcheck: + test: ["CMD-SHELL", "pg_isready -U streamgres -d tickets"] + interval: 2s + timeout: 2s + retries: 30 + + # Creates the schema and loads the dataset, then exits. Runs before + # the engine so the tables are there when it takes its snapshot. + seed: + build: + context: .. + dockerfile: docker/demo/Dockerfile + command: ["seed"] + environment: + DATABASE_URL: postgres://streamgres:streamgres@postgres:5432/tickets + DEMO_SCALE: "${DEMO_SCALE:-1}" + depends_on: + postgres: + condition: service_healthy + + engine: + build: + context: .. + dockerfile: docker/server/Dockerfile + environment: + STREAMGRES_PG_DSN: postgresql://streamgres:streamgres@postgres:5432/tickets + # The app the trigger stack and the publication are made for. + STREAMGRES_APP_ID: demo + STREAMGRES_PUBLICATION: demo_pub + STREAMGRES_S2S_TOKEN: "${STREAMGRES_S2S_TOKEN:-demo-token}" + # The Zero side is not part of this demo; the endpoints are never + # called, but the server wants them set. + STREAMGRES_QUERY_URL: http://127.0.0.1:1/query + STREAMGRES_MUTATE_URL: http://127.0.0.1:1/mutate + STREAMGRES_ADDR: 0.0.0.0:4848 + ports: + - "127.0.0.1:${ENGINE_PORT:-4848}:4848" + # A migration the engine cannot follow (`demo migrate drop`) + # stops it on purpose; restarted, it serves the schema as it is. + restart: on-failure + depends_on: + seed: + condition: service_completed_successfully + healthcheck: + test: ["CMD", "curl", "-fsS", "http://127.0.0.1:4848/health"] + interval: 2s + timeout: 2s + retries: 60 + + demo: + build: + context: .. + dockerfile: docker/demo/Dockerfile + command: ["serve"] + environment: + DATABASE_URL: postgres://streamgres:streamgres@postgres:5432/tickets + STREAMGRES_S2S_URL: ws://engine:4848/zero/s2s/v1 + STREAMGRES_S2S_TOKEN: "${STREAMGRES_S2S_TOKEN:-demo-token}" + DEMO_ADDR: 0.0.0.0:8080 + DEMO_WRITER_MS: "${DEMO_WRITER_MS:-1000}" + DEMO_WRITER_MIX: "${DEMO_WRITER_MIX:-}" + DEMO_VERIFY_MS: "${DEMO_VERIFY_MS:-0}" + DEMO_SCALE: "${DEMO_SCALE:-1}" + ports: + - "127.0.0.1:${DEMO_PORT:-8080}:8080" + depends_on: + engine: + condition: service_healthy diff --git a/demo/postgres-init.sql b/demo/postgres-init.sql new file mode 100644 index 0000000..3dace18 --- /dev/null +++ b/demo/postgres-init.sql @@ -0,0 +1,3 @@ +-- The publication the engine's change feed streams. FOR ALL TABLES +-- needs a superuser, which the compose file's POSTGRES_USER is. +CREATE PUBLICATION demo_pub FOR ALL TABLES; diff --git a/docker/demo/Dockerfile b/docker/demo/Dockerfile new file mode 100644 index 0000000..8d47f49 --- /dev/null +++ b/docker/demo/Dockerfile @@ -0,0 +1,24 @@ +# The ticket-desk demo as an image: `demo serve`, `demo seed` +# and `demo bench` in one binary, built from the locked workspace. +# +# docker build -f docker/demo/Dockerfile -t streamgres-demo . +ARG BASE_REGISTRY=docker.io/library +FROM ${BASE_REGISTRY}/rust:1.89-bookworm AS builder +WORKDIR /build +COPY Cargo.toml Cargo.lock clippy.toml ./ +COPY src ./src +COPY crates ./crates +RUN cargo build --release --locked -p streamgres-demo + +FROM ${BASE_REGISTRY}/debian:bookworm-slim AS runtime +RUN apt-get update \ + && apt-get install -y --no-install-recommends ca-certificates curl libpq5 \ + && rm -rf /var/lib/apt/lists/* \ + && useradd --system --uid 10002 --home-dir /app --shell /usr/sbin/nologin demo \ + && mkdir -p /app && chown demo:demo /app +WORKDIR /app +COPY --from=builder /build/target/release/demo /app/demo +USER demo +EXPOSE 8080 +ENTRYPOINT ["/app/demo"] +CMD ["serve"] diff --git a/docker/server/Dockerfile b/docker/server/Dockerfile index 4455e87..dc107c4 100644 --- a/docker/server/Dockerfile +++ b/docker/server/Dockerfile @@ -9,6 +9,9 @@ FROM ${BASE_REGISTRY}/rust:1.89-bookworm AS builder WORKDIR /build COPY Cargo.toml Cargo.lock clippy.toml ./ COPY src ./src +# The workspace's other crates (the Diesel wrapper) are not part of the +# server, but the locked workspace must resolve to build it. +COPY crates ./crates RUN cargo build --release --locked --bin server FROM ${BASE_REGISTRY}/debian:bookworm-slim AS runtime diff --git a/docs/deploy.md b/docs/deploy.md index 8ac93e9..52343f1 100644 --- a/docs/deploy.md +++ b/docs/deploy.md @@ -280,6 +280,12 @@ depend on what the defaults happen to be: | `STREAMGRES_GROUP_LOG_BYTES` | `262144` | recent updates kept per client group, so a briefly disconnected tab catches up instead of starting over. `0` keeps none | | `STREAMGRES_MAX_MESSAGE_BYTES` | `16777216` | largest message accepted from a client | | `STREAMGRES_BYTES_PER_PART` | `768000` | bytes of serialized row changes per message part (750 KiB) | +| `STREAMGRES_S2S_TOKEN` | none | secret a backend must send as `Authorization: Bearer ` to open the S2S endpoint. Unset, the endpoint takes any connection | +| `STREAMGRES_S2S_MAX_CONNECTIONS` | `64` | S2S connections open at once; past it an upgrade is answered `503`. Each backend process opens `STREAMGRES_S2S_CONNECTIONS` (2), so 64 serves 32 processes. `0` lifts the limit | +| `STREAMGRES_S2S_CONNECTION_SUBSCRIPTIONS` | `10000` | subscriptions one S2S connection may hold; past it a subscribe is refused. The engine announces it, and the client stops short of it. `0` lifts the limit | +| `STREAMGRES_HEARTBEAT_MS` | `500` | how often the engine writes a heartbeat into the WAL (`pg_logical_emit_message`, not transactional) and measures how long its stream takes to bring it back. S2S clients are told that lag and read Postgres past their `STREAMGRES_S2S_MAX_LAG_MS`. Needs `STREAMGRES_PG_DSN` to reach a primary; `0` turns it off | +| `STREAMGRES_S2S_MAX_DISJUNCTS` | `1024` | alternatives an S2S query's filters may expand to (an `AND` of `OR`s multiplies them); past it a subscribe is refused and the backend reads Postgres. The engine expands a filter on the thread every subscriber's changes go through: an `AND` of thirty `OR` pairs would hold them all up and exhaust its memory. `0` lifts the limit | +| `STREAMGRES_S2S_MAX_LIST_VALUES` | `10000` | `IN` and `NOT IN` values an S2S query may hold in all; past it a subscribe is refused. The engine files each value on that same thread, about a second per 150,000. `0` lifts the limit | | `STREAMGRES_LOG` | `info` | `error`, `warn`, `info` or `debug`. `debug` costs throughput | | `STREAMGRES_SLOW_QUERY_MS` | `1000` | queries and pushes slower than this are logged at `warn` | | `STREAMGRES_METRICS_INTERVAL_MS` | `10000` | how often process metrics are sampled | @@ -290,9 +296,27 @@ depend on what the defaults happen to be: - **One port, 4848.** It serves: - the WebSocket at `/sync/v51/connect` + - the server-to-server WebSocket at `/s2s/v1`, for backends + that cache their own reads - `/health`, `/metrics` and `/stats` - Keep `/metrics` and `/stats` private. + Keep `/metrics` and `/stats` private. Do the same for `/s2s/v1`, and + set `STREAMGRES_S2S_TOKEN`. The endpoint carries no user: anyone who + can open it can subscribe to the rows of every table the engine + follows. The server speaks plain `ws://`, so the token crosses the + network as text: terminate TLS in front of it (an ingress, a mesh, + a sidecar) unless the backends share a private network with it. The + two S2S limits multiply: at most `STREAMGRES_S2S_MAX_CONNECTIONS` × + `STREAMGRES_S2S_CONNECTION_SUBSCRIPTIONS` subscriptions in all. Each + S2S connection queues at most 4,096 frames and 64 MiB of them for a + backend that is slow to read; past that it is closed, and its socket + cut off 5 seconds later if it still cannot be flushed. +- **Times.** The engine keeps every time in microseconds, Postgres's own + precision, and the S2S endpoint ships them so. Zero clients see a time + as its whole milliseconds, and a Zero query's filter or page cursor on + a time compares on the millisecond, as the client does. Rows of one + page that fall in the same millisecond are ordered by their + microseconds, then by key, where the client orders them by key alone. - **Health.** `GET /health` (also `/healthz`, `/readyz`, and each of them under the base path) returns `503` until the server is ready. That means the change feed has passed the first read snapshot, and the warm start, if diff --git a/src/bin/server.rs b/src/bin/server.rs index be1f235..faeed2d 100644 --- a/src/bin/server.rs +++ b/src/bin/server.rs @@ -20,6 +20,8 @@ struct MallocConf(*const std::ffi::c_char); unsafe impl Sync for MallocConf {} fn main() { + #[cfg(feature = "console")] + console_subscriber::init(); load_dotenv(".env"); raise_file_limit(); let outcome = streamgres::client::Config::from_env().and_then(streamgres::client::serve); diff --git a/src/client/ast.rs b/src/client/ast.rs index 515a5d4..865c627 100644 --- a/src/client/ast.rs +++ b/src/client/ast.rs @@ -553,25 +553,18 @@ fn simple(op: &str, left: &Operand, right: &Operand, table: &DbTable) -> Result< ">=" => GTE, "IN" => IN, "NOT IN" => NOT_IN, - "IS" => { - return Ok(if literal.is_null() { - Where::is_null(name.as_str()) - } else { - Where::condition(name.as_str(), EQ, literal_value(literal, &column.r#type)?) - }); - } - "IS NOT" => { - return Ok(if literal.is_null() { - Where::is_not_null(name.as_str()) - } else { - Where::condition(name.as_str(), NEQ, literal_value(literal, &column.r#type)?) - }); - } + "IS" if literal.is_null() => return Ok(Where::is_null(name.as_str())), + "IS NOT" if literal.is_null() => return Ok(Where::is_not_null(name.as_str())), + "IS" => EQ, + "IS NOT" => NEQ, "LIKE" | "ILIKE" | "NOT LIKE" | "NOT ILIKE" => { return Err(format!("{op} is not supported (a condition on `{name}`)")); } other => return Err(format!("unknown operator `{other}`")), }; + if column.r#type == ValueType::Timestamp && !literal.is_null() { + return time_condition(name.as_str(), operator, literal); + } let value = if matches!( operator, ComparisonOperator::IN | ComparisonOperator::NOT_IN @@ -591,10 +584,78 @@ fn simple(op: &str, left: &Operand, right: &Operand, table: &DbTable) -> Result< Ok(Where::condition(name.as_str(), operator, value)) } +/// `column op literal` on a time column, as Zero means it. The engine +/// keeps a time in microseconds and Zero sees its whole milliseconds, +/// so the comparison is on the millisecond: `= m` holds for every +/// microsecond of `m`, `> m` from the next millisecond on, `<= m` up to +/// the end of `m`, and a fractional literal compares as a whole number +/// of milliseconds would against it. +fn time_condition(column: &str, op: ComparisonOperator, literal: &Json) -> Result { + let millis_of = |literal: &Json| -> Result, String> { + match literal { + Json::Null => Ok(None), + Json::Number(number) => Ok(number.as_f64()), + Json::String(text) => text + .parse::() + .map(Some) + .map_err(|_| format!("`{text}` is not a number")), + other => Err(format!("{other} is not a time")), + } + }; + let micros = |millis: f64| Value::Int((millis as i64).saturating_mul(1000)); + let at_least = |millis: f64| Where::condition(column, GTE, micros(millis)); + let below = |millis: f64| Where::condition(column, LT, micros(millis)); + let within = |millis: f64| -> Where { + if millis.fract() == 0.0 { + Where::AND(vec![at_least(millis), below(millis + 1.0)]) + } else { + Where::OR(Vec::new()) + } + }; + let outside = |millis: f64| -> Where { + if millis.fract() == 0.0 { + Where::OR(vec![below(millis), at_least(millis + 1.0)]) + } else { + Where::is_not_null(column) + } + }; + if matches!(op, IN | NOT_IN) { + let Json::Array(items) = literal else { + return Err(format!("{op:?} on `{column}` needs a list")); + }; + let mut members = Vec::with_capacity(items.len()); + for item in items { + members.push(millis_of(item)?); + } + if op == NOT_IN && members.iter().any(Option::is_none) { + return Ok(Where::OR(Vec::new())); + } + let members = members.into_iter().flatten(); + return Ok(if op == IN { + Where::OR(members.map(within).collect()) + } else { + Where::AND(members.map(outside).chain([Where::is_not_null(column)]).collect()) + }); + } + let Some(millis) = millis_of(literal)? else { + return Ok(Where::OR(Vec::new())); + }; + Ok(match op { + EQ => within(millis), + NEQ => outside(millis), + GT => at_least(millis.floor() + 1.0), + GTE => at_least(millis.ceil()), + LT => below(millis.ceil()), + LTE => below(millis.floor() + 1.0), + other => return Err(format!("{other:?} does not compare times")), + }) +} + /// A JSON literal as the value a column of type `declared` compares /// against. A JSON column's cells are the text `jsonb` writes, so a literal /// against one, whatever its kind, is that text: the string `x` is `"x"`, -/// a number its digits, an object its canonical form. +/// a number its digits, an object its canonical form. A time literal is +/// the protocol's milliseconds, kept as microseconds. fn literal_value(literal: &Json, declared: &ValueType) -> Result { Ok(match (literal, declared) { (Json::Null, _) => Value::Null, @@ -607,6 +668,10 @@ fn literal_value(literal: &Json, declared: &ValueType) -> Result Value::Float(number.as_f64().unwrap_or(f64::NAN)) } (Json::Number(number), ValueType::String) => Value::String(number.to_string()), + (Json::Number(number), ValueType::Timestamp) => match number.as_i64() { + Some(millis) => Value::Int(millis.saturating_mul(1000)), + None => Value::Float(number.as_f64().unwrap_or(f64::NAN) * 1000.0), + }, (Json::Number(number), _) => match number.as_i64() { Some(int) => Value::Int(int), None => Value::Float(number.as_f64().unwrap_or(f64::NAN)), @@ -614,7 +679,11 @@ fn literal_value(literal: &Json, declared: &ValueType) -> Result (Json::String(text), _) if text.contains('\0') => { return Err("a text value contains a NUL byte".to_owned()); } - (Json::String(text), ValueType::Int | ValueType::Timestamp) => text + (Json::String(text), ValueType::Timestamp) => text + .parse::() + .map(|millis| Value::Int(millis.saturating_mul(1000))) + .map_err(|_| format!("`{text}` is not a number"))?, + (Json::String(text), ValueType::Int) => text .parse::() .map(Value::Int) .map_err(|_| format!("`{text}` is not a number"))?, @@ -644,6 +713,8 @@ fn literal_value(literal: &Json, declared: &ValueType) -> Result /// columns, the appended primary key in particular, and is then compared /// on the prefix it has), earlier columns tied and the column itself /// strictly past the cursor, plus the cursor row itself when inclusive. +/// A time column is compared on its milliseconds, as Zero sees it +/// ([`time_condition`]). fn keyset(bound: &Bound, order: &[OrderBy], table: &DbTable) -> Result { let mut keys = Vec::with_capacity(order.len()); for clause in order { @@ -655,31 +726,35 @@ fn keyset(bound: &Bound, order: &[OrderBy], table: &DbTable) -> Result GT, Order::DESC => LT, }; - let mut conjuncts: Vec = keys[..index] + let mut conjuncts = keys[..index] .iter() - .map(|(earlier, _, earlier_value)| { - Where::condition(*earlier, EQ, earlier_value.clone()) - }) - .collect(); - conjuncts.push(Where::condition(*column, op, value.clone())); + .map(|earlier| earlier.compared(EQ)) + .collect::, String>>()?; + conjuncts.push(key.compared(op)?); branches.push(collapse(conjuncts)); } if !bound.exclusive { branches.push(collapse( keys.iter() - .map(|(column, _, value)| Where::condition(*column, EQ, value.clone())) - .collect(), + .map(|key| key.compared(EQ)) + .collect::, String>>()?, )); } Ok(if branches.len() == 1 { @@ -689,6 +764,28 @@ fn keyset(bound: &Bound, order: &[OrderBy], table: &DbTable) -> Result { + column: &'a str, + direction: Order, + literal: &'a Json, + value: Value, + time: bool, +} + +impl CursorKey<'_> { + /// The column compared with the cursor's value by `op`; a time on + /// its milliseconds. + fn compared(&self, op: ComparisonOperator) -> Result { + if self.time { + time_condition(self.column, op, self.literal) + } else { + Ok(Where::condition(self.column, op, self.value.clone())) + } + } +} + /// One part is itself; several are an `AND`; none is `TRUE`. fn collapse(mut parts: Vec) -> Where { if parts.len() == 1 { @@ -967,13 +1064,51 @@ mod tests { /// A cursor that stops short of the appended primary key is compared /// on the columns it has; one that names no sort column is refused. + /// Zero sees a time as its whole milliseconds: a filter or a cursor + /// on a time holds for every microsecond of the millisecond it + /// names, so rows a few hundred microseconds apart compare equal, as + /// they do on the client, and a page's cursor loses none of them. + #[test] + fn a_time_compares_on_its_millisecond() { + use crate::ivm::evaluate; + use crate::model::RowData; + let message = |id: &str, micros: i64| { + RowData::from(std::collections::HashMap::from([ + (crate::model::ColumnName::from("messageId"), Value::from(id)), + (crate::model::ColumnName::from("createdAt"), Value::Int(micros)), + ])) + }; + let (early, late, before, after) = ( + message("a", 7_000_300), + message("b", 7_000_700), + message("c", 6_999_999), + message("d", 7_001_000), + ); + let filter_of = |json: &str| { + let ast: Ast = serde_json::from_str(json).unwrap(); + translate(&ast, &catalog()).unwrap().query.main_table.filter + }; + let holds = |filter: &Where, row: &RowData| evaluate(filter, row, &mut 0); + let equal = filter_of(r#"{"table": "messages", "where": {"type": "simple", "op": "=", "left": {"type": "column", "name": "createdAt"}, "right": {"type": "literal", "value": 7000}}}"#); + assert!(holds(&equal, &early) && holds(&equal, &late)); + assert!(!holds(&equal, &before) && !holds(&equal, &after)); + let at_most = filter_of(r#"{"table": "messages", "where": {"type": "simple", "op": "<=", "left": {"type": "column", "name": "createdAt"}, "right": {"type": "literal", "value": 7000}}}"#); + assert!(holds(&at_most, &late) && !holds(&at_most, &after)); + let next_page = filter_of(r#"{"table": "messages", "orderBy": [["createdAt", "desc"], ["messageId", "desc"]], "start": {"row": {"createdAt": 7000, "messageId": "b"}, "exclusive": true}, "limit": 5}"#); + assert!(holds(&next_page, &early), "the same millisecond, a lower key"); + assert!(holds(&next_page, &before)); + assert!(!holds(&next_page, &late) && !holds(&next_page, &after)); + } + + /// A cursor compares on the prefix of the order it names, its + /// milliseconds kept as microseconds. #[test] fn cursors_compare_on_their_prefix() { let ast: Ast = serde_json::from_str(r#"{"table": "messages", "orderBy": [["createdAt", "desc"]], "start": {"row": {"createdAt": 1000}, "exclusive": true}, "limit": 5}"#).unwrap(); let translated = translate(&ast, &catalog()).unwrap(); assert_eq!( translated.query.main_table.filter, - Where::condition("createdAt", LT, Value::Int(1000)) + Where::condition("createdAt", LT, Value::Int(1_000_000)) ); let empty: Ast = serde_json::from_str(r#"{"table": "messages", "orderBy": [["createdAt", "desc"]], "start": {"row": {"other": 1}, "exclusive": true}}"#).unwrap(); assert!( diff --git a/src/client/config.rs b/src/client/config.rs index e0099aa..2be07cd 100644 --- a/src/client/config.rs +++ b/src/client/config.rs @@ -13,6 +13,17 @@ use crate::log::{Format, Level}; use crate::model::TableName; use crate::sync::pg::{Keepalive, Settings, threads}; +/// A setting that must not reach a log: its `Debug` shows no value. +#[derive(Clone)] +pub struct Secret(pub String); + +impl std::fmt::Debug for Secret { + /// `Secret(..)`, never the value. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("Secret(..)") + } +} + /// Everything the server reads from the environment. #[derive(Debug, Clone)] pub struct Config { @@ -51,6 +62,11 @@ pub struct Config { /// `STREAMGRES_DDL_PREFIX`: the prefix of that trigger's logical /// messages (`//ddl`, so `zero/0/ddl`). pub ddl_prefix: String, + /// `STREAMGRES_HEARTBEAT_MS`: how often the engine writes a heartbeat + /// into the WAL to measure, end to end, how far behind it is (500; 0 + /// sends none, and the S2S endpoint then reports no lag). Any role + /// may emit one; it needs the primary. + pub heartbeat: Duration, /// `STREAMGRES_SNAPSHOT_ROTATION_MS`: how often a fresh exported /// snapshot (a temporary replication slot) is minted for reads (1000). pub snapshot_rotation: Duration, @@ -132,6 +148,31 @@ pub struct Config { /// `STREAMGRES_MAX_MESSAGE_BYTES`: the largest inbound WebSocket message /// accepted (16 MiB). pub max_message_bytes: usize, + /// `STREAMGRES_S2S_TOKEN`: the secret a backend must present as + /// `Authorization: Bearer ` to open the S2S endpoint. Unset or + /// empty, the endpoint takes any connection, which only a trusted + /// network should be given. + pub s2s_token: Option, + /// `STREAMGRES_S2S_MAX_CONNECTIONS`: how many S2S connections may be + /// open at once; past it an upgrade is answered 503 (64; 0 lifts the + /// limit). A backend process opens `STREAMGRES_S2S_CONNECTIONS` (2). + pub s2s_max_connections: usize, + /// `STREAMGRES_S2S_CONNECTION_SUBSCRIPTIONS`: how many subscriptions + /// one S2S connection may hold, those still registering included; past + /// it a subscribe is refused (10000; 0 lifts the limit). + pub s2s_connection_subscriptions: usize, + /// `STREAMGRES_S2S_MAX_DISJUNCTS`: how many disjuncts an S2S query's + /// filters may normalize to, counted without expanding them; past it + /// a subscribe is refused (1024; 0 lifts the limit). The engine + /// expands a filter into its disjuncts on its own thread, so an `AND` + /// of thirty `OR` pairs would otherwise stall every subscriber and + /// exhaust memory. + pub s2s_max_disjuncts: u64, + /// `STREAMGRES_S2S_MAX_LIST_VALUES`: how many `IN` and `NOT IN` values + /// an S2S query may hold in all; past it a subscribe is refused + /// (10000; 0 lifts the limit). The engine files each value in its + /// routing index on its own thread, about a second per 150,000. + pub s2s_max_list_values: usize, /// `STREAMGRES_BYTES_PER_PART`: how many bytes of serialized row /// operations one poke part carries at most (768000, i.e. 750 KiB). pub bytes_per_part: usize, @@ -321,6 +362,7 @@ impl Config { schemas, slot: threads::slot_name(), slot_cleanup_age: millis("STREAMGRES_SLOT_CLEANUP_AGE_MS", 0)?, + heartbeat: millis("STREAMGRES_HEARTBEAT_MS", 500)?, publication: first(&["STREAMGRES_PUBLICATION"]) .unwrap_or_else(|| "xyne_sync_pub".to_owned()), ddl_trigger, @@ -389,6 +431,34 @@ impl Config { })?, None => 16 * 1024 * 1024, }, + s2s_token: first(&["STREAMGRES_S2S_TOKEN"]).map(Secret), + s2s_max_connections: match first(&["STREAMGRES_S2S_MAX_CONNECTIONS"]) { + Some(text) => text.parse().map_err(|_| { + format!("STREAMGRES_S2S_MAX_CONNECTIONS must be a number, got `{text}`") + })?, + None => 64, + }, + s2s_connection_subscriptions: match first(&["STREAMGRES_S2S_CONNECTION_SUBSCRIPTIONS"]) + { + Some(text) => text.parse().map_err(|_| { + format!( + "STREAMGRES_S2S_CONNECTION_SUBSCRIPTIONS must be a number, got `{text}`" + ) + })?, + None => 10_000, + }, + s2s_max_disjuncts: match first(&["STREAMGRES_S2S_MAX_DISJUNCTS"]) { + Some(text) => text.parse().map_err(|_| { + format!("STREAMGRES_S2S_MAX_DISJUNCTS must be a number, got `{text}`") + })?, + None => 1024, + }, + s2s_max_list_values: match first(&["STREAMGRES_S2S_MAX_LIST_VALUES"]) { + Some(text) => text.parse().map_err(|_| { + format!("STREAMGRES_S2S_MAX_LIST_VALUES must be a number, got `{text}`") + })?, + None => 10_000, + }, bytes_per_part: match first(&["STREAMGRES_BYTES_PER_PART"]) { Some(text) => text.parse().map_err(|_| { format!("STREAMGRES_BYTES_PER_PART must be a number of bytes, got `{text}`") @@ -497,6 +567,7 @@ impl Config { ], ddl_trigger: self.ddl_trigger.clone(), ddl_prefix: self.ddl_prefix.clone(), + heartbeat: self.heartbeat, } } } @@ -627,6 +698,63 @@ mod tests { assert_eq!(config.engine_settings().ddl_trigger, "my_trigger"); } + /// The S2S endpoint takes any connection unless a token is set, and + /// the token never shows in the configuration's `Debug`. + #[test] + fn the_s2s_token_is_optional_and_kept_out_of_logs() { + assert!(config_with(&[]).s2s_token.is_none()); + assert!( + config_with(&[("STREAMGRES_S2S_TOKEN", "")]) + .s2s_token + .is_none() + ); + let config = config_with(&[("STREAMGRES_S2S_TOKEN", "hunter2")]); + assert_eq!( + config.s2s_token.as_ref().map(|token| token.0.as_str()), + Some("hunter2") + ); + assert!(!format!("{config:?}").contains("hunter2")); + } + + /// 64 S2S connections may be open unless told otherwise. + #[test] + fn the_s2s_connection_count_has_a_default() { + assert_eq!(config_with(&[]).s2s_max_connections, 64); + let limited = config_with(&[("STREAMGRES_S2S_MAX_CONNECTIONS", "0")]); + assert_eq!(limited.s2s_max_connections, 0); + assert!(config(&[("STREAMGRES_S2S_MAX_CONNECTIONS", "many")]).is_err()); + } + + /// One S2S connection holds 10000 subscriptions unless told otherwise. + #[test] + fn the_s2s_connection_limit_has_a_default() { + assert_eq!(config_with(&[]).s2s_connection_subscriptions, 10_000); + let limited = config_with(&[("STREAMGRES_S2S_CONNECTION_SUBSCRIPTIONS", "3")]); + assert_eq!(limited.s2s_connection_subscriptions, 3); + assert!(config(&[("STREAMGRES_S2S_CONNECTION_SUBSCRIPTIONS", "x")]).is_err()); + } + + /// An S2S query's size is bounded unless told otherwise: 1024 + /// disjuncts and 10000 list values, 0 lifting either. + #[test] + fn the_s2s_query_size_limits_have_defaults() { + let defaults = config_with(&[]); + assert_eq!((defaults.s2s_max_disjuncts, defaults.s2s_max_list_values), (1024, 10_000)); + let set = config_with(&[ + ("STREAMGRES_S2S_MAX_DISJUNCTS", "0"), + ("STREAMGRES_S2S_MAX_LIST_VALUES", "7"), + ]); + assert_eq!((set.s2s_max_disjuncts, set.s2s_max_list_values), (0, 7)); + assert!(config(&[("STREAMGRES_S2S_MAX_DISJUNCTS", "lots")]).is_err()); + } + + /// The engine beats every 500 ms unless told otherwise; 0 stops it. + #[test] + fn the_heartbeat_has_a_default() { + assert_eq!(config_with(&[]).heartbeat, Duration::from_millis(500)); + assert_eq!(config_with(&[("STREAMGRES_HEARTBEAT_MS", "0")]).heartbeat, Duration::ZERO); + } + /// [`config`], expected to parse. fn config_with(vars: &[(&str, &str)]) -> Config { config(vars).unwrap() diff --git a/src/client/connection.rs b/src/client/connection.rs index 81d10a6..c7d8460 100644 --- a/src/client/connection.rs +++ b/src/client/connection.rs @@ -77,6 +77,8 @@ pub struct AppState { pub warmed: watch::Receiver, pub transforms: Arc, pub shutdown: Shutdown, + /// The server-to-server endpoint's task. + pub s2s: super::s2s::Handle, } /// One desired-query change between the client's message and the group @@ -143,7 +145,9 @@ struct ConnectParams { pub fn router(state: Arc) -> Router { let base = state.config.base_path.clone(); let mut router = - Router::new().route(&format!("{base}/sync/v{{version}}/connect"), get(connect)); + Router::new() + .route(&format!("{base}/sync/v{{version}}/connect"), get(connect)) + .route(&format!("{base}/s2s/v1"), get(super::s2s::connect)); for probe in ["health", "healthz", "readyz"] { router = router .route(&format!("{base}/{probe}"), get(health)) diff --git a/src/client/groups.rs b/src/client/groups.rs index 2f936ab..be31afc 100644 --- a/src/client/groups.rs +++ b/src/client/groups.rs @@ -694,6 +694,7 @@ impl Groups { Event::Refused { sub, reason } => self.refused(sub, &reason), Event::Capped { sub } => self.capped(sub), Event::Heavy { sub, table, rows } => self.heavy(sub, table.as_str(), rows), + Event::Heartbeat(_) => {} Event::Hydrated(subs) => { for sub in subs { let Some((group_id, hash)) = self.by_sub.get(&sub).cloned() else { @@ -1726,7 +1727,7 @@ fn build_poke( Patch::Del(table, key) => { dels += 1; let mut bytes = Vec::with_capacity(64); - wire::write_del(&mut bytes, table.as_str(), key); + wire::write_del(&mut bytes, table.as_str(), key, catalog.table(table.as_str())); Bytes::from(bytes) } }; diff --git a/src/client/mod.rs b/src/client/mod.rs index 85ac6a8..2a81eee 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -57,6 +57,7 @@ pub mod ddl_triggers; pub mod groups; pub mod plan; pub mod protocol; +pub mod s2s; pub mod sampler; pub mod schema; pub mod transform; @@ -110,7 +111,10 @@ async fn join_worker( } /// Run the server with `config` until ctrl-c; returns the reason it could -/// not start. +/// not start. The engine feeds one consumer per group thread and one for +/// the server-to-server endpoint, the last sink. Tasks are named for +/// `tokio-console` when the runtime is built with `--cfg tokio_unstable`, +/// and spawned plainly otherwise. pub fn serve(config: Config) -> Result<(), String> { log::set_level(config.log); log::set_format(config.log_format); @@ -163,7 +167,7 @@ pub fn serve(config: Config) -> Result<(), String> { &settings, engine_catalog, reads_handle, - shards, + shards + 1, engine_stats, engine_shutdown.clone(), ) @@ -206,6 +210,10 @@ pub fn serve(config: Config) -> Result<(), String> { .map_err(|_| "the engine thread ended before it was up".to_owned())??; let (readiness, ready) = tokio::sync::watch::channel(false); + let mut events = events; + let s2s_events = events + .pop() + .ok_or("the engine started without the s2s sink".to_owned())?; let mut requests = Vec::with_capacity(shards); let mut group_threads = Vec::with_capacity(shards); for (shard, events) in events.into_iter().enumerate() { @@ -280,6 +288,24 @@ pub fn serve(config: Config) -> Result<(), String> { let warm_budget = config.warm_start; let metrics_interval = config.metrics_interval; let (warmed_tx, warmed) = tokio::sync::watch::channel(!warm.enabled() || warm_budget.is_zero()); + let (s2s, s2s_requests) = s2s::Handle::new(config.s2s_max_connections); + let s2s_task = s2s::run( + catalog.clone(), + shards, + commands.clone(), + s2s_events, + s2s_requests, + ready.clone(), + config.s2s_connection_subscriptions, + stats.clone(), + ); + #[cfg(tokio_unstable)] + tokio::task::Builder::new() + .name("s2s") + .spawn_on(s2s_task, server.handle()) + .map_err(|error| format!("spawning the s2s task: {error}"))?; + #[cfg(not(tokio_unstable))] + server.spawn(s2s_task); let state = Arc::new(connection::AppState { config, requests, @@ -294,6 +320,7 @@ pub fn serve(config: Config) -> Result<(), String> { warmed, transforms, shutdown: shutdown.clone(), + s2s, }); drop(readiness); if warm.enabled() { diff --git a/src/client/s2s.rs b/src/client/s2s.rs new file mode 100644 index 0000000..c0de5bc --- /dev/null +++ b/src/client/s2s.rs @@ -0,0 +1,1986 @@ +//! The server-to-server endpoint: a backend subscribes the SQL its ORM +//! renders and is kept current on it. +//! +//! A backend keeps its own copy of what it reads (a Diesel wrapper serves +//! repeated reads from a local cache). In rows mode the engine ships the +//! query's rows on `accepted` and each committed transaction's changes +//! to them as a `delta`, so the backend keeps the result without reading +//! Postgres. Otherwise it is sent a `touched` naming every subscription a +//! committed transaction changed the rows of, and reads Postgres again. +//! A query outside what the engine follows is `refused`, and the backend +//! serves it live. The endpoint carries no user: it serves whatever the +//! SQL names, so it belongs to trusted backends, behind its token. +//! +//! One task ([`run`]) owns every S2S subscription: it is one more +//! consumer of the engine's events beside the group threads, with its own +//! sink index, and maps the engine's subscription ids back to the +//! connection and the id the backend chose. Nothing of the Zero protocol +//! (client groups, pokes, mutations, the application server) is involved. +//! +//! # Protocol +//! +//! JSON text frames, each a two-element array of a tag and a body, on a +//! WebSocket at `{base}/s2s/v1`. +//! +//! Admission: +//! - With `STREAMGRES_S2S_TOKEN` set, the upgrade must carry +//! `Authorization: Bearer ` and is answered 401 otherwise. +//! Unset, any connection is taken, so the port belongs on a trusted +//! network. +//! - At most `STREAMGRES_S2S_MAX_CONNECTIONS` connections are open at +//! once; past that the upgrade is answered 503. +//! - A connection holds at most `STREAMGRES_S2S_CONNECTION_SUBSCRIPTIONS` +//! subscriptions, those still registering included; a subscribe past +//! that is refused. The two limits together bound what one engine +//! follows for its backends. +//! +//! Upstream: +//! - `["subscribe", {"id": n, "sql": "...", "binds": [{"oid": 25, "bytes": ""}, {"oid": 20, "bytes": null}], "rows": bool}]` +//! `sql` is what the ORM rendered, with `$1..` placeholders; each bind +//! is the parameter's Postgres type OID and its binary bytes, `null` +//! for SQL `NULL`. `id` is the backend's own, unique per connection. +//! `rows` asks for the rows themselves (below) rather than touches. +//! The SQL may join tables (`FROM ("a" INNER JOIN "b" ON (...))`, as +//! Diesel renders it); the engine follows each table and the backend +//! composes the join from the rows it is sent. +//! - `["unsubscribe", {"id": n}]` +//! - `["ping"]` +//! +//! Downstream: +//! - `["hello", {"protocol": 1, "max_subscriptions": n, "database": +//! {"system": "...", "name": "..."}}]` first; `max_subscriptions` is the +//! per-connection limit, `null` for none, so a client can stop short of +//! it rather than be refused. `database` names what the engine streams +//! (the cluster's `system_identifier` and the database's name), so a +//! backend reading another database does not take these rows for its +//! own. An engine that cannot read them closes the socket instead of +//! greeting it. +//! - `["accepted", {"id": n}]` once the engine follows the query; from +//! that moment every change to its rows is reported, so the backend +//! reads Postgres only after it. In rows mode the body also carries +//! the hydrated result: `"nodes"`, one per table of the query, root +//! first, each with `"table"`, `"key"` (its primary-key columns), +//! `"parent"` (the table it joins to, `null` for the root), +//! `"parent_column"` and `"child_column"` (the join equality), +//! `"inner"` (`false` for a LEFT join) and `"rows"` (JSON objects keyed +//! by column, every column of the row); then `"columns"` (the projected +//! `[table, column]` pairs, `null` for `*`), `"names"` (what each comes +//! out as: its alias or its own name; for a `COUNT(*)` the one output +//! name), `"count"` (a `COUNT(*)`), `"order"` (`[[column, "asc"|"desc"], +//! ...]`, on the root table) and `"limit"` (`null` for none). +//! - `["accepted", ...]` also carries `"offset"`: the engine keeps +//! `limit + offset` rows and the client skips `offset` of them. Under +//! a `COUNT(*)` the engine keeps every row, and `limit` and `offset` +//! apply to the one row the count is. A +//! `bytea` column travels as Postgres's hex text; a bind of an enum +//! type is read as its label; a time travels as microseconds since +//! the Unix epoch, Postgres's own precision (the Zero protocol sees +//! milliseconds), and a timestamp bind is read the same way. +//! - `["refused", {"id": n, "reason": "..."}]` when the engine cannot +//! follow the query (outside the SQL subset, an unknown table, too many +//! rows). The backend serves that query live. +//! - `["touched", {"ids": [n, ...]}]` per transaction that changed the +//! rows of those subscriptions, for subscriptions not in rows mode. +//! - `["delta", {"id": n, "changes": [{"table": t, "put": [row, ...], "del": [key, ...]}, ...]}]` +//! per transaction, for one subscription in rows mode: per table, the +//! rows that entered or changed and the keys of rows that left. One +//! entry per key, its last state in the transaction: a changed row is +//! a `put` of its new image, a row that came and went is a `del`, and +//! no key is in both lists. A backend that stops reading its frames +//! is closed once its queue is full rather than buffered for: more +//! than 4096 frames, or more than 64 MiB of them, waiting. +//! - `["pong"]` +//! - `["lag", {"ms": n}]` right after `hello`, then every 500 ms: how far +//! at most the engine's applied changes trail the database, the age of +//! the newest heartbeat the replication stream has brought back and +//! the engine applied (`null` before the first is sent), so a backend +//! can stop trusting its copy while the engine is behind. +//! - `["error", {"reason": "..."}]` for a frame that is not one of the +//! above (not JSON, an unknown tag, no usable `id`). A `subscribe` +//! whose `id` reads but whose body does not is answered `refused` with +//! that id instead, so no subscription is left waiting. +//! +//! An `id` is never reused on a connection: frames already queued for an +//! unsubscribed id may still arrive after the `unsubscribe`, and would +//! otherwise be taken for the new subscription's. + +use std::collections::{HashMap, HashSet}; +use std::panic::AssertUnwindSafe; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; + +use axum::extract::State; +use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade}; +use axum::http::header::AUTHORIZATION; +use axum::http::{HeaderMap, StatusCode}; +use axum::response::{IntoResponse, Response}; +use base64::Engine as _; +use base64::engine::general_purpose::STANDARD as BASE64; +use futures_util::stream::SplitStream; +use futures_util::{FutureExt, SinkExt, StreamExt}; +use serde::Serialize; +use serde_json::value::RawValue; +use serde_json::{Value as Json, json}; +use tokio::sync::mpsc::error::TrySendError; +use tokio::sync::{mpsc, oneshot, watch}; + +use super::ast::Translated; +use super::connection::AppState; +use super::{plan, wire}; +use crate::ivm::Delta; +use crate::log::{log_debug, log_info, log_warn}; +use crate::model::{ + ComparisonOperator, DataFrameKey, DataFrameOperation, DataFrameRow, MultiTableReadQuery, Order, + OrderBy, SubId, TableName, Value, Where, +}; +use crate::parser::binds::decodes; +use crate::parser::{BindValue, Projection, parse_select}; +use crate::stats::Stats; +use crate::sync::catalog::with_enums; +use crate::sync::{CatalogHandle, Command, Event}; + +/// The S2S protocol version the endpoint's `hello` announces. +pub const PROTOCOL_VERSION: u32 = 1; + +/// How many requests may wait for the task before a connection's reader +/// waits too. +const REQUEST_QUEUE: usize = 1024; + +/// How many frames may wait for a connection's writer. A backend that +/// stops reading would otherwise have every result and delta buffered +/// for it without limit; past this the connection is closed instead. +const OUTGOING_QUEUE: usize = 4096; + +/// How many bytes of frames may wait for a connection's writer: a +/// rows-mode `accepted` can be megabytes, so a count of frames alone +/// does not bound what a backend that stops reading pins. A frame is +/// always let onto an empty queue, so one large result still goes out. +const OUTGOING_BYTES: usize = 64 * 1024 * 1024; + +/// How long a closing connection's writer may take to flush what is +/// queued. A backend that stopped reading never lets it finish, and +/// waiting for it would hold the connection's place and socket forever. +const WRITER_DRAIN: Duration = Duration::from_secs(5); + +/// A connection's outgoing frames, bounded in frames and in bytes, with +/// the bytes still waiting for its writer. +#[derive(Debug, Clone)] +pub struct Outgoing { + frames: mpsc::Sender, + queued: Arc, +} + +impl Outgoing { + /// A queue and the receiving end its writer drains. + fn new() -> (Outgoing, mpsc::Receiver) { + let (frames, receiver) = mpsc::channel(OUTGOING_QUEUE); + let queued = Arc::new(AtomicUsize::new(0)); + (Outgoing { frames, queued }, receiver) + } + + /// Queues `text` for the writer. `false` when the backend is not + /// keeping up (the queue is full in frames or in bytes) and the + /// connection must be closed; a writer already gone is not that. + fn push(&self, text: String) -> bool { + let len = text.len(); + let before = self.queued.fetch_add(len, Ordering::AcqRel); + if before > 0 && before + len > OUTGOING_BYTES { + self.queued.fetch_sub(len, Ordering::AcqRel); + return false; + } + match self.frames.try_send(text) { + Ok(()) => true, + Err(TrySendError::Full(_)) => { + self.queued.fetch_sub(len, Ordering::AcqRel); + false + } + Err(TrySendError::Closed(_)) => { + self.queued.fetch_sub(len, Ordering::AcqRel); + true + } + } + } + + /// The writer's side of the byte count. It holds no sender, so the + /// queue still closes once every [`Outgoing`] is dropped. + fn drained(&self) -> Drained { + Drained(Arc::clone(&self.queued)) + } +} + +/// The bytes still waiting in a connection's [`Outgoing`], as its +/// writer releases them. +struct Drained(Arc); + +impl Drained { + /// The writer sent a frame of `len` bytes. + fn sent(&self, len: usize) { + self.0.fetch_sub(len, Ordering::AcqRel); + } +} + +/// What a connection asks of the task. +#[derive(Debug)] +pub enum Request { + Connected { + conn: u64, + out: Outgoing, + /// Fired by the task to close the connection, with the reason. + close: oneshot::Sender<&'static str>, + /// How many subscriptions the connection holds, registering ones + /// included, kept by the task for the connection's handler. + held: Arc, + }, + /// A query parsed and planned by the connection's handler. + Subscribe(Box), + Unsubscribe { + conn: u64, + id: u64, + }, + Disconnected { + conn: u64, + }, +} + +/// A subscription ready to register. +#[derive(Debug)] +pub struct Subscribe { + conn: u64, + id: u64, + query: MultiTableReadQuery, + /// Present when the backend asked for rows. + shape: Option, +} + +/// Per table, what one transaction did to each key: its new image (an +/// index into the transaction's serialized rows), or `None` once +/// deleted. One entry per key, so a row admitted and evicted in one +/// transaction is a delete, not a put the client would keep. +type Changes = Vec<(TableName, HashMap>)>; + +/// A `delta` frame's body, its rows spliced in as they were serialized. +#[derive(Serialize)] +struct DeltaBody<'a> { + id: u64, + changes: Vec>, +} + +/// One table's part of a `delta`. +#[derive(Serialize)] +struct TableChanges<'a> { + table: &'a str, + put: Vec<&'a RawValue>, + del: Vec, +} + +/// How often every connection is told the engine's lag. +const LAG_EVERY: Duration = Duration::from_millis(500); + +/// How often a refused upgrade of one kind is logged at most; the ones +/// in between are counted into the next line. +const REFUSAL_LOG_EVERY: Duration = Duration::from_secs(60); + +/// The handle connections talk to the task through. +#[derive(Debug)] +pub struct Handle { + requests: mpsc::Sender, + next_conn: AtomicU64, + /// Connections open now, against `max_connections` (0: no limit). + open: Arc, + max_connections: usize, + unauthorized: Throttle, + full: Throttle, + unidentified: Throttle, + /// The database the engine streams, as `hello` names it, once read: + /// it cannot change while the process runs. + database: tokio::sync::OnceCell, +} + +impl Handle { + /// A handle admitting at most `max_connections` at once (0: no + /// limit), and the requests its connections send the task. + pub fn new(max_connections: usize) -> (Handle, mpsc::Receiver) { + let (requests, requests_rx) = mpsc::channel(REQUEST_QUEUE); + ( + Handle { + requests, + next_conn: AtomicU64::new(1), + open: Arc::new(AtomicUsize::new(0)), + max_connections, + unauthorized: Throttle::new(REFUSAL_LOG_EVERY), + full: Throttle::new(REFUSAL_LOG_EVERY), + unidentified: Throttle::new(REFUSAL_LOG_EVERY), + database: tokio::sync::OnceCell::new(), + }, + requests_rx, + ) + } + + /// A place among the open connections, if one is free; it is given + /// back when the [`Admitted`] drops. + fn admit(&self) -> Option { + let limit = self.max_connections; + self.open + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |open| { + (limit == 0 || open < limit).then_some(open + 1) + }) + .ok()?; + Some(Admitted(Arc::clone(&self.open))) + } +} + +/// One open connection's place, held for as long as the connection is. +struct Admitted(Arc); + +impl Drop for Admitted { + /// Gives the connection's place back. + fn drop(&mut self) { + self.0.fetch_sub(1, Ordering::AcqRel); + } +} + +/// Lets a repeated warning through at most once per period. A backend +/// with the wrong token retries every second, and so can anyone who +/// reaches the port, so each refusal cannot be a log line. +#[derive(Debug)] +struct Throttle { + every: Duration, + next: Mutex>, + quiet: AtomicU64, +} + +impl Throttle { + /// A throttle logging at most once per `every`. + fn new(every: Duration) -> Throttle { + Throttle { + every, + next: Mutex::new(None), + quiet: AtomicU64::new(0), + } + } + + /// At `now`: how many were held back since the last one let through, + /// when this one may be logged; `None` when it is counted instead. + fn pass(&self, now: Instant) -> Option { + let mut next = self.next.lock().expect("throttle"); + if next.is_some_and(|next| now < next) { + self.quiet.fetch_add(1, Ordering::Relaxed); + return None; + } + *next = Some(now + self.every); + Some(self.quiet.swap(0, Ordering::Relaxed)) + } +} + +/// Logs a refused upgrade, throttled per kind. +fn refusal(throttle: &Throttle, what: &str) { + if let Some(quiet) = throttle.pass(Instant::now()) { + if quiet == 0 { + log_warn!("s2s connection refused: {what}"); + } else { + log_warn!("s2s connection refused: {what} ({quiet} more since the last warning)"); + } + } +} + +/// One backend connection as the task sees it. +struct Conn { + out: Outgoing, + /// Closes the connection; taken once used. + close: Option>, + /// The backend's ids and the engine's subscriptions they became. + subs: HashMap, + /// The backend's ids still registering, and their tokens. + registering: HashMap, + /// `subs` and `registering` together, for the connection's handler. + held: Arc, +} + +impl Conn { + /// A connection holding nothing yet. + fn new(out: Outgoing, close: oneshot::Sender<&'static str>, held: Arc) -> Conn { + Conn { + out, + close: Some(close), + subs: HashMap::new(), + registering: HashMap::new(), + held, + } + } + + /// Queues a frame for the writer. + fn send(&mut self, message: Json) { + self.send_text(message.to_string()); + } + + /// Queues frame text for the writer. A full queue means the backend + /// is not reading: the connection is closed rather than buffered for. + fn send_text(&mut self, text: String) { + if !self.out.push(text) + && let Some(close) = self.close.take() + { + let _ = close.send("the backend is not reading its frames"); + } + } + + /// How many subscriptions the connection holds, registering ones + /// included. + fn holds(&self) -> usize { + self.subs.len() + self.registering.len() + } + + /// Publishes [`Conn::holds`] to the connection's handler. + fn note_held(&self) { + self.held.store(self.holds(), Ordering::Release); + } +} + +/// What a rows-mode subscription needs beyond its ids: the query's shape +/// for the backend to compose, order and project by, and the rows +/// gathered per table while the engine hydrates. +#[derive(Debug)] +pub struct Shape { + nodes: Vec, + projection: Projection, + order: Vec, + /// The statement's own limit and offset; the tree holds `limit + + /// offset` rows and the client skips and cuts (under a `COUNT(*)`, + /// the one counted row). + limit: u32, + offset: u32, +} + +/// One table of the query: how it joins its parent, and its hydrating +/// rows (a key's latest image, `None` once deleted). +#[derive(Debug)] +struct NodeShape { + table: TableName, + parent: Option, + parent_column: Option, + child_column: Option, + inner: bool, + snapshot: HashMap>, +} + +impl Shape { + /// The shape of `query` for a backend that asked for rows, the + /// statement's own `limit` and `offset` beside it. + pub fn new( + query: &MultiTableReadQuery, + projection: Projection, + limit: u32, + offset: u32, + ) -> Shape { + let mut nodes = Vec::new(); + /// Notes `query`'s node and its joined nodes below it, depth first. + fn walk( + query: &MultiTableReadQuery, + parent: Option, + edge: Option<(&str, &str, bool)>, + nodes: &mut Vec, + ) { + nodes.push(NodeShape { + table: query.main_table.table.clone(), + parent, + parent_column: edge.map(|(p, _, _)| p.to_string()), + child_column: edge.map(|(_, c, _)| c.to_string()), + inner: edge.is_none_or(|(_, _, inner)| inner), + snapshot: HashMap::new(), + }); + let index = nodes.len() - 1; + for join in &query.joins { + walk( + &join.sub, + Some(index), + Some(( + join.main_table_column.as_str(), + join.sub_table_column.as_str(), + join.is_inner, + )), + nodes, + ); + } + } + walk(query, None, None, &mut nodes); + Shape { + nodes, + projection, + order: query.main_table.order_by.clone(), + limit, + offset, + } + } + + /// The node of `table`. + fn node_of(&mut self, table: &TableName) -> Option<&mut NodeShape> { + self.nodes.iter_mut().find(|node| &node.table == table) + } + + /// Folds a hydrating change into its node's snapshot: the key's + /// latest image, or `None` once deleted. + fn absorb(&mut self, table: &TableName, op: &DataFrameOperation) { + let Some(node) = self.node_of(table) else { + return; + }; + let (key, row) = match op { + DataFrameOperation::Add(key, row) => (key, Some(row.clone())), + DataFrameOperation::Delete(key, _) => (key, None), + }; + node.snapshot.insert(key.clone(), row); + } +} + +/// The S2S task's state: every connection, registration and +/// subscription, and the engine's side of them. +struct Task { + catalog: Arc, + sink: usize, + commands: mpsc::Sender>, + conns: HashMap, + /// Registrations in flight, by the token the command carried: the + /// connection, its id, and the shape when rows were asked for. + by_token: HashMap)>, + /// Every live subscription, back to its connection and id. + by_sub: HashMap, + /// Rows-mode subscriptions, by the engine's id. + shapes: HashMap, + /// Registered but not yet hydrated: `accepted` waits for the rows, + /// so the backend's first read lands after the engine's. + hydrating: HashSet, + /// The token the next registration is sent under. + next_token: u64, + /// How many subscriptions one connection may hold, registering ones + /// included (0: no limit). + per_connection: usize, + /// Where the engine's lag is measured ([`crate::sync::heartbeat`]). + stats: Arc, +} + +/// Runs the task: the engine's events on one side, the connections' +/// requests on the other. Returns when both channels are closed. +/// Queries are registered only once the engine's position covers the +/// storage; until then a registration would read a snapshot the feed +/// has not caught up with. +#[allow(clippy::too_many_arguments)] +pub async fn run( + catalog: Arc, + sink: usize, + commands: mpsc::Sender>, + mut events: mpsc::UnboundedReceiver, + mut requests: mpsc::Receiver, + mut ready: watch::Receiver, + per_connection: usize, + stats: Arc, +) { + while !*ready.borrow() { + if ready.changed().await.is_err() { + return; + } + } + log_info!("s2s endpoint serving"); + let mut task = Task { + catalog, + sink, + commands, + conns: HashMap::new(), + by_token: HashMap::new(), + by_sub: HashMap::new(), + shapes: HashMap::new(), + hydrating: HashSet::new(), + next_token: 1, + per_connection, + stats, + }; + let mut lag = tokio::time::interval(LAG_EVERY); + lag.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + loop { + tokio::select! { + _ = lag.tick() => task.report_lag(), + event = events.recv() => match event { + Some(event) => task.event(event).await, + None => return, + }, + request = requests.recv() => match request { + Some(request) => task.request(request).await, + None => return, + }, + } + } +} + +impl Task { + /// Tells every connection how far behind the engine is, from its + /// heartbeats: `null` while it sends none. + fn report_lag(&mut self) { + let frame = lag_frame(&self.stats).to_string(); + for connection in self.conns.values_mut() { + connection.send_text(frame.clone()); + } + } + + /// Handles one connection's request. + async fn request(&mut self, request: Request) { + match request { + Request::Connected { + conn, + out, + close, + held, + } => { + self.conns.insert(conn, Conn::new(out, close, held)); + } + Request::Subscribe(subscribe) => { + let Subscribe { + conn, + id, + query, + shape, + } = *subscribe; + self.subscribe(conn, id, query, shape).await + } + Request::Unsubscribe { conn, id } => self.unsubscribe(conn, id).await, + Request::Disconnected { conn } => self.disconnected(conn).await, + } + } + + /// Ends the connection's subscription `id`. A registration still in + /// flight for the id is not wanted either: `Registered` then finds no + /// owner and unregisters it, and the id is free to subscribe again. + async fn unsubscribe(&mut self, conn: u64, id: u64) { + let Some(connection) = self.conns.get_mut(&conn) else { + return; + }; + let sub = connection.subs.remove(&id); + let registering = connection.registering.remove(&id); + connection.note_held(); + if let Some(sub) = sub { + self.forget(sub); + let _ = self.commands.send(Command::Unregister(sub)).await; + } + if let Some(token) = registering { + self.by_token.remove(&token); + } + } + + /// Ends every subscription of a connection that went away, and lets + /// go of its registrations in flight. + async fn disconnected(&mut self, conn: u64) { + let Some(gone) = self.conns.remove(&conn) else { + return; + }; + let subs: Vec = gone.subs.values().copied().collect(); + for sub in &subs { + self.forget(*sub); + } + for token in gone.registering.values() { + self.by_token.remove(token); + } + if !subs.is_empty() { + let _ = self.commands.send(Command::UnregisterAll(subs)).await; + } + } + + /// Registers a connection's query with the engine, unless the id is + /// taken or the connection holds its limit. + async fn subscribe( + &mut self, + conn: u64, + id: u64, + query: MultiTableReadQuery, + shape: Option, + ) { + let Some(connection) = self.conns.get_mut(&conn) else { + return; + }; + if connection.subs.contains_key(&id) || connection.registering.contains_key(&id) { + connection.send(refused(id, "that id is already subscribed")); + return; + } + if self.per_connection > 0 && connection.holds() >= self.per_connection { + let reason = format!( + "the subscription limit of {} per connection is reached", + self.per_connection + ); + connection.send(refused(id, &reason)); + return; + } + let token = self.next_token; + self.next_token += 1; + connection.registering.insert(id, token); + connection.note_held(); + self.by_token.insert(token, (conn, id, shape)); + let command = Command::Register { + sink: self.sink, + query, + token, + }; + if self.commands.send(command).await.is_err() { + self.by_token.remove(&token); + if let Some(connection) = self.conns.get_mut(&conn) { + connection.registering.remove(&id); + connection.note_held(); + connection.send(refused(id, "the engine is shutting down")); + } + } + } + + /// Handles one engine event. + async fn event(&mut self, event: Event) { + match event { + Event::Registered { + token, + sub, + reads, + updates, + } => self.registered(token, sub, reads, &updates).await, + Event::Hydrated(subs) => { + for sub in subs { + if self.hydrating.remove(&sub) { + self.accept(sub).await; + } + } + } + Event::Refused { sub, reason } => self.drop_sub(sub, &reason).await, + Event::Capped { sub } => { + self.drop_sub(sub, "the query has more rows than the engine follows") + .await; + } + Event::Committed { updates, .. } | Event::Landed { updates } => self.changed(&updates), + Event::Heavy { sub, table, rows } => { + if let Some((conn, id)) = self.by_sub.get(&sub) { + log_warn!("s2s connection {conn}: query {id} holds {rows} rows of {table}"); + } + } + Event::Heartbeat(n) => self.stats.freshness.applied(n), + } + } + + /// The engine registered the query of `token` as `sub`. When the + /// connection went away, or unsubscribed the id, while the + /// registration was in flight, the subscription is not wanted and is + /// unregistered. A twin's rows come with the registration, in + /// `updates`, and seed a rows-mode shape's snapshot. + async fn registered(&mut self, token: u64, sub: SubId, reads: usize, updates: &[Delta]) { + let Some((conn, id, shape)) = self.by_token.remove(&token) else { + let _ = self.commands.send(Command::Unregister(sub)).await; + return; + }; + let Some(connection) = self.conns.get_mut(&conn) else { + let _ = self.commands.send(Command::Unregister(sub)).await; + return; + }; + connection.registering.remove(&id); + connection.subs.insert(id, sub); + connection.note_held(); + self.by_sub.insert(sub, (conn, id)); + if let Some(mut shape) = shape { + for delta in updates { + if delta.targets().any(|target| target.sub == sub) { + shape.absorb(&delta.table, &delta.op); + } + } + self.shapes.insert(sub, shape); + } + if reads == 0 { + self.accept(sub).await; + } else { + self.hydrating.insert(sub); + } + } + + /// `accepted`, with the hydrated rows for a rows-mode subscription. + /// A table that left the catalog meanwhile cannot be served: the + /// subscription is refused, and ended on the engine side too. + async fn accept(&mut self, sub: SubId) { + let Some((conn, id)) = self.by_sub.get(&sub).copied() else { + return; + }; + let Some(connection) = self.conns.get_mut(&conn) else { + return; + }; + let Some(shape) = self.shapes.get_mut(&sub) else { + connection.send(accepted(id)); + return; + }; + let catalog = self.catalog.load(); + let table_names: Vec = shape + .nodes + .iter() + .map(|node| node.table.as_str().to_owned()) + .collect(); + let mut nodes = Vec::with_capacity(shape.nodes.len()); + for node in shape.nodes.iter_mut() { + let Some(table) = catalog.table(node.table.as_str()) else { + self.drop_sub(sub, "a table left the catalog").await; + return; + }; + let rows: Vec = node + .snapshot + .drain() + .filter_map(|(_, row)| row) + .map(|row| wire::row_json(&row, table)) + .collect(); + nodes.push(json!({ + "table": table.name.as_str(), + "key": table.pkey.iter().map(|c| c.as_str()).collect::>(), + "parent": node.parent.map(|p| table_names[p].clone()), + "parent_column": node.parent_column, + "child_column": node.child_column, + "inner": node.inner, + "rows": rows, + })); + } + let (columns, names) = match &shape.projection { + Projection::All => (Json::Null, Json::Null), + Projection::Count { name } => (Json::Null, json!([name])), + Projection::Columns(columns) => ( + Json::Array( + columns + .iter() + .map(|p| json!([p.table.as_str(), p.column.as_str()])) + .collect(), + ), + Json::Array( + columns + .iter() + .map(|p| Json::from(p.name.as_str())) + .collect(), + ), + ), + }; + let order: Vec = shape + .order + .iter() + .map(|clause| { + json!([ + clause.column.as_str(), + match clause.direction { + Order::ASC => "asc", + Order::DESC => "desc", + } + ]) + }) + .collect(); + let limit = (shape.limit != u32::MAX).then_some(shape.limit); + connection.send(json!(["accepted", { + "id": id, + "nodes": nodes, + "columns": columns, + "names": names, + "count": matches!(shape.projection, Projection::Count { .. }), + "order": order, + "limit": limit, + "offset": shape.offset, + }])); + } + + /// Tells each connection what the deltas did to its subscriptions: + /// rows-mode ones get the rows per table, the rest a touch. + /// Every row is serialized once, for the first rows-mode + /// subscription it reaches, and spliced into every frame that + /// carries it. + /// + /// A subscription still hydrating is not told: the backend reads + /// only after `accepted`, which carries the change (rows mode) or + /// lands after it (touch mode). For a rows-mode subscription, per + /// table in order of appearance, the last operation on a key wins: a + /// row that left and came back is a put, one that came and went is + /// a delete, and no key is in both lists. A row of a table that left + /// the catalog mid-flight cannot be kept by the backend, so it goes + /// out as a delete. + fn changed(&mut self, updates: &[Delta]) { + let mut touched: HashMap> = HashMap::new(); + let mut deltas: HashMap = HashMap::new(); + let mut images: Vec> = Vec::new(); + let catalog = self.catalog.load(); + for delta in updates { + let table = catalog.table(delta.table.as_str()); + let mut image: Option> = None; + for target in delta.targets() { + let Some((conn, id)) = self.by_sub.get(&target.sub) else { + continue; + }; + if self.hydrating.contains(&target.sub) { + if let Some(shape) = self.shapes.get_mut(&target.sub) { + shape.absorb(&delta.table, &delta.op); + } + continue; + } + if !self.shapes.contains_key(&target.sub) { + touched.entry(*conn).or_default().insert(*id); + continue; + } + let keys = table_keys(deltas.entry(target.sub).or_default(), &delta.table); + match (&delta.op, table) { + (DataFrameOperation::Add(key, row), Some(table)) => { + let image = *image.get_or_insert_with(|| { + serialized(&wire::row_json(row, table)).map(|raw| { + images.push(raw); + images.len() - 1 + }) + }); + keys.insert(key.clone(), image); + } + (DataFrameOperation::Delete(key, _) | DataFrameOperation::Add(key, _), _) => { + keys.insert(key.clone(), None); + } + } + } + } + for (sub, changes) in deltas { + if let Some((conn, id)) = self.by_sub.get(&sub) + && let Some(connection) = self.conns.get_mut(conn) + && let Some(text) = delta_frame(*id, &changes, &images) + { + connection.send_text(text); + } + } + for (conn, ids) in touched { + if let Some(connection) = self.conns.get_mut(&conn) { + let ids: Vec = ids.into_iter().collect(); + connection.send(json!(["touched", {"ids": ids}])); + } + } + } + + /// The engine stopped following a subscription, or this task cannot + /// serve it: the backend must stop trusting its cache for it, and + /// the engine must stop routing to it. The engine's own `Refused` + /// has already unregistered the subscription, and the extra + /// unregister is a no-op (ids are never reused); `Capped` and this + /// task's own refusals leave it registered, and need it. + async fn drop_sub(&mut self, sub: SubId, reason: &str) { + if let Some((conn, id)) = self.by_sub.get(&sub).copied() + && let Some(connection) = self.conns.get_mut(&conn) + { + connection.subs.remove(&id); + connection.note_held(); + connection.send(refused(id, reason)); + } + self.forget(sub); + let _ = self.commands.send(Command::Unregister(sub)).await; + } + + /// Drops the task's own records of `sub`. + fn forget(&mut self, sub: SubId) { + self.by_sub.remove(&sub); + self.shapes.remove(&sub); + self.hydrating.remove(&sub); + } +} + +/// One table's keys in a subscription's changes, added in order of +/// appearance the first time the table is named. +fn table_keys<'c>( + changes: &'c mut Changes, + table: &TableName, +) -> &'c mut HashMap> { + let index = match changes.iter().position(|(named, _)| named == table) { + Some(index) => index, + None => { + changes.push((table.clone(), HashMap::new())); + changes.len() - 1 + } + }; + &mut changes[index].1 +} + +/// A `delta` frame for subscription `id`, its put rows spliced in from +/// `images`, the transaction's serialized rows. +fn delta_frame(id: u64, changes: &Changes, images: &[Box]) -> Option { + let changes: Vec = changes + .iter() + .map(|(table, keys)| { + let mut put = Vec::new(); + let mut del = Vec::new(); + for (key, row) in keys { + match row { + Some(index) => put.push(&*images[*index]), + None => del.push(wire::key_json(key)), + } + } + TableChanges { + table: table.as_str(), + put, + del, + } + }) + .collect(); + serde_json::to_string(&("delta", DeltaBody { id, changes })).ok() +} + +/// A row image as JSON text, written once and spliced into every frame +/// that carries it. A JSON value always serializes; a row that somehow +/// did not would be sent as a delete, which the backend refills. +fn serialized(image: &Json) -> Option> { + serde_json::value::to_raw_value(image).ok() +} + +/// An `accepted` with no rows. +fn accepted(id: u64) -> Json { + json!(["accepted", {"id": id}]) +} + +/// A `refused` for the backend's `id`. +fn refused(id: u64, reason: &str) -> Json { + json!(["refused", {"id": id, "reason": reason}]) +} + +/// A frame from the connection's own handler. `false` when the queue is +/// full, which means the backend is not reading: the handler then ends +/// the connection, as the task would, rather than drop the frame and +/// leave a subscription unanswered. +fn send(out: &Outgoing, message: Json) -> bool { + out.push(message.to_string()) +} + +/// The route handler: upgrades the socket once the engine is serving. +pub async fn connect( + State(state): State>, + headers: HeaderMap, + upgrade: WebSocketUpgrade, +) -> Response { + if let Some(token) = &state.config.s2s_token + && !authorized(&headers, &token.0) + { + refusal(&state.s2s.unauthorized, "no valid bearer token"); + return (StatusCode::UNAUTHORIZED, "a bearer token is required").into_response(); + } + if !*state.ready.borrow() { + return ( + StatusCode::SERVICE_UNAVAILABLE, + "the engine is not serving yet", + ) + .into_response(); + } + let Some(admitted) = state.s2s.admit() else { + let what = format!( + "{} connections are open, the most allowed", + state.s2s.max_connections + ); + refusal(&state.s2s.full, &what); + return (StatusCode::SERVICE_UNAVAILABLE, "too many S2S connections").into_response(); + }; + let upgrade = upgrade.max_message_size(state.config.max_message_bytes); + upgrade.on_upgrade(move |socket| async move { + handle(socket, state).await; + drop(admitted); + }) +} + +/// Whether the request carries `Authorization: Bearer `. The +/// comparison walks the whole token whatever was presented, so neither +/// a near miss nor a presented length says anything about the token. +fn authorized(headers: &HeaderMap, token: &str) -> bool { + let Some(presented) = headers + .get(AUTHORIZATION) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.split_once(' ')) + .filter(|(scheme, _)| scheme.eq_ignore_ascii_case("bearer")) + .map(|(_, presented)| presented.trim_start_matches(' ')) + else { + return false; + }; + let (presented, token) = (presented.as_bytes(), token.as_bytes()); + let lengths = u8::from(presented.len() != token.len()); + token + .iter() + .enumerate() + .fold(lengths, |differ, (at, byte)| { + differ | (presented.get(at).copied().unwrap_or(0) ^ byte) + }) + == 0 +} + +/// One upgraded connection, from `hello` to close. Its writer gets +/// [`WRITER_DRAIN`] to flush once the session ends and is then cut off, +/// so a backend that stopped reading cannot hold the connection open. +/// Whatever ends the session, a panic in a frame's handling included, +/// the task hears the connection is gone, or its subscriptions would +/// outlive it. +async fn handle(socket: WebSocket, state: Arc) { + let handle = &state.s2s; + let database = match database(&state).await { + Ok(database) => database, + Err(error) => { + refusal( + &handle.unidentified, + &format!("the database it streams cannot be read: {error}"), + ); + return; + } + }; + let conn = handle.next_conn.fetch_add(1, Ordering::Relaxed); + let (mut sink, stream) = socket.split(); + let (out, mut out_rx) = Outgoing::new(); + let (close, closed) = oneshot::channel::<&'static str>(); + let held = Arc::new(AtomicUsize::new(0)); + let drained = out.drained(); + let mut writer = tokio::spawn(async move { + while let Some(text) = out_rx.recv().await { + let len = text.len(); + let sent = sink.send(Message::Text(text.into())).await; + drained.sent(len); + if sent.is_err() { + break; + } + } + let _ = sink.send(Message::Close(None)).await; + }); + let limit = state.config.s2s_connection_subscriptions; + send( + &out, + json!(["hello", { + "protocol": PROTOCOL_VERSION, + "max_subscriptions": (limit > 0).then_some(limit), + "database": database, + }]), + ); + send(&out, lag_frame(&state.stats)); + if handle + .requests + .send(Request::Connected { + conn, + out: out.clone(), + close, + held: Arc::clone(&held), + }) + .await + .is_err() + { + writer.abort(); + return; + } + log_info!("s2s connection {conn} open"); + + let session = serve(conn, stream, closed, &out, &held, &state); + if AssertUnwindSafe(session).catch_unwind().await.is_err() { + log_warn!("s2s connection {conn}: a frame's handling panicked"); + } + let _ = handle.requests.send(Request::Disconnected { conn }).await; + drop(out); + if tokio::time::timeout(WRITER_DRAIN, &mut writer).await.is_err() { + writer.abort(); + log_warn!("s2s connection {conn}: its frames were not read in time, cut off"); + } + log_info!("s2s connection {conn} closed"); +} + +/// The `lag` frame: how far the engine's applied changes trail the +/// database, `null` before the first heartbeat is sent. +fn lag_frame(stats: &Stats) -> Json { + let ms = stats + .freshness + .lag() + .map(|lag| u64::try_from(lag.as_millis()).unwrap_or(u64::MAX)); + json!(["lag", {"ms": ms}]) +} + +/// Reads frames until the backend closes, the task asks to close, or +/// the task is gone. The task enforces the subscription limit; checking +/// it here too keeps a connection at its limit from having more queries +/// planned, which reads Postgres, only to be refused. +async fn serve( + conn: u64, + mut stream: SplitStream, + mut closed: oneshot::Receiver<&'static str>, + out: &Outgoing, + held: &AtomicUsize, + state: &Arc, +) { + let handle = &state.s2s; + let limit = state.config.s2s_connection_subscriptions; + let unread = || log_warn!("s2s connection {conn} closed: the backend is not reading its frames"); + loop { + let frame = tokio::select! { + frame = stream.next() => frame, + reason = &mut closed => { + if let Ok(reason) = reason { + log_warn!("s2s connection {conn} closed: {reason}"); + } + return; + } + }; + let text = match frame { + Some(Ok(Message::Text(text))) => text, + Some(Ok(Message::Ping(_))) + | Some(Ok(Message::Pong(_))) + | Some(Ok(Message::Binary(_))) => continue, + Some(Ok(Message::Close(_))) | Some(Err(_)) | None => return, + }; + match upstream(&text) { + Ok(Frame::Ping) => { + if !send(out, json!(["pong"])) { + unread(); + return; + } + } + Ok(Frame::Unsubscribe { id }) => { + if handle + .requests + .send(Request::Unsubscribe { conn, id }) + .await + .is_err() + { + return; + } + } + Ok(Frame::Subscribe { + id, + sql, + binds, + rows, + }) => { + let prepared = if limit > 0 && held.load(Ordering::Acquire) >= limit { + Err(format!( + "the subscription limit of {limit} per connection is reached" + )) + } else { + prepare(state, conn, id, &sql, &binds, rows).await + }; + let request = match prepared { + Ok(request) => request, + Err(reason) => { + log_debug!("s2s connection {conn}: query {id} refused: {reason}"); + if !send(out, refused(id, &reason)) { + unread(); + return; + } + continue; + } + }; + if handle.requests.send(request).await.is_err() { + return; + } + } + Err(Rejected { id: Some(id), reason }) => { + log_warn!("s2s connection {conn}: query {id} refused: {reason}"); + if !send(out, refused(id, &reason)) { + unread(); + return; + } + } + Err(Rejected { id: None, reason }) => { + log_warn!("s2s connection {conn}: {reason}"); + if !send(out, json!(["error", {"reason": reason}])) { + unread(); + return; + } + } + } + } +} + +/// Parses the SQL against the catalog and plans the query the way a +/// Zero query is planned (which side of each join is read whole, how a +/// page is read), so the engine registers it ready to run. A bind of a +/// type the catalog has not seen may be an enum created since it was +/// loaded: those are looked up first ([`learn_enums`]). A query larger +/// than the engine should take on its own thread is refused before it +/// is planned ([`within_limits`]). +async fn prepare( + state: &AppState, + conn: u64, + id: u64, + sql: &str, + binds: &[BindValue], + rows: bool, +) -> Result { + let mut catalog = state.catalog.load(); + let unknown: Vec = binds + .iter() + .map(|bind| bind.type_oid) + .filter(|oid| !decodes(*oid) && !catalog.is_enum_oid(*oid)) + .collect(); + if !unknown.is_empty() && learn_enums(state, &unknown).await { + catalog = state.catalog.load(); + } + let select = parse_select(sql, binds, &catalog).map_err(|error| error.message)?; + within_limits( + &select.query, + state.config.s2s_max_disjuncts, + state.config.s2s_max_list_values, + )?; + let shape = rows.then(|| Shape::new(&select.query, select.projection, select.limit, select.offset)); + let translated = Translated { + query: select.query, + hidden: HashSet::new(), + uncursored: None, + }; + let planned = plan::plan( + sql, + translated, + state.config.policy(), + &state.plans, + &*state.storage, + ) + .await?; + Ok(Request::Subscribe(Box::new(Subscribe { + conn, + id, + query: planned.query, + shape, + }))) +} + +/// Which database the engine streams, for `hello`: the cluster's +/// `system_identifier` and the database's name, as a backend reads them +/// through its own connection. A backend whose reads go to another +/// database must not take this engine's rows for its own, and checks. +/// Read once and kept, since it cannot change while the process runs. +/// A read that fails is an error, and the connection is closed without +/// a `hello`: a `null` would have the backend take the engine on trust, +/// so it reconnects instead, reading Postgres meanwhile. +async fn database(state: &AppState) -> Result { + state + .s2s + .database + .get_or_try_init(|| async { + let rows = state + .storage + .simple_query(DATABASE_SQL) + .await + .map_err(|error| error.to_string())?; + match rows.first().map(|row| (row.get(0), row.get(1))) { + Some((Some(system), Some(name))) => Ok(json!({"system": system, "name": name})), + _ => Err("the identity query returned no row".to_owned()), + } + }) + .await + .cloned() +} + +/// The identity [`database`] announces. +const DATABASE_SQL: &str = + "SELECT system_identifier::text, current_database()::text FROM pg_control_system()"; + +/// Refuses a query whose filters would normalize to more than +/// `max_disjuncts` disjuncts, or whose `IN` and `NOT IN` lists hold more +/// than `max_list_values` values; a limit of 0 is none. The engine does +/// both on the thread every subscriber's changes go through, so one such +/// query would hold them all up, and enough `OR`s would exhaust memory. +fn within_limits( + query: &MultiTableReadQuery, + max_disjuncts: u64, + max_list_values: usize, +) -> Result<(), String> { + let (disjuncts, values) = query_size(query); + if max_disjuncts > 0 && disjuncts > max_disjuncts { + return Err(format!( + "the filters expand to {disjuncts} alternatives, more than the {max_disjuncts} the engine follows" + )); + } + if max_list_values > 0 && values > max_list_values { + return Err(format!( + "the IN lists hold {values} values, more than the {max_list_values} the engine follows" + )); + } + Ok(()) +} + +/// How large a query is to the engine: the disjuncts its tables' +/// filters normalize to, summed, and the values its `IN` and `NOT IN` +/// lists hold. Counted from the filter trees without expanding them, so +/// a filter whose expansion would not fit in memory is measured cheaply; +/// the disjuncts saturate. +fn query_size(query: &MultiTableReadQuery) -> (u64, usize) { + let (mut disjuncts, mut values) = filter_size(&query.main_table.filter); + for join in &query.joins { + let (more_disjuncts, more_values) = query_size(&join.sub); + disjuncts = disjuncts.saturating_add(more_disjuncts); + values = values.saturating_add(more_values); + } + (disjuncts, values) +} + +/// [`query_size`] for one filter: an `AND` multiplies its parts' +/// disjuncts, an `OR` adds them, as normalizing to an `OR` of `AND`s +/// does. +fn filter_size(filter: &Where) -> (u64, usize) { + match filter { + Where::Condition(condition) => { + let values = match (&condition.comparison_operator, &condition.value) { + (ComparisonOperator::IN | ComparisonOperator::NOT_IN, Value::List(items)) => { + items.len() + } + _ => 0, + }; + (1, values) + } + Where::AND(parts) => parts.iter().map(filter_size).fold((1, 0), |(d, v), (pd, pv)| { + (d.saturating_mul(pd), v.saturating_add(pv)) + }), + Where::OR(parts) => parts.iter().map(filter_size).fold((0, 0), |(d, v), (pd, pv)| { + (d.saturating_add(pd), v.saturating_add(pv)) + }), + } +} + +/// Looks `oids` up among the database's enum types and teaches the +/// catalog the ones that are (with their array types), so a bind of an +/// enum created after the catalog was loaded reads as its label. Whether +/// any was learned; a failed lookup learns nothing, and the bind is then +/// refused as a type the engine does not carry. +async fn learn_enums(state: &AppState, oids: &[u32]) -> bool { + let list: Vec = oids.iter().map(u32::to_string).collect(); + let list = list.join(", "); + let sql = format!( + "SELECT oid::text, typarray::text FROM pg_type WHERE typtype = 'e' AND (oid IN ({list}) OR typarray IN ({list}))" + ); + let Ok(rows) = state.storage.simple_query(&sql).await else { + return false; + }; + let learned: Vec = rows + .iter() + .flat_map(|row| [row.get(0), row.get(1)]) + .flatten() + .filter_map(|text| text.parse::().ok()) + .filter(|oid| *oid != 0) + .collect(); + if learned.is_empty() { + return false; + } + state + .catalog + .update(|catalog| with_enums(catalog, learned.iter().copied())); + true +} + +/// One upstream frame. +#[derive(Debug)] +enum Frame { + Ping, + Subscribe { + id: u64, + sql: String, + binds: Vec, + rows: bool, + }, + Unsubscribe { + id: u64, + }, +} + +/// An upstream frame that could not be read: the `id` it carried, when +/// one could be read, and why. +#[derive(Debug)] +struct Rejected { + id: Option, + reason: String, +} + +impl Rejected { + /// A rejection that names no subscription. + fn frame(reason: impl Into) -> Rejected { + Rejected { + id: None, + reason: reason.into(), + } + } +} + +/// Reads one upstream frame. +fn upstream(text: &str) -> Result { + let frame: Json = serde_json::from_str(text) + .map_err(|error| Rejected::frame(format!("not JSON: {error}")))?; + let (tag, body) = match frame.as_array().map(Vec::as_slice) { + Some([Json::String(tag), body]) => (tag.as_str(), body), + Some([Json::String(tag)]) => (tag.as_str(), &Json::Null), + _ => return Err(Rejected::frame("a frame is a [tag, body] array")), + }; + let id = |body: &Json| { + body.get("id") + .and_then(Json::as_u64) + .ok_or_else(|| Rejected::frame("`id` must be a non-negative integer")) + }; + match tag { + "ping" => Ok(Frame::Ping), + "subscribe" => { + let id = id(body)?; + subscription(body) + .map(|(sql, binds, rows)| Frame::Subscribe { + id, + sql, + binds, + rows, + }) + .map_err(|reason| Rejected { + id: Some(id), + reason, + }) + } + "unsubscribe" => Ok(Frame::Unsubscribe { id: id(body)? }), + other => Err(Rejected::frame(format!("unknown tag `{other}`"))), + } +} + +/// A `subscribe` body past its `id`: the SQL, the binds, and whether +/// rows were asked for. +fn subscription(body: &Json) -> Result<(String, Vec, bool), String> { + let sql = body + .get("sql") + .and_then(Json::as_str) + .ok_or("`sql` must be a string")? + .to_string(); + let binds = body + .get("binds") + .map(|binds| binds.as_array().ok_or("`binds` must be an array")) + .transpose()? + .map(|binds| binds.iter().map(bind).collect::, _>>()) + .transpose()? + .unwrap_or_default(); + let rows = body.get("rows").and_then(Json::as_bool).unwrap_or(false); + Ok((sql, binds, rows)) +} + +/// One bind: its type OID and its base64 bytes, `null` for `NULL`. +fn bind(value: &Json) -> Result { + let type_oid = value + .get("oid") + .and_then(Json::as_u64) + .and_then(|oid| u32::try_from(oid).ok()) + .ok_or("a bind's `oid` must be a 32-bit integer")?; + let bytes = match value.get("bytes") { + None | Some(Json::Null) => None, + Some(Json::String(text)) => Some( + BASE64 + .decode(text) + .map_err(|_| "a bind's `bytes` must be base64")?, + ), + Some(_) => return Err("a bind's `bytes` must be a base64 string or null".to_string()), + }; + Ok(BindValue { type_oid, bytes }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ivm::{Audience, QueryPart, Subs}; + use crate::model::{Catalog, ColumnName, DbColumn, DbTable, Value, ValueType}; + + /// One table, `notes (id, name)`. + fn test_catalog() -> Catalog { + Catalog::new([DbTable::new( + "notes", + ["id"], + vec![ + DbColumn::new("id", ValueType::Int), + DbColumn::new("name", ValueType::String), + ], + )]) + } + + /// A task over one table, `notes (id, name)`, with connection 1 open: + /// the task, the engine's commands, and connection 1's frames. + fn task( + per_connection: usize, + ) -> ( + Task, + mpsc::Receiver>, + mpsc::Receiver, + ) { + let catalog = test_catalog(); + let (commands, commands_rx) = mpsc::channel(64); + let mut task = Task { + catalog: Arc::new(CatalogHandle::new(catalog)), + sink: 0, + commands, + conns: HashMap::new(), + by_token: HashMap::new(), + by_sub: HashMap::new(), + shapes: HashMap::new(), + hydrating: HashSet::new(), + next_token: 1, + per_connection, + stats: Stats::shared_with_metric_prefix("test".to_owned()), + }; + let frames = connect(&mut task, 1); + (task, commands_rx, frames) + } + + /// Opens connection `conn` on the task: its frames. + fn connect(task: &mut Task, conn: u64) -> mpsc::Receiver { + let (close, _closed) = oneshot::channel(); + let (out, frames) = Outgoing::new(); + task.conns + .insert(conn, Conn::new(out, close, Arc::new(AtomicUsize::new(0)))); + frames + } + + /// Connection `conn` subscribes `SELECT * FROM notes` under `id`. + async fn subscribe(task: &mut Task, conn: u64, id: u64, rows: bool) { + let select = + parse_select("SELECT * FROM notes", &[], &task.catalog.load()).expect("parses"); + let shape = rows.then(|| Shape::new(&select.query, select.projection, select.limit, select.offset)); + task.subscribe(conn, id, select.query, shape).await; + } + + /// The tokens of the registrations sent to the engine so far. + fn registered(commands: &mut mpsc::Receiver>) -> Vec { + let mut tokens = Vec::new(); + while let Ok(command) = commands.try_recv() { + if let Command::Register { token, .. } = command { + tokens.push(token); + } + } + tokens + } + + /// The frames sent to a connection so far. + fn frames(out: &mut mpsc::Receiver) -> Vec { + let mut frames = Vec::new(); + while let Ok(text) = out.try_recv() { + frames.push(serde_json::from_str(&text).expect("json")); + } + frames + } + + /// The engine registered the subscription of `token` as `sub`, with + /// nothing to read. + async fn register(task: &mut Task, token: u64, sub: u64) { + task.event(Event::Registered { + token, + sub: SubId(sub), + updates: Vec::new(), + reads: 0, + }) + .await; + } + + /// A delta adding note `id` named `name`, aimed at `subs`. + fn put(subs: &[u64], id: i64, name: &str) -> Delta { + let key = DataFrameKey::from(HashMap::from([(ColumnName::from("id"), Value::Int(id))])); + let row = DataFrameRow::from(HashMap::from([ + (ColumnName::from("id"), Value::Int(id)), + (ColumnName::from("name"), Value::from(name)), + ])); + Delta { + table: TableName::from("notes"), + op: DataFrameOperation::Add(key, row), + audiences: vec![Audience { + part: QueryPart::main(), + subs: Subs::Many(subs.iter().map(|sub| SubId(*sub)).collect()), + }], + } + } + + /// A connection holds as many subscriptions as it is allowed, those + /// the engine is still registering counted, and one more is refused + /// until one goes; other connections have their own allowance. + #[tokio::test] + async fn a_connection_holds_at_most_its_subscription_limit() { + let (mut task, mut commands, mut out) = task(2); + subscribe(&mut task, 1, 1, false).await; + subscribe(&mut task, 1, 2, false).await; + subscribe(&mut task, 1, 3, false).await; + assert_eq!(registered(&mut commands), vec![1, 2]); + assert_eq!( + frames(&mut out), + vec![ + json!(["refused", {"id": 3, "reason": "the subscription limit of 2 per connection is reached"}]) + ] + ); + + register(&mut task, 1, 10).await; + subscribe(&mut task, 1, 4, false).await; + assert!( + registered(&mut commands).is_empty(), + "a registered one still counts" + ); + + task.request(Request::Unsubscribe { conn: 1, id: 1 }).await; + subscribe(&mut task, 1, 5, false).await; + assert_eq!(registered(&mut commands), vec![3]); + + let mut other_out = connect(&mut task, 2); + subscribe(&mut task, 2, 1, false).await; + assert_eq!(registered(&mut commands), vec![4]); + assert!(frames(&mut other_out).is_empty()); + } + + /// With the limit at zero a connection subscribes without one. + #[tokio::test] + async fn a_zero_limit_lifts_it() { + let (mut task, mut commands, mut out) = task(0); + for id in 1..=5 { + subscribe(&mut task, 1, id, false).await; + } + assert_eq!(registered(&mut commands).len(), 5); + assert!(frames(&mut out).is_empty()); + } + + /// A row two rows-mode subscriptions share reaches both, and a row + /// of one transaction never stands in for another's. + #[tokio::test] + async fn each_subscription_gets_the_rows_of_its_own_deltas() { + let (mut task, _commands, mut out) = task(0); + subscribe(&mut task, 1, 1, true).await; + subscribe(&mut task, 1, 2, true).await; + register(&mut task, 1, 10).await; + register(&mut task, 2, 20).await; + frames(&mut out); + + task.changed(&[put(&[10, 20], 1, "one"), put(&[10], 2, "two")]); + let mut deltas: Vec = frames(&mut out); + deltas.sort_by_key(|frame| frame[1]["id"].as_u64()); + let puts = |frame: &Json| -> Vec<(i64, String)> { + let mut rows: Vec<(i64, String)> = frame[1]["changes"][0]["put"] + .as_array() + .expect("puts") + .iter() + .map(|row| { + ( + row["id"].as_i64().unwrap(), + row["name"].as_str().unwrap().to_owned(), + ) + }) + .collect(); + rows.sort(); + rows + }; + assert_eq!(deltas.len(), 2, "{deltas:?}"); + assert_eq!( + puts(&deltas[0]), + vec![(1, "one".to_owned()), (2, "two".to_owned())] + ); + assert_eq!(puts(&deltas[1]), vec![(1, "one".to_owned())]); + } + + /// Every connection hears the engine's lag from its heartbeats: + /// `null` before the first, then the age of the first while none is + /// applied, then the age of the newest one the engine's events said + /// was applied. + #[tokio::test] + async fn every_connection_hears_the_lag() { + let (mut task, _commands, mut out) = task(0); + let mut other = connect(&mut task, 2); + task.report_lag(); + assert_eq!(frames(&mut out), vec![json!(["lag", {"ms": null}])]); + task.stats + .freshness + .sent(1, std::time::Instant::now() - Duration::from_millis(1500)); + task.report_lag(); + for frames in [frames(&mut out), frames(&mut other)] { + let ms = frames[frames.len() - 1][1]["ms"].as_u64().expect("a lag"); + assert!((1500..2500).contains(&ms), "{ms}"); + } + task.stats.freshness.sent(2, std::time::Instant::now()); + task.event(Event::Heartbeat(2)).await; + task.report_lag(); + let ms = frames(&mut out)[0][1]["ms"].as_u64().expect("a lag"); + assert!(ms < 500, "caught up: {ms}"); + } + + /// A queue takes frames until either bound is reached, always takes + /// one frame when empty however large, and frees a frame's bytes + /// once the writer has sent it. + #[test] + fn the_outgoing_queue_is_bounded_in_bytes() { + let (out, mut frames) = Outgoing::new(); + let big = "x".repeat(OUTGOING_BYTES + 1); + assert!(out.push(big.clone()), "an empty queue takes any frame"); + assert!(!out.push("y".to_owned()), "past the byte bound"); + let sent = frames.try_recv().expect("the big frame"); + out.drained().sent(sent.len()); + assert!(out.push("y".to_owned()), "room again once it went out"); + assert_eq!(out.queued.load(Ordering::Acquire), 1); + } + + /// The writer's byte count keeps no sender alive: once every + /// `Outgoing` is gone the writer's queue ends, so a closing + /// connection flushes and sends its close frame without waiting. + #[test] + fn the_writer_queue_closes_with_its_senders() { + let (out, mut frames) = Outgoing::new(); + let drained = out.drained(); + assert!(out.push("x".to_owned())); + drop(out); + assert_eq!(frames.try_recv().as_deref(), Ok("x")); + drained.sent(1); + assert!(matches!( + frames.try_recv(), + Err(mpsc::error::TryRecvError::Disconnected) + )); + } + + /// A subscription touched by several deltas of one transaction is + /// named once, and the connection's handler sees what it holds. + #[tokio::test] + async fn a_touch_names_each_subscription_once() { + let (mut task, _commands, mut out) = task(0); + subscribe(&mut task, 1, 1, false).await; + assert_eq!(task.conns[&1].held.load(Ordering::Acquire), 1); + register(&mut task, 1, 10).await; + assert_eq!(task.conns[&1].held.load(Ordering::Acquire), 1); + frames(&mut out); + task.changed(&[put(&[10], 1, "one"), put(&[10], 2, "two"), put(&[10], 3, "three")]); + assert_eq!(frames(&mut out), vec![json!(["touched", {"ids": [1]}])]); + task.request(Request::Unsubscribe { conn: 1, id: 1 }).await; + assert_eq!(task.conns[&1].held.load(Ordering::Acquire), 0); + } + + /// An id unsubscribed while still registering frees its place at + /// once, and its late registration is unregistered. + #[tokio::test] + async fn an_id_unsubscribed_while_registering_is_let_go() { + let (mut task, mut commands, _out) = task(1); + subscribe(&mut task, 1, 1, false).await; + assert_eq!(registered(&mut commands), vec![1]); + task.request(Request::Unsubscribe { conn: 1, id: 1 }).await; + subscribe(&mut task, 1, 2, false).await; + assert_eq!(registered(&mut commands), vec![2], "the place came back"); + register(&mut task, 1, 10).await; + assert!(matches!( + commands.try_recv(), + Ok(Command::Unregister(SubId(10))) + )); + assert!(task.by_sub.is_empty()); + } + + /// Places are handed out up to the limit and come back when their + /// connections close; a zero limit hands them out without one. + #[test] + fn connections_are_admitted_up_to_the_limit() { + let (handle, _requests) = Handle::new(2); + let first = handle.admit().expect("a place"); + let second = handle.admit().expect("a place"); + assert!(handle.admit().is_none()); + drop(first); + let third = handle.admit().expect("the freed place"); + assert!(handle.admit().is_none()); + drop((second, third)); + assert_eq!(handle.open.load(Ordering::Acquire), 0); + + let (open, _requests) = Handle::new(0); + let places: Vec = (0..100).map(|_| open.admit().expect("a place")).collect(); + assert_eq!(open.open.load(Ordering::Acquire), places.len()); + } + + /// The first warning passes, the next ones within the period are + /// counted, and the first after it says how many were held back. + #[test] + fn a_throttled_warning_passes_once_per_period() { + let throttle = Throttle::new(Duration::from_secs(60)); + let start = Instant::now(); + assert_eq!(throttle.pass(start), Some(0)); + assert_eq!(throttle.pass(start + Duration::from_secs(1)), None); + assert_eq!(throttle.pass(start + Duration::from_secs(59)), None); + assert_eq!(throttle.pass(start + Duration::from_secs(60)), Some(2)); + assert_eq!(throttle.pass(start + Duration::from_secs(61)), None); + assert_eq!(throttle.pass(start + Duration::from_secs(200)), Some(1)); + } + + /// Only the configured token, as a bearer, opens the endpoint. + #[test] + fn only_the_bearer_token_authorizes() { + let with = |value: &str| { + let mut headers = HeaderMap::new(); + headers.insert(AUTHORIZATION, value.parse().unwrap()); + headers + }; + assert!(authorized(&with("Bearer s3cret"), "s3cret")); + assert!(authorized(&with("bearer s3cret"), "s3cret")); + assert!(authorized(&with("BEARER s3cret"), "s3cret")); + assert!(!authorized(&with("Bearer s3cres"), "s3cret")); + assert!(!authorized(&with("Bearer s3cret2"), "s3cret")); + assert!(!authorized(&with("Bearer s3cre"), "s3cret")); + assert!(!authorized(&with("Bearer "), "s3cret")); + assert!(!authorized(&with("Basic s3cret"), "s3cret")); + assert!(!authorized(&with("s3cret"), "s3cret")); + assert!(!authorized(&HeaderMap::new(), "s3cret")); + } + + /// A query's size is counted from its filter trees as their expansion + /// would be: thirty `AND`ed `OR` pairs are 2^30 disjuncts, measured + /// without building one, and refused; an `IN` list counts its values. + #[test] + fn an_oversized_query_is_refused_before_the_engine_sees_it() { + let catalog = test_catalog(); + let pairs = |k: usize| { + let filter = (0..k) + .map(|i| format!(r#"("notes"."id" = {i} OR "notes"."name" = 'n{i}')"#)) + .collect::>() + .join(" AND "); + parse_select(&format!(r#"SELECT * FROM "notes" WHERE {filter}"#), &[], &catalog) + .expect("parses") + .query + }; + assert_eq!(query_size(&pairs(3)), (8, 0)); + assert_eq!(query_size(&pairs(30)).0, 1 << 30); + assert!(within_limits(&pairs(10), 1024, 0).is_ok()); + let refused = within_limits(&pairs(11), 1024, 0).expect_err("2048 disjuncts"); + assert!(refused.contains("2048 alternatives"), "{refused}"); + assert!(within_limits(&pairs(30), 0, 0).is_ok(), "0 lifts the limit"); + + let list = (0..5).map(|i| i.to_string()).collect::>().join(", "); + let listed = parse_select( + &format!(r#"SELECT * FROM "notes" WHERE "notes"."id" IN ({list}) OR "notes"."id" NOT IN ({list})"#), + &[], + &catalog, + ) + .expect("parses") + .query; + assert_eq!(query_size(&listed), (2, 10)); + assert!(within_limits(&listed, 0, 10).is_ok()); + assert!(within_limits(&listed, 0, 9).is_err()); + } + + /// Frames read into requests, and a bad one says which id it was for + /// when it can. + #[test] + fn frames_parse_into_requests() { + let frame = upstream( + r#"["subscribe", {"id": 7, "sql": "SELECT * FROM t WHERE a = $1", "binds": [{"oid": 25, "bytes": "T1BFTg=="}, {"oid": 20, "bytes": null}], "rows": true}]"#, + ) + .expect("parses"); + match frame { + Frame::Subscribe { + id, + sql, + binds, + rows, + } => { + assert_eq!(id, 7); + assert_eq!(sql, "SELECT * FROM t WHERE a = $1"); + assert!(rows); + assert_eq!( + binds, + vec![ + BindValue { + type_oid: 25, + bytes: Some(b"OPEN".to_vec()) + }, + BindValue { + type_oid: 20, + bytes: None + } + ] + ); + } + other => panic!("{other:?}"), + } + assert!(matches!( + upstream(r#"["unsubscribe", {"id": 2}]"#), + Ok(Frame::Unsubscribe { id: 2 }) + )); + assert!(matches!(upstream(r#"["ping"]"#), Ok(Frame::Ping))); + assert!(upstream(r#"["subscribe", {"id": -1}]"#).is_err()); + assert!(upstream(r#"["nope", {}]"#).is_err()); + assert!(upstream("{}").is_err()); + assert!(matches!( + upstream( + r#"["subscribe", {"id": 1, "sql": "x", "binds": [{"oid": 25, "bytes": "%%"}]}]"# + ), + Err(Rejected { id: Some(1), .. }) + )); + assert!(matches!( + upstream(r#"["subscribe", {"id": 7, "sql": 5}]"#), + Err(Rejected { id: Some(7), .. }) + )); + assert!(matches!( + upstream(r#"["subscribe", {"sql": "x"}]"#), + Err(Rejected { id: None, .. }) + )); + } +} diff --git a/src/client/sampler.rs b/src/client/sampler.rs index 1711826..7bfe484 100644 --- a/src/client/sampler.rs +++ b/src/client/sampler.rs @@ -58,7 +58,7 @@ fn run(state: Arc, interval: Duration) { stats.publish_thread_core_cpu(attributed.thread_cores); stats.publish_thread_core(attributed.placement); stats.publish_thread_migrations(attributed.migrations); - if ticks % 6 == 0 { + if ticks.is_multiple_of(6) { let (message, fields) = stats.summary_line(&previous, last_summary.elapsed()); crate::log::event(Level::Info, message, fields); previous = stats.snapshot(); diff --git a/src/client/wire.rs b/src/client/wire.rs index 41ac9ce..e1d1526 100644 --- a/src/client/wire.rs +++ b/src/client/wire.rs @@ -54,8 +54,10 @@ pub fn write_put(out: &mut Vec, table: &DbTable, row: &DataFrameRow) { } /// Append the `rowsPatch` entry deleting `key` from `table`: -/// `{"op":"del","tableName":…,"id":{…}}`. -pub fn write_del(out: &mut Vec, table: &str, key: &DataFrameKey) { +/// `{"op":"del","tableName":…,"id":{…}}`. `declared` is the table's +/// schema when the catalog has it, so a key column that is a time goes +/// out in the protocol's milliseconds like every other cell. +pub fn write_del(out: &mut Vec, table: &str, key: &DataFrameKey, declared: Option<&DbTable>) { out.extend_from_slice(b"{\"op\":\"del\",\"tableName\":"); write_text(out, table); out.extend_from_slice(b",\"id\":{"); @@ -67,14 +69,17 @@ pub fn write_del(out: &mut Vec, table: &str, key: &DataFrameKey) { first = false; write_text(out, column.as_str()); out.push(b':'); - write_value(out, value, None); + let r#type = declared + .and_then(|table| table.column(column.as_str())) + .map(|column| &column.r#type); + write_value(out, value, r#type); } out.extend_from_slice(b"}}"); } /// Append one cell as JSON, in the light of its declared type: a JSON -/// column's text is embedded as it is when it parses, a time as its -/// milliseconds, a list item by item. +/// column's text is embedded as it is when it parses, a time (kept as +/// microseconds) as its milliseconds, a list item by item. pub fn write_value(out: &mut Vec, value: &Value, declared: Option<&ValueType>) { match value { Value::Null => out.extend_from_slice(b"null"), @@ -86,6 +91,9 @@ pub fn write_value(out: &mut Vec, value: &Value, declared: Option<&ValueType } _ => write_text(out, text), }, + Value::Int(int) if matches!(declared, Some(ValueType::Timestamp)) => { + let _ = write!(out, "{}", int.div_euclid(1000)); + } Value::Int(int) => { let _ = write!(out, "{int}"); } @@ -123,7 +131,8 @@ pub fn write_value(out: &mut Vec, value: &Value, declared: Option<&ValueType out.push(b']'); } Value::Map(_) | Value::Set(_) => { - if serde_json::to_writer(&mut *out, &value_json(value, declared)).is_err() { + let json = json_in(value, declared, Unit::Millis); + if serde_json::to_writer(&mut *out, &json).is_err() { out.extend_from_slice(b"null"); } } @@ -156,8 +165,47 @@ pub fn key_json(key: &DataFrameKey) -> Json { Json::Object(object) } -/// One cell as JSON, in the light of its declared type. +/// One cell as JSON, in the light of its declared type, with a time as +/// the microseconds the engine keeps (the server-to-server form; the +/// Zero writer above scales to milliseconds). pub fn value_json(value: &Value, declared: Option<&ValueType>) -> Json { + json_in(value, declared, Unit::Micros) +} + +/// The path a cell goes out on, which decides its unit: the engine's +/// microseconds on the server-to-server path, milliseconds on Zero's. +#[derive(Clone, Copy)] +enum Unit { + Micros, + Millis, +} + +impl Unit { + /// A time the engine keeps in microseconds, in this unit. + fn of(self, micros: i64) -> i64 { + match self { + Unit::Micros => micros, + Unit::Millis => micros.div_euclid(1000), + } + } + + /// A float JSON has no number for (`NaN`, `±Infinity`). The + /// server-to-server path names it, as Postgres spells it, so the + /// backend gets the value back; Zero's protocol has always sent + /// `null`. + fn non_finite(self, float: f64) -> Json { + match self { + Unit::Millis => Json::Null, + Unit::Micros if float.is_nan() => Json::from("NaN"), + Unit::Micros if float > 0.0 => Json::from("Infinity"), + Unit::Micros => Json::from("-Infinity"), + } + } +} + +/// One cell as JSON, with every time in it, nested ones included, in +/// `unit`. +fn json_in(value: &Value, declared: Option<&ValueType>, unit: Unit) -> Json { match value { Value::Null => Json::Null, Value::String(text) => match declared { @@ -166,22 +214,31 @@ pub fn value_json(value: &Value, declared: Option<&ValueType>) -> Json { } _ => Json::String(text.clone()), }, + Value::Int(int) if matches!(declared, Some(ValueType::Timestamp)) => { + Json::from(unit.of(*int)) + } Value::Int(int) => Json::from(*int), Value::Float(float) if whole(*float) => Json::from(*float as i64), - Value::Float(float) => { - serde_json::Number::from_f64(*float).map_or(Json::Null, Json::Number) - } + Value::Float(float) => match serde_json::Number::from_f64(*float) { + Some(number) => Json::Number(number), + None => unit.non_finite(*float), + }, Value::Bool(flag) => Json::Bool(*flag), - Value::Date(date) => date - .and_hms_opt(0, 0, 0) - .map_or(Json::Null, |at| Json::from(at.and_utc().timestamp_millis())), - Value::Datetime(datetime) => Json::from(datetime.and_utc().timestamp_millis()), + Value::Date(date) => date.and_hms_opt(0, 0, 0).map_or(Json::Null, |at| { + Json::from(unit.of(at.and_utc().timestamp_micros())) + }), + Value::Datetime(datetime) => Json::from(unit.of(datetime.and_utc().timestamp_micros())), Value::List(items) => { let inner = match declared { Some(ValueType::List(inner)) => Some(inner.as_ref()), _ => None, }; - Json::Array(items.iter().map(|item| value_json(item, inner)).collect()) + Json::Array( + items + .iter() + .map(|item| json_in(item, inner, unit)) + .collect(), + ) } Value::Map(entries) => { let mut object = Map::with_capacity(entries.len()); @@ -190,14 +247,14 @@ pub fn value_json(value: &Value, declared: Option<&ValueType>) -> Json { Value::String(text) => text.clone(), other => format!("{other:?}"), }; - object.insert(key, value_json(item, None)); + object.insert(key, json_in(item, None, unit)); } Json::Object(object) } Value::Set(set) => Json::Array( set.members() .iter() - .map(|item| value_json(item, None)) + .map(|item| json_in(item, None, unit)) .collect(), ), } @@ -299,12 +356,60 @@ mod tests { ); let written: Json = serde_json::from_slice(&fast).expect("valid JSON"); assert_eq!(written["tableName"], "odd \"table\""); - assert_eq!(written["value"], row_json(&planned, &table)); + let mut expected = row_json(&planned, &table); + for time in ["ts", "day", "at"] { + expected[time] = Json::from(expected[time].as_i64().unwrap() / 1000); + } + expected["nan"] = Json::Null; + assert_eq!( + written["value"], expected, + "a time goes out in milliseconds, a NaN as null" + ); } } + /// Every form of a time goes out in one unit per path: microseconds + /// on the server-to-server one, milliseconds on Zero's. + #[test] + fn a_time_has_one_unit_per_path() { + let at = chrono::NaiveDate::from_ymd_opt(2026, 10, 19) + .unwrap() + .and_hms_micro_opt(1, 2, 3, 456_789) + .unwrap(); + let micros = at.and_utc().timestamp_micros(); + let midnight = at + .date() + .and_hms_opt(0, 0, 0) + .unwrap() + .and_utc() + .timestamp_micros(); + let stamp = Some(&ValueType::Timestamp); + assert_eq!(value_json(&Value::Int(micros), stamp), Json::from(micros)); + assert_eq!(value_json(&Value::Datetime(at), None), Json::from(micros)); + assert_eq!( + value_json(&Value::Date(at.date()), None), + Json::from(midnight) + ); + let zero = |value: &Value, declared| { + let mut out = Vec::new(); + write_value(&mut out, value, declared); + String::from_utf8(out).unwrap() + }; + let millis = (micros / 1000).to_string(); + assert_eq!(zero(&Value::Int(micros), stamp), millis); + assert_eq!(zero(&Value::Datetime(at), None), millis); + assert_eq!( + zero(&Value::Date(at.date()), None), + (midnight / 1000).to_string() + ); + let nested = Value::Map(HashMap::from([(Value::from("at"), Value::Datetime(at))])); + assert_eq!(value_json(&nested, None), serde_json::json!({"at": micros})); + assert_eq!(zero(&nested, None), format!(r#"{{"at":{millis}}}"#)); + } + /// A whole-valued float is written as an integer on both paths, a - /// fractional one as it is, and a NaN as null. + /// fractional one as it is; a NaN or an infinity is `null` to Zero + /// and named to a server-to-server backend. #[test] fn whole_floats_are_written_as_integers() { let mut out = Vec::new(); @@ -322,10 +427,17 @@ mod tests { assert_eq!(value_json(&Value::Float(1.0), None), serde_json::json!(1)); assert_eq!(value_json(&Value::Float(1.5), None), serde_json::json!(1.5)); assert!(!whole(1e17), "beyond 2^53 a float is left as it is"); + assert_eq!(value_json(&Value::Float(f64::NAN), None), serde_json::json!("NaN")); + assert_eq!(value_json(&Value::Float(f64::INFINITY), None), serde_json::json!("Infinity")); + assert_eq!( + value_json(&Value::Float(f64::NEG_INFINITY), None), + serde_json::json!("-Infinity") + ); } /// JSON columns are embedded, times are numbers, and the client's - /// schema filters the columns. + /// schema filters the columns. Keys are escaped, and a key column + /// that is a time goes out in milliseconds too. #[test] fn rows_serialize_for_the_client() { let table = DbTable::new( @@ -346,7 +458,10 @@ mod tests { ])); let json = row_json(&row, &table); assert_eq!(json["meta"]["k"][1], 2); - assert_eq!(json["at"], 1_000); + assert_eq!( + json["at"], 1_000, + "the tree keeps the engine's microseconds" + ); assert_eq!(json["extra"], 7, "every column the row carries is written"); let mut bytes = Vec::new(); @@ -354,14 +469,50 @@ mod tests { let written: Json = serde_json::from_slice(&bytes).expect("valid JSON"); assert_eq!(written["op"], "put"); assert_eq!(written["tableName"], "t"); - assert_eq!(written["value"], json, "the bytes say what the tree says"); + let mut expected = json.clone(); + expected["at"] = Json::from(1); + assert_eq!( + written["value"], expected, + "the bytes say what the tree says, times in milliseconds" + ); let mut bytes = Vec::new(); write_del( &mut bytes, "t", &DataFrameKey::new([("id", Value::from("a\"b"))]), + Some(&table), ); let written: Json = serde_json::from_slice(&bytes).expect("valid JSON"); assert_eq!(written["id"]["id"], "a\"b", "keys are escaped"); + + let timed = DbTable::new( + "timed", + ["at"], + vec![ + DbColumn::new("at", ValueType::Timestamp), + DbColumn::new("v", ValueType::Int), + ], + ); + let mut bytes = Vec::new(); + write_del( + &mut bytes, + "timed", + &DataFrameKey::new([("at", Value::Int(1_000))]), + Some(&timed), + ); + let written: Json = serde_json::from_slice(&bytes).expect("valid JSON"); + assert_eq!(written["id"]["at"], 1); + let mut bytes = Vec::new(); + write_del( + &mut bytes, + "timed", + &DataFrameKey::new([("at", Value::Int(1_000))]), + None, + ); + let written: Json = serde_json::from_slice(&bytes).expect("valid JSON"); + assert_eq!( + written["id"]["at"], 1_000, + "without the schema the number is left alone" + ); } } diff --git a/src/model/schema.rs b/src/model/schema.rs index cffad16..aabd8d5 100644 --- a/src/model/schema.rs +++ b/src/model/schema.rs @@ -333,6 +333,14 @@ fn json_string(out: &mut Vec, text: &str) { #[derive(Debug, Clone, Default)] pub struct Catalog { tables: HashMap, + /// Type OIDs of the database's enums and their array types: a bound + /// parameter of one of these carries its label as text. + enum_oids: std::collections::HashSet, + /// Text-carried columns Postgres orders as the engine compares their + /// text, byte by byte: a `uuid`, a `bytea`'s hex, or a string under + /// the `C` collation. Any other text-carried column (an enum, a + /// string under a linguistic collation, an `inet`) orders otherwise. + ordered_as_text: std::collections::HashSet<(TableName, ColumnName)>, } impl Catalog { @@ -343,9 +351,63 @@ impl Catalog { .into_iter() .map(|table| (table.name.clone(), table)) .collect(), + enum_oids: std::collections::HashSet::new(), + ordered_as_text: std::collections::HashSet::new(), } } + /// The same catalog knowing these type OIDs are enums (or arrays of + /// enums). + pub fn with_enum_oids(mut self, oids: impl IntoIterator) -> Self { + self.enum_oids = oids.into_iter().collect(); + self + } + + /// Whether a bound parameter of type `oid` is an enum label or an + /// array of them, which the wire carries as text. + pub fn is_enum_oid(&self, oid: u32) -> bool { + self.enum_oids.contains(&oid) + } + + /// The enum type OIDs this catalog knows, for a catalog built from + /// this one to carry. + pub fn enum_oids(&self) -> impl Iterator + '_ { + self.enum_oids.iter().copied() + } + + /// The same catalog knowing these text-carried columns order in + /// Postgres as their text does, byte by byte. + pub fn with_ordered_as_text( + mut self, + columns: impl IntoIterator, + ) -> Self { + self.ordered_as_text = columns.into_iter().collect(); + self + } + + /// Whether Postgres orders `table.column` as the engine compares it. + /// Numbers, booleans and times always do; a text-carried column only + /// when the catalog was told so; JSON and lists never do. + pub fn orders_as_engine(&self, table: &TableName, column: &ColumnName) -> bool { + match self + .table(table.as_str()) + .and_then(|held| held.column(column.as_str())) + .map(|declared| &declared.r#type) + { + Some(ValueType::String) => self + .ordered_as_text + .contains(&(table.clone(), column.clone())), + Some(ValueType::Json | ValueType::List(_)) | None => false, + Some(_) => true, + } + } + + /// The text-carried columns this catalog knows order as text, for a + /// catalog built from this one to carry. + pub fn ordered_as_text(&self) -> impl Iterator + '_ { + self.ordered_as_text.iter() + } + /// The schema of table `name`, if the catalog contains it. pub fn table(&self, name: &str) -> Option<&DbTable> { self.tables.get(name) diff --git a/src/model/value.rs b/src/model/value.rs index da6cbe0..270e06a 100644 --- a/src/model/value.rs +++ b/src/model/value.rs @@ -48,9 +48,12 @@ pub enum ValueType { Datetime, List(Box), Map(Box, Box), - /// A point in time carried as milliseconds since the Unix epoch, the - /// form protocol clients see for `timestamp`, `timestamptz` and `date` - /// columns; the cell is a [`Value::Int`]. + /// A point in time kept as microseconds since the Unix epoch, which + /// is Postgres's own precision for `timestamp`, `timestamptz` and + /// `date` columns; the cell is a [`Value::Int`]. The Zero protocol + /// sees milliseconds: its arguments are scaled up on the way in and + /// its rows scaled down on the way out; a server-to-server client + /// gets the microseconds as they are. Timestamp, /// A JSON document carried as its text, the form the engine keeps for /// `json` and `jsonb` columns; the cell is a [`Value::String`] and the diff --git a/src/parser/binds.rs b/src/parser/binds.rs new file mode 100644 index 0000000..63b2e79 --- /dev/null +++ b/src/parser/binds.rs @@ -0,0 +1,520 @@ +//! Bind parameters as a Postgres client sends them, decoded into values. +//! +//! An ORM renders a query with `$n` placeholders and ships each parameter +//! in Postgres's binary format with its type OID. The engine's parser +//! accepts those placeholders ([`super::parse_select`]) and reads +//! the parameter here, by OID, into the [`Value`] the literal would have +//! parsed to; the column's declared type then coerces it the way it +//! coerces a literal, so `$1` and `'OPEN'` take the same path. +//! +//! Only the types a predicate compares on are decoded: integers, floats, +//! numerics, booleans, text-like types, UUIDs, JSON, timestamps and dates. +//! Anything else is refused by OID, never guessed at. + +use crate::model::Value; +use crate::sync::pg::text::json_as_jsonb; + +/// One bind parameter: its type OID and its bytes, `None` for `NULL`. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct BindValue { + pub type_oid: u32, + pub bytes: Option>, +} + +/// Postgres's built-in type OIDs that a predicate can carry. +pub(crate) mod oid { + pub const BOOL: u32 = 16; + pub const BYTEA: u32 = 17; + pub const CHAR: u32 = 18; + pub const NAME: u32 = 19; + pub const INT8: u32 = 20; + pub const INT2: u32 = 21; + pub const INT4: u32 = 23; + pub const TEXT: u32 = 25; + pub const OID: u32 = 26; + pub const JSON: u32 = 114; + pub const FLOAT4: u32 = 700; + pub const FLOAT8: u32 = 701; + pub const BPCHAR: u32 = 1042; + pub const VARCHAR: u32 = 1043; + pub const DATE: u32 = 1082; + pub const TIMESTAMP: u32 = 1114; + pub const TIMESTAMPTZ: u32 = 1184; + pub const NUMERIC: u32 = 1700; + pub const UUID: u32 = 2950; + pub const JSONB: u32 = 3802; + + /// Array types, by element: what `= ANY($1)` binds. + pub const BOOL_ARRAY: u32 = 1000; + pub const INT2_ARRAY: u32 = 1005; + pub const INT4_ARRAY: u32 = 1007; + pub const INT8_ARRAY: u32 = 1016; + pub const TEXT_ARRAY: u32 = 1009; + pub const VARCHAR_ARRAY: u32 = 1015; + pub const BPCHAR_ARRAY: u32 = 1014; + pub const FLOAT4_ARRAY: u32 = 1021; + pub const FLOAT8_ARRAY: u32 = 1022; + pub const NUMERIC_ARRAY: u32 = 1231; + pub const DATE_ARRAY: u32 = 1182; + pub const TIMESTAMP_ARRAY: u32 = 1115; + pub const TIMESTAMPTZ_ARRAY: u32 = 1185; + pub const UUID_ARRAY: u32 = 2951; + pub const JSONB_ARRAY: u32 = 3807; +} + +/// Microseconds from the Unix epoch to Postgres's epoch (2000-01-01). +const PG_EPOCH_MICROS: i64 = 946_684_800_000_000; + +/// Whether an enum bind's bytes are an array rather than one label: an +/// enum type's array OID is as much an enum OID as the scalar's, and the +/// two are told apart by the bytes. A one-dimensional array header is +/// 20 bytes with a dimension count of 0 or 1; a label is text, which +/// never starts with three zero bytes. +pub fn looks_like_array(bytes: &[u8]) -> bool { + bytes.len() >= 12 && bytes[..3] == [0, 0, 0] && bytes[3] <= 1 +} + +/// The value a bind stands for, or why it cannot stand for one. +/// +/// Binary `jsonb` is a version byte followed by the text. A time is +/// microseconds since 2000-01-01, and the engine keeps microseconds +/// since the Unix epoch; `infinity` is `i64::MAX` and `-infinity` +/// `i64::MIN` (for a date, `i32`'s), which fit no shift and which no +/// cache can compare on, so they are refused. +pub fn decode(bind: &BindValue) -> Result { + let Some(bytes) = bind.bytes.as_deref() else { + return Ok(Value::Null); + }; + let bytes_of = |want: usize| -> Result<&[u8], String> { + (bytes.len() == want).then_some(bytes).ok_or_else(|| { + format!( + "a bind of type oid {} has {} bytes, not {want}", + bind.type_oid, + bytes.len() + ) + }) + }; + Ok(match bind.type_oid { + oid::BOOL => Value::Bool(bytes_of(1)?[0] != 0), + oid::INT2 => Value::Int(i16::from_be_bytes(bytes_of(2)?.try_into().unwrap()).into()), + oid::INT4 => Value::Int(i32::from_be_bytes(bytes_of(4)?.try_into().unwrap()).into()), + oid::OID => Value::Int(u32::from_be_bytes(bytes_of(4)?.try_into().unwrap()).into()), + oid::INT8 => Value::Int(i64::from_be_bytes(bytes_of(8)?.try_into().unwrap())), + oid::FLOAT4 => Value::Float(f32::from_be_bytes(bytes_of(4)?.try_into().unwrap()).into()), + oid::FLOAT8 => Value::Float(f64::from_be_bytes(bytes_of(8)?.try_into().unwrap())), + oid::NUMERIC => Value::Float(numeric(bytes)?), + oid::TEXT | oid::VARCHAR | oid::BPCHAR | oid::NAME | oid::CHAR => { + Value::String(utf8(bytes, bind.type_oid)?.to_owned()) + } + oid::JSON => Value::String(json_as_jsonb(utf8(bytes, bind.type_oid)?)), + oid::JSONB => { + let text = bytes + .split_first() + .filter(|(version, _)| **version == 1) + .map(|(_, text)| text) + .ok_or("a jsonb bind is not version 1")?; + Value::String(json_as_jsonb(utf8(text, bind.type_oid)?)) + } + oid::UUID => { + let bytes = bytes_of(16)?; + let hex: String = bytes.iter().map(|byte| format!("{byte:02x}")).collect(); + Value::String(format!( + "{}-{}-{}-{}-{}", + &hex[0..8], + &hex[8..12], + &hex[12..16], + &hex[16..20], + &hex[20..32] + )) + } + oid::TIMESTAMP | oid::TIMESTAMPTZ => { + let micros = i64::from_be_bytes(bytes_of(8)?.try_into().unwrap()); + if micros == i64::MAX || micros == i64::MIN { + return Err("an infinite time bind".to_string()); + } + Value::Int(micros.checked_add(PG_EPOCH_MICROS).ok_or("a time bind out of range")?) + } + oid::DATE => { + let days = i32::from_be_bytes(bytes_of(4)?.try_into().unwrap()); + if days == i32::MAX || days == i32::MIN { + return Err("an infinite date bind".to_string()); + } + Value::Int( + i64::from(days) + .checked_mul(86_400_000_000) + .and_then(|micros| micros.checked_add(PG_EPOCH_MICROS)) + .ok_or("a date bind out of range")?, + ) + } + oid::BYTEA => return Err("a bytea bind cannot be compared on".to_string()), + oid::BOOL_ARRAY + | oid::INT2_ARRAY + | oid::INT4_ARRAY + | oid::INT8_ARRAY + | oid::TEXT_ARRAY + | oid::VARCHAR_ARRAY + | oid::BPCHAR_ARRAY + | oid::FLOAT4_ARRAY + | oid::FLOAT8_ARRAY + | oid::NUMERIC_ARRAY + | oid::DATE_ARRAY + | oid::TIMESTAMP_ARRAY + | oid::TIMESTAMPTZ_ARRAY + | oid::UUID_ARRAY + | oid::JSONB_ARRAY => Value::List(array(bytes, element_of(bind.type_oid))?), + other => return Err(format!("a bind of type oid {other} is not supported")), + }) +} + +/// Whether [`decode`] knows type `oid`: a built-in it reads or refuses +/// by name. Any other OID is an enum, or a type the engine does not +/// carry. +pub fn decodes(type_oid: u32) -> bool { + use oid::*; + matches!( + type_oid, + BOOL | BYTEA + | CHAR + | NAME + | INT8 + | INT2 + | INT4 + | TEXT + | OID + | JSON + | FLOAT4 + | FLOAT8 + | BPCHAR + | VARCHAR + | DATE + | TIMESTAMP + | TIMESTAMPTZ + | NUMERIC + | UUID + | JSONB + | BOOL_ARRAY + | INT2_ARRAY + | INT4_ARRAY + | INT8_ARRAY + | TEXT_ARRAY + | VARCHAR_ARRAY + | BPCHAR_ARRAY + | FLOAT4_ARRAY + | FLOAT8_ARRAY + | NUMERIC_ARRAY + | DATE_ARRAY + | TIMESTAMP_ARRAY + | TIMESTAMPTZ_ARRAY + | UUID_ARRAY + | JSONB_ARRAY + ) +} + +/// The element type of an array type. The array's own header names an +/// element type too, but that is what the client declared, which for +/// an enum remapped to text is the enum's OID; the outer type decides. +fn element_of(array_oid: u32) -> u32 { + match array_oid { + oid::BOOL_ARRAY => oid::BOOL, + oid::INT2_ARRAY => oid::INT2, + oid::INT4_ARRAY => oid::INT4, + oid::INT8_ARRAY => oid::INT8, + oid::TEXT_ARRAY => oid::TEXT, + oid::VARCHAR_ARRAY => oid::VARCHAR, + oid::BPCHAR_ARRAY => oid::BPCHAR, + oid::FLOAT4_ARRAY => oid::FLOAT4, + oid::FLOAT8_ARRAY => oid::FLOAT8, + oid::NUMERIC_ARRAY => oid::NUMERIC, + oid::DATE_ARRAY => oid::DATE, + oid::TIMESTAMP_ARRAY => oid::TIMESTAMP, + oid::TIMESTAMPTZ_ARRAY => oid::TIMESTAMPTZ, + oid::UUID_ARRAY => oid::UUID, + oid::JSONB_ARRAY => oid::JSONB, + other => other, + } +} + +/// Postgres's binary array: dimensions, a has-nulls flag, the element +/// type, then per dimension a length and lower bound, then each element +/// as its length (`-1` for NULL) and bytes. Only one dimension is a list +/// a predicate can use. Elements are decoded as `element_oid`. The +/// declared length is the client's word, so the list grows with the +/// elements actually present (each takes at least its four-byte length) +/// rather than being sized from it up front. +fn array(bytes: &[u8], element_oid: u32) -> Result, String> { + let word = |at: usize| -> Result { + bytes + .get(at..at + 4) + .map(|b| i32::from_be_bytes(b.try_into().unwrap())) + .ok_or_else(|| "an array bind ends inside its header".to_string()) + }; + let dimensions = word(0)?; + if dimensions == 0 { + return Ok(Vec::new()); + } + if dimensions != 1 { + return Err(format!("an array bind has {dimensions} dimensions, not one")); + } + let length = usize::try_from(word(12)?).map_err(|_| "an array bind has a negative length")?; + let mut at = 20; + let mut items = Vec::with_capacity(length.min(bytes.len().saturating_sub(20) / 4)); + for _ in 0..length { + let size = word(at)?; + at += 4; + if size < 0 { + items.push(Value::Null); + continue; + } + let size = size as usize; + let element = bytes + .get(at..at + size) + .ok_or_else(|| "an array bind ends inside an element".to_string())?; + items.push(decode(&BindValue { + type_oid: element_oid, + bytes: Some(element.to_vec()), + })?); + at += size; + } + Ok(items) +} + +/// The bytes as text, or why they are not. +fn utf8(bytes: &[u8], type_oid: u32) -> Result<&str, String> { + std::str::from_utf8(bytes).map_err(|_| format!("a bind of type oid {type_oid} is not UTF-8")) +} + +/// Postgres's binary `numeric`: a header of base-10000 digit count, +/// weight, sign and scale, then the digits. +/// +/// The digits are rendered as text and the float parser does the +/// rounding, which keeps every representable value exact. `weight` is +/// the power of 10000 of the first group: zero or more means that many +/// integer groups follow it, negative means every group is fractional +/// and `-weight - 1` empty groups precede the first. +fn numeric(bytes: &[u8]) -> Result { + if bytes.len() < 8 { + return Err("a numeric bind is shorter than its header".to_string()); + } + let word = |at: usize| u16::from_be_bytes([bytes[at], bytes[at + 1]]); + let ndigits = usize::from(word(0)); + let weight = i32::from(word(2) as i16); + let sign = word(4); + let scale = usize::from(word(6)); + if bytes.len() != 8 + 2 * ndigits { + return Err("a numeric bind's digits do not match its header".to_string()); + } + let negative = match sign { + 0x0000 => false, + 0x4000 => true, + 0xC000 => return Ok(f64::NAN), + _ => return Err("a numeric bind has an unknown sign".to_string()), + }; + let group = |index: usize| word(8 + 2 * index); + let mut text = String::new(); + if ndigits == 0 { + text.push('0'); + } else if weight < 0 { + text.push_str("0."); + text.push_str(&"0000".repeat((-weight - 1) as usize)); + for index in 0..ndigits { + text.push_str(&format!("{:04}", group(index))); + } + } else { + let integer_groups = weight as usize + 1; + for index in 0..ndigits { + if index == integer_groups { + text.push('.'); + } + if index == 0 { + text.push_str(&group(index).to_string()); + } else { + text.push_str(&format!("{:04}", group(index))); + } + } + if ndigits < integer_groups { + text.push_str(&"0000".repeat(integer_groups - ndigits)); + } + } + let _ = scale; + let value: f64 = text + .parse() + .map_err(|_| format!("a numeric bind rendered as `{text}` is not a number"))?; + Ok(if negative { -value } else { value }) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// A bind of `type_oid` holding `bytes`. + fn bind(type_oid: u32, bytes: &[u8]) -> BindValue { + BindValue { + type_oid, + bytes: Some(bytes.to_vec()), + } + } + + /// Integers, floats and booleans take their Postgres widths. + #[test] + fn integers_floats_and_bools() { + assert_eq!(decode(&bind(oid::INT2, &(-7i16).to_be_bytes())), Ok(Value::Int(-7))); + assert_eq!(decode(&bind(oid::INT4, &42i32.to_be_bytes())), Ok(Value::Int(42))); + assert_eq!( + decode(&bind(oid::INT8, &(1i64 << 40).to_be_bytes())), + Ok(Value::Int(1 << 40)) + ); + assert_eq!(decode(&bind(oid::FLOAT8, &2.5f64.to_be_bytes())), Ok(Value::Float(2.5))); + assert_eq!(decode(&bind(oid::FLOAT4, &0.5f32.to_be_bytes())), Ok(Value::Float(0.5))); + assert_eq!(decode(&bind(oid::BOOL, &[1])), Ok(Value::Bool(true))); + assert_eq!(decode(&bind(oid::BOOL, &[0])), Ok(Value::Bool(false))); + } + + /// Text-like types are UTF-8, and a UUID is its hyphenated hex. + #[test] + fn text_like_types_and_uuids() { + assert_eq!( + decode(&bind(oid::TEXT, b"OPEN")), + Ok(Value::String("OPEN".into())) + ); + assert_eq!( + decode(&bind(oid::VARCHAR, "m\u{e9}nu".as_bytes())), + Ok(Value::String("m\u{e9}nu".into())) + ); + let uuid = [ + 0x12, 0x34, 0x56, 0x78, 0x9a, 0xbc, 0xde, 0xf0, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, + 0x77, 0x88, + ]; + assert_eq!( + decode(&bind(oid::UUID, &uuid)), + Ok(Value::String("12345678-9abc-def0-1122-334455667788".into())) + ); + assert!(decode(&bind(oid::TEXT, &[0xff])).is_err()); + } + + /// JSON is normalised as a `jsonb` cell is; a `jsonb` bind needs its + /// version byte. + #[test] + fn json_is_normalised_like_a_jsonb_cell() { + assert_eq!( + decode(&bind(oid::JSON, br#"{"b": 1, "a": [true]}"#)), + Ok(Value::String(json_as_jsonb(r#"{"b": 1, "a": [true]}"#))) + ); + let mut jsonb = vec![1u8]; + jsonb.extend_from_slice(br#"{"k":"v"}"#); + assert_eq!( + decode(&bind(oid::JSONB, &jsonb)), + Ok(Value::String(json_as_jsonb(r#"{"k":"v"}"#))) + ); + assert!(decode(&bind(oid::JSONB, b"{}")).is_err(), "no version byte"); + } + + /// Times and dates become microseconds since the Unix epoch: + /// 2000-01-01 00:00:01.5 is 1,500,000 Postgres microseconds. + #[test] + fn times_become_unix_microseconds() { + let micros = 1_500_000i64; + assert_eq!( + decode(&bind(oid::TIMESTAMPTZ, µs.to_be_bytes())), + Ok(Value::Int(PG_EPOCH_MICROS + 1_500_000)) + ); + assert_eq!( + decode(&bind(oid::TIMESTAMP, &(-1_000i64).to_be_bytes())), + Ok(Value::Int(PG_EPOCH_MICROS - 1_000)) + ); + assert_eq!( + decode(&bind(oid::DATE, &2i32.to_be_bytes())), + Ok(Value::Int(PG_EPOCH_MICROS + 2 * 86_400_000_000)) + ); + } + + /// `numeric` digit groups render to the value they spell: + /// `12345.678` is digits `[1, 2345, 6780]`, weight 1, positive, + /// scale 3; `-0.0042` is `[42]`, weight -1, negative, scale 4; `0` + /// has no digits; `70000` is `[7]`, weight 1. + #[test] + fn numerics_render_through_their_digit_groups() { + let numeric = |words: &[u16]| -> Vec { + words.iter().flat_map(|word| word.to_be_bytes()).collect() + }; + for (words, value) in [ + (&[3u16, 1, 0, 3, 1, 2345, 6780][..], 12345.678), + (&[1, 0xffff, 0x4000, 4, 42][..], -0.0042), + (&[0, 0, 0, 0][..], 0.0), + (&[1, 1, 0, 0, 7][..], 70000.0), + ] { + assert_eq!(decode(&bind(oid::NUMERIC, &numeric(words))), Ok(Value::Float(value))); + } + } + + /// Arrays become lists: a `text[]` of `["OPEN", NULL, "DONE"]`, an + /// `int8[]` of `[7]`, and an empty array, which has zero dimensions. + /// A header cut short is refused. + #[test] + fn arrays_become_lists() { + let header = |words: &[i32]| -> Vec { + words.iter().flat_map(|word| word.to_be_bytes()).collect() + }; + let mut bytes = header(&[1, 1, 25, 3, 1]); + bytes.extend_from_slice(&4i32.to_be_bytes()); + bytes.extend_from_slice(b"OPEN"); + bytes.extend_from_slice(&(-1i32).to_be_bytes()); + bytes.extend_from_slice(&4i32.to_be_bytes()); + bytes.extend_from_slice(b"DONE"); + assert_eq!( + decode(&bind(oid::TEXT_ARRAY, &bytes)), + Ok(Value::List(vec![ + Value::String("OPEN".into()), + Value::Null, + Value::String("DONE".into()) + ])) + ); + let mut bytes = header(&[1, 0, 20, 1, 1, 8]); + bytes.extend_from_slice(&7i64.to_be_bytes()); + assert_eq!(decode(&bind(oid::INT8_ARRAY, &bytes)), Ok(Value::List(vec![Value::Int(7)]))); + assert_eq!( + decode(&bind(oid::TEXT_ARRAY, &header(&[0, 0, 25]))), + Ok(Value::List(vec![])) + ); + assert!(decode(&bind(oid::TEXT_ARRAY, &[0, 0, 0, 2])).is_err()); + } + + /// An array header is the client's word: one claiming two billion + /// elements in twenty bytes is refused without sizing anything from + /// the claim. + #[test] + fn a_huge_declared_length_is_refused_not_allocated() { + let mut bytes = Vec::new(); + for word in [1i32, 0, 20, i32::MAX, 1] { + bytes.extend_from_slice(&word.to_be_bytes()); + } + assert_eq!( + decode(&bind(oid::INT8_ARRAY, &bytes)), + Err("an array bind ends inside its header".to_string()) + ); + } + + /// An `oid` is unsigned: one past `i32::MAX` stays positive. + #[test] + fn an_oid_is_unsigned() { + assert_eq!( + decode(&bind(oid::OID, &3_000_000_000u32.to_be_bytes())), + Ok(Value::Int(3_000_000_000)) + ); + } + + /// A missing value is `NULL`; an unknown type, a `bytea` and a wrong + /// width are refused. + #[test] + fn nulls_and_unknown_types() { + assert_eq!( + decode(&BindValue { + type_oid: oid::TEXT, + bytes: None + }), + Ok(Value::Null) + ); + assert!(decode(&bind(600, b"x")).is_err(), "point"); + assert!(decode(&bind(oid::BYTEA, b"x")).is_err()); + assert!(decode(&bind(oid::INT4, &[1, 2])).is_err(), "wrong width"); + } +} diff --git a/src/parser/mod.rs b/src/parser/mod.rs index 4f24049..2fb5c84 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -8,32 +8,46 @@ //! which columns form the primary key. //! //! ```text -//! query := read | write -//! read := SELECT '*' FROM ident [WHERE filter] -//! [ORDER BY column [ASC | DESC] (, column [ASC | DESC])*] [LIMIT integer] -//! write := INSERT INTO ident '(' ident (',' ident)* ')' VALUES row -//! | UPDATE ident SET assign (',' assign)* WHERE filter -//! | DELETE FROM ident WHERE filter -//! row := '(' value (',' value)* ')' -//! assign := ident '=' value -//! filter := conjunct (OR conjunct)* -//! conjunct := primary (AND primary)* -//! primary := '(' filter ')' | condition | TRUE | FALSE -//! condition := column ('=' | '!=' | '<>' | '>' | '>=' | '<' | '<=') value -//! | column [NOT] IN '(' [value (',' value)*] ')' -//! | column IS [NOT] NULL -//! column := ident -//! value := NULL | TRUE | FALSE | number | string -//! | '[' [value (',' value)*] ']' +//! query := read | write +//! read := SELECT projection FROM from [WHERE filter] +//! [ORDER BY column [ASC | DESC] (, column [ASC | DESC])*] +//! [LIMIT count] [OFFSET count] +//! projection := '*' | COUNT '(' '*' ')' [AS ident] +//! | column [AS ident] (',' column [AS ident])* +//! from := table | '(' from ')' +//! | from [INNER | LEFT [OUTER]] JOIN from ON filter +//! count := integer | bind +//! write := INSERT INTO ident '(' ident (',' ident)* ')' VALUES row +//! | UPDATE ident SET assign (',' assign)* WHERE filter +//! | DELETE FROM ident WHERE filter +//! row := '(' value (',' value)* ')' +//! assign := ident '=' value +//! filter := conjunct (OR conjunct)* +//! conjunct := primary (AND primary)* +//! primary := '(' filter ')' | condition | TRUE | FALSE +//! condition := column ('=' | '!=' | '<>' | '>' | '>=' | '<' | '<=') value +//! | column [NOT] IN '(' [value (',' value)*] ')' +//! | column '=' ANY '(' value ')' | column '!=' ALL '(' value ')' +//! | column IS [NOT] NULL +//! column := ident | ident '.' ident +//! ident := a bare word | a double-quoted name, `""` for a quote +//! value := NULL | TRUE | FALSE | number | string | bind +//! | '[' [value (',' value)*] ']' +//! bind := '$' integer //! ``` //! -//! Keywords are case-insensitive; identifiers are bare words; strings are -//! single-quoted with `''` for an embedded quote; `TRUE` / `FALSE` as a -//! filter parse to the vacuous `AND(vec![])` / `OR(vec![])`. +//! Keywords are case-insensitive; strings are single-quoted with `''` for +//! an embedded quote; `TRUE` / `FALSE` as a filter parse to the vacuous +//! `AND(vec![])` / `OR(vec![])`. A bind is read from the parameters an +//! ORM sends with the statement ([`parse_select`], [`binds`]); `= ANY` +//! and `!= ALL` take an array, as ORMs render an `IN` list. //! -//! # v1 restrictions (each rejected with a clear error, never narrowed silently) +//! # Restrictions (each rejected with a clear error, never narrowed silently) //! -//! - Single table only — `LEFT JOIN` is recognised and refused. +//! - Joins, projections, `COUNT(*)` and `OFFSET` are read by +//! [`parse_select`], the server-to-server entry point, under the rules +//! in [`select`]; [`parse`] and [`parse_read`] take one table and +//! refuse a join or an `OFFSET`. //! - Writes address **exactly one row by primary key**: the `WHERE` of an //! `UPDATE` / `DELETE` must be a conjunction of `pkey_column = value` //! conditions covering the whole primary key. Predicate writes ("set all @@ -59,7 +73,7 @@ use std::collections::{HashMap, HashSet}; use std::fmt; use crate::model::{ - ColumnName, ComparisonOperator, Condition, DataFrameKey, DataFrameRow, DbTable, DeleteQuery, + ColumnName, TableName, ComparisonOperator, Condition, DataFrameKey, DataFrameRow, DbTable, DeleteQuery, InsertQuery, Order, OrderBy, SingleTableReadQuery, UpdateQuery, Value, ValueType, Where, WriteQuery, }; @@ -73,6 +87,11 @@ pub const MAX_NESTING_DEPTH: usize = 128; /// because the parser is their primary consumer. pub use crate::model::Catalog; +pub mod binds; +mod select; +pub use binds::BindValue; +pub use select::Select; + /// A successfully parsed statement — either side of the read/write split. #[derive(Debug, Clone, PartialEq)] pub enum ParsedQuery { @@ -82,10 +101,24 @@ pub enum ParsedQuery { /// Parse one statement of either kind. pub fn parse(sql: &str, catalog: &Catalog) -> Result { + parse_with_binds(sql, &[], catalog) +} + +/// Parse a statement whose `$n` placeholders stand for `binds`, numbered +/// from `$1` as Postgres numbers them. Each is decoded by its type OID +/// ([`binds::decode`]) and then coerced to its column's declared type +/// exactly as a literal is. +pub fn parse_with_binds( + sql: &str, + binds: &[BindValue], + catalog: &Catalog, +) -> Result { let mut parser = Parser { tokens: lex(sql)?, pos: 0, depth: 0, + binds, + times: TimeUnit::Millis, }; let parsed = parser.query(catalog)?; if !matches!(parser.peek(), Tok::Eof) { @@ -96,7 +129,89 @@ pub fn parse(sql: &str, catalog: &Catalog) -> Result { /// Parse a statement that must be a `SELECT`. pub fn parse_read(sql: &str, catalog: &Catalog) -> Result { - match parse(sql, catalog)? { + parse_read_with_binds(sql, &[], catalog) +} + +/// What a `SELECT` asks for: everything, a count of the rows, or named +/// columns in order. The engine follows whole rows whatever is selected; +/// a client that ships rows uses this to shape them. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum Projection { + All, + /// `COUNT(*)`, under the output name it was given (`count` when none). + Count { name: String }, + Columns(Vec), +} + +/// One projected column, the table it is of, and the name it comes out +/// under: its alias, or the column's own name. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Projected { + pub table: TableName, + pub column: ColumnName, + pub name: String, +} + +/// A `SELECT` over one table or joined tables, with `$n` placeholders +/// standing for `binds`: the engine's query tree and the projection. +/// An `ORDER BY` on a column Postgres orders differently from the +/// engine (an enum, a string under a linguistic collation, JSON) is +/// refused: the engine would keep, and the backend return, other rows +/// in another order ([`Catalog::orders_as_engine`]). An enum label +/// travels in the binary format as its text, under the enum's own OID, +/// which the catalog names; an array of enums likewise, under the +/// array's OID. +pub fn parse_select(sql: &str, binds: &[BindValue], catalog: &Catalog) -> Result { + let binds: Vec = binds + .iter() + .map(|bind| match (catalog.is_enum_oid(bind.type_oid), bind.bytes.as_deref()) { + (true, Some(bytes)) if binds::looks_like_array(bytes) => BindValue { + type_oid: binds::oid::TEXT_ARRAY, + bytes: bind.bytes.clone(), + }, + (true, _) => BindValue { + type_oid: binds::oid::TEXT, + bytes: bind.bytes.clone(), + }, + (false, _) => bind.clone(), + }) + .collect(); + let mut parser = Parser { + tokens: lex(sql)?, + pos: 0, + depth: 0, + binds: &binds, + times: TimeUnit::Micros, + }; + if !parser.is_word("SELECT") { + return parser.err(format!("expected SELECT, found {}", parser.describe())); + } + let select = parser.select(catalog)?; + if !matches!(parser.peek(), Tok::Eof) { + return parser.err(format!("unexpected trailing input: {}", parser.describe())); + } + if let Some((table, column, position)) = select + .ordered_by + .iter() + .find(|(table, column, _)| !catalog.orders_as_engine(table, column)) + { + return Err(ParseError { + message: format!( + "ORDER BY `{table}`.`{column}`: Postgres orders this column by rules the engine does not follow (an enum, a collation other than C, JSON)" + ), + position: *position, + }); + } + Ok(select) +} + +/// [`parse_read`] with `$n` placeholders standing for `binds`. +pub fn parse_read_with_binds( + sql: &str, + binds: &[BindValue], + catalog: &Catalog, +) -> Result { + match parse_with_binds(sql, binds, catalog)? { ParsedQuery::Read(read) => Ok(read), ParsedQuery::Write(_) => Err(ParseError { message: "expected a SELECT, found a write statement".into(), @@ -151,6 +266,8 @@ pub fn point_at(source: &str, error: &ParseError) -> String { /// One lexed token. /// /// - `Word`: identifiers and keywords alike; keywords match case-insensitively. +/// - `Quoted`: a double-quoted identifier, never a keyword, case kept. +/// - `Bind`: a `$n` placeholder, numbered from 1. /// - `Str`: the unquoted content of a single-quoted string literal. /// - `Num`: raw numeric text; the parser decides int vs float. /// - `Sym`: operators and punctuation, normalised (`<>` lexes as `!=`). @@ -158,6 +275,8 @@ pub fn point_at(source: &str, error: &ParseError) -> String { #[derive(Debug, Clone, PartialEq)] enum Tok { Word(String), + Quoted(String), + Bind(usize), Str(String), Num(String), Sym(&'static str), @@ -208,14 +327,15 @@ fn lex(source: &str) -> Result, ParseError> { tok: Tok::Num(number), at, }); - } else if c == '\'' { + } else if c == '\'' || c == '"' { + let quote = c; chars.next(); let mut text = String::new(); loop { match chars.next() { - Some((_, '\'')) => { - if peek_at_is(&chars, 0, '\'') { - text.push('\''); + Some((_, c)) if c == quote => { + if peek_at_is(&chars, 0, quote) { + text.push(quote); chars.next(); } else { break; @@ -224,16 +344,41 @@ fn lex(source: &str) -> Result, ParseError> { Some((_, c)) => text.push(c), None => { return Err(ParseError { - message: "unterminated string literal".into(), + message: if quote == '"' { + "unterminated quoted identifier".into() + } else { + "unterminated string literal".into() + }, position: at, }); } } } tokens.push(Spanned { - tok: Tok::Str(text), + tok: if quote == '"' { + Tok::Quoted(text) + } else { + Tok::Str(text) + }, at, }); + } else if c == '$' { + chars.next(); + let mut number = String::new(); + take_digits(&mut chars, &mut number); + let index = number.parse::().ok().filter(|index| *index > 0); + match index { + Some(index) => tokens.push(Spanned { + tok: Tok::Bind(index), + at, + }), + None => { + return Err(ParseError { + message: "expected a parameter number after `$`".into(), + position: at, + }); + } + } } else { chars.next(); let sym: &'static str = match c { @@ -242,6 +387,7 @@ fn lex(source: &str) -> Result, ParseError> { '[' => "[", ']' => "]", ',' => ",", + '.' => ".", '*' => "*", '-' => "-", '=' => "=", @@ -329,13 +475,28 @@ fn next_is_digit(chars: &Chars<'_>, ahead: usize) -> bool { /// - `tokens`: the lexed input, always ending with [`Tok::Eof`]. /// - `pos`: index of the current token. /// - `depth`: current `(` / `[` nesting depth, bounded by [`MAX_NESTING_DEPTH`]. -struct Parser { +/// - `binds`: what `$n` placeholders stand for, `$1` first. +struct Parser<'b> { tokens: Vec, pos: usize, depth: usize, + binds: &'b [BindValue], + /// What a number compared with a time column means, set by the + /// entry point: [`parse`] and [`parse_read`] (SQL by hand, in the + /// Zero protocol's terms) read it as milliseconds; [`parse_select`], + /// the server-to-server path, reads every number there, literal or + /// bound, as microseconds, Postgres's own unit and the engine's. + times: TimeUnit, +} + +/// The unit of a number standing for a time, see [`Parser::times`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum TimeUnit { + Millis, + Micros, } -impl Parser { +impl Parser<'_> { /// The current token. fn peek(&self) -> &Tok { &self.tokens[self.pos].tok @@ -379,6 +540,8 @@ impl Parser { fn describe(&self) -> String { match self.peek() { Tok::Word(word) => format!("`{word}`"), + Tok::Quoted(word) => format!("`\"{word}\"`"), + Tok::Bind(index) => format!("parameter `${index}`"), Tok::Str(text) => format!("string `{text}`"), Tok::Num(number) => format!("number `{number}`"), Tok::Sym(symbol) => format!("`{symbol}`"), @@ -435,7 +598,7 @@ impl Parser { /// A bare-word identifier. fn ident(&mut self) -> Result { match self.peek().clone() { - Tok::Word(word) => { + Tok::Word(word) | Tok::Quoted(word) => { self.bump(); Ok(word) } @@ -443,20 +606,20 @@ impl Parser { } } - /// A table name, resolved against the catalog. - fn table<'c>(&mut self, catalog: &'c Catalog) -> Result<&'c DbTable, ParseError> { - let at = self.at(); - let name = self.ident()?; - catalog.table(&name).ok_or(ParseError { - message: format!("unknown table `{name}`"), - position: at, - }) - } - - /// A column name, validated to exist on `table`. + /// A column, bare or qualified by its table (`t.c`, `"t"."c"`); a + /// qualifier must name the table being read. fn column_name(&mut self, table: &DbTable) -> Result { let at = self.at(); - let name = self.ident()?; + let mut name = self.ident()?; + if self.eat_sym(".") { + if name != table.name.as_str() { + return Err(ParseError { + message: format!("`{name}` is not the table being read, `{}`", table.name), + position: at, + }); + } + name = self.ident()?; + } if table.column(&name).is_some() { Ok(ColumnName::from(name)) } else { @@ -467,10 +630,29 @@ impl Parser { } } - /// Dispatch on the leading keyword to a read or a write. + /// A table name, resolved against the catalog. + fn table<'c>(&mut self, catalog: &'c Catalog) -> Result<&'c DbTable, ParseError> { + let at = self.at(); + let name = self.ident()?; + catalog.table(&name).ok_or(ParseError { + message: format!("unknown table `{name}`"), + position: at, + }) + } + + /// Dispatch on the leading keyword to a read or a write. A join or + /// an `OFFSET` is refused here: only a server-to-server client, + /// through [`parse_select`], skips rows itself. fn query(&mut self, catalog: &Catalog) -> Result { if self.is_word("SELECT") { - Ok(ParsedQuery::Read(self.read(catalog)?)) + let select = self.select(catalog)?; + if !select.query.joins.is_empty() { + return self.err("a join parses through `parse_select`, not `parse`"); + } + if select.offset > 0 { + return self.err("OFFSET is not supported"); + } + Ok(ParsedQuery::Read(select.query.main_table)) } else if self.is_word("INSERT") || self.is_word("UPDATE") || self.is_word("DELETE") { Ok(ParsedQuery::Write(self.write(catalog)?)) } else { @@ -481,80 +663,44 @@ impl Parser { } } - /// `SELECT * FROM …` with optional WHERE / ORDER BY / LIMIT. - /// - /// A missing WHERE subscribes to the whole table (the vacuously true - /// `AND(vec![])`). A missing ORDER BY defaults via [`default_order`], - /// and a missing LIMIT becomes `u32::MAX` — the model makes `limit` - /// mandatory, so absent means "no bound". - fn read(&mut self, catalog: &Catalog) -> Result { - self.expect_word("SELECT")?; - self.expect_sym("*")?; - self.expect_word("FROM")?; - let table = self.table(catalog)?; - - if self.is_word("LEFT") || self.is_word("JOIN") { - return self.err("joins are not supported yet — single-table queries only"); - } - - let filter = if self.eat_word("WHERE") { - self.filter(table)? - } else { - Where::AND(Vec::new()) - }; - - let order_by = if self.eat_word("ORDER") { - self.expect_word("BY")?; - let mut clauses = Vec::new(); - loop { - let column = self.column_name(table)?; - let direction = if self.eat_word("DESC") { - Order::DESC - } else { - self.eat_word("ASC"); - Order::ASC - }; - clauses.push(OrderBy::new(column, direction)); - if !self.eat_sym(",") { - break; + /// Advances past the projection without reading it, up to `FROM`. + fn skip_projection(&mut self) -> Result<(), ParseError> { + let mut depth = 0usize; + loop { + match self.peek() { + Tok::Eof => return self.err("expected FROM"), + Tok::Word(word) if depth == 0 && word.eq_ignore_ascii_case("FROM") => { + return Ok(()); } + Tok::Sym("(") => depth += 1, + Tok::Sym(")") => depth = depth.saturating_sub(1), + _ => {} } - clauses - } else { - vec![default_order(table, self.at())?] - }; + self.bump(); + } + } - let limit = if self.eat_word("LIMIT") { - match self.peek().clone() { - Tok::Num(number) => { - let at = self.at(); - self.bump(); - number.parse::().map_err(|_| ParseError { - message: format!("`{number}` is not a valid limit"), - position: at, - })? - } - _ => return self.err(format!("expected a limit, found {}", self.describe())), - } - } else { - u32::MAX - }; + /// An optional `AS name` or bare `name` after a projected column. + fn alias(&mut self) -> Result, ParseError> { + if self.eat_word("AS") { + return self.ident().map(Some); + } + if !self.is_word("FROM") && matches!(self.peek(), Tok::Word(_) | Tok::Quoted(_)) { + return self.ident().map(Some); + } + Ok(None) + } - Ok(SingleTableReadQuery::new( - table.name.clone(), - filter, - order_by, - limit, - )) + /// `FROM` is next. + fn expect_from(&mut self) -> Result<(), ParseError> { + if self.is_word("FROM") { + Ok(()) + } else { + self.err(format!("expected FROM after the projection, found {}", self.describe())) + } } - /// `INSERT` / `UPDATE` / `DELETE`, enforcing the v1 write restrictions. - /// - /// INSERT coerces each value while it still lines up with a named - /// column; a value/column count mismatch is reported only after the - /// whole row is read. UPDATE must SET every non-pkey column because - /// the engine evaluates predicates against `record` as the complete - /// new row image — a partial SET would corrupt the views. + /// `INSERT`, `UPDATE` or `DELETE` to the end of the statement. fn write(&mut self, catalog: &Catalog) -> Result { if self.eat_word("INSERT") { self.expect_word("INTO")?; @@ -577,7 +723,7 @@ impl Parser { let value_at = self.at(); let value = self.value()?; values.push(match columns.get(values.len()) { - Some(column) => coerce_to_column_type(value, table, column, value_at)?, + Some(column) => coerce_to_column_type(value, table, column, value_at, self.times)?, None => value, }); if !self.eat_sym(",") { @@ -623,7 +769,7 @@ impl Parser { self.expect_sym("=")?; let value_at = self.at(); let value = self.value()?; - let value = coerce_to_column_type(value, table, &column, value_at)?; + let value = coerce_to_column_type(value, table, &column, value_at, self.times)?; assignments.push((column, value)); if !self.eat_sym(",") { break; @@ -719,7 +865,7 @@ impl Parser { position: at, }); } - let value = coerce_to_column_type(value, table, &column, at)?; + let value = coerce_to_column_type(value, table, &column, at, self.times)?; if pkey_value.insert(column.clone(), value).is_some() { return Err(ParseError { message: format!("primary-key column `{column}` appears twice in WHERE"), @@ -753,7 +899,8 @@ impl Parser { /// A parenthesised filter, a `TRUE` / `FALSE` literal, or one condition. /// /// `TRUE` / `FALSE` parse to the identities: an empty `AND` is true, an - /// empty `OR` is false. + /// empty `OR` is false. `= ANY($1)` with an array parameter is how + /// ORMs render an `IN` list: one bind, many values. fn primary(&mut self, table: &DbTable) -> Result { if self.eat_sym("(") { self.descend()?; @@ -789,8 +936,18 @@ impl Parser { })); } let comparison_operator = if self.eat_sym("=") { + if self.eat_word("ANY") { + return Ok(Where::Condition(Condition { + column, + comparison_operator: ComparisonOperator::IN, + value: self.quantified_list("ANY")?, + })); + } ComparisonOperator::EQ } else if self.eat_sym("!=") { + if self.eat_word("ALL") { + return Ok(differs_from_all(column, self.quantified_list("ALL")?)); + } ComparisonOperator::NEQ } else if self.eat_sym(">=") { ComparisonOperator::GTE @@ -826,6 +983,22 @@ impl Parser { })) } + /// The `(value)` after `ANY` / `ALL`: a list literal or a bind that + /// decodes to one. + fn quantified_list(&mut self, quantifier: &str) -> Result { + let at = self.at(); + self.expect_sym("(")?; + let value = self.value()?; + self.expect_sym(")")?; + match value { + Value::List(_) => Ok(value), + other => Err(ParseError { + message: format!("{quantifier} needs an array, found {other:?}"), + position: at, + }), + } + } + /// The parenthesised value list of `IN` / `NOT IN`. fn value_list(&mut self) -> Result { self.expect_sym("(")?; @@ -886,6 +1059,20 @@ impl Parser { self.bump(); Ok(Value::String(text)) } + Tok::Bind(index) => { + self.bump(); + let bind = self.binds.get(index - 1).ok_or(ParseError { + message: format!( + "`${index}` has no parameter; {} were given", + self.binds.len() + ), + position: at, + })?; + binds::decode(bind).map_err(|message| ParseError { + message: format!("`${index}`: {message}"), + position: at, + }) + } Tok::Word(word) if word.eq_ignore_ascii_case("NULL") => { self.bump(); Ok(Value::Null) @@ -1003,12 +1190,13 @@ fn coerce_to_column_type( table: &DbTable, column: &ColumnName, at: usize, + times: TimeUnit, ) -> Result { let declared = table .column(column.as_str()) .map(|candidate| &candidate.r#type) .expect("column names are validated against the table before coercion"); - coerce_to_type(value, declared, column.as_str(), at) + coerce_to_type(value, declared, column.as_str(), at, times) } /// Recursive worker for [`coerce_to_column_type`], matching a value (and @@ -1020,6 +1208,7 @@ fn coerce_to_type( declared: &ValueType, column: &str, at: usize, + times: TimeUnit, ) -> Result { Ok(match (value, declared) { (Value::Null, _) => Value::Null, @@ -1027,7 +1216,10 @@ fn coerce_to_type( (value @ Value::Float(_), ValueType::Float) => value, (value @ Value::String(_), ValueType::String) => value, (value @ Value::Bool(_), ValueType::Bool) => value, - (value @ Value::Int(_), ValueType::Timestamp) => value, + (Value::Int(int), ValueType::Timestamp) => Value::Int(match times { + TimeUnit::Millis => int.saturating_mul(1000), + TimeUnit::Micros => int, + }), (value @ Value::String(_), ValueType::Json) => value, (Value::Int(int), ValueType::Float) => Value::Float(int as f64), (Value::Float(float), ValueType::Int) @@ -1038,7 +1230,7 @@ fn coerce_to_type( (Value::List(items), ValueType::List(inner)) => Value::List( items .into_iter() - .map(|item| coerce_to_type(item, inner, column, at)) + .map(|item| coerce_to_type(item, inner, column, at, times)) .collect::>()?, ), (Value::String(_), ValueType::Date | ValueType::Datetime) => { @@ -1097,13 +1289,593 @@ fn collect_pkey_equalities(filter: &Where, out: &mut Vec<(ColumnName, Value)>) - } } +/// `column != ALL(list)`. Postgres holds it for every row when the list +/// is empty, a `NULL` in the column included, which is no condition at +/// all; the engine's `NOT IN` leaves `NULL` out, as its `NULL` rule +/// does everywhere, so an empty list must not become one. +pub(super) fn differs_from_all(column: ColumnName, list: Value) -> Where { + match &list { + Value::List(items) if items.is_empty() => Where::AND(Vec::new()), + _ => Where::Condition(Condition { + column, + comparison_operator: ComparisonOperator::NOT_IN, + value: list, + }), + } +} + #[cfg(test)] mod tests { use super::*; use crate::model::{DbColumn, ValueType}; + /// A `text` bind. + fn text_bind(text: &str) -> BindValue { + BindValue { + type_oid: 25, + bytes: Some(text.as_bytes().to_vec()), + } + } + + /// An `int8` bind. + fn int8_bind(value: i64) -> BindValue { + BindValue { + type_oid: 20, + bytes: Some(value.to_be_bytes().to_vec()), + } + } + + /// An `int4` bind. + fn int4_bind(value: i32) -> BindValue { + BindValue { + type_oid: 23, + bytes: Some(value.to_be_bytes().to_vec()), + } + } + + /// What Diesel renders for + /// `tickets.filter(status.eq("OPEN")).filter(points.gt(3)).order(id.desc()).first()` + /// parses as its literal twin does. + #[test] + fn a_diesel_query_parses_like_its_literal_twin() { + let diesel = r#"SELECT "tickets"."id", "tickets"."status", "tickets"."priority", "tickets"."points" FROM "tickets" WHERE (("tickets"."status" = $1) AND ("tickets"."points" > $2)) ORDER BY "tickets"."id" DESC LIMIT $3"#; + let bound = parse_read_with_binds( + diesel, + &[text_bind("OPEN"), int4_bind(3), int8_bind(1)], + &catalog(), + ) + .expect("diesel shape parses"); + let literal = parse_read( + "SELECT * FROM tickets WHERE status = 'OPEN' AND points > 3 ORDER BY id DESC LIMIT 1", + &catalog(), + ) + .expect("literal parses"); + assert_eq!(bound, literal); + } + + /// `parse_select` says what is selected: every column, a count under + /// its name, or the columns under their aliases. + #[test] + fn parse_select_tells_what_is_selected() { + let all = parse_select("SELECT * FROM tickets", &[], &catalog()).unwrap().projection; + assert_eq!(all, Projection::All); + let count = parse_select(r#"SELECT COUNT(*) FROM "tickets""#, &[], &catalog()).unwrap().projection; + assert_eq!(count, Projection::Count { name: "count".into() }); + let count = parse_select(r#"SELECT COUNT(*) AS n FROM "tickets""#, &[], &catalog()).unwrap().projection; + assert_eq!(count, Projection::Count { name: "n".into() }); + let Select { query: q, projection: columns, .. } = parse_select( + r#"SELECT "tickets"."points" AS p, "tickets"."id" FROM "tickets" LIMIT $1"#, + &[int8_bind(5)], + &catalog(), + ) + .unwrap(); + assert_eq!( + columns, + Projection::Columns(vec![ + Projected { + table: TableName::from("tickets"), + column: ColumnName::from("points"), + name: "p".into() + }, + Projected { + table: TableName::from("tickets"), + column: ColumnName::from("id"), + name: "id".into() + }, + ]) + ); + assert_eq!(q.main_table.limit, 5); + assert!(parse_select("UPDATE tickets SET status = 'x' WHERE id = 1", &[], &catalog()).is_err()); + } + + /// A projection may alias its columns, with or without `AS`; a + /// literal where a column list belongs is refused. + #[test] + fn projected_columns_may_be_aliased() { + let aliased = parse_read( + r#"SELECT "tickets"."status" AS s, "tickets"."points" p FROM "tickets" WHERE "tickets"."id" = 1"#, + &catalog(), + ) + .expect("aliases parse"); + let plain = parse_read("SELECT * FROM tickets WHERE id = 1", &catalog()).expect("plain"); + assert_eq!(aliased, plain); + let junk = parse_read(r#"SELECT "tickets"."status" 'x' FROM "tickets""#, &catalog()) + .expect_err("a literal where a column list belongs"); + assert!(junk.message.contains("expected FROM"), "{}", junk.message); + } + + /// Authors, their books, and the books' reviews. + fn join_catalog() -> Catalog { + Catalog::new(vec![ + DbTable::new( + "authors", + ["id"], + vec![ + DbColumn::new("id", ValueType::Int), + DbColumn::new("name", ValueType::String), + ], + ), + DbTable::new( + "books", + ["id"], + vec![ + DbColumn::new("id", ValueType::Int), + DbColumn::new("author_id", ValueType::Int), + DbColumn::new("title", ValueType::String), + ], + ), + DbTable::new( + "reviews", + ["id"], + vec![ + DbColumn::new("id", ValueType::Int), + DbColumn::new("book_id", ValueType::Int), + DbColumn::new("stars", ValueType::Int), + ], + ), + ]) + } + + /// `authors.inner_join(books).order((authors::id, books::id))`: the + /// backend composes books under each author by their key, so the + /// clause is met and dropped from the tree's order. A joined table + /// orders only by its key, only ascending, only after the root's + /// columns. + #[test] + fn a_joined_tables_primary_key_may_order_it() { + let sql = r#"SELECT "authors"."id", "books"."title" FROM ("authors" INNER JOIN "books" ON ("books"."author_id" = "authors"."id")) WHERE ("authors"."id" = $1) ORDER BY "authors"."id" ASC, "books"."id" ASC"#; + let select = parse_select(sql, &[int8_bind(1)], &join_catalog()).expect("ordered by the child's key"); + let order: Vec<&str> = select.query.main_table.order_by.iter().map(|clause| clause.column.as_str()).collect(); + assert_eq!(order, vec!["id"], "only the root's order reaches the tree"); + for (sql, expect) in [ + ( + r#"SELECT "authors"."id" FROM ("authors" INNER JOIN "books" ON ("books"."author_id" = "authors"."id")) ORDER BY "books"."title""#, + "primary key", + ), + ( + r#"SELECT "authors"."id" FROM ("authors" INNER JOIN "books" ON ("books"."author_id" = "authors"."id")) ORDER BY "books"."id" DESC"#, + "primary key", + ), + ( + r#"SELECT "authors"."id" FROM ("authors" INNER JOIN "books" ON ("books"."author_id" = "authors"."id")) ORDER BY "books"."id", "authors"."id""#, + "must come before", + ), + ] { + let error = parse_select(sql, &[], &join_catalog()).expect_err(sql); + assert!(error.message.contains(expect), "{}", error.message); + } + } + + /// The composition orders each root row's result rows together and + /// the joined tables depth first by key: an ORDER BY is accepted + /// only where it asks for that, so a key alone, a root order that + /// leaves ties, or siblings named out of order are refused. + #[test] + fn a_joined_order_must_be_the_composition_order() { + let catalog = join_catalog(); + let books_then_reviews = r#"SELECT * FROM (("authors" INNER JOIN "books" ON ("books"."author_id" = "authors"."id")) INNER JOIN "reviews" ON ("reviews"."book_id" = "books"."id"))"#; + let order = |clauses: &str| format!("{books_then_reviews} ORDER BY {clauses}"); + parse_select( + &order(r#""authors"."id", "books"."id", "reviews"."id""#), + &[], + &catalog, + ) + .expect("the composition order"); + for (clauses, expect) in [ + (r#""books"."id""#, "whole primary key"), + (r#""authors"."name", "books"."id""#, "whole primary key"), + (r#""authors"."id", "reviews"."id""#, "in the order the tables are joined"), + (r#""authors"."id", "reviews"."id", "books"."id""#, "in the order the tables are joined"), + ] { + let error = parse_select(&order(clauses), &[], &catalog).expect_err(clauses); + assert!(error.message.contains(expect), "{clauses}: {}", error.message); + } + } + + /// An INNER JOIN written after a LEFT JOIN onto its table would drop + /// the rows the LEFT JOIN keeps; inside its parentheses it does not. + /// A WHERE on any table under a LEFT JOIN is refused, a grandchild + /// included. + #[test] + fn an_inner_join_under_a_left_join_must_be_nested() { + let catalog = join_catalog(); + let flat = r#"SELECT * FROM (("authors" LEFT OUTER JOIN "books" ON ("books"."author_id" = "authors"."id")) INNER JOIN "reviews" ON ("reviews"."book_id" = "books"."id"))"#; + let error = parse_select(flat, &[], &catalog).expect_err("flat"); + assert!(error.message.contains("under a LEFT JOIN"), "{}", error.message); + let nested = r#"SELECT * FROM ("authors" LEFT OUTER JOIN ("books" INNER JOIN "reviews" ON ("reviews"."book_id" = "books"."id")) ON ("books"."author_id" = "authors"."id"))"#; + parse_select(nested, &[], &catalog).expect("nested"); + let error = parse_select(&format!(r#"{nested} WHERE "reviews"."stars" = 5"#), &[], &catalog) + .expect_err("a grandchild under the LEFT JOIN"); + assert!(error.message.contains("LEFT JOIN inner"), "{}", error.message); + } + + /// A COUNT counts every row: its LIMIT and OFFSET cut the one row it + /// returns, not the rows it counts. + #[test] + fn a_count_is_not_cut_by_its_limit() { + let select = parse_select( + r#"SELECT COUNT(*) FROM "tickets" LIMIT $1 OFFSET $2"#, + &[int8_bind(1), int8_bind(1)], + &catalog(), + ) + .expect("a count with a limit"); + assert_eq!(select.query.main_table.limit, u32::MAX); + assert_eq!((select.limit, select.offset), (1, 1)); + } + + /// An ORDER BY on a text column is followed only where Postgres + /// orders it as its text, which the catalog says per column. + #[test] + fn an_order_on_text_needs_a_byte_ordered_column() { + let sql = r#"SELECT * FROM "tickets" ORDER BY "tickets"."status" LIMIT 5"#; + let error = parse_select(sql, &[], &catalog()).expect_err("an unknown collation"); + assert!(error.message.contains("rules the engine does not follow"), "{}", error.message); + let known = catalog().with_ordered_as_text([(TableName::from("tickets"), ColumnName::from("status"))]); + parse_select(sql, &[], &known).expect("a C-collated column"); + parse_select(r#"SELECT * FROM "tickets" ORDER BY "tickets"."points" LIMIT 5"#, &[], &catalog()) + .expect("a number"); + } + + /// Diesel's joins become the engine's tree: + /// `authors.inner_join(books).filter(authors::name.eq("x")).order(authors::id.desc()).limit(5)`; + /// `authors.left_join(books)`; `authors.inner_join(books.inner_join(reviews))` + /// in nested parentheses; an explicit `ON` with a condition on the + /// joined table, which filters it; and an inner join's `WHERE` + /// conjuncts, which go to the table they name. + #[test] + fn diesel_joins_become_the_engines_tree() { + use crate::model::{Driver, Join}; + let sql = r#"SELECT "authors"."id", "books"."title" FROM ("authors" INNER JOIN "books" ON ("books"."author_id" = "authors"."id")) WHERE ("authors"."name" = $1) ORDER BY "authors"."id" DESC LIMIT $2"#; + let select = parse_select(sql, &[text_bind("x"), int8_bind(5)], &join_catalog()).expect("inner join"); + let root = &select.query.main_table; + assert_eq!(root.table.as_str(), "authors"); + assert_eq!(root.limit, 5); + assert_eq!(root.order_by, vec![OrderBy::new("id", Order::DESC)]); + assert_eq!( + root.filter, + Where::Condition(Condition { + column: ColumnName::from("name"), + comparison_operator: ComparisonOperator::EQ, + value: Value::String("x".into()), + }) + ); + assert_eq!(select.query.joins.len(), 1); + let join: &Join = &select.query.joins[0]; + assert_eq!(join.sub.main_table.table.as_str(), "books"); + assert_eq!((join.main_table_column.as_str(), join.sub_table_column.as_str()), ("id", "author_id")); + assert!(join.is_inner); + assert_eq!(join.driver, Driver::Sub, "inner edges start sub-driven; the planner may flip"); + assert_eq!(join.sub.main_table.filter, Where::AND(vec![])); + assert_eq!(join.sub.main_table.limit, u32::MAX); + assert_eq!( + select.projection, + Projection::Columns(vec![ + Projected { table: TableName::from("authors"), column: ColumnName::from("id"), name: "id".into() }, + Projected { table: TableName::from("books"), column: ColumnName::from("title"), name: "title".into() }, + ]) + ); + + let left = parse_select( + r#"SELECT "authors"."name", "books"."title" FROM ("authors" LEFT OUTER JOIN "books" ON ("books"."author_id" = "authors"."id"))"#, + &[], + &join_catalog(), + ) + .expect("left join"); + assert!(!left.query.joins[0].is_inner); + assert_eq!(left.query.joins[0].driver, Driver::Main); + + let three = parse_select( + r#"SELECT "authors"."name", "reviews"."stars" FROM ("authors" INNER JOIN ("books" INNER JOIN "reviews" ON ("reviews"."book_id" = "books"."id")) ON ("books"."author_id" = "authors"."id"))"#, + &[], + &join_catalog(), + ) + .expect("nested join"); + let books = &three.query.joins[0]; + assert_eq!(books.sub.main_table.table.as_str(), "books"); + let reviews = &books.sub.joins[0]; + assert_eq!(reviews.sub.main_table.table.as_str(), "reviews"); + assert_eq!((reviews.main_table_column.as_str(), reviews.sub_table_column.as_str()), ("id", "book_id")); + + let explicit = parse_select( + r#"SELECT "authors"."id" FROM ("authors" INNER JOIN "books" ON (("books"."author_id" = "authors"."id") AND ("books"."title" = $1)))"#, + &[text_bind("t")], + &join_catalog(), + ) + .expect("explicit on"); + assert_eq!( + explicit.query.joins[0].sub.main_table.filter, + Where::Condition(Condition { + column: ColumnName::from("title"), + comparison_operator: ComparisonOperator::EQ, + value: Value::String("t".into()), + }) + ); + + let split = parse_select( + r#"SELECT "authors"."id" FROM ("authors" INNER JOIN "books" ON ("books"."author_id" = "authors"."id")) WHERE (("books"."title" = $1) AND ("authors"."name" = $2))"#, + &[text_bind("t"), text_bind("a")], + &join_catalog(), + ) + .expect("split where"); + assert!(matches!(split.query.main_table.filter, Where::Condition(Condition { ref column, .. }) if column.as_str() == "name")); + assert!(matches!(split.query.joins[0].sub.main_table.filter, Where::Condition(Condition { ref column, .. }) if column.as_str() == "title")); + } + + /// Each join shape the engine cannot follow is refused with its + /// reason. + #[test] + fn join_shapes_the_engine_cannot_follow_are_named() { + let catalog = join_catalog(); + let refuse = |sql: &str, want: &str| { + let error = parse_select(sql, &[text_bind("x")], &catalog).expect_err(sql); + assert!(error.message.contains(want), "{sql}: {}", error.message); + }; + refuse( + r#"SELECT * FROM ("authors" LEFT OUTER JOIN "books" ON ("books"."author_id" = "authors"."id")) WHERE "books"."title" = $1"#, + "would turn its LEFT JOIN inner", + ); + refuse( + r#"SELECT * FROM ("authors" INNER JOIN "books" ON ("books"."author_id" = "authors"."id")) ORDER BY "books"."title""#, + "on a joined table", + ); + refuse( + r#"SELECT * FROM ("authors" INNER JOIN "books" ON ("books"."author_id" = "authors"."id")) WHERE ("books"."title" = $1 OR "authors"."name" = $1)"#, + "OR across tables", + ); + refuse( + r#"SELECT * FROM ("authors" RIGHT JOIN "books" ON ("books"."author_id" = "authors"."id"))"#, + "only INNER and LEFT", + ); + refuse( + r#"SELECT * FROM ("authors" INNER JOIN "books" ON ("books"."title" = $1))"#, + "needs an equality", + ); + refuse( + r#"SELECT "name" FROM ("authors" INNER JOIN "books" ON ("books"."author_id" = "authors"."id"))"#, + "must be qualified", + ); + refuse( + r#"SELECT * FROM ("authors" INNER JOIN "books" ON (("books"."author_id" = "authors"."id") AND ("books"."id" = "authors"."id")))"#, + "compound join keys", + ); + } + + /// A `text[]` of `["OPEN", "DONE"]`, as Diesel binds `eq_any`, makes + /// `= ANY` an `IN` list. + #[test] + fn eq_any_with_an_array_bind_is_an_in_list() { + let mut bytes = Vec::new(); + for word in [1i32, 0, 25, 2, 1] { + bytes.extend_from_slice(&word.to_be_bytes()); + } + for item in ["OPEN", "DONE"] { + bytes.extend_from_slice(&(item.len() as i32).to_be_bytes()); + bytes.extend_from_slice(item.as_bytes()); + } + let array = BindValue { + type_oid: 1009, + bytes: Some(bytes), + }; + let any = parse_read_with_binds( + r#"SELECT * FROM "tickets" WHERE ("tickets"."status" = ANY($1))"#, + std::slice::from_ref(&array), + &catalog(), + ) + .expect("= ANY parses"); + let literal = parse_read( + "SELECT * FROM tickets WHERE status IN ('OPEN', 'DONE')", + &catalog(), + ) + .unwrap(); + assert_eq!(any, literal); + let all = parse_read_with_binds( + r#"SELECT * FROM "tickets" WHERE ("tickets"."status" != ALL($1))"#, + &[array], + &catalog(), + ) + .expect("!= ALL parses"); + let literal = parse_read( + "SELECT * FROM tickets WHERE status NOT IN ('OPEN', 'DONE')", + &catalog(), + ) + .unwrap(); + assert_eq!(all, literal); + let scalar = parse_read_with_binds( + r#"SELECT * FROM "tickets" WHERE "tickets"."status" = ANY($1)"#, + &[text_bind("OPEN")], + &catalog(), + ) + .expect_err("a scalar after ANY"); + assert!(scalar.message.contains("needs an array"), "{}", scalar.message); + } + + /// `!= ALL` of an empty array holds for every row in Postgres, `NULL` + /// included, so it is no condition; `= ANY` of one holds for none. + #[test] + fn an_empty_array_under_all_is_no_condition() { + let mut bytes = Vec::new(); + for word in [0i32, 0, 25] { + bytes.extend_from_slice(&word.to_be_bytes()); + } + let empty = BindValue { + type_oid: 1009, + bytes: Some(bytes), + }; + let all = parse_read_with_binds( + r#"SELECT * FROM "tickets" WHERE ("tickets"."status" != ALL($1))"#, + std::slice::from_ref(&empty), + &catalog(), + ) + .expect("!= ALL parses"); + assert_eq!(all.filter, Where::AND(Vec::new())); + let any = parse_read_with_binds( + r#"SELECT * FROM "tickets" WHERE ("tickets"."status" = ANY($1))"#, + &[empty], + &catalog(), + ) + .expect("= ANY parses"); + assert_eq!( + any.filter, + Where::condition("status", ComparisonOperator::IN, Value::List(Vec::new())) + ); + } + + /// A `COUNT(*)` with an `IN` list of binds reads as the list of + /// literals does. + #[test] + fn a_count_and_an_in_list_of_binds() { + let count = parse_read_with_binds( + r#"SELECT COUNT(*) FROM "tickets" WHERE ("tickets"."status" IN ($1, $2))"#, + &[text_bind("OPEN"), text_bind("DONE")], + &catalog(), + ) + .expect("count parses"); + let literal = parse_read( + "SELECT * FROM tickets WHERE status IN ('OPEN', 'DONE')", + &catalog(), + ) + .expect("literal parses"); + assert_eq!(count, literal); + } + + /// Reads keep a literal as spelled and compare by equality folding; a + /// bind is decoded by its OID and kept the same way, `NULL` included. + #[test] + fn a_bind_takes_the_value_its_literal_would() { + let q = parse_read_with_binds( + r#"SELECT * FROM "tickets" WHERE "tickets"."points" = $1"#, + &[int8_bind(7)], + &catalog(), + ) + .expect("int8 bind"); + assert_eq!( + q.filter, + Where::Condition(Condition { + column: ColumnName::from("points"), + comparison_operator: ComparisonOperator::EQ, + value: Value::Int(7), + }) + ); + let null = parse_read_with_binds( + r#"SELECT * FROM "tickets" WHERE "tickets"."status" = $1"#, + &[BindValue { + type_oid: 25, + bytes: None, + }], + &catalog(), + ) + .expect("a NULL bind"); + assert!(matches!( + null.filter, + Where::Condition(Condition { + value: Value::Null, + .. + }) + )); + } + + /// A bind or a projection that cannot be read is refused with its + /// reason. + #[test] + fn bind_and_projection_mistakes_are_named() { + let missing = parse_read_with_binds( + r#"SELECT * FROM "tickets" WHERE "tickets"."status" = $2"#, + &[text_bind("OPEN")], + &catalog(), + ) + .expect_err("$2 with one bind"); + assert!(missing.message.contains("`$2` has no parameter"), "{}", missing.message); + + let other_table = parse_read( + r#"SELECT "users"."id" FROM "tickets""#, + &catalog(), + ) + .expect_err("a column of another table"); + assert!(other_table.message.contains("not a table of this query"), "{}", other_table.message); + + let unknown = parse_read(r#"SELECT "tickets"."owner" FROM "tickets""#, &catalog()) + .expect_err("an unknown column"); + assert!(unknown.message.contains("unknown column `owner`"), "{}", unknown.message); + + let offset = parse_read(r#"SELECT * FROM "tickets" LIMIT 10 OFFSET 5"#, &catalog()) + .expect_err("offset on the single-table API"); + assert!(offset.message.contains("OFFSET"), "{}", offset.message); + + let unsupported = parse_read_with_binds( + r#"SELECT * FROM "tickets" WHERE "tickets"."status" = $1"#, + &[BindValue { + type_oid: 600, + bytes: Some(vec![0; 16]), + }], + &catalog(), + ) + .expect_err("a point bind"); + assert!(unsupported.message.contains("oid 600"), "{}", unsupported.message); + + let null_limit = parse_read_with_binds( + r#"SELECT * FROM "tickets" LIMIT $1"#, + &[BindValue { + type_oid: 20, + bytes: None, + }], + &catalog(), + ) + .expect_err("a NULL limit"); + assert!(null_limit.message.contains("limit parameter"), "{}", null_limit.message); + } + + /// A quoted identifier keeps its case, spaces and doubled quotes. + #[test] + fn quoted_identifiers_keep_their_case_and_escapes() { + let catalog = Catalog::new(vec![DbTable::new( + "Mixed Case", + ["id"], + vec![ + DbColumn::new("id", ValueType::Int), + DbColumn::new("a\"b", ValueType::String), + ], + )]); + let q = parse_read( + r#"SELECT "Mixed Case"."a""b" FROM "Mixed Case" WHERE "a""b" = 'x'"#, + &catalog, + ) + .expect("quoted names parse"); + assert_eq!(q.table.as_str(), "Mixed Case"); + assert_eq!( + q.filter, + Where::Condition(Condition { + column: ColumnName::from("a\"b"), + comparison_operator: ComparisonOperator::EQ, + value: Value::String("x".into()), + }) + ); + let error = parse_read(r#"SELECT "unterminated FROM t"#, &catalog) + .expect_err("unterminated"); + assert!(error.message.contains("quoted identifier")); + } + /// A one-table catalog (`tickets`, pkey `id`) shared by most tests. - fn catalog() -> Catalog { + pub(super) fn catalog() -> Catalog { Catalog::new(vec![DbTable::new( "tickets", ["id"], @@ -1348,7 +2120,7 @@ mod tests { let cases = [ ( "SELECT * FROM tickets LEFT JOIN tickets", - "joins are not supported", + "appears twice", ), ("SELECT * FROM nope", "unknown table"), ("SELECT * FROM tickets WHERE ghost = 1", "unknown column"), @@ -1554,3 +2326,143 @@ mod tests { assert_eq!(error.position, 30); } } + +#[cfg(test)] +mod s2s_shapes { + use super::*; + use crate::model::Value; + + /// The parser tests' catalog. + fn catalog() -> Catalog { + super::tests::catalog() + } + + /// A one-dimensional array of `labels`, no nulls, lower bound 1, + /// under element type `element`: the binary format Diesel sends. + fn label_array(element: i32, labels: &[&str]) -> Vec { + let mut bytes = Vec::new(); + for word in [1, 0, element, labels.len() as i32, 1] { + bytes.extend(word.to_be_bytes()); + } + for label in labels { + bytes.extend((label.len() as i32).to_be_bytes()); + bytes.extend(label.as_bytes()); + } + bytes + } + + /// An `OFFSET` widens the tree's limit to `limit + offset` and + /// travels beside it for the client to skip; with no `LIMIT` the + /// tree has none. + #[test] + fn an_offset_widens_the_limit_and_travels_separately() { + let select = parse_select( + r#"SELECT "tickets"."id" FROM "tickets" ORDER BY "tickets"."id" ASC LIMIT $1 OFFSET $2"#, + &[ + BindValue { + type_oid: 20, + bytes: Some(10i64.to_be_bytes().to_vec()), + }, + BindValue { + type_oid: 20, + bytes: Some(5i64.to_be_bytes().to_vec()), + }, + ], + &catalog(), + ) + .expect("offset"); + assert_eq!(select.query.main_table.limit, 15, "the engine keeps limit + offset rows"); + assert_eq!(select.offset, 5); + + let no_limit = parse_select(r#"SELECT * FROM "tickets" OFFSET 5"#, &[], &catalog()).expect("offset only"); + assert_eq!(no_limit.query.main_table.limit, u32::MAX); + assert_eq!(no_limit.offset, 5); + } + + /// An enum bind is read as its label once the catalog knows its OID, + /// and refused before. An array of labels under the enum's array OID + /// carries the enum's own OID as its element type, as Diesel sends + /// it. + #[test] + fn an_enum_bind_is_its_label() { + let catalog = catalog().with_enum_oids([777, 778]); + let select = parse_select( + r#"SELECT * FROM "tickets" WHERE "tickets"."status" = $1"#, + &[BindValue { + type_oid: 777, + bytes: Some(b"open".to_vec()), + }], + &catalog, + ) + .expect("an enum bind"); + let shown = format!("{:?}", select.query.main_table.filter); + assert!(shown.contains("open"), "{shown}"); + let unknown = parse_select( + r#"SELECT * FROM "tickets" WHERE "tickets"."status" = $1"#, + &[BindValue { + type_oid: 777, + bytes: Some(b"open".to_vec()), + }], + &super::tests::catalog(), + ) + .expect_err("an unknown oid"); + assert!(unknown.message.contains("oid 777"), "{}", unknown.message); + let bytes = label_array(777, &["open", "done"]); + let any = parse_select( + r#"SELECT * FROM "tickets" WHERE "tickets"."status" = ANY($1)"#, + &[BindValue { + type_oid: 778, + bytes: Some(bytes), + }], + &catalog, + ) + .expect("an enum array bind"); + let shown = format!("{:?}", any.query.main_table.filter); + assert!(shown.contains("open") && shown.contains("done"), "{shown}"); + let _ = Value::Null; + } + + /// An infinite time has no microsecond value: refused, not wrapped. + #[test] + fn an_infinite_time_bind_is_refused_not_wrapped() { + let catalog = catalog(); + for (oid, bytes) in [ + (binds::oid::TIMESTAMP, i64::MAX.to_be_bytes().to_vec()), + (binds::oid::TIMESTAMPTZ, i64::MIN.to_be_bytes().to_vec()), + (binds::oid::DATE, i32::MAX.to_be_bytes().to_vec()), + ] { + let error = parse_select( + r#"SELECT * FROM "tickets" WHERE "tickets"."points" < $1"#, + &[BindValue { + type_oid: oid, + bytes: Some(bytes), + }], + &catalog, + ) + .expect_err("infinity has no microsecond value"); + assert!(error.message.contains("infinite"), "{}", error.message); + } + } + + /// A `LIMIT` and `OFFSET` whose sum would reach the no-limit value + /// are refused rather than registering an unbounded query. + #[test] + fn a_limit_and_offset_that_saturate_are_refused() { + let catalog = catalog(); + let error = parse_select( + r#"SELECT * FROM "tickets" ORDER BY "tickets"."id" LIMIT 10 OFFSET 4294967290"#, + &[], + &catalog, + ) + .expect_err("would register an unbounded query"); + assert!(error.message.contains("more rows"), "{}", error.message); + let fine = parse_select( + r#"SELECT * FROM "tickets" ORDER BY "tickets"."id" LIMIT 10 OFFSET 5"#, + &[], + &catalog, + ) + .expect("a sane pair"); + assert_eq!(fine.query.main_table.limit, 15); + assert_eq!(fine.offset, 5); + } +} diff --git a/src/parser/select.rs b/src/parser/select.rs new file mode 100644 index 0000000..d942d03 --- /dev/null +++ b/src/parser/select.rs @@ -0,0 +1,730 @@ +//! `SELECT` over one table or a tree of joined tables, as an ORM renders +//! it, into the engine's [`MultiTableReadQuery`]. +//! +//! The `FROM` is read as a tree. Diesel writes +//! `FROM ("a" INNER JOIN ("b" INNER JOIN "c" ON (...)) ON (...))`: each +//! `JOIN` links the item on its right to one table on its left through an +//! `ON` equality of one column each side, and that becomes an edge from +//! the left table (the parent) to the right item's root (the child). +//! `INNER` (or a bare `JOIN`) is an inner edge, `LEFT [OUTER]` an outer +//! one. Other conditions in an `ON` go to the filter of the one table +//! they name. `WHERE` is split into its top-level conjuncts and each goes +//! to the one table it names. `ORDER BY` and `LIMIT` belong to the root. +//! +//! The backend composes the join under each root row, the joined tables +//! depth first, each in its primary-key order. An `ORDER BY` naming a +//! joined table is accepted when it asks for exactly that: the root's +//! clauses first, covering its whole primary key (so each root row's +//! result rows come out together, as Postgres orders them), then the +//! joined tables' key columns, ascending, in that same order. A +//! `COUNT(*)` counts every row; its `LIMIT` and `OFFSET` apply to the +//! one row it returns. +//! +//! What is refused, by name: a table appearing twice, `RIGHT`/`FULL`/ +//! `CROSS` joins, a condition naming two tables (other than the edge's +//! equality), an `OR` across tables, a `WHERE` condition on a table +//! under a `LEFT` join (it would turn the join inner), an `INNER JOIN` +//! onto a table under a `LEFT` join written after it rather than inside +//! it (Postgres then drops the rows the `LEFT JOIN` keeps), any other +//! `ORDER BY` on a joined table, an unqualified column in a join (Diesel +//! always qualifies), compound join keys. + +use super::{ + ParseError, Parser, Projected, Projection, Tok, collapse, default_order, differs_from_all, +}; +use crate::model::{ + Catalog, ColumnName, ComparisonOperator, Condition, DbTable, Join, MultiTableReadQuery, + Order, OrderBy, SingleTableReadQuery, TableName, Value, Where, +}; + +/// A parsed `SELECT`: the engine's query tree and what the statement +/// projects. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Select { + pub query: MultiTableReadQuery, + pub projection: Projection, + /// The statement's own `LIMIT` (`u32::MAX` for none). + pub limit: u32, + /// Rows the statement skips (`OFFSET`). The tree's limit is the + /// statement's limit plus this, so the engine keeps enough rows for + /// the client to skip and cut; under a `COUNT(*)` both apply to the + /// one counted row and the tree has no limit. + pub offset: u32, + /// Every column the `ORDER BY` names, with where it was named. + pub ordered_by: Vec<(TableName, ColumnName, usize)>, +} + +/// One table of the `FROM`, while the statement is read. +struct Node<'c> { + table: &'c DbTable, + /// Top-level conjuncts of this table's filter, from `ON` and `WHERE`. + conjuncts: Vec, + /// Edges to child nodes: (child, parent column, child column, inner). + edges: Vec<(usize, ColumnName, ColumnName, bool)>, + /// Whether this node is under a `LEFT` edge (its child, or below + /// one), which forbids a `WHERE` condition on it and an `INNER JOIN` + /// onto it written after the `LEFT JOIN`. + left_joined: bool, +} + +/// A column resolved to a node of the `FROM`. +struct Ref { + node: usize, + column: ColumnName, +} + +impl Parser<'_> { + /// `SELECT` to the end of the statement. The projection is read once + /// the `FROM` is known, so its columns can be checked; the engine + /// follows whole rows whatever is selected, so what it names is + /// validated and kept for the client. + pub(super) fn select(&mut self, catalog: &Catalog) -> Result { + self.expect_word("SELECT")?; + if self.is_word("DISTINCT") { + return self.err("DISTINCT is not supported"); + } + let projection_at = self.pos; + self.skip_projection()?; + self.expect_word("FROM")?; + let mut nodes: Vec> = Vec::new(); + let root = self.join_item(catalog, &mut nodes)?; + debug_assert_eq!(root, 0, "the leftmost table is the root"); + let after_from = self.pos; + self.pos = projection_at; + let projection = self.projection_of(&nodes)?; + self.pos = after_from; + + if self.eat_word("WHERE") { + let scope: Vec = (0..nodes.len()).collect(); + for (node, filter) in self.scoped_filter(&nodes, &scope)? { + if nodes[node].left_joined { + return self.err(format!( + "a WHERE condition on `{}` would turn its LEFT JOIN inner; put it in the ON", + nodes[node].table.name + )); + } + nodes[node].conjuncts.push(filter); + } + } + + let mut ordered_by = Vec::new(); + let order_by = if self.eat_word("ORDER") { + self.expect_word("BY")?; + self.order_by(&nodes, &mut ordered_by)? + } else { + vec![default_order(nodes[0].table, self.at())?] + }; + let limit = self.limit()?; + let offset = self.offset()?; + let kept = if matches!(projection, Projection::Count { .. }) || limit == u32::MAX { + u32::MAX + } else { + match limit.checked_add(offset) { + Some(kept) if kept < u32::MAX => kept, + _ => { + return Err(ParseError { + message: "LIMIT plus OFFSET is more rows than the engine follows".to_string(), + position: self.at(), + }); + } + } + }; + + let query = build(&nodes, 0, order_by, kept); + Ok(Select { + query, + projection, + limit, + offset, + ordered_by, + }) + } + + /// The clauses after `ORDER BY`, as the root's order. A joined + /// table's clause is accepted only where the composition already + /// yields it (see the module's notes): the root's clauses first, + /// naming its whole key, then the joined tables' key columns in + /// composition order, ascending. Every column named is noted in + /// `named`. + fn order_by( + &mut self, + nodes: &[Node<'_>], + named: &mut Vec<(TableName, ColumnName, usize)>, + ) -> Result, ParseError> { + let composed = composition(nodes); + let everywhere: Vec = (0..nodes.len()).collect(); + let mut clauses = Vec::new(); + let mut joined = 0; + loop { + let at = self.at(); + let column = self.column_ref(nodes, &everywhere)?; + let direction = if self.eat_word("DESC") { + Order::DESC + } else { + self.eat_word("ASC"); + Order::ASC + }; + named.push((nodes[column.node].table.name.clone(), column.column.clone(), at)); + if column.node == 0 { + if joined > 0 { + return Err(ParseError { + message: "ORDER BY: the first table's columns must come before a joined table's".to_string(), + position: at, + }); + } + clauses.push(OrderBy::new(column.column, direction)); + } else if direction == Order::ASC + && composed.get(joined) == Some(&(column.node, &column.column)) + { + joined += 1; + } else { + return Err(ParseError { + message: format!( + "ORDER BY `{}` is on a joined table; a joined table orders only by its primary key, ascending, in the order the tables are joined", + nodes[column.node].table.name + ), + position: at, + }); + } + if !self.eat_sym(",") { + break; + } + } + let root = nodes[0].table; + if joined > 0 + && !root + .pkey + .iter() + .all(|key| clauses.iter().any(|clause| clause.column == *key)) + { + return self.err(format!( + "ORDER BY on a joined table needs `{}`'s whole primary key ordered first", + root.name + )); + } + if clauses.is_empty() { + Ok(vec![default_order(root, self.at())?]) + } else { + Ok(clauses) + } + } + + /// `LIMIT n`, `LIMIT $k`, or nothing (`u32::MAX`). + fn limit(&mut self) -> Result { + if !self.eat_word("LIMIT") { + return Ok(u32::MAX); + } + self.count("limit") + } + + /// `OFFSET n`, `OFFSET $k`, or nothing (0). + fn offset(&mut self) -> Result { + if !self.eat_word("OFFSET") { + return Ok(0); + } + self.count("offset") + } + + /// A non-negative integer, literal or bound. + fn count(&mut self, what: &str) -> Result { + let at = self.at(); + match self.peek().clone() { + Tok::Num(number) => { + self.bump(); + number.parse::().map_err(|_| ParseError { + message: format!("`{number}` is not a valid {what}"), + position: at, + }) + } + Tok::Bind(_) => match self.value()? { + Value::Int(limit) if limit >= 0 && limit <= i64::from(u32::MAX) => Ok(limit as u32), + other => Err(ParseError { + message: format!( + "a {what} parameter must be a non-negative integer, found {other:?}" + ), + position: at, + }), + }, + _ => self.err(format!("expected a {what}, found {}", self.describe())), + } + } + + /// A `FROM` item: a table or a parenthesised item, followed by any + /// number of `JOIN`s. Returns the index of the item's root node. + fn join_item<'c>( + &mut self, + catalog: &'c Catalog, + nodes: &mut Vec>, + ) -> Result { + let first = nodes.len(); + let root = self.joined_table(catalog, nodes)?; + loop { + let at = self.at(); + let inner = if self.eat_word("INNER") { + self.expect_word("JOIN")?; + true + } else if self.eat_word("LEFT") { + self.eat_word("OUTER"); + self.expect_word("JOIN")?; + false + } else if self.eat_word("JOIN") { + true + } else if self.is_word("RIGHT") || self.is_word("FULL") || self.is_word("CROSS") { + return self.err("only INNER and LEFT joins are supported"); + } else { + return Ok(root); + }; + let left: Vec = (first..nodes.len()).collect(); + let right_root = self.joined_table(catalog, nodes)?; + let right: Vec = (right_root..nodes.len()).collect(); + self.expect_word("ON")?; + let on_at = self.at(); + let mut edge: Option<(usize, ColumnName, ColumnName)> = None; + let scope: Vec = (first..nodes.len()).collect(); + for term in self.on_terms(nodes, &scope)? { + match term { + OnTerm::Equality(a, b) => { + let (parent, child) = if left.contains(&a.node) && b.node == right_root { + (a, b) + } else if left.contains(&b.node) && a.node == right_root { + (b, a) + } else { + return Err(ParseError { + message: format!( + "a JOIN's ON must link `{}` to a table on its left", + nodes[right_root].table.name + ), + position: on_at, + }); + }; + if edge.is_some() { + return Err(ParseError { + message: "compound join keys are not supported".to_string(), + position: on_at, + }); + } + edge = Some((parent.node, parent.column, child.column)); + } + OnTerm::Condition(node, filter) => { + if !inner && !right.contains(&node) { + return Err(ParseError { + message: format!( + "a LEFT JOIN's ON may only filter `{}`", + nodes[right_root].table.name + ), + position: on_at, + }); + } + nodes[node].conjuncts.push(filter); + } + } + } + let Some((parent, parent_column, child_column)) = edge else { + return Err(ParseError { + message: "a JOIN's ON needs an equality of one column each side".to_string(), + position: at, + }); + }; + if inner && nodes[parent].left_joined { + return Err(ParseError { + message: format!( + "an INNER JOIN onto `{}`, which is under a LEFT JOIN, drops the rows the LEFT JOIN keeps; write it inside the LEFT JOIN's parentheses", + nodes[parent].table.name + ), + position: at, + }); + } + nodes[parent] + .edges + .push((right_root, parent_column, child_column, inner)); + if !inner { + for node in &right { + nodes[*node].left_joined = true; + } + } + } + } + + /// A table, or a parenthesised `FROM` item. + fn joined_table<'c>( + &mut self, + catalog: &'c Catalog, + nodes: &mut Vec>, + ) -> Result { + if self.eat_sym("(") { + self.descend()?; + let root = self.join_item(catalog, nodes)?; + self.expect_sym(")")?; + self.depth -= 1; + return Ok(root); + } + let at = self.at(); + let table = self.table(catalog)?; + if self.is_word("AS") { + return self.err("table aliases are not supported"); + } + if nodes.iter().any(|node| node.table.name == table.name) { + return Err(ParseError { + message: format!("`{}` appears twice; a table may be joined once", table.name), + position: at, + }); + } + nodes.push(Node { + table, + conjuncts: Vec::new(), + edges: Vec::new(), + left_joined: false, + }); + Ok(nodes.len() - 1) + } + + /// The terms of an `ON`: a parenthesised `AND` of equalities between + /// tables and conditions on one table. + fn on_terms(&mut self, nodes: &[Node<'_>], scope: &[usize]) -> Result, ParseError> { + let mut terms = Vec::new(); + self.on_term(nodes, scope, &mut terms)?; + Ok(terms) + } + + /// One term of an `ON`, appended to `terms`: in parentheses, either + /// one equality or a nested `AND` of terms; bare, an equality + /// between tables or a condition on one table. + fn on_term( + &mut self, + nodes: &[Node<'_>], + scope: &[usize], + terms: &mut Vec, + ) -> Result<(), ParseError> { + if self.eat_sym("(") { + self.descend()?; + let before = self.pos; + if let Ok(column) = self.column_ref(nodes, scope) + && self.eat_sym("=") + && let Ok(other) = self.column_ref(nodes, scope) + && self.is_sym(")") + && other.node != column.node + { + self.expect_sym(")")?; + self.depth -= 1; + terms.push(OnTerm::Equality(column, other)); + return Ok(()); + } + self.pos = before; + self.on_term(nodes, scope, terms)?; + while self.eat_word("AND") { + self.on_term(nodes, scope, terms)?; + } + self.expect_sym(")")?; + self.depth -= 1; + return Ok(()); + } + let before = self.pos; + if let Ok(column) = self.column_ref(nodes, scope) + && self.eat_sym("=") + && let Ok(other) = self.column_ref(nodes, scope) + && other.node != column.node + { + terms.push(OnTerm::Equality(column, other)); + return Ok(()); + } + self.pos = before; + let mut filters = self.scoped_filter_primary(nodes, scope)?; + if filters.len() != 1 { + return self.err("an ON condition must name one table"); + } + let (node, filter) = filters.pop().expect("one"); + terms.push(OnTerm::Condition(node, filter)); + Ok(()) + } + + /// A filter whose conditions may name different tables, as a list of + /// (table, filter) pairs: an `AND` of parts naming different tables + /// splits, an `OR` must stay within one table: each alternative must + /// collapse to one table, the same one. + fn scoped_filter( + &mut self, + nodes: &[Node<'_>], + scope: &[usize], + ) -> Result, ParseError> { + let at = self.at(); + let mut alternatives = vec![self.scoped_conjunct(nodes, scope)?]; + while self.eat_word("OR") { + alternatives.push(self.scoped_conjunct(nodes, scope)?); + } + if alternatives.len() == 1 { + return Ok(alternatives.pop().expect("one")); + } + let mut node = None; + let mut parts = Vec::new(); + for alternative in alternatives { + let (n, filter) = one_table(alternative, at)?; + if node.is_some_and(|seen| seen != n) { + return Err(ParseError { + message: "an OR across tables is not supported".to_string(), + position: at, + }); + } + node = Some(n); + parts.push(filter); + } + Ok(vec![(node.expect("at least one"), collapse(parts, Where::OR))]) + } + + /// An `AND` of scoped primaries, as (table, filter) pairs. + fn scoped_conjunct( + &mut self, + nodes: &[Node<'_>], + scope: &[usize], + ) -> Result, ParseError> { + let mut parts = self.scoped_filter_primary(nodes, scope)?; + while self.eat_word("AND") { + parts.extend(self.scoped_filter_primary(nodes, scope)?); + } + Ok(parts) + } + + /// A parenthesised filter, `TRUE` / `FALSE`, or one condition. + fn scoped_filter_primary( + &mut self, + nodes: &[Node<'_>], + scope: &[usize], + ) -> Result, ParseError> { + if self.eat_sym("(") { + self.descend()?; + let filter = self.scoped_filter(nodes, scope)?; + self.expect_sym(")")?; + self.depth -= 1; + return Ok(filter); + } + if self.eat_word("TRUE") { + return Ok(vec![(scope[0], Where::AND(Vec::new()))]); + } + if self.eat_word("FALSE") { + return Ok(vec![(scope[0], Where::OR(Vec::new()))]); + } + let column = self.column_ref(nodes, scope)?; + let condition = self.comparison(column.column)?; + Ok(vec![(column.node, condition)]) + } + + /// The operator and value after a column, as the single-table parser + /// reads them. + fn comparison(&mut self, column: ColumnName) -> Result { + if self.eat_word("IS") { + let comparison_operator = if self.eat_word("NOT") { + ComparisonOperator::IS_NOT + } else { + ComparisonOperator::IS + }; + if !self.eat_word("NULL") { + return self.err(format!( + "expected NULL after IS [NOT], found {}", + self.describe() + )); + } + return Ok(Where::Condition(Condition { + column, + comparison_operator, + value: Value::Null, + })); + } + let comparison_operator = if self.eat_sym("=") { + if self.eat_word("ANY") { + return Ok(Where::Condition(Condition { + column, + comparison_operator: ComparisonOperator::IN, + value: self.quantified_list("ANY")?, + })); + } + ComparisonOperator::EQ + } else if self.eat_sym("!=") { + if self.eat_word("ALL") { + return Ok(differs_from_all(column, self.quantified_list("ALL")?)); + } + ComparisonOperator::NEQ + } else if self.eat_sym(">=") { + ComparisonOperator::GTE + } else if self.eat_sym(">") { + ComparisonOperator::GT + } else if self.eat_sym("<=") { + ComparisonOperator::LTE + } else if self.eat_sym("<") { + ComparisonOperator::LT + } else if self.eat_word("IN") { + return Ok(Where::Condition(Condition { + column, + comparison_operator: ComparisonOperator::IN, + value: self.value_list()?, + })); + } else if self.eat_word("NOT") { + self.expect_word("IN")?; + return Ok(Where::Condition(Condition { + column, + comparison_operator: ComparisonOperator::NOT_IN, + value: self.value_list()?, + })); + } else { + return self.err(format!( + "expected a comparison operator, [NOT] IN or IS [NOT] NULL after a column, found {}", + self.describe() + )); + }; + Ok(Where::Condition(Condition { + column, + comparison_operator, + value: self.value()?, + })) + } + + /// A column, qualified by its table when the scope has several. + fn column_ref(&mut self, nodes: &[Node<'_>], scope: &[usize]) -> Result { + let at = self.at(); + let first = self.ident()?; + let (node, name) = if self.eat_sym(".") { + let name = self.ident()?; + let node = scope + .iter() + .copied() + .find(|&n| nodes[n].table.name.as_str() == first) + .ok_or_else(|| ParseError { + message: format!("`{first}` is not a table of this query"), + position: at, + })?; + (node, name) + } else if scope.len() == 1 { + (scope[0], first) + } else { + return Err(ParseError { + message: format!("`{first}` must be qualified by its table in a join"), + position: at, + }); + }; + if nodes[node].table.column(&name).is_none() { + return Err(ParseError { + message: format!( + "unknown column `{name}` on table `{}`", + nodes[node].table.name + ), + position: at, + }); + } + Ok(Ref { + node, + column: ColumnName::from(name), + }) + } + + /// The projection over the whole `FROM`. + fn projection_of(&mut self, nodes: &[Node<'_>]) -> Result { + if self.eat_sym("*") { + self.expect_from()?; + return Ok(Projection::All); + } + if self.is_word("COUNT") { + self.bump(); + self.expect_sym("(")?; + self.expect_sym("*")?; + self.expect_sym(")")?; + let name = self.alias()?.unwrap_or_else(|| "count".to_string()); + self.expect_from()?; + return Ok(Projection::Count { name }); + } + let scope: Vec = (0..nodes.len()).collect(); + let mut columns = Vec::new(); + loop { + let column = self.column_ref(nodes, &scope)?; + let name = self + .alias()? + .unwrap_or_else(|| column.column.as_str().to_string()); + columns.push(Projected { + table: nodes[column.node].table.name.clone(), + column: column.column, + name, + }); + if !self.eat_sym(",") { + self.expect_from()?; + return Ok(Projection::Columns(columns)); + } + } + } +} + +/// One term of an `ON`. +enum OnTerm { + /// A column of one table equal to a column of another: the join key. + Equality(Ref, Ref), + /// A condition on one table, which filters it. + Condition(usize, Where), +} + +/// A list of (table, filter) pairs as one table's filter, or an error +/// when they name several. +fn one_table(mut parts: Vec<(usize, Where)>, at: usize) -> Result<(usize, Where), ParseError> { + let Some(node) = parts.first().map(|(node, _)| *node) else { + return Err(ParseError { + message: "an empty condition".to_string(), + position: at, + }); + }; + if parts.iter().any(|(n, _)| *n != node) { + return Err(ParseError { + message: "an OR across tables is not supported".to_string(), + position: at, + }); + } + if parts.len() == 1 { + return Ok(parts.pop().expect("one")); + } + Ok(( + node, + Where::AND(parts.into_iter().map(|(_, filter)| filter).collect()), + )) +} + +/// The engine's tree from the nodes, rooted at `index`. +fn build(nodes: &[Node<'_>], index: usize, order_by: Vec, limit: u32) -> MultiTableReadQuery { + let node = &nodes[index]; + let filter = collapse_conjuncts(node.conjuncts.clone()); + let main = SingleTableReadQuery::new(node.table.name.clone(), filter, order_by, limit); + let joins = node + .edges + .iter() + .map(|(child, parent_column, child_column, inner)| { + let child_table = nodes[*child].table; + let child_order = default_order(child_table, 0) + .map(|order| vec![order]) + .unwrap_or_default(); + let sub = build(nodes, *child, child_order, u32::MAX); + if *inner { + Join::inner(sub, parent_column.clone(), child_column.clone()) + } else { + Join::left(sub, parent_column.clone(), child_column.clone()) + } + }) + .collect(); + MultiTableReadQuery::new(main, joins) +} + +/// The joined tables' key columns in the order the composition yields +/// them: the tables depth first from the root, as the tree's joins are +/// built, each table's key columns in turn. +fn composition<'n>(nodes: &'n [Node<'_>]) -> Vec<(usize, &'n ColumnName)> { + /// Appends the key columns of the tables under `index`, depth first. + fn walk<'n>(nodes: &'n [Node<'_>], index: usize, out: &mut Vec<(usize, &'n ColumnName)>) { + for (child, _, _, _) in &nodes[index].edges { + out.extend(nodes[*child].table.pkey.iter().map(|key| (*child, key))); + walk(nodes, *child, out); + } + } + let mut out = Vec::new(); + walk(nodes, 0, &mut out); + out +} + +/// A node's conjuncts as one filter: `AND` of them, no filter for none. +fn collapse_conjuncts(conjuncts: Vec) -> Where { + match conjuncts.len() { + 0 => Where::AND(Vec::new()), + _ => collapse(conjuncts, Where::AND), + } +} diff --git a/src/stats.rs b/src/stats.rs index 4a56823..526a923 100644 --- a/src/stats.rs +++ b/src/stats.rs @@ -371,6 +371,9 @@ pub struct Stats { pub warm_shapes: AtomicU64, pub feed_lsn: AtomicU64, pub feed_last_message_ms: AtomicU64, + /// How current the engine is, from its heartbeats + /// ([`crate::sync::heartbeat`]). + pub freshness: crate::sync::heartbeat::Freshness, pub process_rss_bytes: AtomicU64, engine: Mutex, pushes: Mutex>>, @@ -532,6 +535,7 @@ impl Stats { warm_shapes: AtomicU64::new(0), feed_lsn: AtomicU64::new(0), feed_last_message_ms: AtomicU64::new(0), + freshness: crate::sync::heartbeat::Freshness::default(), process_rss_bytes: AtomicU64::new(0), engine: Mutex::new(Engine::default()), pushes: Mutex::new(HashMap::new()), @@ -1842,7 +1846,7 @@ mod tests { assert_eq!(bucket_of(5), 10); assert_eq!(bucket_of(7), 12); assert_eq!(bucket_of(8), 13); - let used = |index: usize| (index - 1) / 4 >= 2 || (index - 1) % 4 == 0; + let used = |index: usize| (index - 1) / 4 >= 2 || (index - 1).is_multiple_of(4); let mut previous = 0; for index in (1..BUCKETS).filter(|index| used(*index)) { assert!(lower_bound(index) > previous, "at {index}"); diff --git a/src/sync/catalog.rs b/src/sync/catalog.rs index 5b045e3..b40c248 100644 --- a/src/sync/catalog.rs +++ b/src/sync/catalog.rs @@ -48,9 +48,19 @@ impl CatalogHandle { pub fn swap(&self, catalog: Arc) { self.current.store(catalog); } + + /// Replace the catalog with `change` of it, retried on whatever a + /// concurrent swap put in, so neither change is lost to the other. + pub fn update(&self, change: impl Fn(&Catalog) -> Catalog) { + self.current.rcu(|current| Arc::new(change(current))); + } } -/// `catalog` with `table` in it, replacing a table of the same name. +/// `catalog` with `table` in it, replacing a table of the same name, +/// and knowing the same enum types and text orderings: a bound enum +/// label must stay readable after a migration, and a column's ordering +/// does not change with it. A column the migration adds is not known to +/// order as text, so an ordering on it is refused rather than guessed. pub fn with_table(catalog: &Catalog, table: DbTable) -> Catalog { let name = table.name.clone(); Catalog::new( @@ -60,6 +70,16 @@ pub fn with_table(catalog: &Catalog, table: DbTable) -> Catalog { .cloned() .chain(std::iter::once(table)), ) + .with_enum_oids(catalog.enum_oids()) + .with_ordered_as_text(catalog.ordered_as_text().cloned()) +} + +/// `catalog` knowing `oids` are enum types too, as a type created after +/// the catalog was loaded is learned. +pub fn with_enums(catalog: &Catalog, oids: impl IntoIterator) -> Catalog { + catalog + .clone() + .with_enum_oids(catalog.enum_oids().chain(oids)) } /// `table` with `column` added: the same name and key, one column more. @@ -116,4 +136,17 @@ mod tests { assert!(before.table("tickets").unwrap().column("owner").is_none()); assert_eq!(after.tables().count(), 2); } + + /// A migration absorbed at runtime must not make the server forget + /// which bind types are enums: every enum bind after it was refused. + #[test] + fn a_table_added_keeps_the_enum_types() { + let tickets = DbTable::new("tickets", ["id"], vec![DbColumn::new("id", ValueType::Int)]); + let before = Catalog::new(vec![tickets]).with_enum_oids([19098, 19099]); + let users = DbTable::new("users", ["id"], vec![DbColumn::new("id", ValueType::Int)]); + let after = with_table(&before, users); + assert!(after.is_enum_oid(19098) && after.is_enum_oid(19099)); + assert!(!after.is_enum_oid(23)); + assert_eq!(after.tables().count(), 2); + } } diff --git a/src/sync/heartbeat.rs b/src/sync/heartbeat.rs new file mode 100644 index 0000000..95828e4 --- /dev/null +++ b/src/sync/heartbeat.rs @@ -0,0 +1,251 @@ +//! How current the engine is, measured end to end. The engine emits a +//! numbered heartbeat into the database's WAL every so often, as a +//! logical message on a prefix of its own slot's ([`prefix`]); the +//! change feed carries it, in order, behind every transaction committed +//! before it; and the S2S endpoint marks it applied once it has passed +//! on those transactions' changes, which the engine thread hands it +//! first. Everything committed before the newest +//! applied heartbeat was sent is applied, so the engine is behind by at +//! most that heartbeat's age ([`Freshness::lag`]): under one interval +//! while it keeps up, growing from the moment it stops, whatever stops +//! it (a transaction too large to decode quickly, a stalled engine +//! thread, a slow feed), and measured on the engine's own clock. The +//! S2S endpoint passes the lag on, so a backend can stop trusting its +//! cache when the engine falls behind. + +use std::collections::VecDeque; +use std::sync::Mutex; +use std::time::{Duration, Instant}; + +/// The most heartbeats kept waiting to be applied. Past it the second +/// oldest is dropped, so the oldest still dates a lag with none applied, +/// and a heartbeat applied later is dated by an older one kept, which +/// can only overstate the lag. +const WAITING: usize = 1024; + +/// Heartbeats sent and the newest applied. +#[derive(Debug, Default)] +pub struct Freshness { + beats: Mutex, +} + +/// What [`Freshness`] guards. +#[derive(Debug, Default)] +struct Beats { + /// The number the next heartbeat goes out under: numbers keep rising + /// across reconnects, since a number at or below the newest applied + /// is not waited on. + next: u64, + /// The newest heartbeat written. + written: u64, + /// Heartbeats sent and not yet applied, oldest first, with when + /// each was sent; at most [`WAITING`]. + sent: VecDeque<(u64, Instant)>, + /// The newest heartbeat applied. + applied: u64, + /// When the newest applied heartbeat was sent. + fresh: Option, +} + +impl Freshness { + /// The number for the next heartbeat, never one handed out before. + pub fn number(&self) -> u64 { + let mut beats = self.beats.lock().expect("heartbeats"); + beats.next += 1; + beats.next + } + + /// The newest heartbeat written, 0 before the first. + pub fn written(&self) -> u64 { + self.beats.lock().expect("heartbeats").written + } + + /// Records heartbeat `n` as sent at `at`. It may have been + /// applied already. + pub fn sent(&self, n: u64, at: Instant) { + let mut beats = self.beats.lock().expect("heartbeats"); + beats.written = beats.written.max(n); + if n > beats.applied { + if beats.sent.len() >= WAITING { + beats.sent.remove(1); + } + beats.sent.push_back((n, at)); + } else if n == beats.applied { + beats.fresh = beats.fresh.max(Some(at)); + } + } + + /// Records heartbeat `n` as applied, every change committed before it + /// passed on, and every heartbeat before it. + pub fn applied(&self, n: u64) { + let mut beats = self.beats.lock().expect("heartbeats"); + if n <= beats.applied { + return; + } + beats.applied = n; + while let Some(&(sent, at)) = beats.sent.front() { + if sent > n { + break; + } + beats.fresh = beats.fresh.max(Some(at)); + beats.sent.pop_front(); + } + } + + /// How far behind the engine is at most: the age of the newest + /// heartbeat it has applied, or of the first one sent while it has + /// applied none; `None` until the first is sent. + pub fn lag(&self) -> Option { + let beats = self.beats.lock().expect("heartbeats"); + beats + .fresh + .or_else(|| beats.sent.front().map(|&(_, at)| at)) + .map(|at| at.elapsed()) + } +} + +/// The logical-message prefix of the heartbeats of the engine streaming +/// `slot`: its own, so the heartbeats of other engines on the database +/// are not taken for its. +pub fn prefix(slot: &str) -> String { + format!("{slot}/heartbeat") +} + +/// The statement writing one heartbeat on a server of version `version` +/// (`server_version_num`). A message outside a transaction writes no +/// commit record, so nothing would flush it, and logical decoding reads +/// only flushed WAL: on a quiet database the heartbeat would wait for +/// someone else's commit. PostgreSQL 17 flushes it when asked; before +/// that, taking a transaction id gives the statement a commit record, +/// whose flush covers the message. +pub fn statement(version: u32) -> &'static str { + if version >= 170_000 { + "SELECT pg_logical_emit_message(false, $1::text, $2::text, true)" + } else { + "SELECT pg_logical_emit_message(false, $1::text, $2::text), pg_current_xact_id()" + } +} + +/// Emits heartbeats on `client` every `every`, numbered by `freshness`, +/// until the connection fails; each is recorded in `freshness` once it +/// is in the WAL and flushed, as sent when it was asked for. A +/// non-transactional message is decoded at once, in WAL order, whatever +/// transactions are open. +pub async fn emit( + client: &tokio_postgres::Client, + prefix: &str, + every: Duration, + freshness: &Freshness, +) -> Result<(), tokio_postgres::Error> { + let version: u32 = client + .query_one("SELECT current_setting('server_version_num')::int4", &[]) + .await? + .get::<_, i32>(0) + .try_into() + .unwrap_or(0); + let heartbeat = client.prepare(statement(version)).await?; + let mut tick = tokio::time::interval(every); + tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + loop { + tick.tick().await; + let n = freshness.number(); + let at = Instant::now(); + let content = n.to_string(); + client.execute(&heartbeat, &[&prefix, &content]).await?; + freshness.sent(n, at); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// The lag is unknown before the first heartbeat, the age of the + /// first one while none is applied, then the age of the newest + /// applied one; applying a heartbeat applies those before it, and + /// one applied before it is recorded as sent still counts. + #[test] + fn the_lag_is_the_newest_applied_heartbeats_age() { + let freshness = Freshness::default(); + assert_eq!(freshness.lag(), None); + let long_ago = Instant::now() - Duration::from_secs(5); + freshness.sent(1, long_ago); + freshness.sent(2, Instant::now()); + assert!(freshness.lag().unwrap() >= Duration::from_secs(5)); + freshness.applied(1); + assert!( + freshness.lag().unwrap() >= Duration::from_secs(5), + "2 is unapplied" + ); + freshness.applied(2); + assert!(freshness.lag().unwrap() < Duration::from_secs(1)); + freshness.sent(3, long_ago); + freshness.sent(4, Instant::now()); + freshness.applied(4); + assert!( + freshness.lag().unwrap() < Duration::from_secs(1), + "4 applies 3 too" + ); + freshness.applied(3); + assert!( + freshness.lag().unwrap() < Duration::from_secs(1), + "an old number moves nothing" + ); + freshness.applied(5); + let at = Instant::now(); + freshness.sent(5, at); + assert!( + freshness.lag().unwrap() < Duration::from_secs(1), + "applied before recorded" + ); + freshness.sent(4, long_ago); + assert!( + freshness.lag().unwrap() < Duration::from_secs(1), + "an old number is not waited on" + ); + } + + /// Numbers keep rising however often the emitter starts over, so a + /// heartbeat after a reconnect is waited on; only the newest written + /// is remembered as written. + #[test] + fn numbers_never_repeat() { + let freshness = Freshness::default(); + let first: Vec = (0..3).map(|_| freshness.number()).collect(); + assert_eq!(first, vec![1, 2, 3]); + for &n in &first { + freshness.sent(n, Instant::now()); + } + freshness.applied(3); + let after = freshness.number(); + assert_eq!(after, 4, "a reconnect does not start over"); + freshness.sent(after, Instant::now()); + assert_eq!(freshness.beats.lock().unwrap().sent.len(), 1, "and is waited on"); + freshness.applied(after); + assert!(freshness.beats.lock().unwrap().sent.is_empty()); + assert_eq!(freshness.written(), 4); + } + + /// While none is applied the waiting heartbeats stay bounded, the + /// oldest kept, so the lag still dates from it. + #[test] + fn waiting_heartbeats_are_bounded() { + let freshness = Freshness::default(); + let long_ago = Instant::now() - Duration::from_secs(60); + freshness.sent(freshness.number(), long_ago); + for _ in 0..10 * WAITING { + freshness.sent(freshness.number(), Instant::now()); + } + let waiting = freshness.beats.lock().unwrap().sent.len(); + assert_eq!(waiting, WAITING); + assert!(freshness.lag().unwrap() >= Duration::from_secs(60)); + } + + /// PostgreSQL 17 is asked to flush the message; an older server is + /// given a commit record to flush it with. + #[test] + fn the_heartbeat_is_flushed_on_every_version() { + assert!(statement(170_002).ends_with(", true)")); + assert!(statement(160_004).contains("pg_current_xact_id()")); + } +} diff --git a/src/sync/mod.rs b/src/sync/mod.rs index 986848c..48c9c63 100644 --- a/src/sync/mod.rs +++ b/src/sync/mod.rs @@ -15,6 +15,7 @@ //! runtime, the engine's rules, and both drivers are shared. pub mod catalog; +pub mod heartbeat; mod local; pub mod pg; mod runtime; diff --git a/src/sync/pg/catalog.rs b/src/sync/pg/catalog.rs index 15d34ce..aa0b385 100644 --- a/src/sync/pg/catalog.rs +++ b/src/sync/pg/catalog.rs @@ -1,8 +1,12 @@ //! The catalog read from the database itself: every table of the chosen //! schemas that has a primary key, its columns mapped onto the engine's //! value types the way the sync protocol maps them for its clients (times as epoch -//! milliseconds, JSON and arrays as JSON, enums and uuids as strings, -//! `bytea` left out), so a server needs no hand-written schema. Tables of +//! microseconds, JSON and arrays as JSON, enums and uuids as strings, +//! `bytea` as Postgres's hex text), so a server needs no hand-written +//! schema. It also notes which type OIDs are enums, so a bound parameter +//! of an enum type can be read as the text it is, and which text-carried +//! columns Postgres orders byte by byte, as the engine compares them +//! ([`orders_as_text`]). Tables of //! the `public` schema keep their bare names; any other table is named //! `schema.table`, which is also how the change feed matches it. @@ -11,7 +15,7 @@ use std::collections::BTreeMap; use tokio_postgres::Client; use super::StorageError; -use crate::model::{Catalog, DbColumn, DbTable, ValueType}; +use crate::model::{Catalog, ColumnName, DbColumn, DbTable, TableName, ValueType}; /// What the load left out and why, one line per table or column. pub type Notes = Vec; @@ -23,7 +27,8 @@ pub async fn load_catalog( ) -> Result<(Catalog, Notes), StorageError> { let columns = client .query( - "SELECT c.table_schema, c.table_name, c.column_name, c.data_type, c.udt_name + "SELECT c.table_schema, c.table_name, c.column_name, c.data_type, c.udt_name, + c.collation_name FROM information_schema.columns c JOIN information_schema.tables t ON t.table_schema = c.table_schema AND t.table_name = c.table_name @@ -45,6 +50,29 @@ pub async fn load_catalog( &[&schemas], ) .await?; + let enums = client + .query("SELECT oid, typarray FROM pg_type WHERE typtype = 'e'", &[]) + .await?; + let database = client + .query_one( + "SELECT datcollate, to_jsonb(d) ->> 'datlocprovider' + FROM pg_database d WHERE datname = current_database()", + &[], + ) + .await?; + let default_collation = DefaultCollation { + name: database.get(0), + provider: database.get(1), + }; + let enum_oids: Vec = enums + .iter() + .flat_map(|row| { + let oid: u32 = row.get(0); + let array: u32 = row.get(1); + [oid, array] + }) + .filter(|oid| *oid != 0) + .collect(); let mut pkeys: BTreeMap<(String, String), Vec> = BTreeMap::new(); for row in &keys { let schema: String = row.get(0); @@ -53,6 +81,7 @@ pub async fn load_catalog( pkeys.entry((schema, table)).or_default().push(column); } let mut declared: BTreeMap<(String, String), Vec> = BTreeMap::new(); + let mut ordered: Vec<(TableName, ColumnName)> = Vec::new(); let mut notes = Vec::new(); for row in &columns { let schema: String = row.get(0); @@ -60,6 +89,13 @@ pub async fn load_catalog( let column: String = row.get(2); let data_type: String = row.get(3); let udt_name: String = row.get(4); + let collation: Option = row.get(5); + if orders_as_text(&data_type, collation.as_deref(), &default_collation) { + ordered.push(( + TableName::from(table_name(&schema, &table)), + ColumnName::from(column.as_str()), + )); + } match value_type(&data_type, &udt_name) { Some(r#type) => declared .entry((schema, table)) @@ -92,7 +128,63 @@ pub async fn load_catalog( columns, )); } - Ok((Catalog::new(tables), notes)) + Ok(( + Catalog::new(tables) + .with_enum_oids(enum_oids) + .with_ordered_as_text(ordered), + notes, + )) +} + +/// The database's default collation: its locale name, and its provider +/// (`c` libc, `i` ICU, `b` builtin; absent before Postgres 15, libc). +pub struct DefaultCollation { + pub name: String, + pub provider: Option, +} + +impl DefaultCollation { + /// Whether the default orders strings by their bytes: libc's `C` or + /// `POSIX` locales (and their UTF-8 spellings), or the builtin + /// provider, whose locales order by code point, which is UTF-8's + /// byte order. ICU never does. + fn is_bytewise(&self) -> bool { + match self.provider.as_deref() { + Some("i") => false, + Some("b") => true, + _ => is_bytewise_collation(&self.name), + } + } +} + +/// Whether a collation by `name` orders strings by their bytes. +fn is_bytewise_collation(name: &str) -> bool { + matches!( + name, + "C" | "POSIX" | "C.UTF-8" | "C.utf8" | "ucs_basic" | "pg_c_utf8" + ) +} + +/// Whether Postgres orders a column of `data_type` (`information_schema`'s +/// name), under `collation` (`None`: the database's default), as its text +/// compares byte by byte. A `uuid`'s lower-case hex and a `bytea`'s hex +/// do, as do `name` (always `C`) and `time`'s zero-padded text; strings +/// do under a byte-order collation; everything else carried as text (an +/// enum, `citext`, an `inet`, an interval, a time with its zone) orders +/// by rules of its own. +pub fn orders_as_text( + data_type: &str, + collation: Option<&str>, + default: &DefaultCollation, +) -> bool { + match data_type { + "uuid" | "bytea" | "name" | "time without time zone" => true, + "text" | "character varying" | "character" => match collation { + Some(name) => is_bytewise_collation(name), + None => default.is_bytewise(), + }, + _ => false, + } } /// The catalog name of a table: bare in `public`, schema-qualified @@ -122,7 +214,9 @@ pub fn value_type(data_type: &str, udt_name: &str) -> Option { /// The engine's type for a scalar Postgres type name (`information_schema`'s /// `data_type` or `pg_type`'s name alike); `None` for one the wire cannot -/// carry. +/// carry. A `bytea` is the hex text Postgres renders (`\x` and two digits +/// a byte): what the change feed carries, and what a server-to-server +/// client turns back into bytes. pub fn scalar_type(name: &str) -> Option { let name = name.to_ascii_lowercase(); let name = name.split('(').next().unwrap_or("").trim(); @@ -158,7 +252,7 @@ pub fn scalar_type(name: &str) -> Option { | "macaddr" | "xml" | "tsvector" => ValueType::String, - "bytea" => return None, + "bytea" => ValueType::String, _ => return None, }) } @@ -168,7 +262,7 @@ mod tests { use super::*; /// The mapping follows the protocol's: times are numbers, json and arrays are - /// json, enums are strings, bytea is left out. + /// json, enums and bytea are strings, ranges are left out. #[test] fn types_map_like_the_protocol() { assert_eq!( @@ -187,7 +281,8 @@ mod tests { ); assert_eq!(value_type("bigint", "int8"), Some(ValueType::Int)); assert_eq!(value_type("numeric", "numeric"), Some(ValueType::Float)); - assert_eq!(value_type("bytea", "bytea"), None); + assert_eq!(value_type("bytea", "bytea"), Some(ValueType::String)); + assert_eq!(value_type("int4range", "int4range"), None); assert_eq!(table_name("public", "tickets"), "tickets"); assert_eq!(table_name("xyne_0", "clients"), "xyne_0.clients"); } diff --git a/src/sync/pg/ddl.rs b/src/sync/pg/ddl.rs index 244ef46..c3316a4 100644 --- a/src/sync/pg/ddl.rs +++ b/src/sync/pg/ddl.rs @@ -146,7 +146,7 @@ impl PublishedTable { /// The engine's type for a published column, mapped as the startup /// catalog maps `information_schema`: an array is a list of its element's -/// type, an enum is a string, `bytea` is not carried. +/// type, an enum is a string, `bytea` is its hex text. pub fn column_type(column: &PublishedColumn) -> Option { if let Some(element) = column.data_type.strip_suffix("[]") { let element = element.trim_end_matches("[]"); @@ -284,7 +284,7 @@ fn typed(text: &str, declared: &ValueType) -> Option { "f" | "false" | "0" => false, _ => return None, }), - ValueType::Timestamp => Value::Int(text::epoch_millis(text)?), + ValueType::Timestamp => Value::Int(text::epoch_micros(text)?), ValueType::List(inner) => match text::array_literal(text, inner) { Value::List(items) => Value::List(items), _ => return None, @@ -564,7 +564,7 @@ mod tests { &ValueType::Timestamp ), Some(Value::Int( - text::epoch_millis("2026-01-02 03:04:05+00").unwrap() + text::epoch_micros("2026-01-02 03:04:05+00").unwrap() )) ); for expression in [ @@ -610,7 +610,8 @@ mod tests { column_type(&column("text[]", "b")), Some(ValueType::List(Box::new(ValueType::String))) ); - assert_eq!(column_type(&column("bytea", "b")), None); + assert_eq!(column_type(&column("bytea", "b")), Some(ValueType::String)); + assert_eq!(column_type(&column("int4range", "r")), None); } /// A column added with a constant default, a table added, both diff --git a/src/sync/pg/mod.rs b/src/sync/pg/mod.rs index f94d0bf..a5cc61a 100644 --- a/src/sync/pg/mod.rs +++ b/src/sync/pg/mod.rs @@ -113,11 +113,11 @@ pub fn read_row_limit() -> usize { "STREAMGRES_SYNC_JOIN_LIMIT", "XYNE_SYNC_JOIN_LIMIT", ] - .iter() - .find_map(|name| std::env::var(name).ok()) - .and_then(|value| value.trim().parse().ok()) - .filter(|limit| *limit > 0) - .unwrap_or(DEFAULT_READ_ROW_LIMIT) + .iter() + .find_map(|name| std::env::var(name).ok()) + .and_then(|value| value.trim().parse().ok()) + .filter(|limit| *limit > 0) + .unwrap_or(DEFAULT_READ_ROW_LIMIT) } /// How the kernel probes a connection to the database while it is silent, @@ -449,7 +449,7 @@ impl PgStorage { catalog: Arc, runtime: Handle, ) -> Result { - let client = open(&read_config, &runtime).await?; + let client = open(&session_config(&read_config, None), &runtime).await?; let pool = Arc::new(Pool { config, read_config, @@ -907,14 +907,12 @@ impl Pool { /// How a read's connection is opened: `config` (which carries the /// [`Keepalive`] the storage was given, if any, so the kernel probes - /// the connection while it sits idle in the pool), with the read - /// timeout also limiting how long a new connection may take to open. + /// the connection while it sits idle in the pool), the session's + /// time zone ([`session_config`]), and the read timeout as both the + /// time the connection may take to open and the session's + /// `statement_timeout`. fn read_config(&self) -> Config { - let mut config = self.read_config.clone(); - if let Some(limit) = self.timeout() { - config.connect_timeout(limit); - } - config + session_config(&self.read_config, self.timeout()) } /// Run one simple-query batch on a pooled connection, within the read @@ -999,6 +997,42 @@ impl Pool { } } +/// The session a storage connection runs in: `config` with the session's +/// time zone pinned to UTC, string literals pinned to the standard form +/// and `bytea` to hex output, and, under a read `timeout`, that long to +/// open and as the session's `statement_timeout`; the settings follow +/// whatever options the connection string already carries, so they win. +/// +/// The literal form matters for safety: a value from a server-to-server +/// bind reaches the SQL as a quoted literal with only its quotes doubled, +/// which is exact while `standard_conforming_strings` is on (Postgres's +/// default) and would let a backslash escape the quote were a role or +/// database to turn it off. `bytea_output = hex` is the form the feed and +/// the catalog carry a `bytea` in. +/// +/// The zone matters: a time in a filter is spelled as a zone-less +/// timestamp (`TIMESTAMP 'epoch' + N microseconds`, so an index on the +/// column answers it), which PostgreSQL reads in the session's zone when +/// the column is `timestamptz`. The engine keeps every time as UTC +/// microseconds, so the session must read the literal as UTC too, or a +/// range on such a column misses rows by the zone's offset on a server +/// whose default zone is not UTC. +pub(crate) fn session_config(config: &Config, timeout: Option) -> Config { + let mut config = config.clone(); + let mut ours = + "-c TimeZone=UTC -c standard_conforming_strings=on -c bytea_output=hex".to_string(); + if let Some(limit) = timeout { + config.connect_timeout(limit); + ours.push_str(&format!(" -c statement_timeout={}", limit.as_millis())); + } + let options = match config.get_options() { + Some(theirs) => format!("{theirs} {ours}"), + None => ours, + }; + config.options(&options); + config +} + /// Run `work` as a task of `runtime` and wait for its result here. async fn run_on( runtime: &Handle, @@ -1111,7 +1145,7 @@ fn decode_row( /// One column of a result row in its text form, decoded by the cast its /// declared type was read with ([`sql::select_expr`]): a time column -/// arrives as its epoch milliseconds, a JSON column as the text of its +/// arrives as its epoch microseconds, a JSON column as the text of its /// `jsonb` (kept in the standard form, [`text::standard_json`]), a list /// or map column as JSON text, the rest in Postgres's text form for the /// cast type. @@ -1145,6 +1179,29 @@ fn decode_text(text: &str, declared: &ValueType) -> Result mod tests { use super::*; + /// Every storage session runs in UTC, after the connection string's + /// own options so the zone wins, with the statement timeout beside + /// it when a read timeout is set. + #[test] + fn a_storage_session_is_pinned_to_utc() { + let mut config = Config::new(); + assert_eq!( + session_config(&config, None).get_options(), + Some("-c TimeZone=UTC -c standard_conforming_strings=on -c bytea_output=hex") + ); + config.options("-c TimeZone=Asia/Kolkata -c search_path=app"); + assert_eq!( + session_config(&config, Some(Duration::from_millis(250))).get_options(), + Some( + "-c TimeZone=Asia/Kolkata -c search_path=app -c TimeZone=UTC -c standard_conforming_strings=on -c bytea_output=hex -c statement_timeout=250" + ) + ); + assert_eq!( + session_config(&config, Some(Duration::from_millis(250))).get_connect_timeout(), + Some(&Duration::from_millis(250)) + ); + } + #[test] fn base36_keeps_any_snapshot_count_within_thirteen_digits() { assert_eq!(base36(0), "0"); diff --git a/src/sync/pg/replication.rs b/src/sync/pg/replication.rs index 154886a..00491b9 100644 --- a/src/sync/pg/replication.rs +++ b/src/sync/pg/replication.rs @@ -279,6 +279,7 @@ fn session_options(config: &Config) -> String { options.push_str(&format!("-c {setting}={value}")); }; set("idle_in_transaction_session_timeout", "0"); + set("TimeZone", "UTC"); if config.get_keepalives() { set( "tcp_keepalives_idle", @@ -371,7 +372,7 @@ mod tests { .keepalives_retries(3); assert_eq!( session_options(&config), - "-c idle_in_transaction_session_timeout=0 -c tcp_keepalives_idle=30 \ + "-c idle_in_transaction_session_timeout=0 -c TimeZone=UTC -c tcp_keepalives_idle=30 \ -c tcp_keepalives_interval=10 -c tcp_keepalives_count=3" ); @@ -382,14 +383,14 @@ mod tests { config.keepalives(false); assert_eq!( session_options(&config), - "-c idle_in_transaction_session_timeout=0" + "-c idle_in_transaction_session_timeout=0 -c TimeZone=UTC" ); let mut config = Config::new(); config.keepalives_idle(Duration::from_millis(500)); assert_eq!( session_options(&config), - "-c idle_in_transaction_session_timeout=0 -c tcp_keepalives_idle=1" + "-c idle_in_transaction_session_timeout=0 -c TimeZone=UTC -c tcp_keepalives_idle=1" ); } diff --git a/src/sync/pg/sql.rs b/src/sync/pg/sql.rs index 358a9f1..713fcb9 100644 --- a/src/sync/pg/sql.rs +++ b/src/sync/pg/sql.rs @@ -268,24 +268,29 @@ pub fn cast_of(declared: &ValueType) -> &'static str { } /// The SQL reading one column in the form the engine keeps it: a time -/// column as epoch milliseconds, a JSON column as the text of its `jsonb` +/// column as epoch microseconds, a JSON column as the text of its `jsonb` /// (the cast is nothing for a `jsonb` column, and gives a `json` column's /// stored text the one spelling `jsonb` has), an array or map column as /// JSON text, every other cast to its declared type's Postgres form. pub fn select_expr(column: &str, declared: &ValueType) -> String { let quoted = quote_ident(column); match declared { - ValueType::Timestamp => epoch_millis("ed), + ValueType::Timestamp => epoch_micros("ed), ValueType::Json => format!("{quoted}::jsonb::text"), ValueType::List(_) | ValueType::Map(_, _) => format!("to_json({quoted})::text"), other => format!("{quoted}::{}", cast_of(other)), } } -/// `expr` as milliseconds since the epoch; `extract` is numeric, so the -/// rounding to `int8` keeps the millisecond a `timestamp(3)` column has. -fn epoch_millis(expr: &str) -> String { - format!("(extract(epoch from {expr}) * 1000)::int8") +/// `expr` as microseconds since the epoch, as the change feed reads it +/// ([`super::text::epoch_micros`]): `extract` is exact at the column's +/// full precision; `infinity` and `-infinity` are the extreme `int8`s, +/// and a finite time past what `int8` microseconds count saturates one +/// short of them, so no time Postgres holds fails the read. +fn epoch_micros(expr: &str) -> String { + format!( + "CASE WHEN isfinite({expr}) THEN LEAST(GREATEST(extract(epoch from {expr})::numeric * 1000000, -9223372036854775807), 9223372036854775806)::int8 WHEN {expr} > 'epoch' THEN 9223372036854775807 ELSE -9223372036854775808 END" + ) } /// A column as the SQL its conditions compare it by: the bare column @@ -303,7 +308,7 @@ fn column_expr(column: &str, scope: &Scope<'_>) -> String { } /// A literal compared against a column of `declared` type: the engine -/// keeps a time as milliseconds since the epoch, so against a time +/// keeps a time as microseconds since the epoch, so against a time /// column the number becomes the timestamp it names, and the column /// itself stays bare for the index; against a JSON column the text the /// engine holds is the `jsonb` it names. @@ -311,12 +316,16 @@ fn literal_as(value: &Value, declared: Option<&ValueType>) -> String { if let (Value::String(text), Some(ValueType::Json)) = (value, declared) { return format!("{}::jsonb", quote_literal(text)); } - let millis = match (value, declared) { - (Value::Int(millis), Some(ValueType::Timestamp)) => *millis, - (Value::Float(millis), Some(ValueType::Timestamp)) if millis.is_finite() => *millis as i64, + let micros = match (value, declared) { + (Value::Int(micros), Some(ValueType::Timestamp)) => *micros, + (Value::Float(micros), Some(ValueType::Timestamp)) if micros.is_finite() => *micros as i64, _ => return literal(value), }; - format!("(TIMESTAMP 'epoch' + {millis} * INTERVAL '1 millisecond')") + match micros { + super::text::INFINITY_MICROS => "TIMESTAMP 'infinity'".to_owned(), + super::text::NEG_INFINITY_MICROS => "TIMESTAMP '-infinity'".to_owned(), + _ => format!("(TIMESTAMP 'epoch' + {micros} * INTERVAL '1 microsecond')"), + } } /// A double-quoted identifier. @@ -347,7 +356,9 @@ fn render_where(filter: &Where, scope: &Scope<'_>) -> String { } } -/// One leaf, with the engine's `NULL` and list semantics made explicit. +/// One leaf, with the engine's `NULL` and list semantics made explicit: +/// a `NOT IN` of no values holds for every row with a value, as the +/// engine evaluates it, so a read and the routing agree on `NULL`. fn render_condition(condition: &Condition, scope: &Scope<'_>) -> String { use ComparisonOperator::*; if condition.column.is_composite() && condition.comparison_operator != EXISTS { @@ -371,7 +382,11 @@ fn render_condition(condition: &Condition, scope: &Scope<'_>) -> String { .map(|value| literal_as(value, declared)) .collect(); if literals.is_empty() { - return if negated { "TRUE" } else { "FALSE" }.to_owned(); + return if negated { + format!("{column} IS NOT NULL") + } else { + "FALSE".to_owned() + }; } let keyword = if negated { "NOT IN" } else { "IN" }; format!("{column} {keyword} ({})", literals.join(", ")) @@ -539,6 +554,21 @@ mod tests { sql[sql.find(" WHERE ").expect("a WHERE") + 7..].to_owned() } + /// A `NOT IN` of no values keeps every row with a value and none + /// without, as the engine's routing does; an `IN` of none keeps none. + #[test] + fn an_empty_list_renders_as_the_engine_evaluates_it() { + let empty = || Value::List(Vec::new()); + assert_eq!( + comments_where(Condition::new("body", ComparisonOperator::NOT_IN, empty())), + "\"body\" IS NOT NULL" + ); + assert_eq!( + comments_where(Condition::new("body", ComparisonOperator::IN, empty())), + "FALSE" + ); + } + /// A composite key renders as a row constructor compared whole: its /// `IN` lists the set's pairs as tuples, each component typed by its /// own column; `=` compares one tuple, `<>` and `NOT IN` compare one @@ -558,7 +588,7 @@ mod tests { ComparisonOperator::IN, Value::Set(set) )), - "(\"org\", \"at\") IN (('acme', (TIMESTAMP 'epoch' + 1000 * INTERVAL '1 millisecond')))" + "(\"org\", \"at\") IN (('acme', (TIMESTAMP 'epoch' + 1000 * INTERVAL '1 microsecond')))" ); assert_eq!( comments_where(Condition::new( @@ -566,7 +596,7 @@ mod tests { ComparisonOperator::EQ, pair("acme", 1000) )), - "(\"org\", \"at\") = ('acme', (TIMESTAMP 'epoch' + 1000 * INTERVAL '1 millisecond'))" + "(\"org\", \"at\") = ('acme', (TIMESTAMP 'epoch' + 1000 * INTERVAL '1 microsecond'))" ); assert_eq!( comments_where(Condition::new( @@ -574,7 +604,7 @@ mod tests { ComparisonOperator::NEQ, pair("acme", 1000) )), - "((\"org\", \"at\") <> ('acme', (TIMESTAMP 'epoch' + 1000 * INTERVAL '1 millisecond')) AND \"org\" IS NOT NULL AND \"at\" IS NOT NULL)" + "((\"org\", \"at\") <> ('acme', (TIMESTAMP 'epoch' + 1000 * INTERVAL '1 microsecond')) AND \"org\" IS NOT NULL AND \"at\" IS NOT NULL)" ); let with_null = Value::List(vec![Value::Null, Value::Int(1000)]); let narrow = Value::List(vec!["acme".into()]); @@ -589,7 +619,7 @@ mod tests { pair("acme", 1000) ]), )), - "(\"org\", \"at\") IN (('acme', (TIMESTAMP 'epoch' + 1000 * INTERVAL '1 millisecond')))" + "(\"org\", \"at\") IN (('acme', (TIMESTAMP 'epoch' + 1000 * INTERVAL '1 microsecond')))" ); assert_eq!( comments_where(Condition::new( @@ -613,7 +643,7 @@ mod tests { ComparisonOperator::NOT_IN, Value::List(vec![pair("acme", 1000)]) )), - "((\"org\", \"at\") NOT IN (('acme', (TIMESTAMP 'epoch' + 1000 * INTERVAL '1 millisecond'))) AND \"org\" IS NOT NULL AND \"at\" IS NOT NULL)" + "((\"org\", \"at\") NOT IN (('acme', (TIMESTAMP 'epoch' + 1000 * INTERVAL '1 microsecond'))) AND \"org\" IS NOT NULL AND \"at\" IS NOT NULL)" ); assert_eq!( comments_where(Condition::new( @@ -903,8 +933,8 @@ mod tests { } /// A time condition compares the bare column against the timestamp - /// the millisecond literal names, so an index on the column answers - /// it; the projection still returns milliseconds. + /// the microsecond literal names, so an index on the column answers + /// it; the projection still returns microseconds. #[test] fn a_time_condition_compares_the_column_itself() { let table = DbTable::new( @@ -931,20 +961,33 @@ mod tests { let sql = select_sql(&query, &table); assert!( sql.contains( - "\"createdAt\" > (TIMESTAMP 'epoch' + 1750000000000 * INTERVAL '1 millisecond')" + "\"createdAt\" > (TIMESTAMP 'epoch' + 1750000000000 * INTERVAL '1 microsecond')" ), "{sql}" ); assert!( - sql.contains("\"createdAt\" IN ((TIMESTAMP 'epoch' + 1000 * INTERVAL '1 millisecond'), (TIMESTAMP 'epoch' + 2000 * INTERVAL '1 millisecond'))"), + sql.contains("\"createdAt\" IN ((TIMESTAMP 'epoch' + 1000 * INTERVAL '1 microsecond'), (TIMESTAMP 'epoch' + 2000 * INTERVAL '1 microsecond'))"), "{sql}" ); assert!( - sql.contains("(extract(epoch from \"createdAt\") * 1000)::int8"), - "the projection still reads milliseconds: {sql}" + sql.contains("extract(epoch from \"createdAt\")::numeric * 1000000"), + "the projection still reads microseconds: {sql}" ); } + /// The infinities go both ways: read as the extreme microseconds, + /// and a condition on them renders as `infinity` again rather than + /// an interval past the end of time. + #[test] + fn an_infinite_time_reads_and_renders() { + let read = select_expr("at", &ValueType::Timestamp); + assert!(read.contains("WHEN isfinite(\"at\")"), "{read}"); + assert!(read.contains("THEN 9223372036854775807 ELSE -9223372036854775808"), "{read}"); + let timed = Some(&ValueType::Timestamp); + assert_eq!(literal_as(&Value::Int(i64::MAX), timed), "TIMESTAMP 'infinity'"); + assert_eq!(literal_as(&Value::Int(i64::MIN), timed), "TIMESTAMP '-infinity'"); + } + #[test] fn renders_null_tests() { let query = SingleTableReadQuery::new( diff --git a/src/sync/pg/stream.rs b/src/sync/pg/stream.rs index dd29525..8d5df70 100644 --- a/src/sync/pg/stream.rs +++ b/src/sync/pg/stream.rs @@ -126,6 +126,9 @@ pub struct Transaction { pub catalog: Option>, pub committed_at_micros: i64, pub decode: Duration, + /// The engine's own heartbeat this event is ([`crate::sync::heartbeat`]), + /// which carries no write; `None` for a transaction. + pub heartbeat: Option, } /// One published table as the feed described it: the catalog table it @@ -144,6 +147,23 @@ struct Described { plain_json: bool, } +/// A heartbeat of this engine's, as the transaction that carries it to +/// the engine thread behind every transaction before it: no write, and +/// the position the feed is at ([`Feed::absorb`]). Content that is not +/// a number is not one of ours. +fn beat(content: &Bytes) -> Option { + let n = std::str::from_utf8(content).ok()?.parse::().ok()?; + Some(Transaction { + at: Lsn(0), + writes: Vec::new(), + schema: Vec::new(), + catalog: None, + committed_at_micros: 0, + decode: Duration::ZERO, + heartbeat: Some(n), + }) +} + /// The mapping of decoded messages onto catalog writes: the catalog as /// the feed has it (it grows as the trigger's messages arrive, so a row /// after a change is decoded on the shape it was written in), the @@ -161,6 +181,8 @@ pub struct Decoder { /// that already has them, and the engine must still be told. changes: Vec, decode: Duration, + /// The prefix of this engine's heartbeats, when it sends any. + heartbeat: Option, } impl Decoder { @@ -173,6 +195,7 @@ impl Decoder { ddl: None, changes: Vec::new(), decode: Duration::ZERO, + heartbeat: None, } } @@ -225,6 +248,7 @@ impl Decoder { schema, committed_at_micros: commit_time_micros, decode: Duration::ZERO, + heartbeat: None, })) } ReplicationEvent::XLogData { data, .. } => { @@ -235,17 +259,28 @@ impl Decoder { prefix, content, lsn, - .. + transactional, } => { if self.ddl.as_ref().is_some_and(|ddl| ddl.prefix == prefix) { self.migrate(position(lsn), &content)?; } + if !transactional && self.heartbeat.as_deref() == Some(prefix.as_str()) { + return Ok(beat(&content)); + } Ok(None) } ReplicationEvent::KeepAlive { .. } | ReplicationEvent::StoppedAt { .. } => Ok(None), } } + /// The decoder hearing the heartbeats sent on `prefix`. + pub fn with_heartbeat(self, prefix: String) -> Self { + Decoder { + heartbeat: Some(prefix), + ..self + } + } + /// Absorb one message of the trigger, written at `at`: the catalog /// becomes what the message makes of it, every row decoded from here /// on is laid out on that, and the changes wait for the commit that @@ -757,6 +792,7 @@ impl Transport { received: Instant::now(), committed_at_micros: transaction.committed_at_micros, decode: transaction.decode, + heartbeat: transaction.heartbeat, }; if out.send(committed).await.is_err() { return Ok(()); @@ -909,6 +945,15 @@ impl Feed { self.decoder.catalog() } + /// The feed hearing the heartbeats sent on `prefix` + /// ([`crate::sync::heartbeat`]). + pub fn with_heartbeat(self, prefix: String) -> Self { + Feed { + decoder: self.decoder.with_heartbeat(prefix), + ..self + } + } + /// Absorb one raw event: a keepalive or a commit moves the position, /// and a commit also yields the transaction it closes. A keepalive /// that arrives inside a transaction names a location below that @@ -918,8 +963,11 @@ impl Feed { if let ReplicationEvent::KeepAlive { wal_end, .. } = &event { self.progress = self.progress.max(position(*wal_end)); } - let transaction = self.decoder.absorb(event)?; - if let Some(transaction) = &transaction { + let mut transaction = self.decoder.absorb(event)?; + if let Some(transaction) = &mut transaction { + if transaction.heartbeat.is_some() { + transaction.at = self.progress; + } self.progress = self.progress.max(transaction.at); } Ok(transaction) @@ -1189,7 +1237,7 @@ fn convert(raw: &str, declared: &ValueType) -> Result { ValueType::Json => Value::String(super::text::standard_json(raw).into_owned()), ValueType::String | ValueType::Map(_, _) => Value::String(raw.to_owned()), ValueType::List(inner) => super::text::array_literal(raw, inner), - ValueType::Timestamp => Value::Int(super::text::epoch_millis(raw).ok_or_else(unreadable)?), + ValueType::Timestamp => Value::Int(super::text::epoch_micros(raw).ok_or_else(unreadable)?), ValueType::Bool => Value::Bool(match raw { "true" | "t" => true, "false" | "f" => false, @@ -1527,6 +1575,43 @@ mod tests { assert_eq!(feed.progress(), Lsn::parse("0/2A8").unwrap()); } + /// A logical message on `prefix`, transactional or not. + fn logical(prefix: &str, content: &str, transactional: bool) -> ReplicationEvent { + ReplicationEvent::Message { + transactional, + lsn: pgwire_replication::Lsn(0x300), + prefix: prefix.to_owned(), + content: Bytes::from(content.to_owned()), + } + } + + /// A heartbeat of the feed's own is a transaction with no write, at + /// the position the feed is at, carrying its number; another + /// engine's, a transactional message on the prefix, or content that + /// is not a number is not one. + #[test] + fn a_heartbeat_comes_through_as_a_mark() { + let mut feed = Feed::new(catalog()).with_heartbeat("slot_a/heartbeat".to_owned()); + assert!(feed.absorb(keepalive("0/100")).unwrap().is_none()); + let beat = feed + .absorb(logical("slot_a/heartbeat", "7", false)) + .unwrap() + .expect("a heartbeat"); + assert_eq!(beat.heartbeat, Some(7)); + assert!(beat.writes.is_empty()); + assert_eq!(beat.at, Lsn::parse("0/100").unwrap()); + assert_eq!(feed.progress(), Lsn::parse("0/100").unwrap()); + for other in [ + logical("slot_b/heartbeat", "8", false), + logical("slot_a/heartbeat", "9", true), + logical("slot_a/heartbeat", "nine", false), + ] { + assert!(feed.absorb(other).unwrap().is_none()); + } + let mut deaf = Feed::new(catalog()); + assert!(deaf.absorb(logical("slot_a/heartbeat", "7", false)).unwrap().is_none()); + } + /// A JSON cell is decoded into the standard form whichever kind of /// column holds it: a `jsonb` cell loses a padded fraction, a `json` /// cell is rewritten as its `jsonb` would read; both are then the text diff --git a/src/sync/pg/text.rs b/src/sync/pg/text.rs index 0564f51..d120f26 100644 --- a/src/sync/pg/text.rs +++ b/src/sync/pg/text.rs @@ -1,5 +1,5 @@ //! The text forms Postgres uses for the values the engine carries in a -//! form of its own: the epoch-millisecond reading of a time, a JSON array +//! form of its own: the epoch-microsecond reading of a time, a JSON array //! read from `to_json(column)`, and the `{a,b}` array literal the change //! feed delivers, each turned into the [`Value`] the column's declared //! type calls for; and the **standard form** of a JSON value, the one text @@ -193,46 +193,105 @@ pub fn json_as_jsonb(text: &str) -> String { } } -/// Milliseconds since the Unix epoch of a Postgres `timestamp`, -/// `timestamptz` or `date` in its text form (`2026-09-15 10:00:00.123`, -/// with an optional `+05:30` style offset, or `2026-09-15`); a timestamp -/// without a zone is read as UTC, the protocol's convention. `None` when the text -/// is not a time (`infinity` included). -pub fn epoch_millis(text: &str) -> Option { +/// The microseconds the engine keeps for `infinity`: after every time, +/// as Postgres orders it, and Postgres's own binary value for it. +pub const INFINITY_MICROS: i64 = i64::MAX; + +/// The microseconds the engine keeps for `-infinity`. +pub const NEG_INFINITY_MICROS: i64 = i64::MIN; + +/// Microseconds since the Unix epoch of a Postgres `timestamp`, +/// `timestamptz` or `date` in its text form: `2026-09-15 10:00:00.123456`, +/// with an optional `+05:30` (or `+05:53:28`) offset, or `2026-09-15`; a +/// year past 9999, a trailing ` BC`, and `infinity` / `-infinity` +/// ([`INFINITY_MICROS`]). A timestamp without a zone is read as UTC, the +/// protocol's convention. Every time Postgres can hold reads; one past +/// what the microseconds can count (the last thirty millennia of +/// Postgres's range) saturates. `None` only for text that is not a +/// time: the change feed stops on that, so nothing Postgres prints may +/// land there. +pub fn epoch_micros(text: &str) -> Option { let text = text.trim(); - if text.len() == 10 { - let date = chrono::NaiveDate::parse_from_str(text, "%Y-%m-%d").ok()?; - return Some(date.and_hms_opt(0, 0, 0)?.and_utc().timestamp_millis()); + match text { + "infinity" => return Some(INFINITY_MICROS), + "-infinity" => return Some(NEG_INFINITY_MICROS), + _ => {} } - let (naive, offset_seconds) = split_offset(text); - let parsed = chrono::NaiveDateTime::parse_from_str(naive, "%Y-%m-%d %H:%M:%S%.f") - .or_else(|_| chrono::NaiveDateTime::parse_from_str(naive, "%Y-%m-%d %H:%M:%S")) - .or_else(|_| chrono::NaiveDateTime::parse_from_str(naive, "%Y-%m-%dT%H:%M:%S%.f")) - .ok()?; - Some(parsed.and_utc().timestamp_millis() - offset_seconds * 1000) + let (text, bc) = match text.strip_suffix(" BC") { + Some(rest) => (rest, true), + None => (text, false), + }; + let (date, time) = match text.split_once([' ', 'T']) { + Some((date, time)) => (date, Some(time)), + None => (text, None), + }; + let mut parts = date.splitn(3, '-'); + let year: i64 = parts.next()?.parse().ok()?; + let month: i64 = parts.next()?.parse().ok()?; + let day: i64 = parts.next()?.parse().ok()?; + if !(1..=12).contains(&month) || !(1..=31).contains(&day) { + return None; + } + let year = if bc { 1 - year } else { year }; + let (seconds, micros, offset) = match time { + Some(time) => clock(time)?, + None => (0, 0, 0), + }; + let total = i128::from(days_from_civil(year, month, day)) * 86_400_000_000 + + i128::from(seconds - offset) * 1_000_000 + + i128::from(micros); + Some(total.clamp(i128::from(i64::MIN + 1), i128::from(i64::MAX - 1)) as i64) } -/// Split a trailing `+HH`, `+HH:MM` or `+HHMM` zone offset off a timestamp -/// text, returning the rest and the offset in seconds east of UTC. -fn split_offset(text: &str) -> (&str, i64) { - let Some(position) = text.rfind(['+', '-']) else { - return (text, 0); +/// A time of day with its optional zone offset: seconds into the day, +/// the fraction in microseconds, and the offset in seconds east of UTC. +fn clock(text: &str) -> Option<(i64, i64, i64)> { + let (time, offset) = match text.find(['+', '-']) { + Some(at) => { + let (time, zone) = text.split_at(at); + let sign = if zone.starts_with('-') { -1 } else { 1 }; + let mut fields = zone[1..].split(':'); + let mut offset = 0; + for unit in [3600, 60, 1] { + match fields.next() { + Some(field) if field.len() == 4 && unit == 3600 => { + offset += field[..2].parse::().ok()? * 3600 + + field[2..].parse::().ok()? * 60; + } + Some(field) => offset += field.parse::().ok()? * unit, + None => break, + } + } + (time, sign * offset) + } + None => (text, 0), }; - if position < 11 { - return (text, 0); + let (whole, fraction) = time.split_once('.').unwrap_or((time, "")); + let mut fields = whole.split(':'); + let hours: i64 = fields.next()?.parse().ok()?; + let minutes: i64 = fields.next()?.parse().ok()?; + let seconds: i64 = fields.next()?.parse().ok()?; + if fields.next().is_some() || fraction.len() > 6 || !fraction.bytes().all(|b| b.is_ascii_digit()) { + return None; } - let (rest, zone) = text.split_at(position); - let sign = if zone.starts_with('-') { -1 } else { 1 }; - let digits: String = zone[1..].chars().filter(char::is_ascii_digit).collect(); - let (hours, minutes) = match digits.len() { - 2 => (digits.parse::().unwrap_or(0), 0), - 4 => ( - digits[..2].parse::().unwrap_or(0), - digits[2..].parse::().unwrap_or(0), - ), - _ => return (text, 0), + let micros = if fraction.is_empty() { + 0 + } else { + fraction.parse::().ok()? * 10_i64.pow(6 - fraction.len() as u32) }; - (rest, sign * (hours * 3600 + minutes * 60)) + Some((hours * 3600 + minutes * 60 + seconds, micros, offset)) +} + +/// Days from 1970-01-01 to a proleptic Gregorian date, astronomical +/// years (1 BC is year 0), for any year Postgres can hold; Howard +/// Hinnant's `days_from_civil`. +fn days_from_civil(year: i64, month: i64, day: i64) -> i64 { + let year = if month <= 2 { year - 1 } else { year }; + let era = year.div_euclid(400); + let year_of_era = year - era * 400; + let day_of_year = (153 * ((month + 9) % 12) + 2) / 5 + day - 1; + let day_of_era = year_of_era * 365 + year_of_era / 4 - year_of_era / 100 + day_of_year; + era * 146_097 + day_of_era - 719_468 } /// A JSON array text (what `to_json(column)::text` returns for an array @@ -355,12 +414,35 @@ mod tests { /// Timestamps without a zone read as UTC, with a zone are shifted, and /// dates land on UTC midnight. #[test] - fn times_read_as_epoch_milliseconds() { - assert_eq!(epoch_millis("1970-01-01 00:00:01.5"), Some(1_500)); - assert_eq!(epoch_millis("1970-01-01 01:00:00+01"), Some(0)); - assert_eq!(epoch_millis("1970-01-01 05:30:00.25+05:30"), Some(250)); - assert_eq!(epoch_millis("1970-01-02"), Some(86_400_000)); - assert_eq!(epoch_millis("infinity"), None); + fn times_read_as_epoch_microseconds() { + assert_eq!(epoch_micros("1970-01-01 00:00:01.5"), Some(1_500_000)); + assert_eq!(epoch_micros("1970-01-01 00:00:00.123456"), Some(123_456)); + assert_eq!(epoch_micros("1970-01-01 01:00:00+01"), Some(0)); + assert_eq!(epoch_micros("1970-01-01 05:30:00.25+05:30"), Some(250_000)); + assert_eq!(epoch_micros("1970-01-02"), Some(86_400_000_000)); + assert_eq!(epoch_micros("not a time"), None); + } + + /// Every time Postgres prints reads, so the change feed never stops + /// on one: the infinities as the extreme microseconds, BC dates, years + /// past 9999, offsets with seconds, and the instants either side of + /// the epoch; what the microseconds cannot count saturates. + #[test] + fn every_time_postgres_prints_reads() { + assert_eq!(epoch_micros("infinity"), Some(INFINITY_MICROS)); + assert_eq!(epoch_micros("-infinity"), Some(NEG_INFINITY_MICROS)); + assert_eq!(epoch_micros("1969-12-31 23:59:59.999999+00"), Some(-1)); + assert_eq!(epoch_micros("2000-01-01 00:00:00+00"), Some(946_684_800_000_000)); + assert_eq!(epoch_micros("0001-01-01 00:00:00+00"), Some(-62_135_596_800_000_000)); + assert_eq!(epoch_micros("0001-12-31 BC"), Some(-62_135_596_800_000_000 - 86_400_000_000)); + assert_eq!(epoch_micros("4713-01-01 00:00:00+00 BC"), Some(-210_863_520_000_000_000)); + assert_eq!(epoch_micros("12345-06-07 08:09:10"), Some(327_416_976_550_000_000)); + assert_eq!(epoch_micros("294276-12-31 23:59:59.999999+00"), Some(i64::MAX - 1)); + assert_eq!(epoch_micros("1900-01-01 00:00:00+05:53:28"), Some(-2_208_988_800_000_000 - 21_208_000_000)); + assert_eq!(epoch_micros("2026-10-21T10:00:00.5-0330"), epoch_micros("2026-10-21 13:30:00.5")); + for text in ["2026-13-01", "2026-01-01 25:00", "2026-01-01 10:00:00.1234567", "", "x-y-z"] { + assert_eq!(epoch_micros(text), None, "{text}"); + } } /// Array literals split on unquoted commas, unquote, and read NULL. diff --git a/src/sync/pg/threads.rs b/src/sync/pg/threads.rs index 4d4244a..73995e6 100644 --- a/src/sync/pg/threads.rs +++ b/src/sync/pg/threads.rs @@ -68,6 +68,9 @@ use crate::sync::{CatalogHandle, Command, Event, Runtime, Service, Sources, Tran /// - `ddl_trigger`: the event trigger on `ddl_command_end` through which /// schema changes are heard; the server will not serve without it. /// - `ddl_prefix`: the prefix of that trigger's logical messages. +/// - `heartbeat`: how often the engine writes a heartbeat into the WAL +/// to measure how far behind it is ([`crate::sync::heartbeat`]); zero +/// for none. #[derive(Debug, Clone)] pub struct Settings { pub dsn: String, @@ -87,6 +90,7 @@ pub struct Settings { pub watched: Vec, pub ddl_trigger: String, pub ddl_prefix: String, + pub heartbeat: Duration, } /// What every slot a server names for itself starts with. @@ -216,7 +220,8 @@ pub fn spawn_feed( prefix: settings.ddl_prefix.clone(), schemas: settings.schemas.clone(), }; - let mut feed = Feed::with_ddl(catalog, ddl); + let mut feed = Feed::with_ddl(catalog, ddl) + .with_heartbeat(crate::sync::heartbeat::prefix(&settings.slot)); let mut start = Some(start); let mut failures = 0u32; loop { @@ -337,6 +342,41 @@ pub async fn stop_feed(mut feed: FeedHandle) -> Result<(), String> { } } +/// Writes this engine's heartbeats into the WAL of the database `config` +/// names, for as long as the process runs, reconnecting a second after a +/// failure. It runs off the engine thread (on the reads runtime), so a +/// stalled engine thread shows as a growing lag rather than as heartbeats +/// that stop coming. Once heartbeats have been written, a stretch in +/// which they cannot be makes the lag grow, so backends read Postgres +/// until they are written again. Until the first is written the lag is +/// unknown and backends trust the engine, as they would one that sends +/// no lag: a standby cannot write them, and its engine keeps serving. +/// A failure is logged once until heartbeats are written again. +async fn beat(config: tokio_postgres::Config, slot: String, every: Duration, stats: Arc) { + let prefix = crate::sync::heartbeat::prefix(&slot); + let mut warned = false; + let mut written = 0; + loop { + let outcome = match super::open(&config, &tokio::runtime::Handle::current()).await { + Ok(client) => crate::sync::heartbeat::emit(&client, &prefix, every, &stats.freshness) + .await + .map_err(|error| error.to_string()), + Err(error) => Err(error.to_string()), + }; + if stats.freshness.written() != written { + written = stats.freshness.written(); + warned = false; + } + if let Err(error) = outcome + && !warned + { + log_warn!("heartbeats stopped: {error}; retrying every second"); + warned = true; + } + tokio::time::sleep(Duration::from_secs(1)).await; + } +} + /// Bring the engine side up on the current thread's local set: the slot /// made to exist, the storage over `settings` (its reads running on /// `reads`) with its first snapshot minted, the feed thread started at @@ -367,14 +407,18 @@ pub async fn start( .parse() .map_err(|error| format!("connecting the read pool: {error}"))?; settings.keepalive.apply(&mut read_config); - let pg = - PgStorage::connect_configured_with_read_config(config, read_config, catalog.load(), reads) - .await - .map_err(|error| format!("connecting the storage: {error}"))? - .with_rotation(settings.snapshot_rotation) - .with_read_connections(settings.read_connections) - .with_read_idle_connections(settings.read_idle_connections) - .with_read_timeout(settings.read_timeout); + let pg = PgStorage::connect_configured_with_read_config( + config.clone(), + read_config, + catalog.load(), + reads.clone(), + ) + .await + .map_err(|error| format!("connecting the storage: {error}"))? + .with_rotation(settings.snapshot_rotation) + .with_read_connections(settings.read_connections) + .with_read_idle_connections(settings.read_idle_connections) + .with_read_timeout(settings.read_timeout); let pg = Arc::new(pg); let (feed, feed_handle) = spawn_feed( settings.clone(), @@ -391,6 +435,9 @@ pub async fn start( .map_err(|error| format!("warming the memory tables: {error}"))?; log_info!("memory tables warmed with {rows} rows"); } + if !settings.heartbeat.is_zero() { + reads.spawn(beat(config.clone(), settings.slot.clone(), settings.heartbeat, Arc::clone(&stats))); + } let (sinks, events): (Vec<_>, Vec<_>) = (0..consumers.max(1)) .map(|_| mpsc::unbounded_channel()) .unzip(); diff --git a/src/sync/runtime.rs b/src/sync/runtime.rs index 50f3501..6146650 100644 --- a/src/sync/runtime.rs +++ b/src/sync/runtime.rs @@ -271,7 +271,17 @@ impl Runtime { snapshot.at, self.position ); + let read_rows = snapshot.rows.len(); let (rows, worst_read) = self.bring_up(&flight.fetch, snapshot.rows, snapshot.at); + crate::log::log_debug!( + "read {} of {}: {} rows at {}, {} after bringing up to {}", + flight.fetch.id.0, + flight.fetch.query.table, + read_rows, + snapshot.at, + rows.len(), + self.position + ); self.stats.reads_landed += 1; step.updates = self.engine.land(&flight.fetch, &rows, worst_read.as_ref()); self.landed.push(rows); diff --git a/src/sync/service.rs b/src/sync/service.rs index c597759..5f2eacb 100644 --- a/src/sync/service.rs +++ b/src/sync/service.rs @@ -90,6 +90,10 @@ pub struct Transaction { pub received: Instant, pub committed_at_micros: i64, pub decode: Duration, + /// The engine's heartbeat this mark carries, passed to every consumer + /// once every transaction before it is applied + /// ([`crate::sync::heartbeat`]). + pub heartbeat: Option, } impl Transaction { @@ -131,6 +135,7 @@ impl Transaction { received: Instant::now(), committed_at_micros: 0, decode: Duration::ZERO, + heartbeat: None, } } @@ -173,6 +178,10 @@ impl Transaction { /// least half the row limit: a query to report by name before it grows /// into the limit and is refused. Sent once per read, to the owner of /// one subscription of the tree. +/// - `Heartbeat`: the engine's heartbeat `n` was applied +/// ([`crate::sync::heartbeat`]), sent to every consumer behind the +/// events of every transaction before it, so a consumer that has +/// passed on those events can take it as applied. #[derive(Debug)] pub enum Event { Registered { @@ -205,6 +214,7 @@ pub enum Event { table: TableName, rows: u64, }, + Heartbeat(u64), } /// How a step ended: what its deltas travel inside of. @@ -478,6 +488,7 @@ where received, committed_at_micros, decode, + heartbeat, } = transaction; let started = Instant::now(); let mut step = Step::default(); @@ -557,6 +568,11 @@ where routed, }, ); + if let Some(n) = heartbeat { + for sink in &self.sinks { + let _ = sink.send(Event::Heartbeat(n)); + } + } self.reap(); } diff --git a/tests/pg_live.rs b/tests/pg_live.rs index c4af53e..7d07313 100644 --- a/tests/pg_live.rs +++ b/tests/pg_live.rs @@ -12,18 +12,18 @@ use std::rc::Rc; use std::sync::Arc; use std::time::{Duration, Instant}; -use tokio::sync::mpsc; -use tokio::task::{LocalSet, spawn_local}; -use tokio_postgres::{Client, NoTls}; +use streamgres::client::ddl_triggers::trigger_stack_sql; use streamgres::ivm::{Delta, Engine, Fetch, MultiTableIVM, QueryPart}; use streamgres::model::*; -use streamgres::client::ddl_triggers::trigger_stack_sql; use streamgres::sync::pg::ddl::DdlSource; use streamgres::sync::pg::estimate_cache::EstimateCache; use streamgres::sync::pg::{PgStorage, PgStream}; use streamgres::sync::{ CatalogHandle, Command, Event, Lsn, Runtime, Service, Sources, Storage, SubId, Transaction, }; +use tokio::sync::mpsc; +use tokio::task::{LocalSet, spawn_local}; +use tokio_postgres::{Client, NoTls}; /// The database under test, if any. fn dsn() -> Option { @@ -432,6 +432,7 @@ fn service_streams_end_to_end() { catalog: None, received: Instant::now(), committed_at_micros: 0, + heartbeat: None, decode: std::time::Duration::ZERO, })) .await @@ -1847,3 +1848,195 @@ fn unchanged_out_of_line_values_are_never_lost() { cleanup(&dsn, &client, &names).await; }); } + +/// A range on a `timestamptz` column is read right whatever zone the +/// session would otherwise run in: the engine keeps times as UTC +/// microseconds and spells a time in a filter as a zone-less timestamp, +/// which PostgreSQL reads in the session's zone, so the storage pins the +/// session to UTC after whatever the connection string asks for. A +/// `timestamp` column, zone-less on both sides, reads the same either +/// way. The desk's "created between" lookup missed every row on a +/// server whose default zone was not UTC, until each row's next write. +#[test] +fn a_time_range_is_read_right_in_any_session_zone() { + let Some(dsn) = dsn() else { return }; + block_on(async { + let table = "live_zone_events"; + let slot = "xyne_sync_live_zone"; + let publication = "xyne_sync_live_zone_pub"; + let client = admin(&dsn).await; + PgStream::drop_slot(&dsn, slot, publication) + .await + .expect("drop slot"); + client + .batch_execute(&format!( + "DROP TABLE IF EXISTS {table}; + CREATE TABLE {table} (id int8 PRIMARY KEY, at timestamptz NOT NULL, naive timestamp NOT NULL); + CREATE PUBLICATION \"{publication}\" FOR ALL TABLES; + INSERT INTO {table} VALUES + (1, '2026-10-01 08:00:00+00', '2026-10-01 08:00:00'), + (2, '2026-10-01 09:30:00+00', '2026-10-01 09:30:00'), + (3, '2026-10-01 11:30:00+00', '2026-10-01 11:30:00'), + (4, '2026-10-01 12:30:00+00', '2026-10-01 12:30:00');", + )) + .await + .expect("prepare"); + let catalog = Arc::new(Catalog::new(vec![DbTable::new( + table, + ["id"], + vec![ + DbColumn::new("id", ValueType::Int), + DbColumn::new("at", ValueType::Timestamp), + DbColumn::new("naive", ValueType::Timestamp), + ], + )])); + let mut stream = PgStream::open(&dsn, slot, publication, catalog.clone()) + .await + .expect("open stream"); + let mut config: tokio_postgres::Config = dsn.parse().expect("dsn"); + config.options("-c TimeZone=Asia/Kolkata"); + let storage = PgStorage::connect_configured( + config, + catalog.clone(), + tokio::runtime::Handle::current(), + ) + .await + .expect("connect"); + let mut runtime = Runtime::new(MultiTableIVM::new()); + catch_up(&mut runtime, &mut stream, &[&storage]).await; + + let micros = |hour: i64| -> Value { + let day = chrono::NaiveDate::from_ymd_opt(2026, 10, 1).unwrap(); + let at = day.and_hms_opt(hour as u32, 0, 0).unwrap().and_utc(); + Value::Int(at.timestamp_micros()) + }; + for column in ["at", "naive"] { + let query = SingleTableReadQuery::new( + table, + Where::AND(vec![ + Where::condition(column, ComparisonOperator::GTE, micros(9)), + Where::condition(column, ComparisonOperator::LT, micros(12)), + ]), + OrderBy::new("id", Order::ASC), + u32::MAX, + ); + let snapshot = storage.select(&query).await.expect("select"); + let read: BTreeSet = snapshot + .rows + .iter() + .filter_map(|(key, _)| match key.pkey_value.get("id") { + Some(Value::Int(id)) => Some(*id), + _ => None, + }) + .collect(); + assert_eq!( + read, + BTreeSet::from([2, 3]), + "the rows between 09:00 and 12:00 UTC by `{column}`, in a session asked for Asia/Kolkata" + ); + } + let _ = client + .batch_execute(&format!("DROP TABLE IF EXISTS {table};")) + .await; + let _ = PgStream::drop_slot(&dsn, slot, publication).await; + }); +} + +/// Every time PostgreSQL can hold reaches the engine, through the +/// snapshot and through the change feed: `infinity` and `-infinity` as +/// the extreme microseconds (where they order, as PostgreSQL orders +/// them), a BC date and a year past 9999 exactly. A single `infinity` +/// in a published table used to stop the change feed, and the engine +/// with it, on every restart until the row was changed by hand. +#[test] +fn every_time_postgres_holds_reaches_the_engine() { + let Some(dsn) = dsn() else { return }; + block_on(async { + let names = Names::new("infinite"); + let table = names.extra.clone(); + let client = prepare(&dsn, &names).await; + client + .batch_execute(&format!( + "CREATE TABLE {table} (id int8 PRIMARY KEY, at timestamptz, naive timestamp, day date); + INSERT INTO {table} VALUES + (1, 'infinity', 'infinity', 'infinity'), + (2, '-infinity', '-infinity', '-infinity'), + (3, '4713-01-01 00:00:00+00 BC', '4713-01-01 00:00:00 BC', '4713-01-01 BC'), + (4, '12345-06-07 08:09:10+00', '12345-06-07 08:09:10', '12345-06-07');" + )) + .await + .expect("times"); + let catalog = Arc::new(Catalog::new(vec![DbTable::new( + table.as_str(), + ["id"], + vec![ + DbColumn::new("id", ValueType::Int), + DbColumn::new("at", ValueType::Timestamp), + DbColumn::new("naive", ValueType::Timestamp), + DbColumn::new("day", ValueType::Timestamp), + ], + )])); + let mut stream = PgStream::open(&dsn, &names.slot, &names.publication, catalog.clone()) + .await + .expect("open stream"); + let storage = PgStorage::connect(&dsn, catalog.clone()) + .await + .expect("connect"); + let mut runtime = Runtime::new(MultiTableIVM::new()); + catch_up(&mut runtime, &mut stream, &[&storage]).await; + let expected = BTreeMap::from([ + (1, i64::MAX), + (2, i64::MIN), + (3, -210_863_520_000_000_000), + (4, 327_416_976_550_000_000), + ]); + let read = |filter: Where| SingleTableReadQuery::new(table.as_str(), filter, OrderBy::new("id", Order::ASC), u32::MAX); + let snapshot = storage.select(&read(Where::AND(vec![]))).await.expect("select"); + for (key, row) in &snapshot.rows { + let Some(Value::Int(id)) = key.pkey_value.get("id") else { + panic!("an id") + }; + let day = if *id == 4 { 327_416_947_200_000_000 } else { expected[id] }; + assert_eq!(row.data.get("at"), Some(&Value::Int(expected[id])), "at of {id}"); + assert_eq!(row.data.get("naive"), Some(&Value::Int(expected[id])), "naive of {id}"); + assert_eq!(row.data.get("day"), Some(&Value::Int(day)), "day of {id}"); + } + assert_eq!(snapshot.rows.len(), 4); + + let infinite = storage + .select(&read(Where::condition("at", ComparisonOperator::EQ, Value::Int(i64::MAX)))) + .await + .expect("a condition on infinity renders"); + assert_eq!(infinite.rows.len(), 1, "infinity finds its row"); + let after = storage + .select(&read(Where::condition("at", ComparisonOperator::GT, Value::Int(0)))) + .await + .expect("select"); + assert_eq!(after.rows.len(), 2, "infinity and the year 12345 come after 1970"); + + client + .batch_execute(&format!( + "UPDATE {table} SET at = '-infinity', naive = '4713-01-01 00:00:00 BC' WHERE id = 1; + INSERT INTO {table} VALUES (5, 'infinity', '12345-06-07 08:09:10', '-infinity');" + )) + .await + .expect("changes"); + let mut images = Vec::new(); + let deadline = Instant::now() + Duration::from_secs(10); + while images.len() < 2 { + assert!(Instant::now() < deadline, "the feed did not deliver both changes"); + let batch = stream.poll().await.expect("the feed reads every time"); + for (write, _) in batch.writes { + if let Some(image) = write.new_row_image() { + images.push(image.data.clone()); + } + } + } + assert_eq!(images[0].get("at"), Some(&Value::Int(i64::MIN))); + assert_eq!(images[0].get("naive"), Some(&Value::Int(-210_863_520_000_000_000))); + assert_eq!(images[1].get("at"), Some(&Value::Int(i64::MAX))); + assert_eq!(images[1].get("naive"), Some(&Value::Int(327_416_976_550_000_000))); + assert_eq!(images[1].get("day"), Some(&Value::Int(i64::MIN))); + cleanup(&dsn, &client, &names).await; + }); +} diff --git a/tests/s2s_live.rs b/tests/s2s_live.rs new file mode 100644 index 0000000..f3e4c7e --- /dev/null +++ b/tests/s2s_live.rs @@ -0,0 +1,379 @@ +//! The server-to-server endpoint spoken to directly, frame by frame, +//! with no client library in between: the greeting, a subscribe in rows +//! mode and in touch mode, the refusals, a change arriving as a delta +//! and as a touch, an unsubscribe that stops the frames, and a frame +//! the endpoint cannot read. +//! +//! Needs a running engine (`STREAMGRES_S2S_URL`) and the database it +//! streams (`STREAMGRES_PG_DSN`); skips itself otherwise. When the engine +//! was started with `STREAMGRES_S2S_TOKEN`, give the test the same one: +//! it then connects with it, and checks that a missing or wrong token is +//! answered 401. + +use std::collections::BTreeSet; +use std::time::{Duration, Instant}; + +use base64::Engine as _; +use base64::engine::general_purpose::STANDARD as BASE64; +use futures_util::{SinkExt, StreamExt}; +use serde_json::{Value, json}; +use tokio::net::TcpStream; +use tokio_postgres::{Client, NoTls}; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; +use tokio_tungstenite::tungstenite::handshake::client::Request; +use tokio_tungstenite::tungstenite::http::header::AUTHORIZATION; +use tokio_tungstenite::tungstenite::{Error, Message}; +use tokio_tungstenite::{MaybeTlsStream, WebSocketStream}; + +const TABLE: &str = "live_s2s_items"; + +/// The socket to the endpoint. +type Socket = WebSocketStream>; + +/// The endpoint and the database, when both are given. +fn targets() -> Option<(String, String)> { + let url = std::env::var("STREAMGRES_S2S_URL").ok(); + let dsn = std::env::var("STREAMGRES_PG_DSN").ok(); + match (url, dsn) { + (Some(url), Some(dsn)) => Some((url, dsn)), + _ => { + eprintln!("STREAMGRES_S2S_URL or STREAMGRES_PG_DSN not set: s2s live scenario skipped"); + None + } + } +} + +/// The upgrade request for `url`, with `token` as a bearer when given. +fn request(url: &str, token: Option<&str>) -> Request { + let mut request = url.into_client_request().expect("a WebSocket URL"); + if let Some(token) = token { + let value = format!("Bearer {token}").parse().expect("a header value"); + request.headers_mut().insert(AUTHORIZATION, value); + } + request +} + +/// The token the engine was started with, if the test was given one. +fn token() -> Option { + std::env::var("STREAMGRES_S2S_TOKEN") + .ok() + .filter(|token| !token.is_empty()) +} + +/// Connects to the database on its own task. +async fn database(dsn: &str) -> Client { + let (client, connection) = tokio_postgres::connect(dsn, NoTls).await.expect("connect"); + tokio::spawn(async move { + let _ = connection.await; + }); + client +} + +/// The table with three rows, `n` 1 to 3, created if the database lacks +/// it. The table stays: dropping it would stop the engine. +async fn seeded(client: &Client) { + client + .batch_execute(&format!( + "CREATE TABLE IF NOT EXISTS {TABLE} (id bigint PRIMARY KEY, name text NOT NULL, n int NOT NULL); + DELETE FROM {TABLE}; + INSERT INTO {TABLE} (id, name, n) VALUES (1, 'one', 1), (2, 'two', 2), (3, 'three', 3)" + )) + .await + .expect("seed"); +} + +/// An `int4` bind as the endpoint takes it. +fn int4(value: i32) -> Value { + json!({"oid": 23, "bytes": BASE64.encode(value.to_be_bytes())}) +} + +/// Sends one frame. +async fn send(socket: &mut Socket, frame: Value) { + socket + .send(Message::Text(frame.to_string().into())) + .await + .expect("send"); +} + +/// The next text frame within `wait` that is not a `lag`, which the +/// engine sends on its own every 500 ms; `None` on a close or when none +/// came. +async fn next(socket: &mut Socket, wait: Duration) -> Option { + let deadline = Instant::now() + wait; + loop { + let left = deadline.saturating_duration_since(Instant::now()); + let frame = next_any(socket, left).await?; + if tag(&frame) != "lag" { + return Some(frame); + } + } +} + +/// The next text frame within `wait`, a `lag` included; `None` on a +/// close or when none came. +async fn next_any(socket: &mut Socket, wait: Duration) -> Option { + let deadline = Instant::now() + wait; + loop { + let left = deadline.saturating_duration_since(Instant::now()); + let frame = tokio::time::timeout(left, socket.next()).await.ok()??; + match frame.expect("frame") { + Message::Text(text) => return Some(serde_json::from_str(text.as_str()).expect("json")), + Message::Close(_) => return None, + _ => continue, + } + } +} + +/// The next frame, which must arrive. +async fn expect(socket: &mut Socket, what: &str) -> Value { + next(socket, Duration::from_secs(5)) + .await + .unwrap_or_else(|| panic!("no frame within five seconds, waiting for {what}")) +} + +/// The frame's tag. +fn tag(frame: &Value) -> &str { + frame[0].as_str().unwrap_or("") +} + +/// Subscribes `sql` under `id`, with `rows` asked for or not, retrying +/// under fresh ids while the engine has not yet heard of a table the +/// test just created. Returns the `accepted` body and the id it got. +async fn subscribe_once_known( + socket: &mut Socket, + id: &mut u64, + sql: &str, + binds: Vec, + rows: bool, +) -> (u64, Value) { + let deadline = Instant::now() + Duration::from_secs(10); + loop { + let this = *id; + *id += 1; + send( + socket, + json!(["subscribe", {"id": this, "sql": sql, "binds": binds, "rows": rows}]), + ) + .await; + let answer = expect(socket, "a verdict").await; + assert_eq!(answer[1]["id"].as_u64(), Some(this), "{answer}"); + match tag(&answer) { + "accepted" => return (this, answer[1].clone()), + "refused" + if answer[1]["reason"] + .as_str() + .is_some_and(|reason| reason.contains("unknown table")) => + { + assert!( + Instant::now() < deadline, + "the engine never heard of {TABLE}: {answer}" + ); + tokio::time::sleep(Duration::from_millis(100)).await; + } + _ => panic!("{answer}"), + } + } +} + +/// The ids of the rows in a `nodes` entry. +fn ids(node: &Value) -> BTreeSet { + node["rows"] + .as_array() + .expect("rows") + .iter() + .map(|row| row["id"].as_i64().expect("id")) + .collect() +} + +/// Reads frames for up to `wait`, returning them all. +async fn frames_for(socket: &mut Socket, wait: Duration) -> Vec { + let deadline = Instant::now() + wait; + let mut frames = Vec::new(); + while let Some(frame) = next(socket, deadline.saturating_duration_since(Instant::now())).await { + frames.push(frame); + } + frames +} + +/// The endpoint speaks its protocol against a live Postgres: `hello` +/// and `pong`; `accepted` with the rows and the query's shape in rows +/// mode, without them otherwise; `refused` with the id and the reason; +/// a commit's `delta` and `touched`; an unsubscribed id no longer told; +/// and `error` for a frame it cannot read. +#[tokio::test] +async fn the_endpoint_speaks_its_protocol() { + let Some((url, dsn)) = targets() else { + return; + }; + let client = database(&dsn).await; + seeded(&client).await; + let (mut socket, _) = tokio_tungstenite::connect_async(request(&url, token().as_deref())) + .await + .expect("connect"); + let hello = expect(&mut socket, "hello").await; + assert_eq!(tag(&hello), "hello"); + assert_eq!(hello[1]["protocol"], 1); + let lag = next_any(&mut socket, Duration::from_secs(5)) + .await + .expect("a lag frame"); + assert_eq!(tag(&lag), "lag", "the lag comes before any answer: {lag}"); + assert!(lag[1]["ms"].is_u64() || lag[1]["ms"].is_null(), "{lag}"); + assert!( + hello[1]["database"]["system"].is_string() && hello[1]["database"]["name"].is_string(), + "the engine names the database it streams: {hello}" + ); + assert!( + hello[1]["max_subscriptions"].is_u64() || hello[1]["max_subscriptions"].is_null(), + "{hello}" + ); + + send(&mut socket, json!(["ping"])).await; + assert_eq!(expect(&mut socket, "pong").await, json!(["pong"])); + + let mut id = 1; + let sql = + format!(r#"SELECT "id", "name", "n" FROM "{TABLE}" WHERE "n" > $1 ORDER BY "id" ASC"#); + let (rows_id, accepted) = + subscribe_once_known(&mut socket, &mut id, &sql, vec![int4(1)], true).await; + let nodes = accepted["nodes"].as_array().expect("nodes"); + assert_eq!(nodes.len(), 1, "one table"); + assert_eq!(nodes[0]["table"].as_str(), Some(TABLE)); + assert_eq!(nodes[0]["key"], json!(["id"])); + assert_eq!(nodes[0]["parent"], Value::Null); + assert_eq!( + ids(&nodes[0]), + BTreeSet::from([2, 3]), + "the rows with n above 1" + ); + assert_eq!( + accepted["columns"], + json!([[TABLE, "id"], [TABLE, "name"], [TABLE, "n"]]) + ); + assert_eq!(accepted["names"], json!(["id", "name", "n"])); + assert_eq!(accepted["count"], json!(false)); + assert_eq!(accepted["order"], json!([["id", "asc"]])); + assert_eq!(accepted["limit"], Value::Null); + assert_eq!(accepted["offset"], json!(0)); + + let (touch_id, accepted) = + subscribe_once_known(&mut socket, &mut id, &sql, vec![int4(0)], false).await; + assert!( + accepted.get("nodes").is_none(), + "no rows asked for: {accepted}" + ); + + send(&mut socket, json!(["subscribe", {"id": id, "sql": format!(r#"SELECT "id" FROM "{TABLE}" WHERE "name" ILIKE $1"#), "binds": [{"oid": 25, "bytes": BASE64.encode("o%")}], "rows": true}])).await; + let refused = expect(&mut socket, "a refusal").await; + assert_eq!(tag(&refused), "refused", "{refused}"); + assert_eq!(refused[1]["id"].as_u64(), Some(id)); + assert!( + refused[1]["reason"] + .as_str() + .is_some_and(|reason| !reason.is_empty()) + ); + id += 1; + + send(&mut socket, json!(["subscribe", {"id": id, "sql": r#"SELECT "id" FROM "no_such_table_anywhere""#, "binds": [], "rows": false}])).await; + let refused = expect(&mut socket, "a refusal").await; + assert_eq!(tag(&refused), "refused", "{refused}"); + assert!( + refused[1]["reason"] + .as_str() + .is_some_and(|reason| reason.contains("unknown table")), + "{refused}" + ); + + client + .execute( + &format!("INSERT INTO {TABLE} (id, name, n) VALUES (4, 'four', 4)"), + &[], + ) + .await + .expect("insert"); + let mut delta = None; + let mut touched = None; + let deadline = Instant::now() + Duration::from_secs(5); + while (delta.is_none() || touched.is_none()) && Instant::now() < deadline { + let Some(frame) = next(&mut socket, Duration::from_secs(5)).await else { + break; + }; + match tag(&frame) { + "delta" => delta = Some(frame), + "touched" => touched = Some(frame), + _ => panic!("{frame}"), + } + } + let delta = delta.expect("the rows-mode subscription got a delta"); + assert_eq!(delta[1]["id"].as_u64(), Some(rows_id)); + let changes = delta[1]["changes"].as_array().expect("changes"); + assert_eq!(changes.len(), 1); + assert_eq!(changes[0]["table"].as_str(), Some(TABLE)); + assert_eq!( + changes[0]["put"], + json!([{"id": 4, "name": "four", "n": 4}]) + ); + assert_eq!(changes[0]["del"], json!([])); + let touched = touched.expect("the touch-mode subscription was touched"); + assert_eq!(touched[1]["ids"], json!([touch_id])); + + send(&mut socket, json!(["unsubscribe", {"id": rows_id}])).await; + client + .execute( + &format!("UPDATE {TABLE} SET name = 'vier' WHERE id = 4"), + &[], + ) + .await + .expect("update"); + let frames = frames_for(&mut socket, Duration::from_secs(2)).await; + assert_eq!(frames.len(), 1, "only the touch: {frames:?}"); + assert_eq!(frames[0], json!(["touched", {"ids": [touch_id]}])); + + send(&mut socket, json!(["bogus"])).await; + let error = expect(&mut socket, "an error").await; + assert_eq!(tag(&error), "error", "{error}"); + + client + .execute(&format!("DELETE FROM {TABLE} WHERE id = 4"), &[]) + .await + .expect("delete"); + assert_eq!( + expect(&mut socket, "the touch for the delete").await, + json!(["touched", {"ids": [touch_id]}]) + ); + send(&mut socket, json!(["unsubscribe", {"id": touch_id}])).await; + client + .execute( + &format!("UPDATE {TABLE} SET name = 'uno' WHERE id = 1"), + &[], + ) + .await + .expect("update"); + assert!( + frames_for(&mut socket, Duration::from_secs(1)) + .await + .is_empty(), + "nothing is subscribed" + ); + socket.close(None).await.expect("close"); +} + +/// An engine started with a token answers 401 to an upgrade without it +/// or with another one. +#[tokio::test] +async fn a_token_guards_the_endpoint() { + let Some((url, _)) = targets() else { + return; + }; + let Some(token) = token() else { + eprintln!("STREAMGRES_S2S_TOKEN not set: the token check is skipped"); + return; + }; + let wrong = format!("{token}-not"); + for presented in [None, Some(wrong.as_str())] { + match tokio_tungstenite::connect_async(request(&url, presented)).await { + Err(Error::Http(response)) => assert_eq!(response.status(), 401), + Err(other) => panic!("expected a 401, got {other}"), + Ok(_) => panic!("the endpoint took {presented:?}"), + } + } +} diff --git a/tests/service_events.rs b/tests/service_events.rs index 65be438..151587e 100644 --- a/tests/service_events.rs +++ b/tests/service_events.rs @@ -97,6 +97,7 @@ fn shapes(events: &[Event]) -> Vec<&'static str> { Event::Committed { .. } => "committed", Event::Capped { .. } => "capped", Event::Heavy { .. } => "heavy", + Event::Heartbeat(_) => "heartbeat", }) .collect() } @@ -107,9 +108,11 @@ fn updates_of(event: &Event) -> &[streamgres::ivm::Delta] { Event::Registered { updates, .. } | Event::Landed { updates } | Event::Committed { updates, .. } => updates, - Event::Hydrated(_) | Event::Refused { .. } | Event::Capped { .. } | Event::Heavy { .. } => { - &[] - } + Event::Hydrated(_) + | Event::Refused { .. } + | Event::Capped { .. } + | Event::Heavy { .. } + | Event::Heartbeat(_) => &[], } } @@ -309,6 +312,7 @@ fn a_commit_reports_the_position_and_the_watched_writes() { catalog: None, received: std::time::Instant::now(), committed_at_micros: 0, + heartbeat: None, decode: std::time::Duration::ZERO, })) .await @@ -335,6 +339,49 @@ fn a_commit_reports_the_position_and_the_watched_writes() { }); } +/// A heartbeat mark reaches every consumer as its own event, after the +/// events of the transactions before it: a consumer that has passed +/// those on may take the heartbeat as applied. +#[test] +fn a_heartbeat_follows_the_transactions_before_it_to_every_consumer() { + block_on(async { + let storage = Rc::new(MemoryStorage::new()); + let (owner_tx, mut owner) = mpsc::unbounded_channel(); + let (other_tx, mut other) = mpsc::unbounded_channel(); + let (service, commands) = + Service::new(MultiTableIVM::new(), storage, owner_tx.clone()); + let service = service.with_sinks(vec![owner_tx, other_tx]); + spawn_local(service.run()); + + let step = |writes: Vec<_>, at: u64, heartbeat: Option| Transaction { + writes, + at: Lsn(at), + progress: Lsn(at), + watched: Vec::new(), + schema: Vec::new(), + catalog: None, + received: std::time::Instant::now(), + committed_at_micros: 0, + heartbeat, + decode: Duration::ZERO, + }; + commands + .send(Command::Transaction(step(vec![insert(1, "OPEN")], 10, None))) + .await + .expect("send"); + commands + .send(Command::Transaction(step(Vec::new(), 10, Some(3)))) + .await + .expect("send"); + let (owner_seen, other_seen) = (drain(&mut owner).await, drain(&mut other).await); + for seen in [&owner_seen, &other_seen] { + assert_eq!(shapes(seen).first(), Some(&"committed"), "{:?}", shapes(seen)); + assert_eq!(shapes(seen).last(), Some(&"heartbeat"), "{:?}", shapes(seen)); + assert!(matches!(seen.last(), Some(Event::Heartbeat(3)))); + } + }); +} + /// Storage that refuses every read, the way a read past the row budget is. struct Refusing(MemoryStorage); @@ -725,6 +772,7 @@ fn a_serving_consumer_is_not_woken_for_nothing() { catalog: None, received: std::time::Instant::now(), committed_at_micros: 0, + heartbeat: None, decode: Duration::ZERO, } };