mirror of
https://github.com/coder/coder.git
synced 2025-07-13 21:36:50 +00:00
feat: integrate Acquirer for provisioner jobs (#9717)
* chore: add Acquirer to provisionerdserver pkg Signed-off-by: Spike Curtis <spike@coder.com> * code review improvements & fixes Signed-off-by: Spike Curtis <spike@coder.com> * feat: integrate Acquirer for provisioner jobs Signed-off-by: Spike Curtis <spike@coder.com> * Fix imports, whitespace Signed-off-by: Spike Curtis <spike@coder.com> * provisionerdserver always closes; remove poll interval from playwright Signed-off-by: Spike Curtis <spike@coder.com> * post jobs outside transactions Signed-off-by: Spike Curtis <spike@coder.com> * graceful shutdown in test Signed-off-by: Spike Curtis <spike@coder.com> * Mark AcquireJob deprecated Signed-off-by: Spike Curtis <spike@coder.com> * Graceful shutdown on all provisionerd tests Signed-off-by: Spike Curtis <spike@coder.com> * Deprecate, not remove CLI flags Signed-off-by: Spike Curtis <spike@coder.com> --------- Signed-off-by: Spike Curtis <spike@coder.com>
This commit is contained in:
Spike Curtis
GitHub
@ -59,7 +59,7 @@ func TestProvisionerd(t *testing.T) {
|
||||
close(done)
|
||||
})
|
||||
closer := createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
|
||||
return createProvisionerDaemonClient(ctx, t, done, provisionerDaemonTestServer{}), nil
|
||||
return createProvisionerDaemonClient(t, done, provisionerDaemonTestServer{}), nil
|
||||
}, provisionerd.LocalProvisioners{})
|
||||
require.NoError(t, closer.Close())
|
||||
})
|
||||
@ -79,33 +79,6 @@ func TestProvisionerd(t *testing.T) {
|
||||
require.NoError(t, closer.Close())
|
||||
})
|
||||
|
||||
t.Run("AcquireEmptyJob", func(t *testing.T) {
|
||||
// The provisioner daemon is supposed to skip the job acquire if
|
||||
// the job provided is empty. This is to show it successfully
|
||||
// tried to get a job, but none were available.
|
||||
t.Parallel()
|
||||
done := make(chan struct{})
|
||||
t.Cleanup(func() {
|
||||
close(done)
|
||||
})
|
||||
completeChan := make(chan struct{})
|
||||
closer := createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
|
||||
acquireJobAttempt := 0
|
||||
return createProvisionerDaemonClient(ctx, t, done, provisionerDaemonTestServer{
|
||||
acquireJob: func(ctx context.Context, _ *proto.Empty) (*proto.AcquiredJob, error) {
|
||||
if acquireJobAttempt == 1 {
|
||||
close(completeChan)
|
||||
}
|
||||
acquireJobAttempt++
|
||||
return &proto.AcquiredJob{}, nil
|
||||
},
|
||||
updateJob: noopUpdateJob,
|
||||
}), nil
|
||||
}, provisionerd.LocalProvisioners{})
|
||||
require.Condition(t, closedWithin(completeChan, testutil.WaitShort))
|
||||
require.NoError(t, closer.Close())
|
||||
})
|
||||
|
||||
t.Run("CloseCancelsJob", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
done := make(chan struct{})
|
||||
@ -118,9 +91,9 @@ func TestProvisionerd(t *testing.T) {
|
||||
var closerMutex sync.Mutex
|
||||
closerMutex.Lock()
|
||||
closer = createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
|
||||
return createProvisionerDaemonClient(ctx, t, done, provisionerDaemonTestServer{
|
||||
acquireJob: func(ctx context.Context, _ *proto.Empty) (*proto.AcquiredJob, error) {
|
||||
return &proto.AcquiredJob{
|
||||
return createProvisionerDaemonClient(t, done, provisionerDaemonTestServer{
|
||||
acquireJobWithCancel: func(stream proto.DRPCProvisionerDaemon_AcquireJobWithCancelStream) error {
|
||||
err := stream.Send(&proto.AcquiredJob{
|
||||
JobId: "test",
|
||||
Provisioner: "someprovisioner",
|
||||
TemplateSourceArchive: createTar(t, map[string]string{
|
||||
@ -131,7 +104,9 @@ func TestProvisionerd(t *testing.T) {
|
||||
Metadata: &sdkproto.Metadata{},
|
||||
},
|
||||
},
|
||||
}, nil
|
||||
})
|
||||
assert.NoError(t, err)
|
||||
return nil
|
||||
},
|
||||
updateJob: noopUpdateJob,
|
||||
failJob: func(ctx context.Context, job *proto.FailedJob) (*proto.Empty, error) {
|
||||
@ -174,9 +149,9 @@ func TestProvisionerd(t *testing.T) {
|
||||
)
|
||||
|
||||
closer := createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
|
||||
return createProvisionerDaemonClient(ctx, t, done, provisionerDaemonTestServer{
|
||||
acquireJob: func(ctx context.Context, _ *proto.Empty) (*proto.AcquiredJob, error) {
|
||||
return &proto.AcquiredJob{
|
||||
return createProvisionerDaemonClient(t, done, provisionerDaemonTestServer{
|
||||
acquireJobWithCancel: func(stream proto.DRPCProvisionerDaemon_AcquireJobWithCancelStream) error {
|
||||
err := stream.Send(&proto.AcquiredJob{
|
||||
JobId: "test",
|
||||
Provisioner: "someprovisioner",
|
||||
TemplateSourceArchive: createTar(t, map[string]string{
|
||||
@ -187,7 +162,9 @@ func TestProvisionerd(t *testing.T) {
|
||||
Metadata: &sdkproto.Metadata{},
|
||||
},
|
||||
},
|
||||
}, nil
|
||||
})
|
||||
assert.NoError(t, err)
|
||||
return nil
|
||||
},
|
||||
updateJob: noopUpdateJob,
|
||||
failJob: func(ctx context.Context, job *proto.FailedJob) (*proto.Empty, error) {
|
||||
@ -214,9 +191,9 @@ func TestProvisionerd(t *testing.T) {
|
||||
)
|
||||
|
||||
closer := createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
|
||||
return createProvisionerDaemonClient(ctx, t, done, provisionerDaemonTestServer{
|
||||
acquireJob: func(ctx context.Context, _ *proto.Empty) (*proto.AcquiredJob, error) {
|
||||
return &proto.AcquiredJob{
|
||||
return createProvisionerDaemonClient(t, done, provisionerDaemonTestServer{
|
||||
acquireJobWithCancel: func(stream proto.DRPCProvisionerDaemon_AcquireJobWithCancelStream) error {
|
||||
err := stream.Send(&proto.AcquiredJob{
|
||||
JobId: "test",
|
||||
Provisioner: "someprovisioner",
|
||||
TemplateSourceArchive: createTar(t, map[string]string{
|
||||
@ -227,7 +204,9 @@ func TestProvisionerd(t *testing.T) {
|
||||
Metadata: &sdkproto.Metadata{},
|
||||
},
|
||||
},
|
||||
}, nil
|
||||
})
|
||||
assert.NoError(t, err)
|
||||
return nil
|
||||
},
|
||||
updateJob: func(ctx context.Context, update *proto.UpdateJobRequest) (*proto.UpdateJobResponse, error) {
|
||||
completeOnce.Do(func() { close(completeChan) })
|
||||
@ -260,36 +239,27 @@ func TestProvisionerd(t *testing.T) {
|
||||
close(done)
|
||||
})
|
||||
var (
|
||||
didComplete atomic.Bool
|
||||
didLog atomic.Bool
|
||||
didAcquireJob atomic.Bool
|
||||
didReadme atomic.Bool
|
||||
completeChan = make(chan struct{})
|
||||
completeOnce sync.Once
|
||||
didComplete atomic.Bool
|
||||
didLog atomic.Bool
|
||||
didReadme atomic.Bool
|
||||
acq = newAcquireOne(t, &proto.AcquiredJob{
|
||||
JobId: "test",
|
||||
Provisioner: "someprovisioner",
|
||||
TemplateSourceArchive: createTar(t, map[string]string{
|
||||
"test.txt": "content",
|
||||
provisionersdk.ReadmeFile: "# A cool template 😎\n",
|
||||
}),
|
||||
Type: &proto.AcquiredJob_TemplateImport_{
|
||||
TemplateImport: &proto.AcquiredJob_TemplateImport{
|
||||
Metadata: &sdkproto.Metadata{},
|
||||
},
|
||||
},
|
||||
})
|
||||
)
|
||||
|
||||
closer := createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
|
||||
return createProvisionerDaemonClient(ctx, t, done, provisionerDaemonTestServer{
|
||||
acquireJob: func(ctx context.Context, _ *proto.Empty) (*proto.AcquiredJob, error) {
|
||||
if !didAcquireJob.CAS(false, true) {
|
||||
completeOnce.Do(func() { close(completeChan) })
|
||||
return &proto.AcquiredJob{}, nil
|
||||
}
|
||||
|
||||
return &proto.AcquiredJob{
|
||||
JobId: "test",
|
||||
Provisioner: "someprovisioner",
|
||||
TemplateSourceArchive: createTar(t, map[string]string{
|
||||
"test.txt": "content",
|
||||
provisionersdk.ReadmeFile: "# A cool template 😎\n",
|
||||
}),
|
||||
Type: &proto.AcquiredJob_TemplateImport_{
|
||||
TemplateImport: &proto.AcquiredJob_TemplateImport{
|
||||
Metadata: &sdkproto.Metadata{},
|
||||
},
|
||||
},
|
||||
}, nil
|
||||
},
|
||||
return createProvisionerDaemonClient(t, done, provisionerDaemonTestServer{
|
||||
acquireJobWithCancel: acq.acquireWithCancel,
|
||||
updateJob: func(ctx context.Context, update *proto.UpdateJobRequest) (*proto.UpdateJobResponse, error) {
|
||||
if len(update.Logs) > 0 {
|
||||
didLog.Store(true)
|
||||
@ -338,7 +308,7 @@ func TestProvisionerd(t *testing.T) {
|
||||
}),
|
||||
})
|
||||
|
||||
require.Condition(t, closedWithin(completeChan, testutil.WaitShort))
|
||||
require.Condition(t, closedWithin(acq.complete, testutil.WaitShort))
|
||||
require.NoError(t, closer.Close())
|
||||
assert.True(t, didLog.Load(), "should log some updates")
|
||||
assert.True(t, didComplete.Load(), "should complete the job")
|
||||
@ -351,36 +321,26 @@ func TestProvisionerd(t *testing.T) {
|
||||
close(done)
|
||||
})
|
||||
var (
|
||||
didComplete atomic.Bool
|
||||
didLog atomic.Bool
|
||||
didAcquireJob atomic.Bool
|
||||
completeChan = make(chan struct{})
|
||||
completeOnce sync.Once
|
||||
|
||||
metadata = &sdkproto.Metadata{}
|
||||
didComplete atomic.Bool
|
||||
didLog atomic.Bool
|
||||
metadata = &sdkproto.Metadata{}
|
||||
acq = newAcquireOne(t, &proto.AcquiredJob{
|
||||
JobId: "test",
|
||||
Provisioner: "someprovisioner",
|
||||
TemplateSourceArchive: createTar(t, map[string]string{
|
||||
"test.txt": "content",
|
||||
}),
|
||||
Type: &proto.AcquiredJob_TemplateDryRun_{
|
||||
TemplateDryRun: &proto.AcquiredJob_TemplateDryRun{
|
||||
Metadata: metadata,
|
||||
},
|
||||
},
|
||||
})
|
||||
)
|
||||
|
||||
closer := createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
|
||||
return createProvisionerDaemonClient(ctx, t, done, provisionerDaemonTestServer{
|
||||
acquireJob: func(ctx context.Context, _ *proto.Empty) (*proto.AcquiredJob, error) {
|
||||
if !didAcquireJob.CAS(false, true) {
|
||||
completeOnce.Do(func() { close(completeChan) })
|
||||
return &proto.AcquiredJob{}, nil
|
||||
}
|
||||
|
||||
return &proto.AcquiredJob{
|
||||
JobId: "test",
|
||||
Provisioner: "someprovisioner",
|
||||
TemplateSourceArchive: createTar(t, map[string]string{
|
||||
"test.txt": "content",
|
||||
}),
|
||||
Type: &proto.AcquiredJob_TemplateDryRun_{
|
||||
TemplateDryRun: &proto.AcquiredJob_TemplateDryRun{
|
||||
Metadata: metadata,
|
||||
},
|
||||
},
|
||||
}, nil
|
||||
},
|
||||
return createProvisionerDaemonClient(t, done, provisionerDaemonTestServer{
|
||||
acquireJobWithCancel: acq.acquireWithCancel,
|
||||
updateJob: func(ctx context.Context, update *proto.UpdateJobRequest) (*proto.UpdateJobResponse, error) {
|
||||
if len(update.Logs) == 0 {
|
||||
t.Log("provisionerDaemonTestServer: no log messages")
|
||||
@ -420,7 +380,7 @@ func TestProvisionerd(t *testing.T) {
|
||||
}),
|
||||
})
|
||||
|
||||
require.Condition(t, closedWithin(completeChan, testutil.WaitShort))
|
||||
require.Condition(t, closedWithin(acq.complete, testutil.WaitShort))
|
||||
require.NoError(t, closer.Close())
|
||||
assert.True(t, didLog.Load(), "should log some updates")
|
||||
assert.True(t, didComplete.Load(), "should complete the job")
|
||||
@ -433,34 +393,25 @@ func TestProvisionerd(t *testing.T) {
|
||||
close(done)
|
||||
})
|
||||
var (
|
||||
didComplete atomic.Bool
|
||||
didLog atomic.Bool
|
||||
didAcquireJob atomic.Bool
|
||||
completeChan = make(chan struct{})
|
||||
completeOnce sync.Once
|
||||
didComplete atomic.Bool
|
||||
didLog atomic.Bool
|
||||
acq = newAcquireOne(t, &proto.AcquiredJob{
|
||||
JobId: "test",
|
||||
Provisioner: "someprovisioner",
|
||||
TemplateSourceArchive: createTar(t, map[string]string{
|
||||
"test.txt": "content",
|
||||
}),
|
||||
Type: &proto.AcquiredJob_WorkspaceBuild_{
|
||||
WorkspaceBuild: &proto.AcquiredJob_WorkspaceBuild{
|
||||
Metadata: &sdkproto.Metadata{},
|
||||
},
|
||||
},
|
||||
})
|
||||
)
|
||||
|
||||
closer := createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
|
||||
return createProvisionerDaemonClient(ctx, t, done, provisionerDaemonTestServer{
|
||||
acquireJob: func(ctx context.Context, _ *proto.Empty) (*proto.AcquiredJob, error) {
|
||||
if !didAcquireJob.CAS(false, true) {
|
||||
completeOnce.Do(func() { close(completeChan) })
|
||||
return &proto.AcquiredJob{}, nil
|
||||
}
|
||||
|
||||
return &proto.AcquiredJob{
|
||||
JobId: "test",
|
||||
Provisioner: "someprovisioner",
|
||||
TemplateSourceArchive: createTar(t, map[string]string{
|
||||
"test.txt": "content",
|
||||
}),
|
||||
Type: &proto.AcquiredJob_WorkspaceBuild_{
|
||||
WorkspaceBuild: &proto.AcquiredJob_WorkspaceBuild{
|
||||
Metadata: &sdkproto.Metadata{},
|
||||
},
|
||||
},
|
||||
}, nil
|
||||
},
|
||||
return createProvisionerDaemonClient(t, done, provisionerDaemonTestServer{
|
||||
acquireJobWithCancel: acq.acquireWithCancel,
|
||||
updateJob: func(ctx context.Context, update *proto.UpdateJobRequest) (*proto.UpdateJobResponse, error) {
|
||||
if len(update.Logs) != 0 {
|
||||
didLog.Store(true)
|
||||
@ -491,7 +442,7 @@ func TestProvisionerd(t *testing.T) {
|
||||
},
|
||||
}),
|
||||
})
|
||||
require.Condition(t, closedWithin(completeChan, testutil.WaitShort))
|
||||
require.Condition(t, closedWithin(acq.complete, testutil.WaitShort))
|
||||
require.NoError(t, closer.Close())
|
||||
assert.True(t, didLog.Load(), "should log some updates")
|
||||
assert.True(t, didComplete.Load(), "should complete the job")
|
||||
@ -504,35 +455,26 @@ func TestProvisionerd(t *testing.T) {
|
||||
close(done)
|
||||
})
|
||||
var (
|
||||
didComplete atomic.Bool
|
||||
didLog atomic.Bool
|
||||
didAcquireJob atomic.Bool
|
||||
didFail atomic.Bool
|
||||
completeChan = make(chan struct{})
|
||||
completeOnce sync.Once
|
||||
didComplete atomic.Bool
|
||||
didLog atomic.Bool
|
||||
didFail atomic.Bool
|
||||
acq = newAcquireOne(t, &proto.AcquiredJob{
|
||||
JobId: "test",
|
||||
Provisioner: "someprovisioner",
|
||||
TemplateSourceArchive: createTar(t, map[string]string{
|
||||
"test.txt": "content",
|
||||
}),
|
||||
Type: &proto.AcquiredJob_WorkspaceBuild_{
|
||||
WorkspaceBuild: &proto.AcquiredJob_WorkspaceBuild{
|
||||
Metadata: &sdkproto.Metadata{},
|
||||
},
|
||||
},
|
||||
})
|
||||
)
|
||||
|
||||
closer := createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
|
||||
return createProvisionerDaemonClient(ctx, t, done, provisionerDaemonTestServer{
|
||||
acquireJob: func(ctx context.Context, _ *proto.Empty) (*proto.AcquiredJob, error) {
|
||||
if !didAcquireJob.CAS(false, true) {
|
||||
completeOnce.Do(func() { close(completeChan) })
|
||||
return &proto.AcquiredJob{}, nil
|
||||
}
|
||||
|
||||
return &proto.AcquiredJob{
|
||||
JobId: "test",
|
||||
Provisioner: "someprovisioner",
|
||||
TemplateSourceArchive: createTar(t, map[string]string{
|
||||
"test.txt": "content",
|
||||
}),
|
||||
Type: &proto.AcquiredJob_WorkspaceBuild_{
|
||||
WorkspaceBuild: &proto.AcquiredJob_WorkspaceBuild{
|
||||
Metadata: &sdkproto.Metadata{},
|
||||
},
|
||||
},
|
||||
}, nil
|
||||
},
|
||||
return createProvisionerDaemonClient(t, done, provisionerDaemonTestServer{
|
||||
acquireJobWithCancel: acq.acquireWithCancel,
|
||||
updateJob: func(ctx context.Context, update *proto.UpdateJobRequest) (*proto.UpdateJobResponse, error) {
|
||||
if len(update.Logs) != 0 {
|
||||
didLog.Store(true)
|
||||
@ -591,7 +533,7 @@ func TestProvisionerd(t *testing.T) {
|
||||
},
|
||||
}),
|
||||
})
|
||||
require.Condition(t, closedWithin(completeChan, testutil.WaitShort))
|
||||
require.Condition(t, closedWithin(acq.complete, testutil.WaitShort))
|
||||
require.NoError(t, closer.Close())
|
||||
assert.True(t, didLog.Load(), "should log some updates")
|
||||
assert.False(t, didComplete.Load(), "should not complete the job")
|
||||
@ -605,34 +547,25 @@ func TestProvisionerd(t *testing.T) {
|
||||
close(done)
|
||||
})
|
||||
var (
|
||||
didFail atomic.Bool
|
||||
didAcquireJob atomic.Bool
|
||||
completeChan = make(chan struct{})
|
||||
completeOnce sync.Once
|
||||
didFail atomic.Bool
|
||||
acq = newAcquireOne(t, &proto.AcquiredJob{
|
||||
JobId: "test",
|
||||
Provisioner: "someprovisioner",
|
||||
TemplateSourceArchive: createTar(t, map[string]string{
|
||||
"test.txt": "content",
|
||||
}),
|
||||
Type: &proto.AcquiredJob_WorkspaceBuild_{
|
||||
WorkspaceBuild: &proto.AcquiredJob_WorkspaceBuild{
|
||||
Metadata: &sdkproto.Metadata{},
|
||||
},
|
||||
},
|
||||
})
|
||||
)
|
||||
|
||||
closer := createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
|
||||
return createProvisionerDaemonClient(ctx, t, done, provisionerDaemonTestServer{
|
||||
acquireJob: func(ctx context.Context, _ *proto.Empty) (*proto.AcquiredJob, error) {
|
||||
if !didAcquireJob.CAS(false, true) {
|
||||
completeOnce.Do(func() { close(completeChan) })
|
||||
return &proto.AcquiredJob{}, nil
|
||||
}
|
||||
|
||||
return &proto.AcquiredJob{
|
||||
JobId: "test",
|
||||
Provisioner: "someprovisioner",
|
||||
TemplateSourceArchive: createTar(t, map[string]string{
|
||||
"test.txt": "content",
|
||||
}),
|
||||
Type: &proto.AcquiredJob_WorkspaceBuild_{
|
||||
WorkspaceBuild: &proto.AcquiredJob_WorkspaceBuild{
|
||||
Metadata: &sdkproto.Metadata{},
|
||||
},
|
||||
},
|
||||
}, nil
|
||||
},
|
||||
updateJob: noopUpdateJob,
|
||||
return createProvisionerDaemonClient(t, done, provisionerDaemonTestServer{
|
||||
acquireJobWithCancel: acq.acquireWithCancel,
|
||||
updateJob: noopUpdateJob,
|
||||
failJob: func(ctx context.Context, job *proto.FailedJob) (*proto.Empty, error) {
|
||||
didFail.Store(true)
|
||||
return &proto.Empty{}, nil
|
||||
@ -661,7 +594,7 @@ func TestProvisionerd(t *testing.T) {
|
||||
},
|
||||
}),
|
||||
})
|
||||
require.Condition(t, closedWithin(completeChan, testutil.WaitShort))
|
||||
require.Condition(t, closedWithin(acq.complete, testutil.WaitShort))
|
||||
require.NoError(t, closer.Close())
|
||||
assert.True(t, didFail.Load(), "should fail the job")
|
||||
})
|
||||
@ -677,9 +610,9 @@ func TestProvisionerd(t *testing.T) {
|
||||
updateChan := make(chan struct{})
|
||||
completeChan := make(chan struct{})
|
||||
server := createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
|
||||
return createProvisionerDaemonClient(ctx, t, done, provisionerDaemonTestServer{
|
||||
acquireJob: func(ctx context.Context, _ *proto.Empty) (*proto.AcquiredJob, error) {
|
||||
return &proto.AcquiredJob{
|
||||
return createProvisionerDaemonClient(t, done, provisionerDaemonTestServer{
|
||||
acquireJobWithCancel: func(stream proto.DRPCProvisionerDaemon_AcquireJobWithCancelStream) error {
|
||||
err := stream.Send(&proto.AcquiredJob{
|
||||
JobId: "test",
|
||||
Provisioner: "someprovisioner",
|
||||
TemplateSourceArchive: createTar(t, map[string]string{
|
||||
@ -690,7 +623,9 @@ func TestProvisionerd(t *testing.T) {
|
||||
Metadata: &sdkproto.Metadata{},
|
||||
},
|
||||
},
|
||||
}, nil
|
||||
})
|
||||
assert.NoError(t, err)
|
||||
return nil
|
||||
},
|
||||
updateJob: func(ctx context.Context, update *proto.UpdateJobRequest) (*proto.UpdateJobResponse, error) {
|
||||
if len(update.Logs) > 0 {
|
||||
@ -755,9 +690,9 @@ func TestProvisionerd(t *testing.T) {
|
||||
updateChan := make(chan struct{})
|
||||
completeChan := make(chan struct{})
|
||||
server := createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
|
||||
return createProvisionerDaemonClient(ctx, t, done, provisionerDaemonTestServer{
|
||||
acquireJob: func(ctx context.Context, _ *proto.Empty) (*proto.AcquiredJob, error) {
|
||||
return &proto.AcquiredJob{
|
||||
return createProvisionerDaemonClient(t, done, provisionerDaemonTestServer{
|
||||
acquireJobWithCancel: func(stream proto.DRPCProvisionerDaemon_AcquireJobWithCancelStream) error {
|
||||
err := stream.Send(&proto.AcquiredJob{
|
||||
JobId: "test",
|
||||
Provisioner: "someprovisioner",
|
||||
TemplateSourceArchive: createTar(t, map[string]string{
|
||||
@ -768,7 +703,9 @@ func TestProvisionerd(t *testing.T) {
|
||||
Metadata: &sdkproto.Metadata{},
|
||||
},
|
||||
},
|
||||
}, nil
|
||||
})
|
||||
assert.NoError(t, err)
|
||||
return nil
|
||||
},
|
||||
updateJob: func(ctx context.Context, update *proto.UpdateJobRequest) (*proto.UpdateJobResponse, error) {
|
||||
resp := &proto.UpdateJobResponse{}
|
||||
@ -825,6 +762,9 @@ func TestProvisionerd(t *testing.T) {
|
||||
})
|
||||
require.Condition(t, closedWithin(updateChan, testutil.WaitShort))
|
||||
require.Condition(t, closedWithin(completeChan, testutil.WaitShort))
|
||||
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
|
||||
defer cancel()
|
||||
require.NoError(t, server.Shutdown(ctx))
|
||||
require.NoError(t, server.Close())
|
||||
})
|
||||
|
||||
@ -844,12 +784,9 @@ func TestProvisionerd(t *testing.T) {
|
||||
completeOnce sync.Once
|
||||
)
|
||||
server := createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
|
||||
client := createProvisionerDaemonClient(ctx, t, done, provisionerDaemonTestServer{
|
||||
acquireJob: func(ctx context.Context, _ *proto.Empty) (*proto.AcquiredJob, error) {
|
||||
if second.Load() {
|
||||
return &proto.AcquiredJob{}, nil
|
||||
}
|
||||
return &proto.AcquiredJob{
|
||||
client := createProvisionerDaemonClient(t, done, provisionerDaemonTestServer{
|
||||
acquireJobWithCancel: func(stream proto.DRPCProvisionerDaemon_AcquireJobWithCancelStream) error {
|
||||
job := &proto.AcquiredJob{
|
||||
JobId: "test",
|
||||
Provisioner: "someprovisioner",
|
||||
TemplateSourceArchive: createTar(t, map[string]string{
|
||||
@ -860,7 +797,15 @@ func TestProvisionerd(t *testing.T) {
|
||||
Metadata: &sdkproto.Metadata{},
|
||||
},
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
if second.Load() {
|
||||
job = &proto.AcquiredJob{}
|
||||
_, err := stream.Recv()
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
err := stream.Send(job)
|
||||
assert.NoError(t, err)
|
||||
return nil
|
||||
},
|
||||
updateJob: func(ctx context.Context, update *proto.UpdateJobRequest) (*proto.UpdateJobResponse, error) {
|
||||
return &proto.UpdateJobResponse{}, nil
|
||||
@ -908,6 +853,9 @@ func TestProvisionerd(t *testing.T) {
|
||||
}),
|
||||
})
|
||||
require.Condition(t, closedWithin(completeChan, testutil.WaitShort))
|
||||
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
|
||||
defer cancel()
|
||||
require.NoError(t, server.Shutdown(ctx))
|
||||
require.NoError(t, server.Close())
|
||||
})
|
||||
|
||||
@ -927,13 +875,15 @@ func TestProvisionerd(t *testing.T) {
|
||||
completeOnce sync.Once
|
||||
)
|
||||
server := createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
|
||||
client := createProvisionerDaemonClient(ctx, t, done, provisionerDaemonTestServer{
|
||||
acquireJob: func(ctx context.Context, _ *proto.Empty) (*proto.AcquiredJob, error) {
|
||||
client := createProvisionerDaemonClient(t, done, provisionerDaemonTestServer{
|
||||
acquireJobWithCancel: func(stream proto.DRPCProvisionerDaemon_AcquireJobWithCancelStream) error {
|
||||
if second.Load() {
|
||||
completeOnce.Do(func() { close(completeChan) })
|
||||
return &proto.AcquiredJob{}, nil
|
||||
_, err := stream.Recv()
|
||||
assert.NoError(t, err)
|
||||
return nil
|
||||
}
|
||||
return &proto.AcquiredJob{
|
||||
job := &proto.AcquiredJob{
|
||||
JobId: "test",
|
||||
Provisioner: "someprovisioner",
|
||||
TemplateSourceArchive: createTar(t, map[string]string{
|
||||
@ -944,7 +894,10 @@ func TestProvisionerd(t *testing.T) {
|
||||
Metadata: &sdkproto.Metadata{},
|
||||
},
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
err := stream.Send(job)
|
||||
assert.NoError(t, err)
|
||||
return nil
|
||||
},
|
||||
failJob: func(ctx context.Context, job *proto.FailedJob) (*proto.Empty, error) {
|
||||
return nil, yamux.ErrSessionShutdown
|
||||
@ -990,6 +943,10 @@ func TestProvisionerd(t *testing.T) {
|
||||
}),
|
||||
})
|
||||
require.Condition(t, closedWithin(completeChan, testutil.WaitShort))
|
||||
t.Log("completeChan closed")
|
||||
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
|
||||
defer cancel()
|
||||
require.NoError(t, server.Shutdown(ctx))
|
||||
require.NoError(t, server.Close())
|
||||
})
|
||||
|
||||
@ -1006,17 +963,21 @@ func TestProvisionerd(t *testing.T) {
|
||||
completeOnce := sync.Once{}
|
||||
|
||||
server := createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
|
||||
return createProvisionerDaemonClient(ctx, t, done, provisionerDaemonTestServer{
|
||||
acquireJob: func(ctx context.Context, _ *proto.Empty) (*proto.AcquiredJob, error) {
|
||||
return createProvisionerDaemonClient(t, done, provisionerDaemonTestServer{
|
||||
acquireJobWithCancel: func(stream proto.DRPCProvisionerDaemon_AcquireJobWithCancelStream) error {
|
||||
m.Lock()
|
||||
defer m.Unlock()
|
||||
logger.Info(ctx, "provisioner stage: AcquiredJob")
|
||||
if len(ops) > 0 {
|
||||
return &proto.AcquiredJob{}, nil
|
||||
_, err := stream.Recv()
|
||||
assert.NoError(t, err)
|
||||
err = stream.Send(&proto.AcquiredJob{})
|
||||
assert.NoError(t, err)
|
||||
return nil
|
||||
}
|
||||
ops = append(ops, "AcquireJob")
|
||||
|
||||
return &proto.AcquiredJob{
|
||||
err := stream.Send(&proto.AcquiredJob{
|
||||
JobId: "test",
|
||||
Provisioner: "someprovisioner",
|
||||
TemplateSourceArchive: createTar(t, map[string]string{
|
||||
@ -1027,7 +988,9 @@ func TestProvisionerd(t *testing.T) {
|
||||
Metadata: &sdkproto.Metadata{},
|
||||
},
|
||||
},
|
||||
}, nil
|
||||
})
|
||||
assert.NoError(t, err)
|
||||
return nil
|
||||
},
|
||||
updateJob: func(ctx context.Context, update *proto.UpdateJobRequest) (*proto.UpdateJobResponse, error) {
|
||||
m.Lock()
|
||||
@ -1076,6 +1039,9 @@ func TestProvisionerd(t *testing.T) {
|
||||
}),
|
||||
})
|
||||
require.Condition(t, closedWithin(completeChan, testutil.WaitShort))
|
||||
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
|
||||
defer cancel()
|
||||
require.NoError(t, server.Shutdown(ctx))
|
||||
require.NoError(t, server.Close())
|
||||
assert.Equal(t, ops[len(ops)-1], "CompleteJob")
|
||||
assert.Contains(t, ops[0:len(ops)-1], "Log: Cleaning Up | ")
|
||||
@ -1105,12 +1071,14 @@ func createTar(t *testing.T, files map[string]string) []byte {
|
||||
// Creates a provisionerd implementation with the provided dialer and provisioners.
|
||||
func createProvisionerd(t *testing.T, dialer provisionerd.Dialer, connector provisionerd.LocalProvisioners) *provisionerd.Server {
|
||||
server := provisionerd.New(dialer, &provisionerd.Options{
|
||||
Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Named("provisionerd").Leveled(slog.LevelDebug),
|
||||
JobPollInterval: 50 * time.Millisecond,
|
||||
UpdateInterval: 50 * time.Millisecond,
|
||||
Connector: connector,
|
||||
Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Named("provisionerd").Leveled(slog.LevelDebug),
|
||||
UpdateInterval: 50 * time.Millisecond,
|
||||
Connector: connector,
|
||||
})
|
||||
t.Cleanup(func() {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
|
||||
defer cancel()
|
||||
_ = server.Shutdown(ctx)
|
||||
_ = server.Close()
|
||||
})
|
||||
return server
|
||||
@ -1118,7 +1086,7 @@ func createProvisionerd(t *testing.T, dialer provisionerd.Dialer, connector prov
|
||||
|
||||
// Creates a provisionerd protobuf client that's connected
|
||||
// to the server implementation provided.
|
||||
func createProvisionerDaemonClient(ctx context.Context, t *testing.T, done <-chan struct{}, server provisionerDaemonTestServer) proto.DRPCProvisionerDaemonClient {
|
||||
func createProvisionerDaemonClient(t *testing.T, done <-chan struct{}, server provisionerDaemonTestServer) proto.DRPCProvisionerDaemonClient {
|
||||
t.Helper()
|
||||
if server.failJob == nil {
|
||||
// Default to asserting the error from the failure, otherwise
|
||||
@ -1137,7 +1105,7 @@ func createProvisionerDaemonClient(ctx context.Context, t *testing.T, done <-cha
|
||||
err := proto.DRPCRegisterProvisionerDaemon(mux, &server)
|
||||
require.NoError(t, err)
|
||||
srv := drpcserver.New(mux)
|
||||
ctx, cancelFunc := context.WithCancel(ctx)
|
||||
ctx, cancelFunc := context.WithCancel(context.Background())
|
||||
closed := make(chan struct{})
|
||||
go func() {
|
||||
defer close(closed)
|
||||
@ -1148,31 +1116,13 @@ func createProvisionerDaemonClient(ctx context.Context, t *testing.T, done <-cha
|
||||
<-closed
|
||||
select {
|
||||
case <-done:
|
||||
// It's possible to get unlucky since the dialer is run in a retry in a goroutine.
|
||||
// The following can occur:
|
||||
// 1. The provisionerd.connect goroutine checks if we're closed prior to attempting to establish a connection
|
||||
// with coderd, sees that we're not.
|
||||
// 2. A test closes the server.
|
||||
// 3. The provisionerd.connect goroutine calls the dialer to establish a connection. This
|
||||
// function detects that someone has tried to create a client after the test finishes.
|
||||
if ctx.Err() == nil {
|
||||
t.Error("createProvisionerDaemonClient cleanup after test was done!")
|
||||
}
|
||||
t.Error("createProvisionerDaemonClient cleanup after test was done!")
|
||||
default:
|
||||
}
|
||||
})
|
||||
select {
|
||||
case <-done:
|
||||
// It's possible to get unlucky since the dialer is run in a retry in a goroutine.
|
||||
// The following can occur:
|
||||
// 1. The provisionerd.connect goroutine checks if we're closed prior to attempting to establish a connection
|
||||
// with coderd, sees that we're not.
|
||||
// 2. A test closes the server.
|
||||
// 3. The provisionerd.connect goroutine calls the dialer to establish a connection. This
|
||||
// function detects that someone has tried to create a client after the test finishes.
|
||||
if ctx.Err() == nil {
|
||||
t.Error("createProvisionerDaemonClient cleanup after test was done!")
|
||||
}
|
||||
t.Error("called createProvisionerDaemonClient after test was done!")
|
||||
default:
|
||||
}
|
||||
return proto.NewDRPCProvisionerDaemonClient(clientPipe)
|
||||
@ -1235,15 +1185,25 @@ func (p *provisionerTestServer) Apply(s *provisionersdk.Session, r *sdkproto.App
|
||||
// Fulfills the protobuf interface for a ProvisionerDaemon with
|
||||
// passable functions for dynamic functionality.
|
||||
type provisionerDaemonTestServer struct {
|
||||
acquireJob func(ctx context.Context, _ *proto.Empty) (*proto.AcquiredJob, error)
|
||||
commitQuota func(ctx context.Context, com *proto.CommitQuotaRequest) (*proto.CommitQuotaResponse, error)
|
||||
updateJob func(ctx context.Context, update *proto.UpdateJobRequest) (*proto.UpdateJobResponse, error)
|
||||
failJob func(ctx context.Context, job *proto.FailedJob) (*proto.Empty, error)
|
||||
completeJob func(ctx context.Context, job *proto.CompletedJob) (*proto.Empty, error)
|
||||
acquireJobWithCancel func(stream proto.DRPCProvisionerDaemon_AcquireJobWithCancelStream) error
|
||||
commitQuota func(ctx context.Context, com *proto.CommitQuotaRequest) (*proto.CommitQuotaResponse, error)
|
||||
updateJob func(ctx context.Context, update *proto.UpdateJobRequest) (*proto.UpdateJobResponse, error)
|
||||
failJob func(ctx context.Context, job *proto.FailedJob) (*proto.Empty, error)
|
||||
completeJob func(ctx context.Context, job *proto.CompletedJob) (*proto.Empty, error)
|
||||
}
|
||||
|
||||
func (p *provisionerDaemonTestServer) AcquireJob(ctx context.Context, empty *proto.Empty) (*proto.AcquiredJob, error) {
|
||||
return p.acquireJob(ctx, empty)
|
||||
func (*provisionerDaemonTestServer) AcquireJob(context.Context, *proto.Empty) (*proto.AcquiredJob, error) {
|
||||
return nil, xerrors.New("deprecated!")
|
||||
}
|
||||
|
||||
func (p *provisionerDaemonTestServer) AcquireJobWithCancel(stream proto.DRPCProvisionerDaemon_AcquireJobWithCancelStream) error {
|
||||
if p.acquireJobWithCancel != nil {
|
||||
return p.acquireJobWithCancel(stream)
|
||||
}
|
||||
// default behavior is to wait for cancel
|
||||
_, _ = stream.Recv()
|
||||
_ = stream.Send(&proto.AcquiredJob{})
|
||||
return nil
|
||||
}
|
||||
|
||||
func (p *provisionerDaemonTestServer) CommitQuota(ctx context.Context, com *proto.CommitQuotaRequest) (*proto.CommitQuotaResponse, error) {
|
||||
@ -1266,3 +1226,38 @@ func (p *provisionerDaemonTestServer) FailJob(ctx context.Context, job *proto.Fa
|
||||
func (p *provisionerDaemonTestServer) CompleteJob(ctx context.Context, job *proto.CompletedJob) (*proto.Empty, error) {
|
||||
return p.completeJob(ctx, job)
|
||||
}
|
||||
|
||||
// acquireOne provides a function that returns a single provisioner job, then subsequent calls block until canceled.
|
||||
// The complete channel is closed on the 2nd call.
|
||||
type acquireOne struct {
|
||||
t *testing.T
|
||||
mu sync.Mutex
|
||||
job *proto.AcquiredJob
|
||||
called int
|
||||
complete chan struct{}
|
||||
}
|
||||
|
||||
func newAcquireOne(t *testing.T, job *proto.AcquiredJob) *acquireOne {
|
||||
return &acquireOne{
|
||||
t: t,
|
||||
job: job,
|
||||
complete: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
func (a *acquireOne) acquireWithCancel(stream proto.DRPCProvisionerDaemon_AcquireJobWithCancelStream) error {
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
a.called++
|
||||
if a.called == 2 {
|
||||
close(a.complete)
|
||||
}
|
||||
if a.called > 1 {
|
||||
_, _ = stream.Recv()
|
||||
_ = stream.Send(&proto.AcquiredJob{})
|
||||
return nil
|
||||
}
|
||||
err := stream.Send(a.job)
|
||||
assert.NoError(a.t, err)
|
||||
return nil
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user