Skip to content

Commit 35dddfc

Browse files
brandurbrandur
committed
Faster count implementation that's still quite accurate
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× | > Warm local PostgreSQL 18 averages with all rows in `completed`. The planner estimate includes the `last_analyze` metadata lookup and one state-specific `EXPLAIN`; bounded speedup compares the bounded count with the existing exact count. > Warm local PostgreSQL 18 averages with all rows in `completed`. The planner estimate includes the `last_analyze` metadata lookup and one state-specific `EXPLAIN`; cold estimated total is bounded count plus planner estimate, and speedup compares that total with the existing exact count. 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.
1 parent dd0d14c commit 35dddfc

10 files changed

Lines changed: 853 additions & 124 deletions

File tree

‎CHANGELOG.md‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
1111

1212
- Job list: filter jobs matching any of the selected exact tags. [PR #548](https://github.com/riverqueue/riverui/pull/548).
1313

14+
### Changed
15+
16+
- Job state sidebar: keep large counts responsive while preserving useful magnitude with bounded live counts, adaptively cached exact snapshots, and PostgreSQL estimates. [PR #XXX](https://github.com/riverqueue/riverui/pull/XXX).
17+
1418
### Fixed
1519

1620
- Job args: preserve large numeric JSON values exactly when displaying and copying args, while keeping object keys sorted. [Fixes #593](https://github.com/riverqueue/riverui/issues/593). [PR #594](https://github.com/riverqueue/riverui/pull/594).

‎handler_api_endpoint.go‎

Lines changed: 378 additions & 50 deletions
Large diffs are not rendered by default.

‎handler_api_endpoint_test.go‎

Lines changed: 222 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package riverui
33
import (
44
"context"
55
"encoding/json"
6+
"errors"
67
"log/slog"
78
"net/http"
89
"net/http/httptest"
@@ -1054,6 +1055,25 @@ func TestStateAndCountGetEndpoint(t *testing.T) {
10541055
t.Parallel()
10551056

10561057
ctx := context.Background()
1058+
stateCountsFromResponse := func(resp *stateAndCountGetResponse) map[rivertype.JobState]*stateCountResponse {
1059+
return map[rivertype.JobState]*stateCountResponse{
1060+
rivertype.JobStateAvailable: &resp.Available,
1061+
rivertype.JobStateCancelled: &resp.Cancelled,
1062+
rivertype.JobStateCompleted: &resp.Completed,
1063+
rivertype.JobStateDiscarded: &resp.Discarded,
1064+
rivertype.JobStatePending: &resp.Pending,
1065+
rivertype.JobStateRetryable: &resp.Retryable,
1066+
rivertype.JobStateRunning: &resp.Running,
1067+
rivertype.JobStateScheduled: &resp.Scheduled,
1068+
}
1069+
}
1070+
requireExactCounts := func(t *testing.T, resp *stateAndCountGetResponse) {
1071+
t.Helper()
1072+
for state, stateCount := range stateCountsFromResponse(resp) {
1073+
require.Equal(t, stateCountAccuracyExact, stateCount.Accuracy, state)
1074+
require.NotNil(t, stateCount.ObservedAt, state)
1075+
}
1076+
}
10571077

10581078
t.Run("Success", func(t *testing.T) {
10591079
t.Parallel()
@@ -1092,55 +1112,231 @@ func TestStateAndCountGetEndpoint(t *testing.T) {
10921112

10931113
resp, err := apitest.InvokeHandler(ctx, endpoint.Execute, testMountOpts(t), &stateAndCountGetRequest{})
10941114
require.NoError(t, err)
1095-
require.Equal(t, &stateAndCountGetResponse{
1096-
Available: 1,
1097-
Cancelled: 2,
1098-
Completed: 3,
1099-
Discarded: 4,
1100-
Pending: 5,
1101-
Retryable: 6,
1102-
Running: 7,
1103-
Scheduled: 8,
1104-
}, resp)
1115+
requireExactCounts(t, resp)
1116+
require.Equal(t, 1, resp.Available.Count)
1117+
require.Equal(t, 2, resp.Cancelled.Count)
1118+
require.Equal(t, 3, resp.Completed.Count)
1119+
require.Equal(t, 4, resp.Discarded.Count)
1120+
require.Equal(t, 5, resp.Pending.Count)
1121+
require.Equal(t, 6, resp.Retryable.Count)
1122+
require.Equal(t, 7, resp.Running.Count)
1123+
require.Equal(t, 8, resp.Scheduled.Count)
11051124
})
11061125

1107-
t.Run("WithCachedQueryAboveSkipThreshold", func(t *testing.T) {
1126+
t.Run("AtCountMaxIsExact", func(t *testing.T) {
11081127
t.Parallel()
11091128

1110-
endpoint, bundle := setupEndpoint(ctx, t, newStateAndCountGetEndpoint)
1129+
const countMax = 3
1130+
endpoint, bundle := setupEndpoint(ctx, t, func(bundle apibundle.APIBundle[pgx.Tx]) *stateAndCountGetEndpoint[pgx.Tx] {
1131+
endpoint := newStateAndCountGetEndpoint(bundle)
1132+
endpoint.countMax = countMax
1133+
return endpoint
1134+
})
11111135

1112-
const queryCacheSkipThreshold = 3
1113-
for range queryCacheSkipThreshold + 1 {
1136+
for range countMax {
11141137
_ = testfactory.Job(ctx, t, bundle.exec, &testfactory.JobOpts{State: ptrutil.Ptr(rivertype.JobStateAvailable)})
11151138
}
11161139

1117-
_, err := endpoint.queryCacher.RunQuery(ctx)
1140+
resp, err := apitest.InvokeHandler(ctx, endpoint.Execute, testMountOpts(t), &stateAndCountGetRequest{})
11181141
require.NoError(t, err)
1142+
requireExactCounts(t, resp)
1143+
require.Equal(t, countMax, resp.Available.Count)
1144+
})
1145+
1146+
t.Run("WithExactCachedSnapshot", func(t *testing.T) {
1147+
t.Parallel()
1148+
1149+
const countMax = 3
1150+
endpoint, bundle := setupEndpoint(ctx, t, func(bundle apibundle.APIBundle[pgx.Tx]) *stateAndCountGetEndpoint[pgx.Tx] {
1151+
endpoint := newStateAndCountGetEndpoint(bundle)
1152+
endpoint.countMax = countMax
1153+
return endpoint
1154+
})
1155+
1156+
for range countMax + 1 {
1157+
_ = testfactory.Job(ctx, t, bundle.exec, &testfactory.JobOpts{State: ptrutil.Ptr(rivertype.JobStateAvailable)})
1158+
}
1159+
1160+
_, err := endpoint.boundedQueryCacher.RunQuery(ctx)
1161+
require.NoError(t, err)
1162+
_, err = endpoint.exactQueryCacher.RunQuery(ctx)
1163+
require.NoError(t, err)
1164+
1165+
// Once a state is capped, both caches are reused instead of making an
1166+
// exact count part of the request's latency.
1167+
_ = testfactory.Job(ctx, t, bundle.exec, &testfactory.JobOpts{State: ptrutil.Ptr(rivertype.JobStateCancelled), FinalizedAt: ptrutil.Ptr(time.Now())})
11191168

11201169
resp, err := apitest.InvokeHandler(ctx, endpoint.Execute, testMountOpts(t), &stateAndCountGetRequest{})
11211170
require.NoError(t, err)
1122-
require.Equal(t, &stateAndCountGetResponse{
1123-
Available: queryCacheSkipThreshold + 1,
1124-
}, resp)
1171+
require.Equal(t, countMax+1, resp.Available.Count)
1172+
require.Equal(t, stateCountAccuracyExactCached, resp.Available.Accuracy)
1173+
require.NotNil(t, resp.Available.ObservedAt)
1174+
require.Equal(t, 0, resp.Cancelled.Count)
1175+
require.Equal(t, stateCountAccuracyExact, resp.Cancelled.Accuracy)
11251176
})
11261177

1127-
t.Run("WithCachedQueryBelowSkipThreshold", func(t *testing.T) {
1178+
t.Run("WithExactCachedCount", func(t *testing.T) {
11281179
t.Parallel()
11291180

1130-
endpoint, bundle := setupEndpoint(ctx, t, newStateAndCountGetEndpoint)
1181+
const countMax = 3
1182+
endpoint, bundle := setupEndpoint(ctx, t, func(bundle apibundle.APIBundle[pgx.Tx]) *stateAndCountGetEndpoint[pgx.Tx] {
1183+
endpoint := newStateAndCountGetEndpoint(bundle)
1184+
endpoint.countMax = countMax
1185+
return endpoint
1186+
})
11311187

1132-
const queryCacheSkipThreshold = 3
1133-
for range queryCacheSkipThreshold - 1 {
1188+
for range countMax - 1 {
11341189
_ = testfactory.Job(ctx, t, bundle.exec, &testfactory.JobOpts{State: ptrutil.Ptr(rivertype.JobStateAvailable)})
11351190
}
11361191

1137-
_, err := endpoint.queryCacher.RunQuery(ctx)
1192+
_, err := endpoint.boundedQueryCacher.RunQuery(ctx)
11381193
require.NoError(t, err)
11391194

1195+
// An exact cache result is refreshed inline for the latest counts.
1196+
_ = testfactory.Job(ctx, t, bundle.exec, &testfactory.JobOpts{State: ptrutil.Ptr(rivertype.JobStateCancelled), FinalizedAt: ptrutil.Ptr(time.Now())})
1197+
11401198
resp, err := apitest.InvokeHandler(ctx, endpoint.Execute, testMountOpts(t), &stateAndCountGetRequest{})
11411199
require.NoError(t, err)
1142-
require.Equal(t, &stateAndCountGetResponse{
1143-
Available: queryCacheSkipThreshold - 1,
1144-
}, resp)
1200+
requireExactCounts(t, resp)
1201+
require.Equal(t, countMax-1, resp.Available.Count)
1202+
require.Equal(t, 1, resp.Cancelled.Count)
1203+
})
1204+
1205+
t.Run("WithPlannerEstimate", func(t *testing.T) {
1206+
t.Parallel()
1207+
1208+
const countMax = 3
1209+
endpoint, bundle := setupEndpoint(ctx, t, func(bundle apibundle.APIBundle[pgx.Tx]) *stateAndCountGetEndpoint[pgx.Tx] {
1210+
endpoint := newStateAndCountGetEndpoint(bundle)
1211+
endpoint.countMax = countMax
1212+
return endpoint
1213+
})
1214+
1215+
for range countMax + 1 {
1216+
_ = testfactory.Job(ctx, t, bundle.exec, &testfactory.JobOpts{State: ptrutil.Ptr(rivertype.JobStateCompleted), FinalizedAt: ptrutil.Ptr(time.Now())})
1217+
}
1218+
_, err := endpoint.boundedQueryCacher.RunQuery(ctx)
1219+
require.NoError(t, err)
1220+
1221+
observedAt := time.Now().Add(-5 * time.Minute)
1222+
endpoint.estimateCounts = func(_ context.Context, states []rivertype.JobState) (map[rivertype.JobState]stateCountEstimate, error) {
1223+
require.Equal(t, []rivertype.JobState{rivertype.JobStateCompleted}, states)
1224+
return map[rivertype.JobState]stateCountEstimate{
1225+
rivertype.JobStateCompleted: {Count: 1_000_000, ObservedAt: &observedAt},
1226+
}, nil
1227+
}
1228+
1229+
resp, err := apitest.InvokeHandler(ctx, endpoint.Execute, testMountOpts(t), &stateAndCountGetRequest{})
1230+
require.NoError(t, err)
1231+
require.Equal(t, stateCountResponse{
1232+
Accuracy: stateCountAccuracyEstimated,
1233+
Count: 1_000_000,
1234+
ObservedAt: &observedAt,
1235+
}, resp.Completed)
1236+
})
1237+
1238+
t.Run("ReadsPlannerEstimateFromPostgres", func(t *testing.T) {
1239+
t.Parallel()
1240+
1241+
endpoint, bundle := setupEndpoint(ctx, t, newStateAndCountGetEndpoint)
1242+
for range 100 {
1243+
_ = testfactory.Job(ctx, t, bundle.exec, &testfactory.JobOpts{State: ptrutil.Ptr(rivertype.JobStateCompleted), FinalizedAt: ptrutil.Ptr(time.Now())})
1244+
}
1245+
for range 10 {
1246+
_ = testfactory.Job(ctx, t, bundle.exec, &testfactory.JobOpts{State: ptrutil.Ptr(rivertype.JobStateRunning)})
1247+
}
1248+
require.NoError(t, bundle.exec.Exec(ctx, "ANALYZE river_job"))
1249+
1250+
estimates, err := endpoint.queryEstimatedCounts(ctx, []rivertype.JobState{
1251+
rivertype.JobStateCompleted,
1252+
rivertype.JobStateRunning,
1253+
})
1254+
require.NoError(t, err)
1255+
require.Positive(t, estimates[rivertype.JobStateCompleted].Count)
1256+
require.NotNil(t, estimates[rivertype.JobStateCompleted].ObservedAt)
1257+
require.Greater(t, estimates[rivertype.JobStateCompleted].Count, estimates[rivertype.JobStateRunning].Count)
11451258
})
1259+
1260+
t.Run("WithLowerBoundForStaleEstimate", func(t *testing.T) {
1261+
t.Parallel()
1262+
1263+
const countMax = 3
1264+
endpoint, bundle := setupEndpoint(ctx, t, func(bundle apibundle.APIBundle[pgx.Tx]) *stateAndCountGetEndpoint[pgx.Tx] {
1265+
endpoint := newStateAndCountGetEndpoint(bundle)
1266+
endpoint.countMax = countMax
1267+
return endpoint
1268+
})
1269+
1270+
for range countMax + 1 {
1271+
_ = testfactory.Job(ctx, t, bundle.exec, &testfactory.JobOpts{State: ptrutil.Ptr(rivertype.JobStateAvailable)})
1272+
}
1273+
_, err := endpoint.boundedQueryCacher.RunQuery(ctx)
1274+
require.NoError(t, err)
1275+
endpoint.estimateCounts = func(_ context.Context, _ []rivertype.JobState) (map[rivertype.JobState]stateCountEstimate, error) {
1276+
return map[rivertype.JobState]stateCountEstimate{
1277+
rivertype.JobStateAvailable: {Count: countMax - 1},
1278+
}, nil
1279+
}
1280+
1281+
resp, err := apitest.InvokeHandler(ctx, endpoint.Execute, testMountOpts(t), &stateAndCountGetRequest{})
1282+
require.NoError(t, err)
1283+
require.Equal(t, countMax, resp.Available.Count)
1284+
require.Equal(t, stateCountAccuracyLowerBound, resp.Available.Accuracy)
1285+
require.NotNil(t, resp.Available.ObservedAt)
1286+
})
1287+
}
1288+
1289+
func TestAllJobStates(t *testing.T) {
1290+
t.Parallel()
1291+
1292+
// Keep the endpoint's exhaustive response and SQL allowlist synchronized
1293+
// with River when a job state is added or reordered upstream.
1294+
require.Equal(t, rivertype.JobStates(), allJobStates)
1295+
}
1296+
1297+
func TestStateCountExactRefreshPeriod(t *testing.T) {
1298+
t.Parallel()
1299+
1300+
require.Equal(t, stateCountExactRefreshMin, stateCountExactRefreshPeriod(100*time.Millisecond, nil))
1301+
require.Equal(t, 200*time.Second, stateCountExactRefreshPeriod(2*time.Second, nil))
1302+
require.Equal(t, stateCountExactRefreshMax, stateCountExactRefreshPeriod(time.Hour, nil))
1303+
require.Equal(t, stateCountExactRefreshMax, stateCountExactRefreshPeriod(time.Second, errors.New("database busy")))
1304+
}
1305+
1306+
func TestJobStateSQLLiteral(t *testing.T) {
1307+
t.Parallel()
1308+
1309+
expected := map[rivertype.JobState]string{
1310+
rivertype.JobStateAvailable: "'available'",
1311+
rivertype.JobStateCancelled: "'cancelled'",
1312+
rivertype.JobStateCompleted: "'completed'",
1313+
rivertype.JobStateDiscarded: "'discarded'",
1314+
rivertype.JobStatePending: "'pending'",
1315+
rivertype.JobStateRetryable: "'retryable'",
1316+
rivertype.JobStateRunning: "'running'",
1317+
rivertype.JobStateScheduled: "'scheduled'",
1318+
}
1319+
for state, expectedLiteral := range expected {
1320+
literal, err := jobStateSQLLiteral(state)
1321+
require.NoError(t, err)
1322+
require.Equal(t, expectedLiteral, literal)
1323+
}
1324+
1325+
_, err := jobStateSQLLiteral(rivertype.JobState("completed'; DROP TABLE river_job; --"))
1326+
require.EqualError(t, err, `invalid job state for count estimate: "completed'; DROP TABLE river_job; --"`)
1327+
}
1328+
1329+
func TestStateAndCountGetEndpointCustomSchema(t *testing.T) {
1330+
t.Parallel()
1331+
1332+
ctx := context.Background()
1333+
endpoint, bundle := setupEndpointWithCustomSchema(ctx, t, newStateAndCountGetEndpoint)
1334+
jobParams := testfactory.Job_Build(t, &testfactory.JobOpts{State: ptrutil.Ptr(rivertype.JobStateRunning)})
1335+
jobParams.Schema = bundle.client.Schema()
1336+
_, err := bundle.exec.JobInsertFull(ctx, jobParams)
1337+
require.NoError(t, err)
1338+
1339+
resp, err := apitest.InvokeHandler(ctx, endpoint.Execute, testMountOpts(t), &stateAndCountGetRequest{})
1340+
require.NoError(t, err)
1341+
require.Equal(t, 1, resp.Running.Count)
11461342
}

‎internal/querycacher/query_cacher.go‎

Lines changed: 42 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -23,20 +23,37 @@ type QueryCacher[TRes any] struct {
2323
cachedRes TRes
2424
cachedResSet bool
2525
mu sync.RWMutex
26+
nextTickPeriod func(queryDuration time.Duration, queryErr error) time.Duration
2627
runQuery func(ctx context.Context) (TRes, error)
2728
runQueryTestChan chan struct{} // closed when query is run; for testing
2829
tickPeriod time.Duration // constant normally, but settable for testing
2930
}
3031

32+
type QueryCacherOpts struct {
33+
// NextTickPeriod makes the interval adaptive to the cost and result of the
34+
// preceding query. The period starts after the query finishes, so an
35+
// expensive query can never cause this service to run continuously.
36+
NextTickPeriod func(queryDuration time.Duration, queryErr error) time.Duration
37+
}
38+
3139
func NewQueryCacher[TRes any](archetype *baseservice.Archetype, runQuery func(ctx context.Context) (TRes, error)) *QueryCacher[TRes] {
40+
return NewQueryCacherWithOpts(archetype, runQuery, nil)
41+
}
42+
43+
func NewQueryCacherWithOpts[TRes any](archetype *baseservice.Archetype, runQuery func(ctx context.Context) (TRes, error), opts *QueryCacherOpts) *QueryCacher[TRes] {
3244
// +/- 1s random variance to ticker interval. Makes sure that given multiple
3345
// query caches running simultaneously, they all start and are scheduled a
3446
// little differently to make a thundering herd problem less likely.
3547
randomTickVariance := time.Duration(rand.Float64()*float64(2*time.Second)) - 1*time.Second
48+
var nextTickPeriod func(queryDuration time.Duration, queryErr error) time.Duration
49+
if opts != nil {
50+
nextTickPeriod = opts.NextTickPeriod
51+
}
3652

3753
queryCacher := baseservice.Init(archetype, &QueryCacher[TRes]{
38-
runQuery: runQuery,
39-
tickPeriod: 10*time.Second + randomTickVariance,
54+
nextTickPeriod: nextTickPeriod,
55+
runQuery: runQuery,
56+
tickPeriod: 10*time.Second + randomTickVariance,
4057
})
4158

4259
// TODO(brandur): Push this up into baseservice.
@@ -76,7 +93,7 @@ func (s *QueryCacher[TRes]) RunQuery(ctx context.Context) (TRes, error) {
7693
return emptyRes, err
7794
}
7895

79-
s.Logger.DebugContext(ctx, s.Name+": Ran query and cached result", "duration", time.Since(start), "tick_period", s.tickPeriod)
96+
s.Logger.DebugContext(ctx, s.Name+": Ran query and cached result", "duration", time.Since(start))
8097

8198
s.mu.Lock()
8299
s.cachedRes = res
@@ -104,20 +121,36 @@ func (s *QueryCacher[TRes]) Start(ctx context.Context) error {
104121
started()
105122
defer stopped()
106123

107-
// In case a query runs long and exceeds tickPeriod, time.Ticker will
108-
// drop ticks to compensate.
109-
ticker := time.NewTicker(s.tickPeriod)
110-
defer ticker.Stop()
124+
// A timer is reset only after each query finishes. Unlike a ticker, this
125+
// prevents a slow query from leaving a pending tick that starts another
126+
// expensive query immediately.
127+
timer := time.NewTimer(s.tickPeriod)
128+
defer timer.Stop()
111129

112130
for {
113131
select {
114132
case <-ctx.Done():
115133
return
116134

117-
case <-ticker.C:
118-
if _, err := s.RunQuery(ctx); err != nil {
135+
case <-timer.C:
136+
start := time.Now()
137+
_, err := s.RunQuery(ctx)
138+
queryDuration := time.Since(start)
139+
if err != nil {
119140
s.Logger.ErrorContext(ctx, s.Name+": Error running query", "err", err)
120141
}
142+
143+
nextTickPeriod := s.tickPeriod
144+
if s.nextTickPeriod != nil {
145+
nextTickPeriod = s.nextTickPeriod(queryDuration, err)
146+
}
147+
if nextTickPeriod <= 0 {
148+
// A non-positive period would make the service spin. Falling
149+
// back to the base interval is safer than treating bad options
150+
// as permission to continuously query the database.
151+
nextTickPeriod = s.tickPeriod
152+
}
153+
timer.Reset(nextTickPeriod)
121154
}
122155
}
123156
}()

0 commit comments

Comments
 (0)