Skip to content

Commit fdac30b

Browse files
ammarioammario
committed
Complete provisionerd implementation
1 parent 208b9eb commit fdac30b

File tree

7 files changed

+499
-123
lines changed

7 files changed

+499
-123
lines changed

‎provisionerd/proto/provisionerd.pb.go

Lines changed: 251 additions & 97 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎provisionerd/proto/provisionerd.proto

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,12 +106,24 @@ message UpdateJobResponse {
106106
repeated provisioner.ParameterValue parameter_values = 2;
107107
}
108108

109+
message CommitQuotaRequest {
110+
int32 cost = 1;
111+
}
112+
113+
message CommitQuotaResponse {
114+
bool ok = 1;
115+
int32 credits_available = 2;
116+
int32 total_credits = 3;
117+
}
118+
109119
service ProvisionerDaemon {
110120
// AcquireJob requests a job. Implementations should
111121
// hold a lock on the job until CompleteJob() is
112122
// called with the matching ID.
113123
rpc AcquireJob(Empty) returns (AcquiredJob);
114124

125+
rpc CommitQuota(CommitQuotaRequest) returns (CommitQuotaResponse);
126+
115127
// UpdateJob streams periodic updates for a job.
116128
// Implementations should buffer logs so this stream
117129
// is non-blocking.

‎provisionerd/proto/provisionerd_drpc.pb.go

Lines changed: 43 additions & 3 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎provisionerd/provisionerd.go

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -323,6 +323,8 @@ func (p *Server) acquireJob(ctx context.Context) {
323323
ctx,
324324
job,
325325
runner.Options{
326+
Updater: p,
327+
QuotaCommitter: p,
326328
Logger: p.opts.Logger,
327329
Filesystem: p.opts.Filesystem,
328330
WorkDirectory: p.opts.WorkDirectory,
@@ -364,6 +366,17 @@ func (p *Server) clientDoWithRetries(
364366
return nil, ctx.Err()
365367
}
366368

369+
func (p *Server) CommitQuota(ctx context.Context, in *proto.CommitQuotaRequest) (*proto.CommitQuotaResponse, error) {
370+
out, err := p.clientDoWithRetries(ctx, func(ctx context.Context, client proto.DRPCProvisionerDaemonClient) (any, error) {
371+
return client.CommitQuota(ctx, in)
372+
})
373+
if err != nil {
374+
return nil, err
375+
}
376+
// nolint: forcetypeassert
377+
return out.(*proto.CommitQuotaResponse), nil
378+
}
379+
367380
func (p *Server) UpdateJob(ctx context.Context, in *proto.UpdateJobRequest) (*proto.UpdateJobResponse, error) {
368381
out, err := p.clientDoWithRetries(ctx, func(ctx context.Context, client proto.DRPCProvisionerDaemonClient) (any, error) {
369382
return client.UpdateJob(ctx, in)

‎provisionerd/provisionerd_test.go

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -481,6 +481,97 @@ func TestProvisionerd(t *testing.T) {
481481
require.NoError(t, closer.Close())
482482
})
483483

484+
t.Run("WorkspaceBuildQuotaExceeded", func(t *testing.T) {
485+
t.Parallel()
486+
var (
487+
didComplete atomic.Bool
488+
didLog atomic.Bool
489+
didAcquireJob atomic.Bool
490+
didFail atomic.Bool
491+
completeChan = make(chan struct{})
492+
completeOnce sync.Once
493+
)
494+
495+
closer := createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
496+
return createProvisionerDaemonClient(t, provisionerDaemonTestServer{
497+
acquireJob: func(ctx context.Context, _ *proto.Empty) (*proto.AcquiredJob, error) {
498+
if !didAcquireJob.CAS(false, true) {
499+
completeOnce.Do(func() { close(completeChan) })
500+
return &proto.AcquiredJob{}, nil
501+
}
502+
503+
return &proto.AcquiredJob{
504+
JobId: "test",
505+
Provisioner: "someprovisioner",
506+
TemplateSourceArchive: createTar(t, map[string]string{
507+
"test.txt": "content",
508+
}),
509+
Type: &proto.AcquiredJob_WorkspaceBuild_{
510+
WorkspaceBuild: &proto.AcquiredJob_WorkspaceBuild{
511+
Metadata: &sdkproto.Provision_Metadata{},
512+
},
513+
},
514+
}, nil
515+
},
516+
updateJob: func(ctx context.Context, update *proto.UpdateJobRequest) (*proto.UpdateJobResponse, error) {
517+
if len(update.Logs) != 0 {
518+
didLog.Store(true)
519+
}
520+
return &proto.UpdateJobResponse{}, nil
521+
},
522+
completeJob: func(ctx context.Context, job *proto.CompletedJob) (*proto.Empty, error) {
523+
didComplete.Store(true)
524+
return &proto.Empty{}, nil
525+
},
526+
commitQuota: func(ctx context.Context, com *proto.CommitQuotaRequest) (*proto.CommitQuotaResponse, error) {
527+
return &proto.CommitQuotaResponse{
528+
Ok: com.Cost < 20,
529+
}, nil
530+
},
531+
failJob: func(ctx context.Context, job *proto.FailedJob) (*proto.Empty, error) {
532+
didFail.Store(true)
533+
return &proto.Empty{}, nil
534+
},
535+
}), nil
536+
}, provisionerd.Provisioners{
537+
"someprovisioner": createProvisionerClient(t, provisionerTestServer{
538+
provision: func(stream sdkproto.DRPCProvisioner_ProvisionStream) error {
539+
err := stream.Send(&sdkproto.Provision_Response{
540+
Type: &sdkproto.Provision_Response_Log{
541+
Log: &sdkproto.Log{
542+
Level: sdkproto.LogLevel_DEBUG,
543+
Output: "wow",
544+
},
545+
},
546+
})
547+
require.NoError(t, err)
548+
549+
err = stream.Send(&sdkproto.Provision_Response{
550+
Type: &sdkproto.Provision_Response_Complete{
551+
Complete: &sdkproto.Provision_Complete{
552+
Resources: []*sdkproto.Resource{
553+
{
554+
Cost: 10,
555+
},
556+
{
557+
Cost: 15,
558+
},
559+
},
560+
},
561+
},
562+
})
563+
require.NoError(t, err)
564+
return nil
565+
},
566+
}),
567+
})
568+
require.Condition(t, closedWithin(completeChan, testutil.WaitShort))
569+
require.True(t, didLog.Load())
570+
require.True(t, didFail.Load())
571+
require.False(t, didComplete.Load())
572+
require.NoError(t, closer.Close())
573+
})
574+
484575
t.Run("WorkspaceBuildFailComplete", func(t *testing.T) {
485576
t.Parallel()
486577
var (
@@ -1039,6 +1130,7 @@ func (p *provisionerTestServer) Provision(stream sdkproto.DRPCProvisioner_Provis
10391130
// passable functions for dynamic functionality.
10401131
type provisionerDaemonTestServer struct {
10411132
acquireJob func(ctx context.Context, _ *proto.Empty) (*proto.AcquiredJob, error)
1133+
commitQuota func(ctx context.Context, com *proto.CommitQuotaRequest) (*proto.CommitQuotaResponse, error)
10421134
updateJob func(ctx context.Context, update *proto.UpdateJobRequest) (*proto.UpdateJobResponse, error)
10431135
failJob func(ctx context.Context, job *proto.FailedJob) (*proto.Empty, error)
10441136
completeJob func(ctx context.Context, job *proto.CompletedJob) (*proto.Empty, error)
@@ -1047,6 +1139,14 @@ type provisionerDaemonTestServer struct {
10471139
func (p *provisionerDaemonTestServer) AcquireJob(ctx context.Context, empty *proto.Empty) (*proto.AcquiredJob, error) {
10481140
return p.acquireJob(ctx, empty)
10491141
}
1142+
func (p *provisionerDaemonTestServer) CommitQuota(ctx context.Context, com *proto.CommitQuotaRequest) (*proto.CommitQuotaResponse, error) {
1143+
if p.commitQuota == nil {
1144+
return &proto.CommitQuotaResponse{
1145+
Ok: true,
1146+
}, nil
1147+
}
1148+
return p.commitQuota(ctx, com)
1149+
}
10501150

10511151
func (p *provisionerDaemonTestServer) UpdateJob(ctx context.Context, update *proto.UpdateJobRequest) (*proto.UpdateJobResponse, error) {
10521152
return p.updateJob(ctx, update)

‎provisionerd/runner/quota.go

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
package runner
2+
3+
import "github.com/coder/coder/provisionersdk/proto"
4+
5+
func countCost(resources []*proto.Resource) int {
6+
var sum int
7+
for _, r := range resources {
8+
sum += int(r.Cost)
9+
}
10+
return sum
11+
}

0 commit comments

Comments
 (0)