Lookout pruner: repair lease-returned/lease-expired zombie jobs

cmd/lookout/main.go

@@ -117,8 +117,12 @@ func prune(ctx *armadacontext.Context, config configuration.LookoutConfig) {
 	if config.PrunerConfig.ZombieRepairThreshold != nil {
 		zombieRepairThreshold = *config.PrunerConfig.ZombieRepairThreshold
 	}
-	log.Infof("expireAfter: %v, batchSize: %v, timeout: %v, zombieRepairThreshold: %v",
-		config.PrunerConfig.ExpireAfter, config.PrunerConfig.BatchSize, config.PrunerConfig.Timeout, zombieRepairThreshold)
+	leaseReturnedZombieRepairThreshold := 3 * time.Hour
+	if config.PrunerConfig.LeaseReturnedZombieRepairThreshold != nil {
+		leaseReturnedZombieRepairThreshold = *config.PrunerConfig.LeaseReturnedZombieRepairThreshold
+	}
+	log.Infof("expireAfter: %v, batchSize: %v, timeout: %v, zombieRepairThreshold: %v, leaseReturnedZombieRepairThreshold: %v",
+		config.PrunerConfig.ExpireAfter, config.PrunerConfig.BatchSize, config.PrunerConfig.Timeout, zombieRepairThreshold, leaseReturnedZombieRepairThreshold)
 
 	ctxTimeout, cancel := armadacontext.WithTimeout(ctx, config.PrunerConfig.Timeout)
 	defer cancel()
@@ -128,6 +132,7 @@ func prune(ctx *armadacontext.Context, config configuration.LookoutConfig) {
 		config.PrunerConfig.ExpireAfter,
 		config.PrunerConfig.DeduplicationExpireAfter,
 		zombieRepairThreshold,
+		leaseReturnedZombieRepairThreshold,
 		config.PrunerConfig.BatchSize,
 		clock.RealClock{},
 		config.ExperimentalHotColdSplit,

config/lookout/config.yaml

@@ -23,6 +23,7 @@ prunerConfig:
   expireAfter: 1008h # 42 days / 6 weeks
   deduplicationExpireAfter: 168h # 7 days
   zombieRepairThreshold: 15m # grace period before repairing jobs whose latest run is terminal but state is non-terminal
+  leaseReturnedZombieRepairThreshold: 3h # longer grace period before repairing (to failed) jobs whose latest run is lease-returned/lease-expired but state is non-terminal
   timeout: 1h
   batchSize: 1000
   pushgatewayUrl: ""

internal/lookout/configuration/types.go

@@ -51,9 +51,21 @@ type PrunerConfig struct {
 	// If nil, defaults to 15 minutes. Set to 0 explicitly to disable zombie
 	// reconciliation entirely.
 	ZombieRepairThreshold *time.Duration
-	Timeout               time.Duration
-	BatchSize             int
-	Postgres              configuration.PostgresConfig
+	// LeaseReturnedZombieRepairThreshold is the minimum age of a lease-returned
+	// or lease-expired latest run before a zombie job in that state will be
+	// repaired by the pruner (to FAILED). Unlike the other zombie-causing run
+	// states, lease-returned/lease-expired are not unconditionally terminal for
+	// the job: the scheduler is normally expected to follow up with either a
+	// requeue or a failure, and that decision can legitimately take much
+	// longer than ordinary ingester lag. This threshold should therefore
+	// typically be set substantially longer than ZombieRepairThreshold.
+	//
+	// If nil, defaults to 3 hours. Set to 0 explicitly to disable repair of
+	// lease-returned/lease-expired zombies entirely.
+	LeaseReturnedZombieRepairThreshold *time.Duration
+	Timeout                            time.Duration
+	BatchSize                          int
+	Postgres                           configuration.PostgresConfig
 	// PushgatewayUrl is the URL of a Prometheus Pushgateway (or compatible
 	// endpoint) to push pruner metrics to after each run. If empty, no metrics
 	// are pushed.

internal/lookout/pruner/doc.go

@@ -3,9 +3,15 @@
 // The pruner runs periodically and performs three tasks:
 //
 //  1. Reconciles "zombie" jobs whose state column is non-terminal but whose
-//     latest run is in a terminal state. This addresses the residue of a now-
-//     fixed ingester bug. Reconciliation is gated by a configurable grace
-//     period to avoid racing in-flight state transitions and ingester lag.
+//     latest run is in a terminal state, or in a lease-returned/lease-expired
+//     state whose expected follow-up event (requeue or failure) never
+//     arrived. This addresses the residue of a now-fixed ingester bug, and of
+//     job-level events lost by the ingester more generally. Reconciliation is
+//     gated by two independently configurable grace periods to avoid racing
+//     in-flight state transitions and ingester lag:
+//     - a short one for the unconditionally terminal run states
+//     - a longer one for lease-returned/lease-expired, which may still be
+//     legitimately retried
 //
 //  2. Deletes terminal jobs (and their associated run, spec, and error rows)
 //     that are older than a configurable lifetime, in batches.

internal/lookout/pruner/pruner.go

@@ -19,14 +19,15 @@ func PruneDb(
 	jobLifetime time.Duration,
 	deduplicationLifetime time.Duration,
 	zombieRepairThreshold time.Duration,
+	leaseReturnedZombieRepairThreshold time.Duration,
 	batchLimit int,
 	clock clock.Clock,
 	hotColdSplit bool,
 ) error {
 	var result *multierror.Error
 
-	if zombieRepairThreshold > 0 {
-		if _, err := ReconcileZombieJobs(ctx, db, zombieRepairThreshold, batchLimit, clock); err != nil {
+	if zombieRepairThreshold > 0 || leaseReturnedZombieRepairThreshold > 0 {
+		if _, err := ReconcileZombieJobs(ctx, db, zombieRepairThreshold, leaseReturnedZombieRepairThreshold, batchLimit, clock); err != nil {
 			result = multierror.Append(result, err)
 		}
 	}

internal/lookout/pruner/pruner_test.go

@@ -222,7 +222,7 @@ func TestPruneDb(t *testing.T) {
 				dbConn, err := db.Acquire(ctx)
 				assert.NoError(t, err)
 				isHC := isHotColdSchema(ctx, db)
-				err = PruneDb(ctx, dbConn.Conn(), tc.expireAfter, 0, 0, 10, clock.NewFakeClock(baseTime), isHC)
+				err = PruneDb(ctx, dbConn.Conn(), tc.expireAfter, 0, 0, 0, 10, clock.NewFakeClock(baseTime), isHC)
 				assert.NoError(t, err)
 
 				queriedJobIdsPerTable := []map[string]bool{

internal/lookout/pruner/reconcile_zombies.go

@@ -20,6 +20,31 @@ import (
 // parameters) is safe because they are compile-time constants from
 // internal/common/database/lookout, not user input. It also lets PostgreSQL
 // infer mapping.new_state as smallint without explicit casts.
+//
+// LEASE_RETURNED and LEASE_EXPIRED are mapped to job FAILED rather than left
+// unhandled. Ordinarily a lease-returned/expired run is followed by either a
+// JobRequeued or a terminal JobErrors event, so the run itself is not a
+// reliable terminal signal for the job. But if that follow-up event is lost
+// (e.g. dropped by the ingester), the job is stuck showing a non-terminal
+// state forever with no other event to correct it. FAILED is a conservative
+// choice: it may mislabel a job that was actually still being retried, but it
+// stops a permanently stuck job from being reported as active indefinitely.
+//
+// LEASE_RETURNED/LEASE_EXPIRED use a separate, longer cutoff ($2) than the
+// other four run states ($1): unlike the other four, which are unconditionally
+// terminal for the job the moment the run reaches them, a lease-returned or
+// lease-expired run is normally expected to be followed by a legitimate
+// scheduler decision (retry or fail) that can take substantially longer than
+// ordinary ingester lag to arrive.
+//
+// The LEASE_RETURNED/LEASE_EXPIRED branch additionally requires
+// job.last_transition_time <= run.finished. A successful requeue moves
+// job.state to QUEUED (advancing last_transition_time) without changing
+// latest_run_id, which still keeps pointing at the old lease-returned/expired
+// run until the next lease is granted. Without this check, a job that was
+// legitimately requeued and is simply waiting for its next lease -- rather
+// than one whose follow-up event was lost -- would be misdetected as a
+// zombie and incorrectly marked FAILED.
 var reconcileZombiesQuery = fmt.Sprintf(`
 	UPDATE job
 	SET state                        = mapping.new_state,
@@ -32,10 +57,12 @@ var reconcileZombiesQuery = fmt.Sprintf(`
 		       j.state         AS old_state,
 		       j.latest_run_id AS run_id,
 		       CASE r.job_run_state
-		           WHEN %[1]d THEN %[2]d -- run succeeded -> job succeeded
-		           WHEN %[3]d THEN %[4]d -- run failed    -> job failed
-		           WHEN %[5]d THEN %[6]d -- run cancelled -> job cancelled
-		           WHEN %[7]d THEN %[8]d -- run preempted -> job preempted
+		           WHEN %[1]d THEN %[2]d  -- run succeeded      -> job succeeded
+		           WHEN %[3]d THEN %[4]d  -- run failed         -> job failed
+		           WHEN %[5]d THEN %[6]d  -- run cancelled      -> job cancelled
+		           WHEN %[7]d THEN %[8]d  -- run preempted      -> job preempted
+		           WHEN %[13]d THEN %[4]d -- run lease returned -> job failed
+		           WHEN %[14]d THEN %[4]d -- run lease expired  -> job failed
 		           ELSE j.state          -- defensive: a job_run_state that
 		                                 -- passes the IN filter but is not
 		                                 -- listed above (e.g. after a future
@@ -47,10 +74,25 @@ var reconcileZombiesQuery = fmt.Sprintf(`
 		FROM job j
 		JOIN job_run r ON r.run_id = j.latest_run_id
 		WHERE j.state IN (%[9]d, %[10]d, %[11]d, %[12]d)
-		  AND r.job_run_state IN (%[1]d, %[3]d, %[5]d, %[7]d)
 		  AND r.finished IS NOT NULL
-		  AND r.finished < $1
-		LIMIT $2
+		  AND (
+		    (r.job_run_state IN (%[1]d, %[3]d, %[5]d, %[7]d) AND r.finished < $1)
+		    OR
+		    (
+		      r.job_run_state IN (%[13]d, %[14]d)
+		      AND r.finished < $2
+		      -- A JobRequeued event moves job.state to QUEUED and advances
+		      -- last_transition_time, but leaves latest_run_id pointing at the
+		      -- now-stale lease-returned/expired run until the next lease is
+		      -- granted. Without this check, a job that was legitimately
+		      -- requeued and is simply waiting for its next lease would be
+		      -- misdetected as a zombie. Requiring last_transition_time <=
+		      -- finished restricts the match to jobs no later event has
+		      -- touched since the lease was returned/expired.
+		      AND j.last_transition_time <= r.finished
+		    )
+		  )
+		LIMIT $3
 	) AS mapping
 	WHERE job.job_id = mapping.job_id
 	  -- Filter out ELSE-branch rows so they don't (a) get last_transition_time
@@ -66,6 +108,8 @@ var reconcileZombiesQuery = fmt.Sprintf(`
 	lookout.JobLeasedOrdinal,
 	lookout.JobPendingOrdinal,
 	lookout.JobRunningOrdinal,
+	lookout.JobRunLeaseReturnedOrdinal,
+	lookout.JobRunLeaseExpiredOrdinal,
 )
 
 // countZombiesWithNullFinishedQuery counts jobs that match the zombie shape
@@ -76,34 +120,54 @@ var countZombiesWithNullFinishedQuery = fmt.Sprintf(`
 	SELECT COUNT(*)
 	FROM job j
 	JOIN job_run r ON r.run_id = j.latest_run_id
-	WHERE j.state IN (%[5]d, %[6]d, %[7]d, %[8]d)
-	  AND r.job_run_state IN (%[1]d, %[2]d, %[3]d, %[4]d)
+	WHERE j.state IN (%[7]d, %[8]d, %[9]d, %[10]d)
+	  AND r.job_run_state IN (%[1]d, %[2]d, %[3]d, %[4]d, %[5]d, %[6]d)
 	  AND r.finished IS NULL`,
 	lookout.JobRunSucceededOrdinal,
 	lookout.JobRunFailedOrdinal,
 	lookout.JobRunCancelledOrdinal,
 	lookout.JobRunPreemptedOrdinal,
+	lookout.JobRunLeaseReturnedOrdinal,
+	lookout.JobRunLeaseExpiredOrdinal,
 	lookout.JobQueuedOrdinal,
 	lookout.JobLeasedOrdinal,
 	lookout.JobPendingOrdinal,
 	lookout.JobRunningOrdinal,
 )
 
 // ReconcileZombieJobs finds jobs whose state column is non-terminal but whose
-// latest run is in a terminal state and finished more than zombieRepairThreshold
-// ago, and updates job.state (and last_transition_time) to match the run.
-// Returns the number of jobs repaired.
+// latest run is in a terminal state, or in a lease-returned/lease-expired
+// state that never received its expected follow-up event, and updates
+// job.state (and last_transition_time) to match the run. The two cases use
+// separate grace periods: zombieRepairThreshold for the unconditionally
+// terminal run states, and leaseReturnedZombieRepairThreshold (normally much
+// longer) for lease-returned/lease-expired, since those are legitimately
+// followed by a scheduler retry-or-fail decision that can take a while to
+// arrive. A zero threshold disables reconciliation for its run-state group
+// independently of the other. Returns the number of jobs repaired.
 func ReconcileZombieJobs(
 	ctx *armadacontext.Context,
 	db *pgx.Conn,
 	zombieRepairThreshold time.Duration,
+	leaseReturnedZombieRepairThreshold time.Duration,
 	batchLimit int,
 	clock clock.Clock,
 ) (int, error) {
-	cutOffTime := clock.Now().Add(-zombieRepairThreshold)
+	// A zero threshold disables reconciliation for that run-state group. Using
+	// the zero time.Time (year 1) as the cutoff, rather than clock.Now(),
+	// ensures "r.finished < cutoff" can never match instead of matching
+	// everything already finished.
+	cutOffTime := time.Time{}
+	if zombieRepairThreshold > 0 {
+		cutOffTime = clock.Now().Add(-zombieRepairThreshold)
+	}
+	leaseReturnedCutOffTime := time.Time{}
+	if leaseReturnedZombieRepairThreshold > 0 {
+		leaseReturnedCutOffTime = clock.Now().Add(-leaseReturnedZombieRepairThreshold)
+	}
 	totalRepaired := 0
 	for {
-		batchRepaired, err := reconcileZombieBatch(ctx, db, cutOffTime, batchLimit)
+		batchRepaired, err := reconcileZombieBatch(ctx, db, cutOffTime, leaseReturnedCutOffTime, batchLimit)
 		if err != nil {
 			return totalRepaired, err
 		}
@@ -151,6 +215,7 @@ func reconcileZombieBatch(
 	ctx *armadacontext.Context,
 	db *pgx.Conn,
 	cutOffTime time.Time,
+	leaseReturnedCutOffTime time.Time,
 	batchLimit int,
 ) (int, error) {
 	type repair struct {
@@ -168,7 +233,7 @@ func reconcileZombieBatch(
 		IsoLevel:   pgx.ReadCommitted,
 		AccessMode: pgx.ReadWrite,
 	}, func(tx pgx.Tx) error {
-		rows, err := tx.Query(ctx, reconcileZombiesQuery, cutOffTime, batchLimit)
+		rows, err := tx.Query(ctx, reconcileZombiesQuery, cutOffTime, leaseReturnedCutOffTime, batchLimit)
 		if err != nil {
 			return errors.WithStack(err)
 		}