Logging package to add context logger for workers and stores output to metadata (#844)

After adding support for job output, it got me thinking that it might
not be _that_ bad if we put in a way for users to be able to store
limited job-specific logging to job records. Like with output, this
would go in TOAST, and although it could be bad if too much data
accumulated, it shouldn't generally effect performance.

Here, prototype what that would look like. It's implemented as a
middleware in a separate `riverlog` package, which must be installed to
a client or worker to become active. Once it is, logging looks like:

    func (w *LoggingWorker) Work(ctx context.Context, job *river.Job[LoggingArgs]) error {
        riverlog.Logger(ctx).InfoContext(ctx, "Logged from worker")
        return nil
    }

The middleware takes a function that returns a slog handler, which has
the purpose of letting the user select between JSON, text, or something
else, and lets them select options like logging level.

I've left in enough flexibility that we should be able to extend the
interface only slightly to implement a system that sends logs to S3
instead of metadata, but metadata is a good default.

I figure that if we go down this route we can add special handling for
it to River UI to make logs easily accessible for installations even
without a full logging platform.

I also ended up adding support to `Client` and `rivertest.Worker` for
`rivershared.BaseService` in hooks and middleware so that if a hook or
middleware has a base service embedded, it gets its archetype
initialized just like other services in the client. This gives us an
easy way to inherit a client's log instead of a separate one having to
be injected which is _very_ convenient. This shouldn't be used outside
of River core packages, but it should be okay if we make use it.

Also added convenience helpers to make middleware like we already had
for hooks:

* `JobInsertMiddlewareFunc`
* `WorkerMiddlewareFunc `

I put these in because they're great for creating a hook or middleware
that's scoped to a single test case and which is a struct that can embed
something like `baseservice.BaseService`.

Author: brandur

CHANGELOG.md

@@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- 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).
+
 ### Changed
 
 - Client no longer returns an error if stopped before startup could complete (previously, it returned the unexported `ErrShutdown`). [PR #841](https://github.com/riverqueue/river/pull/841).

client.go

@@ -598,6 +598,12 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client
 		}
 	}
 
+	for _, hook := range config.Hooks {
+		if withBaseService, ok := hook.(baseservice.WithBaseService); ok {
+			baseservice.Init(archetype, withBaseService)
+		}
+	}
+
 	client := &Client[TTx]{
 		config:               config,
 		driver:               driver,
@@ -637,6 +643,13 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client
 
 			middleware = append(middleware, workerMiddleware)
 		}
+
+		for _, middleware := range middleware {
+			if withBaseService, ok := middleware.(baseservice.WithBaseService); ok {
+				baseservice.Init(archetype, withBaseService)
+			}
+		}
+
 		client.middlewareLookupGlobal = middlewarelookup.NewMiddlewareLookup(middleware)
 	}
 

client_test.go

@@ -649,6 +649,40 @@ func Test_Client(t *testing.T) {
 		require.True(t, insertBeginHookCalled)
 	})
 
+	t.Run("HookArchetypeInitialized", func(t *testing.T) {
+		t.Parallel()
+
+		_, bundle := setup(t)
+
+		type HookWithBaseService struct {
+			baseservice.BaseService
+			HookInsertBeginFunc
+		}
+
+		var (
+			hook       = &HookWithBaseService{}
+			hookCalled = false
+		)
+		hook.HookInsertBeginFunc = func(ctx context.Context, params *rivertype.JobInsertParams) error {
+			hookCalled = true
+			require.NotEmpty(t, hook.Name) // if name is non-empty, it means the base service was initialized properly
+			return nil
+		}
+
+		AddWorker(bundle.config.Workers, WorkFunc(func(ctx context.Context, job *Job[callbackArgs]) error {
+			return nil
+		}))
+
+		bundle.config.Hooks = []rivertype.Hook{hook}
+		client, err := NewClient(riverpgxv5.New(bundle.dbPool), bundle.config)
+		require.NoError(t, err)
+
+		_, err = client.Insert(ctx, callbackArgs{}, nil)
+		require.NoError(t, err)
+
+		require.True(t, hookCalled)
+	})
+
 	t.Run("WithGlobalWorkBeginHook", func(t *testing.T) {
 		t.Parallel()
 
@@ -2852,6 +2886,41 @@ func Test_Client_InsertManyTx(t *testing.T) {
 		require.JSONEq(t, `{"middleware": "called"}`, string(results[0].Job.Metadata))
 	})
 
+	t.Run("MiddlewareArchetypeInitialized", func(t *testing.T) {
+		t.Parallel()
+
+		_, bundle := setup(t)
+
+		config := newTestConfig(t, nil)
+		config.Queues = nil
+
+		type MiddlewareWithBaseService struct {
+			baseservice.BaseService
+			JobInsertMiddlewareFunc
+		}
+
+		var (
+			middleware       = &MiddlewareWithBaseService{}
+			middlewareCalled bool
+		)
+		middleware.JobInsertMiddlewareFunc = func(ctx context.Context, manyParams []*rivertype.JobInsertParams, doInner func(ctx context.Context) ([]*rivertype.JobInsertResult, error)) ([]*rivertype.JobInsertResult, error) {
+			middlewareCalled = true
+			require.NotEmpty(t, middleware.Name) // if name is non-empty, it means the base service was initialized properly
+			return doInner(ctx)
+		}
+
+		config.Middleware = []rivertype.Middleware{middleware}
+		driver := riverpgxv5.New(nil)
+		client, err := NewClient(driver, config)
+		require.NoError(t, err)
+
+		results, err := client.InsertManyTx(ctx, bundle.tx, []InsertManyParams{{Args: noOpArgs{}}})
+		require.NoError(t, err)
+		require.Len(t, results, 1)
+
+		require.True(t, middlewareCalled)
+	})
+
 	t.Run("WithUniqueOpts", func(t *testing.T) {
 		t.Parallel()
 

common_test.go

@@ -7,6 +7,8 @@ import (
 
 	"github.com/riverqueue/river"
 	"github.com/riverqueue/river/rivershared/riversharedtest"
+	"github.com/riverqueue/river/rivershared/util/sliceutil"
+	"github.com/riverqueue/river/rivertype"
 )
 
 //
@@ -29,7 +31,7 @@ func (w *NoOpWorker) Work(ctx context.Context, job *river.Job[NoOpArgs]) error {
 
 // Wait on the given subscription channel for numJobs. Times out with a panic if
 // jobs take too long to be received.
-func waitForNJobs(subscribeChan <-chan *river.Event, numJobs int) {
+func waitForNJobs(subscribeChan <-chan *river.Event, numJobs int) []*rivertype.JobRow { //nolint:unparam
 	var (
 		timeout  = riversharedtest.WaitTimeout()
 		deadline = time.Now().Add(timeout)
@@ -42,7 +44,7 @@ func waitForNJobs(subscribeChan <-chan *river.Event, numJobs int) {
 			events = append(events, event)
 
 			if len(events) >= numJobs {
-				return
+				return sliceutil.Map(events, func(e *river.Event) *rivertype.JobRow { return e.Job })
 			}
 
 		case <-time.After(time.Until(deadline)):

example_graceful_shutdown_test.go

@@ -169,5 +169,5 @@ func Example_gracefulShutdown() {
 	// Received SIGINT/SIGTERM; initiating soft stop (try to wait for jobs to finish)
 	// Received SIGINT/SIGTERM again; initiating hard stop (cancel everything)
 	// Job cancelled
-	// JobExecutor: Job errored; retrying
+	// jobexecutor.JobExecutor: Job errored; retrying
 }

example_job_cancel_from_client_test.go

@@ -99,5 +99,5 @@ func Example_jobCancelFromClient() {
 	}
 
 	// Output:
-	// JobExecutor: job cancelled remotely
+	// jobexecutor.JobExecutor: job cancelled remotely
 }

hook_defaults_funcs.go

@@ -13,8 +13,8 @@ type HookDefaults struct{}
 
 func (d *HookDefaults) IsHook() bool { return true }
 
-// HookInsertBeginFunc is a convenience helper for implementing HookInsertBegin
-// using a simple function instead of a struct.
+// HookInsertBeginFunc is a convenience helper for implementing
+// rivertype.HookInsertBegin using a simple function instead of a struct.
 type HookInsertBeginFunc func(ctx context.Context, params *rivertype.JobInsertParams) error
 
 func (f HookInsertBeginFunc) InsertBegin(ctx context.Context, params *rivertype.JobInsertParams) error {
@@ -23,8 +23,8 @@ func (f HookInsertBeginFunc) InsertBegin(ctx context.Context, params *rivertype.
 
 func (f HookInsertBeginFunc) IsHook() bool { return true }
 
-// HookWorkBeginFunc is a convenience helper for implementing HookworkBegin
-// using a simple function instead of a struct.
+// 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
 
 func (f HookWorkBeginFunc) WorkBegin(ctx context.Context, job *rivertype.JobRow) error {

middleware_defaults.go

@@ -25,6 +25,16 @@ func (d *JobInsertMiddlewareDefaults) InsertMany(ctx context.Context, manyParams
 	return doInner(ctx)
 }
 
+// JobInsertMiddlewareFunc is a convenience helper for implementing
+// rivertype.JobInsertMiddleware using a simple function instead of a struct.
+type JobInsertMiddlewareFunc func(ctx context.Context, manyParams []*rivertype.JobInsertParams, doInner func(ctx context.Context) ([]*rivertype.JobInsertResult, error)) ([]*rivertype.JobInsertResult, error)
+
+func (f JobInsertMiddlewareFunc) InsertMany(ctx context.Context, manyParams []*rivertype.JobInsertParams, doInner func(ctx context.Context) ([]*rivertype.JobInsertResult, error)) ([]*rivertype.JobInsertResult, error) {
+	return f(ctx, manyParams, doInner)
+}
+
+func (f JobInsertMiddlewareFunc) IsMiddleware() bool { return true }
+
 // WorkerInsertMiddlewareDefaults is an embeddable struct that provides default
 // implementations for the rivertype.WorkerMiddleware. Use of this struct is
 // recommended in case rivertype.WorkerMiddleware is expanded in the future so
@@ -36,3 +46,13 @@ type WorkerMiddlewareDefaults struct{ MiddlewareDefaults }
 func (d *WorkerMiddlewareDefaults) Work(ctx context.Context, job *rivertype.JobRow, doInner func(ctx context.Context) error) error {
 	return doInner(ctx)
 }
+
+// WorkerMiddlewareFunc is a convenience helper for implementing
+// rivertype.WorkerMiddleware using a simple function instead of a struct.
+type WorkerMiddlewareFunc func(ctx context.Context, job *rivertype.JobRow, doInner func(ctx context.Context) error) error
+
+func (f WorkerMiddlewareFunc) Work(ctx context.Context, job *rivertype.JobRow, doInner func(ctx context.Context) error) error {
+	return f(ctx, job, doInner)
+}
+
+func (f WorkerMiddlewareFunc) IsMiddleware() bool { return true }

middleware_defaults_test.go

@@ -1,8 +1,26 @@
 package river
 
-import "github.com/riverqueue/river/rivertype"
+import (
+	"context"
+
+	"github.com/riverqueue/river/rivertype"
+)
 
 var (
 	_ rivertype.JobInsertMiddleware = &JobInsertMiddlewareDefaults{}
 	_ rivertype.WorkerMiddleware    = &WorkerMiddlewareDefaults{}
+
+	_ rivertype.JobInsertMiddleware = JobInsertMiddlewareFunc(func(ctx context.Context, manyParams []*rivertype.JobInsertParams, doInner func(ctx context.Context) ([]*rivertype.JobInsertResult, error)) ([]*rivertype.JobInsertResult, error) {
+		return doInner(ctx)
+	})
+	_ rivertype.Middleware = JobInsertMiddlewareFunc(func(ctx context.Context, manyParams []*rivertype.JobInsertParams, doInner func(ctx context.Context) ([]*rivertype.JobInsertResult, error)) ([]*rivertype.JobInsertResult, error) {
+		return doInner(ctx)
+	})
+
+	_ rivertype.Middleware = WorkerMiddlewareFunc(func(ctx context.Context, job *rivertype.JobRow, doInner func(ctx context.Context) error) error {
+		return doInner(ctx)
+	})
+	_ rivertype.WorkerMiddleware = WorkerMiddlewareFunc(func(ctx context.Context, job *rivertype.JobRow, doInner func(ctx context.Context) error) error {
+		return doInner(ctx)
+	})
 )

recorded_output.go

@@ -10,7 +10,10 @@ import (
 	"github.com/riverqueue/river/rivertype"
 )
 
-const maxOutputSize = 32 * 1024 * 1024 // 32MB
+const (
+	maxOutputSizeMB    = 32
+	maxOutputSizeBytes = maxOutputSizeMB * 1024 * 1024
+)
 
 // RecordOutput records output JSON from a job. The "output" can be any
 // JSON-encodable value and will be stored in the database on the job row after
@@ -55,8 +58,8 @@ func RecordOutput(ctx context.Context, output any) error {
 
 	// Postgres JSONB is limited to 255MB, but it would be a bad idea to get
 	// anywhere close to that limit here.
-	if len(metadataUpdatesBytes) > maxOutputSize {
-		return fmt.Errorf("output is too large: %d bytes (max 32 MB)", len(metadataUpdatesBytes))
+	if len(metadataUpdatesBytes) > maxOutputSizeBytes {
+		return fmt.Errorf("output is too large: %d bytes (max %d MB)", len(metadataUpdatesBytes), maxOutputSizeMB)
 	}
 
 	metadataUpdates[rivertype.MetadataKeyOutput] = json.RawMessage(metadataUpdatesBytes)