EvoBGP/internal/jobs/durable.go

package jobs

import (
	"context"
	"errors"
	"evobgp/internal/logging"
	"evobgp/internal/observability"
	"evobgp/internal/store"
	"sync/atomic"
	"time"
)

const durableLease = time.Minute

// StartDurableQueueLoop claims work only after reserving local capacity.
func StartDurableQueueLoop(ctx context.Context, r *Registry) {
	if r == nil || r.queue == nil {
		return
	}
	log := logging.With("component", "jobs.durable")
	r.durableMu.Lock()
	if r.durableStop {
		r.durableMu.Unlock()
		return
	}
	r.durableLoopWG.Add(2)
	r.durableMu.Unlock()
	go r.reapExpiredJobsLoop(ctx)
	go func() {
		defer r.durableLoopWG.Done()
		ticker := time.NewTicker(time.Second)
		defer ticker.Stop()
		for {
			if ctx.Err() != nil {
				return
			}
			r.durableMu.Lock()
			stopped := r.durableStop
			r.durableMu.Unlock()
			if stopped {
				return
			}
			metricCtx, metricCancel := context.WithTimeout(ctx, 5*time.Second)
			if counts, err := r.queue.CountJobsByStatus(metricCtx); err == nil {
				observability.RecordJobAuditDepth(counts)
			}
			metricCancel()
			free := 0
		reserve:
			for free < cap(r.workerSem) {
				select {
				case r.workerSem <- struct{}{}:
					free++
				default:
					break reserve
				}
			}
			if free > 0 {
				qctx, cancel := context.WithTimeout(ctx, 10*time.Second)
				claimed, err := r.queue.ClaimJobs(qctx, r.ownerID, free, durableLease)
				cancel()
				if err != nil {
					log.Warn("claim failed", "err", err)
				}
				for i := len(claimed); i < free; i++ {
					<-r.workerSem
				}
				observability.RecordJobQueueDepth(len(r.workerSem), cap(r.workerSem))
				for _, c := range claimed {
					r.durableMu.Lock()
					stopped = r.durableStop
					if stopped || ctx.Err() != nil {
						r.durableMu.Unlock()
						releaseCtx, releaseCancel := context.WithTimeout(context.Background(), 10*time.Second)
						if err := r.queue.ReleaseJob(releaseCtx, c.ID, c.OwnerID, c.FencingToken); err != nil {
							log.Warn("release during shutdown failed", "job_id", c.ID, "err", err)
						}
						releaseCancel()
						<-r.workerSem
						continue
					}
					j := jobFromRecord(c)
					r.mu.Lock()
					r.byID[j.ID] = j
					if isRefreshKind(j.Kind) {
						r.inflightRefresh[j.TenantID]++
					}
					r.mu.Unlock()
					r.durableWG.Add(1)
					go r.runClaimed(j)
					r.durableMu.Unlock()
				}
			}
			select {
			case <-ctx.Done():
				return
			case <-ticker.C:
			}
		}
	}()
}

// reapExpiredJobsLoop runs the job_audit reaper every 15s, not on each claim.
func (r *Registry) reapExpiredJobsLoop(ctx context.Context) {
	defer r.durableLoopWG.Done()
	r.reapExpiredJobs(ctx)
	reap := time.NewTicker(15 * time.Second)
	stop := time.NewTicker(time.Second)
	defer reap.Stop()
	defer stop.Stop()
	for {
		if r.durableStopped() || ctx.Err() != nil {
			return
		}
		select {
		case <-ctx.Done():
			return
		case <-stop.C:
			if r.durableStopped() {
				return
			}
		case <-reap.C:
			r.reapExpiredJobs(ctx)
		}
	}
}

func (r *Registry) durableStopped() bool {
	if r == nil {
		return true
	}
	r.durableMu.Lock()
	defer r.durableMu.Unlock()
	return r.durableStop
}

func (r *Registry) reapExpiredJobs(ctx context.Context) {
	if r == nil || r.queue == nil {
		return
	}
	cctx, cancel := context.WithTimeout(ctx, 10*time.Second)
	defer cancel()
	if err := r.queue.ReapExpiredJobs(cctx); err != nil && ctx.Err() == nil {
		logging.With("component", "jobs.durable").Warn("reap expired jobs", "err", err)
	}
}

func (r *Registry) runClaimed(j *Job) {
	log := logging.With("component", "jobs.durable")
	var leaseHealthy atomic.Bool
	leaseHealthy.Store(true)
	defer func() {
		r.mu.Lock()
		delete(r.byID, j.ID)
		r.mu.Unlock()
		<-r.workerSem
		observability.RecordJobQueueDepth(len(r.workerSem), cap(r.workerSem))
		r.durableWG.Done()
	}()
	stop := make(chan struct{})
	done := make(chan struct{})
	go func() {
		defer close(done)
		t := time.NewTicker(durableLease / 4)
		defer t.Stop()
		for {
			select {
			case <-stop:
				return
			case <-t.C:
				ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
				requested, err := r.queue.HeartbeatJob(ctx, j.ID, j.OwnerID, j.FencingToken, durableLease)
				if err == nil {
					err = r.queue.CheckpointJob(ctx, recordFromJob(j))
				}
				cancel()
				if err != nil {
					leaseHealthy.Store(false)
					j.RequestCancel()
					log.Warn("heartbeat failed", "job_id", j.ID, "err", err)
					return
				}
				if requested {
					j.RequestCancel()
				}
			}
		}
	}()
	r.workerStart(j)
	close(stop)
	<-done
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	r.durableMu.Lock()
	stopped := r.durableStop
	r.durableMu.Unlock()
	c := recordFromJob(j)
	if stopped && c.Status != StatusSucceeded && c.Status != StatusFailed {
		if err := r.queue.ReleaseJob(ctx, j.ID, j.OwnerID, j.FencingToken); err != nil && !errors.Is(err, store.ErrJobLeaseLost) {
			log.Warn("release failed", "err", err)
		}
		return
	}
	if !leaseHealthy.Load() {
		// Leave an uncertain lease for expiry/reclaim. This consumes an attempt,
		// whereas a deliberate graceful shutdown releases without using the budget.
		return
	}
	if c.Status != StatusSucceeded && c.Status != StatusFailed && c.Status != StatusCancelled {
		c.Status = StatusFailed
		msg := "worker returned without terminal status"
		c.Error = &msg
	}
	if err := r.queue.FinishJob(ctx, c); err != nil {
		log.Warn("finish failed; lease will retry", "job_id", j.ID, "err", err)
	}
}