Add new HookWorkEnd interface that runs after workers finish

Here, add a new complimentary pair for `HookWorkBegin`: `HookWorkEnd`,
which runs after workers finish, taking in an error result. `HookWorkEnd`
hooks may or may not modify the error result, choosing to suppress an
error on pass it along the stack unchanged.

This is driven by trying to add a new `nilerror` contrib package [1]
that helps detect nil error-compliant structs that return non-nil error
interfaces, which is a common footgun in Go [2].

[1] riverqueue/rivercontrib#25
[2] https://go.dev/doc/faq#nil_error

Author: brandur

CHANGELOG.md

@@ -14,6 +14,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added `river/riverlog` containing middleware that injects a context logger to workers that collates log output and persists it with job metadata. This is paired with a River UI enhancement that shows logs in the UI. [PR #844](https://github.com/riverqueue/river/pull/844).
 - Added `JobInsertMiddlewareFunc` and `WorkerMiddlewareFunc` to easily implement middleware with a function instead of a struct. [PR #844](https://github.com/riverqueue/river/pull/844).
 - Added `Config.Schema` which lets a non-default schema be injected explicitly into a River client that'll be used for all database operations. This may be particularly useful for proxies like PgBouncer that may not respect a schema configured in `search_path`. [PR #848](https://github.com/riverqueue/river/pull/848).
+- Added `rivertype.HookWorkEnd` hook interface that runs after a job has been worked. [PR #XXX](https://github.com/riverqueue/river/pull/XXX).
 
 ### Changed
 

hook_defaults_funcs.go

@@ -32,3 +32,13 @@ func (f HookWorkBeginFunc) WorkBegin(ctx context.Context, job *rivertype.JobRow)
 }
 
 func (f HookWorkBeginFunc) IsHook() bool { return true }
+
+// HookWorkEndFunc is a convenience helper for implementing
+// rivertype.HookworkEnd using a simple function instead of a struct.
+type HookWorkEndFunc func(ctx context.Context, err error) error
+
+func (f HookWorkEndFunc) WorkEnd(ctx context.Context, err error) error {
+	return f(ctx, err)
+}
+
+func (f HookWorkEndFunc) IsHook() bool { return true }

internal/hooklookup/hook_lookup.go

@@ -15,6 +15,7 @@ type HookKind string
 const (
 	HookKindInsertBegin HookKind = "insert_begin"
 	HookKindWorkBegin   HookKind = "work_begin"
+	HookKindWorkEnd     HookKind = "work_end"
 )
 
 //
@@ -88,6 +89,12 @@ func (c *hookLookup) ByHookKind(kind HookKind) []rivertype.Hook {
 				c.hooksByKind[kind] = append(c.hooksByKind[kind], typedHook)
 			}
 		}
+	case HookKindWorkEnd:
+		for _, hook := range c.hooks {
+			if typedHook, ok := hook.(rivertype.HookWorkEnd); ok {
+				c.hooksByKind[kind] = append(c.hooksByKind[kind], typedHook)
+			}
+		}
 	}
 
 	return c.hooksByKind[kind]

internal/hooklookup/hook_lookup_test.go

@@ -22,6 +22,7 @@ func TestHookLookup(t *testing.T) {
 			&testHookInsertAndWorkBegin{},
 			&testHookInsertBegin{},
 			&testHookWorkBegin{},
+			&testHookWorkEnd{},
 		}).(*hookLookup), &testBundle{}
 	}
 
@@ -38,8 +39,11 @@ func TestHookLookup(t *testing.T) {
 			&testHookInsertAndWorkBegin{},
 			&testHookWorkBegin{},
 		}, hookLookup.ByHookKind(HookKindWorkBegin))
+		require.Equal(t, []rivertype.Hook{
+			&testHookWorkEnd{},
+		}, hookLookup.ByHookKind(HookKindWorkEnd))
 
-		require.Len(t, hookLookup.hooksByKind, 2)
+		require.Len(t, hookLookup.hooksByKind, 3)
 
 		// Repeat lookups to make sure we get the same result.
 		require.Equal(t, []rivertype.Hook{
@@ -50,6 +54,9 @@ func TestHookLookup(t *testing.T) {
 			&testHookInsertAndWorkBegin{},
 			&testHookWorkBegin{},
 		}, hookLookup.ByHookKind(HookKindWorkBegin))
+		require.Equal(t, []rivertype.Hook{
+			&testHookWorkEnd{},
+		}, hookLookup.ByHookKind(HookKindWorkEnd))
 	})
 
 	t.Run("Stress", func(t *testing.T) {
@@ -118,6 +125,7 @@ func TestJobHookLookup(t *testing.T) {
 
 		require.Nil(t, jobHookLookup.ByJobArgs(&jobArgsNoHooks{}).ByHookKind(HookKindInsertBegin))
 		require.Nil(t, jobHookLookup.ByJobArgs(&jobArgsNoHooks{}).ByHookKind(HookKindWorkBegin))
+		require.Nil(t, jobHookLookup.ByJobArgs(&jobArgsNoHooks{}).ByHookKind(HookKindWorkEnd))
 		require.Equal(t, []rivertype.Hook{
 			&testHookInsertAndWorkBegin{},
 			&testHookInsertBegin{},
@@ -126,12 +134,16 @@ func TestJobHookLookup(t *testing.T) {
 			&testHookInsertAndWorkBegin{},
 			&testHookWorkBegin{},
 		}, jobHookLookup.ByJobArgs(&jobArgsWithCustomHooks{}).ByHookKind(HookKindWorkBegin))
+		require.Equal(t, []rivertype.Hook{
+			&testHookWorkEnd{},
+		}, jobHookLookup.ByJobArgs(&jobArgsWithCustomHooks{}).ByHookKind(HookKindWorkEnd))
 
 		require.Len(t, jobHookLookup.hookLookupByKind, 2)
 
 		// Repeat lookups to make sure we get the same result.
 		require.Nil(t, jobHookLookup.ByJobArgs(&jobArgsNoHooks{}).ByHookKind(HookKindInsertBegin))
 		require.Nil(t, jobHookLookup.ByJobArgs(&jobArgsNoHooks{}).ByHookKind(HookKindWorkBegin))
+		require.Nil(t, jobHookLookup.ByJobArgs(&jobArgsNoHooks{}).ByHookKind(HookKindWorkEnd))
 		require.Equal(t, []rivertype.Hook{
 			&testHookInsertAndWorkBegin{},
 			&testHookInsertBegin{},
@@ -140,6 +152,9 @@ func TestJobHookLookup(t *testing.T) {
 			&testHookInsertAndWorkBegin{},
 			&testHookWorkBegin{},
 		}, jobHookLookup.ByJobArgs(&jobArgsWithCustomHooks{}).ByHookKind(HookKindWorkBegin))
+		require.Equal(t, []rivertype.Hook{
+			&testHookWorkEnd{},
+		}, jobHookLookup.ByJobArgs(&jobArgsWithCustomHooks{}).ByHookKind(HookKindWorkEnd))
 	})
 
 	t.Run("Stress", func(t *testing.T) {
@@ -195,6 +210,7 @@ func (jobArgsWithCustomHooks) Hooks() []rivertype.Hook {
 		&testHookInsertAndWorkBegin{},
 		&testHookInsertBegin{},
 		&testHookWorkBegin{},
+		&testHookWorkEnd{},
 	}
 }
 
@@ -242,3 +258,15 @@ type testHookWorkBegin struct{ rivertype.Hook }
 func (t *testHookWorkBegin) WorkBegin(ctx context.Context, job *rivertype.JobRow) error {
 	return nil
 }
+
+//
+// testHookWorkEnd
+//
+
+var _ rivertype.HookWorkEnd = &testHookWorkEnd{}
+
+type testHookWorkEnd struct{ rivertype.Hook }
+
+func (t *testHookWorkEnd) WorkEnd(ctx context.Context) error {
+	return nil
+}

internal/jobexecutor/job_executor.go

@@ -208,7 +208,18 @@ func (e *JobExecutor) execute(ctx context.Context) (res *jobExecutorResult) {
 		ctx, cancel := execution.MaybeApplyTimeout(ctx, jobTimeout)
 		defer cancel()
 
-		return e.WorkUnit.Work(ctx)
+		err := e.WorkUnit.Work(ctx)
+
+		{
+			for _, hook := range append(
+				e.HookLookupGlobal.ByHookKind(hooklookup.HookKindWorkEnd),
+				e.WorkUnit.HookLookup(e.HookLookupByJob).ByHookKind(hooklookup.HookKindWorkEnd)...,
+			) {
+				err = hook.(rivertype.HookWorkEnd).WorkEnd(ctx, err) //nolint:forcetypeassert
+			}
+		}
+
+		return err
 	})
 
 	executeFunc := execution.MiddlewareChain(

rivertype/river_type.go

@@ -325,6 +325,33 @@ type HookWorkBegin interface {
 	WorkBegin(ctx context.Context, job *JobRow) error
 }
 
+// HookWorkEnd is an interface to a hook that runs after a job has been worked.
+type HookWorkEnd interface {
+	Hook
+
+	// WorkEnd is invoked after a job has been worked with the error result of
+	// the worked job. It's invoked after any middleware has already run.
+	//
+	// WorkEnd may modify a returned work error or pass it through unchanged.
+	// Each returned error is passed through to the next hook and the final
+	// error result is returned from the job executor:
+	//
+	// 	err := e.WorkUnit.Work(ctx)
+	// 	for _, hook := range hooks {
+	// 		err = hook.(rivertype.HookWorkEnd).WorkEnd(ctx, err)
+	// 	}
+	// 	return err
+	//
+	// If a hook does not want to modify an error result, it should make sure to
+	// return whatever error value it received as its argument whether that
+	// error is nil or not.
+	//
+	// Will not receive a common context related to HookWorkBegin because
+	// WorkBegin doesn't return a context. Middleware should be used for this
+	// sort of shared context instead.
+	WorkEnd(ctx context.Context, error error) error
+}
+
 // Middleware is an arbitrary interface for a struct which will execute some
 // arbitrary code at a predefined step in the job lifecycle.
 //