Skip to content

Commit 4b03d06

Browse files
ammarioammario
committed
provisionerd tests pass!
1 parent b7cd66c commit 4b03d06

File tree

4 files changed

+72
-87
lines changed

4 files changed

+72
-87
lines changed

‎provisioner/echo/serve.go

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -118,11 +118,12 @@ func (e *echo) Provision(stream proto.DRPCProvisioner_ProvisionStream) error {
118118
}
119119

120120
var config *proto.Provision_Config
121-
if msg.GetPlan() != nil {
121+
switch {
122+
case msg.GetPlan() != nil:
122123
config = msg.GetPlan().GetConfig()
123-
} else if msg.GetApply() != nil {
124+
case msg.GetApply() != nil:
124125
config = msg.GetApply().GetConfig()
125-
} else {
126+
default:
126127
// Probably a cancel
127128
return nil
128129
}

‎provisionerd/provisionerd.go

Lines changed: 12 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -322,16 +322,18 @@ func (p *Server) acquireJob(ctx context.Context) {
322322
p.activeJob = runner.New(
323323
ctx,
324324
job,
325-
p,
326-
p.opts.Logger,
327-
p.opts.Filesystem,
328-
p.opts.WorkDirectory,
329-
provisioner,
330-
p.opts.UpdateInterval,
331-
p.opts.ForceCancelInterval,
332-
p.opts.LogBufferInterval,
333-
p.tracer,
334-
p.opts.Metrics.Runner,
325+
runner.Options{
326+
Updater: p,
327+
Logger: p.opts.Logger,
328+
Filesystem: p.opts.Filesystem,
329+
WorkDirectory: p.opts.WorkDirectory,
330+
Provisioner: provisioner,
331+
UpdateInterval: p.opts.UpdateInterval,
332+
ForceCancelInterval: p.opts.ForceCancelInterval,
333+
LogDebounceInterval: p.opts.LogBufferInterval,
334+
Tracer: p.tracer,
335+
Metrics: p.opts.Metrics.Runner,
336+
},
335337
)
336338

337339
go p.activeJob.Run()

‎provisionerd/provisionerd_test.go

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -229,7 +229,7 @@ func TestProvisionerd(t *testing.T) {
229229
didComplete atomic.Bool
230230
didLog atomic.Bool
231231
didAcquireJob atomic.Bool
232-
didDryRun atomic.Bool
232+
didDryRun = atomic.NewBool(true)
233233
didReadme atomic.Bool
234234
completeChan = make(chan struct{})
235235
completeOnce sync.Once
@@ -299,8 +299,8 @@ func TestProvisionerd(t *testing.T) {
299299
provision: func(stream sdkproto.DRPCProvisioner_ProvisionStream) error {
300300
request, err := stream.Recv()
301301
require.NoError(t, err)
302-
if request.GetStart().DryRun {
303-
didDryRun.Store(true)
302+
if request.GetApply() != nil {
303+
didDryRun.Store(false)
304304
}
305305
err = stream.Send(&sdkproto.Provision_Response{
306306
Type: &sdkproto.Provision_Response_Log{

‎provisionerd/runner/runner.go

Lines changed: 53 additions & 71 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,6 @@ type Runner struct {
4242
metrics Metrics
4343
job *proto.AcquiredJob
4444
sender JobUpdater
45-
quotaCommitter QuotaCommitter
4645
logger slog.Logger
4746
filesystem afero.Fs
4847
workDirectory string
@@ -87,13 +86,8 @@ type JobUpdater interface {
8786
CompleteJob(ctx context.Context, in *proto.CompletedJob) error
8887
}
8988

90-
type QuotaCommitter interface {
91-
CommitQuota(ctx context.Context, in *proto.CommitQuotaRequest) (*proto.CommitQuotaResponse, error)
92-
}
93-
9489
type Options struct {
9590
Updater JobUpdater
96-
QuotaCommitter QuotaCommitter
9791
Logger slog.Logger
9892
Filesystem afero.Fs
9993
WorkDirectory string
@@ -121,7 +115,6 @@ func New(
121115
metrics: opts.Metrics,
122116
job: job,
123117
sender: opts.Updater,
124-
quotaCommitter: opts.QuotaCommitter,
125118
logger: opts.Logger.With(slog.F("job_id", job.JobId)),
126119
filesystem: opts.Filesystem,
127120
workDirectory: opts.WorkDirectory,
@@ -678,12 +671,13 @@ func (r *Runner) runTemplateImportProvision(ctx context.Context, values []*sdkpr
678671
}
679672
}()
680673
err = stream.Send(&sdkproto.Provision_Request{
681-
Type: &sdkproto.Provision_Request_Start{
682-
Start: &sdkproto.Provision_Start{
683-
Directory: r.workDirectory,
674+
Type: &sdkproto.Provision_Request_Plan{
675+
Plan: &sdkproto.Provision_Plan{
676+
Config: &sdkproto.Provision_Config{
677+
Directory: r.workDirectory,
678+
Metadata: metadata,
679+
},
684680
ParameterValues: values,
685-
DryRun: true,
686-
Metadata: metadata,
687681
},
688682
},
689683
})
@@ -783,8 +777,8 @@ func (r *Runner) runTemplateDryRun(ctx context.Context) (*proto.CompletedJob, *p
783777
}, nil
784778
}
785779

786-
func (r *Runner) buildWorkspace(ctx context.Context, stage string, dry bool) (
787-
*proto.CompletedJob, *proto.FailedJob,
780+
func (r *Runner) buildWorkspace(ctx context.Context, stage string, req *sdkproto.Provision_Request) (
781+
*sdkproto.Provision_Complete, *proto.FailedJob,
788782
) {
789783
// use the notStopped so that if we attempt to gracefully cancel, the stream
790784
// will still be available for us to send the cancel to the provisioner
@@ -805,17 +799,8 @@ func (r *Runner) buildWorkspace(ctx context.Context, stage string, dry bool) (
805799
})
806800
}
807801
}()
808-
err = stream.Send(&sdkproto.Provision_Request{
809-
Type: &sdkproto.Provision_Request_Start{
810-
Start: &sdkproto.Provision_Start{
811-
Directory: r.workDirectory,
812-
ParameterValues: r.job.GetWorkspaceBuild().ParameterValues,
813-
Metadata: r.job.GetWorkspaceBuild().Metadata,
814-
State: r.job.GetWorkspaceBuild().State,
815-
DryRun: dry,
816-
},
817-
},
818-
})
802+
803+
err = stream.Send(req)
819804
if err != nil {
820805
return nil, r.failedJobf("start provision: %s", err)
821806
}
@@ -864,15 +849,7 @@ func (r *Runner) buildWorkspace(ctx context.Context, stage string, dry bool) (
864849
slog.F("state_length", len(msgType.Complete.State)),
865850
)
866851
// Stop looping!
867-
return &proto.CompletedJob{
868-
JobId: r.job.JobId,
869-
Type: &proto.CompletedJob_WorkspaceBuild_{
870-
WorkspaceBuild: &proto.CompletedJob_WorkspaceBuild{
871-
State: msgType.Complete.State,
872-
Resources: msgType.Complete.Resources,
873-
},
874-
},
875-
}, nil
852+
return msgType.Complete, nil
876853
default:
877854
return nil, r.failedJobf("invalid message type %T received from provisioner", msg.Type)
878855
}
@@ -884,57 +861,62 @@ func (r *Runner) runWorkspaceBuild(ctx context.Context) (*proto.CompletedJob, *p
884861
defer span.End()
885862

886863
var (
887-
stage string
888-
commitQuota bool
864+
applyStage string
889865
)
890866
switch r.job.GetWorkspaceBuild().Metadata.WorkspaceTransition {
891867
case sdkproto.WorkspaceTransition_START:
892-
stage = "Starting workspace"
893-
commitQuota = true
868+
applyStage = "Starting workspace"
894869
case sdkproto.WorkspaceTransition_STOP:
895-
stage = "Stopping workspace"
896-
commitQuota = true
870+
applyStage = "Stopping workspace"
897871
case sdkproto.WorkspaceTransition_DESTROY:
898-
stage = "Destroying workspace"
872+
applyStage = "Destroying workspace"
899873
}
900874

901875
r.queueLog(ctx, &proto.Log{
902876
Source: proto.LogSource_PROVISIONER_DAEMON,
903877
Level: sdkproto.LogLevel_INFO,
904-
Stage: stage,
878+
Stage: applyStage,
905879
CreatedAt: time.Now().UnixMilli(),
906880
})
881+
config := &sdkproto.Provision_Config{
882+
Directory: r.workDirectory,
883+
Metadata: r.job.GetWorkspaceBuild().Metadata,
884+
State: r.job.GetWorkspaceBuild().State,
885+
}
907886

908-
if commitQuota {
909-
stage := "Quota Planning"
910-
completed, failed := r.buildWorkspace(ctx, stage, false)
911-
if failed != nil {
912-
return nil, failed
913-
}
914-
cost := countCost(completed.GetWorkspaceBuild().Resources)
915-
if cost > 0 {
916-
resp, err := r.quotaCommitter.CommitQuota(ctx, &proto.CommitQuotaRequest{
917-
Cost: int32(cost),
918-
})
919-
if err != nil {
920-
return nil, r.failedJobf("commit quota: %+v", err)
921-
}
922-
r.queueLog(ctx, &proto.Log{
923-
Source: proto.LogSource_PROVISIONER,
924-
Level: sdkproto.LogLevel_INFO,
925-
CreatedAt: time.Now().UnixMilli(),
926-
Output: fmt.Sprintf(
927-
"cost: %v\ntotal: %v\n available: %v\n",
928-
cost, resp.TotalCredits, resp.CreditsAvailable,
929-
),
930-
Stage: stage,
931-
})
932-
if !resp.Ok {
933-
return nil, r.failedJobf("insufficient quota")
934-
}
935-
}
887+
completed, failed := r.buildWorkspace(ctx, "Planning infrastructure", &sdkproto.Provision_Request{
888+
Type: &sdkproto.Provision_Request_Plan{
889+
Plan: &sdkproto.Provision_Plan{
890+
Config: config,
891+
ParameterValues: r.job.GetWorkspaceBuild().ParameterValues,
892+
},
893+
},
894+
})
895+
if failed != nil {
896+
return nil, failed
936897
}
937-
return r.buildWorkspace(ctx, stage, false)
898+
899+
completedApply, failed := r.buildWorkspace(ctx, applyStage, &sdkproto.Provision_Request{
900+
Type: &sdkproto.Provision_Request_Apply{
901+
Apply: &sdkproto.Provision_Apply{
902+
Config: config,
903+
Plan: completed.GetPlan(),
904+
},
905+
},
906+
})
907+
if failed != nil {
908+
return nil, failed
909+
}
910+
911+
return &proto.CompletedJob{
912+
JobId: r.job.JobId,
913+
Type: &proto.CompletedJob_WorkspaceBuild_{
914+
WorkspaceBuild: &proto.CompletedJob_WorkspaceBuild{
915+
State: completedApply.GetState(),
916+
Resources: completedApply.GetResources(),
917+
},
918+
},
919+
}, nil
938920
}
939921

940922
func (r *Runner) failedJobf(format string, args ...interface{}) *proto.FailedJob {

0 commit comments

Comments
 (0)