Repository: armadaproject/armada (go)

Issue:
Scheduler: Add `JobRunPreempted` and `JobRunAssigned` event handlers
Currently both `EventSequence_Event_JobRunPreempted` and `EventSequence_Event_JobRunAssigned` are being ignored in the scheduleringester: https://github.com/armadaproject/armada/blob/master/internal/scheduleringester/instructions.go#L110-L111

Since both events would be helpful in the metrics collection, we should handle both events and update the scheduler db to include them.

the changes should include:
- update the scheduler databse to add jobrun-preempted timestamp column, add assigned (or similar) boolean column and add preempted boolean column
- add the corresponding events handlers to the scheduler ingester
- update the unit tests 

Pull request:
Scheduler: handle run_preempted and run_pending events
Fixes #3174 
This should be merge after #3054 

Patch:
diff --git a/internal/scheduler/database/migrations/002_add_runs_timestamps.up.sql b/internal/scheduler/database/migrations/002_add_runs_timestamps.up.sql
index 76300df1..8be7658a 100644
--- a/internal/scheduler/database/migrations/002_add_runs_timestamps.up.sql
+++ b/internal/scheduler/database/migrations/002_add_runs_timestamps.up.sql
@@ -2,3 +2,4 @@ ALTER TABLE runs ADD COLUMN  leased_timestamp timestamptz;
 ALTER TABLE runs ADD COLUMN  pending_timestamp timestamptz;
 ALTER TABLE runs ADD COLUMN  running_timestamp timestamptz;
 ALTER TABLE runs ADD COLUMN  terminated_timestamp timestamptz;
+ALTER TABLE runs ADD COLUMN  preempted_timestamp timestamptz;
diff --git a/internal/scheduler/database/migrations/004_add_preempted_pending.up.sql b/internal/scheduler/database/migrations/004_add_preempted_pending.up.sql
new file mode 100644
index 00000000..3ba0802d
--- /dev/null
+++ b/internal/scheduler/database/migrations/004_add_preempted_pending.up.sql
@@ -0,0 +1,3 @@
+ALTER TABLE runs ADD COLUMN preempted boolean NOT NULL DEFAULT false;
+-- Set to true once the run is sent to an executor, is assigned to a node, and is waiting to start running.
+ALTER TABLE runs ADD COLUMN pending boolean NOT NULL DEFAULT false;
diff --git a/internal/scheduler/database/models.go b/internal/scheduler/database/models.go
index 5143b04e..4600c05a 100644
--- a/internal/scheduler/database/models.go
+++ b/internal/scheduler/database/models.go
@@ -74,5 +74,8 @@ type Run struct {
 	PendingTimestamp    *time.Time `db:"pending_timestamp"`
 	RunningTimestamp    *time.Time `db:"running_timestamp"`
 	TerminatedTimestamp *time.Time `db:"terminated_timestamp"`
+	PreemptedTimestamp  *time.Time `db:"preempted_timestamp"`
 	ScheduledAtPriority *int32     `db:"scheduled_at_priority"`
+	Preempted           bool       `db:"preempted"`
+	Pending             bool       `db:"pending"`
 }
diff --git a/internal/scheduler/database/query.sql.go b/internal/scheduler/database/query.sql.go
index 941e7ba7..3b5d5403 100644
--- a/internal/scheduler/database/query.sql.go
+++ b/internal/scheduler/database/query.sql.go
@@ -427,7 +427,7 @@ func (q *Queries) SelectNewJobs(ctx context.Context, arg SelectNewJobsParams) ([
 }
 
 const selectNewRuns = `-- name: SelectNewRuns :many
-SELECT run_id, job_id, created, job_set, executor, node, cancelled, running, succeeded, failed, returned, run_attempted, serial, last_modified, leased_timestamp, pending_timestamp, running_timestamp, terminated_timestamp, scheduled_at_priority FROM runs WHERE serial > $1 ORDER BY serial LIMIT $2
+SELECT run_id, job_id, created, job_set, executor, node, cancelled, running, succeeded, failed, returned, run_attempted, serial, last_modified, leased_timestamp, pending_timestamp, running_timestamp, terminated_timestamp, preempted_timestamp, scheduled_at_priority, preempted, pending FROM runs WHERE serial > $1 ORDER BY serial LIMIT $2
 `
 
 type SelectNewRunsParams struct {
@@ -463,7 +463,10 @@ func (q *Queries) SelectNewRuns(ctx context.Context, arg SelectNewRunsParams) ([
 			&i.PendingTimestamp,
 			&i.RunningTimestamp,
 			&i.TerminatedTimestamp,
+			&i.PreemptedTimestamp,
 			&i.ScheduledAtPriority,
+			&i.Preempted,
+			&i.Pending,
 		); err != nil {
 			return nil, err
 		}
@@ -476,7 +479,7 @@ func (q *Queries) SelectNewRuns(ctx context.Context, arg SelectNewRunsParams) ([
 }
 
 const selectNewRunsForJobs = `-- name: SelectNewRunsForJobs :many
-SELECT run_id, job_id, created, job_set, executor, node, cancelled, running, succeeded, failed, returned, run_attempted, serial, last_modified, leased_timestamp, pending_timestamp, running_timestamp, terminated_timestamp, scheduled_at_priority FROM runs WHERE serial > $1 AND job_id = ANY($2::text[]) ORDER BY serial
+SELECT run_id, job_id, created, job_set, executor, node, cancelled, running, succeeded, failed, returned, run_attempted, serial, last_modified, leased_timestamp, pending_timestamp, running_timestamp, terminated_timestamp, preempted_timestamp, scheduled_at_priority, preempted, pending FROM runs WHERE serial > $1 AND job_id = ANY($2::text[]) ORDER BY serial
 `
 
 type SelectNewRunsForJobsParams struct {
@@ -512,7 +515,10 @@ func (q *Queries) SelectNewRunsForJobs(ctx context.Context, arg SelectNewRunsFor
 			&i.PendingTimestamp,
 			&i.RunningTimestamp,
 			&i.TerminatedTimestamp,
+			&i.PreemptedTimestamp,
 			&i.ScheduledAtPriority,
+			&i.Preempted,
+			&i.Pending,
 		); err != nil {
 			return nil, err
 		}
diff --git a/internal/scheduleringester/dbops.go b/internal/scheduleringester/dbops.go
index 794bab98..1310b66c 100644
--- a/internal/scheduleringester/dbops.go
+++ b/internal/scheduleringester/dbops.go
@@ -131,6 +131,8 @@ type (
 	MarkRunsSucceeded          map[uuid.UUID]time.Time
 	MarkRunsFailed             map[uuid.UUID]*JobRunFailed
 	MarkRunsRunning            map[uuid.UUID]time.Time
+	MarkRunsPending            map[uuid.UUID]time.Time
+	MarkRunsPreempted          map[uuid.UUID]time.Time
 	InsertJobRunErrors         map[uuid.UUID]*schedulerdb.JobRunError
 	InsertPartitionMarker      struct {
 		markers []*schedulerdb.Marker
@@ -235,6 +237,14 @@ func (a MarkRunsRunning) Merge(b DbOperation) bool {
 	return mergeInMap(a, b)
 }
 
+func (a MarkRunsPending) Merge(b DbOperation) bool {
+	return mergeInMap(a, b)
+}
+
+func (a MarkRunsPreempted) Merge(b DbOperation) bool {
+	return mergeInMap(a, b)
+}
+
 func (a InsertJobRunErrors) Merge(b DbOperation) bool {
 	return mergeInMap(a, b)
 }
@@ -346,6 +356,14 @@ func (a MarkRunsRunning) CanBeAppliedBefore(b DbOperation) bool {
 	return !definesRun(a, b)
 }
 
+func (a MarkRunsPending) CanBeAppliedBefore(b DbOperation) bool {
+	return !definesRun(a, b)
+}
+
+func (a MarkRunsPreempted) CanBeAppliedBefore(b DbOperation) bool {
+	return !definesRun(a, b)
+}
+
 func (a *InsertPartitionMarker) CanBeAppliedBefore(b DbOperation) bool {
 	// Partition markers can never be brought forward
 	return false
diff --git a/internal/scheduleringester/instructions.go b/internal/scheduleringester/instructions.go
index 8204ea8c..f2d751c2 100644
--- a/internal/scheduleringester/instructions.go
+++ b/internal/scheduleringester/instructions.go
@@ -103,12 +103,14 @@ func (c *InstructionConverter) dbOperationsFromEventSequence(es *armadaevents.Ev
 			operationsFromEvent, err = c.handleJobRequeued(event.GetJobRequeued())
 		case *armadaevents.EventSequence_Event_PartitionMarker:
 			operationsFromEvent, err = c.handlePartitionMarker(event.GetPartitionMarker(), *event.Created)
+		case *armadaevents.EventSequence_Event_JobRunPreempted:
+			operationsFromEvent, err = c.handleJobRunPreempted(event.GetJobRunPreempted(), eventTime)
+		case *armadaevents.EventSequence_Event_JobRunAssigned:
+			operationsFromEvent, err = c.handleJobRunAssigned(event.GetJobRunAssigned(), eventTime)
 		case *armadaevents.EventSequence_Event_ReprioritisedJob,
 			*armadaevents.EventSequence_Event_JobDuplicateDetected,
 			*armadaevents.EventSequence_Event_ResourceUtilisation,
-			*armadaevents.EventSequence_Event_StandaloneIngressInfo,
-			*armadaevents.EventSequence_Event_JobRunPreempted,
-			*armadaevents.EventSequence_Event_JobRunAssigned:
+			*armadaevents.EventSequence_Event_StandaloneIngressInfo:
 			// These events can all be safely ignored
 			log.Debugf("Ignoring event type %T", event)
 		default:
@@ -238,6 +240,16 @@ func (c *InstructionConverter) handleJobRunSucceeded(jobRunSucceeded *armadaeven
 	return []DbOperation{MarkRunsSucceeded{runId: successTime}}, nil
 }
 
+func (c *InstructionConverter) handleJobRunPreempted(jobRunPreempted *armadaevents.JobRunPreempted, preemptedTime time.Time) ([]DbOperation, error) {
+	runId := armadaevents.UuidFromProtoUuid(jobRunPreempted.PreemptedRunId)
+	return []DbOperation{MarkRunsPreempted{runId: preemptedTime}}, nil
+}
+
+func (c *InstructionConverter) handleJobRunAssigned(jobRunAssigned *armadaevents.JobRunAssigned, assignedTime time.Time) ([]DbOperation, error) {
+	runId := armadaevents.UuidFromProtoUuid(jobRunAssigned.GetRunId())
+	return []DbOperation{MarkRunsPending{runId: assignedTime}}, nil
+}
+
 func (c *InstructionConverter) handleJobRunErrors(jobRunErrors *armadaevents.JobRunErrors, failureTime time.Time) ([]DbOperation, error) {
 	runId := armadaevents.UuidFromProtoUuid(jobRunErrors.GetRunId())
 	jobId, err := armadaevents.UlidStringFromProtoUuid(jobRunErrors.JobId)
diff --git a/internal/scheduleringester/schedulerdb.go b/internal/scheduleringester/schedulerdb.go
index b57db915..6d4d5854 100644
--- a/internal/scheduleringester/schedulerdb.go
+++ b/internal/scheduleringester/schedulerdb.go
@@ -276,6 +276,32 @@ func (s *SchedulerDb) WriteDbOp(ctx *armadacontext.Context, tx pgx.Tx, op DbOper
 		if _, err := tx.Exec(ctx, sqlStmt, runIds, running, runningTimes); err != nil {
 			return errors.WithStack(err)
 		}
+	case MarkRunsPending:
+		runIds := make([]uuid.UUID, 0, len(o))
+		pendingTimes := make([]interface{}, 0, len(o))
+		pending := make([]bool, 0, len(o))
+		for runId, pendingTime := range o {
+			runIds = append(runIds, runId)
+			pendingTimes = append(pendingTimes, pendingTime)
+			pending = append(pending, true)
+		}
+		sqlStmt := multiColumnRunsUpdateStmt("run_id", "pending", "pending_timestamp")
+		if _, err := tx.Exec(ctx, sqlStmt, runIds, pending, pendingTimes); err != nil {
+			return errors.WithStack(err)
+		}
+	case MarkRunsPreempted:
+		runIds := make([]uuid.UUID, 0, len(o))
+		preemptedTimes := make([]interface{}, 0, len(o))
+		preempted := make([]bool, 0, len(o))
+		for runId, preemptedTime := range o {
+			runIds = append(runIds, runId)
+			preemptedTimes = append(preemptedTimes, preemptedTime)
+			preempted = append(preempted, true)
+		}
+		sqlStmt := multiColumnRunsUpdateStmt("run_id", "preempted", "preempted_timestamp")
+		if _, err := tx.Exec(ctx, sqlStmt, runIds, preempted, preemptedTimes); err != nil {
+			return errors.WithStack(err)
+		}
 	case InsertJobRunErrors:
 		records := make([]any, len(o))
 		i := 0


Test patch:
diff --git a/internal/scheduleringester/instructions_test.go b/internal/scheduleringester/instructions_test.go
index bdb228d8..3f85d96e 100644
--- a/internal/scheduleringester/instructions_test.go
+++ b/internal/scheduleringester/instructions_test.go
@@ -76,6 +76,14 @@ func TestConvertSequence(t *testing.T) {
 			events:   []*armadaevents.EventSequence_Event{f.JobRunSucceeded},
 			expected: []DbOperation{MarkRunsSucceeded{f.RunIdUuid: f.BaseTime}},
 		},
+		"job run pending": {
+			events:   []*armadaevents.EventSequence_Event{f.Assigned},
+			expected: []DbOperation{MarkRunsPending{f.RunIdUuid: f.BaseTime}},
+		},
+		"job run preempted": {
+			events:   []*armadaevents.EventSequence_Event{f.JobPreempted},
+			expected: []DbOperation{MarkRunsPreempted{f.RunIdUuid: f.BaseTime}},
+		},
 		"lease returned": {
 			events: []*armadaevents.EventSequence_Event{f.LeaseReturned},
 			expected: []DbOperation{
@@ -202,7 +210,7 @@ func TestConvertSequence(t *testing.T) {
 			},
 		},
 		"ignored events": {
-			events: []*armadaevents.EventSequence_Event{f.Running, f.JobPreempted, f.JobSucceeded},
+			events: []*armadaevents.EventSequence_Event{f.Running, f.SubmitDuplicate, f.JobSucceeded},
 			expected: []DbOperation{
 				MarkRunsRunning{f.RunIdUuid: f.BaseTime},
 				MarkJobsSucceeded{f.JobIdString: true},
diff --git a/internal/scheduleringester/schedulerdb_test.go b/internal/scheduleringester/schedulerdb_test.go
index 7adfa98a..e3709ea3 100644
--- a/internal/scheduleringester/schedulerdb_test.go
+++ b/internal/scheduleringester/schedulerdb_test.go
@@ -157,6 +157,33 @@ func TestWriteOps(t *testing.T) {
 				jobIds[1]: true,
 			},
 		}},
+		"MarkRunsPending": {Ops: []DbOperation{
+			InsertJobs{
+				jobIds[0]: &schedulerdb.Job{JobID: jobIds[0]},
+				jobIds[1]: &schedulerdb.Job{JobID: jobIds[1]},
+			},
+			InsertRuns{
+				runIds[0]: &JobRunDetails{queue: testQueueName, dbRun: &schedulerdb.Run{JobID: jobIds[0], RunID: runIds[0]}},
+				runIds[1]: &JobRunDetails{queue: testQueueName, dbRun: &schedulerdb.Run{JobID: jobIds[1], RunID: runIds[1]}},
+			},
+			MarkRunsPending{
+				runIds[0]: testfixtures.BaseTime,
+				runIds[1]: testfixtures.BaseTime.Add(time.Hour),
+			},
+		}},
+		"MarkRunsPreempted": {Ops: []DbOperation{
+			InsertJobs{
+				jobIds[0]: &schedulerdb.Job{JobID: jobIds[0]},
+				jobIds[1]: &schedulerdb.Job{JobID: jobIds[1]},
+			},
+			InsertRuns{
+				runIds[0]: &JobRunDetails{queue: testQueueName, dbRun: &schedulerdb.Run{JobID: jobIds[0], RunID: runIds[0]}},
+				runIds[1]: &JobRunDetails{queue: testQueueName, dbRun: &schedulerdb.Run{JobID: jobIds[1], RunID: runIds[1]}},
+			},
+			MarkRunsPreempted{
+				runIds[0]: testfixtures.BaseTime,
+			},
+		}},
 		"MarkJobsFailed": {Ops: []DbOperation{
 			InsertJobs{
 				jobIds[0]: &schedulerdb.Job{JobID: jobIds[0], JobSet: "set1"},
@@ -603,6 +630,58 @@ func assertOpSuccess(t *testing.T, schedulerDb *SchedulerDb, serials map[string]
 			}
 		}
 		assert.Equal(t, len(expected), len(runs))
+	case MarkRunsPending:
+		jobs, err := selectNewJobs(ctx, 0)
+		if err != nil {
+			return errors.WithStack(err)
+		}
+		jobIds := make([]string, 0)
+		for _, job := range jobs {
+			jobIds = append(jobIds, job.JobID)
+		}
+		runs, err := queries.SelectNewRunsForJobs(ctx, schedulerdb.SelectNewRunsForJobsParams{
+			Serial: serials["runs"],
+			JobIds: jobIds,
+		})
+		if err != nil {
+			return errors.WithStack(err)
+		}
+		numChanged := 0
+		for _, run := range runs {
+			if _, ok := expected[run.RunID]; ok {
+				assert.True(t, run.Pending)
+				assert.Equal(t, expected[run.RunID], run.PendingTimestamp.UTC())
+				numChanged++
+			}
+		}
+		assert.Equal(t, numChanged, 2)
+		assert.Equal(t, len(expected), len(runs))
+	case MarkRunsPreempted:
+		jobs, err := selectNewJobs(ctx, 0)
+		if err != nil {
+			return errors.WithStack(err)
+		}
+		jobIds := make([]string, 0)
+		for _, job := range jobs {
+			jobIds = append(jobIds, job.JobID)
+		}
+		runs, err := queries.SelectNewRunsForJobs(ctx, schedulerdb.SelectNewRunsForJobsParams{
+			Serial: serials["runs"],
+			JobIds: jobIds,
+		})
+		if err != nil {
+			return errors.WithStack(err)
+		}
+		numChanged := 0
+		for _, run := range runs {
+			if _, ok := expected[run.RunID]; ok {
+				assert.True(t, run.Preempted)
+				assert.Equal(t, expected[run.RunID], run.PreemptedTimestamp.UTC())
+				numChanged++
+			}
+		}
+		assert.Equal(t, numChanged, 1)
+		assert.Equal(t, len(expected), len(runs))
 	case InsertJobRunErrors:
 		expectedIds := maps.Keys(expected)
 		as, err := queries.SelectRunErrorsById(ctx, expectedIds)