package mill
import (
"context"
"io"
"log/slog"
"runtime"
"slices"
"sync"
"testing"
"time"
"tangled.org/core/notifier"
"tangled.org/core/spindle/models"
"tangled.org/core/spindle/nixcache"
"tangled.org/core/spindle/quota"
millproto "tangled.org/core/spindle/mill/proto"
millv1 "tangled.org/core/spindle/mill/proto/gen"
)
type nixUse struct {
cache string
hashes []string
}
type nixUseStore struct {
*scriptedStore
umu sync.Mutex
uses []nixUse
gate chan struct{}
entered chan struct{}
}
func (s *nixUseStore) RecordNixUse(_ context.Context, cache string, hashes []string) error {
if s.gate != nil {
s.entered <- struct{}{}
<-s.gate
}
s.umu.Lock()
defer s.umu.Unlock()
s.uses = append(s.uses, nixUse{cache, hashes})
return nil
}
func (s *nixUseStore) recorded() []nixUse {
s.umu.Lock()
defer s.umu.Unlock()
return slices.Clone(s.uses)
}
func testRepoCache(t *testing.T) string {
t.Helper()
cache, err := nixcache.CacheName("did:plc:testrepo")
if err != nil {
t.Fatal(err)
}
return cache
}
func placeWithNixCache(t *testing.T, store quota.Store) (*Mill, *quotaCandidate, *RemoteLease) {
t.Helper()
m, _ := quotaMillWith(t, store, Config{NixCaches: true})
cand, lease := placeWorkflow(t, m)
return m, cand, lease
}
func placeWorkflow(t *testing.T, m *Mill) (*quotaCandidate, *RemoteLease) {
t.Helper()
cand := addQuotaCandidate(t, m, "node-a", fixedReport(2, 2048))
slot, err := m.place(context.Background(), "dummy", models.WorkflowId{PipelineId: models.PipelineId("rkey1"), Name: "build"}, testWorkflow("build"))
if err != nil {
t.Fatalf("place: %v", err)
}
t.Cleanup(slot.Release)
m.mu.Lock()
var lease *RemoteLease
for _, l := range m.leases {
lease = l
}
m.mu.Unlock()
return cand, lease
}
func nixObjectRequest(leaseID, cache string) *millv1.QuotaRequest {
return &millv1.QuotaRequest{
RequestId: "req-1",
Operation: millv1.QuotaOperation_QUOTA_OPERATION_RESERVE,
Kind: string(quota.KindNixCache),
Key: quota.NixCacheKey(cache, "abc"),
LeaseId: leaseID,
Resources: map[string]int64{"cache_storage_bytes": 4096},
NixObject: &millv1.NixObject{
Cache: cache,
StoreHash: "abc",
StorePath: "/nix/store/abc-pkg",
NarSize: 4096,
References: []string{"def-dep"},
},
}
}
func captureQuotaReplies(cand *quotaCandidate) *[]*millv1.QuotaResponse {
var replies []*millv1.QuotaResponse
cand.sess.enc = scriptedEncoder(func(msg *millproto.Message) error {
if resp := msg.GetQuotaResp(); resp != nil {
replies = append(replies, resp)
}
return nil
})
return &replies
}
func TestCacheReserveOnAnAtticMillNeedsItsObject(t *testing.T) {
for _, tt := range []struct {
name string
mutate func(*millv1.QuotaRequest)
allowed bool
}{
{"no object", func(req *millv1.QuotaRequest) { req.NixObject = nil }, false},
{"another repos cache", func(req *millv1.QuotaRequest) {
req.NixObject.Cache = "someone-elses"
req.Key = quota.NixCacheKey("someone-elses", "abc")
}, false},
{"another objects key", func(req *millv1.QuotaRequest) { req.Key = quota.NixCacheKey(req.NixObject.Cache, "other") }, false},
{"its own object", func(*millv1.QuotaRequest) {}, true},
} {
t.Run(tt.name, func(t *testing.T) {
store := &scriptedStore{verdicts: []quota.Reservation{
{ID: "res-workflow", Allowed: true, Reason: quota.ReasonWithinLimit},
{ID: "res-cache", Allowed: true, Reason: quota.ReasonWithinLimit},
}}
m, cand, lease := placeWithNixCache(t, store)
replies := captureQuotaReplies(cand)
before, _ := store.snapshot()
req := nixObjectRequest(lease.id, testRepoCache(t))
tt.mutate(req)
m.handleQuotaRequest(cand.sess, req)
if len(*replies) != 1 {
t.Fatalf("replies = %+v, want one", *replies)
}
got := (*replies)[0]
if got.GetAllowed() != tt.allowed || (got.GetError() == "") != tt.allowed {
t.Fatalf("reply = %+v, want allowed=%v", got, tt.allowed)
}
after, _ := store.snapshot()
if !tt.allowed {
if len(after) != len(before) {
t.Fatalf("store reserved %v for a refused request", after[len(before):])
}
return
}
want := "a.NixObject{Cache: testRepoCache(t), StoreHash: "abc", StorePath: "/nix/store/abc-pkg", NarSize: 4096, References: []string{"def-dep"}}
if obj := after[1].NixObject; obj == nil || obj.Cache != want.Cache || obj.StoreHash != want.StoreHash || obj.StorePath != want.StorePath || obj.NarSize != want.NarSize || !slices.Equal(obj.References, want.References) {
t.Fatalf("object = %+v, want %+v", obj, want)
}
})
}
}
func TestNixCacheUsedRecordsOnlyTheLeasesOwnCache(t *testing.T) {
store := &nixUseStore{scriptedStore: &scriptedStore{}}
m, cand, lease := placeWithNixCache(t, store)
m.handleNixCacheUsed(cand.sess, &millv1.NixCacheUsed{LeaseId: lease.id, Cache: testRepoCache(t), StoreHashes: []string{"a", "b"}})
m.handleNixCacheUsed(cand.sess, &millv1.NixCacheUsed{LeaseId: lease.id, Cache: "someone-elses", StoreHashes: []string{"c"}})
m.handleNixCacheUsed(cand.sess, &millv1.NixCacheUsed{LeaseId: "no-such-lease", Cache: testRepoCache(t), StoreHashes: []string{"d"}})
m.cfg.NixCaches = false
m.handleNixCacheUsed(cand.sess, &millv1.NixCacheUsed{LeaseId: lease.id, Cache: testRepoCache(t), StoreHashes: []string{"e"}})
got := store.recorded()
if len(got) != 1 || got[0].cache != testRepoCache(t) || !slices.Equal(got[0].hashes, []string{"a", "b"}) {
t.Fatalf("recorded = %+v, want only the lease's own cache", got)
}
}
func TestRestoredLeaseKeepsItsReposNixCache(t *testing.T) {
store := &nixUseStore{scriptedStore: &scriptedStore{verdicts: []quota.Reservation{
{ID: "res-workflow", Allowed: true, Reason: quota.ReasonWithinLimit},
{ID: "res-cache", Allowed: true, Reason: quota.ReasonWithinLimit},
}}}
m, cand, lease := placeWithNixCache(t, store)
bdb := m.db
n := notifier.New()
restarted := New(slog.New(slog.NewTextHandler(io.Discard, nil)), Config{ReconnectGrace: time.Minute, NixCaches: true})
qm := quota.NewManager(store, 10*time.Millisecond, nil)
t.Cleanup(qm.Close)
restarted.Attach(bdb, &n, qm)
if err := restarted.RestoreState(); err != nil {
t.Fatalf("RestoreState: %v", err)
}
replies := captureQuotaReplies(cand)
restarted.handleQuotaRequest(cand.sess, nixObjectRequest(lease.id, testRepoCache(t)))
restarted.handleQuotaRequest(cand.sess, nixObjectRequest(lease.id, "someone-elses"))
restarted.handleNixCacheUsed(cand.sess, &millv1.NixCacheUsed{LeaseId: lease.id, Cache: testRepoCache(t), StoreHashes: []string{"a"}})
restarted.handleNixCacheUsed(cand.sess, &millv1.NixCacheUsed{LeaseId: lease.id, Cache: "someone-elses", StoreHashes: []string{"b"}})
if len(*replies) != 2 || !(*replies)[0].GetAllowed() || (*replies)[0].GetError() != "" {
t.Fatalf("own cache reply = %+v, want it allowed", *replies)
}
if (*replies)[1].GetAllowed() || (*replies)[1].GetError() == "" {
t.Fatalf("other cache reply = %+v, want a refusal", (*replies)[1])
}
if got := store.recorded(); len(got) != 1 || got[0].cache != testRepoCache(t) || !slices.Equal(got[0].hashes, []string{"a"}) {
t.Fatalf("recorded = %+v, want only the repo's own cache", got)
}
}
func TestNixCacheUsedReportsBeyondTheConcurrencyBoundAreDropped(t *testing.T) {
store := &nixUseStore{scriptedStore: &scriptedStore{}, gate: make(chan struct{}), entered: make(chan struct{}, nixUseConcurrency+1)}
m, cand, lease := placeWithNixCache(t, store)
report := &millv1.NixCacheUsed{LeaseId: lease.id, Cache: testRepoCache(t), StoreHashes: []string{"a"}}
for range nixUseConcurrency {
m.startNixCacheUsed(cand.sess, report)
<-store.entered
}
m.startNixCacheUsed(cand.sess, report)
close(store.gate)
deadline := time.Now().Add(5 * time.Second)
for len(m.nixUseSlots) != 0 {
if time.Now().After(deadline) {
t.Fatal("recording goroutines never finished")
}
runtime.Gosched()
}
if got := store.recorded(); len(got) != nixUseConcurrency {
t.Fatalf("recorded %d reports, want %d: the one over the bound is dropped", len(got), nixUseConcurrency)
}
}