Something went wrong. Try again.
Monorepo for Tangled tangled.org
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194package mill
import ( "context" "slices" "strings" "sync" "testing" "time"
"tangled.org/core/spindle/db" "tangled.org/core/spindle/models" "tangled.org/core/spindle/observability"
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 ittype 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) }}
func TestWorkflowCountsNameTheRepoAndSplitWaiting(t *testing.T) { m, bdb := restoreTestMill(t, Config{ReconnectGrace: time.Minute}) if err := bdb.AddRepo(db.Repo{Knot: "knot", Owner: "did:plc:alice", Rkey: "core", RepoDid: "did:plc:core", Name: "core"}); err != nil { t.Fatal(err) }
m.mu.Lock() for range 2 { lease := newLease(m.nextLeaseIDLocked(), "node-a", "", "dummy") lease.ownerDID, lease.repoDID = "did:plc:alice", "did:plc:core" m.leases[lease.id] = lease } m.mu.Unlock() m.joinPlacement("did:plc:alice", "did:plc:core", "dummy") m.setAdmitting(m.joinPlacement("did:plc:alice", "did:plc:core", "dummy"), true) m.joinPlacement("did:plc:bob", "did:plc:fork", "dummy")
got := m.workflowCounts() slices.SortFunc(got, func(a, b observability.MillWorkflows) int { return strings.Compare(a.State+a.RepoDID, b.State+b.RepoDID) }) want := []observability.MillWorkflows{ {State: "running", OwnerDID: "did:plc:alice", RepoDID: "did:plc:core", Repo: "core", Count: 2}, {State: "waiting_quota", OwnerDID: "did:plc:alice", RepoDID: "did:plc:core", Repo: "core", Count: 1}, {State: "waiting_seat", OwnerDID: "did:plc:alice", RepoDID: "did:plc:core", Repo: "core", Count: 1}, {State: "waiting_seat", OwnerDID: "did:plc:bob", RepoDID: "did:plc:fork", Repo: "did:plc:fork", Count: 1}, } if !slices.Equal(got, want) { t.Fatalf("workflow counts = %+v, want %+v", got, want) }}