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,22 @@ 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.
 var reconcileZombiesQuery = fmt.Sprintf(`
 	UPDATE job
 	SET state                        = mapping.new_state,
@@ -32,10 +48,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 +65,13 @@ 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)
+		  )
+		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 +87,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 +99,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 +194,7 @@ func reconcileZombieBatch(
 	ctx *armadacontext.Context,
 	db *pgx.Conn,
 	cutOffTime time.Time,
+	leaseReturnedCutOffTime time.Time,
 	batchLimit int,
 ) (int, error) {
 	type repair struct {
@@ -168,7 +212,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)
 		}

internal/lookout/pruner/reconcile_zombies_test.go

@@ -156,7 +156,7 @@ func TestReconcileZombieJobs(t *testing.T) {
 				dbConn, err := db.Acquire(ctx)
 				require.NoError(t, err)
 
-				repaired, err := ReconcileZombieJobs(ctx, dbConn.Conn(), tc.zombieRepairThreshold, tc.batchSize, clock.NewFakeClock(baseTime))
+				repaired, err := ReconcileZombieJobs(ctx, dbConn.Conn(), tc.zombieRepairThreshold, tc.zombieRepairThreshold, tc.batchSize, clock.NewFakeClock(baseTime))
 				require.NoError(t, err)
 
 				expectedRepairs := 0
@@ -214,7 +214,7 @@ func TestReconcileZombieJobsLeavesNonZombiesAlone(t *testing.T) {
 		dbConn, err := db.Acquire(ctx)
 		require.NoError(t, err)
 
-		repaired, err := ReconcileZombieJobs(ctx, dbConn.Conn(), 1*time.Hour, 10, clock.NewFakeClock(baseTime))
+		repaired, err := ReconcileZombieJobs(ctx, dbConn.Conn(), 1*time.Hour, 1*time.Hour, 10, clock.NewFakeClock(baseTime))
 		require.NoError(t, err)
 		assert.Equal(t, 0, repaired)
 
@@ -226,6 +226,92 @@ func TestReconcileZombieJobsLeavesNonZombiesAlone(t *testing.T) {
 	assert.NoError(t, err)
 }
 
+// TestReconcileZombieJobsRepairsLeaseReturnedAndExpired verifies that jobs
+// left stuck in a non-terminal state whose latest run is lease-returned or
+// lease-expired -- because the job-level JobRequeued/JobErrors event that
+// should have followed was lost -- are conservatively repaired to FAILED.
+func TestReconcileZombieJobsRepairsLeaseReturnedAndExpired(t *testing.T) {
+	err := lookout.WithLookoutDb(func(db *pgxpool.Pool) error {
+		converter := instructions.NewInstructionConverter(metrics.Get().Metrics, "armadaproject.io/", []string{}, &compress.NoOpCompressor{})
+		store := lookoutdb.NewLookoutDb(db, nil, metrics.Get(), 10, 10)
+
+		ctx, cancel := armadacontext.WithTimeout(armadacontext.Background(), 5*time.Minute)
+		defer cancel()
+
+		leaseReturnedJobId := util.NewULID()
+		leaseReturnedRunId := uuid.NewString()
+		repository.NewJobSimulator(converter, store).
+			Submit("queue", "jobSet", "owner", "namespace", baseTime.Add(-3*time.Hour), &repository.JobOptions{JobId: leaseReturnedJobId}).
+			Lease(leaseReturnedRunId, "cluster", "node", "pool", baseTime.Add(-3*time.Hour)).
+			Pending(leaseReturnedRunId, "cluster", baseTime.Add(-3*time.Hour)).
+			LeaseReturned(leaseReturnedRunId, "lease returned", baseTime.Add(-2*time.Hour)).
+			Build()
+
+		leaseExpiredJobId := util.NewULID()
+		leaseExpiredRunId := uuid.NewString()
+		repository.NewJobSimulator(converter, store).
+			Submit("queue", "jobSet", "owner", "namespace", baseTime.Add(-3*time.Hour), &repository.JobOptions{JobId: leaseExpiredJobId}).
+			Lease(leaseExpiredRunId, "cluster", "node", "pool", baseTime.Add(-3*time.Hour)).
+			Pending(leaseExpiredRunId, "cluster", baseTime.Add(-3*time.Hour)).
+			LeaseExpired(leaseExpiredRunId, baseTime.Add(-2*time.Hour), clock.NewFakeClock(baseTime)).
+			Build()
+
+		dbConn, err := db.Acquire(ctx)
+		require.NoError(t, err)
+
+		// zombieRepairThreshold is set to a value too long to have repaired
+		// these on its own (were it wrongly applied to this run-state group),
+		// proving leaseReturnedZombieRepairThreshold is what governs here.
+		repaired, err := ReconcileZombieJobs(ctx, dbConn.Conn(), 24*time.Hour, 1*time.Hour, 10, clock.NewFakeClock(baseTime))
+		require.NoError(t, err)
+		assert.Equal(t, 2, repaired)
+
+		assert.Equal(t, lookout.JobFailed, readJobState(t, ctx, db, leaseReturnedJobId))
+		assert.Equal(t, lookout.JobFailed, readJobState(t, ctx, db, leaseExpiredJobId))
+
+		return nil
+	})
+	assert.NoError(t, err)
+}
+
+// TestReconcileZombieJobsLeasesReturnedWithinGracePeriodAreLeftAlone verifies
+// that a lease-returned/expired run does not get repaired while still within
+// the grace period, since the job may simply be about to be re-leased.
+func TestReconcileZombieJobsLeasesReturnedWithinGracePeriodAreLeftAlone(t *testing.T) {
+	err := lookout.WithLookoutDb(func(db *pgxpool.Pool) error {
+		converter := instructions.NewInstructionConverter(metrics.Get().Metrics, "armadaproject.io/", []string{}, &compress.NoOpCompressor{})
+		store := lookoutdb.NewLookoutDb(db, nil, metrics.Get(), 10, 10)
+
+		ctx, cancel := armadacontext.WithTimeout(armadacontext.Background(), 5*time.Minute)
+		defer cancel()
+
+		jobId := util.NewULID()
+		runId := uuid.NewString()
+		repository.NewJobSimulator(converter, store).
+			Submit("queue", "jobSet", "owner", "namespace", baseTime.Add(-1*time.Hour), &repository.JobOptions{JobId: jobId}).
+			Lease(runId, "cluster", "node", "pool", baseTime.Add(-1*time.Hour)).
+			Pending(runId, "cluster", baseTime.Add(-1*time.Hour)).
+			LeaseReturned(runId, "lease returned", baseTime.Add(-30*time.Minute)).
+			Build()
+
+		dbConn, err := db.Acquire(ctx)
+		require.NoError(t, err)
+
+		// zombieRepairThreshold is set far shorter than the run's 30-minute
+		// age -- if it were wrongly applied to this run-state group, the job
+		// would get repaired. leaseReturnedZombieRepairThreshold (1 hour) is
+		// what should govern here, and correctly leaves the job alone.
+		repaired, err := ReconcileZombieJobs(ctx, dbConn.Conn(), 1*time.Minute, 1*time.Hour, 10, clock.NewFakeClock(baseTime))
+		require.NoError(t, err)
+		assert.Equal(t, 0, repaired)
+
+		assert.Equal(t, lookout.JobPending, readJobState(t, ctx, db, jobId))
+
+		return nil
+	})
+	assert.NoError(t, err)
+}
+
 // TestReconcileZombieJobsCountsNullFinishedZombies verifies that zombie jobs
 // whose latest run has no finished timestamp are observed via the
 // zombiesSkippedNullFinished metric (and are NOT silently repaired with bogus
@@ -255,7 +341,7 @@ func TestReconcileZombieJobsCountsNullFinishedZombies(t *testing.T) {
 		dbConn, err := db.Acquire(ctx)
 		require.NoError(t, err)
 
-		repaired, err := ReconcileZombieJobs(ctx, dbConn.Conn(), 1*time.Hour, 10, clock.NewFakeClock(baseTime))
+		repaired, err := ReconcileZombieJobs(ctx, dbConn.Conn(), 1*time.Hour, 1*time.Hour, 10, clock.NewFakeClock(baseTime))
 		require.NoError(t, err)
 
 		// Not repaired: the reconciler refuses to touch a row with NULL finished.
@@ -294,7 +380,7 @@ func TestPruneDbRunsZombieReconciliation(t *testing.T) {
 		require.NoError(t, err)
 
 		isHC := isHotColdSchema(ctx, db)
-		err = PruneDb(ctx, dbConn.Conn(), 100*time.Hour, 100*time.Hour, 1*time.Hour, 10, clock.NewFakeClock(baseTime), isHC)
+		err = PruneDb(ctx, dbConn.Conn(), 100*time.Hour, 100*time.Hour, 1*time.Hour, 1*time.Hour, 10, clock.NewFakeClock(baseTime), isHC)
 		require.NoError(t, err)
 
 		assert.Equal(t, lookout.JobSucceeded, readJobState(t, ctx, db, zombie.jobId))