From 0e1f0dae9b6ac4d3bb944d0aa969670fdfbbea4a Mon Sep 17 00:00:00 2001 From: Brandur Date: Mon, 10 Aug 2026 16:07:47 -0500 Subject: [PATCH] Faster count implementation that's still quite accurate MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit I wasn't particularly surprised to pop open our PlanetScale report this morning and see that the count-by-state query used in River UI is now the demo's most expensive query by cumulative time: select state, count(*) from river_job group by state Count: 59,386 · p99: 13,131 ms · Cache hit: 87.5% This has been a known problem for quite some time both in Postgres and in River UI. The demo's now up to 1.6M completed rows, so counts are getting slower by the day. I was having Codex help brainstorm ways that this could be improved, and it came up with what I think is quite a clever strategy that should be very fast with minimum downsides: * The count endpoint starts out with an optimistic query that tries to do a full count by all states, but puts a limit of 10k rows on any particular one. * If only the constrained 10k+ information is available, that's what's shown, but we immediately try to get a full exact count of all rows because even if you have a lot of rows, it's still better to know that you have 10,001 versus 50k versus 200k, versus 5M. This longer count is kicked off in the background, and is refreshed every 1-30 minutes, depending on how long the count is taking. Its results are used when a reasonably fresh cache value is available so we can show users the best available number. Even when a cached value is available, we still prefer a more fresh capped count for states that don't exceed 10k. * In Postgres, if no cached exactly count is available (most commonly right after startup), we use a planner estimate to find a rough number. This value will only be in play for a short time until an exact count is available. The type of count (`exact`, `exact_cached`, `estimated`, `lower_bound`) is communicated o the UI so that it can give context on counts in tooltips. For example, it might show 12.3M, ≈987.7K, or 10K+ depending on the situation, along with source and freshness. I ran a benchmark and you can see that at large numbers doing a bounded count stays orders of magnitude more responsive. This might seem like a small thing, but it keeps the UI more up-to-date and responsive even for very large users, which is very good. | Rows | Table + indexes | Existing exact count | Bounded count | Planner estimate | Bounded speedup | |---:|---:|---:|---:|---:|---:| | 100K | 17 MB | 7.16 ms | 0.93 ms | 0.47 ms | 7.7× | | 1M | 174 MB | 24.45 ms | 1.09 ms | 0.66 ms | 22× | | 10M | 1.7 GB | 203.54 ms | 1.01 ms | 0.59 ms | 201× | I'm sort of hoping that this is a nice compromise for all things -- i.e. fast at small numbers, reasonably fast at large numbers, and still keeps precise numbers so we don't have to get too abstract. The downside is more code complexity, but Codex seems to have done a decent job of implementation (and I tweaked a bunch of stuff for style) and we have pretty good tests. --- .github/workflows/ci.yaml | 19 ++ CHANGELOG.md | 4 + Makefile | 11 +- go.mod | 11 +- go.sum | 50 ++- handler_api_endpoint.go | 291 +++++++++++++++--- handler_api_endpoint_test.go | 235 ++++++++++++-- internal/querycacher/query_cacher.go | 51 ++- internal/querycacher/query_cacher_test.go | 27 ++ .../river_ui_driver_interface.go | 48 +++ .../riveruidrivertest/job_count.go | 168 ++++++++++ .../riveruidrivertest/riveruidrivertest.go | 61 ++++ .../riveruipostgres/internal/dbsqlc/db.go | 25 ++ .../riveruipostgres/internal/dbsqlc/models.go | 74 +++++ .../internal/dbsqlc/river_job.sql | 77 +++++ .../internal/dbsqlc/river_job.sql.go | 182 +++++++++++ .../internal/dbsqlc/river_queries.go | 88 ++++++ .../internal/dbsqlc/river_queries_test.go | 31 ++ .../riveruipostgres/internal/dbsqlc/sqlc.yaml | 26 ++ .../river_ui_postgres_driver.go | 141 +++++++++ .../river_ui_postgres_driver_test.go | 34 ++ .../riveruisqlite/internal/dbsqlc/db.go | 24 ++ .../riveruisqlite/internal/dbsqlc/models.go | 26 ++ .../internal/dbsqlc/river_job.sql | 44 +++ .../internal/dbsqlc/river_job.sql.go | 67 ++++ .../internal/dbsqlc/river_queries.go | 38 +++ .../internal/dbsqlc/river_queries_test.go | 25 ++ .../riveruisqlite/internal/dbsqlc/sqlc.yaml | 16 + .../riveruisqlite/river_ui_sqlite_driver.go | 242 +++++++++++++++ .../river_ui_sqlite_driver_internal_test.go | 89 ++++++ .../river_ui_sqlite_driver_test.go | 42 +++ src/components/JobList.tsx | 4 +- src/components/JobStateFilters.test.tsx | 91 ++++-- src/components/JobStateFilters.tsx | 47 ++- src/services/states.ts | 44 ++- src/utils/jobStateFilterItems.ts | 32 +- 36 files changed, 2357 insertions(+), 128 deletions(-) create mode 100644 internal/riveruidriver/river_ui_driver_interface.go create mode 100644 internal/riveruidriver/riveruidrivertest/job_count.go create mode 100644 internal/riveruidriver/riveruidrivertest/riveruidrivertest.go create mode 100644 internal/riveruidriver/riveruipostgres/internal/dbsqlc/db.go create mode 100644 internal/riveruidriver/riveruipostgres/internal/dbsqlc/models.go create mode 100644 internal/riveruidriver/riveruipostgres/internal/dbsqlc/river_job.sql create mode 100644 internal/riveruidriver/riveruipostgres/internal/dbsqlc/river_job.sql.go create mode 100644 internal/riveruidriver/riveruipostgres/internal/dbsqlc/river_queries.go create mode 100644 internal/riveruidriver/riveruipostgres/internal/dbsqlc/river_queries_test.go create mode 100644 internal/riveruidriver/riveruipostgres/internal/dbsqlc/sqlc.yaml create mode 100644 internal/riveruidriver/riveruipostgres/river_ui_postgres_driver.go create mode 100644 internal/riveruidriver/riveruipostgres/river_ui_postgres_driver_test.go create mode 100644 internal/riveruidriver/riveruisqlite/internal/dbsqlc/db.go create mode 100644 internal/riveruidriver/riveruisqlite/internal/dbsqlc/models.go create mode 100644 internal/riveruidriver/riveruisqlite/internal/dbsqlc/river_job.sql create mode 100644 internal/riveruidriver/riveruisqlite/internal/dbsqlc/river_job.sql.go create mode 100644 internal/riveruidriver/riveruisqlite/internal/dbsqlc/river_queries.go create mode 100644 internal/riveruidriver/riveruisqlite/internal/dbsqlc/river_queries_test.go create mode 100644 internal/riveruidriver/riveruisqlite/internal/dbsqlc/sqlc.yaml create mode 100644 internal/riveruidriver/riveruisqlite/river_ui_sqlite_driver.go create mode 100644 internal/riveruidriver/riveruisqlite/river_ui_sqlite_driver_internal_test.go create mode 100644 internal/riveruidriver/riveruisqlite/river_ui_sqlite_driver_test.go diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index dadbf7bc..af7dc1cb 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -127,6 +127,25 @@ jobs: - name: Run lint run: make lint + sqlc_verify: + name: Verify sqlc generated code + runs-on: ubuntu-latest + timeout-minutes: 2 + + steps: + - name: Checkout repository + uses: actions/checkout@v7.0.1 + + - name: Setup sqlc + uses: sqlc-dev/setup-sqlc@v5 + with: + sqlc-version: "1.31.0" + + - name: Verify sqlc generated code + run: | + echo "Make sure that all sqlc changes are checked in" + make verify/sqlc + js_build_and_test: name: JS Build and Test runs-on: ubuntu-latest diff --git a/CHANGELOG.md b/CHANGELOG.md index c897a5de..b20d9854 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Changed + +- Keep large counts responsive while preserving useful magnitude with bounded live counts, adaptively cached exact snapshots, PostgreSQL planner estimates, and SQLite STAT4 estimates. [PR #655](https://github.com/riverqueue/riverui/pull/655). + ## [v0.19.0] - 2026-08-25 ### Changed diff --git a/Makefile b/Makefile index 981cc638..348faec5 100644 --- a/Makefile +++ b/Makefile @@ -21,6 +21,14 @@ fake_assets: dist: @npm run build +.PHONY: generate +generate: generate/sqlc + +.PHONY: generate/sqlc +generate/sqlc: + cd internal/riveruidriver/riveruipostgres/internal/dbsqlc && sqlc generate + cd internal/riveruidriver/riveruisqlite/internal/dbsqlc && sqlc generate + .PHONY: build build: dist CGO_ENABLED=0 go build @@ -74,4 +82,5 @@ verify: verify/sqlc .PHONY: verify/sqlc verify/sqlc: - cd internal/dbsqlc && sqlc diff + cd internal/riveruidriver/riveruipostgres/internal/dbsqlc && sqlc diff + cd internal/riveruidriver/riveruisqlite/internal/dbsqlc && sqlc diff diff --git a/go.mod b/go.mod index 77f2f6c9..d1595c42 100644 --- a/go.mod +++ b/go.mod @@ -9,14 +9,17 @@ require ( github.com/riverqueue/river v0.45.0 github.com/riverqueue/river/riverdriver v0.45.0 github.com/riverqueue/river/riverdriver/riverpgxv5 v0.45.0 + github.com/riverqueue/river/riverdriver/riversqlite v0.45.0 github.com/riverqueue/river/rivershared v0.45.0 github.com/riverqueue/river/rivertype v0.45.0 github.com/rs/cors v1.11.1 github.com/samber/slog-http v1.12.1 github.com/stretchr/testify v1.12.1 + modernc.org/sqlite v1.56.0 ) require ( + github.com/dustin/go-humanize v1.0.1 // indirect github.com/gabriel-vasile/mimetype v1.4.12 // indirect github.com/go-playground/locales v0.14.1 // indirect github.com/go-playground/universal-translator v0.18.1 // indirect @@ -26,6 +29,9 @@ require ( github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect github.com/jackc/puddle/v2 v2.2.2 // indirect github.com/leodido/go-urn v1.4.0 // indirect + github.com/mattn/go-isatty v0.0.24 // indirect + github.com/ncruces/go-strftime v1.0.0 // indirect + github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect github.com/tidwall/gjson v1.19.0 // indirect github.com/tidwall/match v1.2.0 // indirect github.com/tidwall/pretty v1.2.1 // indirect @@ -36,8 +42,11 @@ require ( go.yaml.in/yaml/v3 v3.0.5 // indirect golang.org/x/crypto v0.52.0 // indirect golang.org/x/sync v0.22.0 // indirect - golang.org/x/sys v0.45.0 // indirect + golang.org/x/sys v0.47.0 // indirect golang.org/x/text v0.41.0 // indirect + modernc.org/libc v1.74.4 // indirect + modernc.org/mathutil v1.7.1 // indirect + modernc.org/memory v1.11.0 // indirect ) retract ( diff --git a/go.sum b/go.sum index e01f71cb..7ede26c4 100644 --- a/go.sum +++ b/go.sum @@ -1,4 +1,6 @@ github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= +github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= github.com/gabriel-vasile/mimetype v1.4.12 h1:e9hWvmLYvtp846tLHam2o++qitpguFiYCKbn0w9jyqw= github.com/gabriel-vasile/mimetype v1.4.12/go.mod h1:d+9Oxyo1wTzWdyVUPMmXFvp4F9tea18J8ufA774AB3s= github.com/go-playground/assert/v2 v2.2.0 h1:JvknZsQTYeFEAhQwI4qEt9cyV5ONwRHC+lYKSsYSR8s= @@ -11,8 +13,12 @@ github.com/go-playground/validator/v10 v10.30.1 h1:f3zDSN/zOma+w6+1Wswgd9fLkdwy0 github.com/go-playground/validator/v10 v10.30.1/go.mod h1:oSuBIQzuJxL//3MelwSLD5hc2Tu889bF0Idm9Dg26cM= github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI= github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= +github.com/google/pprof v0.0.0-20260802141513-ef3492d7dac3 h1:LMLX+LgTNWpfvCBdFebv6EsYotImrt/Ppc5cXIriCSo= +github.com/google/pprof v0.0.0-20260802141513-ef3492d7dac3/go.mod h1:jl5iWTm0/hd5PjEYEOuwAJ57L/CibdZfrqZ5XA5GrCk= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k= +github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM= github.com/jackc/pgerrcode v0.0.0-20250907135507-afb5586c32a6 h1:D/V0gu4zQ3cL2WKeVNVM4r2gLxGGf6McLwgXzRTo2RQ= github.com/jackc/pgerrcode v0.0.0-20250907135507-afb5586c32a6/go.mod h1:a/s9Lp5W7n/DD0VrVoyJ00FbP2ytTPDVOivvn2bMlds= github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM= @@ -25,7 +31,13 @@ github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4= github.com/leodido/go-urn v1.4.0 h1:WT9HwE9SGECu3lg4d/dIA+jxlljEa1/ffXKmRjqdmIQ= github.com/leodido/go-urn v1.4.0/go.mod h1:bvxc+MVxLKB4z00jd1z+Dvzr47oO32F/QSNjSBOlFxI= +github.com/mattn/go-isatty v0.0.24 h1:tGZZoVgT/KiqK1c8ocVLeDS8BSWMRd47J3Lbz7vsReI= +github.com/mattn/go-isatty v0.0.24/go.mod h1:nMCL3Zebbrt45jsMDgnfIwz6ydEQApk5oEI3HqDio6A= +github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w= +github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE= +github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo= github.com/riverqueue/apiframe v0.0.0-20251229202423-2b52ce1c482e h1:OwOgxT3MRpOj5Mp6DhFdZP43FOQOf2hhywAuT5XZCR4= github.com/riverqueue/apiframe v0.0.0-20251229202423-2b52ce1c482e/go.mod h1:O7UmsAMjpMYuToN4au5GNXdmN1gli+5FTldgXqAfaD0= github.com/riverqueue/river v0.45.0 h1:gjp+eYx5sB+sA14URXls6EHdXOTbHRnXGN5u+FvYnH0= @@ -34,6 +46,8 @@ github.com/riverqueue/river/riverdriver v0.45.0 h1:oGSiSw5Pjv6toclmsvcc1VCWhtQXv github.com/riverqueue/river/riverdriver v0.45.0/go.mod h1:s6UignsfjQ4pgQPjEcFpH9mpuNgf30jxKxsPbSxqEHU= github.com/riverqueue/river/riverdriver/riverpgxv5 v0.45.0 h1:6ST4tuudkk2rJrGxmlDDKOi09jI3R/30sd3Csq03QD0= github.com/riverqueue/river/riverdriver/riverpgxv5 v0.45.0/go.mod h1:FgK37hDtuuL/MsqvysdS6kXzsOQiwyK/qV3+9OvpO+g= +github.com/riverqueue/river/riverdriver/riversqlite v0.45.0 h1:3DuGn28+6s/OWOpsjMpsgryLP4Wt3V2AhuWDVM6P+u0= +github.com/riverqueue/river/riverdriver/riversqlite v0.45.0/go.mod h1:eMi2dPTGz/UWRJWlKeMRswYN4VLud4toWEx5BV9Idg8= github.com/riverqueue/river/rivershared v0.45.0 h1:xWEqjaNBhqpE5QpPcCcPTXZNzPajI5PIYcMVdAjH0sM= github.com/riverqueue/river/rivershared v0.45.0/go.mod h1:55trQ+PMQPBrn8Za4J8NeNrkPdncxCIRfktO4Xr26WY= github.com/riverqueue/river/rivertype v0.45.0 h1:AITFM9ZB+kkd/PsWT7YuQ211V/kPCSv4awjSfnHVWMs= @@ -70,11 +84,43 @@ go.yaml.in/yaml/v3 v3.0.5 h1:N6y/pJk8buWs9NY5ERU2HSMfm+IuD/OtfdAnq6kESPw= go.yaml.in/yaml/v3 v3.0.5/go.mod h1:HVTZu1O7/Vkt2N+BFy8Zza+lnLsABggaTM2ZpNIGuKg= golang.org/x/crypto v0.52.0 h1:RMs7fP2rXdep0CftQlK8Uf+kibLm7qkCcradZWYz988= golang.org/x/crypto v0.52.0/go.mod h1:1QgfPxDqh0T2M/elOJtp9RvuR95kVjir0e6/BvEmGbc= +golang.org/x/mod v0.40.0 h1:hUv+3cXcdRHz08UmSiOob7sadHig73uo5bkXxQ/tvUs= +golang.org/x/mod v0.40.0/go.mod h1:0/weTWkPWGBikyTWAX3dkjVztMmBA5hM0DH6BElSupE= golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek= golang.org/x/sync v0.22.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= -golang.org/x/sys v0.45.0 h1:dO4czNzziLiiXplLQgBCEpCvXQ3dnkn0SdaZSYdQ+FY= -golang.org/x/sys v0.45.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= +golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs= +golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= golang.org/x/text v0.41.0 h1:vz/seA0lnX87Othu2f/0L24RcgrXD9/YFTSuGjj3rH8= golang.org/x/text v0.41.0/go.mod h1:jvf1O8ajNzZqhSrQBPbutR/EB83Cc0CFrezNQIwbb5M= +golang.org/x/tools v0.48.0 h1:3+hClM1aLL5mjMKm5ovokw9epgRXPuu2tILgismM6RE= +golang.org/x/tools v0.48.0/go.mod h1:08xX0orndb/F7jJxGDicx061tyd5pcMto75YMAXr6lk= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +modernc.org/cc/v4 v4.29.1 h1:MKgdCV3WykTSPqpVrnxdEDS0HEd2FHpKZDzxzU5LyeI= +modernc.org/cc/v4 v4.29.1/go.mod h1:OnovgIhbbMXMu1aISnJ0wvVD1KnW+cAUJkIrAWh+kVI= +modernc.org/ccgo/v4 v4.34.6 h1:sBgfIwyN0TQ9C5hwIeuqyeAKyMWnbvj2fvpF4L11uzU= +modernc.org/ccgo/v4 v4.34.6/go.mod h1:SZ8YcN9NG7XVsQYdm6jYBvi8PQP1qi+kqB6OhjqI3Fk= +modernc.org/fileutil v1.4.0 h1:j6ZzNTftVS054gi281TyLjHPp6CPHr2KCxEXjEbD6SM= +modernc.org/fileutil v1.4.0/go.mod h1:EqdKFDxiByqxLk8ozOxObDSfcVOv/54xDs/DUHdvCUU= +modernc.org/gc/v2 v2.6.5 h1:nyqdV8q46KvTpZlsw66kWqwXRHdjIlJOhG6kxiV/9xI= +modernc.org/gc/v2 v2.6.5/go.mod h1:YgIahr1ypgfe7chRuJi2gD7DBQiKSLMPgBQe9oIiito= +modernc.org/gc/v3 v3.1.4 h1:2g65LGVSmFQrXeITAw97x7hCRvZFcyE1uDP+7Vng7JI= +modernc.org/gc/v3 v3.1.4/go.mod h1:HFK/6AGESC7Ex+EZJhJ2Gni6cTaYpSMmU/cT9RmlfYY= +modernc.org/goabi0 v0.2.0 h1:HvEowk7LxcPd0eq6mVOAEMai46V+i7Jrj13t4AzuNks= +modernc.org/goabi0 v0.2.0/go.mod h1:CEFRnnJhKvWT1c1JTI3Avm+tgOWbkOu5oPA8eH8LnMI= +modernc.org/libc v1.74.4 h1:fX1Omw4o2/1C2iRkkIsrQTasJQldLhRmuPreXLoWs9k= +modernc.org/libc v1.74.4/go.mod h1:eeQAS9W3sZeKYMFubydxJpII9ybHWshk+7or7bLG9co= +modernc.org/mathutil v1.7.1 h1:GCZVGXdaN8gTqB1Mf/usp1Y/hSqgI2vAGGP4jZMCxOU= +modernc.org/mathutil v1.7.1/go.mod h1:4p5IwJITfppl0G4sUEDtCr4DthTaT47/N3aT6MhfgJg= +modernc.org/memory v1.11.0 h1:o4QC8aMQzmcwCK3t3Ux/ZHmwFPzE6hf2Y5LbkRs+hbI= +modernc.org/memory v1.11.0/go.mod h1:/JP4VbVC+K5sU2wZi9bHoq2MAkCnrt2r98UGeSK7Mjw= +modernc.org/opt v0.2.0 h1:tGyef5ApycA7FSEOMraay9SaTk5zmbx7Tu+cJs4QKZg= +modernc.org/opt v0.2.0/go.mod h1:03fq9lsNfvkYSfxrfUhZCWPk1lm4cq4N+Bh//bEtgns= +modernc.org/sortutil v1.2.1 h1:+xyoGf15mM3NMlPDnFqrteY07klSFxLElE2PVuWIJ7w= +modernc.org/sortutil v1.2.1/go.mod h1:7ZI3a3REbai7gzCLcotuw9AC4VZVpYMjDzETGsSMqJE= +modernc.org/sqlite v1.56.0 h1:/D8e2RfFqoy/Zc6PuC76U28zFwmI/sYx1Kjm4yEn9e0= +modernc.org/sqlite v1.56.0/go.mod h1:yCJ2cmAaIkHQ25oXWrF8H4O1lIfPYPR26yCEDj2P3pQ= +modernc.org/strutil v1.2.1 h1:UneZBkQA+DX2Rp35KcM69cSsNES9ly8mQWD71HKlOA0= +modernc.org/strutil v1.2.1/go.mod h1:EHkiggD70koQxjVdSBM3JKM7k6L0FbGE5eymy9i3B9A= +modernc.org/token v1.1.0 h1:Xl7Ap9dKaEs5kLoOQeQmPWevfnk/DM5qcLcYlA8ys6Y= +modernc.org/token v1.1.0/go.mod h1:UGzOrNV1mAFSEB63lOFHIpNRUVMvYTc6yu1SMY/XTDM= diff --git a/handler_api_endpoint.go b/handler_api_endpoint.go index dce312f1..eaa6e7fe 100644 --- a/handler_api_endpoint.go +++ b/handler_api_endpoint.go @@ -26,6 +26,9 @@ import ( "riverqueue.com/riverui/internal/apibundle" "riverqueue.com/riverui/internal/querycacher" + "riverqueue.com/riverui/internal/riveruidriver" + "riverqueue.com/riverui/internal/riveruidriver/riveruipostgres" + "riverqueue.com/riverui/internal/riveruidriver/riveruisqlite" ) type listResponse[T any] struct { @@ -878,22 +881,48 @@ type stateAndCountGetEndpoint[TTx any] struct { apibundle.APIBundle[TTx] apiendpoint.Endpoint[jobCancelRequest, stateAndCountGetResponse] - queryCacheSkipThreshold int // constant normally, but settable for testing - queryCacher *querycacher.QueryCacher[map[rivertype.JobState]int] + boundedQueryCacher *querycacher.QueryCacher[stateCountSnapshot] + countMax int + estimateCounts func(ctx context.Context, states []rivertype.JobState) (map[rivertype.JobState]stateCountEstimate, error) + exactQueryCacher *querycacher.QueryCacher[stateCountSnapshot] + driver riveruidriver.Driver } -func newStateAndCountGetEndpoint[TTx any](bundle apibundle.APIBundle[TTx]) *stateAndCountGetEndpoint[TTx] { - runQuery := func(ctx context.Context) (map[rivertype.JobState]int, error) { - return dbutil.WithTxV(ctx, bundle.DB, func(ctx context.Context, execTx riverdriver.ExecutorTx) (map[rivertype.JobState]int, error) { - tx := bundle.Driver.UnwrapTx(execTx) +const ( + stateAndCountDefaultMax = 10_000 + + // Two missed maximum-interval refreshes make an estimate preferable to an + // increasingly misleading exact snapshot. + stateCountExactMaxAge = 30 * time.Minute + stateCountExactRefreshMin = 30 * time.Second + stateCountExactRefreshMax = 15 * time.Minute + stateCountExactRefreshCostMul = 50 +) - return bundle.Driver.UnwrapExecutor(tx).JobCountByAllStates(ctx, &riverdriver.JobCountByAllStatesParams{Schema: bundle.Client.Schema()}) - }) +func newStateAndCountGetEndpoint[TTx any](bundle apibundle.APIBundle[TTx]) *stateAndCountGetEndpoint[TTx] { + endpoint := &stateAndCountGetEndpoint[TTx]{ + APIBundle: bundle, + countMax: stateAndCountDefaultMax, + driver: newRiverUIDriver(bundle.Driver.DatabaseName()), } - return &stateAndCountGetEndpoint[TTx]{ - APIBundle: bundle, - queryCacheSkipThreshold: 1_000_000, - queryCacher: querycacher.NewQueryCacher(bundle.Archetype, runQuery), + endpoint.boundedQueryCacher = querycacher.NewQueryCacher(bundle.Archetype, endpoint.queryBoundedCounts) + endpoint.exactQueryCacher = querycacher.NewQueryCacherWithOpts( + bundle.Archetype, + endpoint.queryExactCounts, + &querycacher.QueryCacherOpts{NextTickPeriod: stateCountExactRefreshPeriod}, + ) + endpoint.estimateCounts = endpoint.queryEstimatedCounts + return endpoint +} + +func newRiverUIDriver(databaseName string) riveruidriver.Driver { + switch databaseName { + case riverdriver.DatabaseNamePostgres: + return riveruipostgres.New() + case riverdriver.DatabaseNameSQLite: + return riveruisqlite.New() + default: + panic(fmt.Sprintf("unsupported River UI database %q", databaseName)) } } @@ -905,61 +934,223 @@ func (*stateAndCountGetEndpoint[TTx]) Meta() *apiendpoint.EndpointMeta { } func (a *stateAndCountGetEndpoint[TTx]) SubServices() []startstop.Service { - return []startstop.Service{a.queryCacher} + return []startstop.Service{a.boundedQueryCacher, a.exactQueryCacher} } type stateAndCountGetRequest struct{} -type stateAndCountGetResponse struct { - Available int `json:"available"` - Cancelled int `json:"cancelled"` - Completed int `json:"completed"` - Discarded int `json:"discarded"` - Pending int `json:"pending"` - Retryable int `json:"retryable"` - Running int `json:"running"` - Scheduled int `json:"scheduled"` +type stateCountAccuracy string + +const ( + stateCountAccuracyEstimated stateCountAccuracy = "estimated" // uses database planner/statistics estimate + stateCountAccuracyExact stateCountAccuracy = "exact" // exact + stateCountAccuracyExactCached stateCountAccuracy = "exact_cached" // exact (cached) + stateCountAccuracyLowerBound stateCountAccuracy = "lower_bound" // constrained to stateAndCountDefaultMax +) + +type stateCountResponse struct { + Accuracy stateCountAccuracy `json:"accuracy"` + Count int `json:"count"` + ObservedAt *time.Time `json:"observed_at,omitempty"` } +type stateAndCountGetResponse struct { + Available stateCountResponse `json:"available"` + Cancelled stateCountResponse `json:"cancelled"` + Completed stateCountResponse `json:"completed"` + Discarded stateCountResponse `json:"discarded"` + Pending stateCountResponse `json:"pending"` + Retryable stateCountResponse `json:"retryable"` + Running stateCountResponse `json:"running"` + Scheduled stateCountResponse `json:"scheduled"` +} + +// Execute resolves every state's count from the cheapest sufficiently useful +// source. A bounded index scan gives fresh exact values for small states. Large +// states prefer a recent exact snapshot refreshed adaptively in the background, +// then a database planner/statistics estimate, and finally the bound proven by the index +// scan. Full exact scans are never part of request latency. func (a *stateAndCountGetEndpoint[TTx]) Execute(ctx context.Context, _ *stateAndCountGetRequest) (*stateAndCountGetResponse, error) { - // Counts the total number of jobs in a state and count result. - totalJobs := func(stateAndCountRes map[rivertype.JobState]int) int { - var totalJobs int - for _, count := range stateAndCountRes { - totalJobs += count + countsAreExact := func(snapshot stateCountSnapshot) bool { + for _, count := range snapshot.Counts { + if count > a.countMax { + return false + } + } + return true + } + + // Prefer fresh counts while every state is below the cap. Once any state is + // capped, serve the periodically refreshed result to collapse queries from + // multiple UI clients. Both paths use the same bounded query. + boundedSnapshot, ok := a.boundedQueryCacher.CachedRes() + if !ok || countsAreExact(boundedSnapshot) { + var err error + boundedSnapshot, err = a.queryBoundedCounts(ctx) + if err != nil { + return nil, fmt.Errorf("error getting states and counts: %w", err) + } + } + + cappedStates := make([]rivertype.JobState, 0, len(allJobStates)) + for _, state := range allJobStates { + if boundedSnapshot.Counts[state] > a.countMax { + cappedStates = append(cappedStates, state) + } + } + + var ( + exactSnapshot, hasExactSnapshot = a.exactQueryCacher.CachedRes() + exactSnapshotIsFresh = hasExactSnapshot && time.Since(exactSnapshot.ObservedAt) <= stateCountExactMaxAge + ) + + statesNeedingEstimate := make([]rivertype.JobState, 0, len(cappedStates)) + for _, state := range cappedStates { + if !exactSnapshotIsFresh || exactSnapshot.Counts[state] <= a.countMax { + statesNeedingEstimate = append(statesNeedingEstimate, state) } - return totalJobs } - // Counting jobs can be an expensive operation given a large table, so in - // the presence of such, prefer to use a result that's cached periodically - // instead of querying inline with the API request. In case we don't have a - // cached result yet or there's a relatively small number of job rows, run - // the query directly (in the case of the latter so we present the freshest - // possible information). - stateAndCountRes, ok := a.queryCacher.CachedRes() - if !ok || totalJobs(stateAndCountRes) < a.queryCacheSkipThreshold { + estimates := make(map[rivertype.JobState]stateCountEstimate) + if len(statesNeedingEstimate) > 0 { var err error - stateAndCountRes, err = dbutil.WithTxV(ctx, a.DB, func(ctx context.Context, execTx riverdriver.ExecutorTx) (map[rivertype.JobState]int, error) { - tx := a.Driver.UnwrapTx(execTx) + estimates, err = a.estimateCounts(ctx, statesNeedingEstimate) + if err != nil { + // Estimates are an optional telemetry enhancement. The bounded count + // is still trustworthy, so degrade to a lower bound instead of failing + // the entire sidebar when planner statistics can't be read. + a.Logger.WarnContext(ctx, "Unable to estimate large job counts", "err", err) + estimates = make(map[rivertype.JobState]stateCountEstimate) + } + } + + resolvedCounts := make(map[rivertype.JobState]stateCountResponse, len(allJobStates)) + for _, state := range allJobStates { + boundedCount := boundedSnapshot.Counts[state] + + if boundedCount <= a.countMax { + // The bounded scan reached the end of this state's index range, so the + // value is exact and fresh even if another, larger state was capped. + resolvedCounts[state] = stateCountResponse{ + Accuracy: stateCountAccuracyExact, + Count: boundedCount, + ObservedAt: &boundedSnapshot.ObservedAt, + } + continue + } - return a.Driver.UnwrapExecutor(tx).JobCountByAllStates(ctx, &riverdriver.JobCountByAllStatesParams{Schema: a.Client.Schema()}) + if exactSnapshotIsFresh && exactSnapshot.Counts[state] > a.countMax { + // A recent full scan preserves the useful magnitude for common large + // states. Its timestamp makes the deliberate staleness visible. + resolvedCounts[state] = stateCountResponse{ + Accuracy: stateCountAccuracyExactCached, + Count: exactSnapshot.Counts[state], + ObservedAt: &exactSnapshot.ObservedAt, + } + continue + } + + if estimate, ok := estimates[state]; ok && estimate.Count > a.countMax { + // Planner statistics are cheap and retain an order of magnitude during + // cold start or when the last exact snapshot has become too old. + resolvedCounts[state] = stateCountResponse{ + Accuracy: stateCountAccuracyEstimated, + Count: estimate.Count, + ObservedAt: estimate.ObservedAt, + } + continue + } + + // The bounded scan proves only that there are more than countMax rows. + // Never present a stale planner estimate below that known lower bound. + resolvedCounts[state] = stateCountResponse{ + Accuracy: stateCountAccuracyLowerBound, + Count: a.countMax, + ObservedAt: &boundedSnapshot.ObservedAt, + } + } + + resp := &stateAndCountGetResponse{ + Available: resolvedCounts[rivertype.JobStateAvailable], + Cancelled: resolvedCounts[rivertype.JobStateCancelled], + Completed: resolvedCounts[rivertype.JobStateCompleted], + Discarded: resolvedCounts[rivertype.JobStateDiscarded], + Pending: resolvedCounts[rivertype.JobStatePending], + Retryable: resolvedCounts[rivertype.JobStateRetryable], + Running: resolvedCounts[rivertype.JobStateRunning], + Scheduled: resolvedCounts[rivertype.JobStateScheduled], + } + + return resp, nil +} + +type stateCountSnapshot struct { + Counts map[rivertype.JobState]int + ObservedAt time.Time +} + +type stateCountEstimate = riveruidriver.JobCountEstimateResult + +var allJobStates = []rivertype.JobState{ //nolint:gochecknoglobals + rivertype.JobStateAvailable, + rivertype.JobStateCancelled, + rivertype.JobStateCompleted, + rivertype.JobStateDiscarded, + rivertype.JobStatePending, + rivertype.JobStateRetryable, + rivertype.JobStateRunning, + rivertype.JobStateScheduled, +} + +func (a *stateAndCountGetEndpoint[TTx]) queryBoundedCounts(ctx context.Context) (stateCountSnapshot, error) { + return dbutil.WithTxV(ctx, a.DB, func(ctx context.Context, execTx riverdriver.ExecutorTx) (stateCountSnapshot, error) { + counts, err := a.driver.GetExecutor(execTx).JobCountByAllStatesCapped(ctx, &riveruidriver.JobCountByAllStatesCappedParams{ + Max: a.countMax, + Schema: a.Client.Schema(), }) if err != nil { - return nil, fmt.Errorf("error getting states and counts: %w", err) + return stateCountSnapshot{}, err } + return stateCountSnapshot{Counts: counts, ObservedAt: time.Now()}, nil + }) +} + +func (a *stateAndCountGetEndpoint[TTx]) queryExactCounts(ctx context.Context) (stateCountSnapshot, error) { + return dbutil.WithTxV(ctx, a.DB, func(ctx context.Context, execTx riverdriver.ExecutorTx) (stateCountSnapshot, error) { + counts, err := execTx.JobCountByAllStates(ctx, &riverdriver.JobCountByAllStatesParams{ + Schema: a.Client.Schema(), + }) + if err != nil { + return stateCountSnapshot{}, fmt.Errorf("error counting all jobs by state exactly: %w", err) + } + return stateCountSnapshot{Counts: counts, ObservedAt: time.Now()}, nil + }) +} + +func stateCountExactRefreshPeriod(queryDuration time.Duration, queryErr error) time.Duration { + if queryErr != nil { + // A failed full-table count is likely load-related. Back off to the + // maximum interval instead of repeatedly adding pressure to the database. + return stateCountExactRefreshMax } - return &stateAndCountGetResponse{ - Available: stateAndCountRes[rivertype.JobStateAvailable], - Cancelled: stateAndCountRes[rivertype.JobStateCancelled], - Completed: stateAndCountRes[rivertype.JobStateCompleted], - Discarded: stateAndCountRes[rivertype.JobStateDiscarded], - Pending: stateAndCountRes[rivertype.JobStatePending], - Retryable: stateAndCountRes[rivertype.JobStateRetryable], - Running: stateAndCountRes[rivertype.JobStateRunning], - Scheduled: stateAndCountRes[rivertype.JobStateScheduled], - }, nil + // Target about two percent of wall time for full exact counts. Fast counts + // still wait at least 30 seconds, while the maximum keeps exact telemetry + // reasonably fresh on very large installations. + refreshPeriod := queryDuration * stateCountExactRefreshCostMul + return min(max(refreshPeriod, stateCountExactRefreshMin), stateCountExactRefreshMax) +} + +// queryEstimatedCounts delegates planner telemetry to the database-specific UI +// driver. Drivers without an estimate strategy return no estimates, causing the +// endpoint to retain the lower bound proven by the capped query. +func (a *stateAndCountGetEndpoint[TTx]) queryEstimatedCounts(ctx context.Context, states []rivertype.JobState) (map[rivertype.JobState]stateCountEstimate, error) { + return dbutil.WithTxV(ctx, a.DB, func(ctx context.Context, execTx riverdriver.ExecutorTx) (map[rivertype.JobState]stateCountEstimate, error) { + return a.driver.GetExecutor(execTx).JobCountEstimate(ctx, &riveruidriver.JobCountEstimateParams{ + Schema: a.Client.Schema(), + States: states, + }) + }) } func NewNotFoundJob(jobID int64) *apierror.NotFound { diff --git a/handler_api_endpoint_test.go b/handler_api_endpoint_test.go index df2a133f..79ca473c 100644 --- a/handler_api_endpoint_test.go +++ b/handler_api_endpoint_test.go @@ -3,6 +3,7 @@ package riverui import ( "context" "encoding/json" + "errors" "log/slog" "net/http" "net/http/httptest" @@ -27,6 +28,8 @@ import ( "riverqueue.com/riverui/internal/apibundle" "riverqueue.com/riverui/internal/riverinternaltest/testfactory" + "riverqueue.com/riverui/internal/riveruidriver/riveruipostgres" + "riverqueue.com/riverui/internal/riveruidriver/riveruisqlite" "riverqueue.com/riverui/internal/uicommontest" ) @@ -1053,6 +1056,25 @@ func TestStateAndCountGetEndpoint(t *testing.T) { t.Parallel() ctx := context.Background() + stateCountsFromResponse := func(resp *stateAndCountGetResponse) map[rivertype.JobState]*stateCountResponse { + return map[rivertype.JobState]*stateCountResponse{ + rivertype.JobStateAvailable: &resp.Available, + rivertype.JobStateCancelled: &resp.Cancelled, + rivertype.JobStateCompleted: &resp.Completed, + rivertype.JobStateDiscarded: &resp.Discarded, + rivertype.JobStatePending: &resp.Pending, + rivertype.JobStateRetryable: &resp.Retryable, + rivertype.JobStateRunning: &resp.Running, + rivertype.JobStateScheduled: &resp.Scheduled, + } + } + requireExactCounts := func(t *testing.T, resp *stateAndCountGetResponse) { + t.Helper() + for state, stateCount := range stateCountsFromResponse(resp) { + require.Equal(t, stateCountAccuracyExact, stateCount.Accuracy, state) + require.NotNil(t, stateCount.ObservedAt, state) + } + } t.Run("Success", func(t *testing.T) { t.Parallel() @@ -1091,55 +1113,216 @@ func TestStateAndCountGetEndpoint(t *testing.T) { resp, err := apitest.InvokeHandler(ctx, endpoint.Execute, testMountOpts(t), &stateAndCountGetRequest{}) require.NoError(t, err) - require.Equal(t, &stateAndCountGetResponse{ - Available: 1, - Cancelled: 2, - Completed: 3, - Discarded: 4, - Pending: 5, - Retryable: 6, - Running: 7, - Scheduled: 8, - }, resp) + requireExactCounts(t, resp) + require.Equal(t, 1, resp.Available.Count) + require.Equal(t, 2, resp.Cancelled.Count) + require.Equal(t, 3, resp.Completed.Count) + require.Equal(t, 4, resp.Discarded.Count) + require.Equal(t, 5, resp.Pending.Count) + require.Equal(t, 6, resp.Retryable.Count) + require.Equal(t, 7, resp.Running.Count) + require.Equal(t, 8, resp.Scheduled.Count) }) - t.Run("WithCachedQueryAboveSkipThreshold", func(t *testing.T) { + t.Run("AtCountMaxIsExact", func(t *testing.T) { t.Parallel() - endpoint, bundle := setupEndpoint(ctx, t, newStateAndCountGetEndpoint) + const countMax = 3 + endpoint, bundle := setupEndpoint(ctx, t, func(bundle apibundle.APIBundle[pgx.Tx]) *stateAndCountGetEndpoint[pgx.Tx] { + endpoint := newStateAndCountGetEndpoint(bundle) + endpoint.countMax = countMax + return endpoint + }) + + for range countMax { + _ = testfactory.Job(ctx, t, bundle.exec, &testfactory.JobOpts{State: new(rivertype.JobStateAvailable)}) + } + + resp, err := apitest.InvokeHandler(ctx, endpoint.Execute, testMountOpts(t), &stateAndCountGetRequest{}) + require.NoError(t, err) + requireExactCounts(t, resp) + require.Equal(t, countMax, resp.Available.Count) + }) - const queryCacheSkipThreshold = 3 - for range queryCacheSkipThreshold + 1 { + t.Run("WithExactCachedSnapshot", func(t *testing.T) { + t.Parallel() + + const countMax = 3 + endpoint, bundle := setupEndpoint(ctx, t, func(bundle apibundle.APIBundle[pgx.Tx]) *stateAndCountGetEndpoint[pgx.Tx] { + endpoint := newStateAndCountGetEndpoint(bundle) + endpoint.countMax = countMax + return endpoint + }) + + for range countMax + 1 { _ = testfactory.Job(ctx, t, bundle.exec, &testfactory.JobOpts{State: new(rivertype.JobStateAvailable)}) } - _, err := endpoint.queryCacher.RunQuery(ctx) + _, err := endpoint.boundedQueryCacher.RunQuery(ctx) + require.NoError(t, err) + _, err = endpoint.exactQueryCacher.RunQuery(ctx) require.NoError(t, err) + // Once a state is capped, both caches are reused instead of making an + // exact count part of the request's latency. + _ = testfactory.Job(ctx, t, bundle.exec, &testfactory.JobOpts{State: new(rivertype.JobStateCancelled), FinalizedAt: new(time.Now())}) + resp, err := apitest.InvokeHandler(ctx, endpoint.Execute, testMountOpts(t), &stateAndCountGetRequest{}) require.NoError(t, err) - require.Equal(t, &stateAndCountGetResponse{ - Available: queryCacheSkipThreshold + 1, - }, resp) + require.Equal(t, countMax+1, resp.Available.Count) + require.Equal(t, stateCountAccuracyExactCached, resp.Available.Accuracy) + require.NotNil(t, resp.Available.ObservedAt) + require.Equal(t, 0, resp.Cancelled.Count) + require.Equal(t, stateCountAccuracyExact, resp.Cancelled.Accuracy) }) - t.Run("WithCachedQueryBelowSkipThreshold", func(t *testing.T) { + t.Run("WithExactCachedCount", func(t *testing.T) { t.Parallel() - endpoint, bundle := setupEndpoint(ctx, t, newStateAndCountGetEndpoint) + const countMax = 3 + endpoint, bundle := setupEndpoint(ctx, t, func(bundle apibundle.APIBundle[pgx.Tx]) *stateAndCountGetEndpoint[pgx.Tx] { + endpoint := newStateAndCountGetEndpoint(bundle) + endpoint.countMax = countMax + return endpoint + }) - const queryCacheSkipThreshold = 3 - for range queryCacheSkipThreshold - 1 { + for range countMax - 1 { _ = testfactory.Job(ctx, t, bundle.exec, &testfactory.JobOpts{State: new(rivertype.JobStateAvailable)}) } - _, err := endpoint.queryCacher.RunQuery(ctx) + _, err := endpoint.boundedQueryCacher.RunQuery(ctx) require.NoError(t, err) + // An exact cache result is refreshed inline for the latest counts. + _ = testfactory.Job(ctx, t, bundle.exec, &testfactory.JobOpts{State: new(rivertype.JobStateCancelled), FinalizedAt: new(time.Now())}) + resp, err := apitest.InvokeHandler(ctx, endpoint.Execute, testMountOpts(t), &stateAndCountGetRequest{}) require.NoError(t, err) - require.Equal(t, &stateAndCountGetResponse{ - Available: queryCacheSkipThreshold - 1, - }, resp) + requireExactCounts(t, resp) + require.Equal(t, countMax-1, resp.Available.Count) + require.Equal(t, 1, resp.Cancelled.Count) + }) + + t.Run("WithPlannerEstimate", func(t *testing.T) { + t.Parallel() + + const countMax = 3 + endpoint, bundle := setupEndpoint(ctx, t, func(bundle apibundle.APIBundle[pgx.Tx]) *stateAndCountGetEndpoint[pgx.Tx] { + endpoint := newStateAndCountGetEndpoint(bundle) + endpoint.countMax = countMax + return endpoint + }) + + for range countMax + 1 { + _ = testfactory.Job(ctx, t, bundle.exec, &testfactory.JobOpts{State: new(rivertype.JobStateCompleted), FinalizedAt: new(time.Now())}) + } + _, err := endpoint.boundedQueryCacher.RunQuery(ctx) + require.NoError(t, err) + + observedAt := time.Now().Add(-5 * time.Minute) + endpoint.estimateCounts = func(_ context.Context, states []rivertype.JobState) (map[rivertype.JobState]stateCountEstimate, error) { + require.Equal(t, []rivertype.JobState{rivertype.JobStateCompleted}, states) + return map[rivertype.JobState]stateCountEstimate{ + rivertype.JobStateCompleted: {Count: 1_000_000, ObservedAt: &observedAt}, + }, nil + } + + resp, err := apitest.InvokeHandler(ctx, endpoint.Execute, testMountOpts(t), &stateAndCountGetRequest{}) + require.NoError(t, err) + require.Equal(t, stateCountResponse{ + Accuracy: stateCountAccuracyEstimated, + Count: 1_000_000, + ObservedAt: &observedAt, + }, resp.Completed) + }) + + t.Run("ReadsPlannerEstimateFromPostgres", func(t *testing.T) { + t.Parallel() + + endpoint, bundle := setupEndpoint(ctx, t, newStateAndCountGetEndpoint) + for range 100 { + _ = testfactory.Job(ctx, t, bundle.exec, &testfactory.JobOpts{State: new(rivertype.JobStateCompleted), FinalizedAt: new(time.Now())}) + } + for range 10 { + _ = testfactory.Job(ctx, t, bundle.exec, &testfactory.JobOpts{State: new(rivertype.JobStateRunning)}) + } + require.NoError(t, bundle.exec.Exec(ctx, "ANALYZE river_job")) + + estimates, err := endpoint.queryEstimatedCounts(ctx, []rivertype.JobState{ + rivertype.JobStateCompleted, + rivertype.JobStateRunning, + }) + require.NoError(t, err) + require.Positive(t, estimates[rivertype.JobStateCompleted].Count) + require.NotNil(t, estimates[rivertype.JobStateCompleted].ObservedAt) + require.Greater(t, estimates[rivertype.JobStateCompleted].Count, estimates[rivertype.JobStateRunning].Count) + }) + + t.Run("WithLowerBoundForStaleEstimate", func(t *testing.T) { + t.Parallel() + + const countMax = 3 + endpoint, bundle := setupEndpoint(ctx, t, func(bundle apibundle.APIBundle[pgx.Tx]) *stateAndCountGetEndpoint[pgx.Tx] { + endpoint := newStateAndCountGetEndpoint(bundle) + endpoint.countMax = countMax + return endpoint + }) + + for range countMax + 1 { + _ = testfactory.Job(ctx, t, bundle.exec, &testfactory.JobOpts{State: new(rivertype.JobStateAvailable)}) + } + _, err := endpoint.boundedQueryCacher.RunQuery(ctx) + require.NoError(t, err) + endpoint.estimateCounts = func(_ context.Context, _ []rivertype.JobState) (map[rivertype.JobState]stateCountEstimate, error) { + return map[rivertype.JobState]stateCountEstimate{ + rivertype.JobStateAvailable: {Count: countMax - 1}, + }, nil + } + + resp, err := apitest.InvokeHandler(ctx, endpoint.Execute, testMountOpts(t), &stateAndCountGetRequest{}) + require.NoError(t, err) + require.Equal(t, countMax, resp.Available.Count) + require.Equal(t, stateCountAccuracyLowerBound, resp.Available.Accuracy) + require.NotNil(t, resp.Available.ObservedAt) }) } + +func TestAllJobStates(t *testing.T) { + t.Parallel() + + // Keep the endpoint's exhaustive response synchronized with River when a job + // state is added or reordered upstream. + require.Equal(t, rivertype.JobStates(), allJobStates) +} + +func TestNewRiverUIDriver(t *testing.T) { + t.Parallel() + + require.IsType(t, &riveruipostgres.Driver{}, newRiverUIDriver(riverdriver.DatabaseNamePostgres)) + require.IsType(t, &riveruisqlite.Driver{}, newRiverUIDriver(riverdriver.DatabaseNameSQLite)) + require.PanicsWithValue(t, `unsupported River UI database "mysql"`, func() { newRiverUIDriver("mysql") }) +} + +func TestStateCountExactRefreshPeriod(t *testing.T) { + t.Parallel() + + require.Equal(t, stateCountExactRefreshMin, stateCountExactRefreshPeriod(100*time.Millisecond, nil)) + require.Equal(t, 100*time.Second, stateCountExactRefreshPeriod(2*time.Second, nil)) + require.Equal(t, stateCountExactRefreshMax, stateCountExactRefreshPeriod(time.Hour, nil)) + require.Equal(t, stateCountExactRefreshMax, stateCountExactRefreshPeriod(time.Second, errors.New("database busy"))) +} + +func TestStateAndCountGetEndpointCustomSchema(t *testing.T) { + t.Parallel() + + ctx := context.Background() + endpoint, bundle := setupEndpointWithCustomSchema(ctx, t, newStateAndCountGetEndpoint) + jobParams := testfactory.Job_Build(t, &testfactory.JobOpts{State: new(rivertype.JobStateRunning)}) + jobParams.Schema = bundle.client.Schema() + _, err := bundle.exec.JobInsertFull(ctx, jobParams) + require.NoError(t, err) + + resp, err := apitest.InvokeHandler(ctx, endpoint.Execute, testMountOpts(t), &stateAndCountGetRequest{}) + require.NoError(t, err) + require.Equal(t, 1, resp.Running.Count) +} diff --git a/internal/querycacher/query_cacher.go b/internal/querycacher/query_cacher.go index e7ff0b28..357d471e 100644 --- a/internal/querycacher/query_cacher.go +++ b/internal/querycacher/query_cacher.go @@ -23,20 +23,37 @@ type QueryCacher[TRes any] struct { cachedRes TRes cachedResSet bool mu sync.RWMutex + nextTickPeriod func(queryDuration time.Duration, queryErr error) time.Duration runQuery func(ctx context.Context) (TRes, error) runQueryTestChan chan struct{} // closed when query is run; for testing tickPeriod time.Duration // constant normally, but settable for testing } +type QueryCacherOpts struct { + // NextTickPeriod makes the interval adaptive to the cost and result of the + // preceding query. The period starts after the query finishes, so an + // expensive query can never cause this service to run continuously. + NextTickPeriod func(queryDuration time.Duration, queryErr error) time.Duration +} + func NewQueryCacher[TRes any](archetype *baseservice.Archetype, runQuery func(ctx context.Context) (TRes, error)) *QueryCacher[TRes] { + return NewQueryCacherWithOpts(archetype, runQuery, nil) +} + +func NewQueryCacherWithOpts[TRes any](archetype *baseservice.Archetype, runQuery func(ctx context.Context) (TRes, error), opts *QueryCacherOpts) *QueryCacher[TRes] { // +/- 1s random variance to ticker interval. Makes sure that given multiple // query caches running simultaneously, they all start and are scheduled a // little differently to make a thundering herd problem less likely. randomTickVariance := time.Duration(rand.Float64()*float64(2*time.Second)) - 1*time.Second + var nextTickPeriod func(queryDuration time.Duration, queryErr error) time.Duration + if opts != nil { + nextTickPeriod = opts.NextTickPeriod + } queryCacher := baseservice.Init(archetype, &QueryCacher[TRes]{ - runQuery: runQuery, - tickPeriod: 10*time.Second + randomTickVariance, + nextTickPeriod: nextTickPeriod, + runQuery: runQuery, + tickPeriod: 10*time.Second + randomTickVariance, }) // TODO(brandur): Push this up into baseservice. @@ -76,7 +93,7 @@ func (s *QueryCacher[TRes]) RunQuery(ctx context.Context) (TRes, error) { return emptyRes, err } - s.Logger.DebugContext(ctx, s.Name+": Ran query and cached result", "duration", time.Since(start), "tick_period", s.tickPeriod) + s.Logger.DebugContext(ctx, s.Name+": Ran query and cached result", "duration", time.Since(start)) s.mu.Lock() s.cachedRes = res @@ -104,20 +121,36 @@ func (s *QueryCacher[TRes]) Start(ctx context.Context) error { started() defer stopped() - // In case a query runs long and exceeds tickPeriod, time.Ticker will - // drop ticks to compensate. - ticker := time.NewTicker(s.tickPeriod) - defer ticker.Stop() + // A timer is reset only after each query finishes. Unlike a ticker, this + // prevents a slow query from leaving a pending tick that starts another + // expensive query immediately. + timer := time.NewTimer(s.tickPeriod) + defer timer.Stop() for { select { case <-ctx.Done(): return - case <-ticker.C: - if _, err := s.RunQuery(ctx); err != nil { + case <-timer.C: + start := time.Now() + _, err := s.RunQuery(ctx) + queryDuration := time.Since(start) + if err != nil { s.Logger.ErrorContext(ctx, s.Name+": Error running query", "err", err) } + + nextTickPeriod := s.tickPeriod + if s.nextTickPeriod != nil { + nextTickPeriod = s.nextTickPeriod(queryDuration, err) + } + if nextTickPeriod <= 0 { + // A non-positive period would make the service spin. Falling + // back to the base interval is safer than treating bad options + // as permission to continuously query the database. + nextTickPeriod = s.tickPeriod + } + timer.Reset(nextTickPeriod) } } }() diff --git a/internal/querycacher/query_cacher_test.go b/internal/querycacher/query_cacher_test.go index df813462..67bd0b51 100644 --- a/internal/querycacher/query_cacher_test.go +++ b/internal/querycacher/query_cacher_test.go @@ -115,6 +115,33 @@ func TestQueryCacher(t *testing.T) { }, res) }) + t.Run("UsesAdaptivePeriodAfterQueryFinishes", func(t *testing.T) { + t.Parallel() + + var queryFinishedAt time.Time + nextPeriodCalled := make(chan struct{}) + queryCacher := NewQueryCacherWithOpts( + riversharedtest.BaseServiceArchetype(t), + func(_ context.Context) (int, error) { + queryFinishedAt = time.Now() + return 1, nil + }, + &QueryCacherOpts{ + NextTickPeriod: func(_ time.Duration, queryErr error) time.Duration { + require.NoError(t, queryErr) + require.False(t, queryFinishedAt.IsZero()) + close(nextPeriodCalled) + return time.Hour + }, + }, + ) + queryCacher.tickPeriod = time.Millisecond + + require.NoError(t, queryCacher.Start(ctx)) + t.Cleanup(queryCacher.Stop) + riversharedtest.WaitOrTimeout(t, nextPeriodCalled) + }) + t.Run("StartStopStress", func(t *testing.T) { t.Parallel() diff --git a/internal/riveruidriver/river_ui_driver_interface.go b/internal/riveruidriver/river_ui_driver_interface.go new file mode 100644 index 00000000..0a04c554 --- /dev/null +++ b/internal/riveruidriver/river_ui_driver_interface.go @@ -0,0 +1,48 @@ +// Package riveruidriver defines the database operations that River UI adds on top of +// River's core driver. Database-specific implementations keep SQL dialect +// details out of HTTP endpoints while continuing to use the executor already +// owned by the caller's River client. +package riveruidriver + +import ( + "context" + "time" + + "github.com/riverqueue/river/riverdriver" + "github.com/riverqueue/river/rivertype" +) + +// Driver decorates a River executor with queries used only by River UI. +type Driver interface { + GetExecutor(riverExecutor riverdriver.Executor) Executor +} + +// Executor contains River UI's database operations and embeds River's executor +// so callers retain transaction and core query support through one value. +type Executor interface { + riverdriver.Executor + + // JobCountByAllStatesCapped returns an exact count through Max and Max+1 as + // a sentinel when additional rows exist. + JobCountByAllStatesCapped(ctx context.Context, params *JobCountByAllStatesCappedParams) (map[rivertype.JobState]int, error) + // JobCountEstimate returns inexpensive database estimates where supported. + // Drivers without an estimate strategy return an empty map. + JobCountEstimate(ctx context.Context, params *JobCountEstimateParams) (map[rivertype.JobState]JobCountEstimateResult, error) +} + +type JobCountByAllStatesCappedParams struct { + // Max must be positive and less than math.MaxInt32 so callers get the same + // behavior from every database implementation. + Max int + Schema string +} + +type JobCountEstimateParams struct { + Schema string + States []rivertype.JobState +} + +type JobCountEstimateResult struct { + Count int + ObservedAt *time.Time +} diff --git a/internal/riveruidriver/riveruidrivertest/job_count.go b/internal/riveruidriver/riveruidrivertest/job_count.go new file mode 100644 index 00000000..c65ba697 --- /dev/null +++ b/internal/riveruidriver/riveruidrivertest/job_count.go @@ -0,0 +1,168 @@ +package riveruidrivertest + +import ( + "context" + "math" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/riverqueue/river/riverdriver" + "github.com/riverqueue/river/rivershared/testfactory" + "github.com/riverqueue/river/rivertype" + + "riverqueue.com/riverui/internal/riveruidriver" +) + +func exerciseJobCountByAllStatesCapped( + ctx context.Context, + t *testing.T, + driver riveruidriver.Driver, + executorWithTx func(ctx context.Context, t *testing.T) (riverdriver.Executor, string), +) { + t.Helper() + + t.Run("JobCountByAllStatesCapped", func(t *testing.T) { + t.Parallel() + + t.Run("ValidatesMax", func(t *testing.T) { + t.Parallel() + + executor, schema := executorWithSchema(ctx, t, driver, executorWithTx) + + _, err := executor.JobCountByAllStatesCapped(ctx, &riveruidriver.JobCountByAllStatesCappedParams{ + Max: 0, + Schema: schema, + }) + require.EqualError(t, err, "count max must be positive") + + _, err = executor.JobCountByAllStatesCapped(ctx, &riveruidriver.JobCountByAllStatesCappedParams{ + Max: math.MaxInt32, + Schema: schema, + }) + require.EqualError(t, err, "count max is too large") + }) + + t.Run("CountsEveryStateExactly", func(t *testing.T) { + t.Parallel() + + executor, schema := executorWithSchema(ctx, t, driver, executorWithTx) + expected := make(map[rivertype.JobState]int, len(rivertype.JobStates())) + for stateIndex, state := range rivertype.JobStates() { + count := stateIndex + 1 + expected[state] = count + for range count { + testfactory.Job(ctx, t, executor, &testfactory.JobOpts{Schema: schema, State: &state}) + } + } + + counts, err := executor.JobCountByAllStatesCapped(ctx, &riveruidriver.JobCountByAllStatesCappedParams{ + Max: 100, + Schema: schema, + }) + require.NoError(t, err) + require.Equal(t, expected, counts) + }) + + t.Run("CapsEachStateIndependently", func(t *testing.T) { + t.Parallel() + + executor, schema := executorWithSchema(ctx, t, driver, executorWithTx) + insertJobs(ctx, t, executor, schema, rivertype.JobStateAvailable, 4) + insertJobs(ctx, t, executor, schema, rivertype.JobStateCompleted, 2) + insertJobs(ctx, t, executor, schema, rivertype.JobStateRunning, 1) + + counts, err := executor.JobCountByAllStatesCapped(ctx, &riveruidriver.JobCountByAllStatesCappedParams{ + Max: 2, + Schema: schema, + }) + require.NoError(t, err) + + expected := make(map[rivertype.JobState]int, len(rivertype.JobStates())) + for _, state := range rivertype.JobStates() { + expected[state] = 0 + } + expected[rivertype.JobStateAvailable] = 3 + expected[rivertype.JobStateCompleted] = 2 + expected[rivertype.JobStateRunning] = 1 + require.Equal(t, expected, counts) + }) + }) +} + +func exerciseJobCountEstimate( + ctx context.Context, + t *testing.T, + databaseName string, + driver riveruidriver.Driver, + executorWithTx func(ctx context.Context, t *testing.T) (riverdriver.Executor, string), +) { + t.Helper() + + t.Run("JobCountEstimate", func(t *testing.T) { + t.Parallel() + + t.Run("RejectsInvalidState", func(t *testing.T) { + t.Parallel() + + executor, schema := executorWithSchema(ctx, t, driver, executorWithTx) + invalidState := rivertype.JobState("invalid") + _, err := executor.JobCountEstimate(ctx, &riveruidriver.JobCountEstimateParams{ + Schema: schema, + States: []rivertype.JobState{invalidState}, + }) + require.ErrorContains(t, err, `invalid job state for count estimate: "invalid"`) + }) + + t.Run("ReturnsDatabaseEstimateStrategy", func(t *testing.T) { + t.Parallel() + + executor, schema := executorWithSchema(ctx, t, driver, executorWithTx) + insertJobs(ctx, t, executor, schema, rivertype.JobStateCompleted, 100) + insertJobs(ctx, t, executor, schema, rivertype.JobStateRunning, 10) + + estimateParams := &riveruidriver.JobCountEstimateParams{ + Schema: schema, + States: []rivertype.JobState{ + rivertype.JobStateCompleted, + rivertype.JobStateRunning, + }, + } + + switch databaseName { + case riverdriver.DatabaseNamePostgres: + require.NoError(t, executor.Exec(ctx, "ANALYZE river_job")) + + estimates, err := executor.JobCountEstimate(ctx, estimateParams) + require.NoError(t, err) + require.Len(t, estimates, len(estimateParams.States)) + require.Positive(t, estimates[rivertype.JobStateCompleted].Count) + require.NotNil(t, estimates[rivertype.JobStateCompleted].ObservedAt) + require.Greater(t, estimates[rivertype.JobStateCompleted].Count, estimates[rivertype.JobStateRunning].Count) + case riverdriver.DatabaseNameSQLite: + estimates, err := executor.JobCountEstimate(ctx, estimateParams) + require.NoError(t, err) + require.Empty(t, estimates) + + require.NoError(t, executor.Exec(ctx, "ANALYZE river_job")) + + estimates, err = executor.JobCountEstimate(ctx, estimateParams) + require.NoError(t, err) + require.Len(t, estimates, len(estimateParams.States)) + require.Equal(t, 100, estimates[rivertype.JobStateCompleted].Count) + require.Nil(t, estimates[rivertype.JobStateCompleted].ObservedAt) + require.Equal(t, 10, estimates[rivertype.JobStateRunning].Count) + default: + t.Fatalf("unsupported database %q", databaseName) + } + }) + }) +} + +func insertJobs(ctx context.Context, t *testing.T, executor riverdriver.Executor, schema string, state rivertype.JobState, count int) { + t.Helper() + + for range count { + testfactory.Job(ctx, t, executor, &testfactory.JobOpts{Schema: schema, State: &state}) + } +} diff --git a/internal/riveruidriver/riveruidrivertest/riveruidrivertest.go b/internal/riveruidriver/riveruidrivertest/riveruidrivertest.go new file mode 100644 index 00000000..6917337c --- /dev/null +++ b/internal/riveruidriver/riveruidrivertest/riveruidrivertest.go @@ -0,0 +1,61 @@ +// Package riveruidrivertest provides a shared conformance suite for River UI +// database drivers. +package riveruidrivertest + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/riverqueue/river/riverdriver" + + "riverqueue.com/riverui/internal/riveruidriver" +) + +// Exercise runs River UI's complete driver contract against a concrete driver. +// executorWithTx must return an executor in an isolated, migrated transaction +// and the schema that River UI queries should target. +func Exercise( + ctx context.Context, + t *testing.T, + databaseName string, + driver riveruidriver.Driver, + executorWithTx func(ctx context.Context, t *testing.T) (riverdriver.Executor, string), +) { + t.Helper() + + exerciseGetExecutor(ctx, t, driver, executorWithTx) + exerciseJobCountByAllStatesCapped(ctx, t, driver, executorWithTx) + exerciseJobCountEstimate(ctx, t, databaseName, driver, executorWithTx) +} + +func exerciseGetExecutor( + ctx context.Context, + t *testing.T, + driver riveruidriver.Driver, + executorWithTx func(ctx context.Context, t *testing.T) (riverdriver.Executor, string), +) { + t.Helper() + + t.Run("GetExecutor", func(t *testing.T) { + t.Parallel() + + riverExecutor, _ := executorWithTx(ctx, t) + executor := driver.GetExecutor(riverExecutor) + + require.NoError(t, executor.Exec(ctx, "SELECT 1")) + }) +} + +func executorWithSchema( + ctx context.Context, + t *testing.T, + driver riveruidriver.Driver, + executorWithTx func(ctx context.Context, t *testing.T) (riverdriver.Executor, string), +) (riveruidriver.Executor, string) { + t.Helper() + + riverExecutor, schema := executorWithTx(ctx, t) + return driver.GetExecutor(riverExecutor), schema +} diff --git a/internal/riveruidriver/riveruipostgres/internal/dbsqlc/db.go b/internal/riveruidriver/riveruipostgres/internal/dbsqlc/db.go new file mode 100644 index 00000000..fb54e311 --- /dev/null +++ b/internal/riveruidriver/riveruipostgres/internal/dbsqlc/db.go @@ -0,0 +1,25 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.31.0 + +package dbsqlc + +import ( + "context" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" +) + +type DBTX interface { + Exec(context.Context, string, ...interface{}) (pgconn.CommandTag, error) + Query(context.Context, string, ...interface{}) (pgx.Rows, error) + QueryRow(context.Context, string, ...interface{}) pgx.Row +} + +func New() *Queries { + return &Queries{} +} + +type Queries struct { +} diff --git a/internal/riveruidriver/riveruipostgres/internal/dbsqlc/models.go b/internal/riveruidriver/riveruipostgres/internal/dbsqlc/models.go new file mode 100644 index 00000000..316efb74 --- /dev/null +++ b/internal/riveruidriver/riveruipostgres/internal/dbsqlc/models.go @@ -0,0 +1,74 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.31.0 + +package dbsqlc + +import ( + "database/sql/driver" + "fmt" + "time" +) + +type RiverJobState string + +const ( + RiverJobStateAvailable RiverJobState = "available" + RiverJobStateCancelled RiverJobState = "cancelled" + RiverJobStateCompleted RiverJobState = "completed" + RiverJobStateDiscarded RiverJobState = "discarded" + RiverJobStatePending RiverJobState = "pending" + RiverJobStateRetryable RiverJobState = "retryable" + RiverJobStateRunning RiverJobState = "running" + RiverJobStateScheduled RiverJobState = "scheduled" +) + +func (e *RiverJobState) Scan(src interface{}) error { + switch s := src.(type) { + case []byte: + *e = RiverJobState(s) + case string: + *e = RiverJobState(s) + default: + return fmt.Errorf("unsupported scan type for RiverJobState: %T", src) + } + return nil +} + +type NullRiverJobState struct { + RiverJobState RiverJobState + Valid bool // Valid is true if RiverJobState is not NULL +} + +// Scan implements the Scanner interface. +func (ns *NullRiverJobState) Scan(value interface{}) error { + if value == nil { + ns.RiverJobState, ns.Valid = "", false + return nil + } + ns.Valid = true + return ns.RiverJobState.Scan(value) +} + +// Value implements the driver Valuer interface. +func (ns NullRiverJobState) Value() (driver.Value, error) { + if !ns.Valid { + return nil, nil + } + return string(ns.RiverJobState), nil +} + +type PgStatAllTables struct { + Schemaname string + Relname string + LastAnalyze *time.Time + LastAutoanalyze *time.Time +} + +type RiverJob struct { + ID int64 + Priority int16 + Queue string + State RiverJobState + ScheduledAt time.Time +} diff --git a/internal/riveruidriver/riveruipostgres/internal/dbsqlc/river_job.sql b/internal/riveruidriver/riveruipostgres/internal/dbsqlc/river_job.sql new file mode 100644 index 00000000..146f957a --- /dev/null +++ b/internal/riveruidriver/riveruipostgres/internal/dbsqlc/river_job.sql @@ -0,0 +1,77 @@ +CREATE TYPE river_job_state AS ENUM( + 'available', + 'cancelled', + 'completed', + 'discarded', + 'pending', + 'retryable', + 'running', + 'scheduled' +); + +-- This is a minimal representation with only enough schema for sqlc to enable +-- the River UI queries currently defined below. Expand it toward River's full +-- table definition if future UI queries need additional columns. +CREATE TABLE river_job ( + id bigserial PRIMARY KEY, + priority smallint NOT NULL DEFAULT 1, + queue text NOT NULL DEFAULT 'default', + state river_job_state NOT NULL DEFAULT 'available', + scheduled_at timestamptz NOT NULL DEFAULT now() +); + +-- Minimal catalog view definition needed by sqlc to type-check the statistics +-- query. PostgreSQL provides the real pg_stat_all_tables view at runtime. +CREATE TABLE pg_stat_all_tables ( + schemaname text NOT NULL, + relname text NOT NULL, + last_analyze timestamptz, + last_autoanalyze timestamptz +); + +-- name: JobCountByAllStatesCapped :one +-- Each subquery follows the remaining columns of River's prioritized fetching +-- index so the database can stop its index scan as soon as the limit is met. +SELECT + (SELECT count(*) FROM (SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'available' ORDER BY queue, priority, scheduled_at, id LIMIT @max::int) AS limited_available) AS available, + (SELECT count(*) FROM (SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'cancelled' ORDER BY queue, priority, scheduled_at, id LIMIT @max::int) AS limited_cancelled) AS cancelled, + (SELECT count(*) FROM (SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'completed' ORDER BY queue, priority, scheduled_at, id LIMIT @max::int) AS limited_completed) AS completed, + (SELECT count(*) FROM (SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'discarded' ORDER BY queue, priority, scheduled_at, id LIMIT @max::int) AS limited_discarded) AS discarded, + (SELECT count(*) FROM (SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'pending' ORDER BY queue, priority, scheduled_at, id LIMIT @max::int) AS limited_pending) AS pending, + (SELECT count(*) FROM (SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'retryable' ORDER BY queue, priority, scheduled_at, id LIMIT @max::int) AS limited_retryable) AS retryable, + (SELECT count(*) FROM (SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'running' ORDER BY queue, priority, scheduled_at, id LIMIT @max::int) AS limited_running) AS running, + (SELECT count(*) FROM (SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'scheduled' ORDER BY queue, priority, scheduled_at, id LIMIT @max::int) AS limited_scheduled) AS scheduled; + +-- name: JobCountAnalyzedAt :one +SELECT GREATEST(last_analyze, last_autoanalyze)::timestamptz AS analyzed_at +FROM pg_stat_all_tables +WHERE schemaname = COALESCE(NULLIF(@schema::text, ''), current_schema()) + AND relname = 'river_job'; + +-- These queries intentionally embed a state literal. A parameterized query can +-- eventually receive PostgreSQL's generic prepared plan, which loses the +-- per-state selectivity that makes the estimate useful. + +-- name: JobCountEstimateAvailable :one +EXPLAIN (FORMAT JSON) SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'available'; + +-- name: JobCountEstimateCancelled :one +EXPLAIN (FORMAT JSON) SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'cancelled'; + +-- name: JobCountEstimateCompleted :one +EXPLAIN (FORMAT JSON) SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'completed'; + +-- name: JobCountEstimateDiscarded :one +EXPLAIN (FORMAT JSON) SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'discarded'; + +-- name: JobCountEstimatePending :one +EXPLAIN (FORMAT JSON) SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'pending'; + +-- name: JobCountEstimateRetryable :one +EXPLAIN (FORMAT JSON) SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'retryable'; + +-- name: JobCountEstimateRunning :one +EXPLAIN (FORMAT JSON) SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'running'; + +-- name: JobCountEstimateScheduled :one +EXPLAIN (FORMAT JSON) SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'scheduled'; diff --git a/internal/riveruidriver/riveruipostgres/internal/dbsqlc/river_job.sql.go b/internal/riveruidriver/riveruipostgres/internal/dbsqlc/river_job.sql.go new file mode 100644 index 00000000..fcd66b25 --- /dev/null +++ b/internal/riveruidriver/riveruipostgres/internal/dbsqlc/river_job.sql.go @@ -0,0 +1,182 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.31.0 +// source: river_job.sql + +package dbsqlc + +import ( + "context" + "time" +) + +const jobCountAnalyzedAt = `-- name: JobCountAnalyzedAt :one +SELECT GREATEST(last_analyze, last_autoanalyze)::timestamptz AS analyzed_at +FROM pg_stat_all_tables +WHERE schemaname = COALESCE(NULLIF($1::text, ''), current_schema()) + AND relname = 'river_job' +` + +func (q *Queries) JobCountAnalyzedAt(ctx context.Context, db DBTX, schema string) (time.Time, error) { + row := db.QueryRow(ctx, jobCountAnalyzedAt, schema) + var analyzed_at time.Time + err := row.Scan(&analyzed_at) + return analyzed_at, err +} + +const jobCountByAllStatesCapped = `-- name: JobCountByAllStatesCapped :one +SELECT + (SELECT count(*) FROM (SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'available' ORDER BY queue, priority, scheduled_at, id LIMIT $1::int) AS limited_available) AS available, + (SELECT count(*) FROM (SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'cancelled' ORDER BY queue, priority, scheduled_at, id LIMIT $1::int) AS limited_cancelled) AS cancelled, + (SELECT count(*) FROM (SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'completed' ORDER BY queue, priority, scheduled_at, id LIMIT $1::int) AS limited_completed) AS completed, + (SELECT count(*) FROM (SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'discarded' ORDER BY queue, priority, scheduled_at, id LIMIT $1::int) AS limited_discarded) AS discarded, + (SELECT count(*) FROM (SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'pending' ORDER BY queue, priority, scheduled_at, id LIMIT $1::int) AS limited_pending) AS pending, + (SELECT count(*) FROM (SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'retryable' ORDER BY queue, priority, scheduled_at, id LIMIT $1::int) AS limited_retryable) AS retryable, + (SELECT count(*) FROM (SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'running' ORDER BY queue, priority, scheduled_at, id LIMIT $1::int) AS limited_running) AS running, + (SELECT count(*) FROM (SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'scheduled' ORDER BY queue, priority, scheduled_at, id LIMIT $1::int) AS limited_scheduled) AS scheduled +` + +type JobCountByAllStatesCappedRow struct { + Available int64 + Cancelled int64 + Completed int64 + Discarded int64 + Pending int64 + Retryable int64 + Running int64 + Scheduled int64 +} + +// Each subquery follows the remaining columns of River's prioritized fetching +// index so the database can stop its index scan as soon as the limit is met. +func (q *Queries) JobCountByAllStatesCapped(ctx context.Context, db DBTX, max int32) (*JobCountByAllStatesCappedRow, error) { + row := db.QueryRow(ctx, jobCountByAllStatesCapped, max) + var i JobCountByAllStatesCappedRow + err := row.Scan( + &i.Available, + &i.Cancelled, + &i.Completed, + &i.Discarded, + &i.Pending, + &i.Retryable, + &i.Running, + &i.Scheduled, + ) + return &i, err +} + +const jobCountEstimateAvailable = `-- name: JobCountEstimateAvailable :one + +EXPLAIN (FORMAT JSON) SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'available' +` + +type JobCountEstimateAvailableRow struct { +} + +// These queries intentionally embed a state literal. A parameterized query can +// eventually receive PostgreSQL's generic prepared plan, which loses the +// per-state selectivity that makes the estimate useful. +func (q *Queries) JobCountEstimateAvailable(ctx context.Context, db DBTX) (*JobCountEstimateAvailableRow, error) { + row := db.QueryRow(ctx, jobCountEstimateAvailable) + var i JobCountEstimateAvailableRow + err := row.Scan() + return &i, err +} + +const jobCountEstimateCancelled = `-- name: JobCountEstimateCancelled :one +EXPLAIN (FORMAT JSON) SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'cancelled' +` + +type JobCountEstimateCancelledRow struct { +} + +func (q *Queries) JobCountEstimateCancelled(ctx context.Context, db DBTX) (*JobCountEstimateCancelledRow, error) { + row := db.QueryRow(ctx, jobCountEstimateCancelled) + var i JobCountEstimateCancelledRow + err := row.Scan() + return &i, err +} + +const jobCountEstimateCompleted = `-- name: JobCountEstimateCompleted :one +EXPLAIN (FORMAT JSON) SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'completed' +` + +type JobCountEstimateCompletedRow struct { +} + +func (q *Queries) JobCountEstimateCompleted(ctx context.Context, db DBTX) (*JobCountEstimateCompletedRow, error) { + row := db.QueryRow(ctx, jobCountEstimateCompleted) + var i JobCountEstimateCompletedRow + err := row.Scan() + return &i, err +} + +const jobCountEstimateDiscarded = `-- name: JobCountEstimateDiscarded :one +EXPLAIN (FORMAT JSON) SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'discarded' +` + +type JobCountEstimateDiscardedRow struct { +} + +func (q *Queries) JobCountEstimateDiscarded(ctx context.Context, db DBTX) (*JobCountEstimateDiscardedRow, error) { + row := db.QueryRow(ctx, jobCountEstimateDiscarded) + var i JobCountEstimateDiscardedRow + err := row.Scan() + return &i, err +} + +const jobCountEstimatePending = `-- name: JobCountEstimatePending :one +EXPLAIN (FORMAT JSON) SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'pending' +` + +type JobCountEstimatePendingRow struct { +} + +func (q *Queries) JobCountEstimatePending(ctx context.Context, db DBTX) (*JobCountEstimatePendingRow, error) { + row := db.QueryRow(ctx, jobCountEstimatePending) + var i JobCountEstimatePendingRow + err := row.Scan() + return &i, err +} + +const jobCountEstimateRetryable = `-- name: JobCountEstimateRetryable :one +EXPLAIN (FORMAT JSON) SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'retryable' +` + +type JobCountEstimateRetryableRow struct { +} + +func (q *Queries) JobCountEstimateRetryable(ctx context.Context, db DBTX) (*JobCountEstimateRetryableRow, error) { + row := db.QueryRow(ctx, jobCountEstimateRetryable) + var i JobCountEstimateRetryableRow + err := row.Scan() + return &i, err +} + +const jobCountEstimateRunning = `-- name: JobCountEstimateRunning :one +EXPLAIN (FORMAT JSON) SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'running' +` + +type JobCountEstimateRunningRow struct { +} + +func (q *Queries) JobCountEstimateRunning(ctx context.Context, db DBTX) (*JobCountEstimateRunningRow, error) { + row := db.QueryRow(ctx, jobCountEstimateRunning) + var i JobCountEstimateRunningRow + err := row.Scan() + return &i, err +} + +const jobCountEstimateScheduled = `-- name: JobCountEstimateScheduled :one +EXPLAIN (FORMAT JSON) SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'scheduled' +` + +type JobCountEstimateScheduledRow struct { +} + +func (q *Queries) JobCountEstimateScheduled(ctx context.Context, db DBTX) (*JobCountEstimateScheduledRow, error) { + row := db.QueryRow(ctx, jobCountEstimateScheduled) + var i JobCountEstimateScheduledRow + err := row.Scan() + return &i, err +} diff --git a/internal/riveruidriver/riveruipostgres/internal/dbsqlc/river_queries.go b/internal/riveruidriver/riveruipostgres/internal/dbsqlc/river_queries.go new file mode 100644 index 00000000..a37c7633 --- /dev/null +++ b/internal/riveruidriver/riveruipostgres/internal/dbsqlc/river_queries.go @@ -0,0 +1,88 @@ +package dbsqlc + +import ( + "context" + "fmt" + "time" + + "github.com/jackc/pgx/v5/pgtype" + + "github.com/riverqueue/river/rivertype" +) + +// Row and RowQuerier are the small common subset exposed by every River +// executor. sqlc's generated DBTX includes pgx-specific methods that these +// read-only UI queries don't need, so this adapter keeps the UI driver usable +// with any PostgreSQL River driver, including database/sql implementations. +type Row interface { + Scan(dest ...any) error +} + +type RowQuerier interface { + QueryRow(ctx context.Context, query string, args ...any) Row +} + +func JobCountByAllStatesCappedRiver(ctx context.Context, db RowQuerier, maxRows int32) (*JobCountByAllStatesCappedRow, error) { + row := db.QueryRow(ctx, jobCountByAllStatesCapped, maxRows) + var result JobCountByAllStatesCappedRow + err := row.Scan( + &result.Available, + &result.Cancelled, + &result.Completed, + &result.Discarded, + &result.Pending, + &result.Retryable, + &result.Running, + &result.Scheduled, + ) + return &result, err +} + +func JobCountAnalyzedAtRiver(ctx context.Context, db RowQuerier, schema string) (time.Time, bool, error) { + var analyzedAt pgtype.Timestamptz + if err := db.QueryRow(ctx, jobCountAnalyzedAt, schema).Scan(&analyzedAt); err != nil { + return time.Time{}, false, err + } + if !analyzedAt.Valid { + return time.Time{}, false, nil + } + return analyzedAt.Time, true, nil +} + +func JobCountEstimateRiver(ctx context.Context, db RowQuerier, state rivertype.JobState) ([]byte, error) { + // sqlc validates and generates the EXPLAIN query constants, but can't infer + // EXPLAIN's JSON result column, so scan that one value through this adapter. + query, err := jobCountEstimateQuery(state) + if err != nil { + return nil, err + } + + var rawPlan []byte + if err := db.QueryRow(ctx, query).Scan(&rawPlan); err != nil { + return nil, err + } + return rawPlan, nil +} + +func jobCountEstimateQuery(state rivertype.JobState) (string, error) { + switch state { + case rivertype.JobStateAvailable: + return jobCountEstimateAvailable, nil + case rivertype.JobStateCancelled: + return jobCountEstimateCancelled, nil + case rivertype.JobStateCompleted: + return jobCountEstimateCompleted, nil + case rivertype.JobStateDiscarded: + return jobCountEstimateDiscarded, nil + case rivertype.JobStatePending: + return jobCountEstimatePending, nil + case rivertype.JobStateRetryable: + return jobCountEstimateRetryable, nil + case rivertype.JobStateRunning: + return jobCountEstimateRunning, nil + case rivertype.JobStateScheduled: + return jobCountEstimateScheduled, nil + default: + return "", fmt.Errorf("invalid job state for count estimate: %q", state) + } +} diff --git a/internal/riveruidriver/riveruipostgres/internal/dbsqlc/river_queries_test.go b/internal/riveruidriver/riveruipostgres/internal/dbsqlc/river_queries_test.go new file mode 100644 index 00000000..b7e94119 --- /dev/null +++ b/internal/riveruidriver/riveruipostgres/internal/dbsqlc/river_queries_test.go @@ -0,0 +1,31 @@ +package dbsqlc + +import ( + "testing" + + "github.com/stretchr/testify/require" + + "github.com/riverqueue/river/rivertype" +) + +func TestJobCountEstimateQuery(t *testing.T) { + t.Parallel() + + for _, state := range rivertype.JobStates() { + query, err := jobCountEstimateQuery(state) + require.NoError(t, err) + require.Contains(t, query, "state = '"+string(state)+"'") + require.NotContains(t, query, "$1", "state must remain a literal so PostgreSQL produces a state-specific plan") + } + + _, err := jobCountEstimateQuery(rivertype.JobState("completed'; DROP TABLE river_job; --")) + require.EqualError(t, err, `invalid job state for count estimate: "completed'; DROP TABLE river_job; --"`) +} + +func TestJobCountByAllStatesCappedQueryStates(t *testing.T) { + t.Parallel() + + for _, state := range rivertype.JobStates() { + require.Contains(t, jobCountByAllStatesCapped, "state = '"+string(state)+"'") + } +} diff --git a/internal/riveruidriver/riveruipostgres/internal/dbsqlc/sqlc.yaml b/internal/riveruidriver/riveruipostgres/internal/dbsqlc/sqlc.yaml new file mode 100644 index 00000000..5f1918c7 --- /dev/null +++ b/internal/riveruidriver/riveruipostgres/internal/dbsqlc/sqlc.yaml @@ -0,0 +1,26 @@ +version: "2" +sql: + - engine: "postgresql" + queries: + - river_job.sql + schema: + - river_job.sql + gen: + go: + package: "dbsqlc" + sql_package: "pgx/v5" + out: "." + emit_exact_table_names: true + emit_methods_with_db_argument: true + emit_params_struct_pointers: true + emit_result_struct_pointers: true + + overrides: + - db_type: "timestamptz" + go_type: "time.Time" + + - db_type: "timestamptz" + go_type: + type: "time.Time" + pointer: true + nullable: true diff --git a/internal/riveruidriver/riveruipostgres/river_ui_postgres_driver.go b/internal/riveruidriver/riveruipostgres/river_ui_postgres_driver.go new file mode 100644 index 00000000..007e7228 --- /dev/null +++ b/internal/riveruidriver/riveruipostgres/river_ui_postgres_driver.go @@ -0,0 +1,141 @@ +// Package riveruipostgres implements River UI's database operations for +// PostgreSQL. It decorates an executor from any PostgreSQL River driver rather +// than owning a second pool. +package riveruipostgres + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "math" + "time" + + "github.com/riverqueue/river/riverdriver" + "github.com/riverqueue/river/rivershared/sqlctemplate" + "github.com/riverqueue/river/rivershared/util/dbutil" + "github.com/riverqueue/river/rivertype" + + "riverqueue.com/riverui/internal/riveruidriver" + "riverqueue.com/riverui/internal/riveruidriver/riveruipostgres/internal/dbsqlc" +) + +type Driver struct{} + +// New returns a PostgreSQL River UI driver. Unlike River's primary database +// driver, it doesn't take a pool because River UI only adds queries on top of +// the executor already owned by the caller's River client. +func New() *Driver { + return &Driver{} +} + +func (*Driver) GetExecutor(riverExecutor riverdriver.Executor) riveruidriver.Executor { + return &Executor{ + Executor: riverExecutor, + dbtx: riverExecutorWrapper{executor: riverExecutor}, + } +} + +type Executor struct { + riverdriver.Executor + + dbtx riverExecutorWrapper +} + +var ( + _ riveruidriver.Driver = (*Driver)(nil) + _ riveruidriver.Executor = (*Executor)(nil) +) + +func (e *Executor) JobCountByAllStatesCapped(ctx context.Context, params *riveruidriver.JobCountByAllStatesCappedParams) (map[rivertype.JobState]int, error) { + if params.Max < 1 { + return nil, errors.New("count max must be positive") + } + if params.Max >= math.MaxInt32 { + return nil, errors.New("count max is too large") + } + + // sqlc owns the query while the template context safely injects the optional + // schema identifier. Max+1 lets the caller distinguish an exact count at the + // cap from a state whose index range contains additional rows. + row, err := dbsqlc.JobCountByAllStatesCappedRiver(schemaTemplateParam(ctx, params.Schema), e.dbtx, int32(params.Max+1)) + if err != nil { + return nil, fmt.Errorf("error counting jobs by state: %w", err) + } + + return countsByState(row.Available, row.Cancelled, row.Completed, row.Discarded, row.Pending, row.Retryable, row.Running, row.Scheduled), nil +} + +func (e *Executor) JobCountEstimate(ctx context.Context, params *riveruidriver.JobCountEstimateParams) (map[rivertype.JobState]riveruidriver.JobCountEstimateResult, error) { + // EXPLAIN's Plan Rows comes from PostgreSQL's existing ANALYZE statistics, + // so it retains order-of-magnitude telemetry without reading every matching + // row. Failure to read the timestamp doesn't invalidate the estimate itself. + analyzedAt, analyzedAtValid, _ := dbsqlc.JobCountAnalyzedAtRiver(ctx, e.dbtx, params.Schema) + var observedAt *time.Time + if analyzedAtValid { + observedAt = &analyzedAt + } + + type explainPlan struct { + Plan struct { + Rows int `json:"Plan Rows"` //nolint:tagliatelle // PostgreSQL owns this JSON key. + } `json:"Plan"` //nolint:tagliatelle // PostgreSQL owns this JSON key. + } + + estimates := make(map[rivertype.JobState]riveruidriver.JobCountEstimateResult, len(params.States)) + explainCtx := schemaTemplateParam(ctx, params.Schema) + for _, state := range params.States { + rawPlan, err := dbsqlc.JobCountEstimateRiver(explainCtx, e.dbtx, state) + if err != nil { + return nil, fmt.Errorf("error explaining job count for state %q: %w", state, err) + } + + var plans []explainPlan + if err := json.Unmarshal(rawPlan, &plans); err != nil { + return nil, fmt.Errorf("error decoding job count estimate for state %q: %w", state, err) + } + if len(plans) != 1 { + return nil, fmt.Errorf("expected one job count estimate plan for state %q, got %d", state, len(plans)) + } + + estimates[state] = riveruidriver.JobCountEstimateResult{ + Count: plans[0].Plan.Rows, + ObservedAt: observedAt, + } + } + + return estimates, nil +} + +func countsByState(available, cancelled, completed, discarded, pending, retryable, running, scheduled int64) map[rivertype.JobState]int { + return map[rivertype.JobState]int{ + rivertype.JobStateAvailable: int(available), + rivertype.JobStateCancelled: int(cancelled), + rivertype.JobStateCompleted: int(completed), + rivertype.JobStateDiscarded: int(discarded), + rivertype.JobStatePending: int(pending), + rivertype.JobStateRetryable: int(retryable), + rivertype.JobStateRunning: int(running), + rivertype.JobStateScheduled: int(scheduled), + } +} + +func schemaTemplateParam(ctx context.Context, schema string) context.Context { + if schema != "" { + schema = dbutil.SafeIdentifier(schema) + "." + } + + return sqlctemplate.WithReplacements(ctx, map[string]sqlctemplate.Replacement{ + "schema": {Value: schema, Stable: true}, + }, nil) +} + +type riverExecutorWrapper struct { + executor riverdriver.Executor +} + +func (w riverExecutorWrapper) QueryRow(ctx context.Context, query string, args ...any) dbsqlc.Row { + // River's executor owns the dialect-specific sqlctemplate wrapper, so pass + // the query and context through unchanged and let it consume the template. + return w.executor.QueryRow(ctx, query, args...) +} diff --git a/internal/riveruidriver/riveruipostgres/river_ui_postgres_driver_test.go b/internal/riveruidriver/riveruipostgres/river_ui_postgres_driver_test.go new file mode 100644 index 00000000..a62783e5 --- /dev/null +++ b/internal/riveruidriver/riveruipostgres/river_ui_postgres_driver_test.go @@ -0,0 +1,34 @@ +package riveruipostgres_test + +import ( + "context" + "testing" + + "github.com/riverqueue/river/riverdbtest" + "github.com/riverqueue/river/riverdriver" + "github.com/riverqueue/river/riverdriver/riverpgxv5" + "github.com/riverqueue/river/rivershared/riversharedtest" + + "riverqueue.com/riverui/internal/riveruidriver/riveruidrivertest" + "riverqueue.com/riverui/internal/riveruidriver/riveruipostgres" +) + +func TestDriver(t *testing.T) { + t.Parallel() + + ctx := context.Background() + riverDriver := riverpgxv5.New(riversharedtest.DBPool(ctx, t)) + + riveruidrivertest.Exercise( + ctx, + t, + riverdriver.DatabaseNamePostgres, + riveruipostgres.New(), + func(ctx context.Context, t *testing.T) (riverdriver.Executor, string) { + t.Helper() + + tx, schema := riverdbtest.TestTxPgxDriver(ctx, t, riverDriver, nil) + return riverDriver.UnwrapExecutor(tx), schema + }, + ) +} diff --git a/internal/riveruidriver/riveruisqlite/internal/dbsqlc/db.go b/internal/riveruidriver/riveruisqlite/internal/dbsqlc/db.go new file mode 100644 index 00000000..3bebd3a3 --- /dev/null +++ b/internal/riveruidriver/riveruisqlite/internal/dbsqlc/db.go @@ -0,0 +1,24 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.31.0 + +package dbsqlc + +import ( + "context" + "database/sql" +) + +type DBTX interface { + ExecContext(context.Context, string, ...interface{}) (sql.Result, error) + PrepareContext(context.Context, string) (*sql.Stmt, error) + QueryContext(context.Context, string, ...interface{}) (*sql.Rows, error) + QueryRowContext(context.Context, string, ...interface{}) *sql.Row +} + +func New() *Queries { + return &Queries{} +} + +type Queries struct { +} diff --git a/internal/riveruidriver/riveruisqlite/internal/dbsqlc/models.go b/internal/riveruidriver/riveruisqlite/internal/dbsqlc/models.go new file mode 100644 index 00000000..7f213228 --- /dev/null +++ b/internal/riveruidriver/riveruisqlite/internal/dbsqlc/models.go @@ -0,0 +1,26 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.31.0 + +package dbsqlc + +import ( + "time" +) + +type RiverJob struct { + ID int64 + Priority int64 + Queue string + State string + ScheduledAt time.Time +} + +type SqliteStat4 struct { + Tbl string + Idx string + Neq string + Nlt string + Ndlt string + Sample []byte +} diff --git a/internal/riveruidriver/riveruisqlite/internal/dbsqlc/river_job.sql b/internal/riveruidriver/riveruisqlite/internal/dbsqlc/river_job.sql new file mode 100644 index 00000000..4f5dc25e --- /dev/null +++ b/internal/riveruidriver/riveruisqlite/internal/dbsqlc/river_job.sql @@ -0,0 +1,44 @@ +-- This is a minimal representation with only enough schema for sqlc to enable +-- the River UI queries currently defined below. Expand it toward River's full +-- table definition if future UI queries need additional columns. +CREATE TABLE river_job ( + id integer PRIMARY KEY, + priority integer NOT NULL DEFAULT 1, + queue text NOT NULL DEFAULT 'default', + state text NOT NULL DEFAULT 'available', + scheduled_at timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP +); + +-- Built-in table populated by a full ANALYZE when SQLite is compiled with the +-- non-default SQLITE_ENABLE_STAT4 option. Most SQLite builds omit STAT4, and +-- approximate ANALYZE does not populate it, so this is an uncommon optional +-- optimization path. Declaring its shape here is only for sqlc. +CREATE TABLE sqlite_stat4 ( + tbl text NOT NULL, + idx text NOT NULL, + neq text NOT NULL, + nlt text NOT NULL, + ndlt text NOT NULL, + sample blob NOT NULL +); + +-- name: JobCountByAllStatesCapped :one +-- Each subquery follows the remaining columns of River's prioritized fetching +-- index so the database can stop its index scan as soon as the limit is met. +SELECT + (SELECT count(*) FROM (SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'available' ORDER BY queue, priority, scheduled_at, id LIMIT @max) AS limited_available) AS available, + (SELECT count(*) FROM (SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'cancelled' ORDER BY queue, priority, scheduled_at, id LIMIT @max) AS limited_cancelled) AS cancelled, + (SELECT count(*) FROM (SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'completed' ORDER BY queue, priority, scheduled_at, id LIMIT @max) AS limited_completed) AS completed, + (SELECT count(*) FROM (SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'discarded' ORDER BY queue, priority, scheduled_at, id LIMIT @max) AS limited_discarded) AS discarded, + (SELECT count(*) FROM (SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'pending' ORDER BY queue, priority, scheduled_at, id LIMIT @max) AS limited_pending) AS pending, + (SELECT count(*) FROM (SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'retryable' ORDER BY queue, priority, scheduled_at, id LIMIT @max) AS limited_retryable) AS retryable, + (SELECT count(*) FROM (SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'running' ORDER BY queue, priority, scheduled_at, id LIMIT @max) AS limited_running) AS running, + (SELECT count(*) FROM (SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'scheduled' ORDER BY queue, priority, scheduled_at, id LIMIT @max) AS limited_scheduled) AS scheduled; + +-- name: JobCountEstimateStat4 :one +-- When the uncommon STAT4 data is available, aggregate its small number of +-- encoded index samples into one scalar because River's cross-database +-- executor exposes QueryRow only. +SELECT CAST(coalesce(group_concat(neq || ':' || hex(sample), '|'), '') AS text) AS samples +FROM /* TEMPLATE: schema */sqlite_stat4 +WHERE tbl = 'river_job' AND idx = 'river_job_prioritized_fetching_index'; diff --git a/internal/riveruidriver/riveruisqlite/internal/dbsqlc/river_job.sql.go b/internal/riveruidriver/riveruisqlite/internal/dbsqlc/river_job.sql.go new file mode 100644 index 00000000..fe30a058 --- /dev/null +++ b/internal/riveruidriver/riveruisqlite/internal/dbsqlc/river_job.sql.go @@ -0,0 +1,67 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.31.0 +// source: river_job.sql + +package dbsqlc + +import ( + "context" +) + +const jobCountByAllStatesCapped = `-- name: JobCountByAllStatesCapped :one +SELECT + (SELECT count(*) FROM (SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'available' ORDER BY queue, priority, scheduled_at, id LIMIT ?1) AS limited_available) AS available, + (SELECT count(*) FROM (SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'cancelled' ORDER BY queue, priority, scheduled_at, id LIMIT ?1) AS limited_cancelled) AS cancelled, + (SELECT count(*) FROM (SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'completed' ORDER BY queue, priority, scheduled_at, id LIMIT ?1) AS limited_completed) AS completed, + (SELECT count(*) FROM (SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'discarded' ORDER BY queue, priority, scheduled_at, id LIMIT ?1) AS limited_discarded) AS discarded, + (SELECT count(*) FROM (SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'pending' ORDER BY queue, priority, scheduled_at, id LIMIT ?1) AS limited_pending) AS pending, + (SELECT count(*) FROM (SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'retryable' ORDER BY queue, priority, scheduled_at, id LIMIT ?1) AS limited_retryable) AS retryable, + (SELECT count(*) FROM (SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'running' ORDER BY queue, priority, scheduled_at, id LIMIT ?1) AS limited_running) AS running, + (SELECT count(*) FROM (SELECT 1 FROM /* TEMPLATE: schema */river_job WHERE state = 'scheduled' ORDER BY queue, priority, scheduled_at, id LIMIT ?1) AS limited_scheduled) AS scheduled +` + +type JobCountByAllStatesCappedRow struct { + Available int64 + Cancelled int64 + Completed int64 + Discarded int64 + Pending int64 + Retryable int64 + Running int64 + Scheduled int64 +} + +// Each subquery follows the remaining columns of River's prioritized fetching +// index so the database can stop its index scan as soon as the limit is met. +func (q *Queries) JobCountByAllStatesCapped(ctx context.Context, db DBTX, max int64) (*JobCountByAllStatesCappedRow, error) { + row := db.QueryRowContext(ctx, jobCountByAllStatesCapped, max) + var i JobCountByAllStatesCappedRow + err := row.Scan( + &i.Available, + &i.Cancelled, + &i.Completed, + &i.Discarded, + &i.Pending, + &i.Retryable, + &i.Running, + &i.Scheduled, + ) + return &i, err +} + +const jobCountEstimateStat4 = `-- name: JobCountEstimateStat4 :one +SELECT CAST(coalesce(group_concat(neq || ':' || hex(sample), '|'), '') AS text) AS samples +FROM /* TEMPLATE: schema */sqlite_stat4 +WHERE tbl = 'river_job' AND idx = 'river_job_prioritized_fetching_index' +` + +// When the uncommon STAT4 data is available, aggregate its small number of +// encoded index samples into one scalar because River's cross-database +// executor exposes QueryRow only. +func (q *Queries) JobCountEstimateStat4(ctx context.Context, db DBTX) (string, error) { + row := db.QueryRowContext(ctx, jobCountEstimateStat4) + var samples string + err := row.Scan(&samples) + return samples, err +} diff --git a/internal/riveruidriver/riveruisqlite/internal/dbsqlc/river_queries.go b/internal/riveruidriver/riveruisqlite/internal/dbsqlc/river_queries.go new file mode 100644 index 00000000..ef42302d --- /dev/null +++ b/internal/riveruidriver/riveruisqlite/internal/dbsqlc/river_queries.go @@ -0,0 +1,38 @@ +package dbsqlc + +import "context" + +// Row and RowQuerier avoid coupling these read-only generated queries to +// database/sql's concrete *sql.Row return type. River's executor deliberately +// exposes the same minimal Scan contract across database implementations. +type Row interface { + Scan(dest ...any) error +} + +type RowQuerier interface { + QueryRow(ctx context.Context, query string, args ...any) Row +} + +func JobCountByAllStatesCappedRiver(ctx context.Context, db RowQuerier, maxRows int64) (*JobCountByAllStatesCappedRow, error) { + row := db.QueryRow(ctx, jobCountByAllStatesCapped, maxRows) + var result JobCountByAllStatesCappedRow + err := row.Scan( + &result.Available, + &result.Cancelled, + &result.Completed, + &result.Discarded, + &result.Pending, + &result.Retryable, + &result.Running, + &result.Scheduled, + ) + return &result, err +} + +func JobCountEstimateStat4River(ctx context.Context, db RowQuerier) (string, error) { + var samples string + if err := db.QueryRow(ctx, jobCountEstimateStat4).Scan(&samples); err != nil { + return "", err + } + return samples, nil +} diff --git a/internal/riveruidriver/riveruisqlite/internal/dbsqlc/river_queries_test.go b/internal/riveruidriver/riveruisqlite/internal/dbsqlc/river_queries_test.go new file mode 100644 index 00000000..9baae87d --- /dev/null +++ b/internal/riveruidriver/riveruisqlite/internal/dbsqlc/river_queries_test.go @@ -0,0 +1,25 @@ +package dbsqlc + +import ( + "testing" + + "github.com/stretchr/testify/require" + + "github.com/riverqueue/river/rivertype" +) + +func TestJobCountByAllStatesCappedQueryStates(t *testing.T) { + t.Parallel() + + for _, state := range rivertype.JobStates() { + require.Contains(t, jobCountByAllStatesCapped, "state = '"+string(state)+"'") + } +} + +func TestJobCountEstimateStat4Query(t *testing.T) { + t.Parallel() + + require.Contains(t, jobCountEstimateStat4, "sqlite_stat4") + require.Contains(t, jobCountEstimateStat4, "river_job_prioritized_fetching_index") + require.Contains(t, jobCountEstimateStat4, "group_concat") +} diff --git a/internal/riveruidriver/riveruisqlite/internal/dbsqlc/sqlc.yaml b/internal/riveruidriver/riveruisqlite/internal/dbsqlc/sqlc.yaml new file mode 100644 index 00000000..ad5fb697 --- /dev/null +++ b/internal/riveruidriver/riveruisqlite/internal/dbsqlc/sqlc.yaml @@ -0,0 +1,16 @@ +version: "2" +sql: + - engine: "sqlite" + queries: + - river_job.sql + schema: + - river_job.sql + gen: + go: + package: "dbsqlc" + out: "." + emit_exact_table_names: true + emit_methods_with_db_argument: true + emit_params_struct_pointers: true + emit_pointers_for_null_types: true + emit_result_struct_pointers: true diff --git a/internal/riveruidriver/riveruisqlite/river_ui_sqlite_driver.go b/internal/riveruidriver/riveruisqlite/river_ui_sqlite_driver.go new file mode 100644 index 00000000..47760b12 --- /dev/null +++ b/internal/riveruidriver/riveruisqlite/river_ui_sqlite_driver.go @@ -0,0 +1,242 @@ +// Package riveruisqlite implements River UI's database operations for SQLite. +package riveruisqlite + +import ( + "context" + "encoding/hex" + "errors" + "fmt" + "math" + "strconv" + "strings" + + "github.com/riverqueue/river/riverdriver" + "github.com/riverqueue/river/rivershared/sqlctemplate" + "github.com/riverqueue/river/rivershared/util/dbutil" + "github.com/riverqueue/river/rivertype" + + "riverqueue.com/riverui/internal/riveruidriver" + "riverqueue.com/riverui/internal/riveruidriver/riveruisqlite/internal/dbsqlc" +) + +type Driver struct{} + +// New returns a SQLite River UI driver. Unlike River's primary database +// driver, it doesn't take a pool because River UI only adds queries on top of +// the executor already owned by the caller's River client. +func New() *Driver { + return &Driver{} +} + +func (*Driver) GetExecutor(riverExecutor riverdriver.Executor) riveruidriver.Executor { + return &Executor{ + Executor: riverExecutor, + dbtx: riverExecutorWrapper{executor: riverExecutor}, + } +} + +type Executor struct { + riverdriver.Executor + + dbtx riverExecutorWrapper +} + +var ( + _ riveruidriver.Driver = (*Driver)(nil) + _ riveruidriver.Executor = (*Executor)(nil) +) + +func (e *Executor) JobCountByAllStatesCapped(ctx context.Context, params *riveruidriver.JobCountByAllStatesCappedParams) (map[rivertype.JobState]int, error) { + if params.Max < 1 { + return nil, errors.New("count max must be positive") + } + if params.Max >= math.MaxInt32 { + return nil, errors.New("count max is too large") + } + + row, err := dbsqlc.JobCountByAllStatesCappedRiver(schemaTemplateParam(ctx, params.Schema), e.dbtx, int64(params.Max+1)) + if err != nil { + return nil, fmt.Errorf("error counting jobs by state: %w", err) + } + + return countsByState(row.Available, row.Cancelled, row.Completed, row.Discarded, row.Pending, row.Retryable, row.Running, row.Scheduled), nil +} + +func (e *Executor) JobCountEstimate(ctx context.Context, params *riveruidriver.JobCountEstimateParams) (map[rivertype.JobState]riveruidriver.JobCountEstimateResult, error) { + requestedStates := make(map[rivertype.JobState]struct{}, len(params.States)) + for _, state := range params.States { + switch state { + case rivertype.JobStateAvailable, + rivertype.JobStateCancelled, + rivertype.JobStateCompleted, + rivertype.JobStateDiscarded, + rivertype.JobStatePending, + rivertype.JobStateRetryable, + rivertype.JobStateRunning, + rivertype.JobStateScheduled: + default: + return nil, fmt.Errorf("invalid job state for count estimate: %q", state) + } + requestedStates[state] = struct{}{} + } + + // SQLITE_ENABLE_STAT4 makes a full ANALYZE retain a few encoded samples from + // each index. It is a non-default compile-time feature that most SQLite + // builds omit, and approximate ANALYZE (including the bounded analysis used + // by modern PRAGMA optimize) does not populate it. This is therefore an + // opportunistic, rarely used fast path; the capped query remains the normal + // SQLite behavior. + // + // When STAT4 is available, the first nEq number for River's prioritized + // fetching index is the number of rows matching the sample's leading state. + // Reading those samples is constant-sized and avoids scanning the job index. + stat4Exists, err := e.TableExists(ctx, &riverdriver.TableExistsParams{ + Schema: params.Schema, + Table: "sqlite_stat4", + }) + if err != nil { + return nil, fmt.Errorf("error checking for SQLite STAT4 count estimates: %w", err) + } + if !stat4Exists { + return map[rivertype.JobState]riveruidriver.JobCountEstimateResult{}, nil + } + + samples, err := dbsqlc.JobCountEstimateStat4River(schemaTemplateParam(ctx, params.Schema), e.dbtx) + if err != nil { + return nil, fmt.Errorf("error reading SQLite STAT4 count estimates: %w", err) + } + + estimates, err := parseJobCountEstimateStat4(samples, requestedStates) + if err != nil { + return nil, fmt.Errorf("error decoding SQLite STAT4 count estimates: %w", err) + } + return estimates, nil +} + +func parseJobCountEstimateStat4(samples string, requestedStates map[rivertype.JobState]struct{}) (map[rivertype.JobState]riveruidriver.JobCountEstimateResult, error) { + estimates := make(map[rivertype.JobState]riveruidriver.JobCountEstimateResult, len(requestedStates)) + if samples == "" { + return estimates, nil + } + + for sampleAndCount := range strings.SplitSeq(samples, "|") { + countFields, sampleHex, ok := strings.Cut(sampleAndCount, ":") + if !ok { + return nil, errors.New("STAT4 sample has no count separator") + } + + countText, _, _ := strings.Cut(countFields, " ") + count, err := strconv.Atoi(countText) + if err != nil || count < 0 { + return nil, fmt.Errorf("invalid STAT4 count %q", countText) + } + + sample, err := hex.DecodeString(sampleHex) + if err != nil { + return nil, fmt.Errorf("invalid STAT4 sample encoding: %w", err) + } + stateText, err := sqliteRecordFirstText(sample) + if err != nil { + return nil, fmt.Errorf("invalid STAT4 record: %w", err) + } + state := rivertype.JobState(stateText) + if _, requested := requestedStates[state]; !requested { + continue + } + + if previous, exists := estimates[state]; exists && previous.Count != count { + return nil, fmt.Errorf("inconsistent STAT4 counts for state %q: %d and %d", state, previous.Count, count) + } + estimates[state] = riveruidriver.JobCountEstimateResult{Count: count} + } + + return estimates, nil +} + +// sqliteRecordFirstText decodes the leading text value from SQLite's compact +// record format. sqlite_stat4.sample stores the complete indexed row in this +// format, with the job state first in River's prioritized fetching index. +func sqliteRecordFirstText(record []byte) (string, error) { + headerSize, headerVarintSize, err := sqliteVarint(record) + if err != nil { + return "", err + } + if headerSize > math.MaxInt { + return "", fmt.Errorf("invalid header size %d", headerSize) + } + headerSizeInt := int(headerSize) + if headerSizeInt < headerVarintSize || headerSizeInt > len(record) { + return "", fmt.Errorf("invalid header size %d", headerSize) + } + + serialType, serialTypeSize, err := sqliteVarint(record[headerVarintSize:]) + if err != nil { + return "", err + } + if headerVarintSize+serialTypeSize > headerSizeInt { + return "", errors.New("serial type extends past record header") + } + if serialType < 13 || serialType%2 == 0 { + return "", fmt.Errorf("leading value has non-text serial type %d", serialType) + } + + textSize := (serialType - 13) / 2 + if textSize > math.MaxInt { + return "", errors.New("text value is too large") + } + textSizeInt := int(textSize) + if textSizeInt > len(record)-headerSizeInt { + return "", errors.New("text value extends past record data") + } + return string(record[headerSizeInt : headerSizeInt+textSizeInt]), nil +} + +func sqliteVarint(data []byte) (uint64, int, error) { + var value uint64 + for i := range 8 { + if i >= len(data) { + return 0, 0, errors.New("truncated varint") + } + value = value<<7 | uint64(data[i]&0x7f) + if data[i]&0x80 == 0 { + return value, i + 1, nil + } + } + if len(data) < 9 { + return 0, 0, errors.New("truncated varint") + } + return value<<8 | uint64(data[8]), 9, nil +} + +func countsByState(available, cancelled, completed, discarded, pending, retryable, running, scheduled int64) map[rivertype.JobState]int { + return map[rivertype.JobState]int{ + rivertype.JobStateAvailable: int(available), + rivertype.JobStateCancelled: int(cancelled), + rivertype.JobStateCompleted: int(completed), + rivertype.JobStateDiscarded: int(discarded), + rivertype.JobStatePending: int(pending), + rivertype.JobStateRetryable: int(retryable), + rivertype.JobStateRunning: int(running), + rivertype.JobStateScheduled: int(scheduled), + } +} + +func schemaTemplateParam(ctx context.Context, schema string) context.Context { + if schema != "" { + schema = dbutil.SafeIdentifier(schema) + "." + } + + return sqlctemplate.WithReplacements(ctx, map[string]sqlctemplate.Replacement{ + "schema": {Value: schema, Stable: true}, + }, nil) +} + +type riverExecutorWrapper struct { + executor riverdriver.Executor +} + +func (w riverExecutorWrapper) QueryRow(ctx context.Context, query string, args ...any) dbsqlc.Row { + // River's executor owns the dialect-specific sqlctemplate wrapper, so pass + // the query and context through unchanged and let it consume the template. + return w.executor.QueryRow(ctx, query, args...) +} diff --git a/internal/riveruidriver/riveruisqlite/river_ui_sqlite_driver_internal_test.go b/internal/riveruidriver/riveruisqlite/river_ui_sqlite_driver_internal_test.go new file mode 100644 index 00000000..0acf3b1c --- /dev/null +++ b/internal/riveruidriver/riveruisqlite/river_ui_sqlite_driver_internal_test.go @@ -0,0 +1,89 @@ +package riveruisqlite + +import ( + "encoding/hex" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/riverqueue/river/rivertype" +) + +func TestParseJobCountEstimateStat4(t *testing.T) { + t.Parallel() + + requestedStates := map[rivertype.JobState]struct{}{ + rivertype.JobStateCompleted: {}, + rivertype.JobStateRunning: {}, + } + samples := "100 1 1:" + sqliteRecordTextHex("completed") + + "|10 1 1:" + sqliteRecordTextHex("running") + + "|5 1 1:" + sqliteRecordTextHex("available") + + "|100 1 1:" + sqliteRecordTextHex("completed") + + estimates, err := parseJobCountEstimateStat4(samples, requestedStates) + require.NoError(t, err) + require.Len(t, estimates, 2) + require.Equal(t, 100, estimates[rivertype.JobStateCompleted].Count) + require.Nil(t, estimates[rivertype.JobStateCompleted].ObservedAt) + require.Equal(t, 10, estimates[rivertype.JobStateRunning].Count) +} + +func TestParseJobCountEstimateStat4Empty(t *testing.T) { + t.Parallel() + + estimates, err := parseJobCountEstimateStat4("", map[rivertype.JobState]struct{}{ + rivertype.JobStateCompleted: {}, + }) + require.NoError(t, err) + require.Empty(t, estimates) +} + +func TestParseJobCountEstimateStat4RejectsInconsistentSamples(t *testing.T) { + t.Parallel() + + sample := sqliteRecordTextHex("completed") + _, err := parseJobCountEstimateStat4("100 1:"+sample+"|99 1:"+sample, map[rivertype.JobState]struct{}{ + rivertype.JobStateCompleted: {}, + }) + require.EqualError(t, err, `inconsistent STAT4 counts for state "completed": 100 and 99`) +} + +func TestSQLiteRecordFirstText(t *testing.T) { + t.Parallel() + + record, err := hex.DecodeString(sqliteRecordTextHex("completed")) + require.NoError(t, err) + value, err := sqliteRecordFirstText(record) + require.NoError(t, err) + require.Equal(t, "completed", value) + + for _, testCase := range []struct { + name string + record []byte + }{ + {name: "Empty", record: nil}, + {name: "InvalidHeaderSize", record: []byte{0x03, 0x0f}}, + {name: "TruncatedSerialType", record: []byte{0x02, 0x80}}, + {name: "NonText", record: []byte{0x02, 0x01, 0x01}}, + {name: "TruncatedText", record: append([]byte{0x02, 0x1f}, []byte("short")...)}, + } { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + _, err := sqliteRecordFirstText(testCase.record) + require.Error(t, err) + }) + } +} + +// sqliteRecordTextHex returns a minimal one-column SQLite record containing a +// text value. All River states are short enough for one-byte header varints. +func sqliteRecordTextHex(value string) string { + serialTypes := map[string]byte{ + "available": 0x1f, + "completed": 0x1f, + "running": 0x1b, + } + record := append([]byte{0x02, serialTypes[value]}, []byte(value)...) + return hex.EncodeToString(record) +} diff --git a/internal/riveruidriver/riveruisqlite/river_ui_sqlite_driver_test.go b/internal/riveruidriver/riveruisqlite/river_ui_sqlite_driver_test.go new file mode 100644 index 00000000..68883ec4 --- /dev/null +++ b/internal/riveruidriver/riveruisqlite/river_ui_sqlite_driver_test.go @@ -0,0 +1,42 @@ +package riveruisqlite_test + +import ( + "context" + "database/sql" + "testing" + + _ "modernc.org/sqlite" + + "github.com/riverqueue/river/riverdbtest" + "github.com/riverqueue/river/riverdriver" + "github.com/riverqueue/river/riverdriver/riversqlite" + "github.com/riverqueue/river/rivershared/riversharedtest" + + "riverqueue.com/riverui/internal/riveruidriver/riveruidrivertest" + "riverqueue.com/riverui/internal/riveruidriver/riveruisqlite" +) + +func TestDriver(t *testing.T) { + t.Parallel() + + ctx := context.Background() + + riveruidrivertest.Exercise( + ctx, + t, + riverdriver.DatabaseNameSQLite, + riveruisqlite.New(), + func(ctx context.Context, t *testing.T) (riverdriver.Executor, string) { + t.Helper() + + riverDriver := riversqlite.New(nil) + tx, schema := riverdbtest.TestTx[*sql.Tx](ctx, t, riverDriver, &riverdbtest.TestTxOpts{ + DisableSchemaSharing: true, + ProcurePool: func(ctx context.Context, schema string) (any, string) { + return riversharedtest.DBPoolSQLite(ctx, t, schema), "" + }, + }) + return riverDriver.UnwrapExecutor(tx), schema + }, + ) +} diff --git a/src/components/JobList.tsx b/src/components/JobList.tsx index 8a7ab4f9..ecdd9186 100644 --- a/src/components/JobList.tsx +++ b/src/components/JobList.tsx @@ -482,9 +482,9 @@ const JobList = (props: JobListProps) => { const stateFormatted = state.charAt(0).toUpperCase() + state.slice(1); const jobsInState = useMemo(() => { if (!statesAndCounts) { - return 0; + return BigInt(0); } - return statesAndCounts[state] || 0; + return statesAndCounts[state].count; }, [state, statesAndCounts]); const filterItems = useMemo( diff --git a/src/components/JobStateFilters.test.tsx b/src/components/JobStateFilters.test.tsx index fb76ba35..eecca95b 100644 --- a/src/components/JobStateFilters.test.tsx +++ b/src/components/JobStateFilters.test.tsx @@ -1,3 +1,4 @@ +import { StatesAndCounts } from "@services/states"; import { JobState } from "@services/types"; import { createMemoryHistory, @@ -14,30 +15,29 @@ import { describe, expect, test } from "vitest"; import { defaultValues, jobSearchSchema } from "../routes/jobs/index.schema"; import { JobStateFilters } from "./JobStateFilters"; -const rootRoute = createRootRoute({ - component: () => , -}); - -const jobsRoute = createRoute({ - component: () => , - getParentRoute: () => rootRoute, - path: "/jobs", - search: { - middlewares: [stripSearchParams(defaultValues)], - }, - validateSearch: jobSearchSchema, -}); - -const routeTree = rootRoute.addChildren([jobsRoute]); - -const renderWithLocation = async (location: string) => { +const renderWithLocation = async ( + location: string, + statesAndCounts?: StatesAndCounts, +) => { + const rootRoute = createRootRoute({ + component: () => , + }); + const jobsRoute = createRoute({ + component: () => , + getParentRoute: () => rootRoute, + path: "/jobs", + search: { + middlewares: [stripSearchParams(defaultValues)], + }, + validateSearch: jobSearchSchema, + }); const history = createMemoryHistory({ initialEntries: [location], }); const router = createRouter({ history, - routeTree, + routeTree: rootRoute.addChildren([jobsRoute]), }); await router.load(); @@ -45,6 +45,20 @@ const renderWithLocation = async (location: string) => { return render(); }; +const statesAndCounts = ( + overrides: Partial, +): StatesAndCounts => ({ + available: { accuracy: "exact", count: 0n }, + cancelled: { accuracy: "exact", count: 0n }, + completed: { accuracy: "exact", count: 0n }, + discarded: { accuracy: "exact", count: 0n }, + pending: { accuracy: "exact", count: 0n }, + retryable: { accuracy: "exact", count: 0n }, + running: { accuracy: "exact", count: 0n }, + scheduled: { accuracy: "exact", count: 0n }, + ...overrides, +}); + describe("JobStateFilters", () => { test("only the selected state link is active", async () => { await renderWithLocation(`/jobs?state=${JobState.Discarded}`); @@ -64,4 +78,45 @@ describe("JobStateFilters", () => { const runningLink = await screen.findByRole("link", { name: "Running" }); expect(runningLink).toHaveAttribute("data-status", "active"); }); + + test("shows exact, cached, estimated, and lower-bound telemetry", async () => { + const observedAt = new Date("2026-08-10T12:00:00Z"); + await renderWithLocation( + "/jobs", + statesAndCounts({ + available: { + accuracy: "lower_bound", + count: 10_000n, + observedAt, + }, + completed: { + accuracy: "exact_cached", + count: 12_345_678n, + observedAt, + }, + discarded: { + accuracy: "estimated", + count: 987_654n, + observedAt, + }, + running: { accuracy: "exact", count: 2n, observedAt }, + }), + ); + + expect(await screen.findByText("10K+")).toHaveAttribute( + "title", + expect.stringContaining("useful database estimate"), + ); + expect(screen.getByText("12.3M")).toHaveAttribute( + "title", + expect.stringContaining("12,345,678 jobs (exact snapshot"), + ); + expect(screen.getByText("≈987.7K")).toHaveAttribute( + "title", + expect.stringContaining( + "Approximately 987,654 jobs (database statistics", + ), + ); + expect(screen.getByText("2")).toBeInTheDocument(); + }); }); diff --git a/src/components/JobStateFilters.tsx b/src/components/JobStateFilters.tsx index 757119ed..e888d728 100644 --- a/src/components/JobStateFilters.tsx +++ b/src/components/JobStateFilters.tsx @@ -6,10 +6,54 @@ import React, { useMemo } from "react"; import { Badge } from "./Badge"; +const compactCountFormatter = new Intl.NumberFormat("en-US", { + maximumFractionDigits: 1, + notation: "compact", +}); + type JobStateFiltersProps = { statesAndCounts?: StatesAndCounts; }; +const formatFilterItemCount = ( + item: ReturnType[number], +): string => { + switch (item.accuracy) { + case "estimated": + // The approximation marker prevents a planner estimate from looking + // indistinguishable from an exact snapshot. + return `≈${compactCountFormatter.format(item.count)}`; + case "exact": + // Small exact values are easiest to scan without abbreviation. + return item.count.toString(); + case "exact_cached": + // Compact notation retains the useful order of magnitude in a narrow + // sidebar; the tooltip below keeps the full exact snapshot available. + return compactCountFormatter.format(item.count); + case "lower_bound": + // A plus is the strongest claim supported by the bounded index scan. + return `${compactCountFormatter.format(item.count)}+`; + } +}; + +const filterItemCountTitle = ( + item: ReturnType[number], +): string => { + const fullCount = item.count.toLocaleString("en-US"); + const observedAt = item.observedAt?.toLocaleString(); + + switch (item.accuracy) { + case "estimated": + return `Approximately ${fullCount} jobs (database statistics${observedAt ? ` from ${observedAt}` : ""})`; + case "exact": + return `${fullCount} jobs (exact)`; + case "exact_cached": + return `${fullCount} jobs (exact snapshot${observedAt ? ` from ${observedAt}` : ""})`; + case "lower_bound": + return `At least ${fullCount} jobs; an exact snapshot or useful database estimate is not available yet`; + } +}; + export const JobStateFilters: ( props: JobStateFiltersProps, ) => React.JSX.Element = ({ statesAndCounts }) => { @@ -56,8 +100,9 @@ export const JobStateFilters: ( - {item.count.toString()} + {formatFilterItemCount(item)} ) : null} diff --git a/src/services/states.ts b/src/services/states.ts index 96aafc92..307d6980 100644 --- a/src/services/states.ts +++ b/src/services/states.ts @@ -2,14 +2,41 @@ import type { QueryFunction } from "@tanstack/react-query"; import { API } from "@utils/api"; -import type { JobState, SnakeToCamelCase } from "./types"; +import { JobState } from "./types"; + +export type StateCount = { + accuracy: StateCountAccuracy; + count: bigint; + observedAt?: Date; +}; + +export type StateCountAccuracy = + "estimated" | "exact_cached" | "exact" | "lower_bound"; export type StatesAndCounts = { - [Key in JobState as SnakeToCamelCase]: bigint; + [Key in JobState]: StateCount; }; type CountsByStateKey = ["countsByState"]; +type StatesAndCountsFromAPI = { + [Key in JobState]: { + accuracy: StateCountAccuracy; + count: number; + observed_at?: string; + }; +}; + +const stateCountFromAPI = ( + stateCount: StatesAndCountsFromAPI[JobState], +): StateCount => ({ + accuracy: stateCount.accuracy, + count: BigInt(stateCount.count), + observedAt: stateCount.observed_at + ? new Date(stateCount.observed_at) + : undefined, +}); + export const countsByStateKey = (): CountsByStateKey => { return ["countsByState"]; }; @@ -18,7 +45,16 @@ export const countsByState: QueryFunction< StatesAndCounts, CountsByStateKey > = async ({ signal }) => { - return API.get({ path: "/states" }, { signal }).then( - (response) => response, + return API.get({ path: "/states" }, { signal }).then( + (response) => ({ + available: stateCountFromAPI(response.available), + cancelled: stateCountFromAPI(response.cancelled), + completed: stateCountFromAPI(response.completed), + discarded: stateCountFromAPI(response.discarded), + pending: stateCountFromAPI(response.pending), + retryable: stateCountFromAPI(response.retryable), + running: stateCountFromAPI(response.running), + scheduled: stateCountFromAPI(response.scheduled), + }), ); }; diff --git a/src/utils/jobStateFilterItems.ts b/src/utils/jobStateFilterItems.ts index a934dcb0..e1ca5266 100644 --- a/src/utils/jobStateFilterItems.ts +++ b/src/utils/jobStateFilterItems.ts @@ -1,60 +1,64 @@ -import { StatesAndCounts } from "@services/states"; +import { StateCountAccuracy, StatesAndCounts } from "@services/states"; import { JobState } from "@services/types"; export type JobStateFilterItem = { + accuracy: StateCountAccuracy; count: bigint; name: string; + observedAt?: Date; state: JobState; }; export const jobStateFilterItems: ( statesAndCounts: StatesAndCounts | undefined, ) => JobStateFilterItem[] = (statesAndCounts) => { - const getCount = (state: JobState): bigint => { - if (statesAndCounts) { - return BigInt(statesAndCounts[state]); - } - return BigInt(0); + const getStateCount = (state: JobState) => { + return ( + statesAndCounts?.[state] ?? { + accuracy: "exact" as const, + count: BigInt(0), + } + ); }; return [ { - count: getCount(JobState.Pending), + ...getStateCount(JobState.Pending), name: "Pending", state: JobState.Pending, }, { - count: getCount(JobState.Scheduled), + ...getStateCount(JobState.Scheduled), name: "Scheduled", state: JobState.Scheduled, }, { - count: getCount(JobState.Available), + ...getStateCount(JobState.Available), name: "Available", state: JobState.Available, }, { - count: getCount(JobState.Running), + ...getStateCount(JobState.Running), name: "Running", state: JobState.Running, }, { - count: getCount(JobState.Retryable), + ...getStateCount(JobState.Retryable), name: "Retryable", state: JobState.Retryable, }, { - count: getCount(JobState.Cancelled), + ...getStateCount(JobState.Cancelled), name: "Cancelled", state: JobState.Cancelled, }, { - count: getCount(JobState.Discarded), + ...getStateCount(JobState.Discarded), name: "Discarded", state: JobState.Discarded, }, { - count: getCount(JobState.Completed), + ...getStateCount(JobState.Completed), name: "Completed", state: JobState.Completed, },