Add hook HookQueueStateCount + read middleware as hooks and vice versa

Here, add a new hook called `HookQueueStateCount` which gets invoked to
produce job queue count statistics. We do this by adding a new
maintenance service which like other maintenance services, runs only on
the leader, so we only have one client performing counts at any given
time. Furthermore, in order to not introduce a potential operational
problem without opt-in from River users, the counts only run if a
`HookQueueStateCount` hook/middleware has been added to the client.

The reason we do all this to to implement a feature requested by one of
users: for `otelriver` in contrib to be able to emit queue count
metrics, which seems like a pretty reasonable ask for the package to be
able to do, and something that every River user would likely want access
to in their ops charts.

A slight oddity, but which I think is _probably_ okay, is that the new
hook ideally stays a hook, but the existing `otelriver` middleware is a
middleware. It'd be nice not to have to put `otelriver.Middleware` into
both a client's `Hooks` and `Middleware` configuration, so we modify
client to allow for hooks that middleware and middleware which are
hooks. This lets `otelriver.Middleware` continue doing what it was
already doing, but also to start producing new counts as a hook.

Author: brandur

client.go

@@ -657,6 +657,7 @@ type clientTestSignals struct {
 	periodicJobEnqueuer   *maintenance.PeriodicJobEnqueuerTestSignals
 	queueCleaner          *maintenance.QueueCleanerTestSignals
 	queueMaintainerLeader *maintenance.QueueMaintainerLeaderTestSignals
+	queueStateCounter     *maintenance.QueueStateCounterTestSignals
 	reindexer             *maintenance.ReindexerTestSignals
 }
 
@@ -679,6 +680,9 @@ func (ts *clientTestSignals) Init(tb testutil.TestingTB) {
 	if ts.queueMaintainerLeader != nil {
 		ts.queueMaintainerLeader.Init(tb)
 	}
+	if ts.queueStateCounter != nil {
+		ts.queueStateCounter.Init(tb)
+	}
 	if ts.reindexer != nil {
 		ts.reindexer.Init(tb)
 	}
@@ -759,7 +763,7 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client
 		config:               config,
 		driver:               driver,
 		hookLookupByJob:      hooklookup.NewJobHookLookup(),
-		hookLookupGlobal:     hooklookup.NewHookLookup(config.Hooks),
+		hookLookupGlobal:     nil, // initialized below after cross-referencing with middleware
 		producersByQueueName: make(map[string]*producer),
 		testSignals:          clientTestSignals{},
 		workCancel:           func(cause error) {}, // replaced on start, but here in case StopAndCancel is called before start up
@@ -807,6 +811,26 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client
 			}
 		}
 
+		// Cross-reference hooks and middleware: any middleware that also
+		// implements Hook is added to hooks, and any hook that also implements
+		// Middleware is added to middleware. It may have been better to have
+		// Client only take one combined set, but we only realized that we may
+		// need a combined hook/middleware implementation after that ship sailed.
+		hooks := config.Hooks
+
+		for _, mw := range middleware {
+			if hook, ok := mw.(rivertype.Hook); ok {
+				hooks = append(hooks, hook)
+			}
+		}
+
+		for _, hook := range config.Hooks {
+			if mw, ok := hook.(rivertype.Middleware); ok {
+				middleware = append(middleware, mw)
+			}
+		}
+
+		client.hookLookupGlobal = hooklookup.NewHookLookup(hooks)
 		client.middlewareLookupGlobal = middlewarelookup.NewMiddlewareLookup(middleware)
 	}
 
@@ -961,6 +985,16 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client
 			client.testSignals.queueCleaner = &queueCleaner.TestSignals
 		}
 
+		{
+			queueStateCounter := maintenance.NewQueueStateCounter(archetype, &maintenance.QueueStateCounterConfig{
+				HookLookupGlobal: client.hookLookupGlobal,
+				QueueNames:       maputil.Keys(config.Queues),
+				Schema:           config.Schema,
+			}, driver.GetExecutor())
+			maintenanceServices = append(maintenanceServices, queueStateCounter)
+			client.testSignals.queueStateCounter = &queueStateCounter.TestSignals
+		}
+
 		{
 			var scheduleFunc func(time.Time) time.Time
 			if config.ReindexerSchedule != nil {

client_test.go

@@ -24,6 +24,7 @@ import (
 	"github.com/tidwall/sjson"
 
 	"github.com/riverqueue/river/internal/dbunique"
+	"github.com/riverqueue/river/internal/hooklookup"
 	"github.com/riverqueue/river/internal/jobexecutor"
 	"github.com/riverqueue/river/internal/maintenance"
 	"github.com/riverqueue/river/internal/middlewarelookup"
@@ -7626,6 +7627,24 @@ func Test_NewClient_Validations(t *testing.T) {
 				require.Len(t, client.middlewareLookupGlobal.ByMiddlewareKind(middlewarelookup.MiddlewareKindWorker), 1)
 			},
 		},
+		{
+			name: "Middleware implementing Hook is available in hook lookup",
+			configFunc: func(config *Config) {
+				config.Middleware = []rivertype.Middleware{&middlewareWithHook{}}
+			},
+			validateResult: func(t *testing.T, client *Client[pgx.Tx]) { //nolint:thelper
+				require.Len(t, client.hookLookupGlobal.ByHookKind(hooklookup.HookKindWorkBegin), 1)
+			},
+		},
+		{
+			name: "Hook implementing Middleware is available in middleware lookup",
+			configFunc: func(config *Config) {
+				config.Hooks = []rivertype.Hook{&hookWithMiddleware{}}
+			},
+			validateResult: func(t *testing.T, client *Client[pgx.Tx]) { //nolint:thelper
+				require.Len(t, client.middlewareLookupGlobal.ByMiddlewareKind(middlewarelookup.MiddlewareKindWorker), 1)
+			},
+		},
 		{
 			name: "Middleware not allowed with JobInsertMiddleware",
 			configFunc: func(config *Config) {
@@ -8625,3 +8644,31 @@ func (f JobArgsWithHooksFunc) Hooks() []rivertype.Hook {
 func (JobArgsWithHooksFunc) MarshalJSON() ([]byte, error) { return []byte("{}"), nil }
 
 func (JobArgsWithHooksFunc) UnmarshalJSON([]byte) error { return nil }
+
+// middlewareWithHook is a middleware that also implements HookWorkBegin,
+// used to test cross-pollination from middleware to hooks.
+type middlewareWithHook struct {
+	MiddlewareDefaults
+}
+
+func (m *middlewareWithHook) Work(ctx context.Context, job *rivertype.JobRow, doInner func(ctx context.Context) error) error {
+	return doInner(ctx)
+}
+
+func (m *middlewareWithHook) IsHook() bool { return true }
+
+func (m *middlewareWithHook) WorkBegin(ctx context.Context, job *rivertype.JobRow) error {
+	return nil
+}
+
+// hookWithMiddleware is a hook that also implements WorkerMiddleware,
+// used to test cross-pollination from hooks to middleware.
+type hookWithMiddleware struct {
+	HookDefaults
+}
+
+func (h *hookWithMiddleware) IsMiddleware() bool { return true }
+
+func (h *hookWithMiddleware) Work(ctx context.Context, job *rivertype.JobRow, doInner func(ctx context.Context) error) error {
+	return doInner(ctx)
+}

hook_defaults_funcs.go

@@ -33,6 +33,16 @@ func (f HookPeriodicJobsStartFunc) Start(ctx context.Context, params *rivertype.
 	return f(ctx, params)
 }
 
+// HookQueueStateCountFunc is a convenience helper for implementing
+// rivertype.HookQueueStateCount using a simple function instead of a struct.
+type HookQueueStateCountFunc func(ctx context.Context, params *rivertype.HookQueueStateCountParams)
+
+func (f HookQueueStateCountFunc) IsHook() bool { return true }
+
+func (f HookQueueStateCountFunc) QueueStateCount(ctx context.Context, params *rivertype.HookQueueStateCountParams) {
+	f(ctx, params)
+}
+
 // HookWorkBeginFunc is a convenience helper for implementing
 // rivertype.HookWorkBegin using a simple function instead of a struct.
 type HookWorkBeginFunc func(ctx context.Context, job *rivertype.JobRow) error

internal/hooklookup/hook_lookup.go

@@ -15,6 +15,7 @@ type HookKind string
 const (
 	HookKindInsertBegin       HookKind = "insert_begin"
 	HookKindPeriodicJobsStart HookKind = "periodic_job_start"
+	HookKindQueueStateCount   HookKind = "queue_state_count"
 	HookKindWorkBegin         HookKind = "work_begin"
 	HookKindWorkEnd           HookKind = "work_end"
 )
@@ -90,6 +91,12 @@ func (c *hookLookup) ByHookKind(kind HookKind) []rivertype.Hook {
 				c.hooksByKind[kind] = append(c.hooksByKind[kind], typedHook)
 			}
 		}
+	case HookKindQueueStateCount:
+		for _, hook := range c.hooks {
+			if typedHook, ok := hook.(rivertype.HookQueueStateCount); ok {
+				c.hooksByKind[kind] = append(c.hooksByKind[kind], typedHook)
+			}
+		}
 	case HookKindWorkBegin:
 		for _, hook := range c.hooks {
 			if typedHook, ok := hook.(rivertype.HookWorkBegin); ok {

internal/maintenance/queue_state_counter.go

@@ -0,0 +1,188 @@
+package maintenance
+
+import (
+	"cmp"
+	"context"
+	"errors"
+	"log/slog"
+	"time"
+
+	"github.com/riverqueue/river/internal/hooklookup"
+	"github.com/riverqueue/river/riverdriver"
+	"github.com/riverqueue/river/rivershared/baseservice"
+	"github.com/riverqueue/river/rivershared/riversharedmaintenance"
+	"github.com/riverqueue/river/rivershared/startstop"
+	"github.com/riverqueue/river/rivershared/testsignal"
+	"github.com/riverqueue/river/rivershared/util/testutil"
+	"github.com/riverqueue/river/rivershared/util/timeutil"
+	"github.com/riverqueue/river/rivertype"
+)
+
+const QueueStateCounterIntervalDefault = 10 * time.Second
+
+var jobStateAll = rivertype.JobStates() //nolint:gochecknoglobals
+
+// QueueStateCounterTestSignals are internal signals used exclusively in tests.
+type QueueStateCounterTestSignals struct {
+	CountedOnce testsignal.TestSignal[struct{}] // notifies when a count pass finishes
+}
+
+func (ts *QueueStateCounterTestSignals) Init(tb testutil.TestingTB) {
+	ts.CountedOnce.Init(tb)
+}
+
+type QueueStateCounterConfig struct {
+	// HookLookupGlobal provides access to globally registered hooks.
+	HookLookupGlobal hooklookup.HookLookupInterface
+
+	// Interval is the amount of time between count runs.
+	Interval time.Duration
+
+	// QueueNames is the list of configured queue names. Counts are emitted for
+	// all of these queues even if they have no jobs, with zero counts for all
+	// states.
+	QueueNames []string
+
+	// Schema where River tables are located. Empty string omits schema, causing
+	// Postgres to default to `search_path`.
+	Schema string
+}
+
+func (c *QueueStateCounterConfig) mustValidate() *QueueStateCounterConfig {
+	if c.Interval <= 0 {
+		panic("QueueStateCounterConfig.Interval must be above zero")
+	}
+
+	return c
+}
+
+// QueueStateCounter periodically counts jobs by queue and state, logging the
+// results. This provides visibility into queue health without requiring
+// external monitoring queries. The maintenance service only runs if there is a
+// HookQueueStateCount hook registered that consumes the counts.
+type QueueStateCounter struct {
+	riversharedmaintenance.QueueMaintainerServiceBase
+	startstop.BaseStartStop
+
+	// exported for test purposes
+	Config      *QueueStateCounterConfig
+	TestSignals QueueStateCounterTestSignals
+
+	exec riverdriver.Executor
+}
+
+func NewQueueStateCounter(archetype *baseservice.Archetype, config *QueueStateCounterConfig, exec riverdriver.Executor) *QueueStateCounter {
+	return baseservice.Init(archetype, &QueueStateCounter{
+		Config: (&QueueStateCounterConfig{
+			HookLookupGlobal: config.HookLookupGlobal,
+			Interval:         cmp.Or(config.Interval, QueueStateCounterIntervalDefault),
+			QueueNames:       config.QueueNames,
+			Schema:           config.Schema,
+		}).mustValidate(),
+		exec: exec,
+	})
+}
+
+func (s *QueueStateCounter) Start(ctx context.Context) error {
+	ctx, shouldStart, started, stopped := s.StartInit(ctx)
+	if !shouldStart {
+		return nil
+	}
+
+	s.StaggerStart(ctx)
+
+	go func() {
+		started()
+		defer stopped() // this defer should come first so it's last out
+
+		s.Logger.DebugContext(ctx, s.Name+riversharedmaintenance.LogPrefixRunLoopStarted)
+		defer s.Logger.DebugContext(ctx, s.Name+riversharedmaintenance.LogPrefixRunLoopStopped)
+
+		// If no hooks are registered, there's no one to send counts to, so
+		// start, but don't do anything.
+		if len(s.Config.HookLookupGlobal.ByHookKind(hooklookup.HookKindQueueStateCount)) < 1 {
+			<-ctx.Done()
+			return
+		}
+
+		ticker := timeutil.NewTickerWithInitialTick(ctx, s.Config.Interval)
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case <-ticker.C:
+			}
+
+			if err := s.runOnce(ctx); err != nil {
+				if !errors.Is(err, context.Canceled) {
+					s.Logger.ErrorContext(ctx, s.Name+": Error counting queue states", slog.String("error", err.Error()))
+				}
+			}
+		}
+	}()
+
+	return nil
+}
+
+func (s *QueueStateCounter) runOnce(ctx context.Context) error {
+	ctx, cancelFunc := context.WithTimeout(ctx, riversharedmaintenance.TimeoutDefault)
+	defer cancelFunc()
+
+	rawCounts, err := s.exec.JobCountByQueueAndState(ctx, &riverdriver.JobCountByQueueAndStateParams{
+		QueueNames: s.Config.QueueNames,
+		Schema:     s.Config.Schema,
+	})
+	if err != nil {
+		return err
+	}
+
+	byQueue := s.buildResults(ctx, rawCounts)
+
+	for _, hook := range s.Config.HookLookupGlobal.ByHookKind(hooklookup.HookKindQueueStateCount) {
+		hook.(rivertype.HookQueueStateCount).QueueStateCount(ctx, &rivertype.HookQueueStateCountParams{ //nolint:forcetypeassert
+			ByQueue: byQueue,
+		})
+	}
+
+	s.TestSignals.CountedOnce.Signal(struct{}{})
+
+	return nil
+}
+
+// buildResults converts raw driver counts into results with all configured
+// queues and all job states filled in (zeroed where needed), logging one line
+// per queue.
+func (s *QueueStateCounter) buildResults(ctx context.Context, rawCounts map[string]map[rivertype.JobState]int) map[string]rivertype.HookQueueStateCountResult {
+	// Ensure all configured queues are present, even if they have no jobs.
+	for _, queue := range s.Config.QueueNames {
+		if rawCounts[queue] == nil {
+			rawCounts[queue] = make(map[rivertype.JobState]int)
+		}
+	}
+
+	countsByQueue := make(map[string]rivertype.HookQueueStateCountResult, len(rawCounts))
+
+	for queue, stateCounts := range rawCounts {
+		attrs := make([]slog.Attr, 0, 2+len(jobStateAll))
+		attrs = append(attrs, slog.String("queue", queue))
+		total := 0
+
+		for _, state := range jobStateAll {
+			if _, ok := stateCounts[state]; !ok {
+				stateCounts[state] = 0
+			}
+			total += stateCounts[state]
+			attrs = append(attrs, slog.Int(string(state), stateCounts[state]))
+		}
+
+		attrs = append(attrs, slog.Int("total", total))
+		s.Logger.LogAttrs(ctx, slog.LevelInfo, s.Name+": Queue state counts", attrs...)
+
+		countsByQueue[queue] = rivertype.HookQueueStateCountResult{
+			ByState: stateCounts,
+			Total:   total,
+		}
+	}
+
+	return countsByQueue
+}