Support for client middleware? #159

Activity

1. CGA1123 commented on Jan 18, 2024

   

   CGA1123

   on Jan 18, 2024

   Last edited by CGA1123

   ContributorAuthor

   More actions

   I've ended up with the following interface for the time being to enable instrumenting enqueues and propagating tracing information

   package riverutil
   
   import (
   	"context"
   
   	"github.com/jackc/pgx/v5"
   	"github.com/riverqueue/river"
   	"github.com/riverqueue/river/rivertype"
   )
   
   // Ensure *river.Client satisfies this interface at compile time.
   var _ EnqueueClient[pgx.Tx] = (*river.Client[pgx.Tx])(nil)
   
   // EnqueueClient defines the subset of the *river.Client interface which is
   // available when initialised without a `*pgxpool.Pool`.
   type EnqueueClient[TTx any] interface {
   	InsertTx(context.Context, TTx, river.JobArgs, *river.InsertOpts) (*rivertype.JobRow, error)
   	InsertManyTx(context.Context, TTx, []river.InsertManyParams) (int64, error)
   }

   And the current implementation for instrumented enqueues (doesn't create any new spans, but propagates the SpanContext through, so the work can know where the job came from and correlate through):

   package olly
   
   import (
   	"context"
   	"fmt"
   
   	"github.com/jackc/pgx/v5"
   	"github.com/CGA1123/riverplayground/riverutil"
   	"github.com/riverqueue/river"
   	"github.com/riverqueue/river/riverdriver/riverpgxv5"
   	"github.com/riverqueue/river/rivertype"
   	"go.opentelemetry.io/otel"
   	"go.opentelemetry.io/otel/propagation"
   )
   
   type enqueueClient struct {
   	river *river.Client[pgx.Tx]
   }
   
   func (ec *enqueueClient) InsertTx(ctx context.Context, tx pgx.Tx, j river.JobArgs, opts *river.InsertOpts) (*rivertype.JobRow, error) {
   	opts = propagateRiverTrace(ctx, j, opts)
   
   	row, err := ec.river.InsertTx(ctx, tx, j, opts)
   	if err != nil {
   		return row, err
   	}
   
   	return row, err
   }
   
   func (ec *enqueueClient) InsertManyTx(ctx context.Context, tx pgx.Tx, jobs []river.InsertManyParams) (int64, error) {
   	count, err := ec.river.InsertManyTx(ctx, tx, jobs)
   	if err != nil {
   		return count, err
   	}
   
   	for _, j := range jobs {
   		j.InsertOpts = propagateRiverTrace(ctx, j.Args, j.InsertOpts)
   	}
   
   	return count, err
   }
   
   func propagateRiverTrace(ctx context.Context, j river.JobArgs, opts *river.InsertOpts) *river.InsertOpts {
   	if opts == nil {
   		opts = &river.InsertOpts{}
   	}
   
   	var tags []string
   	if argsWithOpts, ok := j.(river.JobArgsWithInsertOpts); ok {
   		tags = argsWithOpts.InsertOpts().Tags
   	}
   
   	c := propagation.MapCarrier(map[string]string{})
   	otel.GetTextMapPropagator().Inject(ctx, c)
   	ollyTags := make([]string, 0, len(c))
   	for k, v := range c {
   		ollyTags = append(ollyTags, fmt.Sprintf("%s%s:%s", riverTagPrefix, k, v))
   	}
   
   	// This replicates the behaviour of `river.insertParamsFromArgsAndOptions`
   	if opts.Tags == nil {
   		opts.Tags = append(tags, ollyTags...)
   	} else {
   		opts.Tags = append(opts.Tags, ollyTags...)
   	}
   
   	return opts
   }
   
   // EnqueueClient returns a wrapped `*river.Client` with a reduced set of
   // methods exposed.
   //
   // It includes additional observability and propagates relevant metadata at
   // enqueue-time.
   //
   // This function will panic if `river.NewClient` returns an error.
   func EnqueueClient(ctx context.Context) riverutil.EnqueueClient[pgx.Tx] {
   	c, err := river.NewClient(riverpgxv5.New(nil), &river.Config{})
   	if err != nil {
   		panic(err)
   	}
   
   	return &enqueueClient{river: c}
   }
2. riverqueue

   locked and limited conversation to collaborators on Jan 20, 2024
3. brandur

   converted this issue into a discussion #167 on Jan 20, 2024