Provisioning: Reclaim executing jobs (#109228)

* Reclaim Executing Jobs

* Remove FIXME

* Format code

* Fetch job before update

* Add claims for all namespaces

* Remove unused import

* Update pkg/registry/apis/provisioning/jobs/concurrent_driver.go

* Update pkg/registry/apis/provisioning/jobs/concurrent_driver.go

---------

Co-authored-by: Stephanie Hingtgen <stephanie.hingtgen@grafana.com>

pkg/registry/apis/provisioning/jobs/concurrent_driver.go

@@ -11,20 +11,21 @@ import (
 
 // ConcurrentJobDriver manages multiple jobDriver instances for concurrent job processing.
 type ConcurrentJobDriver struct {
-	numDrivers      int
-	jobTimeout      time.Duration
-	cleanupInterval time.Duration
-	jobInterval     time.Duration
-	store           Store
-	repoGetter      RepoGetter
-	historicJobs    History
-	workers         []Worker
+	numDrivers           int
+	jobTimeout           time.Duration
+	cleanupInterval      time.Duration
+	jobInterval          time.Duration
+	leaseRenewalInterval time.Duration
+	store                Store
+	repoGetter           RepoGetter
+	historicJobs         History
+	workers              []Worker
 }
 
 // NewConcurrentJobDriver creates a new concurrent job driver that spawns multiple job drivers.
 func NewConcurrentJobDriver(
 	numDrivers int,
-	jobTimeout, cleanupInterval, jobInterval time.Duration,
+	jobTimeout, cleanupInterval, jobInterval, leaseRenewalInterval time.Duration,
 	store Store,
 	repoGetter RepoGetter,
 	historicJobs History,
@@ -33,19 +34,35 @@ func NewConcurrentJobDriver(
 	if numDrivers <= 0 {
 		return nil, fmt.Errorf("numWorkers must be greater than 0, got %d", numDrivers)
 	}
-	if cleanupInterval < jobTimeout {
-		return nil, fmt.Errorf("the cleanup interval must be larger than the jobTimeout (cleanup:%s < job:%s)",
-			cleanupInterval.String(), jobTimeout.String())
+	// Default lease renewal interval to 1/3 of job timeout, minimum 5 seconds
+	if leaseRenewalInterval <= 0 {
+		leaseRenewalInterval = jobTimeout / 3
 	}
+	if leaseRenewalInterval < 5*time.Second {
+		leaseRenewalInterval = 5 * time.Second
+	}
+	// For lease-based cleanup, run at most every 3-4 lease renewal intervals
+	// to detect expired leases promptly but not too aggressively
+	if cleanupInterval <= 0 {
+		cleanupInterval = leaseRenewalInterval * 3
+	}
+	if cleanupInterval < 30*time.Second {
+		cleanupInterval = 30 * time.Second // Minimum cleanup interval
+	}
+	if cleanupInterval > 5*time.Minute {
+		cleanupInterval = 5 * time.Minute // Maximum cleanup interval
+	}
+
 	return &ConcurrentJobDriver{
-		numDrivers:      numDrivers,
-		jobTimeout:      jobTimeout,
-		cleanupInterval: cleanupInterval,
-		jobInterval:     jobInterval,
-		store:           store,
-		repoGetter:      repoGetter,
-		historicJobs:    historicJobs,
-		workers:         workers,
+		numDrivers:           numDrivers,
+		jobTimeout:           jobTimeout,
+		cleanupInterval:      cleanupInterval,
+		jobInterval:          jobInterval,
+		leaseRenewalInterval: leaseRenewalInterval,
+		store:                store,
+		repoGetter:           repoGetter,
+		historicJobs:         historicJobs,
+		workers:              workers,
 	}, nil
 }
 
@@ -53,9 +70,9 @@ func NewConcurrentJobDriver(
 // This is a blocking function that will run until the context is canceled or an error occurs.
 func (c *ConcurrentJobDriver) Run(ctx context.Context) error {
 	logger := logging.FromContext(ctx).With("logger", "concurrent-job-driver", "num_drivers", c.numDrivers)
-	logger.Info("starting concurrent job driver")
+	logger.Info("starting concurrent job driver with lease-based cleanup", "cleanup_interval", c.cleanupInterval)
 
-	// Set up cleanup ticker - only one cleanup process for all workers
+	// Set up cleanup ticker - runs more frequently with lease-based approach
 	cleanupTicker := time.NewTicker(c.cleanupInterval)
 	defer cleanupTicker.Stop()
 
@@ -96,6 +113,7 @@ func (c *ConcurrentJobDriver) Run(ctx context.Context) error {
 			driver, err := NewJobDriver(
 				c.jobTimeout,
 				c.jobInterval,
+				c.leaseRenewalInterval,
 				c.store,
 				c.repoGetter,
 				c.historicJobs,

pkg/registry/apis/provisioning/jobs/driver.go

@@ -3,8 +3,10 @@ package jobs
 import (
 	"context"
 	"errors"
+	"strings"
 	"time"
 
+	apierrors "k8s.io/apimachinery/pkg/api/errors"
 	"k8s.io/apiserver/pkg/endpoints/request"
 
 	"github.com/grafana/grafana-app-sdk/logging"
@@ -40,6 +42,13 @@ type Store interface {
 
 	// Update saves the job back to the store.
 	Update(ctx context.Context, job *provisioning.Job) (*provisioning.Job, error)
+
+	// RenewLease renews the lease for a claimed job, extending its expiry time.
+	// Returns an error if the lease cannot be renewed (e.g., job was completed or lease expired).
+	RenewLease(ctx context.Context, job *provisioning.Job) error
+
+	// Get retrieves a job by name for conflict resolution.
+	Get(ctx context.Context, name string) (*provisioning.Job, error)
 }
 
 var _ Store = (*persistentStore)(nil)
@@ -54,6 +63,9 @@ type jobDriver struct {
 	// JobInterval is the time between job ticks. This should be relatively low.
 	jobInterval time.Duration
 
+	// LeaseRenewalInterval is how often to renew job leases.
+	leaseRenewalInterval time.Duration
+
 	// Store is the job storage backend.
 	store Store
 	// RepoGetter lets us access repositories to pass to the worker.
@@ -68,19 +80,20 @@ type jobDriver struct {
 }
 
 func NewJobDriver(
-	jobTimeout, jobInterval time.Duration,
+	jobTimeout, jobInterval, leaseRenewalInterval time.Duration,
 	store Store,
 	repoGetter RepoGetter,
 	historicJobs History,
 	workers ...Worker,
 ) (*jobDriver, error) {
 	return &jobDriver{
-		jobTimeout:   jobTimeout,
-		jobInterval:  jobInterval,
-		store:        store,
-		repoGetter:   repoGetter,
-		historicJobs: historicJobs,
-		workers:      workers,
+		jobTimeout:           jobTimeout,
+		jobInterval:          jobInterval,
+		leaseRenewalInterval: leaseRenewalInterval,
+		store:                store,
+		repoGetter:           repoGetter,
+		historicJobs:         historicJobs,
+		workers:              workers,
 	}, nil
 }
 
@@ -153,12 +166,19 @@ func (d *jobDriver) claimAndProcessOneJob(ctx context.Context) error {
 	jobctx, cancel := context.WithTimeout(ctx, d.jobTimeout)
 	defer cancel() // Ensure resources are released when the function returns
 
+	// Set up lease renewal goroutine
+	leaseRenewalCtx, cancelLeaseRenewal := context.WithCancel(jobctx)
+	leaseExpired := make(chan struct{})
+
+	go d.leaseRenewalLoop(leaseRenewalCtx, job, logger, leaseExpired)
+	defer cancelLeaseRenewal()
+
 	recorder := newJobProgressRecorder(d.onProgress(job))
 
-	// Process the job.
+	// Process the job with lease loss detection
 	start := time.Now()
 	job.Status.Started = start.UnixMilli()
-	err = d.processJob(jobctx, job, recorder) // NOTE: We pass in a pointer here such that the job status can be kept in Complete without re-fetching.
+	err = d.processJobWithLeaseCheck(jobctx, job, recorder, leaseExpired)
 	end := time.Now()
 	logger.Debug("job processed", "duration", end.Sub(start), "error", err)
 
@@ -187,6 +207,70 @@ func (d *jobDriver) claimAndProcessOneJob(ctx context.Context) error {
 	return nil
 }
 
+// leaseRenewalLoop continuously renews the lease for a job until the context is cancelled.
+// If lease renewal fails persistently, it signals via the leaseExpired channel.
+func (d *jobDriver) leaseRenewalLoop(ctx context.Context, job *provisioning.Job, logger logging.Logger, leaseExpired chan struct{}) {
+	ticker := time.NewTicker(d.leaseRenewalInterval)
+	defer ticker.Stop()
+
+	logger.Debug("starting lease renewal loop", "renewal_interval", d.leaseRenewalInterval)
+
+	consecutiveFailures := 0
+	maxFailures := 3 // Allow a few failures before giving up
+
+	for {
+		select {
+		case <-ctx.Done():
+			logger.Debug("lease renewal loop stopping")
+			return
+		case <-ticker.C:
+			err := d.store.RenewLease(ctx, job)
+			if err != nil {
+				consecutiveFailures++
+				if apierrors.IsNotFound(err) ||
+					strings.Contains(err.Error(), "job no longer exists") {
+					logger.Error("job no longer exists - lease expired", "error", err)
+					close(leaseExpired)
+					return
+				}
+
+				logger.Warn("failed to renew lease", "error", err, "consecutive_failures", consecutiveFailures)
+
+				if consecutiveFailures >= maxFailures {
+					logger.Error("too many consecutive lease renewal failures - job will be aborted",
+						"consecutive_failures", consecutiveFailures, "max_failures", maxFailures)
+					close(leaseExpired)
+					return
+				}
+			} else {
+				if consecutiveFailures > 0 {
+					logger.Debug("lease renewal recovered", "previous_failures", consecutiveFailures)
+				}
+				consecutiveFailures = 0
+				logger.Debug("lease renewed successfully")
+			}
+		}
+	}
+}
+
+// processJobWithLeaseCheck processes a job but aborts if the lease expires.
+func (d *jobDriver) processJobWithLeaseCheck(ctx context.Context, job *provisioning.Job, recorder JobProgressRecorder, leaseExpired <-chan struct{}) error {
+	// Run the job processing in a goroutine so we can monitor lease expiry
+	resultChan := make(chan error, 1)
+	go func() {
+		resultChan <- d.processJob(ctx, job, recorder)
+	}()
+
+	select {
+	case err := <-resultChan:
+		return err
+	case <-leaseExpired:
+		return apifmt.Errorf("job aborted due to lease expiry")
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+}
+
 func (d *jobDriver) processJob(ctx context.Context, job *provisioning.Job, recorder JobProgressRecorder) error {
 	for _, worker := range d.workers {
 		if !worker.IsSupported(ctx, *job) {
@@ -207,14 +291,42 @@ func (d *jobDriver) processJob(ctx context.Context, job *provisioning.Job, recor
 func (d *jobDriver) onProgress(job *provisioning.Job) ProgressFn {
 	return func(ctx context.Context, status provisioning.JobStatus) error {
 		logging.FromContext(ctx).Debug("job progress", "status", status)
-		job.Status = status
 
-		updated, err := d.store.Update(ctx, job)
-		if err != nil {
-			return apifmt.Errorf("failed to update job: %w", err)
+		const maxRetries = 3
+		for attempt := 0; attempt < maxRetries; attempt++ {
+			// Use the current job for the first attempt, fetch fresh for retries
+			currentJob := job
+			if attempt > 0 {
+				// Fetch the latest version to resolve conflicts
+				latest, err := d.store.Get(ctx, job.GetName())
+				if err != nil {
+					if apierrors.IsNotFound(err) {
+						// Job was completed/deleted, nothing to update
+						return nil
+					}
+					return apifmt.Errorf("failed to fetch job for progress update: %w", err)
+				}
+				currentJob = latest
+			}
+
+			// Update status on the current job
+			currentJob.Status = status
+
+			updated, err := d.store.Update(ctx, currentJob)
+			if err != nil {
+				if apierrors.IsConflict(err) && attempt < maxRetries-1 {
+					// Conflict detected, retry with fresh data
+					logging.FromContext(ctx).Debug("progress update conflict, retrying", "attempt", attempt+1)
+					continue
+				}
+				return apifmt.Errorf("failed to update job progress: %w", err)
+			}
+
+			// Update succeeded, update our local copy
+			*job = *updated
+			return nil
 		}
 
-		*job = *updated
-		return nil
+		return apifmt.Errorf("failed to update job progress after %d attempts", maxRetries)
 	}
 }

pkg/registry/apis/provisioning/jobs/persistentstore.go

@@ -253,6 +253,24 @@ func (s *persistentStore) Update(ctx context.Context, job *provisioning.Job) (*p
 	return updatedJob, nil
 }
 
+// Get retrieves a job by name for conflict resolution.
+func (s *persistentStore) Get(ctx context.Context, name string) (*provisioning.Job, error) {
+	obj, err := s.jobStore.Get(ctx, name, &metav1.GetOptions{})
+	if err != nil {
+		if apierrors.IsNotFound(err) {
+			return nil, apifmt.Errorf("job '%s' not found", name)
+		}
+		return nil, apifmt.Errorf("failed to get job '%s': %w", name, err)
+	}
+
+	job, ok := obj.(*provisioning.Job)
+	if !ok {
+		return nil, apifmt.Errorf("unexpected object type %T", obj)
+	}
+
+	return job, nil
+}
+
 // Complete marks a job as completed and moves it to the historic job store.
 // When in the historic store, there is no more claim on the job.
 func (s *persistentStore) Complete(ctx context.Context, job *provisioning.Job) error {
@@ -279,24 +297,71 @@ func (s *persistentStore) Complete(ctx context.Context, job *provisioning.Job) e
 	return nil
 }
 
-// Cleanup should be called periodically to clean up abandoned jobs.
-// An abandoned job is one that has been claimed by a worker, but the worker has not updated the job in a while.
-func (s *persistentStore) Cleanup(ctx context.Context) error {
-	if err := s.cleanupClaims(ctx); err != nil {
-		return apifmt.Errorf("failed to clean up claims: %w", err)
+// RenewLease renews the lease for a claimed job, extending its expiry time.
+// Returns an error if the lease cannot be renewed (e.g., job was completed or lease expired).
+func (s *persistentStore) RenewLease(ctx context.Context, job *provisioning.Job) error {
+	if job.Labels == nil || job.Labels[LabelJobClaim] == "" {
+		return apifmt.Errorf("job '%s' in '%s' is not claimed", job.GetName(), job.GetNamespace())
 	}
 
+	// Fetch the latest version to avoid conflicts
+	latestObj, err := s.jobStore.Get(ctx, job.GetName(), &metav1.GetOptions{})
+	if err != nil {
+		if apierrors.IsNotFound(err) {
+			return apifmt.Errorf("failed to renew lease for job '%s' in '%s': job no longer exists", job.GetName(), job.GetNamespace())
+		}
+		return apifmt.Errorf("failed to fetch job for lease renewal '%s' in '%s': %w", job.GetName(), job.GetNamespace(), err)
+	}
+
+	latestJob, ok := latestObj.(*provisioning.Job)
+	if !ok {
+		return apifmt.Errorf("unexpected object type %T", latestObj)
+	}
+
+	// Verify we still own the lease
+	if latestJob.Labels == nil || latestJob.Labels[LabelJobClaim] == "" {
+		return apifmt.Errorf("lease lost for job '%s' in '%s': no longer claimed", job.GetName(), job.GetNamespace())
+	}
+
+	// Update the claim timestamp to current time
+	updatedJob := latestJob.DeepCopy()
+	updatedJob.Labels[LabelJobClaim] = strconv.FormatInt(s.clock().UnixMilli(), 10)
+
+	// Update the job in storage with the latest resource version
+	_, _, err = s.jobStore.Update(ctx,
+		updatedJob.GetName(),                      // name
+		rest.DefaultUpdatedObjectInfo(updatedJob), // objInfo
+		failCreation,                              // createValidation
+		nil,                                       // updateValidation
+		false,                                     // forceAllowCreate
+		&metav1.UpdateOptions{},                   // options
+	)
+	if apierrors.IsConflict(err) {
+		return apifmt.Errorf("failed to renew lease for job '%s' in '%s': lease conflict", job.GetName(), job.GetNamespace())
+	}
+	if apierrors.IsNotFound(err) || errors.Is(err, errWouldCreate) {
+		return apifmt.Errorf("failed to renew lease for job '%s' in '%s': job no longer exists", job.GetName(), job.GetNamespace())
+	}
+	if err != nil {
+		return apifmt.Errorf("failed to renew lease for job '%s' in '%s': %w", job.GetName(), job.GetNamespace(), err)
+	}
+
+	// Update the job's claim timestamp and resource version in memory
+	job.Labels[LabelJobClaim] = updatedJob.Labels[LabelJobClaim]
+	job.ResourceVersion = updatedJob.ResourceVersion
 	return nil
 }
 
-// cleanupClaims will clean up abandoned claims.
-// Any claim that is older than the expiry time will have their claims removed.
-//
-// This is only necessary because Kubernetes does not support logical OR in label selectors.
-func (s *persistentStore) cleanupClaims(ctx context.Context) error {
-	// We will list all jobs that have been claimed but not updated in a while.
-	// We will then remove the claim from them.
-	// We will not care about the result of the update, as the job may have been completed in the meantime.
+// Cleanup finds jobs with expired leases and marks them as failed.
+// This replaces the old cleanup mechanism and should be called more frequently.
+func (s *persistentStore) Cleanup(ctx context.Context) error {
+	// Set up provisioning identity to access jobs across all namespaces
+	ctx, _, err := identity.WithProvisioningIdentity(ctx, "*") // "*" grants access to all namespaces
+	if err != nil {
+		return apifmt.Errorf("failed to grant provisioning identity for cleanup: %w", err)
+	}
+
+	// Find jobs with expired leases (older than expiry time)
 	expiry := s.clock().Add(-s.expiry).UnixMilli()
 	requirement, err := labels.NewRequirement(LabelJobClaim, selection.LessThan, []string{strconv.FormatInt(expiry, 10)})
 	if err != nil {
@@ -306,49 +371,42 @@ func (s *persistentStore) cleanupClaims(ctx context.Context) error {
 	timeoutCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
 	jobsObj, err := s.jobStore.List(timeoutCtx, &internalversion.ListOptions{
 		LabelSelector: labels.NewSelector().Add(*requirement),
-		// We don't need to clean up everything all the time. Just do enough such that we have a fair amount of work.
-		Limit: 100,
+		Limit:         100, // Process in batches
 	})
-	cancel() // by the time we have the list, there is no response body to read, so just cancel immediately
+	cancel()
 	if err != nil {
-		return apifmt.Errorf("failed to list jobs: %w", err)
+		return apifmt.Errorf("failed to list jobs with expired leases: %w", err)
 	}
 	jobs, ok := jobsObj.(*provisioning.JobList)
 	if !ok {
 		return apifmt.Errorf("unexpected object type %T", jobsObj)
 	}
 
-	for _, job := range jobs.Items {
-		if job.Labels == nil {
-			job.Labels = make(map[string]string)
-		}
-		delete(job.Labels, LabelJobClaim)
-		job.Status.State = provisioning.JobStatePending
+	// If no jobs found, cleanup is complete
+	if len(jobs.Items) == 0 {
+		return nil
+	}
 
-		// We list jobs from all namespaces. So when we want to update a specific job, we also need its namespace in the context.
+	for _, job := range jobs.Items {
+		// Mark job as failed due to lease expiry and archive it
+		job := job.DeepCopy()
+		job.Status.State = provisioning.JobStateError
+		job.Status.Message = "Job failed due to lease expiry - worker may have crashed or lost connection"
+
+		// Set namespace context for the completion
 		ctx := request.WithNamespace(ctx, job.GetNamespace())
-		// Likewise, we should use the provisioning identity now that we have the namespace we are operating within.
 		ctx, _, err = identity.WithProvisioningIdentity(ctx, job.GetNamespace())
 		if err != nil {
-			// This should never happen, as it is already a valid namespace from the job existing... but better be safe.
 			return apifmt.Errorf("failed to get provisioning identity for '%s': %w", job.GetNamespace(), err)
 		}
 
-		timeoutCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
-		_, _, err := s.jobStore.Update(timeoutCtx,
-			job.GetName(),                       // name
-			rest.DefaultUpdatedObjectInfo(&job), // objInfo
-			failCreation,                        // createValidation
-			nil,                                 // updateValidation
-			false,                               // forceAllowCreate
-			&metav1.UpdateOptions{},             // options
-		)
-		cancel() // we have no response body to read, so just cancel immediately
-		if apierrors.IsConflict(err) || errors.Is(err, errWouldCreate) {
-			continue
-		}
-		if err != nil {
-			return apifmt.Errorf("failed to unclaim job '%s' in '%s': %w", job.GetName(), job.GetNamespace(), err)
+		// Use Complete to properly archive the failed job
+		if err := s.Complete(ctx, job); err != nil {
+			if apierrors.IsNotFound(err) {
+				// Job was already completed/deleted by another process
+				continue
+			}
+			return apifmt.Errorf("failed to complete expired job '%s' in '%s': %w", job.GetName(), job.GetNamespace(), err)
 		}
 	}
 

pkg/registry/apis/provisioning/register.go

@@ -637,8 +637,9 @@ func (b *APIBuilder) GetPostStartHooks() (map[string]genericapiserver.PostStartH
 			driver, err := jobs.NewConcurrentJobDriver(
 				3,              // 3 drivers for now
 				20*time.Minute, // Max time for each job
-				22*time.Minute, // Cleanup any checked out jobs. FIXME: this is slow if things crash/fail!
+				time.Minute,    // Cleanup jobs
 				30*time.Second, // Periodically look for new jobs
+				30*time.Second, // Lease renewal interval
 				b.jobs, b, b.jobHistory,
 				workers...,
 			)