diff --git a/spindle/mill/fair.go b/spindle/mill/fair.go new file mode 100644 index 00000000..c23bb9c1 --- /dev/null +++ b/spindle/mill/fair.go @@ -0,0 +1,86 @@ +package mill + +// a workflow waiting in place for a seat. when a seat frees up, waiters bid +// one at a time, starting with the owner that has the fewest workflows on +// the fleet, so one owner's burst can't take every seat +type placeWaiter struct { + owner string + engine string + seq uint64 + admitting bool + // set when a bid in triedGen found no seat. the waiter then sits out and + // stops blocking the ones behind it until something changes + tried bool + triedGen uint64 +} + +func (w *placeWaiter) triedSince(gen uint64) bool { + return w.tried && w.triedGen == gen +} + +func (m *Mill) joinPlacement(owner, engine string) *placeWaiter { + m.mu.Lock() + defer m.mu.Unlock() + m.waiterSeq++ + w := &placeWaiter{owner: owner, engine: engine, seq: m.waiterSeq} + m.waiters[w] = struct{}{} + return w +} + +func (m *Mill) leavePlacement(w *placeWaiter) { + m.mu.Lock() + defer m.mu.Unlock() + delete(m.waiters, w) + m.wakeWaitersLocked() +} + +// reports whether w may bid now, plus the generation and channel to wait on +func (m *Mill) placementTurn(w *placeWaiter) (bool, uint64, <-chan struct{}) { + m.mu.Lock() + defer m.mu.Unlock() + if w.triedSince(m.changeGen) { + return false, m.changeGen, m.changeCh + } + running := m.runningByOwnerLocked() + for v := range m.waiters { + if v == w || v.engine != w.engine || v.admitting || v.triedSince(m.changeGen) { + continue + } + if running[v.owner] < running[w.owner] || (running[v.owner] == running[w.owner] && v.seq < w.seq) { + return false, m.changeGen, m.changeCh + } + } + return true, m.changeGen, m.changeCh +} + +// the bid found no seat, so let the next waiter try +func (m *Mill) passTurn(w *placeWaiter, gen uint64) { + m.mu.Lock() + defer m.mu.Unlock() + w.tried = true + w.triedGen = gen + m.wakeWaitersLocked() +} + +func (m *Mill) setAdmitting(w *placeWaiter, admitting bool) { + m.mu.Lock() + defer m.mu.Unlock() + w.admitting = admitting + if !admitting { + m.wakeWaitersLocked() + } +} + +// admitting waiters count too because they're about to hold a seat +func (m *Mill) runningByOwnerLocked() map[string]int { + running := make(map[string]int) + for _, lease := range m.leases { + running[lease.ownerDID]++ + } + for w := range m.waiters { + if w.admitting { + running[w.owner]++ + } + } + return running +} diff --git a/spindle/mill/fair_test.go b/spindle/mill/fair_test.go new file mode 100644 index 00000000..5a79011b --- /dev/null +++ b/spindle/mill/fair_test.go @@ -0,0 +1,157 @@ +package mill + +import ( + "context" + "sync" + "testing" + "time" + + "tangled.org/core/spindle/models" + + millproto "tangled.org/core/spindle/mill/proto" + millv1 "tangled.org/core/spindle/mill/proto/gen" +) + +func workflowOwnedBy(wf *models.Workflow, owner string) *models.Workflow { + state := wf.Data.(*millWorkflowState) + md := *state.RawPipeline.TriggerMetadata + repo := *md.Repo + repo.Did = owner + md.Repo = &repo + state.RawPipeline.TriggerMetadata = &md + return wf +} + +// a node that turns bids away until the test frees seats on it +type seatedNode struct { + mu sync.Mutex + free int +} + +func (n *seatedNode) open(m *Mill, seats int) { + n.mu.Lock() + n.free += seats + n.mu.Unlock() + m.notifyChange() +} + +func addSeatedNode(t *testing.T, m *Mill) *seatedNode { + t.Helper() + node := &seatedNode{} + var sess *millSession + sess = addCandidateSession(t, m, "node-a", nil, 0, scriptedEncoder(func(msg *millproto.Message) error { + rs := msg.GetReserveSeat() + if rs == nil { + return nil + } + node.mu.Lock() + accepted := node.free > 0 + if accepted { + node.free-- + } + node.mu.Unlock() + sess.deliver(rs.GetLeaseId(), &millproto.Message{ReserveResult: &millv1.ReserveResult{ + LeaseId: rs.GetLeaseId(), + Accepted: accepted, + }}) + return nil + })) + return node +} + +func runningFor(m *Mill, owner string, n int) { + m.mu.Lock() + defer m.mu.Unlock() + for range n { + lease := newLease(m.nextLeaseIDLocked(), "node-a", "", "dummy") + lease.ownerDID = owner + m.leases[lease.id] = lease + } +} + +func waitForTriedWaiters(t *testing.T, m *Mill, want int) { + t.Helper() + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + m.mu.Lock() + tried := 0 + for w := range m.waiters { + if w.tried { + tried++ + } + } + m.mu.Unlock() + if tried == want { + return + } + time.Sleep(time.Millisecond) + } + t.Fatalf("waiters never all bid, want %d", want) +} + +func placeInBackground(t *testing.T, m *Mill, wf *models.Workflow) <-chan string { + t.Helper() + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + placed := make(chan string, 1) + go func() { + wid := models.WorkflowId{PipelineId: models.PipelineId("p"), Name: wf.Name} + if _, err := m.place(ctx, "dummy", wid, wf); err != nil { + placed <- "error: " + err.Error() + return + } + placed <- wf.Name + }() + return placed +} + +func nextPlaced(t *testing.T, placed ...<-chan string) string { + t.Helper() + cases := make(chan string, len(placed)) + for _, ch := range placed { + go func() { + if name, ok := <-ch; ok { + cases <- name + } + }() + } + select { + case name := <-cases: + return name + case <-time.After(2 * time.Second): + t.Fatal("no workflow was placed") + return "" + } +} + +func TestFreedSeatGoesToTheOwnerWithFewerRunning(t *testing.T) { + m, _ := quotaMill(t, &scriptedStore{}) + node := addSeatedNode(t, m) + runningFor(m, "did:plc:busy", 2) + + busy := placeInBackground(t, m, workflowOwnedBy(testWorkflow("busy"), "did:plc:busy")) + waitForTriedWaiters(t, m, 1) + quiet := placeInBackground(t, m, workflowOwnedBy(testWorkflow("quiet"), "did:plc:quiet")) + waitForTriedWaiters(t, m, 2) + + node.open(m, 1) + if got := nextPlaced(t, busy, quiet); got != "quiet" { + t.Fatalf("first freed seat went to %q, want the owner with nothing running", got) + } +} + +func TestWaiterThatCannotFitDoesNotBlockTheRest(t *testing.T) { + m, _ := quotaMill(t, &scriptedStore{}) + node := addSeatedNode(t, m) + runningFor(m, "did:plc:busy", 2) + + // ranks first but no node has the label it needs + gpu := placeInBackground(t, m, workflowOwnedBy(testWorkflowWithRunsOn("gpu", []string{"gpu"}), "did:plc:quiet")) + busy := placeInBackground(t, m, workflowOwnedBy(testWorkflow("busy"), "did:plc:busy")) + waitForTriedWaiters(t, m, 2) + + node.open(m, 1) + if got := nextPlaced(t, gpu, busy); got != "busy" { + t.Fatalf("freed seat went to %q, want the waiter that fits", got) + } +} diff --git a/spindle/mill/mill.go b/spindle/mill/mill.go index 7d0d0243..455b9d3d 100644 --- a/spindle/mill/mill.go +++ b/spindle/mill/mill.go @@ -92,6 +92,9 @@ type Mill struct { protoStrikes map[string]int pending int changeCh chan struct{} + changeGen uint64 + waiters map[*placeWaiter]struct{} + waiterSeq uint64 cacheDeleteOnce sync.Once cacheDeleteMu sync.Mutex @@ -135,6 +138,7 @@ func New(l *slog.Logger, cfg Config) *Mill { nodeSeqno: make(map[string]uint64), changeCh: make(chan struct{}), quotaLeases: make(map[string]*RemoteLease), + waiters: make(map[*placeWaiter]struct{}), } m.startCacheDeletionWorker() return m @@ -179,10 +183,13 @@ func (m *Mill) ensureCacheSentinel(store storage.Storage) { func (m *Mill) nextLeaseID() string { m.mu.Lock() + defer m.mu.Unlock() + return m.nextLeaseIDLocked() +} + +func (m *Mill) nextLeaseIDLocked() string { m.leaseSeq++ - seq := m.leaseSeq - m.mu.Unlock() - return fmt.Sprintf("%s-%d", tid.TID(), seq) + return fmt.Sprintf("%s-%d", tid.TID(), m.leaseSeq) } func (m *Mill) dropReservation(id string) { @@ -198,6 +205,11 @@ func (m *Mill) notifyChange() { } func (m *Mill) notifyChangeLocked() { + m.changeGen++ + m.wakeWaitersLocked() +} + +func (m *Mill) wakeWaitersLocked() { close(m.changeCh) m.changeCh = make(chan struct{}) } @@ -459,6 +471,13 @@ func (m *Mill) place(ctx context.Context, engineName string, wid models.Workflow m.mu.Unlock() }() + owner := wf.OwnerDID + if id, ok := m.quotaIdentity(wf); ok { + owner = id.OwnerDID + } + w := m.joinPlacement(owner, engineName) + defer m.leavePlacement(w) + for { if err := ctx.Err(); err != nil { placementResult = "cancelled" @@ -468,30 +487,35 @@ func (m *Mill) place(ctx context.Context, engineName string, wid models.Workflow return nil, err } - // capture the channel before bidding - // a change during the bid makes the next wait return immediately - ch := m.currentChangeCh() - - lease, err := m.bid(ctx, engineName, wid, wf) - if err != nil { - return nil, err - } - if lease != nil { - wf.CacheNamespace = lease.cacheNamespace - slot, retry, err := m.admit(ctx, engineName, wid, wf, lease) + // the channel comes from before the bid, so a change during the bid + // makes the next wait return immediately + turn, gen, ch := m.placementTurn(w) + if turn { + lease, err := m.bid(ctx, engineName, wid, wf) if err != nil { - var failure *engine.WorkflowFailure - if errors.As(err, &failure) && failure.Class == engine.FailureClassPolicy { - placementResult = "rejected" - } return nil, err } - if slot != nil { - placementResult = "success" - return slot, nil - } - if retry { - continue + if lease == nil { + m.passTurn(w, gen) + } else { + wf.CacheNamespace = lease.cacheNamespace + m.setAdmitting(w, true) + slot, retry, err := m.admit(ctx, engineName, wid, wf, lease) + m.setAdmitting(w, false) + if err != nil { + var failure *engine.WorkflowFailure + if errors.As(err, &failure) && failure.Class == engine.FailureClassPolicy { + placementResult = "rejected" + } + return nil, err + } + if slot != nil { + placementResult = "success" + return slot, nil + } + if retry { + continue + } } }