engine: refactor rotations to process through job queue

app/initgraphql.go

@@ -42,6 +42,7 @@ func (app *App) initGraphQL(ctx context.Context) error {
 		SWO:                 app.cfg.SWO,
 		APIKeyStore:         app.APIKeyStore,
 		DestReg:             app.DestRegistry,
+		EventBus:            app.EventBus,
 	}
 
 	return nil

engine/engine.go

@@ -95,7 +95,7 @@ func NewEngine(ctx context.Context, db *sql.DB, c *Config) (*Engine, error) {
 		return nil, err
 	}
 
-	rotMgr, err := rotationmanager.NewDB(ctx, db)
+	rotMgr, err := rotationmanager.NewDB(ctx, db, c.RiverDBSQL)
 	if err != nil {
 		return nil, errors.Wrap(err, "rotation management backend")
 	}

engine/rotationmanager/advance.go

@@ -17,20 +17,26 @@ type advance struct {
 }
 
 type rotState struct {
-	rotation.State
-	Version int
+	ShiftStart time.Time
+	Position   int
+	Version    int
 }
 
 // calcAdvance will calculate rotation advancement if it is required. If not, nil is returned
-func calcAdvance(ctx context.Context, t time.Time, rot *rotation.Rotation, state rotState, partCount int) *advance {
+func calcAdvance(ctx context.Context, t time.Time, rot *rotation.Rotation, state rotState, partCount int) (*advance, error) {
 	var mustUpdate bool
 	origPos := state.Position
 
 	// get next shift start time
 	newStart := rot.EndTime(state.ShiftStart)
-	if state.Version == 1 {
+	switch state.Version {
+	case 1:
 		newStart = calcVersion1EndTime(rot, state.ShiftStart)
 		mustUpdate = true
+	case 2:
+		// no-op
+	default:
+		return nil, fmt.Errorf("unknown rotation version (supported: 1,2): %d", state.Version)
 	}
 
 	if state.Position >= partCount {
@@ -48,10 +54,10 @@ func calcAdvance(ctx context.Context, t time.Time, rot *rotation.Rotation, state
 				// If migrating from version 1 to 2 without changing
 				// who's on-call do so silently.
 				silent: state.Version == 1 && state.Position == origPos,
-			}
+			}, nil
 		}
 		// in the future, so nothing to do yet
-		return nil
+		return nil, nil
 	}
 
 	if !newStart.After(t.Add(-15 * time.Minute)) {
@@ -78,5 +84,5 @@ func calcAdvance(ctx context.Context, t time.Time, rot *rotation.Rotation, state
 	return &advance{
 		id:          rot.ID,
 		newPosition: state.Position,
-	}
+	}, nil
 }

engine/rotationmanager/db.go

@@ -4,63 +4,33 @@ import (
 	"context"
 	"database/sql"
 
+	"github.com/riverqueue/river"
 	"github.com/target/goalert/engine/processinglock"
-	"github.com/target/goalert/util"
 )
 
 // DB manages rotations in Postgres.
 type DB struct {
 	lock *processinglock.Lock
 
-	currentTime *sql.Stmt
-
-	lockPart   *sql.Stmt
-	rotate     *sql.Stmt
-	rotateData *sql.Stmt
+	riverDBSQL *river.Client[*sql.Tx]
 }
 
 // Name returns the name of the module.
 func (db *DB) Name() string { return "Engine.RotationManager" }
 
 // NewDB will create a new DB, preparing all statements necessary.
-func NewDB(ctx context.Context, db *sql.DB) (*DB, error) {
+func NewDB(ctx context.Context, db *sql.DB, riverDBSQL *river.Client[*sql.Tx]) (*DB, error) {
 	lock, err := processinglock.NewLock(ctx, db, processinglock.Config{
 		Type:    processinglock.TypeRotation,
 		Version: 2,
 	})
 	if err != nil {
 		return nil, err
 	}
-	p := &util.Prepare{Ctx: ctx, DB: db}
 
 	return &DB{
 		lock: lock,
 
-		currentTime: p.P(`select now()`),
-		lockPart:    p.P(`lock rotation_participants, rotation_state in exclusive mode`),
-		rotate: p.P(`
-			update rotation_state
-			set
-				shift_start = now(),
-				rotation_participant_id = (select id from rotation_participants where rotation_id = $1 and position = $2),
-				version = 2
-			where rotation_id = $1
-		`),
-		rotateData: p.P(`
-			select
-				rot.id,
-				rot."type",
-				rot.start_time,
-				rot.shift_length,
-				rot.time_zone,
-				state.shift_start,
-				state."position",
-				rot.participant_count,
-				state.version
-			from rotations rot
-			join rotation_state state on state.rotation_id = rot.id
-			where $1 or state.rotation_id = $2
-			for update skip locked
-		`),
-	}, p.Err
+		riverDBSQL: riverDBSQL,
+	}, nil
 }

engine/rotationmanager/lookforwork.go

@@ -0,0 +1,63 @@
+package rotationmanager
+
+import (
+	"context"
+	"database/sql"
+	"errors"
+	"fmt"
+	"time"
+
+	"github.com/riverqueue/river"
+	"github.com/target/goalert/gadb"
+)
+
+type LookForWorkArgs struct{}
+
+func (LookForWorkArgs) Kind() string { return "rotation-manager-lfw" }
+
+// lookForWork will schedule jobs for rotations in the entity_updates table.
+func (db *DB) lookForWork(ctx context.Context, j *river.Job[LookForWorkArgs]) error {
+	var hadWork bool
+	err := db.lock.WithTxShared(ctx, func(ctx context.Context, tx *sql.Tx) error {
+		g := gadb.New(tx)
+
+		rotations, err := g.RotMgrFindWork(ctx)
+		if errors.Is(err, sql.ErrNoRows) {
+			// done, no more work
+			return nil
+		}
+		if err != nil {
+			return fmt.Errorf("find work: %w", err)
+		}
+		if len(rotations) == 0 {
+			return nil
+		}
+
+		var params []river.InsertManyParams
+		for _, r := range rotations {
+			params = append(params, river.InsertManyParams{
+				Args: UpdateArgs{RotationID: r},
+				InsertOpts: &river.InsertOpts{
+					Queue:    QueueName,
+					Priority: PriorityEvent,
+				},
+			})
+		}
+
+		_, err = db.riverDBSQL.InsertManyFastTx(ctx, tx, params)
+		if err != nil {
+			return fmt.Errorf("insert many: %w", err)
+		}
+		hadWork = true
+		return nil
+	})
+	if err != nil {
+		return fmt.Errorf("look for work: %w", err)
+	}
+	if !hadWork {
+		return nil
+	}
+
+	// There was work to do, so wait a bit before looking again.
+	return river.JobSnooze(time.Second)
+}

engine/rotationmanager/queries.sql

@@ -0,0 +1,82 @@
+-- name: RotMgrRotationData :one
+-- Get rotation data for a given rotation ID
+SELECT
+    now()::timestamptz AS now,
+    rot.id,
+    rot.type,
+    rot.start_time,
+    rot.shift_length,
+    rot.time_zone,
+    state.position,
+    state.shift_start,
+    state.version,
+    ARRAY (
+        SELECT
+            p.id
+        FROM
+            rotation_participants p
+        WHERE
+            p.rotation_id = rot.id
+        ORDER BY
+            position)::uuid[] AS participants
+    FROM
+        rotations rot
+    LEFT JOIN rotation_state state ON rot.id = state.rotation_id
+WHERE
+    rot.id = @rotation_id;
+
+-- name: RotMgrStart :exec
+-- Start a rotation.
+INSERT INTO rotation_state(rotation_id, position, shift_start, rotation_participant_id)
+SELECT
+    p.rotation_id,
+    0,
+    now(),
+    id
+FROM
+    rotation_participants p
+WHERE
+    p.rotation_id = @rotation_id
+    AND position = 0;
+
+-- name: RotMgrEnd :exec
+-- End a rotation.
+DELETE FROM rotation_state
+WHERE rotation_id = @rotation_id;
+
+-- name: RotMgrUpdate :exec
+-- Update the rotation state.
+UPDATE
+    rotation_state
+SET
+    position = @position,
+    shift_start = now(),
+    rotation_participant_id = @rotation_participant_id
+WHERE
+    rotation_id = @rotation_id;
+
+-- name: RotMgrFindWork :many
+WITH items AS (
+    SELECT
+        id,
+        entity_id
+    FROM
+        entity_updates
+    WHERE
+        entity_type = 'rotation'
+    FOR UPDATE
+        SKIP LOCKED
+    LIMIT 1000
+),
+_delete AS (
+    DELETE FROM entity_updates
+    WHERE id IN (
+            SELECT
+                id
+            FROM
+                items))
+SELECT DISTINCT
+    entity_id
+FROM
+    items;
+

engine/rotationmanager/setup.go

@@ -0,0 +1,54 @@
+package rotationmanager
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/riverqueue/river"
+	"github.com/target/goalert/engine/processinglock"
+	"github.com/target/goalert/event"
+	"github.com/target/goalert/schedule/rotation"
+)
+
+const (
+	QueueName         = "rotation-manager"
+	PriorityScheduled = 1
+	PriorityEvent     = 2
+	PriorityLFW       = 4
+)
+
+var _ processinglock.Setupable = (*DB)(nil) // assert that DB implements processinglock.Setupable
+
+// Setup implements processinglock.Setupable.
+func (db *DB) Setup(ctx context.Context, args processinglock.SetupArgs) error {
+	river.AddWorker(args.Workers, river.WorkFunc(db.updateRotation))
+	river.AddWorker(args.Workers, river.WorkFunc(db.lookForWork))
+
+	event.RegisterJobSource(args.EventBus, func(data rotation.Update) (river.JobArgs, *river.InsertOpts) {
+		return UpdateArgs{RotationID: data.ID}, &river.InsertOpts{
+			Queue:    QueueName,
+			Priority: PriorityEvent,
+		}
+	})
+
+	err := args.River.Queues().Add(QueueName, river.QueueConfig{MaxWorkers: 5})
+	if err != nil {
+		return fmt.Errorf("add queue: %w", err)
+	}
+
+	args.River.PeriodicJobs().AddMany([]*river.PeriodicJob{
+		river.NewPeriodicJob(
+			river.PeriodicInterval(time.Minute),
+			func() (river.JobArgs, *river.InsertOpts) {
+				return LookForWorkArgs{}, &river.InsertOpts{
+					Queue:    QueueName,
+					Priority: PriorityLFW,
+				}
+			},
+			&river.PeriodicJobOpts{RunOnStart: true},
+		),
+	})
+
+	return nil
+}