diff --git a/cmd/spindle-microvm-run/main_linux.go b/cmd/spindle-microvm-run/main_linux.go index 642136d8..414bc78b 100644 --- a/cmd/spindle-microvm-run/main_linux.go +++ b/cmd/spindle-microvm-run/main_linux.go @@ -208,7 +208,7 @@ func runMicroVMRunDev(ctx context.Context, cmd *cli.Command) error { var uploadCache *microvm.UploadCacheProxy if cmd.String("cache-upload-url") != "" { var err error - uploadCache, err = microvm.StartUploadCacheProxy(ctx, vm.CID(), cmd.String("cache-upload-url"), upstreams, filepath.Join(vm.WorkDir(), "upload-cache"), logger) + uploadCache, err = microvm.StartUploadCacheProxy(ctx, vm.CID(), cmd.String("cache-upload-url"), upstreams, filepath.Join(vm.WorkDir(), "upload-cache"), logger, nil, "", "") if err != nil { return fmt.Errorf("start upload cache proxy: %w", err) } diff --git a/spindle/db/db.go b/spindle/db/db.go index 34ca9d5b..764a98eb 100644 --- a/spindle/db/db.go +++ b/spindle/db/db.go @@ -144,14 +144,17 @@ func Make(ctx context.Context, dbPath string) (*DB, error) { ); create table if not exists mill_leases ( - lease_id text primary key, - node_id text not null, - epoch text not null, - engine text not null, - knot text not null, - rkey text not null, - workflow text not null, - state text not null + lease_id text primary key, + node_id text not null, + epoch text not null, + engine text not null, + knot text not null, + rkey text not null, + workflow text not null, + state text not null, + quota_reservation_id text, + owner_did text, + repo_did text ); create table if not exists mill_executor_cursors ( @@ -560,6 +563,28 @@ func runMigrations(_ context.Context, conn *sql.Conn, logger *slog.Logger) error return err } + // persist charged subjects so restarts can distinguish live reservations + if err := orm.RunMigration(conn, logger, "mill-leases-quota-column", func(tx *sql.Tx) error { + for _, column := range []string{"quota_reservation_id", "owner_did", "repo_did"} { + var present int + if err := tx.QueryRow( + `select count(*) from pragma_table_info('mill_leases') where name = ?`, + column, + ).Scan(&present); err != nil { + return err + } + if present != 0 { + continue + } + if _, err := tx.Exec(`alter table mill_leases add column ` + column + ` text`); err != nil { + return err + } + } + return nil + }); err != nil { + return err + } + return nil } @@ -629,3 +654,7 @@ func (d *DB) GetLastTimeUs() (int64, error) { err := row.Scan(&lastTimeUs) return lastTimeUs, err } +func (d *DB) GetRepoOwnerAndDid(knot, rkey string) (owner string, repoDid string, err error) { + err = d.QueryRow("SELECT owner, repo_did FROM repos WHERE knot = ? AND rkey = ?", knot, rkey).Scan(&owner, &repoDid) + return +} diff --git a/spindle/db/mill_state.go b/spindle/db/mill_state.go index 1524e2aa..1cfe5014 100644 --- a/spindle/db/mill_state.go +++ b/spindle/db/mill_state.go @@ -8,16 +8,19 @@ import ( "tangled.org/core/notifier" ) -// enough to rebuild the fencing token and workflow identity after a restart +// recovery uses persisted identity, fencing and quota state type MillLease struct { - LeaseID string - NodeID string - Epoch string - Engine string - Knot string - Rkey string - Workflow string - State string + LeaseID string + NodeID string + Epoch string + Engine string + Knot string + Rkey string + Workflow string + State string + QuotaReservationID string + OwnerDID string + RepoDID string } type ExecutorCursor struct { @@ -41,10 +44,16 @@ type OutboxDeletion struct { func (d *DB) SaveMillLease(l MillLease) error { _, err := d.Exec( `insert into mill_leases ( - lease_id, node_id, epoch, engine, knot, rkey, workflow, state - ) values (?, ?, ?, ?, ?, ?, ?, ?) - on conflict(lease_id) do update set state = excluded.state`, + lease_id, node_id, epoch, engine, knot, rkey, workflow, state, + quota_reservation_id, owner_did, repo_did + ) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + on conflict(lease_id) do update set + state = excluded.state, + quota_reservation_id = excluded.quota_reservation_id, + owner_did = excluded.owner_did, + repo_did = excluded.repo_did`, l.LeaseID, l.NodeID, l.Epoch, l.Engine, l.Knot, l.Rkey, l.Workflow, l.State, + l.QuotaReservationID, l.OwnerDID, l.RepoDID, ) return err } @@ -56,7 +65,8 @@ func (d *DB) DeleteMillLease(leaseID string) error { func (d *DB) ListMillLeases() ([]MillLease, error) { rows, err := d.Query(` - select lease_id, node_id, epoch, engine, knot, rkey, workflow, state + select lease_id, node_id, epoch, engine, knot, rkey, workflow, state, + coalesce(quota_reservation_id, ''), coalesce(owner_did, ''), coalesce(repo_did, '') from mill_leases `) if err != nil { @@ -69,6 +79,7 @@ func (d *DB) ListMillLeases() ([]MillLease, error) { var l MillLease if err := rows.Scan( &l.LeaseID, &l.NodeID, &l.Epoch, &l.Engine, &l.Knot, &l.Rkey, &l.Workflow, &l.State, + &l.QuotaReservationID, &l.OwnerDID, &l.RepoDID, ); err != nil { return nil, err } diff --git a/spindle/db/quota.go b/spindle/db/quota.go index f0b80b32..d1d65c29 100644 --- a/spindle/db/quota.go +++ b/spindle/db/quota.go @@ -395,7 +395,7 @@ func (d *QuotaStore) Reserve(ctx context.Context, req quota.ReserveRequest) (quo return res, nil } -func (d *QuotaStore) BeginPublish(ctx context.Context, reservationID string) error { +func (d *QuotaStore) BeginCommit(ctx context.Context, reservationID string) error { if reservationID == "" { return nil } diff --git a/spindle/db/quota_test.go b/spindle/db/quota_test.go index f3cc7677..2edffc1e 100644 --- a/spindle/db/quota_test.go +++ b/spindle/db/quota_test.go @@ -42,7 +42,7 @@ func TestRepoDedup(t *testing.T) { t.Fatalf("expected allowed reservation, got %v", res) } - err = qs.BeginPublish(ctx, res.ID) + err = qs.BeginCommit(ctx, res.ID) if err != nil { t.Fatal(err) } @@ -117,7 +117,7 @@ func TestOwnerUnionDedup(t *testing.T) { t.Fatal("expected res1 to be allowed") } - if err := qs.BeginPublish(ctx, res1.ID); err != nil { + if err := qs.BeginCommit(ctx, res1.ID); err != nil { t.Fatal(err) } if err := qs.Commit(ctx, res1.ID); err != nil { @@ -377,10 +377,10 @@ func TestIdempotence(t *testing.T) { t.Fatalf("expected identical reservation IDs, got %q and %q", res1.ID, res2.ID) } - if err := qs.BeginPublish(ctx, res1.ID); err != nil { + if err := qs.BeginCommit(ctx, res1.ID); err != nil { t.Fatal(err) } - if err := qs.BeginPublish(ctx, res1.ID); err != nil { + if err := qs.BeginCommit(ctx, res1.ID); err != nil { t.Fatal(err) } @@ -398,7 +398,7 @@ func TestIdempotence(t *testing.T) { t.Fatal(err) } - if err := qs.BeginPublish(ctx, ""); err != nil { + if err := qs.BeginCommit(ctx, ""); err != nil { t.Fatal(err) } if err := qs.Commit(ctx, ""); err != nil { @@ -455,7 +455,7 @@ func TestCommitUpdatesChangedAmountForSameKey(t *testing.T) { if !res.Allowed || res.ID == "" { t.Fatalf("reservation = %+v, want a new allowed reservation", res) } - if err := qs.BeginPublish(ctx, res.ID); err != nil { + if err := qs.BeginCommit(ctx, res.ID); err != nil { t.Fatal(err) } if err := qs.Commit(ctx, res.ID); err != nil { @@ -501,7 +501,7 @@ func TestRecoverUpdatesChangedAmountForPublishingKey(t *testing.T) { if !res.Allowed || res.ID == "" { t.Fatalf("reservation = %+v, want a new allowed reservation", res) } - if err := qs.BeginPublish(ctx, res.ID); err != nil { + if err := qs.BeginCommit(ctx, res.ID); err != nil { t.Fatal(err) } return res @@ -572,7 +572,7 @@ func TestRecoveryPhases(t *testing.T) { if err != nil { t.Fatal(err) } - if err := qs.BeginPublish(ctx, res2.ID); err != nil { + if err := qs.BeginCommit(ctx, res2.ID); err != nil { t.Fatal(err) } @@ -862,7 +862,7 @@ func TestUnequalOwnerClaimsRejected(t *testing.T) { t.Fatalf("expected allowed reservation, got %v", res1) } - if err := qs.BeginPublish(ctx, res1.ID); err != nil { + if err := qs.BeginCommit(ctx, res1.ID); err != nil { t.Fatal(err) } if err := qs.Commit(ctx, res1.ID); err != nil { @@ -947,7 +947,7 @@ func TestUnequalOwnerClaimsCharged(t *testing.T) { t.Fatalf("expected allowed reservation, got %v", res1) } - if err := qs.BeginPublish(ctx, res1.ID); err != nil { + if err := qs.BeginCommit(ctx, res1.ID); err != nil { t.Fatal(err) } if err := qs.Commit(ctx, res1.ID); err != nil { @@ -969,7 +969,7 @@ func TestUnequalOwnerClaimsCharged(t *testing.T) { t.Fatalf("expected allowed unequal claim, got %v", res2) } - if err := qs.BeginPublish(ctx, res2.ID); err != nil { + if err := qs.BeginCommit(ctx, res2.ID); err != nil { t.Fatal(err) } if err := qs.Commit(ctx, res2.ID); err != nil { diff --git a/spindle/engine/engine.go b/spindle/engine/engine.go index 4f38c4f1..ad31f84c 100644 --- a/spindle/engine/engine.go +++ b/spindle/engine/engine.go @@ -22,6 +22,7 @@ import ( "tangled.org/core/spindle/db" "tangled.org/core/spindle/models" "tangled.org/core/spindle/observability" + "tangled.org/core/spindle/quota" "tangled.org/core/spindle/secrets" ) @@ -55,19 +56,19 @@ func writeWfError(db *db.DB, n *notifier.Notifier, l *slog.Logger, wfCtx context l = l.With("wid", wid, "phase", phase) switch { case isCanceled(wfCtx): - l.Info("workflow canceled") + l.InfoContext(wfCtx, "workflow canceled") if dbErr := db.StatusCancelled(wid, "User canceled the workflow", -1, n); dbErr != nil { - l.Error("failed to set workflow status to cancelled", "err", dbErr) + l.ErrorContext(wfCtx, "failed to set workflow status to cancelled", "err", dbErr) } case errors.Is(err, ErrTimedOut) || errors.Is(wfCtx.Err(), context.DeadlineExceeded): - l.Info("workflow timed out") + l.InfoContext(wfCtx, "workflow timed out") if dbErr := db.StatusTimeout(wid, n); dbErr != nil { - l.Error("failed to set workflow status to timeout", "err", dbErr) + l.ErrorContext(wfCtx, "failed to set workflow status to timeout", "err", dbErr) } default: - l.Error("workflow failed", "err", err) + l.ErrorContext(wfCtx, "workflow failed", "err", err) if dbErr := db.StatusFailed(wid, err.Error(), -1, n); dbErr != nil { - l.Error("failed to set workflow status to failed", "err", dbErr) + l.ErrorContext(wfCtx, "failed to set workflow status to failed", "err", dbErr) } } } @@ -83,6 +84,13 @@ type RemoteStatusEngine interface { AuthorsRemoteStatus() } +type WorkflowQuotaReporter interface { + QuotaResources(wf *models.Workflow) quota.Resources +} + +type WorkflowQuotaStoreBinder interface { + BindWorkflowQuotaStore(wf *models.Workflow, store quota.ReservationStore) error +} type metricEngineNamer interface { MetricEngineName() string } @@ -142,8 +150,7 @@ func reportWorkflowStatusError(l *slog.Logger, database *db.DB, n *notifier.Noti } } -func StartWorkflows(l *slog.Logger, vault secrets.Manager, cfg *config.Config, stores *artifactstore.Stores, db *db.DB, n *notifier.Notifier, ctx context.Context, pipeline *models.Pipeline, pipelineId models.PipelineId) { - l.Info("starting all workflows in parallel", "pipeline", pipelineId) +func StartWorkflows(l *slog.Logger, vault secrets.Manager, cfg *config.Config, qm *quota.Manager, stores *artifactstore.Stores, db *db.DB, n *notifier.Notifier, ctx context.Context, pipeline *models.Pipeline, pipelineId models.PipelineId) { var allSecrets []secrets.UnlockedSecret // never pass secrets to pipelines that run untrusted (e.g. fork) code @@ -206,14 +213,29 @@ func StartWorkflows(l *slog.Logger, vault secrets.Manager, cfg *config.Config, s ) if st, err := db.GetStatus(wid); err == nil && models.StatusKind(st.Status).IsFinish() { - l.Info("skipping finished workflow", "wid", wid, "status", st.Status) + wl.Info("skipping finished workflow", "wid", wid, "status", st.Status) return } + + requestedResources := quota.Resources{} + _, hasQuotaReporter := eng.(WorkflowQuotaReporter) + if reporter, ok := eng.(WorkflowQuotaReporter); ok { + for resource, amount := range reporter.QuotaResources(&w) { + if amount != 0 { + requestedResources[resource] = amount + } + } + } + reqWorkflows := requestedResources[quota.ResourceWorkflows] + reqMemoryMiB := requestedResources[quota.ResourceMemoryMiB] + reqVCPUs := requestedResources[quota.ResourceVCPUs] + reqDiskMiB := requestedResources[quota.ResourceDiskMiB] + reqCacheBytes := requestedResources[quota.ResourceCacheStorageBytes] + var err error var wfCtx context.Context = ctx engName := engineName(eng) - // start span for workflow execution var span trace.Span wfCtx, span = observability.Tracer().Start(wfCtx, "workflow.run", trace.WithAttributes( attribute.String(observability.WorkflowEngineKey, engName), @@ -229,6 +251,7 @@ func StartWorkflows(l *slog.Logger, vault secrets.Manager, cfg *config.Config, s if repoDID != "" { attrs = append(attrs, attribute.String(observability.RepoDIDKey, repoDID)) } + attrs = append(attrs, observability.RequestedResourceAttrs(reqWorkflows, reqVCPUs, reqMemoryMiB, reqDiskMiB, reqCacheBytes)...) span.SetAttributes(attrs...) } defer span.End() @@ -281,6 +304,11 @@ func StartWorkflows(l *slog.Logger, vault secrets.Manager, cfg *config.Config, s wl.InfoContext(wfCtx, "workflow finished", "result", resStr, "duration_seconds", duration.Seconds(), + "requested_workflows", reqWorkflows, + "requested_vcpus", reqVCPUs, + "requested_memory_mib", reqMemoryMiB, + "requested_disk_mib", reqDiskMiB, + "requested_cache_bytes", reqCacheBytes, "actual_cpu_usec", cpuUsec, "actual_memory_current_bytes", memoryCurrent, "actual_memory_peak_bytes", memoryPeak, @@ -318,20 +346,20 @@ func StartWorkflows(l *slog.Logger, vault secrets.Manager, cfg *config.Config, s if p, ok := eng.(workflowLoggerProvider); ok { wfLogger = p.WorkflowLogger(wid) } else if fileLogger, err := models.NewFileWorkflowLogger(cfg.Server.LogDir, wid, secretValues); err != nil { - l.Warn("failed to setup step logger; logs will not be persisted", "error", err) + wl.WarnContext(wfCtx, "failed to setup step logger; logs will not be persisted", "error", err) wfLogger = models.NullLogger{} } else { - l.Info("setup step logger; logs will be persisted", "logDir", cfg.Server.LogDir, "wid", wid) + wl.InfoContext(wfCtx, "setup step logger; logs will be persisted", "logDir", cfg.Server.LogDir, "wid", wid) wfLogger = fileLogger var closeOnce sync.Once closeLog = func() { closeOnce.Do(func() { if err := fileLogger.Close(); err != nil { - l.Error("failed to close workflow log", "wid", wid, "err", err) + wl.ErrorContext(wfCtx, "failed to close workflow log", "wid", wid, "err", err) } }) } - defer archiveWorkflowLog(l, stores, db, cfg.Server.LogDir, wid) + defer archiveWorkflowLog(wl, stores, db, cfg.Server.LogDir, wid) defer closeLog() } @@ -352,14 +380,14 @@ func StartWorkflows(l *slog.Logger, vault secrets.Manager, cfg *config.Config, s activeMu.Unlock() }() - l.Info("waiting for slot", "wid", wid) slot := WorkflowSlot(NoopSlot{}) var publishTerminalStatus func() + var quotaLease quota.Lease destroyWorkflow := false slotAcquired := false setTerminalError := func(phase string, workflowErr error) { publishTerminalStatus = func() { - writeWfError(db, n, l, wfCtx, wid, phase, workflowErr) + writeWfError(db, n, wl, wfCtx, wid, phase, workflowErr) } } defer func() { @@ -369,14 +397,38 @@ func StartWorkflows(l *slog.Logger, vault secrets.Manager, cfg *config.Config, s } if destroyWorkflow { if err := eng.DestroyWorkflow(ctx, wid); err != nil { - l.Error("failed to destroy workflow", "wid", wid, "err", err) + wl.ErrorContext(wfCtx, "failed to destroy workflow", "wid", wid, "err", err) } } if slotAcquired { slot.Release() } + if quotaLease != nil { + quotaLease.Release() + } }() + if qm != nil && hasQuotaReporter { + resID := quota.WorkflowReservationID(w.RunID, w.OwnerDID, w.RepoDID, wid.Knot, wid.Rkey, wid.Name) + req := quota.ReserveRequest{ + ID: resID, + Kind: quota.KindWorkflow, + Key: resID, + Identity: quota.Identity{ + OwnerDID: w.OwnerDID, + RepoDID: w.RepoDID, + }, + Resources: requestedResources, + } + quotaLease, err = qm.Acquire(wfCtx, req) + if err != nil { + setTerminalError("acquiring quota", err) + return + } + } + + wl.InfoContext(wfCtx, "waiting for slot", "wid", wid) + if s, ok := eng.(WorkflowSlotter); ok { slot, err = s.AcquireWorkflowSlot(wfCtx, wid, &w, Wait) if err != nil { @@ -389,11 +441,20 @@ func StartWorkflows(l *slog.Logger, vault secrets.Manager, cfg *config.Config, s if !remoteStatus { err = db.StatusRunning(wid, n) if err != nil { - l.Error("failed to set workflow status to running", "wid", wid, "err", err) + wl.ErrorContext(wfCtx, "failed to set workflow status to running", "wid", wid, "err", err) return } } + wl.InfoContext(wfCtx, "workflow started", + "engine", engName, + "requested_workflows", reqWorkflows, + "requested_vcpus", reqVCPUs, + "requested_memory_mib", reqMemoryMiB, + "requested_disk_mib", reqDiskMiB, + "requested_cache_bytes", reqCacheBytes, + ) + err = eng.SetupWorkflow(wfCtx, wid, &w, wfLogger) if err != nil { destroyWorkflow = !isCanceled(wfCtx) @@ -475,7 +536,7 @@ func StartWorkflows(l *slog.Logger, vault secrets.Manager, cfg *config.Config, s if !remoteStatus { publishTerminalStatus = func() { if err := db.StatusSuccess(wid, n); err != nil { - l.Error("failed to set workflow status to success", "wid", wid, "err", err) + wl.ErrorContext(wfCtx, "failed to set workflow status to success", "wid", wid, "err", err) } } } diff --git a/spindle/engine/engine_quota_test.go b/spindle/engine/engine_quota_test.go new file mode 100644 index 00000000..0c122340 --- /dev/null +++ b/spindle/engine/engine_quota_test.go @@ -0,0 +1,258 @@ +package engine + +import ( + "context" + "errors" + "io" + "log/slog" + "sync" + "testing" + "time" + + "tangled.org/core/spindle/config" + "tangled.org/core/spindle/models" + "tangled.org/core/spindle/observability" + "tangled.org/core/spindle/quota" +) + +type spyStore struct { + mu sync.Mutex + reserveCalls []quota.ReserveRequest + releaseCalls []string + reserveErr error +} + +func (s *spyStore) Reserve(ctx context.Context, req quota.ReserveRequest) (quota.Reservation, error) { + s.mu.Lock() + defer s.mu.Unlock() + s.reserveCalls = append(s.reserveCalls, req) + if s.reserveErr != nil { + return quota.Reservation{}, s.reserveErr + } + return quota.Reservation{ + ID: "res-123", + Allowed: true, + Temporary: false, + }, nil +} + +func (s *spyStore) BeginCommit(ctx context.Context, id string) error { return nil } +func (s *spyStore) Commit(ctx context.Context, id string) error { return nil } +func (s *spyStore) Release(ctx context.Context, id string) error { + s.mu.Lock() + defer s.mu.Unlock() + s.releaseCalls = append(s.releaseCalls, id) + return nil +} + +func (s *spyStore) ListLimits(ctx context.Context) ([]quota.Limit, error) { return nil, nil } +func (s *spyStore) ListUsage(ctx context.Context) ([]quota.Usage, error) { return nil, nil } +func (s *spyStore) MetricsSnapshot(ctx context.Context) (quota.MetricsSnapshot, error) { + return quota.MetricsSnapshot{}, nil +} +func (s *spyStore) Recover(ctx context.Context, liveIDs []string) error { return nil } +func (s *spyStore) SetLimit(ctx context.Context, scope quota.Scope, did string, resource string, limit int64) error { + return nil +} +func (s *spyStore) UnsetLimit(ctx context.Context, scope quota.Scope, did string, resource string) error { + return nil +} + +type quotaReporterEngine struct { + *mockEngine + workflows, memory, vcpus, disk int64 +} + +func (e *quotaReporterEngine) QuotaResources(wf *models.Workflow) quota.Resources { + return quota.Resources{ + quota.ResourceWorkflows: e.workflows, + quota.ResourceMemoryMiB: e.memory, + quota.ResourceVCPUs: e.vcpus, + quota.ResourceDiskMiB: e.disk, + } +} + +type noReporterEngine struct { + *mockEngine +} + +func TestStartWorkflows_NoDoubleAcquisitionByNoReporterEngine(t *testing.T) { + testDB := newTestDB(t) + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + cfg := &config.Config{Server: config.Server{LogDir: t.TempDir()}} + + store := &spyStore{} + qm := quota.NewManager(store, 50*time.Millisecond, nil) + defer qm.Close() + + eng := &noReporterEngine{mockEngine: &mockEngine{}} + + wf := models.Workflow{ + Name: "job1", + Steps: []models.Step{mockStep{name: "step1"}}, + OwnerDID: "did:web:alice", + RepoDID: "did:web:alice/repo", + } + + pipeline := &models.Pipeline{ + Workflows: map[models.Engine][]models.Workflow{ + eng: {wf}, + }, + } + pipelineId := models.PipelineId{Knot: "knot", Rkey: "rkey"} + + StartWorkflows(logger, nil, cfg, qm, nil, testDB, nil, context.Background(), pipeline, pipelineId) + + store.mu.Lock() + resCount := len(store.reserveCalls) + store.mu.Unlock() + + if resCount != 0 { + t.Fatalf("expected 0 quota reservations for engine without reporter, got %d", resCount) + } +} + +func TestStartWorkflows_QuotaAcquisitionAndIdentity(t *testing.T) { + testDB := newTestDB(t) + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + cfg := &config.Config{Server: config.Server{LogDir: t.TempDir()}} + + store := &spyStore{} + qm := quota.NewManager(store, 50*time.Millisecond, nil) + defer qm.Close() + + eng := "aReporterEngine{ + mockEngine: &mockEngine{}, + workflows: 1, + memory: 256, + vcpus: 2, + disk: 512, + } + + wf := models.Workflow{ + Name: "job1", + Steps: []models.Step{mockStep{name: "step1"}}, + OwnerDID: "did:web:alice", + RepoDID: "did:web:alice/repo", + } + + pipeline := &models.Pipeline{ + Workflows: map[models.Engine][]models.Workflow{ + eng: {wf}, + }, + } + pipelineId := models.PipelineId{Knot: "knot", Rkey: "rkey"} + + StartWorkflows(logger, nil, cfg, qm, nil, testDB, nil, context.Background(), pipeline, pipelineId) + + store.mu.Lock() + resCount := len(store.reserveCalls) + store.mu.Unlock() + + if resCount != 1 { + t.Fatalf("expected 1 quota reservation call, got %d", resCount) + } + + req := store.reserveCalls[0] + if req.Identity.OwnerDID != "did:web:alice" || req.Identity.RepoDID != "did:web:alice/repo" { + t.Errorf("unexpected identity in reservation: %+v", req.Identity) + } + + expectedID := quota.WorkflowReservationID("", "did:web:alice", "did:web:alice/repo", "knot", "rkey", "job1") + if req.ID != expectedID { + t.Errorf("expected reservation ID %q, got %q", expectedID, req.ID) + } + if req.Key != expectedID { + t.Errorf("expected reservation key %q, got %q", expectedID, req.Key) + } + + if req.Resources[quota.ResourceWorkflows] != 1 || + req.Resources[quota.ResourceMemoryMiB] != 256 || + req.Resources[quota.ResourceVCPUs] != 2 || + req.Resources[quota.ResourceDiskMiB] != 512 { + t.Errorf("unexpected resource vectors: %+v", req.Resources) + } + + engNoDisk := "aReporterEngine{ + mockEngine: &mockEngine{}, + workflows: 1, + memory: 256, + vcpus: 0, + disk: 0, + } + store.mu.Lock() + store.reserveCalls = nil + store.mu.Unlock() + + wf2 := wf + wf2.Name = "job2" + + pipelineNoDisk := &models.Pipeline{ + Workflows: map[models.Engine][]models.Workflow{ + engNoDisk: {wf2}, + }, + } + StartWorkflows(logger, nil, cfg, qm, nil, testDB, nil, context.Background(), pipelineNoDisk, pipelineId) + + store.mu.Lock() + resCount2 := len(store.reserveCalls) + store.mu.Unlock() + + if resCount2 != 1 { + t.Fatalf("expected 1 quota reservation call, got %d", resCount2) + } + req2 := store.reserveCalls[0] + if _, exists := req2.Resources[quota.ResourceVCPUs]; exists { + t.Error("vcpus = 0 should be omitted from resources map") + } + if _, exists := req2.Resources[quota.ResourceDiskMiB]; exists { + t.Error("disk_mib = 0 should be omitted from resources map") + } +} + +func TestStartWorkflows_QuotaFailureRecordsWorkflowFailure(t *testing.T) { + testDB := newTestDB(t) + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + cfg := &config.Config{Server: config.Server{LogDir: t.TempDir()}} + + store := &spyStore{reserveErr: errors.New("quota store unavailable")} + qm := quota.NewManager(store, 50*time.Millisecond, nil) + defer qm.Close() + + eng := "aReporterEngine{ + mockEngine: &mockEngine{}, + workflows: 1, + } + pipeline := &models.Pipeline{ + RepoDid: "did:web:alice/repo", + Workflows: map[models.Engine][]models.Workflow{ + eng: {{ + Name: "job1", + Steps: []models.Step{mockStep{name: "step1"}}, + OwnerDID: "did:web:alice", + RepoDID: "did:web:alice/repo", + }}, + }, + } + metrics := observability.NewMetrics() + ctx := observability.WithMetrics(context.Background(), metrics) + StartWorkflows(logger, nil, cfg, qm, nil, testDB, nil, ctx, pipeline, models.PipelineId{Knot: "knot", Rkey: "rkey"}) + + families, err := metrics.Registry().Gather() + if err != nil { + t.Fatal(err) + } + for _, family := range families { + if family.GetName() != "spindle_workflows_total" { + continue + } + for _, metric := range family.GetMetric() { + for _, label := range metric.GetLabel() { + if label.GetName() == "result" && label.GetValue() == "failure" && metric.GetCounter().GetValue() == 1 { + return + } + } + } + } + t.Fatal("quota acquisition failure was not recorded as a failed workflow") +} diff --git a/spindle/engine/engine_test.go b/spindle/engine/engine_test.go index 66a8fc7d..de7ce8c4 100644 --- a/spindle/engine/engine_test.go +++ b/spindle/engine/engine_test.go @@ -12,10 +12,15 @@ import ( "time" "github.com/bluesky-social/indigo/atproto/syntax" + "go.opentelemetry.io/otel" + sdktrace "go.opentelemetry.io/otel/sdk/trace" + "go.opentelemetry.io/otel/sdk/trace/tracetest" "tangled.org/core/api/tangled" "tangled.org/core/spindle/config" "tangled.org/core/spindle/db" "tangled.org/core/spindle/models" + "tangled.org/core/spindle/observability" + "tangled.org/core/spindle/quota" "tangled.org/core/spindle/secrets" ) @@ -39,7 +44,18 @@ type mockEngine struct { } func (m *mockEngine) InitWorkflow(twf tangled.Pipeline_Workflow, tpl tangled.Pipeline) (*models.Workflow, error) { - return &models.Workflow{}, nil + owner := "did:web:alice" + repo := "did:web:alice/repo" + if md := tpl.TriggerMetadata; md != nil && md.Repo != nil { + owner = md.Repo.Did + if md.Repo.RepoDid != nil { + repo = *md.Repo.RepoDid + } + } + return &models.Workflow{ + OwnerDID: owner, + RepoDID: repo, + }, nil } func (m *mockEngine) SetupWorkflow(ctx context.Context, wid models.WorkflowId, wf *models.Workflow, wfLogger models.WorkflowLogger) error { @@ -121,7 +137,7 @@ func TestStartWorkflows_CollisionRejection(t *testing.T) { } cfg := &config.Config{Server: config.Server{LogDir: t.TempDir()}} - StartWorkflows(logger, nil, cfg, nil, testDB, nil, context.Background(), pipeline, pipelineId) + StartWorkflows(logger, nil, cfg, nil, nil, testDB, nil, context.Background(), pipeline, pipelineId) eng.mu.Lock() setupCalls := append([]models.WorkflowId(nil), eng.setupCalls...) @@ -202,7 +218,7 @@ func TestCancelWorkflow_NotOverwritten(t *testing.T) { cfg := &config.Config{Server: config.Server{LogDir: t.TempDir()}} doneChan := make(chan struct{}) go func() { - StartWorkflows(logger, nil, cfg, nil, testDB, nil, context.Background(), pipeline, pipelineId) + StartWorkflows(logger, nil, cfg, nil, nil, testDB, nil, context.Background(), pipeline, pipelineId) close(doneChan) }() @@ -289,7 +305,7 @@ func TestStartWorkflows_FlushesLogBeforeTerminalStatus(t *testing.T) { done := make(chan struct{}) go func() { - StartWorkflows(logger, vault, cfg, nil, testDB, nil, context.Background(), pipeline, pipelineId) + StartWorkflows(logger, vault, cfg, nil, nil, testDB, nil, context.Background(), pipeline, pipelineId) close(done) }() @@ -348,7 +364,7 @@ func TestSetupTimeout_ReportsTimeout(t *testing.T) { } cfg := &config.Config{Server: config.Server{LogDir: t.TempDir()}} - StartWorkflows(logger, nil, cfg, nil, testDB, nil, context.Background(), pipeline, pipelineId) + StartWorkflows(logger, nil, cfg, nil, nil, testDB, nil, context.Background(), pipeline, pipelineId) st, err := testDB.GetStatus(wid) if err != nil { @@ -362,3 +378,319 @@ func TestSetupTimeout_ReportsTimeout(t *testing.T) { t.Fatalf("expected no steps to run after setup timeout, got %d", len(eng.runStepCalls)) } } + +func TestStartWorkflows_SpanEnrichment(t *testing.T) { + oldProvider := otel.GetTracerProvider() + recorder := tracetest.NewSpanRecorder() + provider := sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(recorder)) + otel.SetTracerProvider(provider) + t.Cleanup(func() { + _ = provider.Shutdown(context.Background()) + otel.SetTracerProvider(oldProvider) + }) + + testDB := newTestDB(t) + logger := slog.New(slog.NewTextHandler(os.Stderr, nil)) + + eng := &mockEngine{} + pipelineId := models.PipelineId{ + Knot: "test-knot", + Rkey: "test-rkey", + } + + pipeline := &models.Pipeline{ + RepoDid: "did:web:test-repo", + Workflows: map[models.Engine][]models.Workflow{ + eng: { + { + Name: "test-workflow", + Steps: []models.Step{ + mockStep{name: "step-1"}, + }, + OwnerDID: "did:web:test-owner", + RepoDID: "did:web:test-repo", + }, + }, + }, + TrustedSource: false, + } + + cfg := &config.Config{Server: config.Server{LogDir: t.TempDir()}} + StartWorkflows(logger, nil, cfg, nil, nil, testDB, nil, context.Background(), pipeline, pipelineId) + + spans := recorder.Ended() + var workflowRunSpan, stepRunSpan sdktrace.ReadOnlySpan + for _, span := range spans { + if span.Name() == "workflow.run" { + workflowRunSpan = span + } else if span.Name() == "step.run" { + stepRunSpan = span + } + } + + if workflowRunSpan == nil { + t.Fatal("workflow.run span not recorded") + } + if stepRunSpan == nil { + t.Fatal("step.run span not recorded") + } + + workflowAttrs := make(map[string]string) + for _, attr := range workflowRunSpan.Attributes() { + workflowAttrs[string(attr.Key)] = attr.Value.AsString() + } + + wid := models.WorkflowId{ + PipelineId: pipelineId, + Name: "test-workflow", + } + + expectedWfAttrs := map[string]string{ + observability.OwnerDIDKey: "did:web:test-owner", + observability.RepoDIDKey: "did:web:test-repo", + observability.PipelineIDKey: pipelineId.AtUri().String(), + observability.WorkflowIDKey: wid.String(), + } + for k, wantVal := range expectedWfAttrs { + if gotVal := workflowAttrs[k]; gotVal != wantVal { + t.Errorf("workflow attribute %q = %q, want %q", k, gotVal, wantVal) + } + } + + stepAttrs := make(map[string]string) + for _, attr := range stepRunSpan.Attributes() { + stepAttrs[string(attr.Key)] = attr.Value.AsString() + } + + expectedStepAttrs := map[string]string{ + observability.OwnerDIDKey: "did:web:test-owner", + observability.RepoDIDKey: "did:web:test-repo", + observability.PipelineIDKey: pipelineId.AtUri().String(), + observability.WorkflowIDKey: wid.String(), + observability.StepNameKey: "step-1", + } + for k, wantVal := range expectedStepAttrs { + if gotVal := stepAttrs[k]; gotVal != wantVal { + t.Errorf("step attribute %q = %q, want %q", k, gotVal, wantVal) + } + } + hasStepIndex := false + for _, attr := range stepRunSpan.Attributes() { + if string(attr.Key) == observability.StepIndexKey && attr.Value.AsInterface() == int64(0) { + hasStepIndex = true + } + } + if !hasStepIndex { + t.Error("step.run span missing step_index attribute or it is not 0") + } +} + +type mockResourceUsageReporter struct { + *mockEngine + workflows, memory, vcpus, disk int64 + usage WorkflowResourceUsage + usageOk bool +} + +func (m *mockResourceUsageReporter) QuotaResources(wf *models.Workflow) quota.Resources { + return quota.Resources{ + quota.ResourceWorkflows: m.workflows, + quota.ResourceMemoryMiB: m.memory, + quota.ResourceVCPUs: m.vcpus, + quota.ResourceDiskMiB: m.disk, + } +} + +func (m *mockResourceUsageReporter) WorkflowResourceUsage(wf *models.Workflow) (WorkflowResourceUsage, bool) { + return m.usage, m.usageOk +} + +func TestStartWorkflows_ResourceTracking(t *testing.T) { + oldProvider := otel.GetTracerProvider() + recorder := tracetest.NewSpanRecorder() + provider := sdktrace.NewTracerProvider( + sdktrace.WithSpanProcessor(recorder), + sdktrace.WithSampler(sdktrace.AlwaysSample()), + ) + otel.SetTracerProvider(provider) + t.Cleanup(func() { + _ = provider.Shutdown(context.Background()) + otel.SetTracerProvider(oldProvider) + }) + + testDB := newTestDB(t) + + mu := &sync.Mutex{} + records := []slog.Record{} + handler := &captureHandler{ + mu: mu, + records: &records, + } + logger := slog.New(handler) + + eng := &mockResourceUsageReporter{ + mockEngine: &mockEngine{}, + workflows: 2, + memory: 1024, + vcpus: 4, + disk: 512, + usage: WorkflowResourceUsage{ + CPUUsec: 12345, + MemoryCurrentBytes: 2048, + MemoryPeakBytes: 4096, + SwapCurrentBytes: 100, + SwapPeakBytes: 200, + PIDsCurrent: 10, + IOReadBytes: 300, + IOWriteBytes: 400, + IOReadOps: 30, + IOWriteOps: 40, + VolumeAllocatedBytes: 1000, + CgroupAvailable: true, + VolumeAvailable: true, + }, + usageOk: true, + } + + pipelineId := models.PipelineId{ + Knot: "test-knot", + Rkey: "test-rkey", + } + + pipeline := &models.Pipeline{ + RepoDid: "did:web:test-repo", + Workflows: map[models.Engine][]models.Workflow{ + eng: { + { + Name: "test-workflow", + Steps: []models.Step{ + mockStep{name: "step-1"}, + }, + OwnerDID: "did:web:test-owner", + RepoDID: "did:web:test-repo", + }, + }, + }, + TrustedSource: false, + } + + cfg := &config.Config{Server: config.Server{LogDir: t.TempDir()}} + + metrics := observability.NewMetrics() + ctx := observability.WithMetrics(context.Background(), metrics) + + StartWorkflows(logger, nil, cfg, nil, nil, testDB, nil, ctx, pipeline, pipelineId) + + handler.mu.Lock() + records = *handler.records + handler.mu.Unlock() + + var startedLog, finishedLog slog.Record + for _, r := range records { + if r.Message == "workflow started" { + startedLog = r + } else if r.Message == "workflow finished" { + finishedLog = r + } + } + + if startedLog.Message == "" { + t.Fatal("workflow started log not found") + } + if finishedLog.Message == "" { + t.Fatal("workflow finished log not found") + } + + getAttr := func(r slog.Record, key string) any { + var val any + r.Attrs(func(a slog.Attr) bool { + if a.Key == key { + val = a.Value.Any() + return false + } + return true + }) + return val + } + + if getAttr(startedLog, "requested_workflows") != int64(2) { + t.Errorf("started log requested_workflows = %v, want 2", getAttr(startedLog, "requested_workflows")) + } + if getAttr(startedLog, "requested_vcpus") != int64(4) { + t.Errorf("started log requested_vcpus = %v, want 4", getAttr(startedLog, "requested_vcpus")) + } + + if getAttr(finishedLog, "result") != "success" { + t.Errorf("finished log result = %v, want success", getAttr(finishedLog, "result")) + } + if getAttr(finishedLog, "actual_cpu_usec") != uint64(12345) { + t.Errorf("finished log actual_cpu_usec = %v, want 12345", getAttr(finishedLog, "actual_cpu_usec")) + } + if getAttr(finishedLog, "actual_memory_peak_bytes") != uint64(4096) { + t.Errorf("finished log actual_memory_peak_bytes = %v, want 4096", getAttr(finishedLog, "actual_memory_peak_bytes")) + } + + spans := recorder.Ended() + var workflowRunSpan sdktrace.ReadOnlySpan + for _, span := range spans { + if span.Name() == "workflow.run" { + workflowRunSpan = span + break + } + } + + if workflowRunSpan == nil { + t.Fatal("workflow.run span not found") + } + + attrs := make(map[string]any) + for _, attr := range workflowRunSpan.Attributes() { + attrs[string(attr.Key)] = attr.Value.AsInterface() + } + + if attrs[observability.ReqWorkflowsKey] != int64(2) { + t.Errorf("span attr %s = %v, want 2", observability.ReqWorkflowsKey, attrs[observability.ReqWorkflowsKey]) + } + if attrs[observability.ReqVCPUsKey] != int64(4) { + t.Errorf("span attr %s = %v, want 4", observability.ReqVCPUsKey, attrs[observability.ReqVCPUsKey]) + } + + if attrs[observability.ActCPUUsecKey] != int64(12345) { + t.Errorf("span attr %s = %v, want 12345", observability.ActCPUUsecKey, attrs[observability.ActCPUUsecKey]) + } + if attrs[observability.ActMemoryPeakBytesKey] != int64(4096) { + t.Errorf("span attr %s = %v, want 4096", observability.ActMemoryPeakBytesKey, attrs[observability.ActMemoryPeakBytesKey]) + } +} + +type captureHandler struct { + attrs []slog.Attr + records *[]slog.Record + mu *sync.Mutex +} + +func (h *captureHandler) Enabled(ctx context.Context, level slog.Level) bool { + return true +} + +func (h *captureHandler) Handle(ctx context.Context, r slog.Record) error { + h.mu.Lock() + defer h.mu.Unlock() + r.AddAttrs(h.attrs...) + *h.records = append(*h.records, r) + return nil +} + +func (h *captureHandler) WithAttrs(attrs []slog.Attr) slog.Handler { + newAttrs := append([]slog.Attr(nil), h.attrs...) + newAttrs = append(newAttrs, attrs...) + return &captureHandler{ + attrs: newAttrs, + records: h.records, + mu: h.mu, + } +} + +func (h *captureHandler) WithGroup(name string) slog.Handler { + return h +} diff --git a/spindle/engine_microvm_linux.go b/spindle/engine_microvm_linux.go index 5c5551b0..a7115a74 100644 --- a/spindle/engine_microvm_linux.go +++ b/spindle/engine_microvm_linux.go @@ -9,8 +9,9 @@ import ( "tangled.org/core/spindle/db" "tangled.org/core/spindle/engines/microvm" "tangled.org/core/spindle/models" + "tangled.org/core/spindle/quota" ) -func newMicrovmEngine(ctx context.Context, cfg *config.Config, d *db.DB) (models.Engine, error) { - return microvm.New(ctx, cfg, d) +func newMicrovmEngine(ctx context.Context, cfg *config.Config, d *db.DB, qStore quota.ReservationStore) (models.Engine, error) { + return microvm.New(ctx, cfg, d, qStore) } diff --git a/spindle/engine_microvm_other.go b/spindle/engine_microvm_other.go index 9da4c9fc..c6de2e95 100644 --- a/spindle/engine_microvm_other.go +++ b/spindle/engine_microvm_other.go @@ -9,8 +9,9 @@ import ( "tangled.org/core/spindle/config" "tangled.org/core/spindle/db" "tangled.org/core/spindle/models" + "tangled.org/core/spindle/quota" ) -func newMicrovmEngine(context.Context, *config.Config, *db.DB) (models.Engine, error) { +func newMicrovmEngine(ctx context.Context, cfg *config.Config, d *db.DB, qStore quota.ReservationStore) (models.Engine, error) { return nil, fmt.Errorf("microvm engine is only supported on Linux") } diff --git a/spindle/engines/dummy/engine.go b/spindle/engines/dummy/engine.go index eedc4a17..2d300751 100644 --- a/spindle/engines/dummy/engine.go +++ b/spindle/engines/dummy/engine.go @@ -10,6 +10,7 @@ import ( "tangled.org/core/api/tangled" "tangled.org/core/spindle/engine" "tangled.org/core/spindle/models" + "tangled.org/core/spindle/quota" "tangled.org/core/spindle/secrets" ) @@ -35,7 +36,7 @@ func (s Step) Name() string { return s.name } func (s Step) Command() string { return s.command } func (s Step) Kind() models.StepKind { return s.kind } -func (e *DummyEngine) InitWorkflow(twf tangled.Pipeline_Workflow, _ tangled.Pipeline) (*models.Workflow, error) { +func (e *DummyEngine) InitWorkflow(twf tangled.Pipeline_Workflow, tpl tangled.Pipeline) (*models.Workflow, error) { dwf := &struct { Steps []struct { Name string `yaml:"name"` @@ -48,9 +49,26 @@ func (e *DummyEngine) InitWorkflow(twf tangled.Pipeline_Workflow, _ tangled.Pipe return nil, err } + var ownerDID, repoDID string + if md := tpl.TriggerMetadata; md != nil && md.Repo != nil { + ownerDID = md.Repo.Did + repoDID = ownerDID + if md.Repo.RepoDid != nil && *md.Repo.RepoDid != "" { + repoDID = *md.Repo.RepoDid + } + } + if ownerDID == "" { + return nil, fmt.Errorf("missing owner DID in pipeline trigger metadata") + } + if repoDID == "" { + return nil, fmt.Errorf("missing repository DID in pipeline trigger metadata") + } + wf := &models.Workflow{ Name: twf.Name, Environment: dwf.Environment, + OwnerDID: ownerDID, + RepoDID: repoDID, } for _, ds := range dwf.Steps { wf.Steps = append(wf.Steps, Step{ @@ -104,3 +122,6 @@ func (e *DummyEngine) RunStep(ctx context.Context, wid models.WorkflowId, w *mod } return nil } +func (e *DummyEngine) QuotaResources(wf *models.Workflow) quota.Resources { + return quota.Resources{quota.ResourceWorkflows: 1} +} diff --git a/spindle/engines/microvm/budget.go b/spindle/engines/microvm/budget.go index a45d834e..7ea9cbbc 100644 --- a/spindle/engines/microvm/budget.go +++ b/spindle/engines/microvm/budget.go @@ -10,6 +10,7 @@ import ( "tangled.org/core/spindle/config" "tangled.org/core/spindle/engine" "tangled.org/core/spindle/models" + "tangled.org/core/spindle/quota" ) // memory buffer for qemu process / slirp4netns itself @@ -91,3 +92,16 @@ func resourcesForImage(spec ImageSpec) Resources { DiskMiB: diskMiB, } } +func (e *Engine) QuotaResources(wf *models.Workflow) quota.Resources { + state, ok := wf.Data.(*workflowState) + if !ok || state == nil { + return quota.Resources{quota.ResourceWorkflows: 1} + } + req := resourcesForImage(state.ImageSpec) + return quota.Resources{ + quota.ResourceWorkflows: 1, + quota.ResourceMemoryMiB: req.MemoryMiB, + quota.ResourceVCPUs: req.VCPUs, + quota.ResourceDiskMiB: req.DiskMiB, + } +} diff --git a/spindle/engines/microvm/engine.go b/spindle/engines/microvm/engine.go index 46f73a05..e72d83d0 100644 --- a/spindle/engines/microvm/engine.go +++ b/spindle/engines/microvm/engine.go @@ -29,6 +29,7 @@ import ( "tangled.org/core/spindle/engine" "tangled.org/core/spindle/models" "tangled.org/core/spindle/observability" + "tangled.org/core/spindle/quota" "tangled.org/core/spindle/secrets" ) @@ -50,6 +51,7 @@ type Engine struct { l *slog.Logger cfg *config.Config db *db.DB + quotaStore quota.ReservationStore agentMu sync.Mutex agent *agentHub scheduler *engine.ResourceScheduler[Resources] @@ -78,7 +80,7 @@ func (s Step) Name() string { return s.name } func (s Step) Command() string { return s.command } func (s Step) Kind() models.StepKind { return s.kind } -func New(ctx context.Context, cfg *config.Config, d *db.DB) (*Engine, error) { +func New(ctx context.Context, cfg *config.Config, d *db.DB, quotaStore quota.ReservationStore) (*Engine, error) { l := log.FromContext(ctx).With("component", "engine.microvm") budget, max, agingThreshold := newVMBudgetConfig(cfg.MicroVMPipelines) l.Info("initialized microVM workflow budget", "budget", budget.String(), "maxWorkflow", max.String(), "agingThreshold", agingThreshold) @@ -96,6 +98,7 @@ func New(ctx context.Context, cfg *config.Config, d *db.DB) (*Engine, error) { l: l, cfg: cfg, db: d, + quotaStore: quotaStore, scheduler: engine.NewResourceScheduler(budget, max, agingThreshold), cgroupParent: cgroupParent, budget: budget, @@ -201,7 +204,11 @@ func (e *Engine) InitWorkflow(twf tangled.Pipeline_Workflow, tpl tangled.Pipelin // of another tenant's identical config repoDid := "" if md := tpl.TriggerMetadata; md != nil && md.Repo != nil { - repoDid = md.Repo.Did + if md.Repo.RepoDid != nil { + repoDid = *md.Repo.RepoDid + } else { + repoDid = md.Repo.Did + } } var err error configKey, err = buildConfigKey(imageSpec, config, repoDid) @@ -231,6 +238,23 @@ func (e *Engine) InitWorkflow(twf tangled.Pipeline_Workflow, tpl tangled.Pipelin return nil, err } + var ownerDID, repoDID string + if md := tpl.TriggerMetadata; md != nil && md.Repo != nil { + ownerDID = md.Repo.Did + repoDID = ownerDID + if md.Repo.RepoDid != nil && *md.Repo.RepoDid != "" { + repoDID = *md.Repo.RepoDid + } + } + if ownerDID == "" { + return nil, fmt.Errorf("missing owner DID in pipeline trigger metadata") + } + if repoDID == "" { + return nil, fmt.Errorf("missing repository DID in pipeline trigger metadata") + } + swf.OwnerDID = ownerDID + swf.RepoDID = repoDID + swf.Data = &workflowState{ ImageSpec: imageSpec, ImageSpecPath: imageSpecPath, @@ -240,10 +264,25 @@ func (e *Engine) InitWorkflow(twf tangled.Pipeline_Workflow, tpl tangled.Pipelin CacheReadURLs: cacheURLs, CacheTrustedPublicKeys: cacheKeys, NixOSToplevelCache: newNixOSToplevelCacheStore(e.db), + OwnerDID: ownerDID, + QuotaStore: e.quotaStore, + RepoDID: repoDID, } return swf, nil } +func (e *Engine) BindWorkflowQuotaStore(wf *models.Workflow, store quota.ReservationStore) error { + state, ok := wf.Data.(*workflowState) + if !ok || state == nil { + return errors.New("workflow state is not initialized") + } + if store == nil { + return errors.New("workflow quota store is nil") + } + state.QuotaStore = store + return nil +} + func (e *Engine) SetupWorkflow(ctx context.Context, wid models.WorkflowId, wf *models.Workflow, wfLogger models.WorkflowLogger) (err error) { l := e.l.With("workflow", wid) setupStep := Step{name: "microVM setup", kind: models.StepKindSystem} @@ -309,9 +348,14 @@ func (e *Engine) SetupWorkflow(ctx context.Context, wid models.WorkflowId, wf *m return err } state.ReadCache = readCache - + if strings.TrimSpace(e.cfg.NixCache.UploadURL) != "" { + if state.OwnerDID == "" || state.RepoDID == "" { + return fmt.Errorf("missing owner or repo identity for upload cache quota") + } + } stagingDir := filepath.Join(workDir, "upload-cache") - uploadCache, err := StartUploadCacheProxy(ctx, cid, e.cfg.NixCache.UploadURL, upstreams, stagingDir, l) + + uploadCache, err := StartUploadCacheProxy(ctx, cid, e.cfg.NixCache.UploadURL, upstreams, stagingDir, l, state.QuotaStore, state.OwnerDID, state.RepoDID) if err != nil { return err } diff --git a/spindle/engines/microvm/engine_test.go b/spindle/engines/microvm/engine_test.go index 5a5aac7b..8a2f5ad0 100644 --- a/spindle/engines/microvm/engine_test.go +++ b/spindle/engines/microvm/engine_test.go @@ -39,8 +39,16 @@ func testEngine(t *testing.T, imageDir string) *Engine { } } +func testPipelineIdentity() tangled.Pipeline { + return tangled.Pipeline{ + TriggerMetadata: &tangled.Pipeline_TriggerMetadata{ + Repo: &tangled.Pipeline_TriggerRepo{Did: "did:web:owner"}, + }, + } +} + func TestNewDefersAgentHubUntilWorkflowSetup(t *testing.T) { - e, err := New(context.Background(), &config.Config{}, nil) + e, err := New(context.Background(), &config.Config{}, nil, nil) if err != nil { t.Fatal(err) } @@ -63,7 +71,7 @@ steps: - name: hello command: hello `, - }, tangled.Pipeline{}) + }, testPipelineIdentity()) if err == nil { t.Fatal("expected error for NixOS config options on a non-NixOS image") } @@ -84,12 +92,12 @@ steps: - name: hello command: echo hello `, - }, tangled.Pipeline{}) + }, testPipelineIdentity()) if err != nil { t.Fatal(err) } - if len(wf.Steps) != 1 { - t.Fatalf("expected exactly the user step, got %d steps", len(wf.Steps)) + if len(wf.Steps) != 2 { + t.Fatalf("expected checkout and user steps, got %d steps", len(wf.Steps)) } state, ok := wf.Data.(*workflowState) if !ok { @@ -116,14 +124,14 @@ steps: - name: hello command: hello `, - }, tangled.Pipeline{}) + }, testPipelineIdentity()) if err != nil { t.Fatal(err) } - if len(wf.Steps) != 2 { - t.Fatalf("expected activation step + user step, got %d steps", len(wf.Steps)) + if len(wf.Steps) != 3 { + t.Fatalf("expected checkout, activation, and user steps, got %d steps", len(wf.Steps)) } - if step, ok := wf.Steps[0].(Step); !ok || step.action != activationStepAction { - t.Fatalf("first step should be the activation step, got %+v", wf.Steps[0]) + if step, ok := wf.Steps[1].(Step); !ok || step.action != activationStepAction { + t.Fatalf("second step should be the activation step, got %+v", wf.Steps[1]) } } diff --git a/spindle/engines/microvm/image.go b/spindle/engines/microvm/image.go index be2e657c..c2de0de0 100644 --- a/spindle/engines/microvm/image.go +++ b/spindle/engines/microvm/image.go @@ -9,6 +9,8 @@ import ( "os" "path/filepath" "strings" + + "tangled.org/core/spindle/config" ) const imageSpecFileName = "spec.json" @@ -188,10 +190,10 @@ func resolveImageSpecPath(base, path string) string { return filepath.Join(base, path) } -func (e *Engine) resolveImage(name string) (ImageSpec, string, string, error) { +func ResolveImageSpec(pipelinesCfg config.MicroVMPipelines, name string) (ImageSpec, string, string, error) { name = strings.TrimSpace(name) if name == "" { - name = strings.TrimSpace(e.cfg.MicroVMPipelines.DefaultImage) + name = strings.TrimSpace(pipelinesCfg.DefaultImage) } if name == "" { return ImageSpec{}, "", "", fmt.Errorf("no image specified in workflow and SPINDLE_MICROVM_PIPELINES_DEFAULT_IMAGE is not set") @@ -200,7 +202,7 @@ func (e *Engine) resolveImage(name string) (ImageSpec, string, string, error) { return ImageSpec{}, "", "", fmt.Errorf("invalid microVM image name %q: must be a plain name, not a path", name) } - imageDir := strings.TrimSpace(e.cfg.MicroVMPipelines.ImageDir) + imageDir := strings.TrimSpace(pipelinesCfg.ImageDir) if imageDir == "" { return ImageSpec{}, "", "", fmt.Errorf("microVM workflows require SPINDLE_MICROVM_PIPELINES_IMAGE_DIR") } @@ -224,6 +226,10 @@ func (e *Engine) resolveImage(name string) (ImageSpec, string, string, error) { return ImageSpec{}, "", "", fmt.Errorf("microVM image %q was not found; looked in: %s", name, strings.Join(candidates, ", ")) } +func (e *Engine) resolveImage(name string) (ImageSpec, string, string, error) { + return ResolveImageSpec(e.cfg.MicroVMPipelines, name) +} + // check if image name is not a path func isPlainImageName(name string) bool { if name == "" || name == "." || name == ".." { diff --git a/spindle/engines/microvm/upload_cache_proxy.go b/spindle/engines/microvm/upload_cache_proxy.go index cc4bf32e..1331a0f8 100644 --- a/spindle/engines/microvm/upload_cache_proxy.go +++ b/spindle/engines/microvm/upload_cache_proxy.go @@ -14,6 +14,8 @@ import ( "time" "github.com/mdlayher/vsock" + "tangled.org/core/spindle/observability" + "tangled.org/core/spindle/quota" ) type UploadCacheBackend interface { @@ -29,7 +31,7 @@ type UploadCacheProxy struct { backend UploadCacheBackend } -func StartUploadCacheProxy(ctx context.Context, cid uint32, uploadURL string, readUpstreams []CacheUpstream, stagingDir string, logger *slog.Logger) (*UploadCacheProxy, error) { +func StartUploadCacheProxy(ctx context.Context, cid uint32, uploadURL string, readUpstreams []CacheUpstream, stagingDir string, logger *slog.Logger, quotaStore quota.ReservationStore, ownerDID, repoDID string) (*UploadCacheProxy, error) { if strings.TrimSpace(uploadURL) == "" { return nil, nil } @@ -44,6 +46,12 @@ func StartUploadCacheProxy(ctx context.Context, cid uint32, uploadURL string, re return nil, err } + var finalBackend UploadCacheBackend = backend + if quotaStore != nil { + metrics := observability.GetMetrics(ctx) + finalBackend = NewStagingWrapper(backend, stagingDir, quotaStore, ownerDID, repoDID, uploadURL, logger, metrics) + } + ln, port, err := listenRandomVsockUploadPort(ctx) if err != nil { return nil, fmt.Errorf("listen for cache upload proxy: %w", err) @@ -52,10 +60,10 @@ func StartUploadCacheProxy(ctx context.Context, cid uint32, uploadURL string, re proxy := &UploadCacheProxy{ port: port, ln: ln, - backend: backend, + backend: finalBackend, } proxy.server = &http.Server{ - Handler: backend, + Handler: finalBackend, Protocols: cacheProxyProtocols(), ReadHeaderTimeout: 30 * time.Second, } diff --git a/spindle/engines/microvm/upload_quota_test.go b/spindle/engines/microvm/upload_quota_test.go new file mode 100644 index 00000000..89248c50 --- /dev/null +++ b/spindle/engines/microvm/upload_quota_test.go @@ -0,0 +1,1019 @@ +//go:build linux + +package microvm + +import ( + "context" + "encoding/json" + "errors" + "log/slog" + "net/http" + "net/http/httptest" + "net/url" + "os" + "path/filepath" + "strings" + "sync" + "testing" + + "tangled.org/core/api/tangled" + "tangled.org/core/spindle/config" + "tangled.org/core/spindle/observability" + "tangled.org/core/spindle/quota" +) + +const ( + validQuotaNarinfo = "StorePath: /nix/store/00000000000000000000000000000000-abc\nURL: nar/abc.nar.xz\nNarHash: sha256:abc\nNarSize: 1000\n" + validQuotaNarinfoPath = "/00000000000000000000000000000000.narinfo" +) + +type mockQuotaStore struct { + mu sync.Mutex + reserveCalls []quota.ReserveRequest + publishCalls []string + commitCalls []string + releaseCalls []string + allowed bool + reason string + beginCommitErr error + releaseErr error + + ReserveFunc func(ctx context.Context, req quota.ReserveRequest) (quota.Reservation, error) + BeginCommitFunc func(ctx context.Context, reservationID string) error + CommitFunc func(ctx context.Context, reservationID string) error +} + +func newMockQuotaStore(allowed bool, reason string) *mockQuotaStore { + return &mockQuotaStore{ + allowed: allowed, + reason: reason, + } +} + +func (m *mockQuotaStore) Reserve(ctx context.Context, req quota.ReserveRequest) (quota.Reservation, error) { + m.mu.Lock() + if m.ReserveFunc != nil { + f := m.ReserveFunc + m.mu.Unlock() + return f(ctx, req) + } + defer m.mu.Unlock() + m.reserveCalls = append(m.reserveCalls, req) + if err := ctx.Err(); err != nil { + return quota.Reservation{}, err + } + if m.allowed { + return quota.Reservation{ + ID: "res-abc", + Allowed: true, + Reason: quota.ReasonWithinLimit, + }, nil + } + return quota.Reservation{ + ID: "", + Allowed: false, + Reason: m.reason, + }, nil +} + +func (m *mockQuotaStore) BeginCommit(ctx context.Context, reservationID string) error { + m.mu.Lock() + if m.BeginCommitFunc != nil { + f := m.BeginCommitFunc + m.mu.Unlock() + return f(ctx, reservationID) + } + defer m.mu.Unlock() + m.publishCalls = append(m.publishCalls, reservationID) + if err := ctx.Err(); err != nil { + return err + } + return m.beginCommitErr +} + +func (m *mockQuotaStore) Commit(ctx context.Context, reservationID string) error { + m.mu.Lock() + if m.CommitFunc != nil { + f := m.CommitFunc + m.mu.Unlock() + return f(ctx, reservationID) + } + defer m.mu.Unlock() + m.commitCalls = append(m.commitCalls, reservationID) + return nil +} + +func (m *mockQuotaStore) Release(ctx context.Context, reservationID string) error { + m.mu.Lock() + defer m.mu.Unlock() + m.releaseCalls = append(m.releaseCalls, reservationID) + return m.releaseErr +} + +func TestQuotaNoBackendPublicationBeforeNarinfo(t *testing.T) { + staging := t.TempDir() + store := newMockQuotaStore(true, "") + + var httpCalled bool + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + httpCalled = true + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + u, _ := url.Parse(srv.URL) + backend := newHTTPUploadProxyBackend(u, nil, slog.Default()) + wrapper := NewStagingWrapper(backend, staging, store, "owner-1", "repo-1", srv.URL, slog.Default(), nil) + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPut, "/nar/abc.nar.xz", strings.NewReader("some-nar-data")) + wrapper.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", rec.Code) + } + if httpCalled { + t.Fatal("http target was called before narinfo upload") + } +} + +func TestQuotaStagingBytesAreBounded(t *testing.T) { + staging := t.TempDir() + store := newMockQuotaStore(true, "") + backend := newHTTPUploadProxyBackend(&url.URL{Scheme: "http", Host: "127.0.0.1:1"}, nil, slog.Default()) + wrapper := NewStagingWrapper(backend, staging, store, "owner-1", "repo-1", "http://127.0.0.1:1", slog.Default(), nil) + sw := wrapper.(*stagingWrapper) + sw.maxStagedBytes = 5 + + rec := httptest.NewRecorder() + wrapper.ServeHTTP(rec, httptest.NewRequest(http.MethodPut, "/nar/a.nar", strings.NewReader("four"))) + if rec.Code != http.StatusOK { + t.Fatalf("first upload status = %d, want 200", rec.Code) + } + + rec = httptest.NewRecorder() + wrapper.ServeHTTP(rec, httptest.NewRequest(http.MethodPut, "/nar/b.nar", strings.NewReader("two"))) + if rec.Code != http.StatusRequestEntityTooLarge { + t.Fatalf("second upload status = %d, want 413", rec.Code) + } + if sw.stagedBytes != 4 { + t.Fatalf("staged bytes = %d, want 4", sw.stagedBytes) + } +} + +func TestQuotaExactCompressedByteReservation(t *testing.T) { + staging := t.TempDir() + store := newMockQuotaStore(true, "") + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + u, _ := url.Parse(srv.URL) + backend := newHTTPUploadProxyBackend(u, nil, slog.Default()) + wrapper := NewStagingWrapper(backend, staging, store, "owner-1", "repo-1", srv.URL, slog.Default(), nil) + + narData := "compressed-bytes" + expectedSize := int64(len(narData)) + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPut, "/nar/abc.nar.xz", strings.NewReader(narData)) + wrapper.ServeHTTP(rec, req) + + narinfoBody := validQuotaNarinfo + rec = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodPut, validQuotaNarinfoPath, strings.NewReader(narinfoBody)) + wrapper.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", rec.Code) + } + + store.mu.Lock() + defer store.mu.Unlock() + if len(store.reserveCalls) != 1 { + t.Fatalf("expected 1 reserve call, got %d", len(store.reserveCalls)) + } + if store.reserveCalls[0].Resources[quota.ResourceCacheStorageBytes] != expectedSize { + t.Errorf("expected reserved size %d, got %d", expectedSize, store.reserveCalls[0].Resources[quota.ResourceCacheStorageBytes]) + } + if store.reserveCalls[0].Key != "nar/abc.nar.xz" { + t.Errorf("reservation key = %q, want canonical nar path", store.reserveCalls[0].Key) + } +} + +func TestQuotaPublishOrder(t *testing.T) { + staging := t.TempDir() + store := newMockQuotaStore(true, "") + + var mu sync.Mutex + var order []string + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + mu.Lock() + defer mu.Unlock() + order = append(order, r.URL.Path) + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + u, _ := url.Parse(srv.URL) + backend := newHTTPUploadProxyBackend(u, nil, slog.Default()) + wrapper := NewStagingWrapper(backend, staging, store, "owner-1", "repo-1", srv.URL, slog.Default(), nil) + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPut, "/nar/abc.nar.xz", strings.NewReader("data")) + wrapper.ServeHTTP(rec, req) + + narinfoBody := validQuotaNarinfo + rec = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodPut, validQuotaNarinfoPath, strings.NewReader(narinfoBody)) + wrapper.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", rec.Code) + } + + mu.Lock() + defer mu.Unlock() + if len(order) != 2 { + t.Fatalf("expected 2 backend calls, got %d", len(order)) + } + if order[0] != "/nar/abc.nar.xz" { + t.Errorf("expected first call to be nar, got %q", order[0]) + } + if order[1] != validQuotaNarinfoPath { + t.Errorf("expected second call to be narinfo, got %q", order[1]) + } +} + +func TestQuotaSkipReturnsSuccessAndCleansStaging(t *testing.T) { + staging := t.TempDir() + store := newMockQuotaStore(false, "user_limit") + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + u, _ := url.Parse(srv.URL) + backend := newHTTPUploadProxyBackend(u, nil, slog.Default()) + wrapper := NewStagingWrapper(backend, staging, store, "owner-1", "repo-1", srv.URL, slog.Default(), nil) + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPut, "/nar/abc.nar.xz", strings.NewReader("some-nar-data")) + wrapper.ServeHTTP(rec, req) + + narinfoBody := validQuotaNarinfo + rec = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodPut, validQuotaNarinfoPath, strings.NewReader(narinfoBody)) + wrapper.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", rec.Code) + } + + narPath := filepath.Join(staging, "nar/abc.nar.xz") + if _, err := os.Stat(narPath); !errors.Is(err, os.ErrNotExist) { + t.Error("staged nar file was not cleaned up after quota skip") + } +} + +func TestQuotaSuccessfulCommit(t *testing.T) { + staging := t.TempDir() + store := newMockQuotaStore(true, "") + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + u, _ := url.Parse(srv.URL) + backend := newHTTPUploadProxyBackend(u, nil, slog.Default()) + wrapper := NewStagingWrapper(backend, staging, store, "owner-1", "repo-1", srv.URL, slog.Default(), nil) + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPut, "/nar/abc.nar.xz", strings.NewReader("data")) + wrapper.ServeHTTP(rec, req) + + narinfoBody := validQuotaNarinfo + rec = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodPut, validQuotaNarinfoPath, strings.NewReader(narinfoBody)) + wrapper.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", rec.Code) + } + + store.mu.Lock() + defer store.mu.Unlock() + if len(store.commitCalls) != 1 || store.commitCalls[0] != "res-abc" { + t.Errorf("expected commit call with res-abc, got %v", store.commitCalls) + } +} + +func TestQuotaHTTPFailureRetention(t *testing.T) { + staging := t.TempDir() + store := newMockQuotaStore(true, "") + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusForbidden) + })) + defer srv.Close() + + u, _ := url.Parse(srv.URL) + backend := newHTTPUploadProxyBackend(u, nil, slog.Default()) + wrapper := NewStagingWrapper(backend, staging, store, "owner-1", "repo-1", srv.URL, slog.Default(), nil) + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPut, "/nar/abc.nar.xz", strings.NewReader("data")) + wrapper.ServeHTTP(rec, req) + + narinfoBody := validQuotaNarinfo + rec = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodPut, validQuotaNarinfoPath, strings.NewReader(narinfoBody)) + wrapper.ServeHTTP(rec, req) + + store.mu.Lock() + if len(store.releaseCalls) != 0 { + store.mu.Unlock() + t.Fatalf("HTTP refusal released reservations: %v", store.releaseCalls) + } + store.mu.Unlock() + sw := wrapper.(*stagingWrapper) + sw.mu.Lock() + defer sw.mu.Unlock() + if len(sw.pendingCommits) != 1 || sw.pendingCommits[0] != "res-abc" { + t.Errorf("pending commits = %v, want [res-abc]", sw.pendingCommits) + } +} + +func TestQuotaUncertainFailureRetention(t *testing.T) { + staging := t.TempDir() + store := newMockQuotaStore(true, "") + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + hj, ok := w.(http.Hijacker) + if ok { + conn, _, _ := hj.Hijack() + conn.Close() + } + })) + defer srv.Close() + + u, _ := url.Parse(srv.URL) + backend := newHTTPUploadProxyBackend(u, nil, slog.Default()) + wrapper := NewStagingWrapper(backend, staging, store, "owner-1", "repo-1", srv.URL, slog.Default(), nil) + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPut, "/nar/abc.nar.xz", strings.NewReader("data")) + wrapper.ServeHTTP(rec, req) + + narinfoBody := validQuotaNarinfo + rec = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodPut, validQuotaNarinfoPath, strings.NewReader(narinfoBody)) + wrapper.ServeHTTP(rec, req) + + store.mu.Lock() + defer store.mu.Unlock() + if len(store.releaseCalls) != 0 { + t.Errorf("expected 0 release calls on uncertain failure, got %v", store.releaseCalls) + } +} + +func TestQuotaIdentityPropagation(t *testing.T) { + staging := t.TempDir() + store := newMockQuotaStore(true, "") + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + u, _ := url.Parse(srv.URL) + backend := newHTTPUploadProxyBackend(u, nil, slog.Default()) + wrapper := NewStagingWrapper(backend, staging, store, "owner-xyz", "repo-123", srv.URL, slog.Default(), nil) + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPut, "/nar/abc.nar.xz", strings.NewReader("data")) + wrapper.ServeHTTP(rec, req) + + narinfoBody := validQuotaNarinfo + rec = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodPut, validQuotaNarinfoPath, strings.NewReader(narinfoBody)) + wrapper.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", rec.Code) + } + + store.mu.Lock() + defer store.mu.Unlock() + if len(store.reserveCalls) != 1 { + t.Fatalf("expected 1 reserve call, got %d", len(store.reserveCalls)) + } + gotOwner := store.reserveCalls[0].Identity.OwnerDID + gotRepo := store.reserveCalls[0].Identity.RepoDID + if gotOwner != "owner-xyz" { + t.Errorf("owner did not propagate: got %q, want owner-xyz", gotOwner) + } + if gotRepo != "repo-123" { + t.Errorf("repo did not propagate: got %q, want repo-123", gotRepo) + } +} + +func TestQuotaTraversalAndSizeGuards(t *testing.T) { + staging := t.TempDir() + store := newMockQuotaStore(true, "") + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + u, _ := url.Parse(srv.URL) + backend := newHTTPUploadProxyBackend(u, nil, slog.Default()) + wrapper := NewStagingWrapper(backend, staging, store, "owner-1", "repo-1", srv.URL, slog.Default(), nil) + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPut, "/nar/../../evil.nar.xz", strings.NewReader("evil")) + wrapper.ServeHTTP(rec, req) + if rec.Code == http.StatusOK { + t.Error("expected traversal nar path to be blocked") + } + + rec = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodPut, "/../evil.narinfo", strings.NewReader("evil")) + wrapper.ServeHTTP(rec, req) + if rec.Code == http.StatusOK { + t.Error("expected traversal narinfo path to be blocked") + } + + oversizedBody := strings.Repeat("a", maxNarinfoSize+10) + rec = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodPut, validQuotaNarinfoPath, strings.NewReader(oversizedBody)) + wrapper.ServeHTTP(rec, req) + if rec.Code == http.StatusOK { + t.Error("expected oversized narinfo body to be blocked") + } +} + +func TestQuotaSecondStageFailure(t *testing.T) { + staging := t.TempDir() + store := newMockQuotaStore(true, "") + + var mu sync.Mutex + var order []string + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + mu.Lock() + defer mu.Unlock() + order = append(order, r.URL.Path) + if strings.HasSuffix(r.URL.Path, ".narinfo") { + w.WriteHeader(http.StatusForbidden) + } else { + w.WriteHeader(http.StatusOK) + } + })) + defer srv.Close() + + u, _ := url.Parse(srv.URL) + backend := newHTTPUploadProxyBackend(u, nil, slog.Default()) + wrapper := NewStagingWrapper(backend, staging, store, "owner-1", "repo-1", srv.URL, slog.Default(), nil) + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPut, "/nar/abc.nar.xz", strings.NewReader("data")) + wrapper.ServeHTTP(rec, req) + + narinfoBody := validQuotaNarinfo + rec = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodPut, validQuotaNarinfoPath, strings.NewReader(narinfoBody)) + wrapper.ServeHTTP(rec, req) + + if rec.Code == http.StatusOK { + t.Fatal("expected publish to fail but got success") + } + + store.mu.Lock() + defer store.mu.Unlock() + + if len(store.releaseCalls) != 0 { + t.Errorf("expected no release calls on second stage failure, got %v", store.releaseCalls) + } +} + +func TestQuotaNixStoreSecondStageFailure(t *testing.T) { + staging := t.TempDir() + store := newMockQuotaStore(true, "") + + runner := &fakeRunner{nextErr: errors.New("nix copy failed")} + backend, err := newNixStoreUploadBackend("ssh-ng://cache-host", staging, nil, slog.Default(), runner) + if err != nil { + t.Fatalf("failed to create nix store backend: %v", err) + } + + wrapper := NewStagingWrapper(backend, staging, store, "owner-1", "repo-1", "", slog.Default(), nil) + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPut, "/nar/abc.nar.xz", strings.NewReader("data")) + wrapper.ServeHTTP(rec, req) + + narinfoBody := validQuotaNarinfo + rec = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodPut, validQuotaNarinfoPath, strings.NewReader(narinfoBody)) + wrapper.ServeHTTP(rec, req) + + if rec.Code == http.StatusOK { + t.Fatal("expected publish to fail but got success") + } + + store.mu.Lock() + defer store.mu.Unlock() + + if len(store.releaseCalls) != 0 { + t.Errorf("expected no release calls on nix import failure, got %v", store.releaseCalls) + } +} + +func TestActivationConfigRepoDid(t *testing.T) { + dir := t.TempDir() + + spec := validImageSpec() + spec.BaseConfigHash = "abcdef" + specData, _ := json.Marshal(spec) + _ = os.WriteFile(filepath.Join(dir, "nixos.json"), specData, 0o644) + + e := &Engine{ + l: slog.Default(), + cfg: &config.Config{ + MicroVMPipelines: config.MicroVMPipelines{ + ImageDir: dir, + DefaultImage: "nixos", + }, + }, + } + + twf := tangled.Pipeline_Workflow{ + Raw: ` +image: nixos +services: + dummy: + command: true +`, + } + + repoDidVal := "did:web:repo-1" + tpl := tangled.Pipeline{ + TriggerMetadata: &tangled.Pipeline_TriggerMetadata{ + Repo: &tangled.Pipeline_TriggerRepo{ + Did: "did:web:owner-1", + RepoDid: &repoDidVal, + }, + }, + } + + wf, err := e.InitWorkflow(twf, tpl) + if err != nil { + t.Fatalf("InitWorkflow failed: %v", err) + } + + state, ok := wf.Data.(*workflowState) + if !ok || state == nil { + t.Fatal("workflowstate not initialized") + } + + expectedKey, err := buildConfigKey(state.ImageSpec, state.Config, repoDidVal) + if err != nil { + t.Fatalf("buildConfigKey failed: %v", err) + } + + if state.ConfigKey != expectedKey { + t.Errorf("ConfigKey mismatch: got %q, want %q", state.ConfigKey, expectedKey) + } +} + +func TestQuotaMetricsRecord(t *testing.T) { + staging := t.TempDir() + store := newMockQuotaStore(true, "") + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + m := observability.NewMetrics() + + u, _ := url.Parse(srv.URL) + backend := newHTTPUploadProxyBackend(u, nil, slog.Default()) + wrapper := NewStagingWrapper(backend, staging, store, "owner-1", "repo-1", srv.URL, slog.Default(), m) + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPut, "/nar/abc.nar.xz", strings.NewReader("data")) + wrapper.ServeHTTP(rec, req) + + narinfoBody := validQuotaNarinfo + rec = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodPut, validQuotaNarinfoPath, strings.NewReader(narinfoBody)) + wrapper.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", rec.Code) + } + + g, err := m.Registry().Gather() + if err != nil { + t.Fatalf("gather failed: %v", err) + } + + var foundDecision, foundUpload, foundBytes bool + for _, metricFamily := range g { + switch metricFamily.GetName() { + case "spindle_quota_decisions_total": + foundDecision = true + case "spindle_cache_uploads_total": + foundUpload = true + case "spindle_cache_upload_bytes_total": + foundBytes = true + } + } + + if !foundDecision { + t.Error("expected to find spindle_quota_decisions_total metric") + } + if !foundUpload { + t.Error("expected to find spindle_cache_uploads_total metric") + } + if !foundBytes { + t.Error("expected to find spindle_cache_upload_bytes_total metric") + } +} + +func TestQuotaCanceledContext(t *testing.T) { + t.Run("canceled_before_begin_publish", func(t *testing.T) { + staging := t.TempDir() + store := newMockQuotaStore(true, "") + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + u, _ := url.Parse(srv.URL) + backend := newHTTPUploadProxyBackend(u, nil, slog.Default()) + wrapper := NewStagingWrapper(backend, staging, store, "owner-1", "repo-1", srv.URL, slog.Default(), nil) + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPut, "/nar/abc.nar.xz", strings.NewReader("data")) + wrapper.ServeHTTP(rec, req) + + ctx, cancel := context.WithCancel(context.Background()) + + store.mu.Lock() + store.ReserveFunc = func(c context.Context, req quota.ReserveRequest) (quota.Reservation, error) { + cancel() + return quota.Reservation{ + ID: "res-abc", + Allowed: true, + Reason: quota.ReasonWithinLimit, + }, nil + } + store.mu.Unlock() + + rec = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodPut, validQuotaNarinfoPath, strings.NewReader(validQuotaNarinfo)) + req = req.WithContext(ctx) + wrapper.ServeHTTP(rec, req) + + store.mu.Lock() + released := len(store.releaseCalls) > 0 + store.mu.Unlock() + + if !released { + t.Fatal("expected reservation to be released if context was canceled before BeginCommit") + } + }) + + t.Run("canceled_during_publish", func(t *testing.T) { + staging := t.TempDir() + store := newMockQuotaStore(true, "") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if strings.HasSuffix(r.URL.Path, ".nar.xz") { + cancel() + return + } + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + u, _ := url.Parse(srv.URL) + backend := newHTTPUploadProxyBackend(u, nil, slog.Default()) + wrapper := NewStagingWrapper(backend, staging, store, "owner-1", "repo-1", srv.URL, slog.Default(), nil) + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPut, "/nar/abc.nar.xz", strings.NewReader("data")) + wrapper.ServeHTTP(rec, req) + + rec = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodPut, validQuotaNarinfoPath, strings.NewReader(validQuotaNarinfo)).WithContext(ctx) + wrapper.ServeHTTP(rec, req) + + store.mu.Lock() + defer store.mu.Unlock() + if len(store.releaseCalls) != 0 { + t.Fatalf("ambiguous publication released reservations: %v", store.releaseCalls) + } + }) +} + +func TestQuotaLocalStagingFailure(t *testing.T) { + staging := t.TempDir() + store := newMockQuotaStore(true, "") + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + u, _ := url.Parse(srv.URL) + backend := newHTTPUploadProxyBackend(u, nil, slog.Default()) + wrapper := NewStagingWrapper(backend, staging, store, "owner-1", "repo-1", srv.URL, slog.Default(), nil) + + narDir := filepath.Join(staging, "nar") + if err := os.MkdirAll(narDir, 0755); err != nil { + t.Fatalf("failed to create nar directory: %v", err) + } + if err := os.WriteFile(filepath.Join(narDir, "abc.nar.xz"), []byte("data"), 0644); err != nil { + t.Fatalf("failed to write nar file: %v", err) + } + + dstDir := filepath.Join(staging, "00000000000000000000000000000000.narinfo") + if err := os.MkdirAll(dstDir, 0755); err != nil { + t.Fatalf("failed to create directory at destination: %v", err) + } + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPut, validQuotaNarinfoPath, strings.NewReader(validQuotaNarinfo)) + wrapper.ServeHTTP(rec, req) + + if rec.Code == http.StatusOK { + t.Fatal("expected request to fail due to local staging error") + } + + store.mu.Lock() + defer store.mu.Unlock() + if len(store.releaseCalls) != 1 || store.releaseCalls[0] != "res-abc" { + t.Fatalf("expected 1 release call on local staging failure, got %v", store.releaseCalls) + } +} + +func TestQuotaBeginCommitFailure(t *testing.T) { + staging := t.TempDir() + store := newMockQuotaStore(true, "") + store.beginCommitErr = errors.New("begin publish failed") + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + u, _ := url.Parse(srv.URL) + backend := newHTTPUploadProxyBackend(u, nil, slog.Default()) + wrapper := NewStagingWrapper(backend, staging, store, "owner-1", "repo-1", srv.URL, slog.Default(), nil) + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPut, "/nar/abc.nar.xz", strings.NewReader("data")) + wrapper.ServeHTTP(rec, req) + + rec = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodPut, validQuotaNarinfoPath, strings.NewReader(validQuotaNarinfo)) + wrapper.ServeHTTP(rec, req) + + if rec.Code != http.StatusInternalServerError { + t.Fatalf("expected 500, got %d", rec.Code) + } + + store.mu.Lock() + defer store.mu.Unlock() + if len(store.releaseCalls) != 1 || store.releaseCalls[0] != "res-abc" { + t.Fatalf("expected release call with res-abc, got %v", store.releaseCalls) + } +} + +func TestQuotaTransientReleaseFailure(t *testing.T) { + staging := t.TempDir() + store := newMockQuotaStore(true, "") + store.releaseErr = errors.New("transient release failure") + store.beginCommitErr = errors.New("begin publish failed") + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + u, _ := url.Parse(srv.URL) + backend := newHTTPUploadProxyBackend(u, nil, slog.Default()) + wrapper := NewStagingWrapper(backend, staging, store, "owner-1", "repo-1", srv.URL, slog.Default(), nil) + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPut, "/nar/abc.nar.xz", strings.NewReader("data")) + wrapper.ServeHTTP(rec, req) + + rec = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodPut, validQuotaNarinfoPath, strings.NewReader(validQuotaNarinfo)) + wrapper.ServeHTTP(rec, req) + + store.mu.Lock() + if len(store.releaseCalls) != 1 || store.releaseCalls[0] != "res-abc" { + store.mu.Unlock() + t.Fatalf("expected 1 release call, got %v", store.releaseCalls) + } + store.mu.Unlock() + + store.mu.Lock() + store.releaseErr = nil + store.mu.Unlock() + + rec = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodGet, "/nix-cache-info", nil) + wrapper.ServeHTTP(rec, req) + + store.mu.Lock() + if len(store.releaseCalls) != 2 { + store.mu.Unlock() + t.Fatalf("expected 2 release calls after retry, got %v", store.releaseCalls) + } + store.mu.Unlock() + + store.mu.Lock() + store.releaseErr = errors.New("transient release failure") + store.mu.Unlock() + + sw := wrapper.(*stagingWrapper) + sw.mu.Lock() + sw.pendingReleases = append(sw.pendingReleases, "res-xyz") + sw.mu.Unlock() + + store.mu.Lock() + store.releaseErr = nil + store.mu.Unlock() + + if err := wrapper.Close(); err != nil { + t.Fatalf("Close failed: %v", err) + } + + store.mu.Lock() + defer store.mu.Unlock() + foundXYZ := false + for _, call := range store.releaseCalls { + if call == "res-xyz" { + foundXYZ = true + break + } + } + if !foundXYZ { + t.Fatal("expected Close to retry pending release for res-xyz") + } +} + +func TestQuotaCommitFailureRetriesOnNextRequest(t *testing.T) { + staging := t.TempDir() + store := newMockQuotaStore(true, "") + + var commitMu sync.Mutex + commitCalls := 0 + store.CommitFunc = func(context.Context, string) error { + commitMu.Lock() + defer commitMu.Unlock() + commitCalls++ + if commitCalls == 1 { + return errors.New("transient commit failure") + } + return nil + } + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + u, _ := url.Parse(srv.URL) + backend := newHTTPUploadProxyBackend(u, nil, slog.Default()) + wrapper := NewStagingWrapper(backend, staging, store, "owner-1", "repo-1", srv.URL, slog.Default(), nil) + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPut, "/nar/abc.nar.xz", strings.NewReader("data")) + wrapper.ServeHTTP(rec, req) + + rec = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodPut, validQuotaNarinfoPath, strings.NewReader(validQuotaNarinfo)) + wrapper.ServeHTTP(rec, req) + if rec.Code != http.StatusOK { + t.Fatalf("narinfo upload status = %d, want 200", rec.Code) + } + + rec = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodGet, "/nix-cache-info", nil) + wrapper.ServeHTTP(rec, req) + + commitMu.Lock() + defer commitMu.Unlock() + if commitCalls != 2 { + t.Fatalf("commit calls = %d, want 2", commitCalls) + } + sw := wrapper.(*stagingWrapper) + sw.mu.Lock() + defer sw.mu.Unlock() + if len(sw.pendingCommits) != 0 { + t.Fatalf("pending commits = %v, want none", sw.pendingCommits) + } +} + +func TestQuotaPublishingStateOrdering(t *testing.T) { + staging := t.TempDir() + store := newMockQuotaStore(true, "") + + var mu sync.Mutex + var trace []string + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + mu.Lock() + defer mu.Unlock() + trace = append(trace, "backend_"+r.URL.Path) + w.WriteHeader(http.StatusOK) + })) + defer srv.Close() + + u, _ := url.Parse(srv.URL) + backend := newHTTPUploadProxyBackend(u, nil, slog.Default()) + wrapper := NewStagingWrapper(backend, staging, store, "owner-1", "repo-1", srv.URL, slog.Default(), nil) + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPut, "/nar/abc.nar.xz", strings.NewReader("data")) + wrapper.ServeHTTP(rec, req) + + store.mu.Lock() + store.BeginCommitFunc = func(ctx context.Context, id string) error { + mu.Lock() + trace = append(trace, "begin_publish") + mu.Unlock() + + expectedPath := filepath.Join(staging, "00000000000000000000000000000000.narinfo") + if _, err := os.Stat(expectedPath); err == nil { + mu.Lock() + trace = append(trace, "local_narinfo_exists") + mu.Unlock() + } + return nil + } + store.ReserveFunc = func(ctx context.Context, req quota.ReserveRequest) (quota.Reservation, error) { + mu.Lock() + trace = append(trace, "reserve") + mu.Unlock() + return quota.Reservation{ + ID: "res-abc", + Allowed: true, + Reason: quota.ReasonWithinLimit, + }, nil + } + store.CommitFunc = func(ctx context.Context, id string) error { + mu.Lock() + trace = append(trace, "commit") + mu.Unlock() + return nil + } + store.mu.Unlock() + + rec = httptest.NewRecorder() + req = httptest.NewRequest(http.MethodPut, validQuotaNarinfoPath, strings.NewReader(validQuotaNarinfo)) + wrapper.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("expected 200, got %d", rec.Code) + } + + mu.Lock() + defer mu.Unlock() + + expectedTrace := []string{ + "reserve", + "begin_publish", + "local_narinfo_exists", + "backend_/nar/abc.nar.xz", + "backend_" + validQuotaNarinfoPath, + "commit", + } + + if len(trace) != len(expectedTrace) { + t.Fatalf("expected trace length %d, got %d. Trace: %v", len(expectedTrace), len(trace), trace) + } + + for i, v := range expectedTrace { + if trace[i] != v { + t.Errorf("trace mismatch at step %d: got %q, want %q", i, trace[i], v) + } + } +} diff --git a/spindle/engines/microvm/upload_staging.go b/spindle/engines/microvm/upload_staging.go new file mode 100644 index 00000000..9d6b46f1 --- /dev/null +++ b/spindle/engines/microvm/upload_staging.go @@ -0,0 +1,488 @@ +//go:build linux + +package microvm + +import ( + "bytes" + "context" + "errors" + "fmt" + "io" + "log/slog" + "net" + "net/http" + "net/url" + "os" + "path/filepath" + "strings" + "sync" + "time" + + "tangled.org/core/spindle/observability" + "tangled.org/core/spindle/quota" +) + +type stagingWrapper struct { + backend UploadCacheBackend + stagingDir string + quotaStore quota.ReservationStore + ownerDID string + repoDID string + uploadURL string + logger *slog.Logger + metrics *observability.Metrics + stageMu sync.Mutex + stagedNARs map[string]int64 + stagedBytes int64 + maxStagedBytes int64 + mu sync.Mutex + pendingReleases []string + pendingCommits []string +} + +func NewStagingWrapper(backend UploadCacheBackend, stagingDir string, quotaStore quota.ReservationStore, ownerDID, repoDID string, uploadURL string, logger *slog.Logger, metrics *observability.Metrics) UploadCacheBackend { + return &stagingWrapper{ + backend: backend, + stagingDir: stagingDir, + quotaStore: quotaStore, + ownerDID: ownerDID, + repoDID: repoDID, + uploadURL: uploadURL, + logger: logger, + metrics: metrics, + stagedNARs: make(map[string]int64), + maxStagedBytes: int64(maxNarUploadSize), + } +} + +func (s *stagingWrapper) ServeHTTP(w http.ResponseWriter, r *http.Request) { + cleanupErr := s.cleanupPending() + + relPath, err := normalizeUploadCachePath(r.URL.Path) + if err != nil { + s.logger.Warn("refusing upload cache request with unsafe path", "path", r.URL.Path, "error", err) + http.Error(w, "invalid path", http.StatusBadRequest) + return + } + if cleanupErr != nil && r.Method == http.MethodPut { + s.logger.Warn("cache upload blocked while quota lifecycle updates remain pending", "error", cleanupErr) + http.Error(w, "quota lifecycle update pending", http.StatusServiceUnavailable) + return + } + + if r.Method != http.MethodPut { + s.backend.ServeHTTP(w, r) + return + } + + if relPath == "nix-cache-info" { + s.backend.ServeHTTP(w, r) + return + } + + s.stageMu.Lock() + defer s.stageMu.Unlock() + + if isNarObjectPath(relPath) { + s.handlePutNar(w, r, relPath) + return + } + + if isNarinfoObjectPath(relPath) { + s.handlePutNarinfo(w, r, relPath) + return + } + + http.Error(w, "not found", http.StatusNotFound) +} + +func (s *stagingWrapper) handlePutNar(w http.ResponseWriter, r *http.Request, relPath string) { + name := strings.TrimPrefix(relPath, "nar/") + dst, err := s.stagingObjectPath(relPath) + if err != nil { + s.logger.Warn("refusing nar upload with unsafe path", "name", name, "error", err) + http.Error(w, "invalid nar path", http.StatusBadRequest) + return + } + + oldSize := s.stagedNARs[dst] + remaining := s.maxStagedBytes - (s.stagedBytes - oldSize) + if remaining <= 0 { + http.Error(w, "workflow staging limit reached", http.StatusRequestEntityTooLarge) + return + } + r.Body = http.MaxBytesReader(w, r.Body, remaining) + + var copyErr error + written, err := writeFileAtomic(dst, ".tmp-nar", func(f *os.File) (int64, error) { + n, err := io.Copy(f, r.Body) + copyErr = err + return n, err + }) + if err != nil { + s.logger.Warn("stage nar upload failed", "name", name, "error", err) + var maxErr *http.MaxBytesError + if errors.As(err, &maxErr) { + http.Error(w, "nar too large", http.StatusRequestEntityTooLarge) + return + } + if copyErr != nil { + http.Error(w, "upload failed", http.StatusBadRequest) + return + } + http.Error(w, "internal error", http.StatusInternalServerError) + return + } + s.stagedBytes += written - oldSize + s.stagedNARs[dst] = written + + s.logger.Debug("staged nar", "name", name, "bytes", written) + w.WriteHeader(http.StatusOK) +} + +func (s *stagingWrapper) handlePutNarinfo(w http.ResponseWriter, r *http.Request, relPath string) { + body, err := io.ReadAll(io.LimitReader(r.Body, maxNarinfoSize+1)) + if err != nil { + s.logger.Warn("read narinfo body failed", "path", relPath, "error", err) + http.Error(w, "upload failed", http.StatusBadRequest) + return + } + if len(body) > maxNarinfoSize { + s.logger.Warn("narinfo body exceeds maximum size", "path", relPath, "bytes", len(body)) + http.Error(w, "narinfo too large", http.StatusBadRequest) + return + } + + info, err := parseNarinfo(bytes.NewReader(body)) + if err != nil { + s.logger.Warn("refusing narinfo upload with invalid body", "path", relPath, "error", err) + http.Error(w, "invalid narinfo: "+err.Error(), http.StatusBadRequest) + return + } + storePathHash, _, err := parseStorePath(info.StorePath) + if err != nil { + s.logger.Warn("refusing narinfo upload with invalid store path", "path", relPath, "store_path", info.StorePath, "error", err) + http.Error(w, "invalid StorePath", http.StatusBadRequest) + return + } + fileHash := strings.TrimSuffix(filepath.Base(relPath), ".narinfo") + if fileHash != storePathHash { + s.logger.Warn("refusing narinfo upload with mismatched filename hash", "path", relPath, "store_path", info.StorePath) + http.Error(w, "narinfo filename does not match StorePath hash", http.StatusBadRequest) + return + } + if !isNarObjectPath(info.URL) { + s.logger.Warn("narinfo references invalid nar URL", "path", relPath, "url", info.URL) + http.Error(w, "invalid nar URL", http.StatusBadRequest) + return + } + + narPath, err := s.stagingObjectPath(info.URL) + if err != nil { + s.logger.Warn("narinfo references unsafe nar URL", "path", relPath, "url", info.URL, "error", err) + http.Error(w, "invalid nar URL", http.StatusBadRequest) + return + } + narFi, err := os.Stat(narPath) + if err != nil { + s.logger.Warn("narinfo references missing nar", "path", relPath, "url", info.URL, "error", err) + http.Error(w, "referenced nar does not exist", http.StatusBadRequest) + return + } + narSize := narFi.Size() + + req := quota.ReserveRequest{ + Kind: quota.KindNixCache, + Key: info.URL, + Identity: quota.Identity{ + OwnerDID: s.ownerDID, + RepoDID: s.repoDID, + }, + Resources: quota.Resources{ + quota.ResourceCacheStorageBytes: narSize, + }, + } + + backendName := "http" + if _, ok := s.backend.(*NixStoreUploadBackend); ok { + backendName = "nix_store" + } + + resID, err := s.quotaStore.Reserve(r.Context(), req) + if err != nil { + s.logger.Error("quota reservation failed", "error", err) + http.Error(w, "quota reservation error", http.StatusInternalServerError) + return + } + + metrics := s.metrics + if metrics != nil { + metrics.RecordQuotaDecision(string(quota.KindNixCache), quota.ResourceCacheStorageBytes, resID.Allowed, resID.Temporary, resID.Reason) + } + + if !resID.Allowed { + s.removeStagedNar(narPath) + s.logger.Info("cache upload skipped after quota rejection", + "reason", string(resID.Reason), + "bytes", narSize, + ) + if metrics != nil { + metrics.RecordCacheUpload(backendName, "quota_skipped") + metrics.RecordCacheUploadBytes(backendName, "quota_skipped", narSize) + } + w.WriteHeader(http.StatusOK) + return + } + + dst, err := s.stagingObjectPath(relPath) + if err != nil { + s.logger.Warn("refusing narinfo upload with unsafe path", "path", relPath, "error", err) + s.abortCleanup(resID.ID, narPath, "") + http.Error(w, "invalid path", http.StatusBadRequest) + return + } + if _, err := writeNarinfoFile(dst, body); err != nil { + s.logger.Warn("stage narinfo upload failed", "path", relPath, "error", err) + s.abortCleanup(resID.ID, narPath, dst) + http.Error(w, "internal error", http.StatusInternalServerError) + return + } + + nixBackend, usesNixStore := s.backend.(*NixStoreUploadBackend) + var fNar *os.File + if !usesNixStore { + fNar, err = os.Open(narPath) + if err != nil { + s.logger.Warn("open staged nar failed", "path", narPath, "error", err) + s.abortCleanup(resID.ID, narPath, dst) + http.Error(w, "internal error", http.StatusInternalServerError) + return + } + defer fNar.Close() + } + + if err := s.quotaStore.BeginCommit(r.Context(), resID.ID); err != nil { + s.logger.Error("quota publishing transition failed", "error", err) + s.abortCleanup(resID.ID, narPath, dst) + http.Error(w, "internal error", http.StatusInternalServerError) + return + } + + var pubErr error + var narPublished bool + if usesNixStore { + narPublished = true + pubErr = nixBackend.importStorePath(r.Context(), info.StorePath) + } else { + pubErr = s.publishHTTP(r.Context(), info.URL, fNar, narSize) + if pubErr == nil { + narPublished = true + pubErr = s.publishHTTP(r.Context(), relPath, bytes.NewReader(body), int64(len(body))) + } + } + + if pubErr != nil { + s.logger.Warn("backend publish failed", "error", pubErr) + if !narPublished && !s.isUncertain(pubErr) { + s.abortCleanup(resID.ID, narPath, dst) + } else { + s.removeStagedNar(narPath) + _ = os.Remove(dst) + s.trackCommit(resID.ID) + } + if metrics != nil { + metrics.RecordCacheUpload(backendName, "failed") + metrics.RecordCacheUploadBytes(backendName, "failed", narSize) + } + http.Error(w, "publish failed: "+pubErr.Error(), http.StatusBadGateway) + return + } + + err = s.quotaStore.Commit(r.Context(), resID.ID) + if err != nil { + s.logger.Error("quota commit failed, retaining publishing reservation", "reservation_id", resID.ID, "error", err) + s.trackCommit(resID.ID) + } + + s.removeStagedNar(narPath) + _ = os.Remove(dst) + + if metrics != nil { + metrics.RecordCacheUpload(backendName, "published") + metrics.RecordCacheUploadBytes(backendName, "published", narSize) + } + + w.WriteHeader(http.StatusOK) +} + +var errHTTPPublishUncertain = errors.New("http upload result is uncertain") + +func (s *stagingWrapper) publishHTTP(ctx context.Context, relPath string, body io.Reader, contentLength int64) error { + targetURL, err := url.Parse(s.uploadURL) + if err != nil { + return err + } + + u := *targetURL + u.Path = strings.TrimSuffix(u.Path, "/") + "/" + relPath + + req, err := http.NewRequestWithContext(ctx, http.MethodPut, u.String(), body) + if err != nil { + return err + } + + req.Header.Set("Content-Type", "application/octet-stream") + if contentLength > 0 { + req.ContentLength = contentLength + } + + if targetURL.User != nil { + password, _ := targetURL.User.Password() + req.SetBasicAuth(targetURL.User.Username(), password) + } + + client := &http.Client{ + Transport: proxyTransport, + } + + resp, err := client.Do(req) + if err != nil { + return err + } + defer resp.Body.Close() + + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 1024)) + return fmt.Errorf("%w: status %d: %s", errHTTPPublishUncertain, resp.StatusCode, string(respBody)) + } + + return nil +} + +func (s *stagingWrapper) isUncertain(err error) bool { + if err == nil { + return false + } + if errors.Is(err, errHTTPPublishUncertain) { + return true + } + if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) { + return true + } + var netErr net.Error + if errors.As(err, &netErr) { + return true + } + return false +} + +func (s *stagingWrapper) stagingObjectPath(relPath string) (string, error) { + if !isNarObjectPath(relPath) && !isNarinfoObjectPath(relPath) { + return "", fmt.Errorf("invalid cache object path %q", relPath) + } + + local, err := filepath.Localize(relPath) + if err != nil { + return "", fmt.Errorf("unsafe cache object path %q: %w", relPath, err) + } + + return filepath.Join(s.stagingDir, local), nil +} + +func (s *stagingWrapper) removeStagedNar(path string) { + if size, ok := s.stagedNARs[path]; ok { + s.stagedBytes -= size + delete(s.stagedNARs, path) + } + _ = os.Remove(path) +} + +func (s *stagingWrapper) abortCleanup(resID string, narPath, narinfoPath string) { + if narPath != "" { + s.removeStagedNar(narPath) + } + if narinfoPath != "" { + _ = os.Remove(narinfoPath) + } + if resID != "" { + s.releaseOrTrack(resID) + } +} + +func (s *stagingWrapper) trackCommit(resID string) { + if resID == "" { + return + } + s.mu.Lock() + s.pendingCommits = append(s.pendingCommits, resID) + s.mu.Unlock() +} + +func (s *stagingWrapper) releaseOrTrack(resID string) { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + err := s.quotaStore.Release(ctx, resID) + if err != nil { + s.logger.Error("quota reservation release failed", "reservation_id", resID, "error", err) + s.mu.Lock() + s.pendingReleases = append(s.pendingReleases, resID) + s.mu.Unlock() + } +} + +func (s *stagingWrapper) cleanupPending() error { + s.mu.Lock() + if len(s.pendingCommits) == 0 && len(s.pendingReleases) == 0 { + s.mu.Unlock() + return nil + } + commits := s.pendingCommits + releases := s.pendingReleases + s.pendingCommits = nil + s.pendingReleases = nil + s.mu.Unlock() + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + var failedCommits []string + for _, id := range commits { + if err := s.quotaStore.Commit(ctx, id); err != nil { + s.logger.Warn("quota commit retry failed", "reservation_id", id, "error", err) + failedCommits = append(failedCommits, id) + } else { + s.logger.Info("quota commit retry succeeded", "reservation_id", id) + } + } + + var failedReleases []string + for _, id := range releases { + if err := s.quotaStore.Release(ctx, id); err != nil { + s.logger.Warn("quota reservation release retry failed", "reservation_id", id, "error", err) + failedReleases = append(failedReleases, id) + } else { + s.logger.Info("quota reservation release retry succeeded", "reservation_id", id) + } + } + + s.mu.Lock() + s.pendingCommits = append(s.pendingCommits, failedCommits...) + s.pendingReleases = append(s.pendingReleases, failedReleases...) + s.mu.Unlock() + + var errs []error + if len(failedCommits) > 0 { + errs = append(errs, fmt.Errorf("%d quota commits remain pending", len(failedCommits))) + } + if len(failedReleases) > 0 { + errs = append(errs, fmt.Errorf("%d quota releases remain pending", len(failedReleases))) + } + return errors.Join(errs...) +} + +func (s *stagingWrapper) Close() error { + cleanupErr := s.cleanupPending() + removeErr := os.RemoveAll(s.stagingDir) + backendErr := s.backend.Close() + return errors.Join(cleanupErr, removeErr, backendErr) +} diff --git a/spindle/engines/microvm/vm.go b/spindle/engines/microvm/vm.go index 9399b3de..d3902af1 100644 --- a/spindle/engines/microvm/vm.go +++ b/spindle/engines/microvm/vm.go @@ -25,6 +25,7 @@ import ( "tangled.org/core/spindle/engine" "tangled.org/core/spindle/models" + "tangled.org/core/spindle/quota" ) const ( @@ -182,6 +183,9 @@ type workflowState struct { DNSProxy *DNSProxy WorkDir string NixOSToplevelCache nixosToplevelCacheStore + OwnerDID string + RepoDID string + QuotaStore quota.ReservationStore StartedAt time.Time // when the VM booted, for the max-lifetime cap ResourceUsage *engine.WorkflowResourceUsage ResourceUsageAvailable bool diff --git a/spindle/engines/nixery/engine.go b/spindle/engines/nixery/engine.go index fb5232c3..b89013f2 100644 --- a/spindle/engines/nixery/engine.go +++ b/spindle/engines/nixery/engine.go @@ -24,6 +24,7 @@ import ( "tangled.org/core/spindle/config" "tangled.org/core/spindle/engine" "tangled.org/core/spindle/models" + "tangled.org/core/spindle/quota" "tangled.org/core/spindle/secrets" ) @@ -107,6 +108,23 @@ func (e *Engine) InitWorkflow(twf tangled.Pipeline_Workflow, tpl tangled.Pipelin sstep.kind = models.StepKindUser swf.Steps = append(swf.Steps, sstep) } + var ownerDID, repoDID string + if md := tpl.TriggerMetadata; md != nil && md.Repo != nil { + ownerDID = md.Repo.Did + repoDID = ownerDID + if md.Repo.RepoDid != nil && *md.Repo.RepoDid != "" { + repoDID = *md.Repo.RepoDid + } + } + if ownerDID == "" { + return nil, fmt.Errorf("missing owner DID in pipeline trigger metadata") + } + if repoDID == "" { + return nil, fmt.Errorf("missing repository DID in pipeline trigger metadata") + } + swf.OwnerDID = ownerDID + swf.RepoDID = repoDID + swf.Name = twf.Name swf.Environment = dwf.Environment addl.image = workflowImage(dwf.Dependencies, e.cfg.NixeryPipelines.Nixery) @@ -533,3 +551,9 @@ func (e *Engine) drainCleanups(wid models.WorkflowId) []cleanupFunc { func networkName(wid models.WorkflowId) string { return fmt.Sprintf("workflow-network-%s", wid) } +func (e *Engine) QuotaResources(wf *models.Workflow) quota.Resources { + return quota.Resources{ + quota.ResourceWorkflows: 1, + quota.ResourceMemoryMiB: e.cfg.NixeryPipelines.MaxJobMemoryMB, + } +} diff --git a/spindle/mill/auth_test.go b/spindle/mill/auth_test.go index c78f851c..4e06a057 100644 --- a/spindle/mill/auth_test.go +++ b/spindle/mill/auth_test.go @@ -117,7 +117,7 @@ func TestOnAttemptResultIgnoresForeignLease(t *testing.T) { n := notifier.New() m := New(l, Config{ReconnectGrace: time.Minute}) - m.Attach(bdb, &n) + m.Attach(bdb, &n, testQuotaManager(t, bdb)) foreign := newLease("lease-foreign", "node-a", "inc-a", "dummy") m.mu.Lock() @@ -134,7 +134,7 @@ func TestOnAttemptResultIgnoresForeignLease(t *testing.T) { LeaseId: foreign.id, Payload: &millv1.Event_AttemptResult{ AttemptResult: &millv1.AttemptResult{ - Status: millv1.TerminalStatus_SUCCESS, + Status: millv1.TerminalStatus_TERMINAL_STATUS_SUCCESS, }, }, }, @@ -160,7 +160,7 @@ func TestOnAttemptResultIgnoresAbsentLease(t *testing.T) { n := notifier.New() m := New(l, Config{ReconnectGrace: time.Minute}) - m.Attach(bdb, &n) + m.Attach(bdb, &n, testQuotaManager(t, bdb)) // the reporting node owns this bystander lease, proving an absent-lease // stream does not spill onto another lease @@ -180,7 +180,7 @@ func TestOnAttemptResultIgnoresAbsentLease(t *testing.T) { LeaseId: "lease-nonexistent", Payload: &millv1.Event_AttemptResult{ AttemptResult: &millv1.AttemptResult{ - Status: millv1.TerminalStatus_SUCCESS, + Status: millv1.TerminalStatus_TERMINAL_STATUS_SUCCESS, }, }, }, @@ -208,7 +208,7 @@ func TestOnAttemptResultDeliversOwnedLease(t *testing.T) { n := notifier.New() m := New(l, Config{ReconnectGrace: time.Minute}) - m.Attach(bdb, &n) + m.Attach(bdb, &n, testQuotaManager(t, bdb)) owned := newLease("lease-owned", "node-b", "inc-b", "dummy") m.mu.Lock() @@ -225,7 +225,7 @@ func TestOnAttemptResultDeliversOwnedLease(t *testing.T) { LeaseId: owned.id, Payload: &millv1.Event_AttemptResult{ AttemptResult: &millv1.AttemptResult{ - Status: millv1.TerminalStatus_SUCCESS, + Status: millv1.TerminalStatus_TERMINAL_STATUS_SUCCESS, }, }, }, @@ -236,8 +236,8 @@ func TestOnAttemptResultDeliversOwnedLease(t *testing.T) { if !ok { t.Fatal("attempt-result on an owned lease was not delivered") } - if got := res.GetStatus(); got != millv1.TerminalStatus_SUCCESS { - t.Fatalf("delivered terminal status = %v, want %v", got, millv1.TerminalStatus_SUCCESS) + if got := res.GetStatus(); got != millv1.TerminalStatus_TERMINAL_STATUS_SUCCESS { + t.Fatalf("delivered terminal status = %v, want %v", got, millv1.TerminalStatus_TERMINAL_STATUS_SUCCESS) } if owned.getState() != leaseDone { t.Fatal("owned lease was not sealed after its terminal was delivered") @@ -256,7 +256,7 @@ func TestOnStatusEventOwnership(t *testing.T) { n := notifier.New() m := New(l, Config{ReconnectGrace: time.Minute}) - m.Attach(bdb, &n) + m.Attach(bdb, &n, testQuotaManager(t, bdb)) foreign := newLease("lease-foreign", "node-x", "inc-x", "dummy") foreign.wid = models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "foreign"}, Name: "build"} @@ -280,7 +280,7 @@ func TestOnStatusEventOwnership(t *testing.T) { LeaseId: foreign.id, Payload: &millv1.Event_StatusEvent{ StatusEvent: &millv1.StatusEvent{ - Status: millv1.NonterminalStatus_RUNNING, + Status: millv1.NonterminalStatus_NONTERMINAL_STATUS_RUNNING, }, }, }, @@ -298,7 +298,7 @@ func TestOnStatusEventOwnership(t *testing.T) { LeaseId: owned.id, Payload: &millv1.Event_StatusEvent{ StatusEvent: &millv1.StatusEvent{ - Status: millv1.NonterminalStatus_RUNNING, + Status: millv1.NonterminalStatus_NONTERMINAL_STATUS_RUNNING, }, }, }, @@ -325,7 +325,7 @@ func setupTestServer(t *testing.T, authorizedLabels []string) (*Mill, *db.DB, *h m := New(discardLogger(), Config{ ReconnectGrace: time.Minute, }) - m.Attach(bdb, &n) + m.Attach(bdb, &n, testQuotaManager(t, bdb)) const secret = "test-secret" registerTestExecutor(t, bdb, "dev-node", HashToken(secret), authorizedLabels) diff --git a/spindle/mill/engine.go b/spindle/mill/engine.go index 0f87c235..2fd35045 100644 --- a/spindle/mill/engine.go +++ b/spindle/mill/engine.go @@ -40,9 +40,19 @@ func NewEngine(name string, mill *Mill) *Engine { // and marks pending normally. the real InitWorkflow runs exactly once, on // the executor inside ReserveSeat, and commit reuses that workflow func (e *Engine) InitWorkflow(twf tangled.Pipeline_Workflow, tpl tangled.Pipeline) (*models.Workflow, error) { + var ownerDID, repoDID string + if md := tpl.TriggerMetadata; md != nil && md.Repo != nil { + ownerDID = md.Repo.Did + repoDID = ownerDID + if md.Repo.RepoDid != nil && *md.Repo.RepoDid != "" { + repoDID = *md.Repo.RepoDid + } + } return &models.Workflow{ Name: twf.Name, Environment: map[string]string{}, + OwnerDID: ownerDID, + RepoDID: repoDID, Steps: []models.Step{remoteStep{}}, Data: &millWorkflowState{ RawWorkflow: twf, diff --git a/spindle/mill/executor/executor.go b/spindle/mill/executor/executor.go index ba33523c..16255818 100644 --- a/spindle/mill/executor/executor.go +++ b/spindle/mill/executor/executor.go @@ -15,6 +15,7 @@ import ( "github.com/bluesky-social/indigo/atproto/syntax" "github.com/gorilla/websocket" + "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/codes" "go.opentelemetry.io/otel/trace" "tangled.org/core/api/tangled" @@ -28,6 +29,7 @@ import ( millv1 "tangled.org/core/spindle/mill/proto/gen" "tangled.org/core/spindle/models" "tangled.org/core/spindle/observability" + "tangled.org/core/spindle/quota" ) const ( @@ -75,6 +77,7 @@ type Executor struct { snapshotMu sync.Mutex nextSeqno uint64 + quotaClient *QuotaClient lifecycleCtx context.Context jobsWG sync.WaitGroup @@ -94,6 +97,7 @@ type reservation struct { cancel context.CancelFunc ttlTimer *time.Timer stopTail func() + runDone chan struct{} traceParent trace.SpanContext } @@ -236,11 +240,17 @@ func (e *Executor) runSession(ctx context.Context) error { e.sessionCancel = cancelSession e.enc = enc e.connMu.Unlock() + if e.quotaClient != nil { + e.quotaClient.SetSendFn(e.sendQuota) + } defer func() { e.connMu.Lock() e.enc = nil e.sessionCancel = nil e.connMu.Unlock() + if e.quotaClient != nil { + e.quotaClient.OnDisconnect() + } }() readErr := make(chan error, 1) @@ -296,11 +306,31 @@ func (e *Executor) dispatch(ctx context.Context, msg *millproto.Message) { e.handleCancel(msg.GetCancelAttempt().GetLeaseId()) case msg.GetAck() != nil: e.handleAck(msg.GetAck()) + case msg.GetQuotaResp() != nil: + if e.quotaClient != nil { + e.quotaClient.HandleResponse(msg.GetQuotaResp().GetRequestId(), msg.GetQuotaResp()) + } default: e.l.Warn("unhandled incoming message", "type", fmt.Sprintf("%T", msg)) } } +func (e *Executor) SetQuotaClient(c *QuotaClient) { + e.mu.Lock() + defer e.mu.Unlock() + e.quotaClient = c +} + +func (e *Executor) sendQuota(msg *millproto.Message) error { + e.connMu.Lock() + enc := e.enc + e.connMu.Unlock() + if enc == nil { + return errors.New("not connected to mill") + } + e.send(msg) + return nil +} func (e *Executor) sendReject(leaseID string, reason string, class millv1.RejectClass) { e.send(&millproto.Message{ReserveResult: &millv1.ReserveResult{ @@ -336,6 +366,13 @@ func (e *Executor) handleReserve(ctx context.Context, rs *millv1.ReserveSeat) { } reserveCtx := observability.ExtractFromTraceparentAndTracestate(parentCtx, rs.GetTraceparent(), rs.GetTracestate()) reserveCtx, span := observability.Tracer().Start(reserveCtx, "executor.assignment") + if span.IsRecording() { + span.SetAttributes( + attribute.String(observability.LeaseIDKey, rs.GetLeaseId()), + attribute.String(observability.ExecutorNodeIDKey, e.nodeID), + attribute.String(observability.PipelineIDKey, (&models.PipelineId{Knot: rs.GetKnot(), Rkey: rs.GetRkey()}).AtUri().String()), + ) + } accepted := false defer func() { if accepted { @@ -373,6 +410,11 @@ func (e *Executor) handleReserve(ctx context.Context, rs *millv1.ReserveSeat) { reject("bad workflow json", millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE) return } + pipelineId := models.PipelineId{Knot: rs.GetKnot(), Rkey: rs.GetRkey()} + wid := models.WorkflowId{PipelineId: pipelineId, Name: twf.Name} + if span.IsRecording() { + span.SetAttributes(attribute.String(observability.WorkflowIDKey, wid.String())) + } var tpl tangled.Pipeline if err := json.Unmarshal([]byte(rs.GetRawPipelineJson()), &tpl); err != nil { reject("bad pipeline json", millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE) @@ -382,15 +424,19 @@ func (e *Executor) handleReserve(ctx context.Context, rs *millv1.ReserveSeat) { reject("pipeline missing trigger metadata", millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE) return } - - pipelineId := models.PipelineId{Knot: rs.GetKnot(), Rkey: rs.GetRkey()} - wid := models.WorkflowId{PipelineId: pipelineId, Name: twf.Name} - wf, err := realEngine.InitWorkflow(twf, tpl) if err != nil { reject("init workflow: "+err.Error(), millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE) return } + if e.quotaClient != nil { + if binder, ok := realEngine.(engine.WorkflowQuotaStoreBinder); ok { + if err := binder.BindWorkflowQuotaStore(wf, e.quotaClient.ForLease(rs.GetLeaseId())); err != nil { + reject("bind workflow quota store: "+err.Error(), millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE) + return + } + } + } if validator, ok := realEngine.(engine.WorkflowPlacementValidator); ok { if err := validator.ValidateWorkflowPlacement(wf); err != nil { reject("validate workflow placement: "+err.Error(), millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE) @@ -416,6 +462,11 @@ func (e *Executor) handleReserve(ctx context.Context, rs *millv1.ReserveSeat) { if tpl.TriggerMetadata != nil && tpl.TriggerMetadata.Repo != nil && tpl.TriggerMetadata.Repo.RepoDid != nil { repoDid, _ = syntax.ParseDID(*tpl.TriggerMetadata.Repo.RepoDid) } + if span.IsRecording() { + if repoDid.String() != "" { + span.SetAttributes(attribute.String(observability.RepoDIDKey, repoDid.String())) + } + } res := &reservation{ leaseID: rs.GetLeaseId(), @@ -444,14 +495,30 @@ func (e *Executor) handleReserve(ctx context.Context, rs *millv1.ReserveSeat) { e.mu.Unlock() e.send(&millproto.Message{ReserveResult: &millv1.ReserveResult{ - LeaseId: rs.GetLeaseId(), - Accepted: true, + LeaseId: rs.GetLeaseId(), + Accepted: true, + QuotaResources: reportedResources(realEngine, wf), }}) accepted = true e.pushSnapshotLocked() e.snapshotMu.Unlock() } +func reportedResources(eng models.Engine, wf *models.Workflow) map[string]int64 { + reporter, ok := eng.(engine.WorkflowQuotaReporter) + if !ok { + return nil + } + resources := reporter.QuotaResources(wf) + if err := quota.ValidateResources(resources); err != nil { + return nil + } + if len(resources) == 0 { + return nil + } + return maps.Clone(resources) +} + func (e *Executor) handleCommit(ctx context.Context, cl *millv1.CommitLease) { e.mu.Lock() res := e.active[cl.GetLeaseId()] @@ -482,9 +549,25 @@ func (e *Executor) handleCommit(ctx context.Context, cl *millv1.CommitLease) { } commitCtx := observability.ExtractFromTraceparentAndTracestate(parentCtx, cl.GetTraceparent(), cl.GetTracestate()) runCtx, runSpan := observability.Tracer().Start(commitCtx, "executor.run") + if runSpan.IsRecording() { + attrs := []attribute.KeyValue{ + attribute.String(observability.WorkflowIDKey, res.wid.String()), + attribute.String(observability.LeaseIDKey, res.leaseID), + attribute.String(observability.ExecutorNodeIDKey, e.nodeID), + attribute.String(observability.PipelineIDKey, res.wid.PipelineId.AtUri().String()), + } + if res.repoDid.String() != "" { + attrs = append(attrs, attribute.String(observability.RepoDIDKey, res.repoDid.String())) + } + if res.wf != nil && res.wf.OwnerDID != "" { + attrs = append(attrs, attribute.String(observability.OwnerDIDKey, res.wf.OwnerDID)) + } + runSpan.SetAttributes(attrs...) + } jobCtx, cancel := context.WithCancel(runCtx) res.cancel = cancel + res.runDone = make(chan struct{}) e.mu.Unlock() vault := newMemVault(cl.GetSecrets()) @@ -500,10 +583,14 @@ func (e *Executor) handleCommit(ctx context.Context, cl *millv1.CommitLease) { e.jobsWG.Add(1) go func() { defer e.jobsWG.Done() + defer close(res.runDone) defer runSpan.End() - engine.StartWorkflows(e.l, vault, e.cfg, nil, e.db, e.n, jobCtx, pipeline, res.wid.PipelineId) + el := e.l.With( + "lease_id", res.leaseID, + "node_id", e.nodeID, + ) + engine.StartWorkflows(el, vault, e.cfg, nil, nil, e.db, e.n, jobCtx, pipeline, res.wid.PipelineId) }() - e.sendCommitted(cl.GetLeaseId()) } diff --git a/spindle/mill/executor/observe.go b/spindle/mill/executor/observe.go index 8ef595fb..d11caa2c 100644 --- a/spindle/mill/executor/observe.go +++ b/spindle/mill/executor/observe.go @@ -74,8 +74,13 @@ func (e *Executor) finishJob(res *reservation, st *tangled.PipelineStatus) error return nil } cancelled := res.cancelled + runDone := res.runDone e.mu.Unlock() + if runDone != nil { + <-runDone + } + // the log tail finalizes first so all log lines precede the terminal event if res.stopTail != nil { res.stopTail() diff --git a/spindle/mill/executor/outbox.go b/spindle/mill/executor/outbox.go index 983b125b..e210e1de 100644 --- a/spindle/mill/executor/outbox.go +++ b/spindle/mill/executor/outbox.go @@ -115,7 +115,7 @@ func (e *Executor) appendStatus(leaseID string, st *tangled.PipelineStatus) erro } exit, errStr := parseStatusExitAndError(st) payload := &millv1.Event_StatusEvent{StatusEvent: &millv1.StatusEvent{ - Status: millv1.NonterminalStatus_RUNNING, + Status: millv1.NonterminalStatus_NONTERMINAL_STATUS_RUNNING, Error: errStr, ExitCode: exit, }} @@ -130,13 +130,13 @@ func (e *Executor) appendTerminalWithArtifact(leaseID, status string, st *tangle var terminalStatus millv1.TerminalStatus switch status { case string(models.StatusKindSuccess): - terminalStatus = millv1.TerminalStatus_SUCCESS + terminalStatus = millv1.TerminalStatus_TERMINAL_STATUS_SUCCESS case string(models.StatusKindFailed): - terminalStatus = millv1.TerminalStatus_FAILED + terminalStatus = millv1.TerminalStatus_TERMINAL_STATUS_FAILED case string(models.StatusKindTimeout): - terminalStatus = millv1.TerminalStatus_TIMEOUT + terminalStatus = millv1.TerminalStatus_TERMINAL_STATUS_TIMEOUT case string(models.StatusKindCancelled): - terminalStatus = millv1.TerminalStatus_CANCELLED + terminalStatus = millv1.TerminalStatus_TERMINAL_STATUS_CANCELLED default: return fmt.Errorf("unsupported terminal status %q", status) } diff --git a/spindle/mill/executor/quota_client.go b/spindle/mill/executor/quota_client.go new file mode 100644 index 00000000..5f9a14fa --- /dev/null +++ b/spindle/mill/executor/quota_client.go @@ -0,0 +1,166 @@ +package executor + +import ( + "context" + "errors" + "fmt" + "maps" + "sync" + + "github.com/google/uuid" + "tangled.org/core/spindle/quota" + + millproto "tangled.org/core/spindle/mill/proto" + millv1 "tangled.org/core/spindle/mill/proto/gen" +) + +type QuotaClient struct { + mu sync.Mutex + pending map[string]chan *millv1.QuotaResponse + sendFn func(*millproto.Message) error +} + +func NewQuotaClient() *QuotaClient { + return &QuotaClient{ + pending: make(map[string]chan *millv1.QuotaResponse), + } +} + +func (c *QuotaClient) SetSendFn(sendFn func(*millproto.Message) error) { + c.mu.Lock() + defer c.mu.Unlock() + c.sendFn = sendFn +} + +func (c *QuotaClient) OnDisconnect() { + c.mu.Lock() + oldPending := c.pending + c.pending = make(map[string]chan *millv1.QuotaResponse) + c.sendFn = nil + c.mu.Unlock() + + for _, ch := range oldPending { + select { + case ch <- nil: + default: + } + } +} + +func (c *QuotaClient) HandleResponse(reqID string, resp *millv1.QuotaResponse) { + c.mu.Lock() + ch, ok := c.pending[reqID] + c.mu.Unlock() + if ok && resp != nil { + select { + case ch <- resp: + default: + } + } +} + +func (c *QuotaClient) ForLease(leaseID string) quota.ReservationStore { + return &leaseQuotaStore{client: c, leaseID: leaseID} +} + +type leaseQuotaStore struct { + client *QuotaClient + leaseID string +} + +var _ quota.ReservationStore = (*leaseQuotaStore)(nil) + +func (s *leaseQuotaStore) Reserve(ctx context.Context, req quota.ReserveRequest) (quota.Reservation, error) { + if err := quota.Validate(req); err != nil { + return quota.Reservation{}, err + } + if s.leaseID == "" { + return quota.Reservation{}, errors.New("empty workflow lease") + } + + resp, err := s.client.roundTrip(ctx, &millv1.QuotaRequest{ + Operation: millv1.QuotaOperation_QUOTA_OPERATION_RESERVE, + LeaseId: s.leaseID, + Kind: string(req.Kind), + Key: req.Key, + Resources: maps.Clone(req.Resources), + }) + if err != nil { + return quota.Reservation{}, err + } + if resp.GetAllowed() && resp.GetTemporary() { + return quota.Reservation{}, errors.New("mill allowed a reservation and deferred it at once") + } + return quota.Reservation{ + ID: resp.GetReservationId(), + Allowed: resp.GetAllowed(), + Temporary: resp.GetTemporary(), + Reason: resp.GetReason(), + Resource: resp.GetResource(), + }, nil +} + +func (s *leaseQuotaStore) BeginCommit(ctx context.Context, reservationID string) error { + return s.transition(ctx, millv1.QuotaOperation_QUOTA_OPERATION_BEGIN_COMMIT, reservationID) +} + +func (s *leaseQuotaStore) Commit(ctx context.Context, reservationID string) error { + return s.transition(ctx, millv1.QuotaOperation_QUOTA_OPERATION_COMMIT, reservationID) +} + +func (s *leaseQuotaStore) Release(ctx context.Context, reservationID string) error { + return s.transition(ctx, millv1.QuotaOperation_QUOTA_OPERATION_RELEASE, reservationID) +} + +func (s *leaseQuotaStore) transition(ctx context.Context, operation millv1.QuotaOperation, reservationID string) error { + if reservationID == "" { + return nil + } + if s.leaseID == "" { + return fmt.Errorf("reservation %q is not associated with a live lease", reservationID) + } + _, err := s.client.roundTrip(ctx, &millv1.QuotaRequest{ + Operation: operation, + LeaseId: s.leaseID, + ReservationId: reservationID, + }) + return err +} + +func (c *QuotaClient) roundTrip(ctx context.Context, req *millv1.QuotaRequest) (*millv1.QuotaResponse, error) { + reqID := uuid.NewString() + req.RequestId = reqID + ch := make(chan *millv1.QuotaResponse, 1) + + c.mu.Lock() + send := c.sendFn + if send == nil { + c.mu.Unlock() + return nil, errors.New("quota client not connected") + } + c.pending[reqID] = ch + c.mu.Unlock() + + defer func() { + c.mu.Lock() + delete(c.pending, reqID) + c.mu.Unlock() + }() + + if err := send(&millproto.Message{QuotaReq: req}); err != nil { + return nil, fmt.Errorf("send quota request: %w", err) + } + + select { + case <-ctx.Done(): + return nil, ctx.Err() + case resp := <-ch: + if resp == nil { + return nil, errors.New("quota client disconnected") + } + if resp.GetError() != "" { + return nil, errors.New(resp.GetError()) + } + return resp, nil + } +} diff --git a/spindle/mill/executor/quota_client_test.go b/spindle/mill/executor/quota_client_test.go new file mode 100644 index 00000000..1c25f5eb --- /dev/null +++ b/spindle/mill/executor/quota_client_test.go @@ -0,0 +1,704 @@ +package 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() +} diff --git a/spindle/mill/executor/reserved.go b/spindle/mill/executor/reserved.go index 607e9e85..d3309da9 100644 --- a/spindle/mill/executor/reserved.go +++ b/spindle/mill/executor/reserved.go @@ -8,11 +8,9 @@ import ( "tangled.org/core/spindle/engine" "tangled.org/core/spindle/models" + "tangled.org/core/spindle/quota" ) -// wraps a real engine so StartWorkflows gets the slot ReserveSeat already -// acquired, not a second one. everything else delegates, the execution -// path runs exactly like standalone type reservedEngine struct { models.Engine slot engine.WorkflowSlot @@ -48,6 +46,14 @@ func (e *reservedEngine) AcquireWorkflowSlot(ctx context.Context, wid models.Wor return slot, nil } +func (e *reservedEngine) QuotaResources(wf *models.Workflow) quota.Resources { + reporter, ok := e.Engine.(engine.WorkflowQuotaReporter) + if !ok { + return nil + } + return reporter.QuotaResources(wf) +} + func (e *reservedEngine) WorkflowResourceUsage(wf *models.Workflow) (engine.WorkflowResourceUsage, bool) { reporter, ok := e.Engine.(engine.WorkflowResourceUsageReporter) if !ok { diff --git a/spindle/mill/executor/reserved_test.go b/spindle/mill/executor/reserved_test.go index 1c4c1847..da8b5438 100644 --- a/spindle/mill/executor/reserved_test.go +++ b/spindle/mill/executor/reserved_test.go @@ -24,6 +24,7 @@ import ( millproto "tangled.org/core/spindle/mill/proto" millv1 "tangled.org/core/spindle/mill/proto/gen" "tangled.org/core/spindle/models" + "tangled.org/core/spindle/quota" "tangled.org/core/spindle/secrets" ) @@ -118,6 +119,20 @@ func (e *fakeEngine) AcquireWorkflowSlot(ctx context.Context, wid models.Workflo return &fakeSlot{}, nil } +type resourceReportingEngine struct { + *fakeEngine + resources quota.Resources + usage engine.WorkflowResourceUsage +} + +func (e *resourceReportingEngine) QuotaResources(*models.Workflow) quota.Resources { + return e.resources +} + +func (e *resourceReportingEngine) WorkflowResourceUsage(*models.Workflow) (engine.WorkflowResourceUsage, bool) { + return e.usage, true +} + type fakeStep struct{} func (fakeStep) Name() string { return "test" } @@ -149,7 +164,7 @@ func TestNewFailsWhenOutboxCannotInitialize(t *testing.T) { } n := notifier.New() cfg := &config.Config{} - if _, err := New(cfg, nil, d, &n, slog.New(slog.NewTextHandler(io.Discard, nil))); err == nil { + if _, err := New(cfg, nil, d, &n, slog.New(slog.NewTextHandler(io.Discard, nil)), nil); err == nil { t.Fatal("New succeeded with an unavailable outbox database") } } @@ -175,6 +190,38 @@ func TestReservedEngineHandsBackHeldSlotOnce(t *testing.T) { } } +func TestReservedEngineForwardsResourceReporting(t *testing.T) { + inner := &resourceReportingEngine{ + fakeEngine: &fakeEngine{}, + resources: quota.Resources{ + quota.ResourceWorkflows: 1, + quota.ResourceVCPUs: 4, + }, + usage: engine.WorkflowResourceUsage{ + CPUUsec: 1234, + MemoryPeakBytes: 4096, + CgroupAvailable: true, + VolumeAllocatedBytes: 8192, + VolumeAvailable: true, + }, + } + re := newReservedEngine(inner, &fakeSlot{}) + wf := &models.Workflow{} + + resources := re.(engine.WorkflowQuotaReporter).QuotaResources(wf) + if resources[quota.ResourceWorkflows] != 1 || resources[quota.ResourceVCPUs] != 4 { + t.Fatalf("QuotaResources() = %v, want workflow and vCPU resources", resources) + } + + usage, ok := re.(engine.WorkflowResourceUsageReporter).WorkflowResourceUsage(wf) + if !ok { + t.Fatal("WorkflowResourceUsage() unavailable") + } + if usage != inner.usage { + t.Fatalf("WorkflowResourceUsage() = %+v, want %+v", usage, inner.usage) + } +} + func TestHandleCommitIsIdempotent(t *testing.T) { enc := newCaptureEncoder() e := testExecutor(t) @@ -232,7 +279,7 @@ func TestHandleCancelFinalizesExpiredReservation(t *testing.T) { if err := proto.Unmarshal(rows[0].Payload, &entry); err != nil { t.Fatal(err) } - if got := entry.GetAttemptResult().GetStatus(); got != millv1.TerminalStatus_CANCELLED { + if got := entry.GetAttemptResult().GetStatus(); got != millv1.TerminalStatus_TERMINAL_STATUS_CANCELLED { t.Fatalf("cancel terminal = %v, want CANCELLED", got) } } @@ -404,11 +451,64 @@ func TestFinishJobReportsCancelledReservationAsCancelled(t *testing.T) { t.Fatal(err) } got := entry.GetAttemptResult().GetStatus() - if got != millv1.TerminalStatus_CANCELLED { + if got != millv1.TerminalStatus_TERMINAL_STATUS_CANCELLED { t.Fatalf("terminal status = %v, want CANCELLED", got) } } +func TestFinishJobWaitsForEngineCleanup(t *testing.T) { + d := testDB(t) + res := &reservation{ + leaseID: "lease-1", + wid: models.WorkflowId{ + PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, + Name: "build", + }, + runDone: make(chan struct{}), + } + e := &Executor{ + db: d, + l: slog.New(slog.NewTextHandler(io.Discard, nil)), + active: map[string]*reservation{res.leaseID: res}, + maxOutboxBytes: 10 * 1024 * 1024, + } + if err := e.initOutbox(); err != nil { + t.Fatal(err) + } + + done := make(chan error, 1) + go func() { + done <- e.finishJob(res, &tangled.PipelineStatus{ + Pipeline: string(res.wid.PipelineId.AtUri()), + Workflow: res.wid.Name, + Status: string(models.StatusKindSuccess), + }) + }() + + select { + case err := <-done: + t.Fatalf("finishJob returned before engine cleanup: %v", err) + case <-time.After(50 * time.Millisecond): + } + rows, err := d.ListOutboxRows() + if err != nil { + t.Fatal(err) + } + if len(rows) != 0 { + t.Fatalf("terminal rows before engine cleanup = %d, want 0", len(rows)) + } + + close(res.runDone) + select { + case err := <-done: + if err != nil { + t.Fatal(err) + } + case <-time.After(time.Second): + t.Fatal("finishJob did not resume after engine cleanup") + } +} + func TestReplayRejectsMalformedOutboxRow(t *testing.T) { d := testDB(t) e := &Executor{ diff --git a/spindle/mill/integration_test.go b/spindle/mill/integration_test.go index 07983519..6af984f9 100644 --- a/spindle/mill/integration_test.go +++ b/spindle/mill/integration_test.go @@ -34,7 +34,7 @@ func TestEndToEndDummyJob(t *testing.T) { } bn := notifier.New() mill := New(l, Config{LogDir: millDir, ReconnectGrace: time.Minute, BidTimeout: 2 * time.Second}) - mill.Attach(bdb, &bn) + mill.Attach(bdb, &bn, testQuotaManager(t, bdb)) registerTestExecutor(t, bdb, "exec-1", HashToken("test-token"), nil) srv := httptest.NewServer(http.HandlerFunc(mill.HandleExecutorConn)) @@ -69,7 +69,7 @@ func TestEndToEndDummyJob(t *testing.T) { Name: "build", Raw: "steps:\n - name: hello\n command: echo hi\n", } - wf, err := be.InitWorkflow(twf, tangled.Pipeline{TriggerMetadata: &tangled.Pipeline_TriggerMetadata{}}) + wf, err := be.InitWorkflow(twf, testPipeline()) if err != nil { t.Fatalf("InitWorkflow: %v", err) } @@ -140,7 +140,7 @@ func TestExecutorConfiguredLabelsAreStoredOnSession(t *testing.T) { } bn := notifier.New() mill := New(l, Config{LogDir: millDir, ReconnectGrace: time.Minute, BidTimeout: 2 * time.Second}) - mill.Attach(bdb, &bn) + mill.Attach(bdb, &bn, testQuotaManager(t, bdb)) registerTestExecutor(t, bdb, "exec-labels", HashToken("test-token"), []string{"linux", "arm64", "gpu"}) srv := httptest.NewServer(http.HandlerFunc(mill.HandleExecutorConn)) @@ -187,7 +187,7 @@ func TestEndToEndDummyJobUsesRequiredLabelsAcrossExecutors(t *testing.T) { } bn := notifier.New() mill := New(l, Config{LogDir: millDir, ReconnectGrace: time.Minute, BidTimeout: 2 * time.Second}) - mill.Attach(bdb, &bn) + mill.Attach(bdb, &bn, testQuotaManager(t, bdb)) registerTestExecutor(t, bdb, "exec-x86", HashToken("token-x86"), []string{"linux/amd64", "kvm"}) registerTestExecutor(t, bdb, "exec-arm", HashToken("token-arm"), []string{"linux/arm64", "kvm"}) @@ -236,7 +236,7 @@ func TestEndToEndDummyJobUsesRequiredLabelsAcrossExecutors(t *testing.T) { RunsOn: []string{"linux/arm64"}, Raw: "steps:\n - name: hello\n command: echo hi\n", } - wf, err := be.InitWorkflow(twf, tangled.Pipeline{TriggerMetadata: &tangled.Pipeline_TriggerMetadata{}}) + wf, err := be.InitWorkflow(twf, testPipeline()) if err != nil { t.Fatalf("InitWorkflow: %v", err) } diff --git a/spindle/mill/lease.go b/spindle/mill/lease.go index dee1332d..b63fd956 100644 --- a/spindle/mill/lease.go +++ b/spindle/mill/lease.go @@ -4,6 +4,7 @@ import ( "sync" "tangled.org/core/spindle/models" + "tangled.org/core/spindle/quota" millv1 "tangled.org/core/spindle/mill/proto/gen" ) @@ -32,12 +33,23 @@ const ( ) // mill-side fencing token for one placed job +type remoteQuota struct { + key string + resources quota.Resources + committing bool +} + type RemoteLease struct { id string nodeID string epoch string engine string wid models.WorkflowId // job this lease carries, set once placed + // identity comes from the mill's pipeline record + ownerDID string + repoDID string + // keep the winning resources if quota waiting causes a re-bid + resources quota.Resources // restored after a mill restart. no RunStep waits on it, so terminals // and death are authored directly. set before publication, never mutated orphaned bool @@ -46,6 +58,12 @@ type RemoteLease struct { cleanedUp bool cleanupRetry bool + // quotaLease is nil when only its persisted id survived a restart + quotaLease quota.Lease + quotaID string + // executor-side reservations stay attached until commit or release + remoteQuotas map[string]remoteQuota + mu sync.Mutex state leaseState cancel bool @@ -58,12 +76,13 @@ type RemoteLease struct { func newLease(id, nodeID, epoch, engine string) *RemoteLease { return &RemoteLease{ - id: id, - nodeID: nodeID, - epoch: epoch, - engine: engine, - state: leaseReserved, - terminal: make(chan *millv1.AttemptResult, 1), + id: id, + nodeID: nodeID, + epoch: epoch, + engine: engine, + state: leaseReserved, + terminal: make(chan *millv1.AttemptResult, 1), + remoteQuotas: make(map[string]remoteQuota), } } @@ -145,7 +164,7 @@ func (l *RemoteLease) cleanupReady() bool { func (l *RemoteLease) deliverCancelled(reason string) { l.deliverTerminal(&millv1.AttemptResult{ - Status: millv1.TerminalStatus_CANCELLED, + Status: millv1.TerminalStatus_TERMINAL_STATUS_CANCELLED, Error: reason, }) } diff --git a/spindle/mill/mill.go b/spindle/mill/mill.go index f0d7e9f8..58bb3bf6 100644 --- a/spindle/mill/mill.go +++ b/spindle/mill/mill.go @@ -6,6 +6,7 @@ import ( "errors" "fmt" "log/slog" + "maps" "os" "path/filepath" "slices" @@ -13,12 +14,14 @@ import ( "sync" "time" + "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/codes" "tangled.org/core/notifier" "tangled.org/core/spindle/db" "tangled.org/core/spindle/engine" "tangled.org/core/spindle/models" "tangled.org/core/spindle/observability" + "tangled.org/core/spindle/quota" "tangled.org/core/spindle/secrets" "tangled.org/core/tid" @@ -35,8 +38,7 @@ const ( defaultQuarantineStrikes = 3 ) -// marks session-ending errors that are the executor's fault, enough of them -// in a row gets the node quarantined +// repeated protocol errors quarantine the node var errProtocolViolation = errors.New("executor protocol violation") func protoErrf(format string, args ...any) error { @@ -60,17 +62,20 @@ type Mill struct { l *slog.Logger cfg Config - db *db.DB - n *notifier.Notifier - metrics *observability.Metrics - mu sync.Mutex - sessions map[string]*millSession - leases map[string]*RemoteLease - reservations map[string]*RemoteLease - nodeSeqno map[string]uint64 - protoStrikes map[string]int - pending int - changeCh chan struct{} // closed and replaced to wake placement waiters + db *db.DB + n *notifier.Notifier + metrics *observability.Metrics + qm *quota.Manager + mu sync.Mutex + quotaLifecycleMu sync.Mutex + sessions map[string]*millSession + leases map[string]*RemoteLease + reservations map[string]*RemoteLease + nodeSeqno map[string]uint64 + protoStrikes map[string]int + quotaLeases map[string]*RemoteLease + pending int + changeCh chan struct{} leaseSeq uint64 } @@ -106,13 +111,15 @@ func New(l *slog.Logger, cfg Config) *Mill { nodeSeqno: make(map[string]uint64), protoStrikes: make(map[string]int), changeCh: make(chan struct{}), + quotaLeases: make(map[string]*RemoteLease), } } -func (m *Mill) Attach(d *db.DB, n *notifier.Notifier) { +func (m *Mill) Attach(d *db.DB, n *notifier.Notifier, qm *quota.Manager) { m.mu.Lock() m.db = d m.n = n + m.qm = qm m.mu.Unlock() } @@ -148,8 +155,7 @@ func (m *Mill) currentChangeCh() <-chan struct{} { } func (m *Mill) attachSession(sess *millSession) (uint64, bool) { - // read before taking mu so the db never runs under the mill lock, this - // also picks up operator cursor resets without a mill restart + // load db state before taking the mill lock var dbCursor uint64 haveCursor := false if m.db != nil { @@ -163,8 +169,6 @@ func (m *Mill) attachSession(sess *millSession) (uint64, bool) { if old := m.sessions[sess.nodeID]; old != nil { if old.live(m.cfg.ReconnectGrace) { - // a second live session for the same identity is a hijack - // attempt, reject it return 0, false } if old.graceTimer != nil { @@ -215,10 +219,7 @@ func (m *Mill) detachSession(sess *millSession) { m.notifyChange() } -// an executor whose sessions keep dying on protocol errors is stuck, eg. an -// unrecoverable stream gap, and its reconnects keep resetting grace so its -// leases never fail. after enough strikes quarantine it until an operator -// fixes its stream. clean deaths reset the count +// repeated protocol errors quarantine the node until its stream is repaired func (m *Mill) noteSessionError(sess *millSession, err error) { if !errors.Is(err, errProtocolViolation) { m.mu.Lock() @@ -370,8 +371,8 @@ func (m *Mill) place(ctx context.Context, engineName string, wid models.Workflow return nil, err } - // grab the channel before bidding. a change mid-bid closes it, so - // the wait below re-bids right away + // 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) @@ -380,24 +381,20 @@ func (m *Mill) place(ctx context.Context, engineName string, wid models.Workflow return nil, err } if lease != nil { - lease.wid = wid - if err := m.persistLease(lease, leaseRowReserved); err != nil { - m.releaseRemote(lease) - m.mu.Lock() - delete(m.reservations, lease.id) - m.mu.Unlock() + slot, retry, err := m.admit(ctx, engineName, wid, wf, lease) + if err != nil { m.metrics.RecordMillPlacementResult("error") - return nil, fmt.Errorf("persist reserved mill lease: %w", err) + return nil, err } - m.mu.Lock() - delete(m.reservations, lease.id) - m.leases[lease.id] = lease - if st, ok := wf.Data.(*millWorkflowState); ok && st != nil { - st.Lease = lease + if slot != nil { + m.metrics.RecordMillPlacementResult("success") + return slot, nil + } + // quota churn handed the seat back, bid again without waiting + // on a fleet change that already happened + if retry { + continue } - m.mu.Unlock() - m.metrics.RecordMillPlacementResult("success") - return &millSlot{fleet: m, lease: lease}, nil } // no executor available. wait for a change or ctx @@ -414,6 +411,149 @@ func (m *Mill) place(ctx context.Context, engineName string, wid models.Workflow } } +func (m *Mill) admit(ctx context.Context, engineName string, wid models.WorkflowId, wf *models.Workflow, lease *RemoteLease) (engine.WorkflowSlot, bool, error) { + req, ok := m.quotaRequest(wid, wf, lease.resources) + if !ok { + return m.publish(wid, wf, lease, nil) + } + + qlease, res, err := m.qm.TryAcquire(ctx, req) + switch { + case err != nil: + m.releaseSeat(lease) + return nil, false, fmt.Errorf("acquire workflow quota: %w", err) + + case res.Allowed: + return m.publish(wid, wf, lease, qlease) + + case !res.Temporary: + m.releaseSeat(lease) + return nil, false, fmt.Errorf("%w: workflow quota denied: %s", engine.ErrWorkflowFailed, res.Reason) + } + + // don't hold fleet capacity while waiting for quota + m.releaseSeat(lease) + + qlease, err = m.qm.Acquire(ctx, req) + if err != nil { + return nil, false, fmt.Errorf("wait for workflow quota: %w", err) + } + + // a re-bid must match the reserved resources + next, err := m.bid(ctx, engineName, wid, wf) + if err != nil { + _ = m.releaseAbandonedQuota(qlease) + return nil, false, err + } + if next == nil { + _ = m.releaseAbandonedQuota(qlease) + return nil, false, nil + } + if !maps.Equal(next.resources, req.Resources) { + m.l.Warn("re-bid reported different resources than the quota reserved, retrying placement", + "lease", next.id, "node", next.nodeID, "engine", engineName) + m.releaseSeat(next) + _ = m.releaseAbandonedQuota(qlease) + return nil, true, nil + } + return m.publish(wid, wf, next, qlease) +} + +func (m *Mill) releaseAbandonedQuota(qlease quota.Lease) error { + if qlease == nil { + return nil + } + err := qlease.ReleaseWithError() + if err != nil { + dummy := &RemoteLease{ + id: "abandoned-" + qlease.ID(), + quotaID: qlease.ID(), + state: leaseDone, + cleanedUp: false, + } + dummy.finishMu.Lock() + m.scheduleCleanupLocked(dummy) + dummy.finishMu.Unlock() + m.l.Error("abandoned quota release failed, retry scheduled", "quota_id", qlease.ID(), "err", err) + } + return err +} + +// teardown releases the attached quota exactly once +func (m *Mill) publish(wid models.WorkflowId, wf *models.Workflow, lease *RemoteLease, qlease quota.Lease) (engine.WorkflowSlot, bool, error) { + lease.wid = wid + if id, ok := m.quotaIdentity(wf); ok { + lease.ownerDID = id.OwnerDID + lease.repoDID = id.RepoDID + } + if qlease != nil { + lease.quotaLease = qlease + lease.quotaID = qlease.ID() + } + if err := m.persistLease(lease, leaseRowReserved); err != nil { + m.releaseSeat(lease) + // release an unrecorded reservation + releaseErr := m.releaseQuotaLocked(lease) + if releaseErr != nil { + lease.finishMu.Lock() + m.scheduleCleanupLocked(lease) + lease.finishMu.Unlock() + } + return nil, false, errors.Join(fmt.Errorf("persist reserved mill lease: %w", err), releaseErr) + } + m.mu.Lock() + delete(m.reservations, lease.id) + m.leases[lease.id] = lease + if st, ok := wf.Data.(*millWorkflowState); ok && st != nil { + st.Lease = lease + } + m.mu.Unlock() + return &millSlot{fleet: m, lease: lease}, false, nil +} + +func (m *Mill) releaseSeat(lease *RemoteLease) { + m.dropReservation(lease.id) + m.releaseRemote(lease) +} + +func (m *Mill) quotaRequest(wid models.WorkflowId, wf *models.Workflow, resources quota.Resources) (quota.ReserveRequest, bool) { + if m.qm == nil || len(resources) == 0 { + return quota.ReserveRequest{}, false + } + id, ok := m.quotaIdentity(wf) + if !ok { + m.l.Warn("placing workflow without quota because the pipeline has no repo identity", "workflow", wid.String()) + return quota.ReserveRequest{}, false + } + resID := quota.WorkflowReservationID(wf.RunID, id.OwnerDID, id.RepoDID, wid.Knot, wid.Rkey, wid.Name) + return quota.ReserveRequest{ + ID: resID, + Kind: quota.KindWorkflow, + Key: resID, + Identity: id, + Resources: resources, + }, true +} + +func (m *Mill) quotaIdentity(wf *models.Workflow) (quota.Identity, bool) { + st, ok := wf.Data.(*millWorkflowState) + if !ok || st == nil || st.RawPipeline.TriggerMetadata == nil { + return quota.Identity{}, false + } + repo := st.RawPipeline.TriggerMetadata.Repo + if repo == nil { + return quota.Identity{}, false + } + id := quota.Identity{OwnerDID: repo.Did, RepoDID: repo.Did} + if repo.RepoDid != nil && *repo.RepoDid != "" { + id.RepoDID = *repo.RepoDid + } + if !strings.HasPrefix(id.OwnerDID, "did:") || !strings.HasPrefix(id.RepoDID, "did:") { + return quota.Identity{}, false + } + return id, true +} + func (m *Mill) bid(ctx context.Context, engineName string, wid models.WorkflowId, wf *models.Workflow) (*RemoteLease, error) { rawPipeline, rawWorkflow, err := marshalJob(wf) if err != nil { @@ -450,6 +590,32 @@ func (m *Mill) bid(ctx context.Context, engineName string, wid models.WorkflowId m.reservations[leaseID] = lease m.mu.Unlock() bidCtx, span := observability.Tracer().Start(bidCtx, "mill.placement.bid") + if span.IsRecording() { + var ownerDID, repoDID string + if wf != nil { + ownerDID = wf.OwnerDID + repoDID = wf.RepoDID + } + if (ownerDID == "" || repoDID == "") && wf != nil { + if id, ok := m.quotaIdentity(wf); ok { + ownerDID = id.OwnerDID + repoDID = id.RepoDID + } + } + attrs := []attribute.KeyValue{ + attribute.String(observability.WorkflowIDKey, wid.String()), + attribute.String(observability.LeaseIDKey, leaseID), + attribute.String(observability.ExecutorNodeIDKey, sess.nodeID), + attribute.String(observability.PipelineIDKey, wid.PipelineId.AtUri().String()), + } + if ownerDID != "" { + attrs = append(attrs, attribute.String(observability.OwnerDIDKey, ownerDID)) + } + if repoDID != "" { + attrs = append(attrs, attribute.String(observability.RepoDIDKey, repoDID)) + } + span.SetAttributes(attrs...) + } accepted := false defer func() { if accepted { @@ -494,6 +660,15 @@ func (m *Mill) bid(ctx context.Context, engineName string, wid models.WorkflowId results <- bidResult{} return } + // invalid resource claims are protocol errors + if err := quota.ValidateResources(rr.GetQuotaResources()); err != nil { + m.l.WarnContext(bidCtx, "executor reported invalid workflow resources", "node", sess.nodeID, "lease", leaseID, "err", err) + m.dropReservation(leaseID) + m.releaseRemote(lease) + results <- bidResult{} + return + } + lease.resources = maps.Clone(rr.GetQuotaResources()) accepted = true results <- bidResult{sess: sess, lease: lease, rank: rank} } @@ -649,6 +824,21 @@ func (m *Mill) commitAndWait(ctx context.Context, wf *models.Workflow, unlocked lease := st.Lease ctx, span := observability.Tracer().Start(ctx, "mill.commit") + if span.IsRecording() { + attrs := []attribute.KeyValue{ + attribute.String(observability.WorkflowIDKey, lease.wid.String()), + attribute.String(observability.LeaseIDKey, lease.id), + attribute.String(observability.ExecutorNodeIDKey, lease.nodeID), + attribute.String(observability.PipelineIDKey, lease.wid.PipelineId.AtUri().String()), + } + if lease.ownerDID != "" { + attrs = append(attrs, attribute.String(observability.OwnerDIDKey, lease.ownerDID)) + } + if lease.repoDID != "" { + attrs = append(attrs, attribute.String(observability.RepoDIDKey, lease.repoDID)) + } + span.SetAttributes(attrs...) + } defer func() { if err != nil { span.SetStatus(codes.Error, "commit failed") @@ -766,11 +956,11 @@ func (m *Mill) waitCommitRetry(ctx context.Context, lease *RemoteLease) (bool, e func terminalError(status millv1.TerminalStatus) error { switch status { - case millv1.TerminalStatus_SUCCESS: + case millv1.TerminalStatus_TERMINAL_STATUS_SUCCESS: return nil - case millv1.TerminalStatus_TIMEOUT: + case millv1.TerminalStatus_TERMINAL_STATUS_TIMEOUT: return engine.ErrTimedOut - case millv1.TerminalStatus_CANCELLED: + case millv1.TerminalStatus_TERMINAL_STATUS_CANCELLED: return engine.ErrWorkflowCanceled default: return engine.ErrWorkflowFailed @@ -836,6 +1026,7 @@ func (m *Mill) cleanupLeaseLocked(lease *RemoteLease) error { return nil } if m.db != nil { + // delete the row before releasing the reservation if err := m.db.DeleteMillLease(lease.id); err != nil { m.scheduleCleanupLocked(lease) return err @@ -844,11 +1035,67 @@ func (m *Mill) cleanupLeaseLocked(lease *RemoteLease) error { m.mu.Lock() delete(m.leases, lease.id) m.mu.Unlock() + + if err := m.releaseQuotaLocked(lease); err != nil { + m.scheduleCleanupLocked(lease) + return err + } lease.cleanedUp = true m.notifyChange() return nil } +func (m *Mill) releaseQuotaLocked(lease *RemoteLease) error { + if m.qm == nil { + return nil + } + lease.mu.Lock() + workflowID := lease.quotaID + if workflowID == "" && lease.quotaLease != nil { + workflowID = lease.quotaLease.ID() + } + remoteQuotas := maps.Clone(lease.remoteQuotas) + lease.mu.Unlock() + + ctx := context.Background() + var releaseErr error + if workflowID != "" { + if err := m.qm.Release(ctx, workflowID); err != nil { + releaseErr = errors.Join(releaseErr, fmt.Errorf("release workflow quota: %w", err)) + } else { + lease.mu.Lock() + if lease.quotaID == workflowID { + lease.quotaLease = nil + lease.quotaID = "" + } + lease.mu.Unlock() + } + } + for id, rq := range remoteQuotas { + if rq.committing { + if err := m.qm.Commit(ctx, id); err != nil { + releaseErr = errors.Join(releaseErr, fmt.Errorf("commit remote quota %s: %w", id, err)) + continue + } + } else { + if err := m.qm.Release(ctx, id); err != nil { + releaseErr = errors.Join(releaseErr, fmt.Errorf("release remote quota %s: %w", id, err)) + continue + } + } + lease.mu.Lock() + delete(lease.remoteQuotas, id) + lease.mu.Unlock() + + m.quotaLifecycleMu.Lock() + if m.quotaLeases[id] == lease { + delete(m.quotaLeases, id) + } + m.quotaLifecycleMu.Unlock() + } + return releaseErr +} + func (m *Mill) scheduleCleanupLocked(lease *RemoteLease) { if lease.cleanedUp || lease.cleanupRetry { return @@ -1014,13 +1261,13 @@ func (m *Mill) onEventBatch(sess *millSession, batch *millv1.EventBatch) error { ar := entry.GetAttemptResult() statusStr := "success" switch ar.Status { - case millv1.TerminalStatus_SUCCESS: + case millv1.TerminalStatus_TERMINAL_STATUS_SUCCESS: statusStr = "success" - case millv1.TerminalStatus_FAILED: + case millv1.TerminalStatus_TERMINAL_STATUS_FAILED: statusStr = "failed" - case millv1.TerminalStatus_TIMEOUT: + case millv1.TerminalStatus_TERMINAL_STATUS_TIMEOUT: statusStr = "timeout" - case millv1.TerminalStatus_CANCELLED: + case millv1.TerminalStatus_TERMINAL_STATUS_CANCELLED: statusStr = "cancelled" default: return protoErrf("unsupported terminal status %v", ar.Status) @@ -1282,3 +1529,251 @@ func (m *Mill) RegisterMetrics(metrics *observability.Metrics) { ) } +func (m *Mill) authorizeAndLockQuota(s *millSession, reservationID, requestedLeaseID string) (*RemoteLease, error) { + // snapshot the current owner before locking its teardown path + m.quotaLifecycleMu.Lock() + lease, ok := m.quotaLeases[reservationID] + m.quotaLifecycleMu.Unlock() + + if !ok || lease == nil { + return nil, fmt.Errorf("quota reservation %q is not attached to a live lease", reservationID) + } + + // keep teardown from changing the lease during authorization + lease.finishMu.Lock() + + // serialize the final ownership check with quota phase transitions + m.quotaLifecycleMu.Lock() + + if m.quotaLeases[reservationID] != lease { + m.quotaLifecycleMu.Unlock() + lease.finishMu.Unlock() + return nil, fmt.Errorf("quota reservation %q is not attached to a live lease", reservationID) + } + if lease.id != requestedLeaseID { + m.quotaLifecycleMu.Unlock() + lease.finishMu.Unlock() + return nil, fmt.Errorf("quota reservation %q belongs to another workflow", reservationID) + } + if lease.nodeID != s.nodeID { + m.quotaLifecycleMu.Unlock() + lease.finishMu.Unlock() + return nil, fmt.Errorf("quota reservation %q belongs to another executor", reservationID) + } + if lease.epoch != s.epoch { + m.quotaLifecycleMu.Unlock() + lease.finishMu.Unlock() + return nil, fmt.Errorf("quota reservation %q has stale epoch %q (current: %q)", reservationID, lease.epoch, s.epoch) + } + + lease.mu.Lock() + rq, hasQuota := lease.remoteQuotas[reservationID] + active := hasQuota && (lease.state != leaseDone || rq.committing) + lease.mu.Unlock() + if !active { + m.quotaLifecycleMu.Unlock() + lease.finishMu.Unlock() + return nil, fmt.Errorf("quota reservation %q is not active", reservationID) + } + + // the caller releases both locks after its phase transition + return lease, nil +} + +func (m *Mill) handleQuotaRequest(s *millSession, req *millv1.QuotaRequest) { + var ( + res quota.Reservation + err error + ) + switch req.GetOperation() { + case millv1.QuotaOperation_QUOTA_OPERATION_RESERVE: + res, err = m.reserveRemoteQuota(s, req) + case millv1.QuotaOperation_QUOTA_OPERATION_BEGIN_COMMIT, + millv1.QuotaOperation_QUOTA_OPERATION_COMMIT, + millv1.QuotaOperation_QUOTA_OPERATION_RELEASE: + err = m.transitionRemoteQuota(s, req) + default: + err = fmt.Errorf("unsupported quota operation %q", req.GetOperation()) + } + + resp := &millv1.QuotaResponse{ + RequestId: req.GetRequestId(), + ReservationId: res.ID, + Allowed: res.Allowed, + Temporary: res.Temporary, + Reason: res.Reason, + Resource: string(res.Resource), + } + if err != nil { + resp.Error = err.Error() + } + _ = s.send(&millproto.Message{QuotaResp: resp}) +} + +func (m *Mill) reserveRemoteQuota(s *millSession, req *millv1.QuotaRequest) (quota.Reservation, error) { + if m.qm == nil { + return quota.Reservation{}, errors.New("quota manager not attached") + } + if req.GetReservationId() != "" { + return quota.Reservation{}, errors.New("reserve operation cannot name a reservation") + } + kind, resources, err := remoteQuotaPolicy(req) + if err != nil { + return quota.Reservation{}, err + } + + m.mu.Lock() + lease := m.leases[req.GetLeaseId()] + m.mu.Unlock() + if lease == nil || lease.nodeID != s.nodeID || lease.epoch != s.epoch { + return quota.Reservation{}, fmt.Errorf("no live lease %q on this node with matching epoch", req.GetLeaseId()) + } + lease.mu.Lock() + id := quota.Identity{OwnerDID: lease.ownerDID, RepoDID: lease.repoDID} + lease.mu.Unlock() + if id.OwnerDID == "" || id.RepoDID == "" { + return quota.Reservation{}, errors.New("lease has no charged subject") + } + + lease.finishMu.Lock() + defer lease.finishMu.Unlock() + if lease.cleanedUp || lease.getState() == leaseDone { + return quota.Reservation{}, fmt.Errorf("lease %q ended before quota reservation", lease.id) + } + + m.quotaLifecycleMu.Lock() + defer m.quotaLifecycleMu.Unlock() + reservationKey := string(kind) + "\x00" + req.GetKey() + lease.mu.Lock() + var existingID string + mismatch := false + for resID, rq := range lease.remoteQuotas { + if rq.key == reservationKey { + existingID = resID + mismatch = !maps.Equal(rq.resources, resources) + break + } + } + lease.mu.Unlock() + if existingID != "" { + // a repeat is only idempotent for the exact same charge; anything + // else is an executor trying to inflate an existing reservation + if mismatch { + return quota.Reservation{}, fmt.Errorf("cache key %q already reserved with different resources", req.GetKey()) + } + return quota.Reservation{ + ID: existingID, + Allowed: true, + Reason: quota.ReasonWithinLimit, + }, nil + } + + qlease, res, err := m.qm.TryAcquire(context.Background(), quota.ReserveRequest{ + Kind: kind, + Key: req.GetKey(), + Identity: id, + Resources: resources, + }) + if err != nil { + return quota.Reservation{}, err + } + if qlease == nil || res.ID == "" { + return res, nil + } + + existingLease := m.quotaLeases[res.ID] + if existingLease != nil && existingLease != lease { + return quota.Reservation{}, fmt.Errorf("quota reservation %q is already owned by another workflow", res.ID) + } + m.quotaLeases[res.ID] = lease + lease.mu.Lock() + rq := lease.remoteQuotas[res.ID] + rq.key = reservationKey + rq.resources = resources + lease.remoteQuotas[res.ID] = rq + lease.mu.Unlock() + return res, nil +} + +func remoteQuotaPolicy(req *millv1.QuotaRequest) (quota.Kind, quota.Resources, error) { + kind := quota.Kind(req.GetKind()) + if kind != quota.KindNixCache && kind != quota.KindGenericCache { + return "", nil, fmt.Errorf("unsupported remote quota kind %q", req.GetKind()) + } + resources := req.GetResources() + if len(resources) == 0 { + return "", nil, errors.New("empty remote quota resources") + } + if err := quota.ValidateResources(resources); err != nil { + return "", nil, fmt.Errorf("invalid remote quota resources: %w", err) + } + cacheBytes, ok := resources[quota.ResourceCacheStorageBytes] + if !ok || cacheBytes <= 0 || len(resources) != 1 { + return "", nil, errors.New("remote cache quota requires one positive cache_storage_bytes resource") + } + return kind, maps.Clone(resources), nil +} + +func (m *Mill) transitionRemoteQuota(s *millSession, req *millv1.QuotaRequest) error { + reservationID := req.GetReservationId() + if reservationID == "" { + return errors.New("quota transition requires a reservation") + } + if req.GetKind() != "" || req.GetKey() != "" || req.GetResources() != nil { + return errors.New("quota transition cannot include reservation parameters") + } + + // successful terminal transitions are idempotent + if req.GetOperation() != millv1.QuotaOperation_QUOTA_OPERATION_BEGIN_COMMIT { + m.quotaLifecycleMu.Lock() + _, ok := m.quotaLeases[reservationID] + m.quotaLifecycleMu.Unlock() + if !ok { + return nil + } + } + + lease, err := m.authorizeAndLockQuota(s, reservationID, req.GetLeaseId()) + if err != nil { + return err + } + defer lease.finishMu.Unlock() + defer m.quotaLifecycleMu.Unlock() + + if m.qm == nil { + return errors.New("quota manager not attached") + } + switch req.GetOperation() { + case millv1.QuotaOperation_QUOTA_OPERATION_BEGIN_COMMIT: + if err := m.qm.BeginCommit(context.Background(), reservationID); err != nil { + return err + } + lease.mu.Lock() + rq := lease.remoteQuotas[reservationID] + rq.committing = true + lease.remoteQuotas[reservationID] = rq + lease.mu.Unlock() + case millv1.QuotaOperation_QUOTA_OPERATION_COMMIT: + if err := m.qm.Commit(context.Background(), reservationID); err != nil { + return err + } + m.forgetQuotaReservationLocked(reservationID) + case millv1.QuotaOperation_QUOTA_OPERATION_RELEASE: + if err := m.qm.Release(context.Background(), reservationID); err != nil { + return err + } + m.forgetQuotaReservationLocked(reservationID) + } + return nil +} + +func (m *Mill) forgetQuotaReservationLocked(reservationID string) { + lease := m.quotaLeases[reservationID] + delete(m.quotaLeases, reservationID) + if lease == nil { + return + } + lease.mu.Lock() + delete(lease.remoteQuotas, reservationID) + lease.mu.Unlock() +} diff --git a/spindle/mill/mill_test.go b/spindle/mill/mill_test.go index 8a29a88b..928a50c0 100644 --- a/spindle/mill/mill_test.go +++ b/spindle/mill/mill_test.go @@ -29,11 +29,22 @@ func testWorkflow(name string) *models.Workflow { Steps: []models.Step{remoteStep{}}, Data: &millWorkflowState{ RawWorkflow: tangled.Pipeline_Workflow{Name: name}, - RawPipeline: tangled.Pipeline{TriggerMetadata: &tangled.Pipeline_TriggerMetadata{}}, + RawPipeline: testPipeline(), }, } } +func TestEngineInitWorkflowCarriesRepositoryIdentity(t *testing.T) { + m := &Mill{l: slog.New(slog.NewTextHandler(io.Discard, nil))} + wf, err := NewEngine("dummy", m).InitWorkflow(tangled.Pipeline_Workflow{Name: "build"}, testPipeline()) + if err != nil { + t.Fatal(err) + } + if wf.OwnerDID != "did:plc:testowner" || wf.RepoDID != "did:plc:testrepo" { + t.Fatalf("workflow identity = %q/%q, want source owner and repository", wf.OwnerDID, wf.RepoDID) + } +} + func testWorkflowWithRunsOn(name string, runsOn []string) *models.Workflow { wf := testWorkflow(name) wf.Data.(*millWorkflowState).RawWorkflow.RunsOn = runsOn @@ -185,7 +196,7 @@ func TestCommitRetriesAfterSessionCloseBeforeCommitted(t *testing.T) { LeaseId: leaseID, Payload: &millv1.Event_AttemptResult{ AttemptResult: &millv1.AttemptResult{ - Status: millv1.TerminalStatus_SUCCESS, + Status: millv1.TerminalStatus_TERMINAL_STATUS_SUCCESS, }, }, }, @@ -224,7 +235,7 @@ func TestDestroyRunningLeaseDoesNotDropCancelledTerminal(t *testing.T) { } lease.deliverTerminal(&millv1.AttemptResult{ - Status: millv1.TerminalStatus_CANCELLED, + Status: millv1.TerminalStatus_TERMINAL_STATUS_CANCELLED, }) res := <-lease.terminal if err := terminalError(res.Status); !errors.Is(err, engine.ErrWorkflowCanceled) { @@ -396,7 +407,7 @@ func TestCancelledRunningLeaseSurvivesReleaseForReconnectReplay(t *testing.T) { LeaseId: lease.id, Payload: &millv1.Event_AttemptResult{ AttemptResult: &millv1.AttemptResult{ - Status: millv1.TerminalStatus_CANCELLED, + Status: millv1.TerminalStatus_TERMINAL_STATUS_CANCELLED, }, }, }, @@ -543,7 +554,7 @@ func TestGapsAndDuplicates(t *testing.T) { Seqno: 0, LeaseId: owned.id, Payload: &millv1.Event_StatusEvent{ - StatusEvent: &millv1.StatusEvent{Status: millv1.NonterminalStatus_RUNNING}, + StatusEvent: &millv1.StatusEvent{Status: millv1.NonterminalStatus_NONTERMINAL_STATUS_RUNNING}, }, }, }, @@ -559,7 +570,7 @@ func TestGapsAndDuplicates(t *testing.T) { Seqno: 2, LeaseId: owned.id, Payload: &millv1.Event_StatusEvent{ - StatusEvent: &millv1.StatusEvent{Status: millv1.NonterminalStatus_RUNNING}, + StatusEvent: &millv1.StatusEvent{Status: millv1.NonterminalStatus_NONTERMINAL_STATUS_RUNNING}, }, }, }, @@ -598,7 +609,7 @@ func TestAtomicBatchRollback(t *testing.T) { Seqno: 1, LeaseId: owned.id, Payload: &millv1.Event_StatusEvent{ - StatusEvent: &millv1.StatusEvent{Status: millv1.NonterminalStatus_RUNNING}, + StatusEvent: &millv1.StatusEvent{Status: millv1.NonterminalStatus_NONTERMINAL_STATUS_RUNNING}, }, }, }, @@ -645,7 +656,7 @@ func TestTerminalBeforeACK(t *testing.T) { Seqno: 1, LeaseId: owned.id, Payload: &millv1.Event_AttemptResult{ - AttemptResult: &millv1.AttemptResult{Status: millv1.TerminalStatus_SUCCESS}, + AttemptResult: &millv1.AttemptResult{Status: millv1.TerminalStatus_TERMINAL_STATUS_SUCCESS}, }, }, }, diff --git a/spindle/mill/proto/gen/mill.pb.go b/spindle/mill/proto/gen/mill.pb.go index 2783aa99..e8aa399e 100644 --- a/spindle/mill/proto/gen/mill.pb.go +++ b/spindle/mill/proto/gen/mill.pb.go @@ -75,18 +75,18 @@ type NonterminalStatus int32 const ( NonterminalStatus_NONTERMINAL_STATUS_UNSPECIFIED NonterminalStatus = 0 - NonterminalStatus_RUNNING NonterminalStatus = 1 + NonterminalStatus_NONTERMINAL_STATUS_RUNNING NonterminalStatus = 1 ) // Enum value maps for NonterminalStatus. var ( NonterminalStatus_name = map[int32]string{ 0: "NONTERMINAL_STATUS_UNSPECIFIED", - 1: "RUNNING", + 1: "NONTERMINAL_STATUS_RUNNING", } NonterminalStatus_value = map[string]int32{ "NONTERMINAL_STATUS_UNSPECIFIED": 0, - "RUNNING": 1, + "NONTERMINAL_STATUS_RUNNING": 1, } ) @@ -121,27 +121,27 @@ type TerminalStatus int32 const ( TerminalStatus_TERMINAL_STATUS_UNSPECIFIED TerminalStatus = 0 - TerminalStatus_SUCCESS TerminalStatus = 1 - TerminalStatus_FAILED TerminalStatus = 2 - TerminalStatus_TIMEOUT TerminalStatus = 3 - TerminalStatus_CANCELLED TerminalStatus = 4 + TerminalStatus_TERMINAL_STATUS_SUCCESS TerminalStatus = 1 + TerminalStatus_TERMINAL_STATUS_FAILED TerminalStatus = 2 + TerminalStatus_TERMINAL_STATUS_TIMEOUT TerminalStatus = 3 + TerminalStatus_TERMINAL_STATUS_CANCELLED TerminalStatus = 4 ) // Enum value maps for TerminalStatus. var ( TerminalStatus_name = map[int32]string{ 0: "TERMINAL_STATUS_UNSPECIFIED", - 1: "SUCCESS", - 2: "FAILED", - 3: "TIMEOUT", - 4: "CANCELLED", + 1: "TERMINAL_STATUS_SUCCESS", + 2: "TERMINAL_STATUS_FAILED", + 3: "TERMINAL_STATUS_TIMEOUT", + 4: "TERMINAL_STATUS_CANCELLED", } TerminalStatus_value = map[string]int32{ "TERMINAL_STATUS_UNSPECIFIED": 0, - "SUCCESS": 1, - "FAILED": 2, - "TIMEOUT": 3, - "CANCELLED": 4, + "TERMINAL_STATUS_SUCCESS": 1, + "TERMINAL_STATUS_FAILED": 2, + "TERMINAL_STATUS_TIMEOUT": 3, + "TERMINAL_STATUS_CANCELLED": 4, } ) @@ -172,6 +172,61 @@ func (TerminalStatus) EnumDescriptor() ([]byte, []int) { return file_spindle_mill_v1_mill_proto_rawDescGZIP(), []int{2} } +type QuotaOperation int32 + +const ( + QuotaOperation_QUOTA_OPERATION_UNSPECIFIED QuotaOperation = 0 + QuotaOperation_QUOTA_OPERATION_RESERVE QuotaOperation = 1 + QuotaOperation_QUOTA_OPERATION_BEGIN_COMMIT QuotaOperation = 2 + QuotaOperation_QUOTA_OPERATION_COMMIT QuotaOperation = 3 + QuotaOperation_QUOTA_OPERATION_RELEASE QuotaOperation = 4 +) + +// Enum value maps for QuotaOperation. +var ( + QuotaOperation_name = map[int32]string{ + 0: "QUOTA_OPERATION_UNSPECIFIED", + 1: "QUOTA_OPERATION_RESERVE", + 2: "QUOTA_OPERATION_BEGIN_COMMIT", + 3: "QUOTA_OPERATION_COMMIT", + 4: "QUOTA_OPERATION_RELEASE", + } + QuotaOperation_value = map[string]int32{ + "QUOTA_OPERATION_UNSPECIFIED": 0, + "QUOTA_OPERATION_RESERVE": 1, + "QUOTA_OPERATION_BEGIN_COMMIT": 2, + "QUOTA_OPERATION_COMMIT": 3, + "QUOTA_OPERATION_RELEASE": 4, + } +) + +func (x QuotaOperation) Enum() *QuotaOperation { + p := new(QuotaOperation) + *p = x + return p +} + +func (x QuotaOperation) String() string { + return protoimpl.X.EnumStringOf(x.Descriptor(), protoreflect.EnumNumber(x)) +} + +func (QuotaOperation) Descriptor() protoreflect.EnumDescriptor { + return file_spindle_mill_v1_mill_proto_enumTypes[3].Descriptor() +} + +func (QuotaOperation) Type() protoreflect.EnumType { + return &file_spindle_mill_v1_mill_proto_enumTypes[3] +} + +func (x QuotaOperation) Number() protoreflect.EnumNumber { + return protoreflect.EnumNumber(x) +} + +// Deprecated: Use QuotaOperation.Descriptor instead. +func (QuotaOperation) EnumDescriptor() ([]byte, []int) { + return file_spindle_mill_v1_mill_proto_rawDescGZIP(), []int{3} +} + // executor identity, sent on connect type Hello struct { state protoimpl.MessageState `protogen:"open.v1"` @@ -523,13 +578,14 @@ func (x *ReserveSeat) GetTracestate() string { } type ReserveResult struct { - state protoimpl.MessageState `protogen:"open.v1"` - LeaseId string `protobuf:"bytes,1,opt,name=lease_id,json=leaseId,proto3" json:"lease_id,omitempty"` - Accepted bool `protobuf:"varint,2,opt,name=accepted,proto3" json:"accepted,omitempty"` - RejectReason string `protobuf:"bytes,3,opt,name=reject_reason,json=rejectReason,proto3" json:"reject_reason,omitempty"` - RejectClass RejectClass `protobuf:"varint,4,opt,name=reject_class,json=rejectClass,proto3,enum=spindle.mill.v1.RejectClass" json:"reject_class,omitempty"` - unknownFields protoimpl.UnknownFields - sizeCache protoimpl.SizeCache + state protoimpl.MessageState `protogen:"open.v1"` + LeaseId string `protobuf:"bytes,1,opt,name=lease_id,json=leaseId,proto3" json:"lease_id,omitempty"` + Accepted bool `protobuf:"varint,2,opt,name=accepted,proto3" json:"accepted,omitempty"` + RejectReason string `protobuf:"bytes,3,opt,name=reject_reason,json=rejectReason,proto3" json:"reject_reason,omitempty"` + RejectClass RejectClass `protobuf:"varint,4,opt,name=reject_class,json=rejectClass,proto3,enum=spindle.mill.v1.RejectClass" json:"reject_class,omitempty"` + QuotaResources map[string]int64 `protobuf:"bytes,5,rep,name=quota_resources,json=quotaResources,proto3" json:"quota_resources,omitempty" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"varint,2,opt,name=value"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache } func (x *ReserveResult) Reset() { @@ -590,6 +646,13 @@ func (x *ReserveResult) GetRejectClass() RejectClass { return RejectClass_REJECT_CLASS_UNSPECIFIED } +func (x *ReserveResult) GetQuotaResources() map[string]int64 { + if x != nil { + return x.QuotaResources + } + return nil +} + // a single unlocked secret type Secret struct { state protoimpl.MessageState `protogen:"open.v1"` @@ -1352,6 +1415,8 @@ type Message struct { EventBatch *EventBatch `protobuf:"bytes,11,opt,name=event_batch,json=eventBatch,proto3" json:"event_batch,omitempty"` Ack *Ack `protobuf:"bytes,12,opt,name=ack,proto3" json:"ack,omitempty"` LiveLog *LiveLog `protobuf:"bytes,13,opt,name=live_log,json=liveLog,proto3" json:"live_log,omitempty"` + QuotaReq *QuotaRequest `protobuf:"bytes,14,opt,name=quota_req,json=quotaReq,proto3" json:"quota_req,omitempty"` + QuotaResp *QuotaResponse `protobuf:"bytes,15,opt,name=quota_resp,json=quotaResp,proto3" json:"quota_resp,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -1477,6 +1542,208 @@ func (x *Message) GetLiveLog() *LiveLog { return nil } +func (x *Message) GetQuotaReq() *QuotaRequest { + if x != nil { + return x.QuotaReq + } + return nil +} + +func (x *Message) GetQuotaResp() *QuotaResponse { + if x != nil { + return x.QuotaResp + } + return nil +} + +// asks the mill to perform one reservation lifecycle operation +// the mill resolves identity from lease_id so the executor cannot choose who pays +type QuotaRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + RequestId string `protobuf:"bytes,1,opt,name=request_id,json=requestId,proto3" json:"request_id,omitempty"` + Operation QuotaOperation `protobuf:"varint,2,opt,name=operation,proto3,enum=spindle.mill.v1.QuotaOperation" json:"operation,omitempty"` + LeaseId string `protobuf:"bytes,3,opt,name=lease_id,json=leaseId,proto3" json:"lease_id,omitempty"` + ReservationId string `protobuf:"bytes,4,opt,name=reservation_id,json=reservationId,proto3" json:"reservation_id,omitempty"` + Kind string `protobuf:"bytes,5,opt,name=kind,proto3" json:"kind,omitempty"` + Key string `protobuf:"bytes,6,opt,name=key,proto3" json:"key,omitempty"` + Resources map[string]int64 `protobuf:"bytes,7,rep,name=resources,proto3" json:"resources,omitempty" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"varint,2,opt,name=value"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *QuotaRequest) Reset() { + *x = QuotaRequest{} + mi := &file_spindle_mill_v1_mill_proto_msgTypes[20] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *QuotaRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*QuotaRequest) ProtoMessage() {} + +func (x *QuotaRequest) ProtoReflect() protoreflect.Message { + mi := &file_spindle_mill_v1_mill_proto_msgTypes[20] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use QuotaRequest.ProtoReflect.Descriptor instead. +func (*QuotaRequest) Descriptor() ([]byte, []int) { + return file_spindle_mill_v1_mill_proto_rawDescGZIP(), []int{20} +} + +func (x *QuotaRequest) GetRequestId() string { + if x != nil { + return x.RequestId + } + return "" +} + +func (x *QuotaRequest) GetOperation() QuotaOperation { + if x != nil { + return x.Operation + } + return QuotaOperation_QUOTA_OPERATION_UNSPECIFIED +} + +func (x *QuotaRequest) GetLeaseId() string { + if x != nil { + return x.LeaseId + } + return "" +} + +func (x *QuotaRequest) GetReservationId() string { + if x != nil { + return x.ReservationId + } + return "" +} + +func (x *QuotaRequest) GetKind() string { + if x != nil { + return x.Kind + } + return "" +} + +func (x *QuotaRequest) GetKey() string { + if x != nil { + return x.Key + } + return "" +} + +func (x *QuotaRequest) GetResources() map[string]int64 { + if x != nil { + return x.Resources + } + return nil +} + +type QuotaResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + RequestId string `protobuf:"bytes,1,opt,name=request_id,json=requestId,proto3" json:"request_id,omitempty"` + ReservationId string `protobuf:"bytes,2,opt,name=reservation_id,json=reservationId,proto3" json:"reservation_id,omitempty"` + Allowed bool `protobuf:"varint,3,opt,name=allowed,proto3" json:"allowed,omitempty"` + Reason string `protobuf:"bytes,4,opt,name=reason,proto3" json:"reason,omitempty"` + Error string `protobuf:"bytes,5,opt,name=error,proto3" json:"error,omitempty"` + // a denial the fair queue may still grant later + // callers must not treat it as a permanent refusal + Temporary bool `protobuf:"varint,6,opt,name=temporary,proto3" json:"temporary,omitempty"` + Resource string `protobuf:"bytes,7,opt,name=resource,proto3" json:"resource,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *QuotaResponse) Reset() { + *x = QuotaResponse{} + mi := &file_spindle_mill_v1_mill_proto_msgTypes[21] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *QuotaResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*QuotaResponse) ProtoMessage() {} + +func (x *QuotaResponse) ProtoReflect() protoreflect.Message { + mi := &file_spindle_mill_v1_mill_proto_msgTypes[21] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use QuotaResponse.ProtoReflect.Descriptor instead. +func (*QuotaResponse) Descriptor() ([]byte, []int) { + return file_spindle_mill_v1_mill_proto_rawDescGZIP(), []int{21} +} + +func (x *QuotaResponse) GetRequestId() string { + if x != nil { + return x.RequestId + } + return "" +} + +func (x *QuotaResponse) GetReservationId() string { + if x != nil { + return x.ReservationId + } + return "" +} + +func (x *QuotaResponse) GetAllowed() bool { + if x != nil { + return x.Allowed + } + return false +} + +func (x *QuotaResponse) GetReason() string { + if x != nil { + return x.Reason + } + return "" +} + +func (x *QuotaResponse) GetError() string { + if x != nil { + return x.Error + } + return "" +} + +func (x *QuotaResponse) GetTemporary() bool { + if x != nil { + return x.Temporary + } + return false +} + +func (x *QuotaResponse) GetResource() string { + if x != nil { + return x.Resource + } + return "" +} + var File_spindle_mill_v1_mill_proto protoreflect.FileDescriptor const file_spindle_mill_v1_mill_proto_rawDesc = "" + @@ -1515,12 +1782,16 @@ const file_spindle_mill_v1_mill_proto_rawDesc = "" + "\vtraceparent\x18\b \x01(\tR\vtraceparent\x12\x1e\n" + "\n" + "tracestate\x18\t \x01(\tR\n" + - "tracestate\"\xbf\x01\n" + + "tracestate\"\x81\x03\n" + "\rReserveResult\x12\"\n" + "\blease_id\x18\x01 \x01(\tB\a\xbaH\x04r\x02\x10\x01R\aleaseId\x12\x1a\n" + "\baccepted\x18\x02 \x01(\bR\baccepted\x12#\n" + "\rreject_reason\x18\x03 \x01(\tR\frejectReason\x12I\n" + - "\freject_class\x18\x04 \x01(\x0e2\x1c.spindle.mill.v1.RejectClassB\b\xbaH\x05\x82\x01\x02\x10\x01R\vrejectClass\"0\n" + + "\freject_class\x18\x04 \x01(\x0e2\x1c.spindle.mill.v1.RejectClassB\b\xbaH\x05\x82\x01\x02\x10\x01R\vrejectClass\x12}\n" + + "\x0fquota_resources\x18\x05 \x03(\v22.spindle.mill.v1.ReserveResult.QuotaResourcesEntryB \xbaH\x1d\x9a\x01\x1a\x10\x10\"\x06r\x04 \x01(@*\x0e\"\f\x18\x80\x80\x80\x80\x80\x80\x80\x80@(\x00R\x0equotaResources\x1aA\n" + + "\x13QuotaResourcesEntry\x12\x10\n" + + "\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" + + "\x05value\x18\x02 \x01(\x03R\x05value:\x028\x01\"0\n" + "\x06Secret\x12\x10\n" + "\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" + "\x05value\x18\x02 \x01(\tR\x05value\"\xa6\x01\n" + @@ -1569,7 +1840,7 @@ const file_spindle_mill_v1_mill_proto_rawDesc = "" + "\x06events\x18\x02 \x03(\v2\x16.spindle.mill.v1.EventB\b\xbaH\x05\x92\x01\x02\b\x01R\x06events\"D\n" + "\x03Ack\x12\x1d\n" + "\x05epoch\x18\x01 \x01(\tB\a\xbaH\x04r\x02\x10\x01R\x05epoch\x12\x1e\n" + - "\vup_to_seqno\x18\x02 \x01(\x04R\tupToSeqno\"\xb8\a\n" + + "\vup_to_seqno\x18\x02 \x01(\x04R\tupToSeqno\"\xd0\b\n" + "\aMessage\x12,\n" + "\x05hello\x18\x01 \x01(\v2\x16.spindle.mill.v1.HelloR\x05hello\x12/\n" + "\x06resume\x18\x02 \x01(\v2\x17.spindle.mill.v1.ResumeR\x06resume\x12B\n" + @@ -1586,7 +1857,10 @@ const file_spindle_mill_v1_mill_proto_rawDesc = "" + "\vevent_batch\x18\v \x01(\v2\x1b.spindle.mill.v1.EventBatchR\n" + "eventBatch\x12&\n" + "\x03ack\x18\f \x01(\v2\x14.spindle.mill.v1.AckR\x03ack\x123\n" + - "\blive_log\x18\r \x01(\v2\x18.spindle.mill.v1.LiveLogR\aliveLog:\xa5\x01\xbaH\xa1\x01\"\x9e\x01\n" + + "\blive_log\x18\r \x01(\v2\x18.spindle.mill.v1.LiveLogR\aliveLog\x12:\n" + + "\tquota_req\x18\x0e \x01(\v2\x1d.spindle.mill.v1.QuotaRequestR\bquotaReq\x12=\n" + + "\n" + + "quota_resp\x18\x0f \x01(\v2\x1e.spindle.mill.v1.QuotaResponseR\tquotaResp:\xbc\x01\xbaH\xb8\x01\"\xb5\x01\n" + "\x05hello\n" + "\x06resume\n" + "\rnode_snapshot\n" + @@ -1600,21 +1874,54 @@ const file_spindle_mill_v1_mill_proto_rawDesc = "" + "cancel_ack\n" + "\vevent_batch\n" + "\x03ack\n" + - "\blive_log\x10\x01*f\n" + + "\blive_log\n" + + "\tquota_req\n" + + "\n" + + "quota_resp\x10\x01J\x04\b\x10\x10\x16\"\xc2\x03\n" + + "\fQuotaRequest\x12)\n" + + "\n" + + "request_id\x18\x01 \x01(\tB\n" + + "\xbaH\ar\x05\x10\x01\x18\x80\x01R\trequestId\x12I\n" + + "\toperation\x18\x02 \x01(\x0e2\x1f.spindle.mill.v1.QuotaOperationB\n" + + "\xbaH\a\x82\x01\x04\x10\x01 \x00R\toperation\x12%\n" + + "\blease_id\x18\x03 \x01(\tB\n" + + "\xbaH\ar\x05\x10\x01\x18\x80\x01R\aleaseId\x12/\n" + + "\x0ereservation_id\x18\x04 \x01(\tB\b\xbaH\x05r\x03\x18\x80\x01R\rreservationId\x12\x1c\n" + + "\x04kind\x18\x05 \x01(\tB\b\xbaH\x05r\x03\x18\x80\x01R\x04kind\x12\x1a\n" + + "\x03key\x18\x06 \x01(\tB\b\xbaH\x05r\x03\x18\x80\x04R\x03key\x12l\n" + + "\tresources\x18\a \x03(\v2,.spindle.mill.v1.QuotaRequest.ResourcesEntryB \xbaH\x1d\x9a\x01\x1a\x10\x10\"\x06r\x04 \x01(@*\x0e\"\f\x18\x80\x80\x80\x80\x80\x80\x80\x80@(\x00R\tresources\x1a<\n" + + "\x0eResourcesEntry\x12\x10\n" + + "\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" + + "\x05value\x18\x02 \x01(\x03R\x05value:\x028\x01\"\x8a\x02\n" + + "\rQuotaResponse\x12)\n" + + "\n" + + "request_id\x18\x01 \x01(\tB\n" + + "\xbaH\ar\x05\x10\x01\x18\x80\x01R\trequestId\x12/\n" + + "\x0ereservation_id\x18\x02 \x01(\tB\b\xbaH\x05r\x03\x18\x80\x01R\rreservationId\x12\x18\n" + + "\aallowed\x18\x03 \x01(\bR\aallowed\x12\x1f\n" + + "\x06reason\x18\x04 \x01(\tB\a\xbaH\x04r\x02\x18@R\x06reason\x12\x1f\n" + + "\x05error\x18\x05 \x01(\tB\t\xbaH\x06r\x04(\x80\x80\x04R\x05error\x12\x1c\n" + + "\ttemporary\x18\x06 \x01(\bR\ttemporary\x12#\n" + + "\bresource\x18\a \x01(\tB\a\xbaH\x04r\x02(@R\bresource*f\n" + "\vRejectClass\x12\x1c\n" + "\x18REJECT_CLASS_UNSPECIFIED\x10\x00\x12\x1a\n" + "\x16REJECT_CLASS_TRANSIENT\x10\x01\x12\x1d\n" + - "\x19REJECT_CLASS_INCOMPATIBLE\x10\x02*D\n" + + "\x19REJECT_CLASS_INCOMPATIBLE\x10\x02*W\n" + "\x11NonterminalStatus\x12\"\n" + - "\x1eNONTERMINAL_STATUS_UNSPECIFIED\x10\x00\x12\v\n" + - "\aRUNNING\x10\x01*f\n" + + "\x1eNONTERMINAL_STATUS_UNSPECIFIED\x10\x00\x12\x1e\n" + + "\x1aNONTERMINAL_STATUS_RUNNING\x10\x01*\xa6\x01\n" + "\x0eTerminalStatus\x12\x1f\n" + - "\x1bTERMINAL_STATUS_UNSPECIFIED\x10\x00\x12\v\n" + - "\aSUCCESS\x10\x01\x12\n" + - "\n" + - "\x06FAILED\x10\x02\x12\v\n" + - "\aTIMEOUT\x10\x03\x12\r\n" + - "\tCANCELLED\x10\x04B0Z.tangled.org/core/spindle/mill/proto/gen;millv1b\x06proto3" + "\x1bTERMINAL_STATUS_UNSPECIFIED\x10\x00\x12\x1b\n" + + "\x17TERMINAL_STATUS_SUCCESS\x10\x01\x12\x1a\n" + + "\x16TERMINAL_STATUS_FAILED\x10\x02\x12\x1b\n" + + "\x17TERMINAL_STATUS_TIMEOUT\x10\x03\x12\x1d\n" + + "\x19TERMINAL_STATUS_CANCELLED\x10\x04*\xa9\x01\n" + + "\x0eQuotaOperation\x12\x1f\n" + + "\x1bQUOTA_OPERATION_UNSPECIFIED\x10\x00\x12\x1b\n" + + "\x17QUOTA_OPERATION_RESERVE\x10\x01\x12 \n" + + "\x1cQUOTA_OPERATION_BEGIN_COMMIT\x10\x02\x12\x1a\n" + + "\x16QUOTA_OPERATION_COMMIT\x10\x03\x12\x1b\n" + + "\x17QUOTA_OPERATION_RELEASE\x10\x04B0Z.tangled.org/core/spindle/mill/proto/gen;millv1b\x06proto3" var ( file_spindle_mill_v1_mill_proto_rawDescOnce sync.Once @@ -1628,65 +1935,75 @@ func file_spindle_mill_v1_mill_proto_rawDescGZIP() []byte { return file_spindle_mill_v1_mill_proto_rawDescData } -var file_spindle_mill_v1_mill_proto_enumTypes = make([]protoimpl.EnumInfo, 3) -var file_spindle_mill_v1_mill_proto_msgTypes = make([]protoimpl.MessageInfo, 22) +var file_spindle_mill_v1_mill_proto_enumTypes = make([]protoimpl.EnumInfo, 4) +var file_spindle_mill_v1_mill_proto_msgTypes = make([]protoimpl.MessageInfo, 26) var file_spindle_mill_v1_mill_proto_goTypes = []any{ (RejectClass)(0), // 0: spindle.mill.v1.RejectClass (NonterminalStatus)(0), // 1: spindle.mill.v1.NonterminalStatus (TerminalStatus)(0), // 2: spindle.mill.v1.TerminalStatus - (*Hello)(nil), // 3: spindle.mill.v1.Hello - (*Resume)(nil), // 4: spindle.mill.v1.Resume - (*EngineAvailability)(nil), // 5: spindle.mill.v1.EngineAvailability - (*NodeSnapshot)(nil), // 6: spindle.mill.v1.NodeSnapshot - (*ReserveSeat)(nil), // 7: spindle.mill.v1.ReserveSeat - (*ReserveResult)(nil), // 8: spindle.mill.v1.ReserveResult - (*Secret)(nil), // 9: spindle.mill.v1.Secret - (*CommitLease)(nil), // 10: spindle.mill.v1.CommitLease - (*Committed)(nil), // 11: spindle.mill.v1.Committed - (*ReleaseLease)(nil), // 12: spindle.mill.v1.ReleaseLease - (*CancelAttempt)(nil), // 13: spindle.mill.v1.CancelAttempt - (*CancelAck)(nil), // 14: spindle.mill.v1.CancelAck - (*StatusEvent)(nil), // 15: spindle.mill.v1.StatusEvent - (*LogArtifact)(nil), // 16: spindle.mill.v1.LogArtifact - (*AttemptResult)(nil), // 17: spindle.mill.v1.AttemptResult - (*LiveLog)(nil), // 18: spindle.mill.v1.LiveLog - (*Event)(nil), // 19: spindle.mill.v1.Event - (*EventBatch)(nil), // 20: spindle.mill.v1.EventBatch - (*Ack)(nil), // 21: spindle.mill.v1.Ack - (*Message)(nil), // 22: spindle.mill.v1.Message - nil, // 23: spindle.mill.v1.EngineAvailability.LoadEntry - nil, // 24: spindle.mill.v1.NodeSnapshot.EnginesEntry + (QuotaOperation)(0), // 3: spindle.mill.v1.QuotaOperation + (*Hello)(nil), // 4: spindle.mill.v1.Hello + (*Resume)(nil), // 5: spindle.mill.v1.Resume + (*EngineAvailability)(nil), // 6: spindle.mill.v1.EngineAvailability + (*NodeSnapshot)(nil), // 7: spindle.mill.v1.NodeSnapshot + (*ReserveSeat)(nil), // 8: spindle.mill.v1.ReserveSeat + (*ReserveResult)(nil), // 9: spindle.mill.v1.ReserveResult + (*Secret)(nil), // 10: spindle.mill.v1.Secret + (*CommitLease)(nil), // 11: spindle.mill.v1.CommitLease + (*Committed)(nil), // 12: spindle.mill.v1.Committed + (*ReleaseLease)(nil), // 13: spindle.mill.v1.ReleaseLease + (*CancelAttempt)(nil), // 14: spindle.mill.v1.CancelAttempt + (*CancelAck)(nil), // 15: spindle.mill.v1.CancelAck + (*StatusEvent)(nil), // 16: spindle.mill.v1.StatusEvent + (*LogArtifact)(nil), // 17: spindle.mill.v1.LogArtifact + (*AttemptResult)(nil), // 18: spindle.mill.v1.AttemptResult + (*LiveLog)(nil), // 19: spindle.mill.v1.LiveLog + (*Event)(nil), // 20: spindle.mill.v1.Event + (*EventBatch)(nil), // 21: spindle.mill.v1.EventBatch + (*Ack)(nil), // 22: spindle.mill.v1.Ack + (*Message)(nil), // 23: spindle.mill.v1.Message + (*QuotaRequest)(nil), // 24: spindle.mill.v1.QuotaRequest + (*QuotaResponse)(nil), // 25: spindle.mill.v1.QuotaResponse + nil, // 26: spindle.mill.v1.EngineAvailability.LoadEntry + nil, // 27: spindle.mill.v1.NodeSnapshot.EnginesEntry + nil, // 28: spindle.mill.v1.ReserveResult.QuotaResourcesEntry + nil, // 29: spindle.mill.v1.QuotaRequest.ResourcesEntry } var file_spindle_mill_v1_mill_proto_depIdxs = []int32{ - 23, // 0: spindle.mill.v1.EngineAvailability.load:type_name -> spindle.mill.v1.EngineAvailability.LoadEntry - 24, // 1: spindle.mill.v1.NodeSnapshot.engines:type_name -> spindle.mill.v1.NodeSnapshot.EnginesEntry + 26, // 0: spindle.mill.v1.EngineAvailability.load:type_name -> spindle.mill.v1.EngineAvailability.LoadEntry + 27, // 1: spindle.mill.v1.NodeSnapshot.engines:type_name -> spindle.mill.v1.NodeSnapshot.EnginesEntry 0, // 2: spindle.mill.v1.ReserveResult.reject_class:type_name -> spindle.mill.v1.RejectClass - 9, // 3: spindle.mill.v1.CommitLease.secrets:type_name -> spindle.mill.v1.Secret - 1, // 4: spindle.mill.v1.StatusEvent.status:type_name -> spindle.mill.v1.NonterminalStatus - 2, // 5: spindle.mill.v1.AttemptResult.status:type_name -> spindle.mill.v1.TerminalStatus - 16, // 6: spindle.mill.v1.AttemptResult.log_artifact:type_name -> spindle.mill.v1.LogArtifact - 15, // 7: spindle.mill.v1.Event.status_event:type_name -> spindle.mill.v1.StatusEvent - 17, // 8: spindle.mill.v1.Event.attempt_result:type_name -> spindle.mill.v1.AttemptResult - 19, // 9: spindle.mill.v1.EventBatch.events:type_name -> spindle.mill.v1.Event - 3, // 10: spindle.mill.v1.Message.hello:type_name -> spindle.mill.v1.Hello - 4, // 11: spindle.mill.v1.Message.resume:type_name -> spindle.mill.v1.Resume - 6, // 12: spindle.mill.v1.Message.node_snapshot:type_name -> spindle.mill.v1.NodeSnapshot - 7, // 13: spindle.mill.v1.Message.reserve_seat:type_name -> spindle.mill.v1.ReserveSeat - 8, // 14: spindle.mill.v1.Message.reserve_result:type_name -> spindle.mill.v1.ReserveResult - 10, // 15: spindle.mill.v1.Message.commit_lease:type_name -> spindle.mill.v1.CommitLease - 11, // 16: spindle.mill.v1.Message.committed:type_name -> spindle.mill.v1.Committed - 12, // 17: spindle.mill.v1.Message.release_lease:type_name -> spindle.mill.v1.ReleaseLease - 13, // 18: spindle.mill.v1.Message.cancel_attempt:type_name -> spindle.mill.v1.CancelAttempt - 14, // 19: spindle.mill.v1.Message.cancel_ack:type_name -> spindle.mill.v1.CancelAck - 20, // 20: spindle.mill.v1.Message.event_batch:type_name -> spindle.mill.v1.EventBatch - 21, // 21: spindle.mill.v1.Message.ack:type_name -> spindle.mill.v1.Ack - 18, // 22: spindle.mill.v1.Message.live_log:type_name -> spindle.mill.v1.LiveLog - 5, // 23: spindle.mill.v1.NodeSnapshot.EnginesEntry.value:type_name -> spindle.mill.v1.EngineAvailability - 24, // [24:24] is the sub-list for method output_type - 24, // [24:24] is the sub-list for method input_type - 24, // [24:24] is the sub-list for extension type_name - 24, // [24:24] is the sub-list for extension extendee - 0, // [0:24] is the sub-list for field type_name + 28, // 3: spindle.mill.v1.ReserveResult.quota_resources:type_name -> spindle.mill.v1.ReserveResult.QuotaResourcesEntry + 10, // 4: spindle.mill.v1.CommitLease.secrets:type_name -> spindle.mill.v1.Secret + 1, // 5: spindle.mill.v1.StatusEvent.status:type_name -> spindle.mill.v1.NonterminalStatus + 2, // 6: spindle.mill.v1.AttemptResult.status:type_name -> spindle.mill.v1.TerminalStatus + 17, // 7: spindle.mill.v1.AttemptResult.log_artifact:type_name -> spindle.mill.v1.LogArtifact + 16, // 8: spindle.mill.v1.Event.status_event:type_name -> spindle.mill.v1.StatusEvent + 18, // 9: spindle.mill.v1.Event.attempt_result:type_name -> spindle.mill.v1.AttemptResult + 20, // 10: spindle.mill.v1.EventBatch.events:type_name -> spindle.mill.v1.Event + 4, // 11: spindle.mill.v1.Message.hello:type_name -> spindle.mill.v1.Hello + 5, // 12: spindle.mill.v1.Message.resume:type_name -> spindle.mill.v1.Resume + 7, // 13: spindle.mill.v1.Message.node_snapshot:type_name -> spindle.mill.v1.NodeSnapshot + 8, // 14: spindle.mill.v1.Message.reserve_seat:type_name -> spindle.mill.v1.ReserveSeat + 9, // 15: spindle.mill.v1.Message.reserve_result:type_name -> spindle.mill.v1.ReserveResult + 11, // 16: spindle.mill.v1.Message.commit_lease:type_name -> spindle.mill.v1.CommitLease + 12, // 17: spindle.mill.v1.Message.committed:type_name -> spindle.mill.v1.Committed + 13, // 18: spindle.mill.v1.Message.release_lease:type_name -> spindle.mill.v1.ReleaseLease + 14, // 19: spindle.mill.v1.Message.cancel_attempt:type_name -> spindle.mill.v1.CancelAttempt + 15, // 20: spindle.mill.v1.Message.cancel_ack:type_name -> spindle.mill.v1.CancelAck + 21, // 21: spindle.mill.v1.Message.event_batch:type_name -> spindle.mill.v1.EventBatch + 22, // 22: spindle.mill.v1.Message.ack:type_name -> spindle.mill.v1.Ack + 19, // 23: spindle.mill.v1.Message.live_log:type_name -> spindle.mill.v1.LiveLog + 24, // 24: spindle.mill.v1.Message.quota_req:type_name -> spindle.mill.v1.QuotaRequest + 25, // 25: spindle.mill.v1.Message.quota_resp:type_name -> spindle.mill.v1.QuotaResponse + 3, // 26: spindle.mill.v1.QuotaRequest.operation:type_name -> spindle.mill.v1.QuotaOperation + 29, // 27: spindle.mill.v1.QuotaRequest.resources:type_name -> spindle.mill.v1.QuotaRequest.ResourcesEntry + 6, // 28: spindle.mill.v1.NodeSnapshot.EnginesEntry.value:type_name -> spindle.mill.v1.EngineAvailability + 29, // [29:29] is the sub-list for method output_type + 29, // [29:29] is the sub-list for method input_type + 29, // [29:29] is the sub-list for extension type_name + 29, // [29:29] is the sub-list for extension extendee + 0, // [0:29] is the sub-list for field type_name } func init() { file_spindle_mill_v1_mill_proto_init() } @@ -1703,8 +2020,8 @@ func file_spindle_mill_v1_mill_proto_init() { File: protoimpl.DescBuilder{ GoPackagePath: reflect.TypeOf(x{}).PkgPath(), RawDescriptor: unsafe.Slice(unsafe.StringData(file_spindle_mill_v1_mill_proto_rawDesc), len(file_spindle_mill_v1_mill_proto_rawDesc)), - NumEnums: 3, - NumMessages: 22, + NumEnums: 4, + NumMessages: 26, NumExtensions: 0, NumServices: 0, }, diff --git a/spindle/mill/proto/protocol.go b/spindle/mill/proto/protocol.go index 2e1d1402..b9db8cbb 100644 --- a/spindle/mill/proto/protocol.go +++ b/spindle/mill/proto/protocol.go @@ -17,7 +17,7 @@ import ( ) const ( - ProtocolVersion = 1 + ProtocolVersion = 3 // generous vs agentproto's 1 MiB. a ReserveSeat carries the raw pipeline and // workflow JSON, and streamed log lines can be chunky MaxMessageBytes = 8 * 1024 * 1024 diff --git a/spindle/mill/proto/protocol_test.go b/spindle/mill/proto/protocol_test.go index 0e0b9637..b1d860e6 100644 --- a/spindle/mill/proto/protocol_test.go +++ b/spindle/mill/proto/protocol_test.go @@ -3,9 +3,11 @@ package millproto import ( "bytes" "encoding/binary" + "strings" "testing" millv1 "tangled.org/core/spindle/mill/proto/gen" + "tangled.org/core/spindle/quota" ) func TestEncodeDecodeRoundTrip(t *testing.T) { @@ -129,7 +131,7 @@ func TestValidationRules(t *testing.T) { LeaseId: "lease-1", Payload: &millv1.Event_StatusEvent{ StatusEvent: &millv1.StatusEvent{ - Status: millv1.NonterminalStatus_RUNNING, + Status: millv1.NonterminalStatus_NONTERMINAL_STATUS_RUNNING, }, }, }, @@ -185,7 +187,7 @@ func TestValidationRules(t *testing.T) { LeaseId: "lease-1", Payload: &millv1.Event_StatusEvent{ StatusEvent: &millv1.StatusEvent{ - Status: millv1.NonterminalStatus_RUNNING, + Status: millv1.NonterminalStatus_NONTERMINAL_STATUS_RUNNING, }, }, }, @@ -205,7 +207,7 @@ func TestValidationRules(t *testing.T) { LeaseId: "lease-1", Payload: &millv1.Event_AttemptResult{ AttemptResult: &millv1.AttemptResult{ - Status: millv1.TerminalStatus_SUCCESS, + Status: millv1.TerminalStatus_TERMINAL_STATUS_SUCCESS, }, }, }, @@ -214,6 +216,138 @@ func TestValidationRules(t *testing.T) { }, wantErr: false, }, + { + name: "valid quota reserve request", + msg: &Message{ + QuotaReq: &millv1.QuotaRequest{ + RequestId: "req-1", + Operation: millv1.QuotaOperation_QUOTA_OPERATION_RESERVE, + LeaseId: "lease-1", + Kind: "nix_cache", + Key: "key-1", + Resources: map[string]int64{"cache_storage_bytes": 100}, + }, + }, + wantErr: false, + }, + { + name: "invalid quota request - empty request_id", + msg: &Message{ + QuotaReq: &millv1.QuotaRequest{ + RequestId: "", + Operation: millv1.QuotaOperation_QUOTA_OPERATION_RESERVE, + LeaseId: "lease-1", + }, + }, + wantErr: true, + }, + { + name: "invalid quota request - unspecified operation", + msg: &Message{ + QuotaReq: &millv1.QuotaRequest{ + RequestId: "req-1", + Operation: millv1.QuotaOperation_QUOTA_OPERATION_UNSPECIFIED, + LeaseId: "lease-1", + }, + }, + wantErr: true, + }, + { + name: "invalid quota request - empty lease_id", + msg: &Message{ + QuotaReq: &millv1.QuotaRequest{ + RequestId: "req-1", + Operation: millv1.QuotaOperation_QUOTA_OPERATION_RESERVE, + LeaseId: "", + }, + }, + wantErr: true, + }, + { + name: "invalid quota request - negative resource amount", + msg: &Message{ + QuotaReq: &millv1.QuotaRequest{ + RequestId: "req-1", + Operation: millv1.QuotaOperation_QUOTA_OPERATION_RESERVE, + LeaseId: "lease-1", + Resources: map[string]int64{"cache_storage_bytes": -1}, + }, + }, + wantErr: true, + }, + { + name: "valid quota transition request with reservation_id", + msg: &Message{ + QuotaReq: &millv1.QuotaRequest{ + RequestId: "req-1", + Operation: millv1.QuotaOperation_QUOTA_OPERATION_BEGIN_COMMIT, + LeaseId: "lease-1", + ReservationId: "res-1", + }, + }, + wantErr: false, + }, + { + name: "valid reserve result reporting workflow resources", + msg: &Message{ + ReserveResult: &millv1.ReserveResult{ + LeaseId: "lease-1", + Accepted: true, + QuotaResources: map[string]int64{ + "workflows": 1, + "memory_mib": 2048, + "vcpus": 2, + "disk_mib": 8192, + }, + }, + }, + wantErr: false, + }, + { + name: "valid reserve result reporting custom resources", + msg: &Message{ + ReserveResult: &millv1.ReserveResult{ + LeaseId: "lease-1", + Accepted: true, + QuotaResources: map[string]int64{ + "gpu_count": 4, + }, + }, + }, + wantErr: false, + }, + { + name: "invalid reserve result - negative resource amount", + msg: &Message{ + ReserveResult: &millv1.ReserveResult{ + LeaseId: "lease-1", + Accepted: true, + QuotaResources: map[string]int64{"memory_mib": -1}, + }, + }, + wantErr: true, + }, + { + name: "valid quota response with resource field shape", + msg: &Message{ + QuotaResp: &millv1.QuotaResponse{ + RequestId: "req-1", + ReservationId: "res-1", + Allowed: true, + Resource: "cache_storage_bytes", + }, + }, + wantErr: false, + }, + { + name: "invalid quota response - empty request_id", + msg: &Message{ + QuotaResp: &millv1.QuotaResponse{ + RequestId: "", + }, + }, + wantErr: true, + }, } for _, tc := range tests { @@ -225,3 +359,28 @@ func TestValidationRules(t *testing.T) { }) } } + +func TestQuotaResourceMapBounds(t *testing.T) { + tooMany := make(map[string]int64) + for i := range quota.MaxResourcePairs + 1 { + tooMany[strings.Repeat("x", i+1)] = 1 + } + for _, resources := range []map[string]int64{ + {"": 1}, + {strings.Repeat("é", 33): 1}, + {"large": quota.MaxResourceAmount + 1}, + tooMany, + } { + msg := &Message{QuotaReq: &millv1.QuotaRequest{ + RequestId: "req-1", + Operation: millv1.QuotaOperation_QUOTA_OPERATION_RESERVE, + LeaseId: "lease-1", + Kind: "custom", + Key: "key", + Resources: resources, + }} + if err := validator.Validate(msg); err == nil { + t.Fatalf("expected invalid resources: %v", resources) + } + } +} diff --git a/spindle/mill/proto/spindle/mill/v1/mill.proto b/spindle/mill/proto/spindle/mill/v1/mill.proto index 13d9c9c1..ed93751b 100644 --- a/spindle/mill/proto/spindle/mill/v1/mill.proto +++ b/spindle/mill/proto/spindle/mill/v1/mill.proto @@ -62,6 +62,11 @@ message ReserveResult { bool accepted = 2; string reject_reason = 3; RejectClass reject_class = 4 [(buf.validate.field).enum.defined_only = true]; + map quota_resources = 5 [(buf.validate.field).map = { + max_pairs: 16, + keys: {string: {min_bytes: 1, max_bytes: 64}}, + values: {int64: {gte: 0, lte: 4611686018427387904}} + }]; } // a single unlocked secret @@ -99,15 +104,15 @@ message CancelAck { enum NonterminalStatus { NONTERMINAL_STATUS_UNSPECIFIED = 0; - RUNNING = 1; + NONTERMINAL_STATUS_RUNNING = 1; } enum TerminalStatus { TERMINAL_STATUS_UNSPECIFIED = 0; - SUCCESS = 1; - FAILED = 2; - TIMEOUT = 3; - CANCELLED = 4; + TERMINAL_STATUS_SUCCESS = 1; + TERMINAL_STATUS_FAILED = 2; + TERMINAL_STATUS_TIMEOUT = 3; + TERMINAL_STATUS_CANCELLED = 4; } message StatusEvent { @@ -170,7 +175,8 @@ message Message { fields: [ "hello", "resume", "node_snapshot", "reserve_seat", "reserve_result", "commit_lease", "committed", "release_lease", "cancel_attempt", - "cancel_ack", "event_batch", "ack", "live_log" + "cancel_ack", "event_batch", "ack", "live_log", + "quota_req", "quota_resp" ], required: true }; @@ -189,5 +195,44 @@ message Message { Ack ack = 12; LiveLog live_log = 13; + QuotaRequest quota_req = 14; + QuotaResponse quota_resp = 15; + + reserved 16 to 21; +} + +enum QuotaOperation { + QUOTA_OPERATION_UNSPECIFIED = 0; + QUOTA_OPERATION_RESERVE = 1; + QUOTA_OPERATION_BEGIN_COMMIT = 2; + QUOTA_OPERATION_COMMIT = 3; + QUOTA_OPERATION_RELEASE = 4; +} + +message QuotaRequest { + string request_id = 1 [(buf.validate.field).string = {min_len: 1, max_len: 128}]; + QuotaOperation operation = 2 [(buf.validate.field).enum = { + defined_only: true + not_in: 0 + }]; + string lease_id = 3 [(buf.validate.field).string = {min_len: 1, max_len: 128}]; + string reservation_id = 4 [(buf.validate.field).string.max_len = 128]; + string kind = 5 [(buf.validate.field).string.max_len = 128]; + string key = 6 [(buf.validate.field).string.max_len = 512]; + map resources = 7 [(buf.validate.field).map = { + max_pairs: 16, + keys: {string: {min_bytes: 1, max_bytes: 64}}, + values: {int64: {gte: 0, lte: 4611686018427387904}} + }]; +} + +message QuotaResponse { + string request_id = 1 [(buf.validate.field).string = {min_len: 1, max_len: 128}]; + string reservation_id = 2 [(buf.validate.field).string.max_len = 128]; + bool allowed = 3; + string reason = 4 [(buf.validate.field).string.max_len = 64]; + string error = 5 [(buf.validate.field).string.max_bytes = 65536]; + bool temporary = 6; + string resource = 7 [(buf.validate.field).string.max_bytes = 64]; } diff --git a/spindle/mill/quota_admission_test.go b/spindle/mill/quota_admission_test.go new file mode 100644 index 00000000..fa9df45a --- /dev/null +++ b/spindle/mill/quota_admission_test.go @@ -0,0 +1,1245 @@ +package mill + +import ( + "context" + "errors" + "io" + "log/slog" + "path/filepath" + "strings" + "sync" + "testing" + "time" + + "tangled.org/core/notifier" + "tangled.org/core/spindle/db" + "tangled.org/core/spindle/models" + "tangled.org/core/spindle/quota" + + millproto "tangled.org/core/spindle/mill/proto" + millv1 "tangled.org/core/spindle/mill/proto/gen" +) + +type scriptedStore struct { + mu sync.Mutex + verdicts []quota.Reservation + reserved []quota.ReserveRequest + released []string + commits []string + seq int + releaseErr error + releaseFailCount int + commitErr error + commitFailCount int +} + +func (s *scriptedStore) Reserve(_ context.Context, req quota.ReserveRequest) (quota.Reservation, error) { + s.mu.Lock() + defer s.mu.Unlock() + s.reserved = append(s.reserved, req) + var res quota.Reservation + if s.seq < len(s.verdicts) { + res = s.verdicts[s.seq] + } else if len(s.verdicts) > 0 { + res = s.verdicts[len(s.verdicts)-1] + } else { + res = quota.Reservation{ID: "res-1", Allowed: true, Reason: quota.ReasonUnlimited} + } + s.seq++ + return res, nil +} + +func (s *scriptedStore) BeginCommit(context.Context, string) error { return nil } + +func (s *scriptedStore) Commit(_ context.Context, id string) error { + s.mu.Lock() + defer s.mu.Unlock() + s.commits = append(s.commits, id) + if id == "res-cache-1" && s.commitFailCount > 0 { + s.commitFailCount-- + return s.commitErr + } + return nil +} + +func (s *scriptedStore) Release(_ context.Context, id string) error { + s.mu.Lock() + defer s.mu.Unlock() + s.released = append(s.released, id) + if id == "res-cache-1" && s.releaseFailCount > 0 { + s.releaseFailCount-- + return s.releaseErr + } + return nil +} + +func (s *scriptedStore) Recover(context.Context, []string) error { return nil } +func (s *scriptedStore) SetLimit(context.Context, string, string, int64) error { + return nil +} +func (s *scriptedStore) GetLimit(context.Context, string, string) (*quota.Limit, error) { + return nil, nil +} + +func (s *scriptedStore) UnsetLimit(context.Context, string, string) error { + return nil +} +func (s *scriptedStore) ListLimits(context.Context) ([]quota.Limit, error) { return nil, nil } +func (s *scriptedStore) ListUsage(context.Context) ([]quota.Usage, error) { return nil, nil } +func (s *scriptedStore) MetricsSnapshot(context.Context) (quota.MetricsSnapshot, error) { + return quota.MetricsSnapshot{}, nil +} + +func (s *scriptedStore) snapshot() ([]quota.ReserveRequest, []string) { + s.mu.Lock() + defer s.mu.Unlock() + return append([]quota.ReserveRequest(nil), s.reserved...), append([]string(nil), s.released...) +} + +type quotaCandidate struct { + sess *millSession + mu sync.Mutex + reserved int + released int +} + +func (c *quotaCandidate) counts() (int, int) { + c.mu.Lock() + defer c.mu.Unlock() + return c.reserved, c.released +} + +func addQuotaCandidate(t *testing.T, m *Mill, nodeID string, report func(attempt int) map[string]int64) *quotaCandidate { + t.Helper() + cand := "aCandidate{} + var sess *millSession + sess = addCandidateSession(t, m, nodeID, nil, 0, scriptedEncoder(func(msg *millproto.Message) error { + if rs := msg.GetReserveSeat(); rs != nil { + cand.mu.Lock() + cand.reserved++ + attempt := cand.reserved + cand.mu.Unlock() + sess.deliver(rs.GetLeaseId(), &millproto.Message{ReserveResult: &millv1.ReserveResult{ + LeaseId: rs.GetLeaseId(), + Accepted: true, + QuotaResources: report(attempt), + }}) + return nil + } + if msg.GetReleaseLease() != nil { + cand.mu.Lock() + cand.released++ + cand.mu.Unlock() + } + return nil + })) + cand.sess = sess + return cand +} + +func cacheReservationID(t *testing.T, lease *RemoteLease) string { + t.Helper() + lease.mu.Lock() + defer lease.mu.Unlock() + if len(lease.remoteQuotas) != 1 { + t.Fatalf("remote reservation IDs = %v, want exactly one", lease.remoteQuotas) + } + for id := range lease.remoteQuotas { + return id + } + return "" +} + +func quotaMill(t *testing.T, store quota.Store) (*Mill, *db.DB) { + t.Helper() + bdb, err := db.Make(context.Background(), filepath.Join(t.TempDir(), "mill.db")) + if err != nil { + t.Fatalf("db.Make: %v", err) + } + t.Cleanup(func() { bdb.Close() }) + n := notifier.New() + m := New(slog.New(slog.NewTextHandler(io.Discard, nil)), Config{ + BidTimeout: time.Second, + ReconnectGrace: time.Minute, + }) + qm := quota.NewManager(store, 10*time.Millisecond, nil) + t.Cleanup(qm.Close) + m.Attach(bdb, &n, qm) + return m, bdb +} + +func fixedReport(vcpus, memoryMiB int64) func(int) map[string]int64 { + return func(int) map[string]int64 { + return map[string]int64{"workflows": 1, "vcpus": vcpus, "memory_mib": memoryMiB} + } +} + +func TestPlaceChargesReportedResourcesToPipelineIdentity(t *testing.T) { + store := &scriptedStore{} + m, _ := quotaMill(t, store) + addQuotaCandidate(t, m, "node-a", fixedReport(2, 2048)) + + wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "knot.test", Rkey: "rkey1"}, Name: "build"} + slot, err := m.place(context.Background(), "dummy", wid, testWorkflow("build")) + if err != nil { + t.Fatalf("place: %v", err) + } + defer slot.Release() + + reserved, _ := store.snapshot() + if len(reserved) != 1 { + t.Fatalf("store saw %d reservations, want 1", len(reserved)) + } + req := reserved[0] + if req.Kind != quota.KindWorkflow { + t.Errorf("kind = %q, want workflow", req.Kind) + } + // charge the pipeline identity, not executor input + if req.Identity.OwnerDID != "did:plc:testowner" || req.Identity.RepoDID != "did:plc:testrepo" { + t.Errorf("identity = %+v, want the pipeline's owner and repo", req.Identity) + } + want := quota.Resources{ + quota.ResourceWorkflows: 1, + quota.ResourceVCPUs: 2, + quota.ResourceMemoryMiB: 2048, + } + if len(req.Resources) != len(want) { + t.Fatalf("resources = %v, want %v", req.Resources, want) + } + for res, amount := range want { + if req.Resources[res] != amount { + t.Errorf("resources[%s] = %d, want %d", res, req.Resources[res], amount) + } + } +} + +func TestPlaceReleasesSeatBeforeWaitingOnQuota(t *testing.T) { + store := &scriptedStore{verdicts: []quota.Reservation{ + {Allowed: false, Temporary: true, Reason: quota.ReasonRepoLimit, Resource: "vcpus"}, + {ID: "res-queued", Allowed: true, Reason: quota.ReasonWithinLimit}, + }} + m, _ := quotaMill(t, store) + cand := addQuotaCandidate(t, m, "node-a", fixedReport(2, 2048)) + + wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "knot.test", Rkey: "rkey1"}, Name: "build"} + slot, err := m.place(context.Background(), "dummy", wid, testWorkflow("build")) + if err != nil { + t.Fatalf("place: %v", err) + } + defer slot.Release() + + reserved, released := cand.counts() + if reserved != 2 { + t.Errorf("executor was asked %d times, want 2 (deferred bid then re-bid)", reserved) + } + if released != 1 { + t.Errorf("executor seat releases = %d, want exactly 1 while quota waited", released) + } + + m.mu.Lock() + liveLeases := len(m.leases) + liveReservations := len(m.reservations) + m.mu.Unlock() + if liveLeases != 1 || liveReservations != 0 { + t.Errorf("leases=%d reservations=%d, want one published lease and no dangling seat", liveLeases, liveReservations) + } + + if _, storeReleased := store.snapshot(); len(storeReleased) != 0 { + t.Errorf("quota releases during a successful placement = %v, want none", storeReleased) + } +} + +func TestPlaceFailsAndReleasesSeatOnPermanentDenial(t *testing.T) { + store := &scriptedStore{verdicts: []quota.Reservation{ + {Allowed: false, Temporary: false, Reason: "request_exceeds_limit", Resource: "memory_mib"}, + }} + m, _ := quotaMill(t, store) + cand := addQuotaCandidate(t, m, "node-a", fixedReport(2, 2048)) + + wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "knot.test", Rkey: "rkey1"}, Name: "build"} + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + slot, err := m.place(ctx, "dummy", wid, testWorkflow("build")) + if err == nil { + slot.Release() + t.Fatal("place succeeded despite a permanent quota denial") + } + if !strings.Contains(err.Error(), "request_exceeds_limit") { + t.Errorf("place error = %v, want the denial reason", err) + } + if _, released := cand.counts(); released != 1 { + t.Errorf("executor seat releases = %d, want 1", released) + } + m.mu.Lock() + live := len(m.leases) + len(m.reservations) + m.mu.Unlock() + if live != 0 { + t.Errorf("mill kept %d lease records after a permanent denial", live) + } +} + +func TestPlaceRejectsRebidWithDifferentResources(t *testing.T) { + store := &scriptedStore{verdicts: []quota.Reservation{ + {Allowed: false, Temporary: true, Reason: quota.ReasonRepoLimit, Resource: "vcpus"}, + {ID: "res-queued", Allowed: true, Reason: quota.ReasonWithinLimit}, + {ID: "res-second", Allowed: true, Reason: quota.ReasonWithinLimit}, + }} + m, _ := quotaMill(t, store) + cand := addQuotaCandidate(t, m, "node-a", func(attempt int) map[string]int64 { + if attempt >= 2 { + return map[string]int64{"workflows": 1, "vcpus": 8, "memory_mib": 2048} + } + return map[string]int64{"workflows": 1, "vcpus": 2, "memory_mib": 2048} + }) + + wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "knot.test", Rkey: "rkey1"}, Name: "build"} + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + slot, err := m.place(ctx, "dummy", wid, testWorkflow("build")) + if err != nil { + t.Fatalf("place: %v", err) + } + defer slot.Release() + + _, released := store.snapshot() + if len(released) == 0 { + t.Fatal("mismatched re-bid kept its quota reservation") + } + if released[0] != "res-queued" { + t.Errorf("released %v, want the mismatched reservation res-queued first", released) + } + if _, seatReleases := cand.counts(); seatReleases < 2 { + t.Errorf("executor seat releases = %d, want at least 2 (deferred bid and mismatched re-bid)", seatReleases) + } +} + +func TestSlotReleaseReleasesQuotaExactlyOnce(t *testing.T) { + store := &scriptedStore{verdicts: []quota.Reservation{ + {ID: "res-1", Allowed: true, Reason: quota.ReasonWithinLimit}, + }} + m, _ := quotaMill(t, store) + addQuotaCandidate(t, m, "node-a", fixedReport(2, 2048)) + + wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "knot.test", Rkey: "rkey1"}, Name: "build"} + slot, err := m.place(context.Background(), "dummy", wid, testWorkflow("build")) + if err != nil { + t.Fatalf("place: %v", err) + } + + // teardown must release once + slot.Release() + slot.Release() + _, released := store.snapshot() + if len(released) != 1 || released[0] != "res-1" { + t.Fatalf("quota releases = %v, want exactly [res-1]", released) + } +} + +func TestFailedLeaseAfterDisconnectReleasesQuotaOnce(t *testing.T) { + store := &scriptedStore{verdicts: []quota.Reservation{ + {ID: "res-1", Allowed: true, Reason: quota.ReasonWithinLimit}, + }} + m, _ := quotaMill(t, store) + cand := addQuotaCandidate(t, m, "node-a", fixedReport(2, 2048)) + + wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "knot.test", Rkey: "rkey1"}, Name: "build"} + slot, err := m.place(context.Background(), "dummy", wid, testWorkflow("build")) + if err != nil { + t.Fatalf("place: %v", err) + } + + m.mu.Lock() + var lease *RemoteLease + for _, l := range m.leases { + lease = l + } + m.mu.Unlock() + if lease == nil { + t.Fatal("no published lease") + } + + m.detachSession(cand.sess) + m.failLeasesAfterGrace(cand.sess) + + _, released := store.snapshot() + if len(released) != 1 { + t.Fatalf("quota releases after executor loss = %v, want exactly one", released) + } + slot.Release() + if _, after := store.snapshot(); len(after) != 1 { + t.Fatalf("quota releases after slot unwind = %v, want still exactly one", after) + } +} + +func TestCacheReserveChargesTheLeaseSubject(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, _ := quotaMill(t, store) + cand := addQuotaCandidate(t, m, "node-a", fixedReport(2, 2048)) + + wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "knot.test", Rkey: "rkey1"}, Name: "build"} + slot, err := m.place(context.Background(), "dummy", wid, testWorkflow("build")) + if err != nil { + t.Fatalf("place: %v", err) + } + defer slot.Release() + + m.mu.Lock() + var leaseID string + for id := range m.leases { + leaseID = id + } + m.mu.Unlock() + + m.handleQuotaRequest(cand.sess, &millv1.QuotaRequest{ + RequestId: "req-1", + Operation: millv1.QuotaOperation_QUOTA_OPERATION_RESERVE, + Kind: string(quota.KindNixCache), + Key: "narhash", + LeaseId: leaseID, + Resources: map[string]int64{"cache_storage_bytes": 4096}, + }) + + reserved, _ := store.snapshot() + if len(reserved) != 2 { + t.Fatalf("store saw %d reservations, want workflow then cache", len(reserved)) + } + cache := reserved[1] + if cache.Kind != quota.KindNixCache || cache.Key != "narhash" { + t.Errorf("cache reservation = %+v, want a nix_cache charge for narhash", cache) + } + if cache.Identity.OwnerDID != "did:plc:testowner" || cache.Identity.RepoDID != "did:plc:testrepo" { + t.Errorf("cache identity = %+v, want the lease's subject", cache.Identity) + } + if cache.Resources[quota.ResourceCacheStorageBytes] != 4096 { + t.Errorf("cache resources = %v, want 4096 cache bytes", cache.Resources) + } +} + +func TestCacheReserveRejectsUnknownLeaseAndWorkflowKind(t *testing.T) { + store := &scriptedStore{} + m, _ := quotaMill(t, store) + cand := addQuotaCandidate(t, m, "node-a", fixedReport(2, 2048)) + + var replies []*millv1.QuotaResponse + sess := cand.sess + sess.enc = scriptedEncoder(func(msg *millproto.Message) error { + if resp := msg.GetQuotaResp(); resp != nil { + replies = append(replies, resp) + } + return nil + }) + + m.handleQuotaRequest(sess, &millv1.QuotaRequest{ + RequestId: "req-1", + Operation: millv1.QuotaOperation_QUOTA_OPERATION_RESERVE, + Kind: string(quota.KindNixCache), + Key: "narhash", + LeaseId: "nope", + Resources: map[string]int64{"cache_storage_bytes": 10}, + }) + m.handleQuotaRequest(sess, &millv1.QuotaRequest{ + RequestId: "req-2", + Operation: millv1.QuotaOperation_QUOTA_OPERATION_RESERVE, + Kind: string(quota.KindWorkflow), + Key: "narhash", + LeaseId: "nope", + Resources: map[string]int64{"cache_storage_bytes": 10}, + }) + m.handleQuotaRequest(sess, &millv1.QuotaRequest{ + RequestId: "req-3", + Operation: millv1.QuotaOperation_QUOTA_OPERATION_RESERVE, + Kind: string(quota.KindNixCache), + Key: "narhash", + LeaseId: "nope", + Resources: map[string]int64{"gpu_slices": 1}, + }) + + if len(replies) != 3 { + t.Fatalf("got %d replies, want 3", len(replies)) + } + for i, resp := range replies { + if resp.GetAllowed() || resp.GetError() == "" { + t.Errorf("reply %d = %+v, want a refusal with an error", i, resp) + } + } + if !strings.Contains(replies[2].GetError(), "cache_storage_bytes") { + t.Fatalf("non-cache resource rejection = %q", replies[2].GetError()) + } + if reserved, _ := store.snapshot(); len(reserved) != 0 { + t.Errorf("store was asked to reserve %v, want nothing", reserved) + } +} + +func TestCacheReserveReplayRequiresSameResources(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, _ := quotaMill(t, store) + cand := addQuotaCandidate(t, m, "node-a", fixedReport(2, 2048)) + + wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "knot.test", Rkey: "rkey1"}, Name: "build"} + slot, err := m.place(context.Background(), "dummy", wid, testWorkflow("build")) + if err != nil { + t.Fatalf("place: %v", err) + } + defer slot.Release() + + m.mu.Lock() + var leaseID string + for id := range m.leases { + leaseID = id + } + m.mu.Unlock() + + reserve := func(requestID string, bytes int64) (quota.Reservation, error) { + return m.reserveRemoteQuota(cand.sess, &millv1.QuotaRequest{ + RequestId: requestID, + Operation: millv1.QuotaOperation_QUOTA_OPERATION_RESERVE, + Kind: string(quota.KindNixCache), + Key: "narhash", + LeaseId: leaseID, + Resources: map[string]int64{"cache_storage_bytes": bytes}, + }) + } + + first, err := reserve("req-1", 4096) + if err != nil || !first.Allowed { + t.Fatalf("first reserve = %+v, %v", first, err) + } + + // an executor retry with the same charge replays idempotently + retry, err := reserve("req-2", 4096) + if err != nil || !retry.Allowed || retry.ID != first.ID { + t.Fatalf("idempotent retry = %+v, %v, want reservation %q", retry, err, first.ID) + } + + // but a repeat with an inflated charge must not ride the existing reservation + if inflated, err := reserve("req-3", 1<<40); err == nil || inflated.Allowed { + t.Fatalf("inflated replay = %+v, %v, want a refusal", inflated, err) + } + + if reserved, _ := store.snapshot(); len(reserved) != 2 { + t.Fatalf("store saw %d reservations, want workflow then one cache charge", len(reserved)) + } +} + +func TestCacheReserveDuplicateQueuesUntilFirstSettles(t *testing.T) { + // identical charges dedup to one reservation row, so the manager makes a + // concurrent second reserve wait instead of handing out the same id; the + // executor re-polls and lands once the first workflow settles + store := &scriptedStore{verdicts: []quota.Reservation{ + {ID: "res-wf-a", Allowed: true, Reason: quota.ReasonWithinLimit}, + {ID: "res-wf-b", Allowed: true, Reason: quota.ReasonWithinLimit}, + {ID: "res-cache", Allowed: true, Reason: quota.ReasonWithinLimit}, + // the deduped row is still active, so the second reserve sees the + // same id and the manager turns it into a queued verdict + {ID: "res-cache", Allowed: true, Reason: quota.ReasonWithinLimit}, + // after the first commit the content is allocated, so a retry is + // allowed without a fresh reservation + {Allowed: true, Reason: quota.ReasonWithinLimit}, + }} + m, _ := quotaMill(t, store) + candA := addQuotaCandidate(t, m, "node-a", fixedReport(2, 2048)) + candB := addQuotaCandidate(t, m, "node-b", fixedReport(2, 2048)) + + sessFor := func(l *RemoteLease) *millSession { + if l.nodeID == "node-a" { + return candA.sess + } + return candB.sess + } + placeLease := func(name string) *RemoteLease { + wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "knot.test", Rkey: "rkey1"}, Name: name} + slot, err := m.place(context.Background(), "dummy", wid, testWorkflow(name)) + if err != nil { + t.Fatalf("place %s: %v", name, err) + } + t.Cleanup(slot.Release) + return slot.(*millSlot).lease + } + leaseA := placeLease("build-a") + leaseB := placeLease("build-b") + + reserve := func(l *RemoteLease, reqID string) (quota.Reservation, error) { + return m.reserveRemoteQuota(sessFor(l), &millv1.QuotaRequest{ + RequestId: reqID, + Operation: millv1.QuotaOperation_QUOTA_OPERATION_RESERVE, + Kind: string(quota.KindNixCache), + Key: "narhash", + LeaseId: l.id, + Resources: map[string]int64{"cache_storage_bytes": 4096}, + }) + } + + first, err := reserve(leaseA, "req-a") + if err != nil || !first.Allowed || first.ID == "" { + t.Fatalf("first reserve = %+v, %v", first, err) + } + + dup, err := reserve(leaseB, "req-b") + if err != nil { + t.Fatalf("duplicate reserve errored instead of waiting: %v", err) + } + if dup.Allowed || !dup.Temporary { + t.Fatalf("duplicate reserve = %+v, want a temporary queued verdict", dup) + } + + if err := m.transitionRemoteQuota(sessFor(leaseA), &millv1.QuotaRequest{ + RequestId: "commit-a", + Operation: millv1.QuotaOperation_QUOTA_OPERATION_COMMIT, + LeaseId: leaseA.id, + ReservationId: first.ID, + }); err != nil { + t.Fatalf("commit: %v", err) + } + + retry, err := reserve(leaseB, "req-b2") + if err != nil || !retry.Allowed { + t.Fatalf("retry after settle = %+v, %v", retry, err) + } + if retry.ID == first.ID { + t.Fatalf("retry reused settled reservation %q", retry.ID) + } + + m.quotaLifecycleMu.Lock() + _, tracked := m.quotaLeases[first.ID] + m.quotaLifecycleMu.Unlock() + if tracked { + t.Fatal("settled reservation is still tracked") + } +} + +func TestBidRejectsOutOfRangeReportedResources(t *testing.T) { + store := &scriptedStore{} + m, _ := quotaMill(t, store) + // two workflows for one seat is impossible, so the bid must die rather + // than become a quota charge + cand := addQuotaCandidate(t, m, "node-a", func(int) map[string]int64 { + return map[string]int64{"workflows": -1, "vcpus": 2} + }) + + wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "knot.test", Rkey: "rkey1"}, Name: "build"} + ctx, cancel := context.WithTimeout(context.Background(), 300*time.Millisecond) + defer cancel() + slot, err := m.place(ctx, "dummy", wid, testWorkflow("build")) + if err == nil { + slot.Release() + t.Fatal("place accepted a bid claiming two workflows for one seat") + } + if !errors.Is(err, context.DeadlineExceeded) { + t.Errorf("place error = %v, want the placement to keep waiting for a sane bid", err) + } + if reserved, _ := store.snapshot(); len(reserved) != 0 { + t.Errorf("store was asked to charge %v, want nothing", reserved) + } + if _, released := cand.counts(); released == 0 { + t.Error("the invalid bid's seat was never released") + } +} + +func TestLiveQuotaReservationIDsSurviveRestart(t *testing.T) { + store := &scriptedStore{verdicts: []quota.Reservation{ + {ID: "res-workflow", Allowed: true, Reason: quota.ReasonWithinLimit}, + }} + m, bdb := quotaMill(t, store) + addQuotaCandidate(t, m, "node-a", fixedReport(2, 2048)) + + wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "knot.test", Rkey: "rkey1"}, Name: "build"} + if _, err := m.place(context.Background(), "dummy", wid, testWorkflow("build")); err != nil { + t.Fatalf("place: %v", err) + } + + if ids := m.LiveQuotaReservationIDs(); len(ids) != 1 || quota.StorageReservationID(ids[0]) != "res-workflow" { + t.Fatalf("live ids = %v, want one fenced ID for res-workflow", ids) + } + + // a fresh mill over the same db is the restart. the charge is only + // recoverable because the lease row carries its id + n := notifier.New() + restarted := New(slog.New(slog.NewTextHandler(io.Discard, nil)), Config{ReconnectGrace: time.Minute}) + 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) + } + + ids := restarted.LiveQuotaReservationIDs() + if len(ids) != 1 || quota.StorageReservationID(ids[0]) != "res-workflow" { + t.Fatalf("restored live ids = %v, want one fenced ID for res-workflow", ids) + } + + restarted.mu.Lock() + var lease *RemoteLease + for _, l := range restarted.leases { + lease = l + } + restarted.mu.Unlock() + if lease.ownerDID != "did:plc:testowner" || lease.repoDID != "did:plc:testrepo" { + t.Errorf("restored subject = %q/%q, want the persisted owner and repo", lease.ownerDID, lease.repoDID) + } + + // the restored lease is an orphan nobody reclaims, and failing it hands + // the charge back exactly once + reason := "executor did not reconnect after mill restart" + if err := restarted.finishOrphan(lease, string(models.StatusKindFailed), &reason, nil); err != nil { + t.Fatalf("finishOrphan: %v", err) + } + if _, released := store.snapshot(); len(released) != 1 || released[0] != "res-workflow" { + t.Fatalf("releases after orphan sweep = %v, want exactly [res-workflow]", released) + } +} + +func TestPersistedLeaseCarriesQuotaReservation(t *testing.T) { + store := &scriptedStore{verdicts: []quota.Reservation{ + {ID: "res-workflow", Allowed: true, Reason: quota.ReasonWithinLimit}, + }} + m, bdb := quotaMill(t, store) + addQuotaCandidate(t, m, "node-a", fixedReport(2, 2048)) + + wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "knot.test", Rkey: "rkey1"}, Name: "build"} + slot, err := m.place(context.Background(), "dummy", wid, testWorkflow("build")) + if err != nil { + t.Fatalf("place: %v", err) + } + defer slot.Release() + + rows, err := bdb.ListMillLeases() + if err != nil { + t.Fatalf("ListMillLeases: %v", err) + } + if len(rows) != 1 { + t.Fatalf("persisted %d leases, want 1", len(rows)) + } + if quota.StorageReservationID(rows[0].QuotaReservationID) != "res-workflow" { + t.Errorf("persisted reservation id = %q, want a fenced ID for res-workflow", rows[0].QuotaReservationID) + } + if rows[0].OwnerDID != "did:plc:testowner" || rows[0].RepoDID != "did:plc:testrepo" { + t.Errorf("persisted subject = %q/%q, want the pipeline's owner and repo", rows[0].OwnerDID, rows[0].RepoDID) + } +} + +func TestUnchargeablePlacementStillRuns(t *testing.T) { + store := &scriptedStore{} + m, _ := quotaMill(t, store) + // an engine with no quota reporter charges nothing, the fleet must still + // place its work + addQuotaCandidate(t, m, "node-a", func(int) map[string]int64 { return nil }) + + wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "knot.test", Rkey: "rkey1"}, Name: "build"} + slot, err := m.place(context.Background(), "dummy", wid, testWorkflow("build")) + if err != nil { + t.Fatalf("place: %v", err) + } + defer slot.Release() + + if reserved, _ := store.snapshot(); len(reserved) != 0 { + t.Errorf("store was asked to charge %v, want nothing", reserved) + } + if ids := m.LiveQuotaReservationIDs(); len(ids) != 0 { + t.Errorf("live ids = %v, want none", ids) + } +} + +func TestCacheReservationDuplicateIDRace(t *testing.T) { + store := &scriptedStore{ + verdicts: []quota.Reservation{ + {ID: "res-workflow", Allowed: true, Reason: quota.ReasonWithinLimit}, + {ID: "res-cache-1", Allowed: true, Reason: quota.ReasonWithinLimit}, + {ID: "res-cache-1", Allowed: true, Reason: quota.ReasonWithinLimit}, + }, + } + m, _ := quotaMill(t, store) + cand := addQuotaCandidate(t, m, "node-a", fixedReport(2, 2048)) + + wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "knot.test", Rkey: "rkey1"}, Name: "build"} + slot, err := m.place(context.Background(), "dummy", wid, testWorkflow("build")) + if err != nil { + t.Fatalf("place: %v", err) + } + + m.mu.Lock() + var leaseID string + for id := range m.leases { + leaseID = id + } + lease := m.leases[leaseID] + m.mu.Unlock() + + var replies []*millproto.Message + cand.sess.enc = scriptedEncoder(func(msg *millproto.Message) error { + replies = append(replies, msg) + return nil + }) + + m.handleQuotaRequest(cand.sess, &millv1.QuotaRequest{ + RequestId: "req-1", + Operation: millv1.QuotaOperation_QUOTA_OPERATION_RESERVE, + Kind: string(quota.KindNixCache), + Key: "narhash", + LeaseId: leaseID, + Resources: map[string]int64{"cache_storage_bytes": 4096}, + }) + reservationID := cacheReservationID(t, lease) + + m.handleQuotaRequest(cand.sess, &millv1.QuotaRequest{ + RequestId: "req-2", + Operation: millv1.QuotaOperation_QUOTA_OPERATION_BEGIN_COMMIT, + ReservationId: reservationID, + LeaseId: leaseID, + }) + + m.handleQuotaRequest(cand.sess, &millv1.QuotaRequest{ + RequestId: "req-3", + Operation: millv1.QuotaOperation_QUOTA_OPERATION_RESERVE, + Kind: string(quota.KindNixCache), + Key: "narhash", + LeaseId: leaseID, + Resources: map[string]int64{"cache_storage_bytes": 4096}, + }) + + lease.mu.Lock() + rq, ok := lease.remoteQuotas[reservationID] + lease.mu.Unlock() + + if !ok { + t.Errorf("reservation %q was not found in remoteQuotas", reservationID) + } + if !rq.committing { + t.Errorf("publishing ID was not tracked as committing in remoteQuotas") + } + + slot.Release() + + _, released := store.snapshot() + for _, rel := range released { + if rel == "res-cache-1" { + t.Errorf("publishing row was released during terminal cleanup") + } + } +} + +func TestCacheReservationCommitAfterTeardown(t *testing.T) { + store := &scriptedStore{ + verdicts: []quota.Reservation{ + {ID: "res-workflow", Allowed: true, Reason: quota.ReasonWithinLimit}, + {ID: "res-cache-1", Allowed: true, Reason: quota.ReasonWithinLimit}, + }, + } + m, _ := quotaMill(t, store) + cand := addQuotaCandidate(t, m, "node-a", fixedReport(2, 2048)) + + wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "knot.test", Rkey: "rkey1"}, Name: "build"} + slot, err := m.place(context.Background(), "dummy", wid, testWorkflow("build")) + if err != nil { + t.Fatalf("place: %v", err) + } + + m.mu.Lock() + var leaseID string + for id := range m.leases { + leaseID = id + } + lease := m.leases[leaseID] + m.mu.Unlock() + + var replies []*millproto.Message + cand.sess.enc = scriptedEncoder(func(msg *millproto.Message) error { + replies = append(replies, msg) + return nil + }) + + m.handleQuotaRequest(cand.sess, &millv1.QuotaRequest{ + RequestId: "req-1", + Operation: millv1.QuotaOperation_QUOTA_OPERATION_RESERVE, + Kind: string(quota.KindNixCache), + Key: "narhash", + LeaseId: leaseID, + Resources: map[string]int64{"cache_storage_bytes": 4096}, + }) + reservationID := cacheReservationID(t, lease) + + m.handleQuotaRequest(cand.sess, &millv1.QuotaRequest{ + RequestId: "req-2", + Operation: millv1.QuotaOperation_QUOTA_OPERATION_BEGIN_COMMIT, + ReservationId: reservationID, + LeaseId: leaseID, + }) + + slot.Release() + + m.mu.Lock() + _, inLeases := m.leases[leaseID] + m.mu.Unlock() + + if inLeases { + t.Errorf("lease still in m.leases after teardown") + } + + store.mu.Lock() + committedBefore := len(store.commits) > 0 && store.commits[0] == "res-cache-1" + store.mu.Unlock() + + if !committedBefore { + t.Errorf("store did not receive teardown commit for res-cache-1") + } + + m.handleQuotaRequest(cand.sess, &millv1.QuotaRequest{ + RequestId: "req-3", + Operation: millv1.QuotaOperation_QUOTA_OPERATION_COMMIT, + ReservationId: reservationID, + LeaseId: leaseID, + }) + + for _, reply := range replies { + if resp := reply.GetQuotaResp(); resp != nil && resp.GetRequestId() == "req-3" { + if resp.GetError() != "" { + t.Errorf("commit failed after teardown: %s", resp.GetError()) + } + } + } + + m.quotaLifecycleMu.Lock() + _, inQuotaLeasesAfter := m.quotaLeases[reservationID] + m.quotaLifecycleMu.Unlock() + + if inQuotaLeasesAfter { + t.Errorf("reservation was not dropped from m.quotaLeases after commit finished") + } +} + +func TestCacheReservationFailedReleaseRetry(t *testing.T) { + store := &scriptedStore{ + verdicts: []quota.Reservation{ + {ID: "res-workflow", Allowed: true, Reason: quota.ReasonWithinLimit}, + {ID: "res-cache-1", Allowed: true, Reason: quota.ReasonWithinLimit}, + }, + releaseErr: errors.New("db error"), + releaseFailCount: 1, + } + m, _ := quotaMill(t, store) + cand := addQuotaCandidate(t, m, "node-a", fixedReport(2, 2048)) + + wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "knot.test", Rkey: "rkey1"}, Name: "build"} + slot, err := m.place(context.Background(), "dummy", wid, testWorkflow("build")) + if err != nil { + t.Fatalf("place: %v", err) + } + defer slot.Release() + + m.mu.Lock() + var leaseID string + for id := range m.leases { + leaseID = id + } + lease := m.leases[leaseID] + m.mu.Unlock() + + m.handleQuotaRequest(cand.sess, &millv1.QuotaRequest{ + RequestId: "req-1", + Operation: millv1.QuotaOperation_QUOTA_OPERATION_RESERVE, + Kind: string(quota.KindNixCache), + Key: "narhash", + LeaseId: leaseID, + Resources: map[string]int64{"cache_storage_bytes": 4096}, + }) + reservationID := cacheReservationID(t, lease) + + err = m.cleanupLease(lease) + if err == nil { + t.Errorf("cleanupLease succeeded despite failed release") + } + + lease.mu.Lock() + _, retained := lease.remoteQuotas[reservationID] + retryScheduled := lease.cleanupRetry + lease.mu.Unlock() + + if !retained { + t.Errorf("res-cache-1 was not retained in lease.remoteQuotas on failure") + } + if !retryScheduled { + t.Errorf("retry was not scheduled on failure") + } + + err = m.cleanupLease(lease) + if err != nil { + t.Errorf("cleanupLease failed on retry: %v", err) + } + + lease.mu.Lock() + _, retainedAfter := lease.remoteQuotas[reservationID] + lease.mu.Unlock() + + if retainedAfter { + t.Errorf("res-cache-1 was not removed from lease.remoteQuotas after successful retry") + } +} + +func TestCacheReservationStaleEpochRejected(t *testing.T) { + store := &scriptedStore{ + verdicts: []quota.Reservation{ + {ID: "res-workflow", Allowed: true, Reason: quota.ReasonWithinLimit}, + {ID: "res-cache-1", Allowed: true, Reason: quota.ReasonWithinLimit}, + }, + } + m, _ := quotaMill(t, store) + cand := addQuotaCandidate(t, m, "node-a", fixedReport(2, 2048)) + + wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "knot.test", Rkey: "rkey1"}, Name: "build"} + slot, err := m.place(context.Background(), "dummy", wid, testWorkflow("build")) + if err != nil { + t.Fatalf("place: %v", err) + } + defer slot.Release() + + m.mu.Lock() + var leaseID string + for id := range m.leases { + leaseID = id + } + lease := m.leases[leaseID] + m.mu.Unlock() + + var replies []*millproto.Message + cand.sess.enc = scriptedEncoder(func(msg *millproto.Message) error { + replies = append(replies, msg) + return nil + }) + + m.handleQuotaRequest(cand.sess, &millv1.QuotaRequest{ + RequestId: "req-1", + Operation: millv1.QuotaOperation_QUOTA_OPERATION_RESERVE, + Kind: string(quota.KindNixCache), + Key: "narhash", + LeaseId: leaseID, + Resources: map[string]int64{"cache_storage_bytes": 4096}, + }) + reservationID := cacheReservationID(t, lease) + + staleSess := newSession(cand.sess.nodeID, "stale-epoch-999", cand.sess.labels, cand.sess.enc, cand.sess.l) + + m.handleQuotaRequest(cand.sess, &millv1.QuotaRequest{ + RequestId: "req-wrong-lease", + Operation: millv1.QuotaOperation_QUOTA_OPERATION_BEGIN_COMMIT, + ReservationId: reservationID, + LeaseId: "another-workflow", + }) + + m.handleQuotaRequest(staleSess, &millv1.QuotaRequest{ + RequestId: "req-2", + Operation: millv1.QuotaOperation_QUOTA_OPERATION_BEGIN_COMMIT, + ReservationId: reservationID, + LeaseId: leaseID, + }) + + var foundError, foundOwnerError bool + for _, reply := range replies { + if resp := reply.GetQuotaResp(); resp != nil { + if strings.Contains(resp.GetError(), "stale epoch") { + foundError = true + } + if strings.Contains(resp.GetError(), "another workflow") { + foundOwnerError = true + } + } + } + if !foundError { + t.Errorf("stale epoch BeginCommit was not rejected with 'stale epoch' error") + } + if !foundOwnerError { + t.Errorf("wrong-lease BeginCommit was not rejected with ownership error") + } +} +func TestCacheReservationTeardownCommitRetry(t *testing.T) { + store := &scriptedStore{ + verdicts: []quota.Reservation{ + {ID: "res-workflow", Allowed: true, Reason: quota.ReasonWithinLimit}, + {ID: "res-cache-1", Allowed: true, Reason: quota.ReasonWithinLimit}, + }, + commitErr: errors.New("db error"), + commitFailCount: 1, + } + m, _ := quotaMill(t, store) + cand := addQuotaCandidate(t, m, "node-a", fixedReport(2, 2048)) + + wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "knot.test", Rkey: "rkey1"}, Name: "build"} + slot, err := m.place(context.Background(), "dummy", wid, testWorkflow("build")) + if err != nil { + t.Fatalf("place: %v", err) + } + + m.mu.Lock() + var leaseID string + for id := range m.leases { + leaseID = id + } + lease := m.leases[leaseID] + m.mu.Unlock() + + m.handleQuotaRequest(cand.sess, &millv1.QuotaRequest{ + RequestId: "req-1", + Operation: millv1.QuotaOperation_QUOTA_OPERATION_RESERVE, + Kind: string(quota.KindNixCache), + Key: "narhash", + LeaseId: leaseID, + Resources: map[string]int64{"cache_storage_bytes": 4096}, + }) + reservationID := cacheReservationID(t, lease) + + m.handleQuotaRequest(cand.sess, &millv1.QuotaRequest{ + RequestId: "req-2", + Operation: millv1.QuotaOperation_QUOTA_OPERATION_BEGIN_COMMIT, + ReservationId: reservationID, + LeaseId: leaseID, + }) + + slot.Release() + + lease.mu.Lock() + rq, retained := lease.remoteQuotas[reservationID] + lease.mu.Unlock() + + if !retained || !rq.committing { + t.Errorf("res-cache-1 was not retained as committing in lease.remoteQuotas on commit failure") + } + + m.quotaLifecycleMu.Lock() + _, inQuotaLeases := m.quotaLeases[reservationID] + m.quotaLifecycleMu.Unlock() + if !inQuotaLeases { + t.Errorf("reservation was not kept in m.quotaLeases on commit failure") + } + + store.mu.Lock() + store.commitFailCount = 0 + store.mu.Unlock() + + err = m.cleanupLease(lease) + if err != nil { + t.Errorf("cleanupLease failed on retry: %v", err) + } + + lease.mu.Lock() + _, retainedAfter := lease.remoteQuotas[reservationID] + lease.mu.Unlock() + + if retainedAfter { + t.Errorf("res-cache-1 was not removed from lease.remoteQuotas after successful retry") + } + + m.quotaLifecycleMu.Lock() + _, inQuotaLeasesAfter := m.quotaLeases[reservationID] + m.quotaLifecycleMu.Unlock() + if inQuotaLeasesAfter { + t.Errorf("reservation was not removed from m.quotaLeases after successful retry") + } +} + +func TestBlockedSendDoesNotFreezeMill(t *testing.T) { + store := &scriptedStore{ + verdicts: []quota.Reservation{ + {ID: "res-workflow-A", Allowed: true, Reason: quota.ReasonWithinLimit}, + {ID: "res-cache-A", Allowed: true, Reason: quota.ReasonWithinLimit}, + {ID: "res-workflow-B", Allowed: true, Reason: quota.ReasonWithinLimit}, + {ID: "res-cache-B", Allowed: true, Reason: quota.ReasonWithinLimit}, + }, + } + m, _ := quotaMill(t, store) + + candA := addQuotaCandidate(t, m, "node-a", fixedReport(2, 2048)) + widA := models.WorkflowId{PipelineId: models.PipelineId{Knot: "knot.test", Rkey: "rkeyA"}, Name: "buildA"} + slotA, err := m.place(context.Background(), "dummy", widA, testWorkflow("buildA")) + if err != nil { + t.Fatalf("place A: %v", err) + } + defer slotA.Release() + + m.mu.Lock() + var leaseIDA string + for id, lease := range m.leases { + if lease.nodeID == "node-a" { + leaseIDA = id + break + } + } + m.mu.Unlock() + + candB := addQuotaCandidate(t, m, "node-b", fixedReport(2, 2048)) + widB := models.WorkflowId{PipelineId: models.PipelineId{Knot: "knot.test", Rkey: "rkeyB"}, Name: "buildB"} + slotB, err := m.place(context.Background(), "dummy", widB, testWorkflow("buildB")) + if err != nil { + t.Fatalf("place B: %v", err) + } + defer slotB.Release() + + m.mu.Lock() + var leaseIDB string + for id, lease := range m.leases { + if lease.nodeID == "node-b" { + leaseIDB = id + break + } + } + m.mu.Unlock() + + // hold one response write open while an unrelated session completes + blockedCh := make(chan struct{}) + sendStarted := make(chan struct{}) + var blockOnce sync.Once + candA.sess.enc = scriptedEncoder(func(msg *millproto.Message) error { + blockOnce.Do(func() { + close(sendStarted) + <-blockedCh + }) + return nil + }) + + aDone := make(chan struct{}) + go func() { + m.handleQuotaRequest(candA.sess, &millv1.QuotaRequest{ + RequestId: "req-A", + Operation: millv1.QuotaOperation_QUOTA_OPERATION_RESERVE, + Kind: string(quota.KindNixCache), + Key: "narhash-A", + LeaseId: leaseIDA, + Resources: map[string]int64{"cache_storage_bytes": 4096}, + }) + close(aDone) + }() + <-sendStarted + + candB.sess.enc = scriptedEncoder(func(msg *millproto.Message) error { + return nil + }) + + doneCh := make(chan struct{}) + go func() { + m.handleQuotaRequest(candB.sess, &millv1.QuotaRequest{ + RequestId: "req-B", + Operation: millv1.QuotaOperation_QUOTA_OPERATION_RESERVE, + Kind: string(quota.KindNixCache), + Key: "narhash-B", + LeaseId: leaseIDB, + Resources: map[string]int64{"cache_storage_bytes": 4096}, + }) + close(doneCh) + }() + + select { + case <-doneCh: + case <-time.After(2 * time.Second): + t.Fatal("unrelated cache reserve blocked on another session's response") + } + + close(blockedCh) + select { + case <-aDone: + case <-time.After(2 * time.Second): + t.Fatal("blocked response did not finish after release") + } +} diff --git a/spindle/mill/quota_test_helper_test.go b/spindle/mill/quota_test_helper_test.go new file mode 100644 index 00000000..8bd76c1f --- /dev/null +++ b/spindle/mill/quota_test_helper_test.go @@ -0,0 +1,46 @@ +package mill + +import ( + "testing" + "time" + + "tangled.org/core/api/tangled" + "tangled.org/core/spindle/db" + "tangled.org/core/spindle/quota" +) + +// fair manager over the real store with no limits configured, so tests keep +// the quota path live without any subject hitting a ceiling +func testQuotaManager(t *testing.T, bdb *db.DB) *quota.Manager { + t.Helper() + qm := quota.NewManager(db.NewQuotaStore(bdb, quota.Defaults{}), 20*time.Millisecond, nil) + t.Cleanup(qm.Close) + return qm +} + +// manager whose repo limits are exactly the given resources, used to force a +// denial without touching the store's internals +func testLimitedQuotaManager(t *testing.T, bdb *db.DB, defaults quota.Resources) *quota.Manager { + t.Helper() + qm := quota.NewManager(db.NewQuotaStore(bdb, quota.Defaults{ + quota.ScopeRepo: defaults, + quota.ScopeUser: defaults, + }), 20*time.Millisecond, nil) + t.Cleanup(qm.Close) + return qm +} + +// pipeline metadata carrying the identity every engine now demands, and the +// only place the mill is allowed to learn who to charge +func testPipeline() tangled.Pipeline { + repoDID := "did:plc:testrepo" + return tangled.Pipeline{ + TriggerMetadata: &tangled.Pipeline_TriggerMetadata{ + Repo: &tangled.Pipeline_TriggerRepo{ + Knot: "knot.test", + Did: "did:plc:testowner", + RepoDid: &repoDID, + }, + }, + } +} diff --git a/spindle/mill/restore.go b/spindle/mill/restore.go index de77c943..de752a3a 100644 --- a/spindle/mill/restore.go +++ b/spindle/mill/restore.go @@ -2,11 +2,15 @@ package mill import ( "fmt" + "maps" + "slices" + "time" + "tangled.org/core/spindle/db" + "tangled.org/core/spindle/models" + millproto "tangled.org/core/spindle/mill/proto" millv1 "tangled.org/core/spindle/mill/proto/gen" - "tangled.org/core/spindle/models" - "time" ) const ( @@ -18,18 +22,43 @@ func (m *Mill) persistLease(lease *RemoteLease, state string) error { if m.db == nil { return nil } + lease.mu.Lock() + quotaID := lease.quotaID + lease.mu.Unlock() return m.db.SaveMillLease(db.MillLease{ - LeaseID: lease.id, - NodeID: lease.nodeID, - Epoch: lease.epoch, - Engine: lease.engine, - Knot: lease.wid.Knot, - Rkey: lease.wid.Rkey, - Workflow: lease.wid.Name, - State: state, + LeaseID: lease.id, + NodeID: lease.nodeID, + Epoch: lease.epoch, + Engine: lease.engine, + Knot: lease.wid.Knot, + Rkey: lease.wid.Rkey, + Workflow: lease.wid.Name, + State: state, + QuotaReservationID: quotaID, + OwnerDID: lease.ownerDID, + RepoDID: lease.repoDID, }) } +// reservation ids the fleet still has live leases for. startup recovery +// keeps these and reclaims every other row +func (m *Mill) LiveQuotaReservationIDs() []string { + m.mu.Lock() + leases := slices.Collect(maps.Values(m.leases)) + m.mu.Unlock() + + var ids []string + for _, lease := range leases { + lease.mu.Lock() + if lease.quotaID != "" { + ids = append(ids, lease.quotaID) + } + ids = append(ids, slices.Collect(maps.Keys(lease.remoteQuotas))...) + lease.mu.Unlock() + } + return ids +} + func (m *Mill) RestoreState() error { if m.db == nil { return nil @@ -53,6 +82,11 @@ func (m *Mill) RestoreState() error { PipelineId: models.PipelineId{Knot: r.Knot, Rkey: r.Rkey}, Name: r.Workflow, } + // the charge outlives the mill process. only the id survives, so + // release goes through the manager by id rather than a held lease + lease.quotaID = r.QuotaReservationID + lease.ownerDID = r.OwnerDID + lease.repoDID = r.RepoDID // restored leases start as orphans, an executor must reclaim it via // its first snapshot, or the sweep will fail it lease.orphaned = true @@ -208,13 +242,13 @@ func (m *Mill) finishOrphan(lease *RemoteLease, status string, errMsg *string, e func mapTerminalStatusString(s string) millv1.TerminalStatus { switch s { case "success": - return millv1.TerminalStatus_SUCCESS + return millv1.TerminalStatus_TERMINAL_STATUS_SUCCESS case "failed": - return millv1.TerminalStatus_FAILED + return millv1.TerminalStatus_TERMINAL_STATUS_FAILED case "timeout": - return millv1.TerminalStatus_TIMEOUT + return millv1.TerminalStatus_TERMINAL_STATUS_TIMEOUT case "cancelled": - return millv1.TerminalStatus_CANCELLED + return millv1.TerminalStatus_TERMINAL_STATUS_CANCELLED default: return millv1.TerminalStatus_TERMINAL_STATUS_UNSPECIFIED } diff --git a/spindle/mill/restore_test.go b/spindle/mill/restore_test.go index 4694173c..62e24ae2 100644 --- a/spindle/mill/restore_test.go +++ b/spindle/mill/restore_test.go @@ -22,7 +22,7 @@ func restoreTestMill(t *testing.T, cfg Config) (*Mill, *db.DB) { t.Cleanup(func() { bdb.Close() }) n := notifier.New() m := New(discardLogger(), cfg) - m.Attach(bdb, &n) + m.Attach(bdb, &n, testQuotaManager(t, bdb)) return m, bdb } @@ -30,7 +30,7 @@ func restoredMill(t *testing.T, bdb *db.DB, cfg Config) *Mill { t.Helper() n := notifier.New() m := New(discardLogger(), cfg) - m.Attach(bdb, &n) + m.Attach(bdb, &n, testQuotaManager(t, bdb)) if err := m.RestoreState(); err != nil { t.Fatalf("RestoreState: %v", err) } @@ -106,7 +106,7 @@ func TestOrphanTerminalAuthorsStatusRow(t *testing.T) { LeaseId: "lease-1", Payload: &millv1.Event_AttemptResult{ AttemptResult: &millv1.AttemptResult{ - Status: millv1.TerminalStatus_SUCCESS, + Status: millv1.TerminalStatus_TERMINAL_STATUS_SUCCESS, }, }, }, @@ -280,7 +280,7 @@ func TestAckSeqnoPersistsCursor(t *testing.T) { LeaseId: owned.id, Payload: &millv1.Event_StatusEvent{ StatusEvent: &millv1.StatusEvent{ - Status: millv1.NonterminalStatus_RUNNING, + Status: millv1.NonterminalStatus_NONTERMINAL_STATUS_RUNNING, }, }, }, @@ -333,7 +333,7 @@ func TestOrphanTerminalFailureKeepsLeaseAndSeqnoRetryable(t *testing.T) { LeaseId: "lease-1", Payload: &millv1.Event_AttemptResult{ AttemptResult: &millv1.AttemptResult{ - Status: millv1.TerminalStatus_SUCCESS, + Status: millv1.TerminalStatus_TERMINAL_STATUS_SUCCESS, }, }, }, diff --git a/spindle/mill/session.go b/spindle/mill/session.go index 4a942ea1..fc99752b 100644 --- a/spindle/mill/session.go +++ b/spindle/mill/session.go @@ -153,6 +153,8 @@ func (s *millSession) dispatch(m *Mill, msg *millproto.Message) error { m.onCancelAck(s, msg.GetCancelAck()) case msg.GetLiveLog() != nil: return m.onLiveLog(s, msg.GetLiveLog()) + case msg.GetQuotaReq() != nil: + go m.handleQuotaRequest(s, msg.GetQuotaReq()) default: s.l.Warn("session received unexpected message", "node", s.nodeID) } diff --git a/spindle/models/pipeline.go b/spindle/models/pipeline.go index 49657650..e7a49fe5 100644 --- a/spindle/models/pipeline.go +++ b/spindle/models/pipeline.go @@ -31,4 +31,5 @@ type Workflow struct { Environment map[string]string OwnerDID string RepoDID string + RunID string } diff --git a/spindle/quota/manager.go b/spindle/quota/manager.go index a6c151f1..8affafec 100644 --- a/spindle/quota/manager.go +++ b/spindle/quota/manager.go @@ -823,7 +823,7 @@ func (m *Manager) Commit(ctx context.Context, reservationID string) error { return err } -func (m *Manager) BeginPublish(ctx context.Context, reservationID string) error { +func (m *Manager) BeginCommit(ctx context.Context, reservationID string) error { m.lifecycleMu.Lock() defer m.lifecycleMu.Unlock() origID, _ := parseID(reservationID) @@ -837,5 +837,5 @@ func (m *Manager) BeginPublish(ctx context.Context, reservationID string) error return errors.New("quota reservation ownership changed") } m.mu.Unlock() - return m.store.BeginPublish(ctx, origID) + return m.store.BeginCommit(ctx, origID) } diff --git a/spindle/quota/manager_test.go b/spindle/quota/manager_test.go index b3e1bc74..61714c9a 100644 --- a/spindle/quota/manager_test.go +++ b/spindle/quota/manager_test.go @@ -294,7 +294,7 @@ func (m *mockStore) splitRest(r string) []string { return []string{res, kind, key} } -func (m *mockStore) BeginPublish(ctx context.Context, id string) error { +func (m *mockStore) BeginCommit(ctx context.Context, id string) error { m.mu.Lock() defer m.mu.Unlock() if _, ok := m.reservations[id]; ok { @@ -1139,10 +1139,10 @@ func TestManagerFencesDuplicateReservationOwners(t *testing.T) { if quota.StorageReservationID(res1.ID) != quota.StorageReservationID(res3.ID) { t.Fatalf("fenced IDs refer to different storage reservations: %q and %q", res1.ID, res3.ID) } - if err := mgr.BeginPublish(context.Background(), res1.ID); err == nil { + if err := mgr.BeginCommit(context.Background(), res1.ID); err == nil { t.Fatal("superseded owner began publication") } - if err := mgr.BeginPublish(context.Background(), res3.ID); err != nil { + if err := mgr.BeginCommit(context.Background(), res3.ID); err != nil { t.Fatalf("current owner could not begin publication: %v", err) } diff --git a/spindle/quota/quota.go b/spindle/quota/quota.go index c44b0003..c76ca702 100644 --- a/spindle/quota/quota.go +++ b/spindle/quota/quota.go @@ -82,7 +82,7 @@ type Lease interface { type ReservationStore interface { Reserve(ctx context.Context, req ReserveRequest) (Reservation, error) - BeginPublish(ctx context.Context, reservationID string) error + BeginCommit(ctx context.Context, reservationID string) error Commit(ctx context.Context, reservationID string) error Release(ctx context.Context, reservationID string) error } diff --git a/spindle/server.go b/spindle/server.go index 09a30827..86308735 100644 --- a/spindle/server.go +++ b/spindle/server.go @@ -49,6 +49,7 @@ import ( "tangled.org/core/spindle/mill/executor" "tangled.org/core/spindle/models" "tangled.org/core/spindle/observability" + "tangled.org/core/spindle/quota" "tangled.org/core/spindle/secrets" "tangled.org/core/spindle/xrpc" "tangled.org/core/tid" @@ -72,6 +73,7 @@ type executorClient interface { Connect(context.Context) Drain(context.Context) error RegisterMetrics(*observability.Metrics) + SetQuotaClient(*executor.QuotaClient) } type Spindle struct { @@ -98,13 +100,14 @@ type Spindle struct { store artifactstore.Store stores *artifactstore.Stores reader artifactstore.Reader + qm *quota.Manager // set only when this spindle hosts the mill or joins one as an executor mill *mill.Mill exec executorClient } // New creates a new Spindle server with the provided configuration and engines. -func New(ctx context.Context, cfg *config.Config, d *db.DB, engines map[string]models.Engine) (*Spindle, error) { +func New(ctx context.Context, cfg *config.Config, d *db.DB, engines map[string]models.Engine, qm *quota.Manager) (*Spindle, error) { logger := log.FromContext(ctx) metrics := observability.GetMetrics(ctx) if metrics == nil { @@ -139,6 +142,7 @@ func New(ctx context.Context, cfg *config.Config, d *db.DB, engines map[string]m rootCtx: lifecycleCtx, rootCancel: cancelLifecycle, jobWake: make(chan struct{}, 1), + qm: qm, } metrics.AttachDB(ctx, d) diskFallback := "" @@ -556,10 +560,34 @@ func Run(ctx context.Context) error { return fmt.Errorf("failed to setup db: %w", err) } + var qs *db.QuotaStore + var qm *quota.Manager + if cfg.Role != config.RoleExecutor { + defaults, err := cfg.ToDefaults() + if err != nil { + return fmt.Errorf("invalid quota configuration: %w", err) + } + for scope, resources := range defaults { + for resource, limit := range resources { + metrics.SetQuotaDefaultLimit(string(scope), string(resource), limit) + } + } + qs = db.NewQuotaStore(d, defaults) + qm = quota.NewManager(qs, 50*time.Millisecond, metrics.QuotaObserver()) + defer qm.Close() + + if cfg.Role != config.RoleMill { + if err := qs.Recover(ctx, nil); err != nil { + return fmt.Errorf("failed to recover quota reservations: %w", err) + } + } + } + logger := log.FromContext(ctx) var engines map[string]models.Engine var m *mill.Mill + var quotaClient *executor.QuotaClient if cfg.Role == config.RoleMill { // on a mill host, engines place jobs on executors instead of running @@ -580,7 +608,14 @@ func Run(ctx context.Context) error { if err != nil { return err } - microvmEng, err := newMicrovmEngine(ctx, cfg, d) + var engineQStore quota.ReservationStore + if cfg.Role == config.RoleExecutor { + quotaClient = executor.NewQuotaClient() + engineQStore = nil + } else { + engineQStore = qs + } + microvmEng, err := newMicrovmEngine(ctx, cfg, d, engineQStore) if err != nil { return err } @@ -591,26 +626,67 @@ func Run(ctx context.Context) error { } } - s, err := New(ctx, cfg, d, engines) + s, err := New(ctx, cfg, d, engines, qm) if err != nil { return err } + if cfg.Role == config.RoleMill || cfg.Role == config.RoleStandalone { + // flattens the store snapshot into bounded label tuples, the collector + // drops anything outside the bounded scope, resource and status sets + metrics.SetQuotaLoader(func() (observability.QuotaSnapshot, error) { + snap, err := qs.MetricsSnapshot(context.Background()) + if err != nil { + return observability.QuotaSnapshot{}, err + } + var out observability.QuotaSnapshot + for scope, resources := range snap.Usage { + for resource, used := range resources { + out.Usage = append(out.Usage, observability.QuotaUsage{ + Scope: string(scope), + Resource: resource, + Used: used, + }) + } + } + for scope, resources := range snap.Subjects { + for resource, statuses := range resources { + for status, count := range statuses { + out.Subjects = append(out.Subjects, observability.QuotaSubjectCount{ + Scope: string(scope), + Resource: string(resource), + Status: status, + Count: count, + }) + } + } + } + return out, nil + }) + } if m != nil { // the engines built above hold the mill, but the mill's db and // notifier only exist after New, so attach them here - m.Attach(s.DB(), s.Notifier()) + m.Attach(s.DB(), s.Notifier(), qm) s.mill = m s.mill.RegisterMetrics(s.metrics) if err := m.RestoreState(); err != nil { return fmt.Errorf("restoring mill state: %w", err) } + liveIDs := m.LiveQuotaReservationIDs() + for i := range liveIDs { + liveIDs[i] = quota.StorageReservationID(liveIDs[i]) + } + if err := qs.Recover(ctx, liveIDs); err != nil { + return fmt.Errorf("failed to recover mill quota reservations: %w", err) + } } if cfg.Role == config.RoleExecutor { s.exec, err = executor.New(cfg, engines, s.DB(), s.Notifier(), log.SubLogger(logger, "executor"), s.store) if err != nil { return err } + s.exec.SetQuotaClient(quotaClient) s.exec.RegisterMetrics(s.metrics) } @@ -1266,6 +1342,7 @@ func (s *Spindle) runJob(ctx context.Context, job *db.JobRow) { }, fmt.Sprintf("init workflow: %s", err), -1, s.n) continue } + ewf.RunID = fmt.Sprintf("%d", job.Id) if ewf.Environment == nil { ewf.Environment = make(map[string]string) @@ -1274,7 +1351,7 @@ func (s *Spindle) runJob(ctx context.Context, job *db.JobRow) { workflows[eng] = append(workflows[eng], *ewf) } - engine.StartWorkflows(l, s.vault, s.cfg, s.stores, s.db, s.n, jobCtx, &models.Pipeline{ + engine.StartWorkflows(l, s.vault, s.cfg, s.qm, s.stores, s.db, s.n, jobCtx, &models.Pipeline{ RepoDid: syntax.DID(job.RepoDid), Workflows: workflows, TrustedSource: trustedSource, diff --git a/spindle/server_test.go b/spindle/server_test.go index 0bb41bac..07c62542 100644 --- a/spindle/server_test.go +++ b/spindle/server_test.go @@ -13,6 +13,7 @@ import ( kgit "tangled.org/core/knotserver/git" "tangled.org/core/spindle/config" "tangled.org/core/spindle/db" + "tangled.org/core/spindle/mill/executor" "tangled.org/core/spindle/models" "tangled.org/core/spindle/observability" ) @@ -80,7 +81,7 @@ func TestExecutorRoleBuildsMinimalSpindle(t *testing.T) { cfg.ArtifactStores.Disk.Dir = t.TempDir() cfg.Mill.ArtifactStore = "disk" - s, err := New(ctx, cfg, d, map[string]models.Engine{}) + s, err := New(ctx, cfg, d, map[string]models.Engine{}, nil) if err != nil { t.Fatalf("New() error = %v", err) } @@ -131,6 +132,8 @@ func (e *drainTestExecutor) Drain(ctx context.Context) error { func (e *drainTestExecutor) RegisterMetrics(*observability.Metrics) {} +func (e *drainTestExecutor) SetQuotaClient(*executor.QuotaClient) {} + func TestExecutorShutdownDrainsBeforeDisconnecting(t *testing.T) { exec := &drainTestExecutor{ connected: make(chan struct{}), @@ -265,7 +268,7 @@ func TestNewOwnsLifecycleContextBeforeStart(t *testing.T) { cfg.ArtifactStores.Disk.Dir = t.TempDir() cfg.Mill.ArtifactStore = "disk" - s, err := New(ctx, cfg, d, map[string]models.Engine{}) + s, err := New(ctx, cfg, d, map[string]models.Engine{}, nil) if err != nil { t.Fatal(err) }