Add periodic updates

kylecarbs · kylecarbs · commit 79875964d08f · 2022-02-01T23:13:31.000Z
diff --git a/provisionerd/provisionerd.go b/provisionerd/provisionerd.go
@@ -32,16 +32,20 @@ type Provisioners map[string]sdkproto.DRPCProvisionerClient
 type Options struct {
 	Logger slog.Logger
 
-	PollInterval  time.Duration
-	Provisioners  Provisioners
-	WorkDirectory string
+	UpdateInterval time.Duration
+	PollInterval   time.Duration
+	Provisioners   Provisioners
+	WorkDirectory  string
 }
 
 // New creates and starts a provisioner daemon.
 func New(clientDialer Dialer, opts *Options) io.Closer {
 	if opts.PollInterval == 0 {
 		opts.PollInterval = 5 * time.Second
 	}
+	if opts.UpdateInterval == 0 {
+		opts.UpdateInterval = 5 * time.Second
+	}
 	ctx, ctxCancel := context.WithCancel(context.Background())
 	daemon := &provisionerDaemon{
 		clientDialer: clientDialer,
@@ -192,6 +196,24 @@ func (p *provisionerDaemon) isRunningJob() bool {
 }
 
 func (p *provisionerDaemon) runJob(ctx context.Context) {
+	go func() {
+		ticker := time.NewTicker(p.opts.UpdateInterval)
+		defer ticker.Stop()
+		select {
+		case <-p.closed:
+			return
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+			err := p.updateStream.Send(&proto.JobUpdate{
+				JobId: p.acquiredJob.JobId,
+			})
+			if err != nil {
+				p.cancelActiveJob(fmt.Sprintf("send periodic update: %s", err))
+				return
+			}
+		}
+	}()
 	go func() {
 		select {
 		case <-p.closed:
diff --git a/provisionerd/provisionerd_test.go b/provisionerd/provisionerd_test.go
@@ -153,6 +153,48 @@ func TestProvisionerd(t *testing.T) {
 		require.NoError(t, closer.Close())
 	})
 
+	t.Run("RunningPeriodicUpdate", func(t *testing.T) {
+		t.Parallel()
+		completeChan := make(chan struct{})
+		closer := createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
+			return createProvisionerDaemonClient(t, provisionerDaemonTestServer{
+				acquireJob: func(ctx context.Context, _ *proto.Empty) (*proto.AcquiredJob, error) {
+					return &proto.AcquiredJob{
+						JobId:       "test",
+						Provisioner: "someprovisioner",
+						ProjectSourceArchive: createTar(t, map[string]string{
+							"test.txt": "content",
+						}),
+						Type: &proto.AcquiredJob_ProjectImport_{
+							ProjectImport: &proto.AcquiredJob_ProjectImport{},
+						},
+					}, nil
+				},
+				updateJob: func(stream proto.DRPCProvisionerDaemon_UpdateJobStream) error {
+					for {
+						_, err := stream.Recv()
+						if err != nil {
+							return err
+						}
+						close(completeChan)
+					}
+				},
+				cancelJob: func(ctx context.Context, job *proto.CancelledJob) (*proto.Empty, error) {
+					return &proto.Empty{}, nil
+				},
+			}), nil
+		}, provisionerd.Provisioners{
+			"someprovisioner": createProvisionerClient(t, provisionerTestServer{
+				parse: func(request *sdkproto.Parse_Request, stream sdkproto.DRPCProvisioner_ParseStream) error {
+					<-stream.Context().Done()
+					return nil
+				},
+			}),
+		})
+		<-completeChan
+		require.NoError(t, closer.Close())
+	})
+
 	t.Run("ProjectImport", func(t *testing.T) {
 		t.Parallel()
 		var (
@@ -331,10 +373,11 @@ 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, provisioners provisionerd.Provisioners) io.Closer {
 	closer := provisionerd.New(dialer, &provisionerd.Options{
-		Logger:        slogtest.Make(t, nil).Named("provisionerd").Leveled(slog.LevelDebug),
-		PollInterval:  50 * time.Millisecond,
-		Provisioners:  provisioners,
-		WorkDirectory: t.TempDir(),
+		Logger:         slogtest.Make(t, nil).Named("provisionerd").Leveled(slog.LevelDebug),
+		PollInterval:   50 * time.Millisecond,
+		UpdateInterval: 50 * time.Millisecond,
+		Provisioners:   provisioners,
+		WorkDirectory:  t.TempDir(),
 	})
 	t.Cleanup(func() {
 		_ = closer.Close()
