diff --git a/Cargo.toml b/Cargo.toml index 04ccfa3b98d27..0b24b58786bee 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -262,6 +262,9 @@ exclude = [ # All WASM crates are split into their own workspace to avoid needles cache # invalidations for the core Mz crates. "misc/wasm/*", + # The subscribe SDK is a standalone, relocatable crate with its own + # workspace and dependency set, kept out of the core workspace. + "misc/subscribe-sdk/*", # Ignore any Rust dependencies that python packages might pull in. "misc/python/venv/*", # The `src/*/fuzz` cargo-fuzz crates need no entry here: each sets diff --git a/misc/subscribe-sdk/.gitignore b/misc/subscribe-sdk/.gitignore new file mode 100644 index 0000000000000..60f1f538e0f54 --- /dev/null +++ b/misc/subscribe-sdk/.gitignore @@ -0,0 +1,11 @@ +# Rust build output. +target/ + +# Python caches, builds, and virtualenvs. +__pycache__/ +*.pyc +.pytest_cache/ +.mypy_cache/ +.ruff_cache/ +*.egg-info/ +.venv/ diff --git a/misc/subscribe-sdk/README.md b/misc/subscribe-sdk/README.md new file mode 100644 index 0000000000000..e0d763d25c164 --- /dev/null +++ b/misc/subscribe-sdk/README.md @@ -0,0 +1,99 @@ +# Materialize Subscribe SDK + +Correct-by-construction clients for consuming Materialize +[`SUBSCRIBE`](https://materialize.com/docs/sql/subscribe/) and building external +sinks on top of it. + +Consuming `SUBSCRIBE` durably is subtle, and the failure modes are invisible in +testing. These SDKs encode the protocol so callers cannot get it wrong. + +> **Status: initial scaffolding.** This directory currently contains the tested +> protocol core and the client entry points for Rust and Python. It is expected +> to graduate to a standalone repository; it lives here under `misc/` as its own +> isolated workspace so relocating it is a move, not a rewrite. The full design +> lives in the Subscribe SDK design doc under `doc/developer/design/` (companion +> PR). + +## Layout + +| Path | What | +| --------- | ----------------------------------------------------------------- | +| `rust/` | `mz-subscribe` crate (its own Cargo workspace) | +| `python/` | `materialize-subscribe` package | + +Each language is a native implementation. They share one behavior, verified by +mirrored test suites, and (deliberately) produce **byte-identical resume +tokens**: a token minted by one SDK resumes correctly in the other. + +## The protocol these encode + +The durable-subscription protocol, and where each piece lives: + +1. Subscribe `WITH (PROGRESS, SNAPSHOT = true)`. — the statement builder always + requests `PROGRESS`. +2. Buffer updates until a progress message proves a timestamp is closed. — the + batcher. +3. Apply the closed batch and persist its frontier atomically. — the consumer, + using the batch's resume token. +4. On restart, resume `WITH (PROGRESS, SNAPSHOT = false) AS OF frontier - 1`. — + the resume token owns the `- 1`; callers never compute a timestamp. +5. Keep `RETAIN HISTORY` on the subscribed object wider than downtime, and + handle the compaction-horizon error otherwise. — a typed error. + +Applying each batch and persisting its token in one transaction yields +exactly-once **state**: after any crash, resuming from the last token neither +drops nor duplicates data. + +## Layers and cohorts + +The client is three layers, each usable on its own. + +1. The raw decoded stream (`subscribe_raw`): timestamped changes and progress + markers, before any consistency policy. This is the composable substrate. +2. The consistency engine: buffer changes and release everything below a + frontier, consolidated to the net effect within the released batch. +3. The closed unit handed to the consumer: a consistent batch for one view, or a + consistent moment for a cohort of many. + +A **cohort** subscribes to several views at once and releases them only up to +the *minimum* closed frontier across all of them, which is the latest instant +final for every view together. A downstream store therefore moves from one +cross-view-consistent state to the next, never a mix of a newer view with a +staler one. A single view is the cohort of one: the same engine, with one +member. + +The transport drains each subscription continuously into a bounded client-side +buffer. Materialize buffers unread subscribe output without bound, so a consumer +that falls behind fails loud (a buffer-overflow error) rather than pushing that +cost onto the server. The client falls over, never the server. + +A cohort carries a second bound. Consistency means a stalled member pins the +joint frontier, so its peers' changes would buffer without limit while waiting +for it. A lag budget caps the total buffered across the cohort and fails loud +(a cohort-lag error) when a laggard falls too far behind, rather than growing +memory unbounded. + +## Building and testing + +Rust: + +``` +cd rust +cargo test +cargo clippy --all-targets -- -D warnings +``` + +Python (from `python/`): + +``` +python -m venv .venv && . .venv/bin/activate +pip install -e '.[dev,client]' +pytest +mypy materialize_subscribe +ruff check . +``` + +The unit suites need no database; they exercise the protocol core against +synthetic streams. Live smoke tests are gated behind a connection string +(`MZ_SUBSCRIBE_TEST_DSN` for Rust; a running Materialize for Python) and are not +run by default. diff --git a/misc/subscribe-sdk/rust/Cargo.lock b/misc/subscribe-sdk/rust/Cargo.lock new file mode 100644 index 0000000000000..e2c5e8091db37 --- /dev/null +++ b/misc/subscribe-sdk/rust/Cargo.lock @@ -0,0 +1,818 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "async-trait" +version = "0.1.89" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9035ad2d096bed7955a320ee7e2230574d28fd3c3a0f186cbea1ff3c7eed5dbb" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "base64" +version = "0.22.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" + +[[package]] +name = "bitflags" +version = "2.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8" + +[[package]] +name = "block-buffer" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d2f6c7dbe95a6ed67ad9f18e57daf93a2f034c524b99fd2b76d18fdfeb6660aa" +dependencies = [ + "hybrid-array", +] + +[[package]] +name = "bumpalo" +version = "3.20.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" + +[[package]] +name = "byteorder" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" + +[[package]] +name = "bytes" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ae3f5d315924270530207e2a68396c3cc547f6dca3fbdca317cfb1a51edb593" + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "chacha20" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d524456ba66e72eb8b115ff89e01e497f8e6d11d78b70b1aa13c0fbd97540a81" +dependencies = [ + "cfg-if", + "cpufeatures", + "rand_core", +] + +[[package]] +name = "cmov" +version = "0.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c9ea0ac24bc397ab3c98583a3c9ba74fa56b09a4449bbe172b9b1ddb016027a" + +[[package]] +name = "const-oid" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6ef517f0926dd24a1582492c791b6a4818a4d94e789a334894aa15b0d12f55c" + +[[package]] +name = "cpufeatures" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b2a41393f66f16b0823bb79094d54ac5fbd34ab292ddafb9a0456ac9f87d201" +dependencies = [ + "libc", +] + +[[package]] +name = "crypto-common" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce6e4c961d6cd6c9a86db418387425e8bdeaf05b3c8bc1411e6dca4c252f1453" +dependencies = [ + "hybrid-array", +] + +[[package]] +name = "ctutils" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d5515a3834141de9eafb9717ad39eea8247b5674e6066c404e8c4b365d2a29e" +dependencies = [ + "cmov", +] + +[[package]] +name = "digest" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f1dd6dbb5841937940781866fa1281a1ff7bd3bf827091440879f9994983d5c2" +dependencies = [ + "block-buffer", + "const-oid", + "crypto-common", + "ctutils", +] + +[[package]] +name = "fallible-iterator" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4443176a9f2c162692bd3d352d745ef9413eec5782a80d8fd6f8a1ac692a07f7" + +[[package]] +name = "futures-channel" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "07bbe89c50d7a535e539b8c17bc0b49bdb77747034daa8087407d655f3f7cc1d" +dependencies = [ + "futures-core", + "futures-sink", +] + +[[package]] +name = "futures-core" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" + +[[package]] +name = "futures-sink" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c39754e157331b013978ec91992bde1ac089843443c49cbc7f46150b0fad0893" + +[[package]] +name = "futures-task" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "037711b3d59c33004d3856fbdc83b99d4ff37a24768fa1be9ce3538a1cde4393" + +[[package]] +name = "futures-util" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6" +dependencies = [ + "futures-core", + "futures-sink", + "futures-task", + "pin-project-lite", + "slab", +] + +[[package]] +name = "getrandom" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" +dependencies = [ + "cfg-if", + "libc", + "r-efi", + "rand_core", +] + +[[package]] +name = "hmac" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6303bc9732ae41b04cb554b844a762b4115a61bfaa81e3e83050991eeb56863f" +dependencies = [ + "digest", +] + +[[package]] +name = "hybrid-array" +version = "0.4.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "818356c5132c1fede50f837ca96afbe78ff42413047f4abb886217845e1b6c8c" +dependencies = [ + "typenum", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "js-sys" +version = "0.3.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53b44bfcdb3f8d5837a46dae1ca9660a837176eee74a28b229bc626816589102" +dependencies = [ + "cfg-if", + "futures-util", + "wasm-bindgen", +] + +[[package]] +name = "libc" +version = "0.2.186" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" + +[[package]] +name = "libredox" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c943259e342f1e06ff2da7a83eabdfe7f92ce10262688dbf1895ff0b3e6e4652" +dependencies = [ + "libc", +] + +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + +[[package]] +name = "log" +version = "0.4.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" + +[[package]] +name = "md-5" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69b6441f590336821bb897fb28fc622898ccceb1d6cea3fde5ea86b090c4de98" +dependencies = [ + "cfg-if", + "digest", +] + +[[package]] +name = "memchr" +version = "2.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "88904434abc2901f197fe8cc55f0445e7ded921dba5911dad2e2b39b48e663c4" + +[[package]] +name = "mio" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "02bd0af71c67b473010cbbc60715ee815645a4dc942899111f494b4b737d6fda" +dependencies = [ + "libc", + "wasi 0.11.1+wasi-snapshot-preview1", + "windows-sys", +] + +[[package]] +name = "mz-subscribe" +version = "0.1.0" +dependencies = [ + "base64", + "serde", + "serde_json", + "thiserror", + "tokio", + "tokio-postgres", +] + +[[package]] +name = "objc2-core-foundation" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a180dd8642fa45cdb7dd721cd4c11b1cadd4929ce112ebd8b9f5803cc79d536" +dependencies = [ + "bitflags", +] + +[[package]] +name = "objc2-system-configuration" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7216bd11cbda54ccabcab84d523dc93b858ec75ecfb3a7d89513fa22464da396" +dependencies = [ + "objc2-core-foundation", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "parking_lot" +version = "0.12.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + +[[package]] +name = "percent-encoding" +version = "2.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" + +[[package]] +name = "phf" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +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-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "postgres-protocol" +version = "0.6.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08808e3c483c46e999108051c78334f473d5adb59d78bb80a1268c7e6aa6c514" +dependencies = [ + "base64", + "byteorder", + "bytes", + "fallible-iterator", + "hmac", + "md-5", + "memchr", + "rand", + "sha2", + "stringprep", +] + +[[package]] +name = "postgres-types" +version = "0.2.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "851ca9db4932932d69f3ea811b1abe63087a0f740a47692619dd40d4899b68be" +dependencies = [ + "bytes", + "fallible-iterator", + "postgres-protocol", +] + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfbc457d0c7a0759a614551b11a6409e5951f6c7537be1f1b7682b9ae9230368" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "r-efi" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" + +[[package]] +name = "rand" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7f5fa3a058cd35567ef9bfa5e75732bee0f9e4c55fa90477bef2dfcdbc4be80" +dependencies = [ + "chacha20", + "getrandom", + "rand_core", +] + +[[package]] +name = "rand_core" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" + +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + +[[package]] +name = "rustversion" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "serde" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "serde_json" +version = "1.0.150" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8014e44b4736ed0538adeecded0fce2a272f22dc9578a7eb6b2d9993c74cfb9" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "sha2" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "446ba717509524cb3f22f17ecc096f10f4822d76ab5c0b9822c5f9c284e825f4" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "siphasher" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ee5873ec9cce0195efcb7a4e9507a04cd49aec9c83d0389df45b1ef7ba2e649" + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "smallvec" +version = "1.15.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90" + +[[package]] +name = "socket2" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52d1cfed4120b4d927bf7c0f86d2087a4a7d6027c906d9f9d525a80573b9be51" +dependencies = [ + "libc", + "windows-sys", +] + +[[package]] +name = "stringprep" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b4df3d392d81bd458a8a621b8bffbd2302a12ffe288a9d931670948749463b1" +dependencies = [ + "unicode-bidi", + "unicode-normalization", + "unicode-properties", +] + +[[package]] +name = "syn" +version = "2.0.118" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b9ae57f904213ebb649ce6895b8a66c66f0203b9319718f69a5612a065b1422" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tinyvec" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3e61e67053d25a4e82c844e8424039d9745781b3fc4f32b8d55ed50f5f667ef3" +dependencies = [ + "tinyvec_macros", +] + +[[package]] +name = "tinyvec_macros" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" + +[[package]] +name = "tokio" +version = "1.52.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fc7f01b389ac15039e4dc9531aa973a135d7a4135281b12d7c1bc79fd57fffe" +dependencies = [ + "bytes", + "libc", + "mio", + "pin-project-lite", + "socket2", + "tokio-macros", + "windows-sys", +] + +[[package]] +name = "tokio-macros" +version = "2.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "385a6cb71ab9ab790c5fe8d67f1645e6c450a7ce006a33de03daa956cf70a496" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tokio-postgres" +version = "0.7.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a528f7d280f6d5b9cd149635c8705b0dd049754bc67d81d31fa25169a93809d3" +dependencies = [ + "async-trait", + "byteorder", + "bytes", + "fallible-iterator", + "futures-channel", + "futures-util", + "log", + "parking_lot", + "percent-encoding", + "phf", + "pin-project-lite", + "postgres-protocol", + "postgres-types", + "rand", + "socket2", + "tokio", + "tokio-util", + "whoami", +] + +[[package]] +name = "tokio-util" +version = "0.7.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ae9cec805b01e8fc3fd2fe289f89149a9b66dd16786abd8b19cfa7b48cb0098" +dependencies = [ + "bytes", + "futures-core", + "futures-sink", + "pin-project-lite", + "tokio", +] + +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + +[[package]] +name = "unicode-bidi" +version = "0.3.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c1cb5db39152898a79168971543b1cb5020dff7fe43c8dc468b0885f5e29df5" + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "unicode-normalization" +version = "0.1.25" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5fd4f6878c9cb28d874b009da9e8d183b5abc80117c40bbd187a1fde336be6e8" +dependencies = [ + "tinyvec", +] + +[[package]] +name = "unicode-properties" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7df058c713841ad818f1dc5d3fd88063241cc61f49f5fbea4b951e8cf5a8d71d" + +[[package]] +name = "wasi" +version = "0.11.1+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" + +[[package]] +name = "wasi" +version = "0.14.7+wasi-0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "883478de20367e224c0090af9cf5f9fa85bed63a95c1abf3afc5c083ebc06e8c" +dependencies = [ + "wasip2", +] + +[[package]] +name = "wasip2" +version = "1.0.4+wasi-0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b67efb37e106e55ce722a510d6b5f9c17f083e5fc79afc2badeb12cc313d9487" +dependencies = [ + "wit-bindgen", +] + +[[package]] +name = "wasite" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "66fe902b4a6b8028a753d5424909b764ccf79b7a209eac9bf97e59cda9f71a42" +dependencies = [ + "wasi 0.14.7+wasi-0.2.4", +] + +[[package]] +name = "wasm-bindgen" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b067c0c11094aef6b7a801c1e34a26affafdf3d051dba08456b868789aaf9a4" +dependencies = [ + "cfg-if", + "once_cell", + "rustversion", + "wasm-bindgen-macro", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-macro" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "167ce5e579f6bcf889c4f7175a8a5a585de84e8ff93976ce393efa5f2837aab1" +dependencies = [ + "quote", + "wasm-bindgen-macro-support", +] + +[[package]] +name = "wasm-bindgen-macro-support" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3997c7839262f4ef12cf90b818d6340c18e80f263f1a94bf157d0ec4420380e" +dependencies = [ + "bumpalo", + "proc-macro2", + "quote", + "syn", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-shared" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc1b4cb0cc549fcf58d7dfc081778139b3d283a081644e833e84682ad71cea24" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "web-sys" +version = "0.3.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8622dcb61c0bcc9fffa6938bed81210af2da9a7e4a1a834b2e37a59b6dfb6141" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "whoami" +version = "2.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "998767ef88740d1f5b0682a9c53c24431453923962269c2db68ee43788c5a40d" +dependencies = [ + "libc", + "libredox", + "objc2-system-configuration", + "wasite", + "web-sys", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "wit-bindgen" +version = "0.57.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" + +[[package]] +name = "zmij" +version = "1.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" diff --git a/misc/subscribe-sdk/rust/Cargo.toml b/misc/subscribe-sdk/rust/Cargo.toml new file mode 100644 index 0000000000000..ddf5b095e9123 --- /dev/null +++ b/misc/subscribe-sdk/rust/Cargo.toml @@ -0,0 +1,31 @@ +# This crate is its own workspace so it stays isolated from the core +# Materialize workspace (its dependency set and lockfile are independent) and +# can be relocated to a standalone repository without churn. +[workspace] + +[package] +name = "mz-subscribe" +version = "0.1.0" +edition = "2021" +rust-version = "1.94.1" +description = "Correct-by-construction client for consuming Materialize SUBSCRIBE." +license = "Apache-2.0" +publish = false + +[dependencies] +base64 = "0.22.1" +serde = { version = "1.0.219", features = ["derive"] } +serde_json = "1.0.150" +thiserror = "2.0.18" +tokio = { version = "1.44", features = ["macros", "rt-multi-thread", "sync"] } +tokio-postgres = "0.7.15" + +[dev-dependencies] +tokio = { version = "1.44", features = ["macros", "rt-multi-thread", "test-util"] } + +[lints.rust] +missing_debug_implementations = "warn" +unused_qualifications = "warn" + +[lints.clippy] +all = { level = "warn", priority = -1 } \ No newline at end of file diff --git a/misc/subscribe-sdk/rust/README.md b/misc/subscribe-sdk/rust/README.md new file mode 100644 index 0000000000000..e2220763eca9a --- /dev/null +++ b/misc/subscribe-sdk/rust/README.md @@ -0,0 +1,67 @@ +# mz-subscribe + +A correct-by-construction Rust client for consuming Materialize +[`SUBSCRIBE`](https://materialize.com/docs/sql/subscribe/). + +```rust +use mz_subscribe::{Subscribe, SubscribeClient}; + +let client = SubscribeClient::connect("postgres://materialize@localhost:6875").await?; +let mut stream = client + .subscribe(Subscribe::object("winning_bids").envelope_upsert(["id"])) + .await?; + +while let Some(batch) = stream.next().await? { + for change in &batch.updates { + // apply `change` to your sink + } + // Persist `batch.resume_token` atomically with the applied changes for + // exactly-once state; hand it back to `client.resume` on restart. +} +``` + +* The unit of consumption is a `ConsistentBatch`, never a bare row. +* Resuming takes an opaque `ResumeToken`; the `AS OF frontier - 1` arithmetic + lives inside it. +* Failures surface as a small, typed `SubscribeError`. +* For a lower-level stream of decoded changes and progress markers, use + `subscribe_raw`. The batcher is built on it. + +## Cohorts + +To stay consistent across several views at once, a `Cohort` releases every +member only up to their shared minimum frontier, so each moment is a genuine +cross-view snapshot. + +```rust +use mz_subscribe::{Cohort, Subscribe}; + +let mut cohort = Cohort::connect( + "postgres://materialize@localhost:6875", + vec![ + ("orders".to_string(), Subscribe::object("orders")), + ("inventory".to_string(), Subscribe::object("inventory")), + ], +) +.await?; + +while let Some(moment) = cohort.next().await? { + for view in &moment.views { + // view.name, view.updates: this view's changes at the joint frontier + } + // moment.resume_token checkpoints the whole cohort at once. +} +``` + +See the [top-level README](../README.md) for the protocol this encodes and the +guarantees it provides. + +## Development + +``` +cargo test +cargo clippy --all-targets -- -D warnings +``` + +Set `MZ_SUBSCRIBE_TEST_DSN` to run the ignored live smoke test against a +Materialize instance. diff --git a/misc/subscribe-sdk/rust/clippy.toml b/misc/subscribe-sdk/rust/clippy.toml new file mode 100644 index 0000000000000..1b86631178074 --- /dev/null +++ b/misc/subscribe-sdk/rust/clippy.toml @@ -0,0 +1,23 @@ +# Copyright Materialize, Inc. and contributors. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License in the LICENSE file at the +# root of this repository, or online at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# This crate is a standalone, relocatable SDK, so it deliberately does not +# inherit the core workspace's clippy policy (which steers internal code toward +# `mz_postgres_util`). It targets standard PostgreSQL wire compatibility via +# upstream `tokio-postgres`, so the empty disallowed lists below override the +# root `clippy.toml` for this directory. +disallowed-methods = [] +disallowed-macros = [] +disallowed-types = [] diff --git a/misc/subscribe-sdk/rust/src/batch.rs b/misc/subscribe-sdk/rust/src/batch.rs new file mode 100644 index 0000000000000..3acd8e04deff9 --- /dev/null +++ b/misc/subscribe-sdk/rust/src/batch.rs @@ -0,0 +1,334 @@ +// Copyright Materialize, Inc. and contributors. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License in the LICENSE file at the +// root of this repository, or online at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! The progress-driven batcher: the state machine that turns a stream of +//! changes and progress markers into *consistent, closed* batches. +//! +//! This is the part every hand-rolled subscribe consumer gets wrong. The rule +//! it enforces: +//! +//! * Buffer changes as they arrive. +//! * When a progress marker advances the frontier to `F`, every change with a +//! timestamp strictly below `F` is now final. Emit exactly those as one +//! batch, tagged with frontier `F` and a resume token for `F`. +//! * Never emit a change before its timestamp is closed, and never emit the +//! same timestamp twice. +//! +//! A consumer that applies each emitted batch and persists its token atomically +//! gets exactly-once *state*: after any crash, resuming from the last token +//! neither drops nor duplicates data. +//! +//! A single view is the one-member case of the [`crate::cohort`] engine: the +//! release frontier is the view's own progress frontier. Both share the +//! [`ReleaseBuffer`] buffer-and-consolidate core. + +use crate::engine::ReleaseBuffer; +use crate::envelope::{Change, Envelope, StreamMessage}; +use crate::error::SubscribeError; +use crate::token::ResumeToken; + +/// A batch of changes for a contiguous, now-closed range of timestamps. +/// +/// Every change in `updates` has a timestamp strictly below `frontier`, and no +/// future batch will contain a timestamp below `frontier`. `resume_token` +/// checkpoints exactly this position. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ConsistentBatch { + /// The finalized changes, consolidated to their net effect within this + /// batch and ordered by first appearance. + pub updates: Vec, + /// The closed frontier: everything below it is present in this or an + /// earlier batch. + pub frontier: u64, + /// The checkpoint for this position. Persist it atomically with the effects + /// of `updates` to get exactly-once state. + pub resume_token: ResumeToken, + /// Whether this batch carries the initial snapshot. Only ever true once, + /// and only for a subscription started with a snapshot. + pub is_snapshot: bool, +} + +impl ConsistentBatch { + /// Whether the batch carries no changes. Empty batches still advance the + /// frontier and carry a fresh token, which lets a consumer keep its + /// checkpoint moving during idle periods so it does not age out of the + /// source's retained history. + pub fn is_empty(&self) -> bool { + self.updates.is_empty() + } +} + +/// Accumulates [`StreamMessage`]s and emits [`ConsistentBatch`]es as timestamps +/// close. +/// +/// Feed every decoded message to [`Batcher::push`]. It returns `Some(batch)` +/// when a progress marker closes one or more timestamps, and `None` otherwise. +#[derive(Debug)] +pub struct Batcher { + fingerprint: String, + with_snapshot: bool, + /// The buffer-and-consolidate core, shared with the cohort engine. + buffer: ReleaseBuffer, + /// The last frontier emitted, or `None` before the first progress marker. + last_frontier: Option, + /// The first frontier seen, which is the subscription's effective `AS OF`. + /// The snapshot is the data closed just after this point. + first_frontier: Option, + /// Whether the snapshot batch has been emitted yet. + snapshot_emitted: bool, +} + +impl Batcher { + /// Creates a batcher for a subscription with the given query `fingerprint` + /// and `envelope` (which selects how a batch is consolidated). + /// `with_snapshot` must match the subscription: `true` for an initial + /// subscribe, `false` for a resume. + pub fn new(fingerprint: impl Into, envelope: Envelope, with_snapshot: bool) -> Self { + Batcher { + fingerprint: fingerprint.into(), + with_snapshot, + buffer: ReleaseBuffer::new(envelope), + last_frontier: None, + first_frontier: None, + snapshot_emitted: false, + } + } + + /// Feeds one decoded message. Returns a batch when the frontier advances. + pub fn push( + &mut self, + message: StreamMessage, + ) -> Result, SubscribeError> { + match message { + StreamMessage::Data { timestamp, change } => { + self.buffer.push_data(timestamp, change)?; + Ok(None) + } + StreamMessage::Progress { frontier } => self.advance(frontier), + } + } + + fn advance(&mut self, frontier: u64) -> Result, SubscribeError> { + if self.last_frontier == Some(frontier) { + // No advance: nothing new is closed. A regression, by contrast, is + // caught by `observe_progress` below. + return Ok(None); + } + self.buffer.observe_progress(frontier)?; + + if self.first_frontier.is_none() { + self.first_frontier = Some(frontier); + } + + let updates = self.buffer.release_below(frontier); + + // The snapshot is the data closed once we advance past the first + // frontier (the effective `AS OF`). Mark the first such batch. + // + // NOTE: this relies on the server emitting a leading `Progress(as_of)` + // before any snapshot rows (it does, unconditionally, when `PROGRESS` + // is set). That progress sets `first_frontier`; the snapshot rows sit + // at `as_of` and are closed by the next, higher progress. A transport + // that delivered snapshot data before any progress would mistag it. + let is_snapshot = + self.with_snapshot && !self.snapshot_emitted && Some(frontier) != self.first_frontier; + if is_snapshot { + self.snapshot_emitted = true; + } + + self.last_frontier = Some(frontier); + + Ok(Some(ConsistentBatch { + updates, + frontier, + resume_token: ResumeToken::new(frontier, self.fingerprint.clone()), + is_snapshot, + })) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn insert(id: &str) -> Change { + Change::Diff { + row: vec![Some(id.to_string())], + diff: 1, + } + } + + fn data(timestamp: u64, change: Change) -> StreamMessage { + StreamMessage::Data { timestamp, change } + } + + fn progress(frontier: u64) -> StreamMessage { + StreamMessage::Progress { frontier } + } + + #[test] + fn buffers_until_the_timestamp_closes() { + let mut b = Batcher::new("fp", Envelope::Diff, false); + // Data at t=5 is buffered, not emitted. + assert_eq!(b.push(data(5, insert("a"))).unwrap(), None); + assert_eq!(b.push(data(5, insert("b"))).unwrap(), None); + // Progress to 6 closes t=5: both changes emitted together. + let batch = b.push(progress(6)).unwrap().expect("batch"); + assert_eq!(batch.updates, vec![insert("a"), insert("b")]); + assert_eq!(batch.frontier, 6); + assert_eq!(batch.resume_token.frontier(), 6); + } + + #[test] + fn closes_a_timestamp_with_mixed_retractions_and_multiplicities() { + // A single timestamp can hold inserts, retractions, and multiplicities + // beyond one. Distinct rows all survive consolidation, in first-seen + // order. + let mut b = Batcher::new("fp", Envelope::Diff, false); + let insert3 = Change::Diff { + row: vec![Some("a".into())], + diff: 3, + }; + let retract = Change::Diff { + row: vec![Some("b".into())], + diff: -1, + }; + b.push(data(5, insert3.clone())).unwrap(); + b.push(data(5, retract.clone())).unwrap(); + b.push(data(5, insert("c"))).unwrap(); + let batch = b.push(progress(6)).unwrap().expect("batch"); + assert_eq!(batch.updates, vec![insert3, retract, insert("c")]); + } + + #[test] + fn batch_is_consolidated_to_its_net_effect() { + // Repeated changes to the same row within a batch collapse: two inserts + // and a retract of `a` net to +1, and `b` inserted then retracted + // disappears entirely. + let mut b = Batcher::new("fp", Envelope::Diff, false); + b.push(data(5, insert("a"))).unwrap(); + b.push(data(5, insert("a"))).unwrap(); + b.push(data(6, insert("b"))).unwrap(); + b.push(data( + 6, + Change::Diff { + row: vec![Some("a".into())], + diff: -1, + }, + )) + .unwrap(); + b.push(data( + 6, + Change::Diff { + row: vec![Some("b".into())], + diff: -1, + }, + )) + .unwrap(); + let batch = b.push(progress(7)).unwrap().expect("batch"); + assert_eq!( + batch.updates, + vec![Change::Diff { + row: vec![Some("a".into())], + diff: 1, + }] + ); + } + + #[test] + fn data_at_the_frontier_stays_open() { + let mut b = Batcher::new("fp", Envelope::Diff, false); + b.push(data(5, insert("a"))).unwrap(); + b.push(data(6, insert("b"))).unwrap(); + // Progress to 6 closes t=5 only; the t=6 change is still open. + let batch = b.push(progress(6)).unwrap().expect("batch"); + assert_eq!(batch.updates, vec![insert("a")]); + // Progress to 7 now closes t=6. + let batch = b.push(progress(7)).unwrap().expect("batch"); + assert_eq!(batch.updates, vec![insert("b")]); + } + + #[test] + fn empty_batches_still_advance_the_token() { + let mut b = Batcher::new("fp", Envelope::Diff, false); + let batch = b.push(progress(10)).unwrap().expect("batch"); + assert!(batch.is_empty()); + assert_eq!(batch.frontier, 10); + assert_eq!(batch.resume_token.frontier(), 10); + } + + #[test] + fn duplicate_frontier_emits_nothing() { + let mut b = Batcher::new("fp", Envelope::Diff, false); + assert!(b.push(progress(10)).unwrap().is_some()); + // A repeated progress at the same frontier is not an advance. + assert_eq!(b.push(progress(10)).unwrap(), None); + } + + #[test] + fn snapshot_batch_is_the_first_past_the_initial_frontier() { + // Initial subscribe: progress(as_of) then snapshot data at as_of, then + // progress advances past it. + let mut b = Batcher::new("fp", Envelope::Diff, true); + // First progress reveals as_of = 100; leading batch is not the snapshot. + let leading = b.push(progress(100)).unwrap().expect("batch"); + assert!(!leading.is_snapshot); + assert!(leading.is_empty()); + // Snapshot rows arrive at t=100. + b.push(data(100, insert("snap1"))).unwrap(); + b.push(data(100, insert("snap2"))).unwrap(); + // Advancing past 100 closes the snapshot; this batch is the snapshot. + let snap = b.push(progress(101)).unwrap().expect("batch"); + assert!(snap.is_snapshot); + assert_eq!(snap.updates, vec![insert("snap1"), insert("snap2")]); + // Later batches are not snapshots. + b.push(data(101, insert("live"))).unwrap(); + let live = b.push(progress(102)).unwrap().expect("batch"); + assert!(!live.is_snapshot); + } + + #[test] + fn resume_never_marks_a_snapshot() { + // A resumed subscription (with_snapshot = false) never flags a snapshot. + let mut b = Batcher::new("fp", Envelope::Diff, false); + b.push(progress(100)).unwrap(); + b.push(data(100, insert("x"))).unwrap(); + let batch = b.push(progress(101)).unwrap().expect("batch"); + assert!(!batch.is_snapshot); + } + + #[test] + fn late_data_below_the_frontier_is_a_protocol_error() { + let mut b = Batcher::new("fp", Envelope::Diff, false); + b.push(progress(10)).unwrap(); + let err = b.push(data(5, insert("late"))).unwrap_err(); + assert!(matches!(err, SubscribeError::Protocol(_)), "{err:?}"); + } + + #[test] + fn regressing_frontier_is_a_protocol_error() { + let mut b = Batcher::new("fp", Envelope::Diff, false); + b.push(progress(10)).unwrap(); + let err = b.push(progress(9)).unwrap_err(); + assert!(matches!(err, SubscribeError::Protocol(_)), "{err:?}"); + } + + #[test] + fn token_carries_the_fingerprint() { + let mut b = Batcher::new("my-fingerprint", Envelope::Diff, false); + let batch = b.push(progress(1)).unwrap().expect("batch"); + assert_eq!(batch.resume_token.fingerprint(), "my-fingerprint"); + } +} diff --git a/misc/subscribe-sdk/rust/src/client.rs b/misc/subscribe-sdk/rust/src/client.rs new file mode 100644 index 0000000000000..6930474ef48f7 --- /dev/null +++ b/misc/subscribe-sdk/rust/src/client.rs @@ -0,0 +1,480 @@ +// Copyright Materialize, Inc. and contributors. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License in the LICENSE file at the +// root of this repository, or online at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! The pgwire transport: connect to Materialize and drive a subscription. +//! +//! The transport is deliberately thin. It runs the `DECLARE`/`FETCH` loop and +//! reads columns in their text encoding via the simple-query protocol, then +//! hands every row to the tested [`Decoder`]. All protocol judgement lives in +//! the decoder and the consistency engine, not here. +//! +//! A background task drains the cursor continuously into a *bounded* buffer, +//! rather than fetching only when the consumer asks. This is deliberate: +//! Materialize buffers a subscription's unread output in `environmentd` without +//! bound, so a consumer that stops fetching pushes an unbounded cost onto the +//! server. Draining continuously keeps that buffer on the client, where it is +//! bounded and fails loud ([`SubscribeError::BufferOverflow`]) if the consumer +//! cannot keep up. The client falls over, never the server. +//! +//! TLS is not yet wired, so this connects over plaintext (suitable for the +//! emulator or a `sslmode=disable` local server). A TLS connector is the next +//! step for this module. + +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex}; + +use tokio::sync::mpsc; +use tokio_postgres::{NoTls, SimpleQueryMessage}; + +use crate::batch::{Batcher, ConsistentBatch}; +use crate::envelope::{Datum, Decoder, Envelope, StreamMessage}; +use crate::error::SubscribeError; +use crate::statement::Subscribe; +use crate::token::ResumeToken; + +/// The cursor name used for the subscription. One subscription owns its +/// connection, so a fixed name is safe. +const CURSOR: &str = "mz_subscribe_cursor"; + +/// Default client-side buffer capacity, in decoded messages. A slow consumer +/// that falls this far behind trips [`SubscribeError::BufferOverflow`]. +/// +/// TODO: make this configurable per subscription once we have real workloads to +/// size it against. +pub(crate) const DEFAULT_BUFFER_CAPACITY: usize = 1 << 16; + +/// A connection to Materialize for consuming subscriptions. +/// +/// One connection backs one subscription at a time. Transaction-mode connection +/// poolers (such as PgBouncer) break the `DECLARE`/`FETCH` cursor and are not +/// supported. +#[derive(Debug)] +pub struct SubscribeClient { + client: Arc, + buffer_capacity: usize, +} + +impl SubscribeClient { + /// Connects to Materialize using a libpq-style connection string. + /// + /// Connects over plaintext for now (see the module docs). The background + /// connection task is spawned onto the current Tokio runtime. + pub async fn connect(conninfo: &str) -> Result { + let (client, connection) = tokio_postgres::connect(conninfo, NoTls) + .await + .map_err(classify_error)?; + // The connection must be driven for the client to make progress. + tokio::spawn(async move { + // A connection error here surfaces on the next client call as a + // closed connection, which `classify_error` maps to `Transient`. + let _ = connection.await; + }); + Ok(SubscribeClient { + client: Arc::new(client), + buffer_capacity: DEFAULT_BUFFER_CAPACITY, + }) + } + + /// Overrides the client-side buffer capacity, in decoded messages. Larger + /// buffers tolerate burstier consumers at the cost of memory. + pub fn buffer_capacity(mut self, messages: usize) -> Self { + self.buffer_capacity = messages.max(1); + self + } + + /// Starts a new subscription, taking the initial snapshot. + pub async fn subscribe(&self, subscribe: Subscribe) -> Result { + let batcher = Batcher::new( + subscribe.fingerprint(), + subscribe.envelope().clone(), + /* with_snapshot */ true, + ); + let raw = self + .open_raw(&subscribe, subscribe.to_sql_initial()) + .await?; + Ok(BatchStream { raw, batcher }) + } + + /// Resumes a subscription from `token`, skipping the snapshot and starting + /// exactly where the token left off. + /// + /// Fails with [`SubscribeError::SchemaMismatch`] if `subscribe` does not + /// match the query the token was taken against, since resuming would + /// otherwise mix incompatible results. + pub async fn resume( + &self, + subscribe: Subscribe, + token: &ResumeToken, + ) -> Result { + check_fingerprint(&subscribe, token)?; + let batcher = Batcher::new( + subscribe.fingerprint(), + subscribe.envelope().clone(), + /* with_snapshot */ false, + ); + let raw = self + .open_raw(&subscribe, subscribe.to_sql_resume(token)) + .await?; + Ok(BatchStream { raw, batcher }) + } + + /// Starts a subscription and hands back the *raw* decoded stream: the + /// timestamped changes and progress markers, before any batching. + /// + /// This is the composable substrate. Most callers want [`subscribe`], which + /// layers consistent batching on top. Reach for the raw stream to build a + /// different consistency policy, or to feed a [`crate::Cohort`]-style + /// multi-view engine of your own. + /// + /// [`subscribe`]: SubscribeClient::subscribe + pub async fn subscribe_raw(&self, subscribe: Subscribe) -> Result { + self.open_raw(&subscribe, subscribe.to_sql_initial()).await + } + + /// Resumes a raw stream from `token`. See [`subscribe_raw`] and [`resume`]. + /// + /// [`subscribe_raw`]: SubscribeClient::subscribe_raw + /// [`resume`]: SubscribeClient::resume + pub async fn resume_raw( + &self, + subscribe: Subscribe, + token: &ResumeToken, + ) -> Result { + check_fingerprint(&subscribe, token)?; + self.open_raw(&subscribe, subscribe.to_sql_resume(token)) + .await + } + + async fn open_raw( + &self, + subscribe: &Subscribe, + subscribe_sql: String, + ) -> Result { + start_cursor(&self.client, &subscribe_sql).await?; + let (tx, rx) = mpsc::channel(self.buffer_capacity); + let status = Arc::new(Mutex::new(None)); + let cancel = Arc::new(AtomicBool::new(false)); + tokio::spawn(run_drain( + Arc::clone(&self.client), + subscribe.envelope().clone(), + subscribe.is_bounded(), + tx, + |message| message, + cancel, + Arc::clone(&status), + )); + Ok(RawStream { rx, status }) + } +} + +/// Rejects a resume whose query shape no longer matches the checkpoint. +fn check_fingerprint(subscribe: &Subscribe, token: &ResumeToken) -> Result<(), SubscribeError> { + let fingerprint = subscribe.fingerprint(); + if fingerprint != token.fingerprint() { + return Err(SubscribeError::SchemaMismatch { + expected: token.fingerprint().to_string(), + actual: fingerprint, + }); + } + Ok(()) +} + +/// The raw decoded stream for one subscription: [`StreamMessage`]s in arrival +/// order, before any batching or consistency policy is applied. +/// +/// This is layer one, the composable substrate the batcher and cohort are built +/// on. Pull messages with [`RawStream::next`]. +#[derive(Debug)] +pub struct RawStream { + rx: mpsc::Receiver, + status: Arc>>, +} + +impl RawStream { + /// Returns the next decoded message, or `None` when the stream ends (only a + /// bounded subscription ends on its own). Surfaces the terminal error if the + /// drain failed. + pub async fn next(&mut self) -> Result, SubscribeError> { + match self.rx.recv().await { + Some(message) => Ok(Some(message)), + None => take_terminal_status(&self.status), + } + } +} + +/// A stream of [`ConsistentBatch`]es for one subscription. +/// +/// Pull batches with [`BatchStream::next`]. Each returned batch is complete: a +/// consumer that applies it and persists its `resume_token` atomically achieves +/// exactly-once state. This is the raw stream plus the [`Batcher`] consistency +/// engine. +#[derive(Debug)] +pub struct BatchStream { + raw: RawStream, + batcher: Batcher, +} + +impl BatchStream { + /// Returns the next consistent batch, or `None` when a bounded subscription + /// (one with `UP TO`) has terminated. + /// + /// Consumes raw messages until a timestamp closes. Idle periods still yield + /// empty batches as the frontier advances. + pub async fn next(&mut self) -> Result, SubscribeError> { + loop { + match self.raw.next().await? { + Some(message) => { + if let Some(batch) = self.batcher.push(message)? { + return Ok(Some(batch)); + } + } + None => return Ok(None), + } + } + } +} + +/// Reads and clears the terminal status of a finished drain: an error if it +/// failed, or a clean end otherwise. +fn take_terminal_status( + status: &Mutex>, +) -> Result, SubscribeError> { + match status.lock().expect("status mutex poisoned").take() { + Some(error) => Err(error), + None => Ok(None), + } +} + +/// The `FETCH` statement, with a short timeout so an idle unbounded subscription +/// still returns periodically to check for cancellation and emit progress. +pub(crate) fn fetch_sql() -> String { + format!("FETCH ALL {CURSOR} WITH (timeout = '1s')") +} + +/// Connects a fresh connection and opens the subscription cursor on it. Used for +/// cohort members, which each own a connection. +pub(crate) async fn connect_cursor( + conninfo: &str, + subscribe_sql: &str, +) -> Result, SubscribeError> { + let (client, connection) = tokio_postgres::connect(conninfo, NoTls) + .await + .map_err(classify_error)?; + tokio::spawn(async move { + let _ = connection.await; + }); + let client = Arc::new(client); + start_cursor(&client, subscribe_sql).await?; + Ok(client) +} + +/// Opens the subscription cursor inside a transaction on an existing connection. +/// +/// A cursor must live inside a transaction, which also holds the read so the +/// frontier does not advance out from under a slow reader mid-batch. If +/// `DECLARE` fails after `BEGIN`, this rolls back so the connection is not left +/// in an open transaction. +async fn start_cursor( + client: &tokio_postgres::Client, + subscribe_sql: &str, +) -> Result<(), SubscribeError> { + client.simple_query("BEGIN").await.map_err(classify_error)?; + if let Err(error) = client + .simple_query(&format!("DECLARE {CURSOR} CURSOR FOR {subscribe_sql}")) + .await + { + let _ = client.simple_query("ROLLBACK").await; + return Err(classify_error(error)); + } + Ok(()) +} + +/// The background drain: fetch, decode, and forward messages into `tx` until the +/// subscription ends, the consumer leaves, or a sibling drain fails. +/// +/// On any terminal error it records the first error in `status` and sets +/// `cancel` so sibling drains (in a cohort) stop too. It always tries to +/// `ROLLBACK` on the way out, releasing the cursor so the server stops buffering +/// output no one will read. +pub(crate) async fn run_drain( + client: Arc, + envelope: Envelope, + bounded: bool, + tx: mpsc::Sender, + wrap: W, + cancel: Arc, + status: Arc>>, +) where + T: Send + 'static, + W: Fn(StreamMessage) -> T + Send + 'static, +{ + if let Err(error) = drain_loop(&client, &envelope, bounded, &tx, &wrap, &cancel).await { + let mut slot = status.lock().expect("status mutex poisoned"); + if slot.is_none() { + *slot = Some(error); + } + cancel.store(true, Ordering::SeqCst); + } + // Best-effort cursor release. If the connection is already broken this is a + // no-op; if it is healthy it stops the server buffering for a dead reader. + let _ = client.simple_query("ROLLBACK").await; +} + +async fn drain_loop( + client: &tokio_postgres::Client, + envelope: &Envelope, + bounded: bool, + tx: &mpsc::Sender, + wrap: &W, + cancel: &AtomicBool, +) -> Result<(), SubscribeError> +where + W: Fn(StreamMessage) -> T, +{ + let fetch = fetch_sql(); + let mut decoder: Option = None; + loop { + // Stop promptly when a sibling failed or the consumer walked away. The + // fetch timeout bounds how long either takes to notice. + if cancel.load(Ordering::SeqCst) || tx.is_closed() { + return Ok(()); + } + + let messages = client.simple_query(&fetch).await.map_err(classify_error)?; + let mut saw_row = false; + for message in messages { + if let SimpleQueryMessage::Row(row) = message { + saw_row = true; + if decoder.is_none() { + let columns: Vec = + row.columns().iter().map(|c| c.name().to_string()).collect(); + decoder = Some(Decoder::new(&columns, envelope)?); + } + let decoder = decoder.as_ref().expect("decoder set above"); + let raw: Vec = (0..row.len()) + .map(|i| row.get(i).map(str::to_string)) + .collect(); + let decoded = decoder.decode(&raw)?; + match tx.try_send(wrap(decoded)) { + Ok(()) => {} + // The consumer dropped the receiver: nothing more to do. + Err(mpsc::error::TrySendError::Closed(_)) => return Ok(()), + // The buffer is full: the consumer fell behind. Fail loud + // rather than block the drain (which would backpressure the + // server into unbounded buffering). + Err(mpsc::error::TrySendError::Full(_)) => { + return Err(SubscribeError::BufferOverflow); + } + } + } + } + + // A fetch that returned no rows on a bounded subscription means it has + // drained. An unbounded one just idled; loop and fetch again. + if !saw_row && bounded { + return Ok(()); + } + } +} + +/// Classifies a `tokio_postgres` error into the SDK's taxonomy. +/// +/// A server error carries a DB payload (a SQLSTATE and message); anything +/// without one is a connection-level failure and therefore retryable. This +/// discriminator mirrors the Python SDK, so the two agree on what is +/// retryable. +fn classify_error(error: tokio_postgres::Error) -> SubscribeError { + match error.as_db_error() { + Some(db) => classify_server_message(db.message()), + // No DB payload means a connection-level failure: retryable. + None => SubscribeError::Transient(error.to_string()), + } +} + +/// Maps a server error message to a typed error. Pure, so it is unit-tested and +/// kept in lockstep with the Python SDK. +/// +/// The compaction-horizon case is matched by message text because Materialize +/// does not yet expose a dedicated SQLSTATE for it. That match lives *only* +/// here, and both the current and the historical wording are handled so the +/// classification survives server upgrades. +fn classify_server_message(message: &str) -> SubscribeError { + if is_compaction_horizon(message) { + SubscribeError::CompactionHorizon + } else if message.contains("dropped") && message.contains("dependenc") { + SubscribeError::DependencyDropped(message.to_string()) + } else { + SubscribeError::Fatal(message.to_string()) + } +} + +fn is_compaction_horizon(message: &str) -> bool { + // Current wording followed by the pre-#34712 wording. + message.contains("could not find a valid timestamp") + || message.contains("not valid for all inputs") +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn classifies_compaction_horizon_wordings() { + assert!(is_compaction_horizon( + "could not find a valid timestamp for the query" + )); + assert!(is_compaction_horizon( + "Timestamp (123) is not valid for all inputs" + )); + assert!(!is_compaction_horizon("some unrelated error")); + } + + #[test] + fn classifies_server_messages() { + assert!(matches!( + classify_server_message("could not find a valid timestamp for the query"), + SubscribeError::CompactionHorizon + )); + assert!(matches!( + classify_server_message("dependency winning_bids was dropped"), + SubscribeError::DependencyDropped(_) + )); + assert!(matches!( + classify_server_message("column \"x\" does not exist"), + SubscribeError::Fatal(_) + )); + } + + /// End-to-end smoke test against a real Materialize. Ignored by default; + /// run with `MZ_SUBSCRIBE_TEST_DSN` set to exercise the transport. + #[tokio::test] + #[ignore = "requires a running Materialize; set MZ_SUBSCRIBE_TEST_DSN"] + async fn live_subscribe_smoke() { + let dsn = std::env::var("MZ_SUBSCRIBE_TEST_DSN").expect("MZ_SUBSCRIBE_TEST_DSN"); + let client = SubscribeClient::connect(&dsn).await.expect("connect"); + let mut stream = client + .subscribe(Subscribe::query("SELECT 1 AS n")) + .await + .expect("subscribe"); + // The first non-empty batch should carry the snapshot. + loop { + let batch = stream.next().await.expect("batch").expect("not terminated"); + if !batch.is_empty() { + assert!(batch.is_snapshot); + break; + } + } + } +} diff --git a/misc/subscribe-sdk/rust/src/cohort.rs b/misc/subscribe-sdk/rust/src/cohort.rs new file mode 100644 index 0000000000000..ebf4d03b0032f --- /dev/null +++ b/misc/subscribe-sdk/rust/src/cohort.rs @@ -0,0 +1,529 @@ +// Copyright Materialize, Inc. and contributors. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License in the LICENSE file at the +// root of this repository, or online at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Multi-view consistency: subscribe to several views at once and observe them +//! at one shared, consistent moment. +//! +//! Materialize gives every object a place on one logical timeline, so an +//! `mz_timestamp` means the same instant across views. A cohort exploits that: +//! it runs one `SUBSCRIBE` per view (each on its own connection) and releases +//! every view only up to the *minimum* closed frontier across all of them. That +//! minimum is the latest instant that is final for every view at once, so the +//! changes it hands back form a genuine cross-view snapshot, never a mix of a +//! newer view with a staler one. +//! +//! This is the same engine as the single-view [`crate::batch::Batcher`], +//! generalized. A single view is the cohort of one: its minimum frontier is its +//! only frontier. The cost of the guarantee is that a laggard member holds the +//! joint moment back, and the leading members' changes buffer in memory until +//! it catches up. +//! +//! Dynamic membership (adding, dropping, merging, or splitting members of a live +//! cohort) is out of scope here. A cohort's membership is fixed for its life. + +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex}; + +use tokio::sync::mpsc; + +use crate::client::{connect_cursor, run_drain, DEFAULT_BUFFER_CAPACITY}; +use crate::engine::ReleaseBuffer; +use crate::envelope::{Change, Envelope, StreamMessage}; +use crate::error::SubscribeError; +use crate::statement::Subscribe; +use crate::token::CohortToken; + +/// Default cohort lag budget: the most changes the cohort buffers across all +/// members while waiting for the slowest to advance the joint frontier. +/// Exceeding it trips [`SubscribeError::CohortLagExceeded`] rather than growing +/// memory without bound. +/// +/// TODO: make this configurable per cohort once there are workloads to size it. +pub(crate) const DEFAULT_COHORT_LAG_BUDGET: usize = 1 << 20; + +/// The changes one view contributes to a [`CohortMoment`]. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ViewChanges { + /// The view's name, as given when the cohort was created. + pub name: String, + /// The view's changes for this moment, consolidated to their net effect. + pub updates: Vec, +} + +/// A consistent moment across every view in a cohort. +/// +/// Every view's `updates` are closed at the same `frontier`, so applying the +/// whole moment moves a downstream store from one cross-view-consistent state to +/// the next. Persist `resume_token` atomically with the effects for +/// exactly-once state across the cohort. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct CohortMoment { + /// Per-view changes, in the order the cohort's members were given. + pub views: Vec, + /// The joint closed frontier shared by every view. + pub frontier: u64, + /// The checkpoint for this joint position. + pub resume_token: CohortToken, + /// Whether this moment completes the cohort's initial snapshot. True once, + /// for the moment that closes the last member's snapshot. + pub is_snapshot: bool, +} + +impl CohortMoment { + /// Whether no view carries any change. Empty moments still advance the + /// frontier and carry a fresh token. + pub fn is_empty(&self) -> bool { + self.views.iter().all(|v| v.updates.is_empty()) + } +} + +/// One member's buffered state within a cohort. +#[derive(Debug)] +struct CohortMember { + name: String, + fingerprint: String, + buffer: ReleaseBuffer, + /// The member's own first frontier, its effective `AS OF`. + first_frontier: Option, +} + +/// The pure state machine behind a cohort: N per-member buffers released at +/// their shared minimum frontier. Kept separate from the transport so it can be +/// exhaustively unit-tested without a server. +#[derive(Debug)] +pub(crate) struct CohortEngine { + members: Vec, + with_snapshot: bool, + /// The last joint frontier released, or `None` before the first release. + last_release: Option, + /// Whether the joint snapshot has been completed yet. + snapshot_complete: bool, + /// The most changes the cohort may buffer across all members while waiting + /// for the slowest to advance the joint frontier. See + /// [`SubscribeError::CohortLagExceeded`]. + lag_budget: usize, +} + +impl CohortEngine { + /// Builds an engine from each member's `(name, fingerprint, envelope)`, in + /// member order. `with_snapshot` is `true` for a fresh cohort, `false` for a + /// resume. `lag_budget` bounds the total buffered across members. + pub(crate) fn new( + members: Vec<(String, String, Envelope)>, + with_snapshot: bool, + lag_budget: usize, + ) -> Self { + let members = members + .into_iter() + .map(|(name, fingerprint, envelope)| CohortMember { + name, + fingerprint, + buffer: ReleaseBuffer::new(envelope), + first_frontier: None, + }) + .collect(); + CohortEngine { + members, + with_snapshot, + last_release: None, + snapshot_complete: false, + lag_budget, + } + } + + /// Buffers a change from member `idx`. + /// + /// Fails with [`SubscribeError::CohortLagExceeded`] if the cohort is already + /// holding `lag_budget` changes. This is the laggard backstop: a member that + /// stalls pins the joint frontier, so its peers' changes would otherwise + /// buffer without bound. The client fails loud instead. + pub(crate) fn push_data( + &mut self, + idx: usize, + timestamp: u64, + change: Change, + ) -> Result<(), SubscribeError> { + let buffered: usize = self.members.iter().map(|m| m.buffer.pending()).sum(); + if buffered >= self.lag_budget { + return Err(SubscribeError::CohortLagExceeded { + buffered, + limit: self.lag_budget, + }); + } + self.members[idx].buffer.push_data(timestamp, change) + } + + /// Records progress from member `idx`, and emits a joint moment if the + /// minimum frontier across all members advanced. + pub(crate) fn push_progress( + &mut self, + idx: usize, + frontier: u64, + ) -> Result, SubscribeError> { + let member = &mut self.members[idx]; + member.buffer.observe_progress(frontier)?; + if member.first_frontier.is_none() { + member.first_frontier = Some(frontier); + } + + // The joint frontier is the minimum across members, and is defined only + // once every member has reported at least one progress. Until then a + // silent member could still emit at an earlier timestamp, so nothing is + // safe to release. + let release = match self + .members + .iter() + .map(|m| m.buffer.frontier()) + .collect::>>() + { + Some(frontiers) => frontiers.into_iter().min().expect("cohort is non-empty"), + None => return Ok(None), + }; + // Emit only when the joint frontier actually advances. + if let Some(last) = self.last_release { + if release <= last { + return Ok(None); + } + } + + // Every member has reported, so every `first_frontier` is set. The joint + // snapshot is complete once the release frontier passes the last + // member's `AS OF`. + let max_first = self + .members + .iter() + .map(|m| m.first_frontier.expect("every member reported")) + .max() + .expect("cohort is non-empty"); + + let views = self + .members + .iter_mut() + .map(|m| ViewChanges { + name: m.name.clone(), + updates: m.buffer.release_below(release), + }) + .collect(); + + let was_complete = self.snapshot_complete; + if !was_complete && release > max_first { + self.snapshot_complete = true; + } + let is_snapshot = self.with_snapshot && !was_complete && self.snapshot_complete; + + self.last_release = Some(release); + let fingerprints = self.members.iter().map(|m| m.fingerprint.clone()).collect(); + + Ok(Some(CohortMoment { + views, + frontier: release, + resume_token: CohortToken::new(release, fingerprints), + is_snapshot, + })) + } +} + +/// A live cohort of subscriptions, yielding consistent cross-view moments. +/// +/// Each member runs its own `SUBSCRIBE` on its own connection, so a cohort of N +/// views holds N connections. All members drain into one shared bounded buffer, +/// so if the consumer falls behind, the cohort fails loud with +/// [`SubscribeError::BufferOverflow`] rather than letting the server buffer. +#[derive(Debug)] +pub struct Cohort { + rx: mpsc::Receiver<(usize, StreamMessage)>, + status: Arc>>, + engine: CohortEngine, +} + +impl Cohort { + /// Starts a fresh cohort over the given `members`, each a `(name, + /// subscription)` pair. All members connect with `conninfo` and take their + /// initial snapshot. Names are how a member's changes are identified in each + /// [`CohortMoment`], and need not match the object name. + pub async fn connect( + conninfo: &str, + members: Vec<(String, Subscribe)>, + ) -> Result { + Self::start(conninfo, members, None).await + } + + /// Resumes a cohort from `token`, skipping snapshots and reconstructing the + /// exact joint cut. Fails with [`SubscribeError::SchemaMismatch`] if the + /// members no longer match the token (different shapes or order). + pub async fn resume( + conninfo: &str, + members: Vec<(String, Subscribe)>, + token: &CohortToken, + ) -> Result { + Self::start(conninfo, members, Some(token)).await + } + + async fn start( + conninfo: &str, + members: Vec<(String, Subscribe)>, + token: Option<&CohortToken>, + ) -> Result { + if members.is_empty() { + return Err(SubscribeError::Protocol( + "a cohort needs at least one member".to_string(), + )); + } + if let Some(token) = token { + let fingerprints: Vec = members.iter().map(|(_, s)| s.fingerprint()).collect(); + if fingerprints != token.members() { + return Err(SubscribeError::SchemaMismatch { + expected: token.members().join(","), + actual: fingerprints.join(","), + }); + } + } + + let (tx, rx) = mpsc::channel(DEFAULT_BUFFER_CAPACITY); + let status = Arc::new(Mutex::new(None)); + let cancel = Arc::new(AtomicBool::new(false)); + + let mut specs = Vec::with_capacity(members.len()); + for (idx, (name, subscribe)) in members.iter().enumerate() { + let sql = match token { + None => subscribe.to_sql_initial(), + Some(token) => subscribe.to_sql_resume_at(token.as_of()), + }; + let client = match connect_cursor(conninfo, &sql).await { + Ok(client) => client, + Err(error) => { + // Tear down members already started before propagating. + cancel.store(true, Ordering::SeqCst); + return Err(error); + } + }; + let index = idx; + tokio::spawn(run_drain( + client, + subscribe.envelope().clone(), + subscribe.is_bounded(), + tx.clone(), + move |message| (index, message), + Arc::clone(&cancel), + Arc::clone(&status), + )); + specs.push(( + name.clone(), + subscribe.fingerprint(), + subscribe.envelope().clone(), + )); + } + // Drop the original sender so the channel closes once every member's + // drain has exited. + drop(tx); + + Ok(Cohort { + rx, + status, + engine: CohortEngine::new(specs, token.is_none(), DEFAULT_COHORT_LAG_BUDGET), + }) + } + + /// Returns the next consistent cross-view moment, or `None` when every + /// member has terminated (only bounded members do). + pub async fn next(&mut self) -> Result, SubscribeError> { + loop { + match self.rx.recv().await { + Some((idx, StreamMessage::Data { timestamp, change })) => { + self.engine.push_data(idx, timestamp, change)?; + } + Some((idx, StreamMessage::Progress { frontier })) => { + if let Some(moment) = self.engine.push_progress(idx, frontier)? { + return Ok(Some(moment)); + } + } + None => { + return match self.status.lock().expect("status mutex poisoned").take() { + Some(error) => Err(error), + None => Ok(None), + }; + } + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn diff(id: &str, diff: i64) -> Change { + Change::Diff { + row: vec![Some(id.to_string())], + diff, + } + } + + fn engine(members: &[&str]) -> CohortEngine { + engine_with_budget(members, DEFAULT_COHORT_LAG_BUDGET) + } + + fn engine_with_budget(members: &[&str], lag_budget: usize) -> CohortEngine { + let specs = members + .iter() + .map(|name| (name.to_string(), format!("fp-{name}"), Envelope::Diff)) + .collect(); + CohortEngine::new(specs, true, lag_budget) + } + + #[test] + fn no_release_until_every_member_reports() { + let mut e = engine(&["a", "b"]); + // Only member 0 has reported progress; the joint frontier is undefined. + assert_eq!(e.push_progress(0, 10).unwrap(), None); + // Once member 1 reports, the minimum (5) is known and a moment emits. + let moment = e.push_progress(1, 5).unwrap().expect("moment"); + assert_eq!(moment.frontier, 5); + } + + #[test] + fn releases_at_the_minimum_frontier() { + let mut e = engine(&["a", "b"]); + // Member `a` has data at t=3 and t=7; member `b` at t=4. + e.push_data(0, 3, diff("a3", 1)).unwrap(); + e.push_data(0, 7, diff("a7", 1)).unwrap(); + e.push_data(1, 4, diff("b4", 1)).unwrap(); + // `a` closes to 10, `b` only to 6. The joint frontier is min(10, 6) = 6. + e.push_progress(0, 10).unwrap(); + let moment = e.push_progress(1, 6).unwrap().expect("moment"); + assert_eq!(moment.frontier, 6); + // Below 6: a3 and b4 released; a7 held back for the laggard. + assert_eq!(moment.views[0].updates, vec![diff("a3", 1)]); + assert_eq!(moment.views[1].updates, vec![diff("b4", 1)]); + } + + #[test] + fn laggard_holds_the_joint_moment() { + let mut e = engine(&["fast", "slow"]); + e.push_data(0, 5, diff("x", 1)).unwrap(); + // The fast member races ahead, but the slow member has not reported, so + // nothing releases. + assert_eq!(e.push_progress(0, 100).unwrap(), None); + // The slow member finally reports a low frontier; the joint frontier + // follows it, not the fast member. + let moment = e.push_progress(1, 6).unwrap().expect("moment"); + assert_eq!(moment.frontier, 6); + assert_eq!(moment.views[0].updates, vec![diff("x", 1)]); + assert!(moment.views[1].updates.is_empty()); + } + + #[test] + fn moment_carries_a_token_at_the_joint_frontier() { + let mut e = engine(&["a", "b"]); + e.push_progress(0, 8).unwrap(); + let moment = e.push_progress(1, 5).unwrap().expect("moment"); + assert_eq!(moment.resume_token.frontier(), 5); + assert_eq!(moment.resume_token.members(), ["fp-a", "fp-b"]); + } + + #[test] + fn per_view_changes_are_consolidated() { + let mut e = engine(&["a", "b"]); + // Member `a` inserts then retracts the same row within the window. + e.push_data(0, 3, diff("x", 1)).unwrap(); + e.push_data(0, 4, diff("x", -1)).unwrap(); + e.push_progress(0, 10).unwrap(); + let moment = e.push_progress(1, 6).unwrap().expect("moment"); + // `x` nets to zero and disappears. + assert!(moment.views[0].updates.is_empty()); + } + + #[test] + fn snapshot_completes_once_the_last_member_passes_its_as_of() { + // Members reveal different as-ofs: `a` at 5, `b` at 8. + let mut e = engine(&["a", "b"]); + // Leading moment at min(5, 8) = 5 is not yet the completed snapshot. + let leading = e.push_progress(0, 5).unwrap(); + assert!(leading.is_none()); // b has not reported yet + let leading = e.push_progress(1, 8).unwrap().expect("moment"); + assert_eq!(leading.frontier, 5); + assert!(!leading.is_snapshot); + // Advance the joint frontier past the last as_of (8). That moment + // completes the snapshot. + e.push_progress(0, 9).unwrap(); // a -> 9, min still 8 + let done = e.push_progress(1, 9).unwrap().expect("moment"); + assert_eq!(done.frontier, 9); + assert!(done.is_snapshot); + // Subsequent moments are not snapshots. + e.push_progress(0, 11).unwrap(); + let live = e.push_progress(1, 11).unwrap().expect("moment"); + assert!(!live.is_snapshot); + } + + #[test] + fn single_member_cohort_matches_a_batcher() { + // A cohort of one is the single-view case: the joint frontier is that + // member's own frontier, and the snapshot completes on the first advance + // past its as_of. + let mut e = engine(&["only"]); + let leading = e.push_progress(0, 100).unwrap().expect("moment"); + assert_eq!(leading.frontier, 100); + assert!(!leading.is_snapshot); + e.push_data(0, 100, diff("snap", 1)).unwrap(); + let snap = e.push_progress(0, 101).unwrap().expect("moment"); + assert!(snap.is_snapshot); + assert_eq!(snap.views[0].updates, vec![diff("snap", 1)]); + } + + #[test] + fn protocol_errors_surface_from_members() { + let mut e = engine(&["a", "b"]); + e.push_progress(0, 10).unwrap(); + // Data below a member's own frontier is a protocol violation. + assert!(e.push_data(0, 5, diff("late", 1)).is_err()); + } + + #[test] + fn a_laggard_past_the_lag_budget_fails_loud() { + // Budget of two buffered changes. The `slow` member never reports, so + // the joint frontier never advances and `fast`'s changes pile up. + let mut e = engine_with_budget(&["fast", "slow"], 2); + e.push_data(0, 1, diff("a", 1)).unwrap(); + e.push_data(0, 2, diff("b", 1)).unwrap(); + // The third change would exceed the budget: fail loud rather than buffer + // without bound behind the laggard. + let err = e.push_data(0, 3, diff("c", 1)).unwrap_err(); + assert!( + matches!(err, SubscribeError::CohortLagExceeded { limit: 2, .. }), + "{err:?}" + ); + } + + #[test] + fn buffered_changes_free_the_budget_once_released() { + // Budget of two. Once the joint frontier advances and drains the buffer, + // there is room to buffer again. + let mut e = engine_with_budget(&["a", "b"], 2); + e.push_data(0, 1, diff("x", 1)).unwrap(); + e.push_progress(0, 2).unwrap(); + let moment = e.push_progress(1, 2).unwrap().expect("moment"); + // Joint frontier 2 released x (at t=1); the buffer is empty again. + assert_eq!(moment.frontier, 2); + assert_eq!(moment.views[0].updates, vec![diff("x", 1)]); + // Room to buffer again (data at or above each member's frontier of 2). + e.push_data(0, 2, diff("y", 1)).unwrap(); + e.push_data(1, 2, diff("z", 1)).unwrap(); + // Two buffered again, so the next change trips the budget. + assert!(matches!( + e.push_data(0, 2, diff("w", 1)).unwrap_err(), + SubscribeError::CohortLagExceeded { .. } + )); + } +} diff --git a/misc/subscribe-sdk/rust/src/engine.rs b/misc/subscribe-sdk/rust/src/engine.rs new file mode 100644 index 0000000000000..aba56d5d412e4 --- /dev/null +++ b/misc/subscribe-sdk/rust/src/engine.rs @@ -0,0 +1,274 @@ +// Copyright Materialize, Inc. and contributors. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License in the LICENSE file at the +// root of this repository, or online at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! The consistency engine: buffer timestamped changes, then release everything +//! strictly below a frontier as one consolidated batch. +//! +//! This one mechanism underlies both the single-view batcher and the multi-view +//! cohort. A single view is the one-member case: its release frontier is its own +//! progress frontier. A cohort of N views feeds N buffers and releases every +//! buffer below the *minimum* frontier across them, which is the latest +//! timestamp closed for every member at once. +//! +//! Whatever the release frontier, the released set is *consolidated* before it +//! leaves the engine: within one batch the net effect per row is computed and +//! rows that cancel to nothing are dropped. A batch is therefore a clean net +//! delta at its frontier, not a replay of the intra-window churn. Consolidation +//! is scoped to a single batch, never across batches, so a consumer that applies +//! batches in order still sees every intermediate settled state. + +use std::collections::HashMap; + +use crate::envelope::{Change, Envelope, Row}; +use crate::error::SubscribeError; + +/// A per-member buffer of timestamped changes awaiting release. +/// +/// Holds changes until a frontier closes them, enforcing the two ordering +/// invariants the server guarantees: data never arrives below the frontier, and +/// the frontier never regresses. A violation is a [`SubscribeError::Protocol`], +/// since it means the stream is not what the protocol promises. +#[derive(Debug)] +pub(crate) struct ReleaseBuffer { + envelope: Envelope, + /// Buffered changes with their timestamp, in arrival order. + buffered: Vec<(u64, Change)>, + /// The highest progress frontier observed, or `None` before the first. + frontier: Option, +} + +impl ReleaseBuffer { + pub(crate) fn new(envelope: Envelope) -> Self { + ReleaseBuffer { + envelope, + buffered: Vec::new(), + frontier: None, + } + } + + /// The highest frontier observed so far. + pub(crate) fn frontier(&self) -> Option { + self.frontier + } + + /// The number of changes currently buffered, before consolidation. Used by + /// the cohort to bound how much a laggard can make its peers buffer. + pub(crate) fn pending(&self) -> usize { + self.buffered.len() + } + + /// Buffers a change. Errors if its timestamp is below the observed frontier, + /// which the server promises never happens. + pub(crate) fn push_data( + &mut self, + timestamp: u64, + change: Change, + ) -> Result<(), SubscribeError> { + if let Some(frontier) = self.frontier { + if timestamp < frontier { + return Err(SubscribeError::Protocol(format!( + "change at timestamp {timestamp} arrived after the frontier \ + advanced to {frontier}" + ))); + } + } + self.buffered.push((timestamp, change)); + Ok(()) + } + + /// Records a progress frontier. Errors if it regresses. + pub(crate) fn observe_progress(&mut self, frontier: u64) -> Result<(), SubscribeError> { + if let Some(last) = self.frontier { + if frontier < last { + return Err(SubscribeError::Protocol(format!( + "frontier went backwards from {last} to {frontier}" + ))); + } + } + self.frontier = Some(frontier); + Ok(()) + } + + /// Drains and consolidates every buffered change with a timestamp strictly + /// below `release_frontier`. Changes at or above it stay buffered. + pub(crate) fn release_below(&mut self, release_frontier: u64) -> Vec { + let mut closed = Vec::new(); + let mut still_open = Vec::with_capacity(self.buffered.len()); + for (timestamp, change) in self.buffered.drain(..) { + if timestamp < release_frontier { + closed.push(change); + } else { + still_open.push((timestamp, change)); + } + } + self.buffered = still_open; + consolidate(&self.envelope, closed) + } +} + +/// Collapses a batch of changes to its net effect, preserving first-appearance +/// order so output is deterministic and close to arrival order. +/// +/// The strategy depends on the envelope: diff changes sum their multiplicities +/// per row, upsert changes keep only the last operation per key. +pub(crate) fn consolidate(envelope: &Envelope, changes: Vec) -> Vec { + match envelope { + Envelope::Diff => consolidate_diff(changes), + Envelope::Upsert { .. } => consolidate_upsert(changes), + } +} + +/// Sums signed multiplicities per row, dropping rows whose net diff is zero. +fn consolidate_diff(changes: Vec) -> Vec { + let mut acc: Vec<(Row, i64)> = Vec::new(); + let mut index: HashMap = HashMap::new(); + for change in changes { + // A diff-envelope stream only ever carries `Change::Diff`; any other + // variant is silently ignored rather than misreported. + if let Change::Diff { row, diff } = change { + match index.get(&row) { + Some(&i) => acc[i].1 += diff, + None => { + index.insert(row.clone(), acc.len()); + acc.push((row, diff)); + } + } + } + } + acc.into_iter() + .filter(|(_, diff)| *diff != 0) + .map(|(row, diff)| Change::Diff { row, diff }) + .collect() +} + +/// Keeps the last operation observed for each key, in first-seen key order. A +/// window that upserts then deletes a key nets to the delete, and vice versa. +fn consolidate_upsert(changes: Vec) -> Vec { + let mut order: Vec = Vec::new(); + let mut last: HashMap = HashMap::new(); + for change in changes { + let key = match &change { + Change::Upsert { key, .. } | Change::Delete { key } | Change::KeyViolation { key } => { + key.clone() + } + // An upsert-envelope stream never carries a bare diff. + Change::Diff { .. } => continue, + }; + if last.insert(key.clone(), change).is_none() { + order.push(key); + } + } + order + .into_iter() + .map(|key| { + last.remove(&key) + .expect("key recorded in order was inserted") + }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + fn diff(id: &str, diff: i64) -> Change { + Change::Diff { + row: vec![Some(id.to_string())], + diff, + } + } + + fn upsert(id: &str, value: &str) -> Change { + Change::Upsert { + key: vec![Some(id.to_string())], + value: vec![Some(value.to_string())], + } + } + + fn delete(id: &str) -> Change { + Change::Delete { + key: vec![Some(id.to_string())], + } + } + + #[test] + fn diff_consolidation_sums_multiplicities_per_row() { + let out = consolidate_diff(vec![diff("a", 3), diff("b", -1), diff("a", 2)]); + // `a` sums to 5 in first-seen order, `b` stays -1. + assert_eq!(out, vec![diff("a", 5), diff("b", -1)]); + } + + #[test] + fn diff_consolidation_drops_net_zero_rows() { + let out = consolidate_diff(vec![diff("a", 1), diff("b", 1), diff("a", -1)]); + // `a` cancels out and disappears; `b` remains. + assert_eq!(out, vec![diff("b", 1)]); + } + + #[test] + fn diff_consolidation_preserves_first_seen_order() { + let out = consolidate_diff(vec![diff("c", 1), diff("a", 1), diff("b", 1)]); + assert_eq!(out, vec![diff("c", 1), diff("a", 1), diff("b", 1)]); + } + + #[test] + fn upsert_consolidation_keeps_last_op_per_key() { + let out = consolidate_upsert(vec![upsert("1", "v1"), upsert("2", "x"), upsert("1", "v2")]); + // Key 1 keeps its latest value; key 2 unchanged; first-seen key order. + assert_eq!(out, vec![upsert("1", "v2"), upsert("2", "x")]); + } + + #[test] + fn upsert_consolidation_collapses_upsert_then_delete() { + let out = consolidate_upsert(vec![upsert("1", "v"), delete("1")]); + assert_eq!(out, vec![delete("1")]); + } + + #[test] + fn release_holds_data_at_or_above_the_frontier() { + let mut buf = ReleaseBuffer::new(Envelope::Diff); + buf.push_data(5, diff("a", 1)).unwrap(); + buf.push_data(6, diff("b", 1)).unwrap(); + // Releasing below 6 closes t=5 only. + assert_eq!(buf.release_below(6), vec![diff("a", 1)]); + // t=6 is still buffered until a higher frontier. + assert_eq!(buf.release_below(7), vec![diff("b", 1)]); + } + + #[test] + fn release_consolidates_across_the_window() { + let mut buf = ReleaseBuffer::new(Envelope::Diff); + buf.push_data(5, diff("a", 1)).unwrap(); + buf.push_data(6, diff("a", -1)).unwrap(); + buf.observe_progress(7).unwrap(); + // Insert then retract of the same row across two timestamps nets to + // nothing in one batch. + assert!(buf.release_below(7).is_empty()); + } + + #[test] + fn data_below_the_frontier_is_rejected() { + let mut buf = ReleaseBuffer::new(Envelope::Diff); + buf.observe_progress(10).unwrap(); + assert!(buf.push_data(5, diff("late", 1)).is_err()); + } + + #[test] + fn regressing_frontier_is_rejected() { + let mut buf = ReleaseBuffer::new(Envelope::Diff); + buf.observe_progress(10).unwrap(); + assert!(buf.observe_progress(9).is_err()); + } +} diff --git a/misc/subscribe-sdk/rust/src/envelope.rs b/misc/subscribe-sdk/rust/src/envelope.rs new file mode 100644 index 0000000000000..81830d23e392e --- /dev/null +++ b/misc/subscribe-sdk/rust/src/envelope.rs @@ -0,0 +1,478 @@ +// Copyright Materialize, Inc. and contributors. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License in the LICENSE file at the +// root of this repository, or online at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Turning raw `SUBSCRIBE` rows into typed changes and progress markers. +//! +//! Decoding is by column *name*, not position, because the column layout +//! differs between envelopes (the upsert envelope interleaves key, `mz_state`, +//! and value columns) and because it makes the decoder robust to the exact +//! ordering the server chooses. +//! +//! In this initial version a value is its text encoding, with SQL `NULL` +//! represented as `None`. Typed decoding (numbers, timestamps, arrays) is +//! future work; the text form is unambiguous and matches what the pgwire +//! simple-query protocol delivers. + +use crate::error::SubscribeError; + +/// A single column value: its text encoding, or `None` for SQL `NULL`. +pub type Datum = Option; + +/// A row of [`Datum`]s. +pub type Row = Vec; + +/// How the server shapes the update stream. Chosen when the subscription is +/// created and fixed for its lifetime. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum Envelope { + /// Raw differential updates: each change carries a signed `mz_diff` + /// multiplicity. + Diff, + /// Server-side upsert compaction keyed by the given columns: each change is + /// an insert/update, a delete, or a key violation. + Upsert { + /// The names of the key columns. + key: Vec, + }, +} + +/// A decoded change to the subscribed relation. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum Change { + /// A differential update. `diff` is the signed multiplicity: positive means + /// `diff` copies were inserted, negative means `|diff|` were retracted. The + /// magnitude is preserved rather than collapsed to +/-1. + Diff { + /// The affected row. + row: Row, + /// The signed multiplicity (never zero). + diff: i64, + }, + /// An upsert-envelope insert or update. `value` is the new full row for + /// `key`. + Upsert { + /// The key columns. + key: Row, + /// The value columns. + value: Row, + }, + /// An upsert-envelope delete: `key` no longer has a value. + Delete { + /// The key columns. + key: Row, + }, + /// An upsert-envelope key violation: more than one live value was observed + /// for `key`. Best-effort per the server, so treat it as a signal, not a + /// guarantee. + KeyViolation { + /// The offending key columns. + key: Row, + }, +} + +/// A decoded message from the stream: either a data change at a timestamp, or a +/// progress marker advancing the frontier. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum StreamMessage { + /// A change occurring at `timestamp`. + Data { + /// The `mz_timestamp` of the change. + timestamp: u64, + /// The change itself. + change: Change, + }, + /// A progress marker: no more changes will arrive with a timestamp strictly + /// less than `frontier`. + Progress { + /// The advanced frontier. + frontier: u64, + }, +} + +const COL_TIMESTAMP: &str = "mz_timestamp"; +const COL_PROGRESSED: &str = "mz_progressed"; +const COL_DIFF: &str = "mz_diff"; +const COL_STATE: &str = "mz_state"; + +const STATE_UPSERT: &str = "upsert"; +const STATE_DELETE: &str = "delete"; +const STATE_KEY_VIOLATION: &str = "key_violation"; + +/// Decodes raw `SUBSCRIBE` rows into [`StreamMessage`]s. +/// +/// Built once from the result's column names (which the SDK always requests +/// `WITH (PROGRESS)`), then reused for every row. The layout of the metadata +/// columns is resolved up front so per-row decoding is a handful of indexed +/// reads. +#[derive(Debug)] +pub struct Decoder { + timestamp_idx: usize, + progressed_idx: usize, + kind: DecoderKind, +} + +#[derive(Debug)] +enum DecoderKind { + Diff { + diff_idx: usize, + /// Payload column indices, in relation order. + payload_idxs: Vec, + }, + Upsert { + state_idx: usize, + key_idxs: Vec, + value_idxs: Vec, + }, +} + +impl Decoder { + /// Builds a decoder from the result `columns` and the subscription + /// `envelope`. Fails if a required metadata column is missing, which would + /// indicate the subscription was not created by this SDK. + pub fn new(columns: &[String], envelope: &Envelope) -> Result { + let find = |name: &str| columns.iter().position(|c| c == name); + let require = |name: &str| { + find(name).ok_or_else(|| { + SubscribeError::Protocol(format!("result is missing the {name} column")) + }) + }; + + let timestamp_idx = require(COL_TIMESTAMP)?; + let progressed_idx = require(COL_PROGRESSED)?; + + let kind = match envelope { + Envelope::Diff => { + let diff_idx = require(COL_DIFF)?; + let meta = [timestamp_idx, progressed_idx, diff_idx]; + let payload_idxs = (0..columns.len()).filter(|i| !meta.contains(i)).collect(); + DecoderKind::Diff { + diff_idx, + payload_idxs, + } + } + Envelope::Upsert { key } => { + let state_idx = require(COL_STATE)?; + let mut key_idxs = Vec::with_capacity(key.len()); + for key_col in key { + key_idxs.push(find(key_col).ok_or_else(|| { + SubscribeError::Protocol(format!( + "upsert key column {key_col} is not in the result" + )) + })?); + } + let meta: Vec = [timestamp_idx, progressed_idx, state_idx] + .into_iter() + .chain(key_idxs.iter().copied()) + .collect(); + let value_idxs = (0..columns.len()).filter(|i| !meta.contains(i)).collect(); + DecoderKind::Upsert { + state_idx, + key_idxs, + value_idxs, + } + } + }; + + Ok(Decoder { + timestamp_idx, + progressed_idx, + kind, + }) + } + + /// Decodes a single raw row. + pub fn decode(&self, row: &[Datum]) -> Result { + let timestamp = self.parse_timestamp(row)?; + + if self.is_progress(row)? { + return Ok(StreamMessage::Progress { + frontier: timestamp, + }); + } + + let change = match &self.kind { + DecoderKind::Diff { + diff_idx, + payload_idxs, + } => { + let diff = self.parse_diff(row, *diff_idx)?; + Change::Diff { + row: project(row, payload_idxs), + diff, + } + } + DecoderKind::Upsert { + state_idx, + key_idxs, + value_idxs, + } => { + let key = project(row, key_idxs); + let state = cell(row, *state_idx)?; + match state.as_deref() { + Some(STATE_UPSERT) => Change::Upsert { + key, + value: project(row, value_idxs), + }, + Some(STATE_DELETE) => Change::Delete { key }, + Some(STATE_KEY_VIOLATION) => Change::KeyViolation { key }, + other => { + return Err(SubscribeError::Protocol(format!( + "unexpected mz_state value: {other:?}" + ))); + } + } + } + }; + + Ok(StreamMessage::Data { timestamp, change }) + } + + fn is_progress(&self, row: &[Datum]) -> Result { + // The pgwire text encoding of a boolean is "t" / "f". + match cell(row, self.progressed_idx)?.as_deref() { + Some("t") => Ok(true), + Some("f") | None => Ok(false), + Some(other) => Err(SubscribeError::Protocol(format!( + "unexpected mz_progressed value: {other:?}" + ))), + } + } + + fn parse_timestamp(&self, row: &[Datum]) -> Result { + let raw = cell(row, self.timestamp_idx)? + .as_deref() + .ok_or_else(|| SubscribeError::Protocol("mz_timestamp was NULL".to_string()))?; + raw.parse::().map_err(|e| { + SubscribeError::Protocol(format!("mz_timestamp {raw:?} is not a u64: {e}")) + }) + } + + fn parse_diff(&self, row: &[Datum], diff_idx: usize) -> Result { + let raw = cell(row, diff_idx)? + .as_deref() + .ok_or_else(|| SubscribeError::Protocol("mz_diff was NULL".to_string()))?; + raw.parse::() + .map_err(|e| SubscribeError::Protocol(format!("mz_diff {raw:?} is not an i64: {e}"))) + } +} + +/// Reads a cell, erroring if the row is shorter than the layout expects. +fn cell(row: &[Datum], idx: usize) -> Result<&Datum, SubscribeError> { + row.get(idx).ok_or_else(|| { + SubscribeError::Protocol(format!( + "row has {} columns but column {idx} was expected", + row.len() + )) + }) +} + +/// Projects the given column indices out of a row, in index order. +fn project(row: &[Datum], idxs: &[usize]) -> Row { + idxs.iter() + .map(|&i| row.get(i).cloned().flatten()) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + fn cols(names: &[&str]) -> Vec { + names.iter().map(|s| s.to_string()).collect() + } + + fn datum(s: &str) -> Datum { + Some(s.to_string()) + } + + #[test] + fn diff_decoder_reads_progress() { + let columns = cols(&["mz_timestamp", "mz_progressed", "mz_diff", "id", "name"]); + let decoder = Decoder::new(&columns, &Envelope::Diff).unwrap(); + // A progress row: mz_progressed = t, payload columns NULL. + let row = vec![datum("100"), datum("t"), None, None, None]; + assert_eq!( + decoder.decode(&row).unwrap(), + StreamMessage::Progress { frontier: 100 } + ); + } + + #[test] + fn diff_decoder_preserves_multiplicity_and_sign() { + let columns = cols(&["mz_timestamp", "mz_progressed", "mz_diff", "id"]); + let decoder = Decoder::new(&columns, &Envelope::Diff).unwrap(); + + let insert = vec![datum("7"), datum("f"), datum("3"), datum("42")]; + assert_eq!( + decoder.decode(&insert).unwrap(), + StreamMessage::Data { + timestamp: 7, + change: Change::Diff { + row: vec![datum("42")], + diff: 3, + }, + } + ); + + let retract = vec![datum("7"), datum("f"), datum("-2"), datum("42")]; + assert_eq!( + decoder.decode(&retract).unwrap(), + StreamMessage::Data { + timestamp: 7, + change: Change::Diff { + row: vec![datum("42")], + diff: -2, + }, + } + ); + } + + #[test] + fn upsert_decoder_handles_all_states() { + // Layout mirrors the server's upsert envelope: key first, then + // mz_state, mz_timestamp, value columns, mz_progressed. + let columns = cols(&[ + "id", + "mz_state", + "mz_timestamp", + "name", + "email", + "mz_progressed", + ]); + let envelope = Envelope::Upsert { + key: vec!["id".to_string()], + }; + let decoder = Decoder::new(&columns, &envelope).unwrap(); + + let upsert = vec![ + datum("1"), + datum("upsert"), + datum("50"), + datum("alice"), + datum("a@x"), + datum("f"), + ]; + assert_eq!( + decoder.decode(&upsert).unwrap(), + StreamMessage::Data { + timestamp: 50, + change: Change::Upsert { + key: vec![datum("1")], + value: vec![datum("alice"), datum("a@x")], + }, + } + ); + + let delete = vec![ + datum("1"), + datum("delete"), + datum("51"), + None, + None, + datum("f"), + ]; + assert_eq!( + decoder.decode(&delete).unwrap(), + StreamMessage::Data { + timestamp: 51, + change: Change::Delete { + key: vec![datum("1")], + }, + } + ); + + let violation = vec![ + datum("1"), + datum("key_violation"), + datum("52"), + None, + None, + datum("f"), + ]; + assert_eq!( + decoder.decode(&violation).unwrap(), + StreamMessage::Data { + timestamp: 52, + change: Change::KeyViolation { + key: vec![datum("1")], + }, + } + ); + } + + #[test] + fn upsert_decoder_reads_progress() { + let columns = cols(&["id", "mz_state", "mz_timestamp", "name", "mz_progressed"]); + let envelope = Envelope::Upsert { + key: vec!["id".to_string()], + }; + let decoder = Decoder::new(&columns, &envelope).unwrap(); + let row = vec![None, None, datum("200"), None, datum("t")]; + assert_eq!( + decoder.decode(&row).unwrap(), + StreamMessage::Progress { frontier: 200 } + ); + } + + #[test] + fn multi_column_key_is_projected_in_order() { + let columns = cols(&[ + "region", + "id", + "mz_state", + "mz_timestamp", + "total", + "mz_progressed", + ]); + let envelope = Envelope::Upsert { + key: vec!["region".to_string(), "id".to_string()], + }; + let decoder = Decoder::new(&columns, &envelope).unwrap(); + let row = vec![ + datum("us"), + datum("9"), + datum("upsert"), + datum("5"), + datum("100"), + datum("f"), + ]; + assert_eq!( + decoder.decode(&row).unwrap(), + StreamMessage::Data { + timestamp: 5, + change: Change::Upsert { + key: vec![datum("us"), datum("9")], + value: vec![datum("100")], + }, + } + ); + } + + #[test] + fn missing_metadata_column_is_rejected() { + let columns = cols(&["mz_timestamp", "mz_diff", "id"]); + let err = Decoder::new(&columns, &Envelope::Diff).unwrap_err(); + assert!(matches!(err, SubscribeError::Protocol(_)), "{err:?}"); + } + + #[test] + fn unknown_key_column_is_rejected() { + let columns = cols(&["id", "mz_state", "mz_timestamp", "mz_progressed"]); + let envelope = Envelope::Upsert { + key: vec!["nonexistent".to_string()], + }; + assert!(Decoder::new(&columns, &envelope).is_err()); + } +} diff --git a/misc/subscribe-sdk/rust/src/error.rs b/misc/subscribe-sdk/rust/src/error.rs new file mode 100644 index 0000000000000..655fb08fbcdb8 --- /dev/null +++ b/misc/subscribe-sdk/rust/src/error.rs @@ -0,0 +1,143 @@ +// Copyright Materialize, Inc. and contributors. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License in the LICENSE file at the +// root of this repository, or online at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! The error taxonomy consumers match on to decide how to recover. + +/// An error surfaced while subscribing. +/// +/// Variants are grouped by the recovery a caller should attempt. +/// [`SubscribeError::is_retryable`] captures the common "reconnect and resume" +/// versus "give up" split so callers rarely need to match every variant. +#[derive(Debug, thiserror::Error)] +#[non_exhaustive] +pub enum SubscribeError { + /// The requested resume point is older than the source's retained history, + /// so the gap cannot be filled. Recover by re-snapshotting or by accepting + /// a gap, per the caller's policy. Never retryable as-is. + /// + /// Widen `RETAIN HISTORY` on the subscribed object to make a larger resume + /// window available. + #[error( + "resume point is older than the subscribed object's retained history \ + (compaction horizon); re-snapshot or widen RETAIN HISTORY" + )] + CompactionHorizon, + + /// A dependency of the subscribed object (the object itself, or its + /// cluster) was dropped, so the stream ended. Not retryable against the + /// same object. + #[error("a dependency of the subscribed object was dropped: {0}")] + DependencyDropped(String), + + /// A resumed subscription's query fingerprint does not match the one the + /// checkpoint was taken against. Resuming would silently mix incompatible + /// results, so it is refused. + #[error("resumed query does not match the checkpoint (expected {expected}, got {actual})")] + SchemaMismatch { + /// Fingerprint recorded in the resume token. + expected: String, + /// Fingerprint of the query being resumed. + actual: String, + }, + + /// A transient connection or protocol interruption. Retryable: reconnect + /// and resume from the last token. + #[error("transient connection error: {0}")] + Transient(String), + + /// The consumer fell too far behind and the client-side buffer filled. The + /// client drains the server continuously so the server never buffers + /// unboundedly, which means a slow consumer must fail here rather than push + /// that cost onto Materialize. Recover by consuming faster, widening the + /// buffer, or resuming from the last token once caught up. + #[error( + "client-side buffer overflowed: the consumer fell behind the subscription. \ + Consume faster or widen the buffer, then resume from the last token" + )] + BufferOverflow, + + /// A cohort member lagged so far behind that the changes buffered for its + /// peers, waiting on it, exceeded the lag budget. Consistency requires + /// holding a leading member's changes until the slowest catches up, so a + /// stalled member would otherwise grow memory without bound. Recover by + /// investigating the slow view, widening the budget, or resuming from the + /// last token once it recovers. + #[error( + "cohort lag budget exceeded: {buffered} changes buffered waiting on a \ + lagging member (budget {limit}). Investigate the slow view or widen the budget" + )] + CohortLagExceeded { + /// Number of changes buffered across the cohort when the budget was hit. + buffered: usize, + /// The configured lag budget. + limit: usize, + }, + + /// A resume token could not be decoded. + #[error("invalid resume token: {0}")] + InvalidToken(String), + + /// The server sent a stream that violates the subscribe protocol (e.g. a + /// non-monotonic timestamp). Indicates a bug, not a recoverable condition. + #[error("subscribe protocol violation: {0}")] + Protocol(String), + + /// A non-retryable server error: an authentication failure, a malformed + /// query, or a dataflow error that poisons the stream (Materialize repeats + /// such an error until the data or view changes, so retrying will not help). + #[error("fatal error: {0}")] + Fatal(String), +} + +impl SubscribeError { + /// Whether reconnecting and resuming from the last token is a sensible + /// automatic response. Only [`SubscribeError::Transient`] qualifies; every + /// other variant needs an explicit policy decision or is a bug. + pub fn is_retryable(&self) -> bool { + matches!(self, SubscribeError::Transient(_)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn only_transient_is_retryable() { + assert!(SubscribeError::Transient("reset".into()).is_retryable()); + assert!(!SubscribeError::CompactionHorizon.is_retryable()); + assert!(!SubscribeError::DependencyDropped("v".into()).is_retryable()); + assert!(!SubscribeError::SchemaMismatch { + expected: "a".into(), + actual: "b".into(), + } + .is_retryable()); + assert!(!SubscribeError::Fatal("auth".into()).is_retryable()); + assert!(!SubscribeError::BufferOverflow.is_retryable()); + assert!(!SubscribeError::CohortLagExceeded { + buffered: 3, + limit: 2 + } + .is_retryable()); + } + + #[test] + fn display_is_actionable() { + // The compaction-horizon message must point the user at the fix, since + // it is the error operators hit most and least expect. + let msg = SubscribeError::CompactionHorizon.to_string(); + assert!(msg.contains("RETAIN HISTORY"), "{msg}"); + } +} diff --git a/misc/subscribe-sdk/rust/src/lib.rs b/misc/subscribe-sdk/rust/src/lib.rs new file mode 100644 index 0000000000000..5489f6c0e4ad9 --- /dev/null +++ b/misc/subscribe-sdk/rust/src/lib.rs @@ -0,0 +1,96 @@ +// Copyright Materialize, Inc. and contributors. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License in the LICENSE file at the +// root of this repository, or online at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! A correct-by-construction client for consuming Materialize +//! [`SUBSCRIBE`](https://materialize.com/docs/sql/subscribe/). +//! +//! Consuming `SUBSCRIBE` durably requires a specific protocol: buffer until a +//! progress message closes a timestamp, apply the closed batch and persist its +//! frontier atomically, and resume with `SNAPSHOT = false AS OF frontier - 1`. +//! Every step has a failure mode that is invisible in testing. This crate +//! encodes the protocol so callers cannot get it wrong: +//! +//! * The unit of consumption is a [`ConsistentBatch`], never a bare row. You +//! physically cannot observe a half-closed timestamp. +//! * Resuming takes an opaque [`ResumeToken`]. The `- 1` boundary arithmetic +//! lives inside it, so you never compute a timestamp. +//! * Failures surface as a small, typed [`SubscribeError`] taxonomy. +//! +//! # Example +//! +//! ```no_run +//! use mz_subscribe::{Subscribe, SubscribeClient}; +//! +//! # async fn run() -> Result<(), mz_subscribe::SubscribeError> { +//! let client = SubscribeClient::connect("postgres://materialize@localhost:6875").await?; +//! +//! // Subscribe to a view with the upsert envelope, keyed by `id`. +//! let mut stream = client +//! .subscribe(Subscribe::object("winning_bids").envelope_upsert(["id"])) +//! .await?; +//! +//! while let Some(batch) = stream.next().await? { +//! for change in &batch.updates { +//! // Apply `change` to your sink. +//! let _ = change; +//! } +//! // Persist `batch.resume_token` atomically with the applied changes to +//! // get exactly-once state. On restart, hand it back to `client.resume`. +//! let _token = batch.resume_token.encode(); +//! } +//! # Ok(()) +//! # } +//! ``` +//! +//! # Guarantees +//! +//! Applying each batch and persisting its token in the same transaction yields +//! exactly-once *state*: after any crash, resuming from the last token neither +//! drops nor duplicates data. For side-effecting sinks (webhooks, queues) that +//! cannot transact with the checkpoint, use the token for at-least-once +//! delivery and deduplicate downstream. +//! +//! # Layers +//! +//! The client is three layers, each usable on its own: +//! +//! 1. [`RawStream`] (via [`SubscribeClient::subscribe_raw`]): the decoded +//! timestamped changes and progress markers, before any consistency policy. +//! This is the composable substrate. +//! 2. The consistency engine: buffer and release everything below a frontier, +//! consolidated. A [`BatchStream`] is a raw stream plus this engine. +//! 3. [`ConsistentBatch`] (one view) or [`CohortMoment`] (a [`Cohort`] of N +//! views): the closed unit handed to the consumer. +//! +//! A [`Cohort`] releases N views at their shared minimum frontier, so a +//! multi-view sink sees one cross-view-consistent moment at a time. A single +//! view is the cohort of one. + +mod batch; +mod client; +mod cohort; +mod engine; +mod envelope; +mod error; +mod statement; +mod token; + +pub use batch::{Batcher, ConsistentBatch}; +pub use client::{BatchStream, RawStream, SubscribeClient}; +pub use cohort::{Cohort, CohortMoment, ViewChanges}; +pub use envelope::{Change, Datum, Decoder, Envelope, Row, StreamMessage}; +pub use error::SubscribeError; +pub use statement::{Relation, Subscribe}; +pub use token::{CohortToken, ResumeToken}; diff --git a/misc/subscribe-sdk/rust/src/statement.rs b/misc/subscribe-sdk/rust/src/statement.rs new file mode 100644 index 0000000000000..abc2992dcd832 --- /dev/null +++ b/misc/subscribe-sdk/rust/src/statement.rs @@ -0,0 +1,311 @@ +// Copyright Materialize, Inc. and contributors. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License in the LICENSE file at the +// root of this repository, or online at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Building the `SUBSCRIBE` statement, including the initial-versus-resume +//! `SNAPSHOT` and `AS OF` handling that must line up with the resume token. + +use crate::envelope::Envelope; +use crate::token::ResumeToken; + +/// What a subscription reads from. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum Relation { + /// A catalog object (table, view, materialized view, source, or index). + /// Passed through verbatim, so a qualified name must already be a valid + /// SQL object reference (e.g. `my_db.public.orders`). + Object(String), + /// An arbitrary `SELECT`. Wrapped in parentheses in the statement. + Query(String), +} + +/// A subscription specification: what to read, how to shape it, and where to +/// stop. Converted to SQL by [`Subscribe::to_sql_initial`] and +/// [`Subscribe::to_sql_resume`]. +/// +/// `PROGRESS` is always requested. It is the primitive the batcher relies on to +/// know when a timestamp is closed, and requesting it unconditionally keeps the +/// column layout stable. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Subscribe { + relation: Relation, + envelope: Envelope, + up_to: Option, +} + +impl Subscribe { + /// Subscribes to a catalog object with the default (diff) envelope. + pub fn object(name: impl Into) -> Self { + Subscribe { + relation: Relation::Object(name.into()), + envelope: Envelope::Diff, + up_to: None, + } + } + + /// Subscribes to the result of a `SELECT` with the default (diff) envelope. + pub fn query(sql: impl Into) -> Self { + Subscribe { + relation: Relation::Query(sql.into()), + envelope: Envelope::Diff, + up_to: None, + } + } + + /// Switches to the upsert envelope keyed by the given columns. The server + /// collapses diffs into per-key inserts, deletes, and key violations. + pub fn envelope_upsert(mut self, key: I) -> Self + where + I: IntoIterator, + S: Into, + { + self.envelope = Envelope::Upsert { + key: key.into_iter().map(Into::into).collect(), + }; + self + } + + /// Bounds the subscription: it terminates once the frontier reaches + /// `timestamp` (exclusive), yielding a deterministic window. + pub fn up_to(mut self, timestamp: u64) -> Self { + self.up_to = Some(timestamp); + self + } + + /// The envelope this subscription uses. Needed to build a matching + /// [`crate::envelope::Decoder`]. + pub fn envelope(&self) -> &Envelope { + &self.envelope + } + + /// Whether the subscription is bounded by an `UP TO` and will therefore + /// terminate on its own. + pub fn is_bounded(&self) -> bool { + self.up_to.is_some() + } + + /// A stable identifier for this subscription's shape (relation plus + /// envelope), embedded in resume tokens so a changed query is caught on + /// resume rather than silently mixing incompatible results. + /// + /// Stability across processes and SDK versions matters, so this uses a + /// fixed FNV-1a hash rather than [`std::hash`], whose output is not + /// guaranteed stable. + pub fn fingerprint(&self) -> String { + let mut h = Fnv1a::new(); + match &self.relation { + Relation::Object(name) => { + h.update(b"object:"); + h.update(name.as_bytes()); + } + Relation::Query(sql) => { + h.update(b"query:"); + h.update(sql.as_bytes()); + } + } + match &self.envelope { + Envelope::Diff => h.update(b"|diff"), + Envelope::Upsert { key } => { + h.update(b"|upsert:"); + for k in key { + h.update(k.as_bytes()); + h.update(b","); + } + } + } + format!("{:016x}", h.finish()) + } + + /// The SQL for an initial subscription: takes the snapshot and lets the + /// server choose the `AS OF`. + pub fn to_sql_initial(&self) -> String { + self.to_sql(true, None) + } + + /// The SQL for resuming from `token`: no snapshot, and `AS OF` set to + /// `token.as_of()` so emission begins exactly where the token left off. + pub fn to_sql_resume(&self, token: &ResumeToken) -> String { + self.to_sql_resume_at(token.as_of()) + } + + /// The SQL for resuming at an explicit `as_of`: no snapshot, emitting only + /// timestamps strictly greater than `as_of`. Used by a cohort, where every + /// member resumes at the same joint `AS OF`. + pub fn to_sql_resume_at(&self, as_of: u64) -> String { + self.to_sql(false, Some(as_of)) + } + + fn to_sql(&self, snapshot: bool, as_of: Option) -> String { + let relation = match &self.relation { + Relation::Object(name) => name.clone(), + Relation::Query(sql) => format!("({sql})"), + }; + + let mut sql = format!("SUBSCRIBE {relation}"); + + if let Envelope::Upsert { key } = &self.envelope { + let keys = key + .iter() + .map(|k| quote_ident(k)) + .collect::>() + .join(", "); + sql.push_str(&format!(" ENVELOPE UPSERT (KEY ({keys}))")); + } + + sql.push_str(&format!(" WITH (PROGRESS, SNAPSHOT = {snapshot})")); + + if let Some(as_of) = as_of { + sql.push_str(&format!(" AS OF {as_of}")); + } + if let Some(up_to) = self.up_to { + sql.push_str(&format!(" UP TO {up_to}")); + } + sql + } +} + +/// Double-quotes a SQL identifier, escaping embedded quotes, so a key column +/// name cannot break out of the `KEY (...)` clause. +fn quote_ident(ident: &str) -> String { + format!("\"{}\"", ident.replace('"', "\"\"")) +} + +/// A small, stable FNV-1a 64-bit hash. Stable output is the point: fingerprints +/// are persisted in resume tokens and compared across processes. +struct Fnv1a(u64); + +impl Fnv1a { + fn new() -> Self { + Fnv1a(0xcbf2_9ce4_8422_2325) + } + + fn update(&mut self, bytes: &[u8]) { + for &b in bytes { + self.0 ^= u64::from(b); + self.0 = self.0.wrapping_mul(0x0000_0100_0000_01b3); + } + } + + fn finish(&self) -> u64 { + self.0 + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn initial_takes_snapshot_and_omits_as_of() { + let sql = Subscribe::object("orders").to_sql_initial(); + assert_eq!(sql, "SUBSCRIBE orders WITH (PROGRESS, SNAPSHOT = true)"); + } + + #[test] + fn resume_disables_snapshot_and_sets_as_of_to_frontier_minus_one() { + let sub = Subscribe::object("orders"); + let token = ResumeToken::new(100, sub.fingerprint()); + let sql = sub.to_sql_resume(&token); + assert_eq!( + sql, + "SUBSCRIBE orders WITH (PROGRESS, SNAPSHOT = false) AS OF 99" + ); + } + + #[test] + fn query_relation_is_parenthesized() { + let sql = Subscribe::query("SELECT id FROM orders WHERE paid").to_sql_initial(); + assert_eq!( + sql, + "SUBSCRIBE (SELECT id FROM orders WHERE paid) WITH (PROGRESS, SNAPSHOT = true)" + ); + } + + #[test] + fn upsert_envelope_precedes_with_clause() { + let sql = Subscribe::object("orders") + .envelope_upsert(["id"]) + .to_sql_initial(); + assert_eq!( + sql, + "SUBSCRIBE orders ENVELOPE UPSERT (KEY (\"id\")) WITH (PROGRESS, SNAPSHOT = true)" + ); + } + + #[test] + fn multi_column_key_and_up_to() { + let sql = Subscribe::object("sales") + .envelope_upsert(["region", "id"]) + .up_to(500) + .to_sql_initial(); + assert_eq!( + sql, + "SUBSCRIBE sales ENVELOPE UPSERT (KEY (\"region\", \"id\")) \ + WITH (PROGRESS, SNAPSHOT = true) UP TO 500" + ); + } + + #[test] + fn key_identifiers_are_quote_escaped() { + let sql = Subscribe::object("t") + .envelope_upsert(["weird\"name"]) + .to_sql_initial(); + assert!(sql.contains("KEY (\"weird\"\"name\")"), "{sql}"); + } + + #[test] + fn fingerprint_matches_cross_language_vectors() { + // These exact values are also asserted by the Python SDK, so both + // compute the same fingerprint (and therefore interchangeable tokens) + // for the same subscription shape. + assert_eq!( + Subscribe::object("orders").fingerprint(), + "f44d7c35ad6f569c" + ); + assert_eq!( + Subscribe::object("orders") + .envelope_upsert(["id"]) + .fingerprint(), + "a9392428ac0ed8c9" + ); + } + + #[test] + fn fingerprint_is_stable_and_distinguishes_shape() { + let a = Subscribe::object("orders"); + // Stable across constructions. + assert_eq!(a.fingerprint(), Subscribe::object("orders").fingerprint()); + // Different relation, envelope, or key changes the fingerprint. + assert_ne!(a.fingerprint(), Subscribe::object("other").fingerprint()); + assert_ne!( + a.fingerprint(), + Subscribe::object("orders") + .envelope_upsert(["id"]) + .fingerprint() + ); + assert_ne!( + Subscribe::object("orders") + .envelope_upsert(["id"]) + .fingerprint(), + Subscribe::object("orders") + .envelope_upsert(["id", "region"]) + .fingerprint() + ); + // UP TO does not change the shape, so it does not change the fingerprint. + assert_eq!( + a.fingerprint(), + Subscribe::object("orders").up_to(9).fingerprint() + ); + } +} diff --git a/misc/subscribe-sdk/rust/src/token.rs b/misc/subscribe-sdk/rust/src/token.rs new file mode 100644 index 0000000000000..18886215da824 --- /dev/null +++ b/misc/subscribe-sdk/rust/src/token.rs @@ -0,0 +1,245 @@ +// Copyright Materialize, Inc. and contributors. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License in the LICENSE file at the +// root of this repository, or online at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! The opaque checkpoint that makes resuming gap-free. + +use base64::engine::general_purpose::URL_SAFE_NO_PAD; +use base64::Engine; +use serde::{Deserialize, Serialize}; + +use crate::error::SubscribeError; + +/// Current on-the-wire token format. Bumped only if the encoded shape changes, +/// so a stored token from an older SDK is rejected with a clear error rather +/// than misinterpreted. +const TOKEN_FORMAT: u32 = 1; + +/// An opaque, serializable checkpoint marking how far a subscription has been +/// durably consumed. +/// +/// A token records a *closed* frontier: every update with a timestamp strictly +/// below [`ResumeToken::frontier`] has been delivered. Resuming re-subscribes +/// with `SNAPSHOT = false AS OF frontier - 1`, which the server interprets as +/// "emit updates with timestamp strictly greater than `frontier - 1`", i.e. +/// timestamp `>= frontier`. That is exactly the set not yet delivered, so the +/// resume is gap-free and overlap-free. +/// +/// The `- 1` is the single most error-prone part of the subscribe protocol. +/// It lives here, in [`ResumeToken::as_of`], so callers never compute it. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct ResumeToken { + format: u32, + frontier: u64, + fingerprint: String, +} + +impl ResumeToken { + /// Builds a token for a closed `frontier` of the query identified by + /// `fingerprint`. Crate-internal: tokens originate from the batcher, not + /// from user code. + pub(crate) fn new(frontier: u64, fingerprint: impl Into) -> Self { + ResumeToken { + format: TOKEN_FORMAT, + frontier, + fingerprint: fingerprint.into(), + } + } + + /// The closed frontier: every update below this timestamp is delivered. + pub fn frontier(&self) -> u64 { + self.frontier + } + + /// The fingerprint of the query this token was taken against. Compared on + /// resume to catch a changed query (see [`SubscribeError::SchemaMismatch`]). + pub fn fingerprint(&self) -> &str { + &self.fingerprint + } + + /// The `AS OF` value for a gap-free `SNAPSHOT = false` resume. + /// + /// Saturates at zero so a token for the very first frontier resumes from + /// the beginning rather than underflowing. + pub fn as_of(&self) -> u64 { + self.frontier.saturating_sub(1) + } + + /// Encodes the token to a compact, URL-safe string for durable storage. + /// The encoding is opaque: callers persist and return the bytes unchanged. + pub fn encode(&self) -> String { + // `to_string` on this fixed struct cannot fail, so the expect is a true + // invariant rather than error handling. + let json = serde_json::to_vec(self).expect("ResumeToken always serializes"); + URL_SAFE_NO_PAD.encode(json) + } + + /// Decodes a token previously produced by [`ResumeToken::encode`]. + pub fn decode(encoded: &str) -> Result { + let bytes = URL_SAFE_NO_PAD + .decode(encoded.as_bytes()) + .map_err(|e| SubscribeError::InvalidToken(format!("not valid base64: {e}")))?; + let token: ResumeToken = serde_json::from_slice(&bytes) + .map_err(|e| SubscribeError::InvalidToken(format!("malformed token: {e}")))?; + if token.format != TOKEN_FORMAT { + return Err(SubscribeError::InvalidToken(format!( + "unsupported token format {} (this SDK understands {})", + token.format, TOKEN_FORMAT + ))); + } + Ok(token) + } +} + +/// An opaque, serializable checkpoint for a whole cohort of subscriptions. +/// +/// A cohort is released at one shared frontier: the minimum closed frontier +/// across its members. This token records that single joint frontier plus each +/// member's fingerprint, in member order. Resuming re-subscribes every member +/// with `SNAPSHOT = false AS OF frontier - 1`, reconstructing the exact joint +/// cut, and refuses if the members no longer match. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct CohortToken { + format: u32, + frontier: u64, + members: Vec, +} + +impl CohortToken { + /// Builds a token for a cohort closed at `frontier` whose members have the + /// given `fingerprints`, in member order. Crate-internal: tokens originate + /// from the cohort engine, not user code. + pub(crate) fn new(frontier: u64, fingerprints: Vec) -> Self { + CohortToken { + format: TOKEN_FORMAT, + frontier, + members: fingerprints, + } + } + + /// The joint closed frontier: every member has delivered every update below + /// this timestamp. + pub fn frontier(&self) -> u64 { + self.frontier + } + + /// The member fingerprints, in the order the cohort was created. Compared on + /// resume to catch a changed cohort (see [`SubscribeError::SchemaMismatch`]). + pub fn members(&self) -> &[String] { + &self.members + } + + /// The `AS OF` value every member uses for a gap-free `SNAPSHOT = false` + /// resume. Saturates at zero, like [`ResumeToken::as_of`]. + pub fn as_of(&self) -> u64 { + self.frontier.saturating_sub(1) + } + + /// Encodes the token to a compact, URL-safe string for durable storage. + pub fn encode(&self) -> String { + let json = serde_json::to_vec(self).expect("CohortToken always serializes"); + URL_SAFE_NO_PAD.encode(json) + } + + /// Decodes a token previously produced by [`CohortToken::encode`]. + pub fn decode(encoded: &str) -> Result { + let bytes = URL_SAFE_NO_PAD + .decode(encoded.as_bytes()) + .map_err(|e| SubscribeError::InvalidToken(format!("not valid base64: {e}")))?; + let token: CohortToken = serde_json::from_slice(&bytes) + .map_err(|e| SubscribeError::InvalidToken(format!("malformed cohort token: {e}")))?; + if token.format != TOKEN_FORMAT { + return Err(SubscribeError::InvalidToken(format!( + "unsupported token format {} (this SDK understands {})", + token.format, TOKEN_FORMAT + ))); + } + Ok(token) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn as_of_is_frontier_minus_one_and_saturates() { + assert_eq!(ResumeToken::new(0, "q").as_of(), 0); + assert_eq!(ResumeToken::new(1, "q").as_of(), 0); + assert_eq!(ResumeToken::new(5, "q").as_of(), 4); + assert_eq!(ResumeToken::new(u64::MAX, "q").as_of(), u64::MAX - 1); + } + + #[test] + fn encode_decode_round_trips() { + let token = ResumeToken::new(42, "fingerprint-abc"); + let decoded = ResumeToken::decode(&token.encode()).expect("round trip"); + assert_eq!(decoded, token); + assert_eq!(decoded.frontier(), 42); + assert_eq!(decoded.fingerprint(), "fingerprint-abc"); + } + + #[test] + fn encoding_is_stable_and_cross_language() { + // This exact string is also asserted by the Python SDK's test suite, so + // a token minted by either SDK decodes in the other. If this changes, + // the two SDKs have diverged and tokens are no longer interchangeable. + let encoded = ResumeToken::new(42, "x").encode(); + assert_eq!( + encoded, + "eyJmb3JtYXQiOjEsImZyb250aWVyIjo0MiwiZmluZ2VycHJpbnQiOiJ4In0" + ); + } + + #[test] + fn decode_rejects_garbage() { + assert!(ResumeToken::decode("!!! not base64 !!!").is_err()); + assert!(ResumeToken::decode(&URL_SAFE_NO_PAD.encode(b"not json")).is_err()); + } + + #[test] + fn cohort_token_round_trips_and_carries_members() { + let token = CohortToken::new(42, vec!["fp-a".to_string(), "fp-b".to_string()]); + let decoded = CohortToken::decode(&token.encode()).expect("round trip"); + assert_eq!(decoded, token); + assert_eq!(decoded.frontier(), 42); + assert_eq!(decoded.as_of(), 41); + assert_eq!(decoded.members(), ["fp-a", "fp-b"]); + } + + #[test] + fn cohort_token_encoding_is_stable_and_cross_language() { + // Asserted byte-for-byte by the Python SDK too, so a cohort token minted + // by either SDK decodes in the other. + let encoded = CohortToken::new(7, vec!["x".to_string(), "y".to_string()]).encode(); + assert_eq!( + encoded, + "eyJmb3JtYXQiOjEsImZyb250aWVyIjo3LCJtZW1iZXJzIjpbIngiLCJ5Il19" + ); + } + + #[test] + fn decode_rejects_unknown_format() { + // A token encoded with a future format must be rejected, not + // silently misread. + let future = serde_json::json!({ + "format": TOKEN_FORMAT + 1, + "frontier": 1, + "fingerprint": "q", + }); + let encoded = URL_SAFE_NO_PAD.encode(serde_json::to_vec(&future).unwrap()); + let err = ResumeToken::decode(&encoded).unwrap_err(); + assert!(matches!(err, SubscribeError::InvalidToken(_)), "{err:?}"); + } +}