Skip to content

Commit 41e938a

Browse files
brandurbrandur
committed
Add UnknownConfigure function for durable periodic jobs
Here, add an `UnknownConfigure` function to the pilot for use with rescuing and reconfiguring unknown durable periodic jobs. Full explanation in the counterpart pull request.
1 parent edec589 commit 41e938a

16 files changed

Lines changed: 468 additions & 299 deletions

‎client.go‎

Lines changed: 100 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import (
99
"log/slog"
1010
"os"
1111
"regexp"
12+
"slices"
1213
"strings"
1314
"sync"
1415
"time"
@@ -900,11 +901,32 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client
900901
}
901902

902903
{
904+
pilotUnknownConfigure := client.pilot.PeriodicJobUnknownConfigure()
905+
903906
periodicJobEnqueuer, err := maintenance.NewPeriodicJobEnqueuer(archetype, &maintenance.PeriodicJobEnqueuerConfig{
904907
AdvisoryLockPrefix: config.AdvisoryLockPrefix,
905908
Insert: client.insertMany,
906909
Pilot: client.pilot,
907910
Schema: config.Schema,
911+
UnknownConfigure: func(job *riverpilot.PeriodicJob) *maintenance.UnknownConfigureResult {
912+
if pilotUnknownConfigure == nil {
913+
return nil
914+
}
915+
916+
unknownConfigureRes := pilotUnknownConfigure(job)
917+
918+
if unknownConfigureRes == nil {
919+
return nil
920+
}
921+
922+
return &maintenance.UnknownConfigureResult{
923+
JobConstructor: func() (*rivertype.JobInsertParams, error) {
924+
jobArgs, insertOpts := unknownConfigureRes.JobConstructor()
925+
return insertParamsFromConfigArgsAndOptions(archetype, config, jobArgs, insertOpts)
926+
},
927+
Schedule: unknownConfigureRes.Schedule.Next,
928+
}
929+
},
908930
}, driver.GetExecutor())
909931
if err != nil {
910932
return nil, err
@@ -1510,6 +1532,14 @@ func (c *Client[TTx]) ID() string {
15101532
return c.config.ID
15111533
}
15121534

1535+
// Regular expression to which the format of tags must comply. Mainly, no
1536+
// special characters, and with hyphens in the middle.
1537+
//
1538+
// A key property here (in case this is relaxed in the future) is that commas
1539+
// must never be allowed because they're used as a delimiter during batch job
1540+
// insertion for the `riverdatabasesql` driver.
1541+
var tagRE = regexp.MustCompile(`\A[\w][\w\-]+[\w]\z`)
1542+
15131543
func insertParamsFromConfigArgsAndOptions(archetype *baseservice.Archetype, config *Config, args JobArgs, insertOpts *InsertOpts) (*rivertype.JobInsertParams, error) {
15141544
encodedArgs, err := json.Marshal(args)
15151545
if err != nil {
@@ -1562,11 +1592,11 @@ func insertParamsFromConfigArgsAndOptions(archetype *baseservice.Archetype, conf
15621592
var uniqueOpts UniqueOpts
15631593
if !config.Test.DisableUniqueEnforcement {
15641594
uniqueOpts = insertOpts.UniqueOpts
1565-
if uniqueOpts.isEmpty() {
1595+
if uniqueOptsIsEmpty(&uniqueOpts) {
15661596
uniqueOpts = jobInsertOpts.UniqueOpts
15671597
}
15681598
}
1569-
if err := uniqueOpts.validate(); err != nil {
1599+
if err := uniqueOptsValidate(&uniqueOpts); err != nil {
15701600
return nil, err
15711601
}
15721602

@@ -1587,7 +1617,7 @@ func insertParamsFromConfigArgsAndOptions(archetype *baseservice.Archetype, conf
15871617
State: rivertype.JobStateAvailable,
15881618
Tags: tags,
15891619
}
1590-
if !uniqueOpts.isEmpty() {
1620+
if !uniqueOptsIsEmpty(&uniqueOpts) {
15911621
internalUniqueOpts := (*dbunique.UniqueOpts)(&uniqueOpts)
15921622
insertParams.UniqueKey, err = dbunique.UniqueKey(archetype.Time, internalUniqueOpts, insertParams)
15931623
if err != nil {
@@ -2709,3 +2739,70 @@ func defaultClientIDWithHost(startedAt time.Time, host string) string {
27092739

27102740
return host + "_" + strings.Replace(startedAt.Format(rfc3339Compact), ".", "_", 1)
27112741
}
2742+
2743+
// uniqueOptsIsEmpty returns true for an empty, uninitialized options struct.
2744+
//
2745+
// This is required because we can't check against `UniqueOpts{}` because slices
2746+
// aren't comparable. Unfortunately it makes things a little more brittle
2747+
// comparatively because any new options must also be considered here for things
2748+
// to work.
2749+
//
2750+
// This is an unexported function in `river` so that it doesn't have
2751+
// to be exported from `rivertype` and doesn't become part of the public API.
2752+
func uniqueOptsIsEmpty(opts *rivertype.UniqueOpts) bool {
2753+
return !opts.ByArgs &&
2754+
opts.ByPeriod == time.Duration(0) &&
2755+
!opts.ByQueue &&
2756+
opts.ByState == nil
2757+
}
2758+
2759+
var jobStateAll = rivertype.JobStates() //nolint:gochecknoglobals
2760+
2761+
var requiredV3states = []rivertype.JobState{ //nolint:gochecknoglobals
2762+
rivertype.JobStateAvailable,
2763+
rivertype.JobStatePending,
2764+
rivertype.JobStateRunning,
2765+
rivertype.JobStateScheduled,
2766+
}
2767+
2768+
// uniqueOptsValidate validates the given rivertype.UniqueOpts.
2769+
//
2770+
// This is a function instance of an instance function so that it doesn't have
2771+
// to be exported from `rivertype` and doesn't become part of the public API.
2772+
func uniqueOptsValidate(opts *rivertype.UniqueOpts) error {
2773+
if uniqueOptsIsEmpty(opts) {
2774+
return nil
2775+
}
2776+
2777+
if opts.ByPeriod != time.Duration(0) && opts.ByPeriod < 1*time.Second {
2778+
return errors.New("UniqueOpts.ByPeriod should not be less than 1 second")
2779+
}
2780+
2781+
// Job states are typed, but since the underlying type is a string, users
2782+
// can put anything they want in there.
2783+
for _, state := range opts.ByState {
2784+
// This could be turned to a map lookup, but last I checked the speed
2785+
// difference for tiny slice sizes is negligible, and map lookup might
2786+
// even be slower.
2787+
if !slices.Contains(jobStateAll, state) {
2788+
return fmt.Errorf("UniqueOpts.ByState contains invalid state %q", state)
2789+
}
2790+
}
2791+
2792+
// Skip required states validation if no custom states were provided.
2793+
if len(opts.ByState) == 0 {
2794+
return nil
2795+
}
2796+
2797+
var missingStates []string
2798+
for _, state := range requiredV3states {
2799+
if !slices.Contains(opts.ByState, state) {
2800+
missingStates = append(missingStates, string(state))
2801+
}
2802+
}
2803+
if len(missingStates) > 0 {
2804+
return fmt.Errorf("UniqueOpts.ByState must contain all required states, missing: %s", strings.Join(missingStates, ", "))
2805+
}
2806+
2807+
return nil
2808+
}

‎client_test.go‎

Lines changed: 156 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,8 +34,10 @@ import (
3434
"github.com/riverqueue/river/riverdriver"
3535
"github.com/riverqueue/river/riverdriver/riverpgxv5"
3636
"github.com/riverqueue/river/rivershared/baseservice"
37+
"github.com/riverqueue/river/rivershared/riverpilot"
3738
"github.com/riverqueue/river/rivershared/riversharedmaintenance"
3839
"github.com/riverqueue/river/rivershared/riversharedtest"
40+
"github.com/riverqueue/river/rivershared/startstop"
3941
"github.com/riverqueue/river/rivershared/startstoptest"
4042
"github.com/riverqueue/river/rivershared/testfactory"
4143
"github.com/riverqueue/river/rivershared/util/dbutil"
@@ -5173,6 +5175,67 @@ func Test_Client_Maintenance(t *testing.T) {
51735175
require.Len(t, jobs, 1, "Expected to find exactly one job of kind: "+(periodicJobArgs{}).Kind())
51745176
})
51755177

5178+
t.Run("PeriodicJobEnqueuerUnknownConfigureFromPilotNilResult", func(t *testing.T) {
5179+
t.Parallel()
5180+
5181+
var (
5182+
dbPool = riversharedtest.DBPool(ctx, t)
5183+
config = newTestConfig(t, "")
5184+
pluginDriver = newDriverWithPlugin(t, dbPool)
5185+
pluginPilot = &TestPilotWithUnknownConfigure{}
5186+
)
5187+
pluginDriver.pilot = pluginPilot
5188+
5189+
var unknownJobConfigureCalled bool
5190+
pluginPilot.PeriodicJobUnknownConfigureFunc = func(job *riverpilot.PeriodicJob) *riverpilot.UnknownConfigureResult {
5191+
unknownJobConfigureCalled = true
5192+
return nil
5193+
}
5194+
5195+
client, err := NewClient(pluginDriver, config)
5196+
require.NoError(t, err)
5197+
5198+
svc := maintenance.GetService[*maintenance.PeriodicJobEnqueuer](client.queueMaintainer)
5199+
svc.Config.UnknownConfigure(&riverpilot.PeriodicJob{})
5200+
require.True(t, unknownJobConfigureCalled)
5201+
})
5202+
5203+
t.Run("PeriodicJobEnqueuerUnknownConfigureFromPilotNonNilResult", func(t *testing.T) {
5204+
t.Parallel()
5205+
5206+
var (
5207+
dbPool = riversharedtest.DBPool(ctx, t)
5208+
config = newTestConfig(t, "")
5209+
pluginDriver = newDriverWithPlugin(t, dbPool)
5210+
pluginPilot = &TestPilotWithUnknownConfigure{}
5211+
)
5212+
pluginDriver.pilot = pluginPilot
5213+
5214+
var (
5215+
jobConstructorCalled bool
5216+
unknownJobConfigureCalled bool
5217+
)
5218+
pluginPilot.PeriodicJobUnknownConfigureFunc = func(job *riverpilot.PeriodicJob) *riverpilot.UnknownConfigureResult {
5219+
unknownJobConfigureCalled = true
5220+
return &riverpilot.UnknownConfigureResult{
5221+
JobConstructor: func() (rivertype.JobArgs, *rivertype.InsertOpts) {
5222+
jobConstructorCalled = true
5223+
return &noOpArgs{}, &rivertype.InsertOpts{}
5224+
},
5225+
Schedule: cron.Every(time.Minute),
5226+
}
5227+
}
5228+
5229+
client, err := NewClient(pluginDriver, config)
5230+
require.NoError(t, err)
5231+
5232+
svc := maintenance.GetService[*maintenance.PeriodicJobEnqueuer](client.queueMaintainer)
5233+
unknownConfigureRes := svc.Config.UnknownConfigure(&riverpilot.PeriodicJob{})
5234+
require.True(t, unknownJobConfigureCalled)
5235+
unknownConfigureRes.JobConstructor()
5236+
require.True(t, jobConstructorCalled)
5237+
})
5238+
51765239
t.Run("QueueCleaner", func(t *testing.T) {
51775240
t.Parallel()
51785241

@@ -8173,3 +8236,96 @@ func (f JobArgsWithHooksFunc) Hooks() []rivertype.Hook {
81738236
func (JobArgsWithHooksFunc) MarshalJSON() ([]byte, error) { return []byte("{}"), nil }
81748237

81758238
func (JobArgsWithHooksFunc) UnmarshalJSON([]byte) error { return nil }
8239+
8240+
var _ pilotPlugin = &TestPilotWithUnknownConfigure{}
8241+
8242+
type TestPilotWithUnknownConfigure struct {
8243+
riverpilot.StandardPilot
8244+
PeriodicJobUnknownConfigureFunc func(job *riverpilot.PeriodicJob) *riverpilot.UnknownConfigureResult
8245+
}
8246+
8247+
func (p *TestPilotWithUnknownConfigure) PeriodicJobUnknownConfigure() func(job *riverpilot.PeriodicJob) *riverpilot.UnknownConfigureResult {
8248+
return p.PeriodicJobUnknownConfigureFunc
8249+
}
8250+
8251+
func (p *TestPilotWithUnknownConfigure) PluginServices() []startstop.Service { return nil }
8252+
8253+
func (p *TestPilotWithUnknownConfigure) PluginMaintenanceServices() []startstop.Service { return nil }
8254+
8255+
func TestTagRE(t *testing.T) {
8256+
t.Parallel()
8257+
8258+
require.Regexp(t, tagRE, "aaa")
8259+
require.Regexp(t, tagRE, "_aaa")
8260+
require.Regexp(t, tagRE, "aaa_")
8261+
require.Regexp(t, tagRE, "777")
8262+
require.Regexp(t, tagRE, "my-tag")
8263+
require.Regexp(t, tagRE, "my_tag")
8264+
require.Regexp(t, tagRE, "my-longer-tag")
8265+
require.Regexp(t, tagRE, "my_longer_tag")
8266+
require.Regexp(t, tagRE, "My_Capitalized_Tag")
8267+
require.Regexp(t, tagRE, "ALL_CAPS")
8268+
require.Regexp(t, tagRE, "1_2_3")
8269+
8270+
require.NotRegexp(t, tagRE, "a")
8271+
require.NotRegexp(t, tagRE, "aa")
8272+
require.NotRegexp(t, tagRE, "-aaa")
8273+
require.NotRegexp(t, tagRE, "aaa-")
8274+
require.NotRegexp(t, tagRE, "special@characters$banned")
8275+
require.NotRegexp(t, tagRE, "commas,never,allowed")
8276+
}
8277+
8278+
func TestUniqueOptsIsEmpty(t *testing.T) {
8279+
t.Parallel()
8280+
8281+
require.True(t, uniqueOptsIsEmpty(&UniqueOpts{}))
8282+
require.False(t, uniqueOptsIsEmpty(&UniqueOpts{ByArgs: true}))
8283+
require.False(t, uniqueOptsIsEmpty(&UniqueOpts{ByPeriod: 1 * time.Nanosecond}))
8284+
require.False(t, uniqueOptsIsEmpty(&UniqueOpts{ByQueue: true}))
8285+
require.False(t, uniqueOptsIsEmpty(&UniqueOpts{ByState: []rivertype.JobState{rivertype.JobStateAvailable}}))
8286+
}
8287+
8288+
func TestUniqueOptsValidate(t *testing.T) {
8289+
t.Parallel()
8290+
8291+
require.NoError(t, uniqueOptsValidate(&UniqueOpts{}))
8292+
require.NoError(t, uniqueOptsValidate(&UniqueOpts{
8293+
ByArgs: true,
8294+
ByPeriod: 1 * time.Second,
8295+
ByQueue: true,
8296+
}))
8297+
8298+
require.EqualError(t, uniqueOptsValidate(&UniqueOpts{ByPeriod: 1 * time.Millisecond}), "UniqueOpts.ByPeriod should not be less than 1 second")
8299+
require.EqualError(t, uniqueOptsValidate(&UniqueOpts{ByState: []rivertype.JobState{rivertype.JobState("invalid")}}), `UniqueOpts.ByState contains invalid state "invalid"`)
8300+
8301+
requiredStates := []rivertype.JobState{
8302+
rivertype.JobStateAvailable,
8303+
rivertype.JobStatePending,
8304+
rivertype.JobStateRunning,
8305+
rivertype.JobStateScheduled,
8306+
}
8307+
8308+
for _, state := range requiredStates {
8309+
// Test with each state individually removed from requiredStates to ensure
8310+
// it's validated.
8311+
8312+
// Create a copy of requiredStates without the current state
8313+
var testStates []rivertype.JobState
8314+
for _, s := range requiredStates {
8315+
if s != state {
8316+
testStates = append(testStates, s)
8317+
}
8318+
}
8319+
8320+
// Test validation
8321+
require.EqualError(t, uniqueOptsValidate(&UniqueOpts{ByState: testStates}), "UniqueOpts.ByState must contain all required states, missing: "+string(state))
8322+
}
8323+
8324+
// test with more than one required state missing:
8325+
require.EqualError(t, uniqueOptsValidate(&UniqueOpts{ByState: []rivertype.JobState{
8326+
rivertype.JobStateAvailable,
8327+
rivertype.JobStateScheduled,
8328+
}}), "UniqueOpts.ByState must contain all required states, missing: pending, running")
8329+
8330+
require.NoError(t, uniqueOptsValidate(&UniqueOpts{ByState: rivertype.JobStates()}))
8331+
}

0 commit comments

Comments
 (0)