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

@@ -204,11 +204,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

@@ -23,6 +23,16 @@ func (f HookInsertBeginFunc) InsertBegin(ctx context.Context, params *rivertype.
 
 func (f HookInsertBeginFunc) IsHook() 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) 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

@@ -11,9 +11,15 @@ var (
 	_ rivertype.Hook            = HookInsertBeginFunc(func(ctx context.Context, params *rivertype.JobInsertParams) error { return nil })
 	_ rivertype.HookInsertBegin = 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.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.Hook          = HookWorkBeginFunc(func(ctx context.Context, job *rivertype.JobRow) error { return nil })
 	_ rivertype.HookWorkBegin = HookWorkBeginFunc(func(ctx context.Context, job *rivertype.JobRow) error { return nil })
+
+	_ rivertype.Hook        = HookWorkEndFunc(func(ctx context.Context, job *rivertype.JobRow, err error) error { return err })
+	_ rivertype.HookWorkEnd = HookWorkEndFunc(func(ctx context.Context, job *rivertype.JobRow, err error) error { return err })
 )

internal/hooklookup/hook_lookup.go

@@ -14,6 +14,7 @@ type HookKind string
 
 const (
 	HookKindInsertBegin       HookKind = "insert_begin"
+	HookKindMetricEmit        HookKind = "metric_emit"
 	HookKindPeriodicJobsStart HookKind = "periodic_job_start"
 	HookKindWorkBegin         HookKind = "work_begin"
 	HookKindWorkEnd           HookKind = "work_end"
@@ -84,6 +85,12 @@ func (c *hookLookup) ByHookKind(kind HookKind) []rivertype.Hook {
 				c.hooksByKind[kind] = append(c.hooksByKind[kind], typedHook)
 			}
 		}
+	case HookKindMetricEmit:
+		for _, hook := range c.hooks {
+			if typedHook, ok := hook.(rivertype.HookMetricEmit); ok {
+				c.hooksByKind[kind] = append(c.hooksByKind[kind], typedHook)
+			}
+		}
 	case HookKindPeriodicJobsStart:
 		for _, hook := range c.hooks {
 			if typedHook, ok := hook.(rivertype.HookPeriodicJobsStart); ok {

internal/hooklookup/hook_lookup_test.go

@@ -21,6 +21,7 @@ func TestHookLookup(t *testing.T) {
 		return NewHookLookup([]rivertype.Hook{ //nolint:forcetypeassert
 			&testHookInsertAndWorkBegin{},
 			&testHookInsertBegin{},
+			&testHookMetricEmit{},
 			&testHookWorkBegin{},
 			&testHookWorkEnd{},
 		}).(*hookLookup), &testBundle{}
@@ -35,6 +36,9 @@ func TestHookLookup(t *testing.T) {
 			&testHookInsertAndWorkBegin{},
 			&testHookInsertBegin{},
 		}, hookLookup.ByHookKind(HookKindInsertBegin))
+		require.Equal(t, []rivertype.Hook{
+			&testHookMetricEmit{},
+		}, hookLookup.ByHookKind(HookKindMetricEmit))
 		require.Equal(t, []rivertype.Hook{
 			&testHookInsertAndWorkBegin{},
 			&testHookWorkBegin{},
@@ -43,13 +47,16 @@ func TestHookLookup(t *testing.T) {
 			&testHookWorkEnd{},
 		}, hookLookup.ByHookKind(HookKindWorkEnd))
 
-		require.Len(t, hookLookup.hooksByKind, 3)
+		require.Len(t, hookLookup.hooksByKind, 4)
 
 		// Repeat lookups to make sure we get the same result.
 		require.Equal(t, []rivertype.Hook{
 			&testHookInsertAndWorkBegin{},
 			&testHookInsertBegin{},
 		}, hookLookup.ByHookKind(HookKindInsertBegin))
+		require.Equal(t, []rivertype.Hook{
+			&testHookMetricEmit{},
+		}, hookLookup.ByHookKind(HookKindMetricEmit))
 		require.Equal(t, []rivertype.Hook{
 			&testHookInsertAndWorkBegin{},
 			&testHookWorkBegin{},
@@ -75,6 +82,7 @@ func TestHookLookup(t *testing.T) {
 		}
 
 		parallelLookupLoop(HookKindInsertBegin)
+		parallelLookupLoop(HookKindMetricEmit)
 		parallelLookupLoop(HookKindWorkBegin)
 		parallelLookupLoop(HookKindInsertBegin)
 		parallelLookupLoop(HookKindWorkBegin)
@@ -100,6 +108,7 @@ func TestEmptyHookLookup(t *testing.T) {
 		hookLookup, _ := setup(t)
 
 		require.Nil(t, hookLookup.ByHookKind(HookKindInsertBegin))
+		require.Nil(t, hookLookup.ByHookKind(HookKindMetricEmit))
 		require.Nil(t, hookLookup.ByHookKind(HookKindWorkBegin))
 	})
 }
@@ -241,6 +250,17 @@ func (t *testHookInsertBegin) InsertBegin(ctx context.Context, params *rivertype
 	return nil
 }
 
+//
+// testHookMetricEmit
+//
+
+var _ rivertype.HookMetricEmit = &testHookMetricEmit{}
+
+type testHookMetricEmit struct{ rivertype.Hook }
+
+func (t *testHookMetricEmit) MetricEmit(ctx context.Context, params *rivertype.HookMetricEmitParams) {
+}
+
 //
 // testHookWorkBegin
 //

producer.go

@@ -181,15 +181,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.
@@ -233,7 +234,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,
@@ -247,6 +248,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
@@ -743,7 +748,33 @@ func (p *producer) maybeCancelJob(ctx context.Context, id int64) {
 	executor.Cancel(ctx)
 }
 
+func (p *producer) metricEmitHooksFromLookup() []rivertype.HookMetricEmit {
+	hookLookup := p.config.HookLookupGlobal
+	if hookLookup == nil {
+		return nil
+	}
+
+	hooks := hookLookup.ByHookKind(hooklookup.HookKindMetricEmit)
+	if len(hooks) < 1 {
+		return nil
+	}
+
+	metricEmitHooks := make([]rivertype.HookMetricEmit, len(hooks))
+	for i, hook := range hooks {
+		metricEmitHooks[i] = hook.(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
@@ -757,6 +788,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,
@@ -771,9 +807,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) {