Implement Quotas v3 (#5012)

* provisioner/terraform: add cost to resource_metadata * provisionerd/runner: use Options struct * Complete provisionerd implementation * Add quota_allowance to groups * Combine Quota and RBAC licenses * Add Opts to InTx
2026-09-24 15:04:27 +08:00 · 2022-11-14 17:57:33 +00:00
parent 3fb7892c07
commit 97dbd4dc5d
101 changed files with 1575 additions and 845 deletions
@@ -0,0 +1,11 @@
+package runner
+
+import "github.com/coder/coder/provisionersdk/proto"
+
+func sumDailyCost(resources []*proto.Resource) int {
+	var sum int
+	for _, r := range resources {
+		sum += int(r.DailyCost)
+	}
+	return sum
+}
@@ -42,6 +42,7 @@ type Runner struct {
 	metrics             Metrics
 	job                 *proto.AcquiredJob
 	sender              JobUpdater
+	quotaCommitter      QuotaCommitter
 	logger              slog.Logger
 	filesystem          afero.Fs
 	workDirectory       string
@@ -85,9 +86,13 @@ type JobUpdater interface {
 	FailJob(ctx context.Context, in *proto.FailedJob) error
 	CompleteJob(ctx context.Context, in *proto.CompletedJob) error
 }
+type QuotaCommitter interface {
+	CommitQuota(ctx context.Context, in *proto.CommitQuotaRequest) (*proto.CommitQuotaResponse, error)
+}

 type Options struct {
 	Updater             JobUpdater
+	QuotaCommitter      QuotaCommitter
 	Logger              slog.Logger
 	Filesystem          afero.Fs
 	WorkDirectory       string
@@ -115,6 +120,7 @@ func New(
 		metrics:             opts.Metrics,
 		job:                 job,
 		sender:              opts.Updater,
+		quotaCommitter:      opts.QuotaCommitter,
 		logger:              opts.Logger.With(slog.F("job_id", job.JobId)),
 		filesystem:          opts.Filesystem,
 		workDirectory:       opts.WorkDirectory,
@@ -843,7 +849,7 @@ func (r *Runner) buildWorkspace(ctx context.Context, stage string, req *sdkproto
 			}

 			r.logger.Debug(context.Background(), "provision complete no error")
-			r.logger.Info(context.Background(), "provision successful; marked job as complete",
+			r.logger.Info(context.Background(), "provision successful",
 				slog.F("resource_count", len(msgType.Complete.Resources)),
 				slog.F("resources", msgType.Complete.Resources),
 				slog.F("state_length", len(msgType.Complete.State)),
@@ -856,35 +862,81 @@ func (r *Runner) buildWorkspace(ctx context.Context, stage string, req *sdkproto
 	}
 }

+func (r *Runner) commitQuota(ctx context.Context, resources []*sdkproto.Resource) *proto.FailedJob {
+	cost := sumDailyCost(resources)
+	if cost == 0 {
+		return nil
+	}
+
+	const stage = "Commit quota"
+
+	resp, err := r.quotaCommitter.CommitQuota(ctx, &proto.CommitQuotaRequest{
+		JobId:     r.job.JobId,
+		DailyCost: int32(cost),
+	})
+	if err != nil {
+		r.queueLog(ctx, &proto.Log{
+			Source:    proto.LogSource_PROVISIONER,
+			Level:     sdkproto.LogLevel_ERROR,
+			CreatedAt: time.Now().UnixMilli(),
+			Output:    fmt.Sprintf("Failed to commit quota: %+v", err),
+			Stage:     stage,
+		})
+		return r.failedJobf("commit quota: %+v", err)
+	}
+	for _, line := range []string{
+		fmt.Sprintf("Build cost       —   %v", cost),
+		fmt.Sprintf("Budget           —   %v", resp.Budget),
+		fmt.Sprintf("Credits consumed —   %v", resp.CreditsConsumed),
+	} {
+		r.queueLog(ctx, &proto.Log{
+			Source:    proto.LogSource_PROVISIONER,
+			Level:     sdkproto.LogLevel_INFO,
+			CreatedAt: time.Now().UnixMilli(),
+			Output:    line,
+			Stage:     stage,
+		})
+	}
+
+	if !resp.Ok {
+		r.queueLog(ctx, &proto.Log{
+			Source:    proto.LogSource_PROVISIONER,
+			Level:     sdkproto.LogLevel_WARN,
+			CreatedAt: time.Now().UnixMilli(),
+			Output:    "This build would exceed your quota. Failing.",
+			Stage:     stage,
+		})
+		return r.failedJobf("insufficient quota")
+	}
+	return nil
+}
+
 func (r *Runner) runWorkspaceBuild(ctx context.Context) (*proto.CompletedJob, *proto.FailedJob) {
 	ctx, span := r.startTrace(ctx, tracing.FuncName())
 	defer span.End()

 	var (
-		applyStage string
+		applyStage  string
+		commitQuota bool
 	)
 	switch r.job.GetWorkspaceBuild().Metadata.WorkspaceTransition {
 	case sdkproto.WorkspaceTransition_START:
 		applyStage = "Starting workspace"
+		commitQuota = true
 	case sdkproto.WorkspaceTransition_STOP:
 		applyStage = "Stopping workspace"
+		commitQuota = true
 	case sdkproto.WorkspaceTransition_DESTROY:
 		applyStage = "Destroying workspace"
 	}

-	r.queueLog(ctx, &proto.Log{
-		Source:    proto.LogSource_PROVISIONER_DAEMON,
-		Level:     sdkproto.LogLevel_INFO,
-		Stage:     applyStage,
-		CreatedAt: time.Now().UnixMilli(),
-	})
 	config := &sdkproto.Provision_Config{
 		Directory: r.workDirectory,
 		Metadata:  r.job.GetWorkspaceBuild().Metadata,
 		State:     r.job.GetWorkspaceBuild().State,
 	}

-	completed, failed := r.buildWorkspace(ctx, "Planning infrastructure", &sdkproto.Provision_Request{
+	completedPlan, failed := r.buildWorkspace(ctx, "Planning infrastructure", &sdkproto.Provision_Request{
 		Type: &sdkproto.Provision_Request_Plan{
 			Plan: &sdkproto.Provision_Plan{
 				Config:          config,
@@ -895,18 +947,34 @@ func (r *Runner) runWorkspaceBuild(ctx context.Context) (*proto.CompletedJob, *p
 	if failed != nil {
 		return nil, failed
 	}
+	r.flushQueuedLogs(ctx)
+	if commitQuota {
+		failed = r.commitQuota(ctx, completedPlan.GetResources())
+		r.flushQueuedLogs(ctx)
+		if failed != nil {
+			return nil, failed
+		}
+	}
+
+	r.queueLog(ctx, &proto.Log{
+		Source:    proto.LogSource_PROVISIONER_DAEMON,
+		Level:     sdkproto.LogLevel_INFO,
+		Stage:     applyStage,
+		CreatedAt: time.Now().UnixMilli(),
+	})

 	completedApply, failed := r.buildWorkspace(ctx, applyStage, &sdkproto.Provision_Request{
 		Type: &sdkproto.Provision_Request_Apply{
 			Apply: &sdkproto.Provision_Apply{
 				Config: config,
-				Plan:   completed.GetPlan(),
+				Plan:   completedPlan.GetPlan(),
 			},
 		},
 	})
 	if failed != nil {
 		return nil, failed
 	}
+	r.flushQueuedLogs(ctx)

 	return &proto.CompletedJob{
 		JobId: r.job.JobId,