Add new HookMetricEmit that can emit arbitrary metrics

I was thinking about observability in River the other day, and how if
one were trying to track River performance degradation due to dead
tuples, how you might do so. Surprisingly, I don't think it's really
possible under our current framework -- the best proxy for it is time to
lock jobs, a metric which we don't provide any way of extracting out of
River.

Here, add a new hook called `HookMetricEmit`. The idea behind it is that it
may receive an arbitrary metric that River starts to emit, and handle it
by passing it onto a metrics system of choice, like OTEL, DataDog, or
whatever.

Within `HookMetricEmit`, add two specific metrics to get us started:

* `JobGetAvailableDurationMetric`: Time to lock available jobs. This
  will start to show significant degradation in case of dead tuples.

* `JobGetAvailableCountMetric`: Number of jobs locked while getting
  available. This is somewhat useful in seeing how well batch locking is
  working out in a program (i.e. are we locking one job at a time or
  100?).

Author: brandur

client.go

@@ -226,11 +226,11 @@ type Config struct {
 	// lifecycle (see rivertype.Hook), installed globally.
 	//
 	// The effect of hooks in this list will depend on the specific hook
-	// interfaces they implement, so for example implementing
-	// rivertype.HookInsertBegin will cause the hook to be invoked before a job
-	// is inserted, or implementing rivertype.HookWorkBegin will cause it to be
-	// invoked before a job is worked. Hook structs may implement multiple hook
-	// interfaces.
+	// interfaces they implement. rivertype.HookInsertBegin will cause the hook
+	// to be invoked before a job is inserted. rivertype.HookMetricEmit will
+	// cause the hook to be invoked when River emits a metric. Implementing
+	// rivertype.HookWorkBegin will cause it to be invoked before a job is
+	// worked. Hook structs may implement multiple hook interfaces.
 	//
 	// Order in this list is significant. A hook that appears first will be
 	// entered before a hook that appears later. For any particular phase, order

hook_defaults_funcs.go

@@ -27,6 +27,18 @@ func (f HookInsertBeginFunc) IsHook() bool { return true }
 
 func (f HookInsertBeginFunc) IsPlugin() bool { return true }
 
+// HookMetricEmitFunc is a convenience helper for implementing
+// rivertype.HookMetricEmit using a simple function instead of a struct.
+type HookMetricEmitFunc func(ctx context.Context, params *rivertype.HookMetricEmitParams)
+
+func (f HookMetricEmitFunc) IsHook() bool { return true }
+
+func (f HookMetricEmitFunc) IsPlugin() bool { return true }
+
+func (f HookMetricEmitFunc) MetricEmit(ctx context.Context, params *rivertype.HookMetricEmitParams) {
+	f(ctx, params)
+}
+
 // HookPeriodicJobsStartFunc is a convenience helper for implementing
 // rivertype.HookPeriodicJobsStart using a simple function instead of a struct.
 type HookPeriodicJobsStartFunc func(ctx context.Context, params *rivertype.HookPeriodicJobsStartParams) error

hook_defaults_funcs_test.go

@@ -12,6 +12,10 @@ var (
 	_ rivertype.HookInsertBegin = HookInsertBeginFunc(func(ctx context.Context, params *rivertype.JobInsertParams) error { return nil })
 	_ rivertype.Plugin          = HookInsertBeginFunc(func(ctx context.Context, params *rivertype.JobInsertParams) error { return nil })
 
+	_ rivertype.Hook           = HookMetricEmitFunc(func(ctx context.Context, params *rivertype.HookMetricEmitParams) {})
+	_ rivertype.HookMetricEmit = HookMetricEmitFunc(func(ctx context.Context, params *rivertype.HookMetricEmitParams) {})
+	_ rivertype.Plugin         = HookMetricEmitFunc(func(ctx context.Context, params *rivertype.HookMetricEmitParams) {})
+
 	_ rivertype.Hook                  = HookPeriodicJobsStartFunc(func(ctx context.Context, params *rivertype.HookPeriodicJobsStartParams) error { return nil })
 	_ rivertype.HookPeriodicJobsStart = HookPeriodicJobsStartFunc(func(ctx context.Context, params *rivertype.HookPeriodicJobsStartParams) error { return nil })
 	_ rivertype.Plugin                = HookPeriodicJobsStartFunc(func(ctx context.Context, params *rivertype.HookPeriodicJobsStartParams) error { return nil })

internal/pluginlookup/plugin_lookup.go

@@ -15,6 +15,7 @@ type PluginKind string
 
 const (
 	PluginKindHookInsertBegin       PluginKind = "hook_insert_begin"
+	PluginKindHookMetricEmit        PluginKind = "hook_metric_emit"
 	PluginKindHookPeriodicJobsStart PluginKind = "hook_periodic_jobs_start"
 	PluginKindHookWorkBegin         PluginKind = "hook_work_begin"
 	PluginKindHookWorkEnd           PluginKind = "hook_work_end"
@@ -85,6 +86,9 @@ func NewPluginLookup(plugins []any) PluginLookupInterface {
 		if _, ok := plugin.(rivertype.HookInsertBegin); ok {
 			pluginsByKind[PluginKindHookInsertBegin] = append(pluginsByKind[PluginKindHookInsertBegin], plugin)
 		}
+		if _, ok := plugin.(rivertype.HookMetricEmit); ok {
+			pluginsByKind[PluginKindHookMetricEmit] = append(pluginsByKind[PluginKindHookMetricEmit], plugin)
+		}
 		if _, ok := plugin.(rivertype.HookPeriodicJobsStart); ok {
 			pluginsByKind[PluginKindHookPeriodicJobsStart] = append(pluginsByKind[PluginKindHookPeriodicJobsStart], plugin)
 		}

internal/pluginlookup/plugin_lookup_test.go

@@ -30,6 +30,7 @@ func TestEmptyPluginLookup(t *testing.T) {
 		pluginLookup, _ := setup(t)
 
 		require.Nil(t, pluginLookup.ByKind(PluginKindHookInsertBegin))
+		require.Nil(t, pluginLookup.ByKind(PluginKindHookMetricEmit))
 		require.Nil(t, pluginLookup.ByKind(PluginKindHookWorkBegin))
 		require.Nil(t, pluginLookup.ByKind(PluginKindMiddlewareJobInsert))
 		require.Nil(t, pluginLookup.ByKind(PluginKindMiddlewareWorker))
@@ -199,6 +200,7 @@ func TestPluginLookup(t *testing.T) {
 		lookup, isPluginLookup := NewPluginLookup([]any{
 			&testHookInsertAndWorkBegin{},
 			&testHookInsertBegin{},
+			&testHookMetricEmit{},
 			&testHookWorkBegin{},
 			&testHookWorkEnd{},
 			&testMiddlewareJobInsertAndWorker{},
@@ -219,6 +221,9 @@ func TestPluginLookup(t *testing.T) {
 			&testHookInsertAndWorkBegin{},
 			&testHookInsertBegin{},
 		}, pluginLookup.ByKind(PluginKindHookInsertBegin))
+		require.Equal(t, []any{
+			&testHookMetricEmit{},
+		}, pluginLookup.ByKind(PluginKindHookMetricEmit))
 		require.Equal(t, []any{
 			&testHookInsertAndWorkBegin{},
 			&testHookWorkBegin{},
@@ -236,13 +241,16 @@ func TestPluginLookup(t *testing.T) {
 			&testMiddlewareWorker{},
 		}, pluginLookup.ByKind(PluginKindMiddlewareWorker))
 
-		require.Len(t, pluginLookup.pluginsByKind, 5)
+		require.Len(t, pluginLookup.pluginsByKind, 6)
 
 		// Repeat lookups to make sure we get the same result.
 		require.Equal(t, []any{
 			&testHookInsertAndWorkBegin{},
 			&testHookInsertBegin{},
 		}, pluginLookup.ByKind(PluginKindHookInsertBegin))
+		require.Equal(t, []any{
+			&testHookMetricEmit{},
+		}, pluginLookup.ByKind(PluginKindHookMetricEmit))
 		require.Equal(t, []any{
 			&testHookInsertAndWorkBegin{},
 			&testHookWorkBegin{},
@@ -276,6 +284,7 @@ func TestPluginLookup(t *testing.T) {
 		}
 
 		parallelLookupLoop(PluginKindHookInsertBegin)
+		parallelLookupLoop(PluginKindHookMetricEmit)
 		parallelLookupLoop(PluginKindHookWorkBegin)
 		parallelLookupLoop(PluginKindHookInsertBegin)
 		parallelLookupLoop(PluginKindHookWorkBegin)
@@ -387,6 +396,19 @@ func (t *testHookInsertBegin) InsertBegin(ctx context.Context, params *rivertype
 	return nil
 }
 
+//
+// testHookMetricEmit
+//
+
+var _ rivertype.HookMetricEmit = &testHookMetricEmit{}
+
+type testHookMetricEmit struct{ rivertype.Hook }
+
+func (t *testHookMetricEmit) IsPlugin() bool { return true }
+
+func (t *testHookMetricEmit) MetricEmit(ctx context.Context, params *rivertype.HookMetricEmitParams) {
+}
+
 //
 // testHookMiddlewarePlugin
 //

producer.go

@@ -191,15 +191,16 @@ type producer struct {
 	// Jobs which are currently being worked. Only used by main goroutine.
 	activeJobs map[int64]*jobexecutor.JobExecutor
 
-	completer    jobcompleter.JobCompleter
-	config       *producerConfig
-	id           atomic.Int64 // atomic because it's written at startup and read during shutdown
-	exec         riverdriver.Executor
-	errorHandler jobexecutor.ErrorHandler
-	fetchLimiter *chanutil.DebouncedChan
-	state        riverpilot.ProducerState
-	pilot        riverpilot.Pilot
-	workers      *Workers
+	completer       jobcompleter.JobCompleter
+	config          *producerConfig
+	id              atomic.Int64 // atomic because it's written at startup and read during shutdown
+	exec            riverdriver.Executor
+	errorHandler    jobexecutor.ErrorHandler
+	fetchLimiter    *chanutil.DebouncedChan
+	metricEmitHooks []rivertype.HookMetricEmit // memoized hooks of type HookMetricEmit for reuse in dispatchWork
+	state           riverpilot.ProducerState
+	pilot           riverpilot.Pilot
+	workers         *Workers
 
 	// Receives job IDs to cancel. Written by notifier goroutine, only read from
 	// main goroutine.
@@ -243,7 +244,7 @@ func newProducer(archetype *baseservice.Archetype, exec riverdriver.Executor, pi
 		errorHandler = &errorHandlerAdapter{config.ErrorHandler}
 	}
 
-	return baseservice.Init(archetype, &producer{
+	producer := baseservice.Init(archetype, &producer{
 		activeJobs:     make(map[int64]*jobexecutor.JobExecutor),
 		cancelCh:       make(chan int64, 1000),
 		completer:      config.Completer,
@@ -257,6 +258,10 @@ func newProducer(archetype *baseservice.Archetype, exec riverdriver.Executor, pi
 		retryPolicy:    config.RetryPolicy,
 		workers:        config.Workers,
 	})
+
+	producer.metricEmitHooks = producer.metricEmitHooksFromLookup()
+
+	return producer
 }
 
 // Start starts the producer. It backgrounds a goroutine which is stopped when
@@ -785,7 +790,33 @@ func (p *producer) maybeCancelJob(ctx context.Context, id int64) {
 	executor.Cancel(ctx)
 }
 
+func (p *producer) metricEmitHooksFromLookup() []rivertype.HookMetricEmit {
+	pluginLookup := p.config.PluginLookupGlobal
+	if pluginLookup == nil {
+		return nil
+	}
+
+	plugins := pluginLookup.ByKind(pluginlookup.PluginKindHookMetricEmit)
+	if len(plugins) < 1 {
+		return nil
+	}
+
+	metricEmitHooks := make([]rivertype.HookMetricEmit, len(plugins))
+	for i, plugin := range plugins {
+		metricEmitHooks[i] = plugin.(rivertype.HookMetricEmit) //nolint:forcetypeassert
+	}
+
+	return metricEmitHooks
+}
+
 func (p *producer) dispatchWork(workCtx context.Context, count int, fetchResultCh chan<- producerFetchResult) {
+	// When a queue is paused, innerFetchLoop dispatches with count zero so it can
+	// continue servicing state changes without attempting to lock jobs or emit metrics.
+	if count <= 0 {
+		fetchResultCh <- producerFetchResult{}
+		return
+	}
+
 	// This intentionally removes any deadlines or cancellation from the parent
 	// context because we don't want it to get cancelled if the producer is asked
 	// to shut down. In that situation, we want to finish fetching any jobs we are
@@ -799,6 +830,11 @@ func (p *producer) dispatchWork(workCtx context.Context, count int, fetchResultC
 	// rarely hit, but exists to protect against degenerate cases.
 	const maxAttemptedBy = 100
 
+	var startedAt time.Time
+	if len(p.metricEmitHooks) > 0 {
+		startedAt = time.Now()
+	}
+
 	jobs, err := p.pilot.JobGetAvailable(ctx, p.exec, p.state, &riverdriver.JobGetAvailableParams{
 		ClientID:       p.config.ClientID,
 		MaxAttemptedBy: maxAttemptedBy,
@@ -813,9 +849,30 @@ func (p *producer) dispatchWork(workCtx context.Context, count int, fetchResultC
 		return
 	}
 
+	if len(p.metricEmitHooks) > 0 {
+		p.emitMetric(ctx, &rivertype.HookMetricEmitParams{
+			Metric: &rivertype.JobGetAvailableDurationMetric{
+				Duration: time.Since(startedAt),
+				Queue:    p.config.Queue,
+			},
+		})
+		p.emitMetric(ctx, &rivertype.HookMetricEmitParams{
+			Metric: &rivertype.JobGetAvailableCountMetric{
+				Count: len(jobs),
+				Queue: p.config.Queue,
+			},
+		})
+	}
+
 	fetchResultCh <- producerFetchResult{jobs: jobs}
 }
 
+func (p *producer) emitMetric(ctx context.Context, params *rivertype.HookMetricEmitParams) {
+	for _, hook := range p.metricEmitHooks {
+		hook.MetricEmit(ctx, params)
+	}
+}
+
 // Periodically logs an informational log line giving some insight into the
 // current state of the producer.
 func (p *producer) heartbeatLogLoop(ctx context.Context, wg *sync.WaitGroup) {