Faster count implementation that's still quite accurate

CHANGELOG.md

@@ -11,6 +11,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Job list: filter jobs matching any of the selected exact tags. [PR #548](https://github.com/riverqueue/riverui/pull/548).
 
+### Changed
+
+- 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).
+
 ### Fixed
 
 - 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_test.go

@@ -3,6 +3,7 @@ package riverui
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"log/slog"
 	"net/http"
 	"net/http/httptest"
@@ -1054,6 +1055,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()
@@ -1092,55 +1112,231 @@ 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
+		})
 
-		const queryCacheSkipThreshold = 3
-		for range queryCacheSkipThreshold + 1 {
+		for range countMax {
 			_ = testfactory.Job(ctx, t, bundle.exec, &testfactory.JobOpts{State: ptrutil.Ptr(rivertype.JobStateAvailable)})
 		}
 
-		_, err := endpoint.queryCacher.RunQuery(ctx)
+		resp, err := apitest.InvokeHandler(ctx, endpoint.Execute, testMountOpts(t), &stateAndCountGetRequest{})
 		require.NoError(t, err)
+		requireExactCounts(t, resp)
+		require.Equal(t, countMax, resp.Available.Count)
+	})
+
+	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: ptrutil.Ptr(rivertype.JobStateAvailable)})
+		}
+
+		_, 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: ptrutil.Ptr(rivertype.JobStateCancelled), FinalizedAt: ptrutil.Ptr(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: ptrutil.Ptr(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: ptrutil.Ptr(rivertype.JobStateCancelled), FinalizedAt: ptrutil.Ptr(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: ptrutil.Ptr(rivertype.JobStateCompleted), FinalizedAt: ptrutil.Ptr(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: ptrutil.Ptr(rivertype.JobStateCompleted), FinalizedAt: ptrutil.Ptr(time.Now())})
+		}
+		for range 10 {
+			_ = testfactory.Job(ctx, t, bundle.exec, &testfactory.JobOpts{State: ptrutil.Ptr(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: ptrutil.Ptr(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 and SQL allowlist synchronized
+	// with River when a job state is added or reordered upstream.
+	require.Equal(t, rivertype.JobStates(), allJobStates)
+}
+
+func TestStateCountExactRefreshPeriod(t *testing.T) {
+	t.Parallel()
+
+	require.Equal(t, stateCountExactRefreshMin, stateCountExactRefreshPeriod(100*time.Millisecond, nil))
+	require.Equal(t, 200*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 TestJobStateSQLLiteral(t *testing.T) {
+	t.Parallel()
+
+	expected := map[rivertype.JobState]string{
+		rivertype.JobStateAvailable: "'available'",
+		rivertype.JobStateCancelled: "'cancelled'",
+		rivertype.JobStateCompleted: "'completed'",
+		rivertype.JobStateDiscarded: "'discarded'",
+		rivertype.JobStatePending:   "'pending'",
+		rivertype.JobStateRetryable: "'retryable'",
+		rivertype.JobStateRunning:   "'running'",
+		rivertype.JobStateScheduled: "'scheduled'",
+	}
+	for state, expectedLiteral := range expected {
+		literal, err := jobStateSQLLiteral(state)
+		require.NoError(t, err)
+		require.Equal(t, expectedLiteral, literal)
+	}
+
+	_, err := jobStateSQLLiteral(rivertype.JobState("completed'; DROP TABLE river_job; --"))
+	require.EqualError(t, err, `invalid job state for count estimate: "completed'; DROP TABLE river_job; --"`)
+}
+
+func TestStateAndCountGetEndpointCustomSchema(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	endpoint, bundle := setupEndpointWithCustomSchema(ctx, t, newStateAndCountGetEndpoint)
+	jobParams := testfactory.Job_Build(t, &testfactory.JobOpts{State: ptrutil.Ptr(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)
 }

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)
 			}
 		}
 	}()