Improve the performance of creating FairSchedulingAlgoContext

[JamesMurkin]

When creating the fairshare scheduling context, we loop through all jobs of the jobdb and then filter out the ones relevant to the current pool

We then use these jobs to calculate relevant info about the pool

• Which jobs are on which nodes
• The demand of each queue (and subsequently the fairshare/adjusted fairshare)

This is simple but quite inefficient if you have a lot (millions) of jobs in the system and many pools (most of which only a small fraction of jobs interact with)

• As a result most of the scheduling time can be taken up by this data shuffling

Now instead of loading all jobs, we'll load on relevant jobs:

• All leased jobs (needed to calculate which jobs are on which nodes)
  • Possibly this could be further improved by storing them by node but for now this is simple
• All queued jobs against pools we are processing (current pool + away pools)

This should have a significant impact, especially in the case we have 1 pool with most jobs queued against it and many smaller pools

There is larger reworks we could do here to make it even more efficient, however for now this should give us most the benefit for a minor change

• Example would be to calculate demand for each pool + jobs on each node upfront. Incrementally update this as pools are processed rather than recalculate it entirely

[greptile-apps]

Greptile Summary

This PR significantly improves the performance of FairSchedulingAlgoContext creation by replacing an unconditional txn.GetAll() call (which iterated every job in the database) with a targeted fetch of only the jobs relevant to the pool being scheduled: all currently leased jobs, all terminal jobs (for short-job penalty calculation), and queued jobs scoped to the home/away pools under consideration.

Key changes:

• New Leased() predicate on Job — !queued && !terminal && LatestRun() != nil — cleanly encapsulates the leased-job concept used throughout the new indexes.
• Two new secondary indexes in JobDb/Txn (leasedJobs, terminalJobs) maintained incrementally in Upsert and delete, following the same immutable-set pattern already used for unvalidatedJobs.
• GetAllLeasedJobs(), GetAllTerminalJobs(), GetQueuedJobsByPool(pool) — new query methods that expose the new indexes and the existing jobsByPoolAndQueue map.
• Generic UniqueBy utility in common/slices to deduplicate the assembled allJobs slice (necessary because a job queued for multiple pools in allPools would otherwise appear more than once).
• Previously flagged review issues (copy-paste bug assigning newLeasedJobs to txn.unvalidatedJobs, and over-allocation in GetQueuedJobsByPool) are both correctly addressed in this revision.
• The optimization is logically sound: queued jobs for pools outside allPools were already no-ops inside calculateJobSchedulingInfo (filtered by pool membership), so excluding them at load time produces identical scheduling results.
• Minor: terminal jobs are loaded globally rather than per-pool; the short-job penalty filter inside calculateJobSchedulingInfo already gates by jobPool == currentPool, so correctness is unaffected, but in systems with many completed jobs this remains an optimisation opportunity noted in the PR description.

Confidence Score: 5/5

Safe to merge — the optimization is logically equivalent to the previous full-scan approach, previously flagged bugs are fixed, and the new indexes are well-tested.

All remaining findings are P2: a missing capacity hint in UniqueBy and a slightly inaccurate code comment. Neither affects correctness or runtime behaviour. The previously raised P0/P1 concerns (copy-paste bug in delete, over-allocation in GetQueuedJobsByPool) are both resolved in this revision.

No files require special attention.

Important Files Changed

Filename	Overview
internal/scheduler/jobdb/jobdb.go	Adds leasedJobs and terminalJobs indexes maintained in Upsert and delete; previously flagged bugs (copy-paste in delete, over-allocation in GetQueuedJobsByPool) are both addressed in this revision.
internal/scheduler/scheduling/scheduling_algo.go	Replaces txn.GetAll() with targeted leased+terminal+queued-per-pool job loading; logically sound — queued jobs for irrelevant pools were already skipped inside calculateJobSchedulingInfo, so excluding them at load time produces identical results.
internal/scheduler/jobdb/job.go	Adds Leased() predicate (!queued && !terminal && LatestRun() != nil); correct and consistent with how the scheduler treats leased jobs.
internal/common/slices/slices.go	Adds generic UniqueBy helper that deduplicates a slice by a key function; minor missing capacity hint on the result slice.
internal/scheduler/jobdb/jobdb_test.go	Comprehensive lifecycle and deletion tests for leasedJobs, terminalJobs, and GetQueuedJobsByPool; good coverage of edge cases including terminal+queued combo jobs.
internal/common/slices/slices_test.go	Good test coverage for UniqueBy covering nil, empty, consecutive/non-consecutive duplicates, and custom key function.
internal/scheduler/jobdb/job_test.go	Adds TestJob_Leased covering queued, terminal (succeeded/failed/cancelled), and no-run states; all meaningful branches of the new Leased() predicate are exercised.

Sequence Diagram

sequenceDiagram
    participant SA as FairSchedulingAlgo
    participant TXN as JobDb Txn
    participant CSI as calculateJobSchedulingInfo

    SA->>TXN: GetAllLeasedJobs()
    TXN-->>SA: leased []Job

    SA->>TXN: GetAllTerminalJobs()
    TXN-->>SA: terminal []Job

    loop for each pool in allPools
        SA->>TXN: GetQueuedJobsByPool(pool)
        TXN-->>SA: queued []Job
    end

    SA->>SA: UniqueBy(allJobs, job.Id())

    SA->>CSI: calculateJobSchedulingInfo(allJobs, ...)
    CSI-->>SA: jobSchedulingInfo (demand, allocation, shortJobPenalty, jobsByPool)

    SA->>SA: build nodeDb from leasedJobs (jobsByPool) using currentPoolJobs + otherPoolsJobs
    SA-->>SA: FairSchedulingAlgoContext

Loading

_{Reviews (8): Last reviewed commit: "Merge branch 'master' into improve_fsctx..." | Re-trigger Greptile}