Something went wrong. Try again.
Monorepo for Tangled
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705package executor
import ( "context" "errors" "fmt" "math/rand/v2" "sync" "testing" "time"
"tangled.org/core/spindle/quota"
millproto "tangled.org/core/spindle/mill/proto" millv1 "tangled.org/core/spindle/mill/proto/gen")
var testIdentity = quota.Identity{ OwnerDID: "did:web:owner", RepoDID: "did:web:repo",}
func validReserveRequest(kind quota.Kind, key string, resource string, amount int64) quota.ReserveRequest { return quota.ReserveRequest{ Kind: kind, Key: key, Identity: testIdentity, Resources: quota.Resources{ resource: amount, }, }}
func TestQuotaClientReserveAndTransitions(t *testing.T) { client := NewQuotaClient() leaseID := "lease-abc-123" store := client.ForLease(leaseID)
var ( mu sync.Mutex sent []*millproto.Message )
client.SetSendFn(func(msg *millproto.Message) error { mu.Lock() sent = append(sent, msg) mu.Unlock()
req := msg.GetQuotaReq() if req == nil { t.Error("expected QuotaReq message") return nil }
go func() { switch req.GetOperation() { case millv1.QuotaOperation_QUOTA_OPERATION_RESERVE: client.HandleResponse(req.GetRequestId(), &millv1.QuotaResponse{ RequestId: req.GetRequestId(), ReservationId: "res-123", Allowed: true, Reason: "within_limits", Resource: quota.ResourceCacheStorageBytes, }) case millv1.QuotaOperation_QUOTA_OPERATION_BEGIN_COMMIT, millv1.QuotaOperation_QUOTA_OPERATION_COMMIT, millv1.QuotaOperation_QUOTA_OPERATION_RELEASE: client.HandleResponse(req.GetRequestId(), &millv1.QuotaResponse{ RequestId: req.GetRequestId(), Allowed: true, }) } }() return nil })
ctx := context.Background()
res, err := store.Reserve(ctx, validReserveRequest(quota.KindNixCache, "nar-hash-1", quota.ResourceCacheStorageBytes, 2048)) if err != nil { t.Fatalf("Reserve() unexpected error: %v", err) } if res.ID != "res-123" { t.Errorf("res.ID = %q, want %q", res.ID, "res-123") } if !res.Allowed { t.Error("res.Allowed = false, want true") } if res.Temporary { t.Error("res.Temporary = true, want false") } if res.Reason != "within_limits" { t.Errorf("res.Reason = %q, want %q", res.Reason, "within_limits") } if res.Resource != quota.ResourceCacheStorageBytes { t.Errorf("res.Resource = %q, want %q", res.Resource, quota.ResourceCacheStorageBytes) }
mu.Lock() if len(sent) != 1 { t.Fatalf("sent %d messages, want 1", len(sent)) } reqMsg := sent[0].GetQuotaReq() if reqMsg.GetOperation() != millv1.QuotaOperation_QUOTA_OPERATION_RESERVE { t.Errorf("operation = %v, want RESERVE", reqMsg.GetOperation()) } if reqMsg.GetLeaseId() != leaseID { t.Errorf("lease_id = %q, want %q", reqMsg.GetLeaseId(), leaseID) } if reqMsg.GetKind() != string(quota.KindNixCache) { t.Errorf("kind = %q, want %q", reqMsg.GetKind(), quota.KindNixCache) } if reqMsg.GetKey() != "nar-hash-1" { t.Errorf("key = %q, want %q", reqMsg.GetKey(), "nar-hash-1") } if reqMsg.GetResources()["cache_storage_bytes"] != 2048 { t.Errorf("resources.cache_storage_bytes = %d, want 2048", reqMsg.GetResources()["cache_storage_bytes"]) } mu.Unlock()
if err := store.BeginCommit(ctx, "res-123"); err != nil { t.Fatalf("BeginCommit() unexpected error: %v", err) }
mu.Lock() if len(sent) != 2 { t.Fatalf("sent %d messages, want 2", len(sent)) } bcMsg := sent[1].GetQuotaReq() if bcMsg.GetOperation() != millv1.QuotaOperation_QUOTA_OPERATION_BEGIN_COMMIT { t.Errorf("operation = %v, want BEGIN_COMMIT", bcMsg.GetOperation()) } if bcMsg.GetLeaseId() != leaseID { t.Errorf("lease_id = %q, want %q", bcMsg.GetLeaseId(), leaseID) } if bcMsg.GetReservationId() != "res-123" { t.Errorf("reservation_id = %q, want %q", bcMsg.GetReservationId(), "res-123") } mu.Unlock()
if err := store.Commit(ctx, "res-123"); err != nil { t.Fatalf("Commit() unexpected error: %v", err) }
mu.Lock() if len(sent) != 3 { t.Fatalf("sent %d messages, want 3", len(sent)) } cMsg := sent[2].GetQuotaReq() if cMsg.GetOperation() != millv1.QuotaOperation_QUOTA_OPERATION_COMMIT { t.Errorf("operation = %v, want COMMIT", cMsg.GetOperation()) } if cMsg.GetLeaseId() != leaseID { t.Errorf("lease_id = %q, want %q", cMsg.GetLeaseId(), leaseID) } if cMsg.GetReservationId() != "res-123" { t.Errorf("reservation_id = %q, want %q", cMsg.GetReservationId(), "res-123") } mu.Unlock()
if err := store.Release(ctx, "res-123"); err != nil { t.Fatalf("Release() unexpected error: %v", err) }
mu.Lock() if len(sent) != 4 { t.Fatalf("sent %d messages, want 4", len(sent)) } rMsg := sent[3].GetQuotaReq() if rMsg.GetOperation() != millv1.QuotaOperation_QUOTA_OPERATION_RELEASE { t.Errorf("operation = %v, want RELEASE", rMsg.GetOperation()) } if rMsg.GetLeaseId() != leaseID { t.Errorf("lease_id = %q, want %q", rMsg.GetLeaseId(), leaseID) } if rMsg.GetReservationId() != "res-123" { t.Errorf("reservation_id = %q, want %q", rMsg.GetReservationId(), "res-123") } mu.Unlock()
if err := store.BeginCommit(ctx, ""); err != nil { t.Errorf("BeginCommit(\"\") error = %v", err) } if err := store.Commit(ctx, ""); err != nil { t.Errorf("Commit(\"\") error = %v", err) } if err := store.Release(ctx, ""); err != nil { t.Errorf("Release(\"\") error = %v", err) }
mu.Lock() if len(sent) != 4 { t.Errorf("sent %d messages after empty transitions, want 4", len(sent)) } mu.Unlock()}
func TestQuotaClientGenericValidation(t *testing.T) { client := NewQuotaClient() store := client.ForLease("lease-1")
client.SetSendFn(func(msg *millproto.Message) error { t.Error("sendFn must not be called for invalid requests") return nil })
ctx := context.Background()
invalidRequests := []struct { name string req quota.ReserveRequest }{ { name: "invalid kind", req: quota.ReserveRequest{ Kind: quota.Kind("invalid-kind"), Key: "key-1", Identity: testIdentity, Resources: quota.Resources{quota.ResourceCacheStorageBytes: 100}, }, }, { name: "empty key", req: quota.ReserveRequest{ Kind: quota.KindNixCache, Key: "", Identity: testIdentity, Resources: quota.Resources{quota.ResourceCacheStorageBytes: 100}, }, }, { name: "empty resources", req: quota.ReserveRequest{ Kind: quota.KindNixCache, Key: "key-1", Identity: testIdentity, Resources: nil, }, }, { name: "negative resource amount", req: quota.ReserveRequest{ Kind: quota.KindNixCache, Key: "key-1", Identity: testIdentity, Resources: quota.Resources{quota.ResourceCacheStorageBytes: -50}, }, }, { name: "empty owner DID", req: quota.ReserveRequest{ Kind: quota.KindNixCache, Key: "key-1", Identity: quota.Identity{OwnerDID: "", RepoDID: "did:web:repo"}, Resources: quota.Resources{quota.ResourceCacheStorageBytes: 100}, }, }, { name: "empty repo DID", req: quota.ReserveRequest{ Kind: quota.KindNixCache, Key: "key-1", Identity: quota.Identity{OwnerDID: "did:web:owner", RepoDID: ""}, Resources: quota.Resources{quota.ResourceCacheStorageBytes: 100}, }, }, { name: "invalid owner DID prefix", req: quota.ReserveRequest{ Kind: quota.KindNixCache, Key: "key-1", Identity: quota.Identity{OwnerDID: "not-a-did", RepoDID: "did:web:repo"}, Resources: quota.Resources{quota.ResourceCacheStorageBytes: 100}, }, }, { name: "invalid repo DID prefix", req: quota.ReserveRequest{ Kind: quota.KindNixCache, Key: "key-1", Identity: quota.Identity{OwnerDID: "did:web:owner", RepoDID: "not-a-did"}, Resources: quota.Resources{quota.ResourceCacheStorageBytes: 100}, }, }, }
for _, tc := range invalidRequests { t.Run(tc.name, func(t *testing.T) { _, err := store.Reserve(ctx, tc.req) if err == nil { t.Fatalf("Reserve() expected validation error for %s", tc.name) } }) }
validKinds := []quota.Kind{quota.KindWorkflow, quota.KindNixCache, quota.KindGenericCache} for _, k := range validKinds { t.Run("valid kind "+string(k), func(t *testing.T) { called := false client.SetSendFn(func(msg *millproto.Message) error { called = true go client.HandleResponse(msg.GetQuotaReq().GetRequestId(), &millv1.QuotaResponse{ RequestId: msg.GetQuotaReq().GetRequestId(), ReservationId: "res-ok", Allowed: true, }) return nil }) req := validReserveRequest(k, "key-valid", "custom_units", 100) res, err := store.Reserve(ctx, req) if err != nil { t.Fatalf("unexpected validation failure for kind %s: %v", k, err) } if !called || !res.Allowed { t.Fatalf("expected request to be sent and allowed for kind %s", k) } }) }}
func TestQuotaClientResourceResponsePropagation(t *testing.T) { ctx := context.Background()
t.Run("resource roundtrip", func(t *testing.T) { client := NewQuotaClient() store := client.ForLease("lease-1")
client.SetSendFn(func(msg *millproto.Message) error { go client.HandleResponse(msg.GetQuotaReq().GetRequestId(), &millv1.QuotaResponse{ RequestId: msg.GetQuotaReq().GetRequestId(), ReservationId: "res-mem", Allowed: true, Reason: "within_limits", Resource: quota.ResourceMemoryMiB, }) return nil })
res, err := store.Reserve(ctx, validReserveRequest(quota.KindWorkflow, "wf-1", quota.ResourceMemoryMiB, 4096)) if err != nil { t.Fatalf("Reserve() error = %v", err) } if res.Resource != quota.ResourceMemoryMiB { t.Errorf("res.Resource = %q, want %q", res.Resource, quota.ResourceMemoryMiB) } if res.ID != "res-mem" { t.Errorf("res.ID = %q, want %q", res.ID, "res-mem") } })
t.Run("temporary denial", func(t *testing.T) { client := NewQuotaClient() store := client.ForLease("lease-1")
client.SetSendFn(func(msg *millproto.Message) error { go client.HandleResponse(msg.GetQuotaReq().GetRequestId(), &millv1.QuotaResponse{ RequestId: msg.GetQuotaReq().GetRequestId(), Allowed: false, Temporary: true, Reason: "repo_queue_full", }) return nil })
res, err := store.Reserve(ctx, validReserveRequest(quota.KindNixCache, "key-1", quota.ResourceCacheStorageBytes, 500)) if err != nil { t.Fatalf("Reserve() error = %v", err) } if res.Allowed { t.Error("res.Allowed = true, want false") } if !res.Temporary { t.Error("res.Temporary = false, want true") } if res.Reason != "repo_queue_full" { t.Errorf("res.Reason = %q, want %q", res.Reason, "repo_queue_full") } })
t.Run("permanent denial", func(t *testing.T) { client := NewQuotaClient() store := client.ForLease("lease-1")
client.SetSendFn(func(msg *millproto.Message) error { go client.HandleResponse(msg.GetQuotaReq().GetRequestId(), &millv1.QuotaResponse{ RequestId: msg.GetQuotaReq().GetRequestId(), Allowed: false, Temporary: false, Reason: "user_limit_exceeded", }) return nil })
res, err := store.Reserve(ctx, validReserveRequest(quota.KindNixCache, "key-1", quota.ResourceCacheStorageBytes, 500)) if err != nil { t.Fatalf("Reserve() error = %v", err) } if res.Allowed { t.Error("res.Allowed = true, want false") } if res.Temporary { t.Error("res.Temporary = true, want false") } if res.Reason != "user_limit_exceeded" { t.Errorf("res.Reason = %q, want %q", res.Reason, "user_limit_exceeded") } })
t.Run("allowed and temporary contradictory response", func(t *testing.T) { client := NewQuotaClient() store := client.ForLease("lease-1")
client.SetSendFn(func(msg *millproto.Message) error { go client.HandleResponse(msg.GetQuotaReq().GetRequestId(), &millv1.QuotaResponse{ RequestId: msg.GetQuotaReq().GetRequestId(), ReservationId: "res-bad", Allowed: true, Temporary: true, }) return nil })
_, err := store.Reserve(ctx, validReserveRequest(quota.KindNixCache, "key-1", quota.ResourceCacheStorageBytes, 500)) if err == nil { t.Fatal("expected error for allowed+temporary response") } })
t.Run("mill error response", func(t *testing.T) { client := NewQuotaClient() store := client.ForLease("lease-1")
client.SetSendFn(func(msg *millproto.Message) error { go client.HandleResponse(msg.GetQuotaReq().GetRequestId(), &millv1.QuotaResponse{ RequestId: msg.GetQuotaReq().GetRequestId(), Error: "quota store database failure", }) return nil })
_, err := store.Reserve(ctx, validReserveRequest(quota.KindNixCache, "key-1", quota.ResourceCacheStorageBytes, 500)) if err == nil || err.Error() != "quota store database failure" { t.Fatalf("expected error %q, got %v", "quota store database failure", err) } })}
func TestQuotaClientLeaseBinding(t *testing.T) { client := NewQuotaClient() ctx := context.Background()
var ( mu sync.Mutex seenLeases []string )
client.SetSendFn(func(msg *millproto.Message) error { req := msg.GetQuotaReq() mu.Lock() seenLeases = append(seenLeases, req.GetLeaseId()) mu.Unlock()
go client.HandleResponse(req.GetRequestId(), &millv1.QuotaResponse{ RequestId: req.GetRequestId(), ReservationId: "res-" + req.GetLeaseId(), Allowed: true, }) return nil })
s1 := client.ForLease("lease-alpha") s2 := client.ForLease("lease-beta")
_, err := s1.Reserve(ctx, validReserveRequest(quota.KindNixCache, "key-a", quota.ResourceCacheStorageBytes, 100)) if err != nil { t.Fatalf("s1.Reserve() error = %v", err) }
_, err = s2.Reserve(ctx, validReserveRequest(quota.KindNixCache, "key-b", quota.ResourceCacheStorageBytes, 200)) if err != nil { t.Fatalf("s2.Reserve() error = %v", err) }
mu.Lock() if len(seenLeases) != 2 || seenLeases[0] != "lease-alpha" || seenLeases[1] != "lease-beta" { t.Fatalf("seen leases = %v, want [lease-alpha, lease-beta]", seenLeases) } mu.Unlock()
sEmpty := client.ForLease("") if _, err := sEmpty.Reserve(ctx, validReserveRequest(quota.KindNixCache, "key-c", quota.ResourceCacheStorageBytes, 100)); err == nil { t.Error("sEmpty.Reserve() expected error") } if err := sEmpty.BeginCommit(ctx, "res-1"); err == nil { t.Error("sEmpty.BeginCommit() expected error") } if err := sEmpty.Commit(ctx, "res-1"); err == nil { t.Error("sEmpty.Commit() expected error") } if err := sEmpty.Release(ctx, "res-1"); err == nil { t.Error("sEmpty.Release() expected error") }}
func TestQuotaClientDisconnect(t *testing.T) { t.Run("fails closed when not connected", func(t *testing.T) { client := NewQuotaClient() store := client.ForLease("lease-1") ctx := context.Background()
if _, err := store.Reserve(ctx, validReserveRequest(quota.KindNixCache, "key-1", quota.ResourceCacheStorageBytes, 100)); err == nil { t.Error("Reserve() expected error when not connected") } if err := store.BeginCommit(ctx, "res-1"); err == nil { t.Error("BeginCommit() expected error when not connected") } if err := store.Commit(ctx, "res-1"); err == nil { t.Error("Commit() expected error when not connected") } if err := store.Release(ctx, "res-1"); err == nil { t.Error("Release() expected error when not connected") } })
t.Run("in-flight unblocked by OnDisconnect", func(t *testing.T) { client := NewQuotaClient() store := client.ForLease("lease-1") sent := make(chan struct{})
client.SetSendFn(func(msg *millproto.Message) error { close(sent) return nil })
done := make(chan error, 1) go func() { _, err := store.Reserve(context.Background(), validReserveRequest(quota.KindNixCache, "key-1", quota.ResourceCacheStorageBytes, 100)) done <- err }()
<-sent client.OnDisconnect()
select { case err := <-done: if err == nil || err.Error() != "quota client disconnected" { t.Fatalf("expected %q error, got %v", "quota client disconnected", err) } case <-time.After(2 * time.Second): t.Fatal("timed out waiting for Reserve() to return after OnDisconnect") } })
t.Run("late duplicate responses", func(t *testing.T) { client := NewQuotaClient() client.HandleResponse("unknown-id", &millv1.QuotaResponse{ RequestId: "unknown-id", ReservationId: "res-late", Allowed: true, }) })
t.Run("reconnect works", func(t *testing.T) { client := NewQuotaClient() store := client.ForLease("lease-1")
client.SetSendFn(func(msg *millproto.Message) error { return errors.New("temporary failure") })
_, err := store.Reserve(context.Background(), validReserveRequest(quota.KindNixCache, "key-1", quota.ResourceCacheStorageBytes, 100)) if err == nil { t.Fatal("expected error on failed send") }
client.OnDisconnect()
client.SetSendFn(func(msg *millproto.Message) error { go client.HandleResponse(msg.GetQuotaReq().GetRequestId(), &millv1.QuotaResponse{ RequestId: msg.GetQuotaReq().GetRequestId(), ReservationId: "res-reconnected", Allowed: true, }) return nil })
res, err := store.Reserve(context.Background(), validReserveRequest(quota.KindNixCache, "key-1", quota.ResourceCacheStorageBytes, 100)) if err != nil { t.Fatalf("Reserve() after reconnect error = %v", err) } if res.ID != "res-reconnected" { t.Errorf("res.ID = %q, want %q", res.ID, "res-reconnected") } })
t.Run("response vs disconnect race", func(t *testing.T) { for range 20 { client := NewQuotaClient() store := client.ForLease("lease-1")
client.SetSendFn(func(msg *millproto.Message) error { return nil })
done := make(chan error, 1) go func() { _, err := store.Reserve(context.Background(), validReserveRequest(quota.KindNixCache, "key-1", quota.ResourceCacheStorageBytes, 100)) done <- err }()
time.Sleep(time.Duration(rand.IntN(5)) * time.Millisecond) client.OnDisconnect()
select { case err := <-done: if err == nil { t.Fatal("expected an error on disconnect") } case <-time.After(2 * time.Second): t.Fatal("timed out waiting for Reserve() during disconnect race") } } })}
func TestQuotaClientContextCancellation(t *testing.T) { client := NewQuotaClient() store := client.ForLease("lease-1") sent := make(chan struct{})
client.SetSendFn(func(msg *millproto.Message) error { close(sent) return nil })
ctx, cancel := context.WithCancel(context.Background()) done := make(chan error, 1)
go func() { _, err := store.Reserve(ctx, validReserveRequest(quota.KindNixCache, "key-1", quota.ResourceCacheStorageBytes, 100)) done <- err }()
<-sent cancel()
select { case err := <-done: if !errors.Is(err, context.Canceled) { t.Fatalf("expected context.Canceled, got %v", err) } case <-time.After(2 * time.Second): t.Fatal("timed out waiting for Reserve() to return on context cancellation") }
client.mu.Lock() pendingCount := len(client.pending) client.mu.Unlock() if pendingCount != 0 { t.Errorf("pending map has %d entries, want 0", pendingCount) }}
func TestQuotaClientConcurrentCorrelations(t *testing.T) { client := NewQuotaClient() store := client.ForLease("lease-1")
client.SetSendFn(func(msg *millproto.Message) error { req := msg.GetQuotaReq() go func(reqID, key string) { time.Sleep(time.Duration(rand.IntN(10)) * time.Millisecond) client.HandleResponse(reqID, &millv1.QuotaResponse{ RequestId: reqID, ReservationId: "res-" + key, Allowed: true, Reason: "within_limits", Resource: quota.ResourceCacheStorageBytes, }) }(req.GetRequestId(), req.GetKey()) return nil })
const count = 20 var wg sync.WaitGroup wg.Add(count)
for i := range count { go func(idx int) { defer wg.Done() key := fmt.Sprintf("key-%d", idx) res, err := store.Reserve(context.Background(), validReserveRequest(quota.KindNixCache, key, quota.ResourceCacheStorageBytes, int64(idx+1)*100)) if err != nil { t.Errorf("goroutine %d Reserve() error: %v", idx, err) return } expectedID := "res-" + key if res.ID != expectedID { t.Errorf("goroutine %d res.ID = %q, want %q", idx, res.ID, expectedID) } }(i) }
wg.Wait()}