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>
2026-09-24 15:04:27 +08:00 · 2023-09-19 10:25:57 +04:00
parent 6cf531bfef
commit 375c70d141
41 changed files with 1497 additions and 1096 deletions
@@ -15,14 +15,13 @@ import (

 	"cdr.dev/slog"
 	"github.com/coder/coder/v2/coderd/database"
-	"github.com/coder/coder/v2/coderd/database/dbauthz"
 	"github.com/coder/coder/v2/coderd/database/dbtime"
+	"github.com/coder/coder/v2/coderd/database/provisionerjobs"
 	"github.com/coder/coder/v2/coderd/database/pubsub"
 )

 const (
-	EventJobPosted = "provisioner_job_posted"
-	dbMaxBackoff   = 10 * time.Second
+	dbMaxBackoff = 10 * time.Second
 	// backPollDuration is the period for the backup polling described in Acquirer comment
 	backupPollDuration = 30 * time.Second
 )
@@ -106,8 +105,6 @@ func (a *Acquirer) AcquireJob(
 	}
 	// buffer of 1 so that cancel doesn't deadlock while writing to the channel
 	clearance := make(chan struct{}, 1)
-	//nolint:gocritic // Provisionerd has specific authz rules.
-	principal := dbauthz.AsProvisionerd(ctx)
 	for {
 		a.want(pt, tags, clearance)
 		select {
@@ -122,7 +119,7 @@ func (a *Acquirer) AcquireJob(
 			return database.ProvisionerJob{}, err
 		case <-clearance:
 			logger.Debug(ctx, "got clearance to call database")
-			job, err := a.store.AcquireProvisionerJob(principal, database.AcquireProvisionerJobParams{
+			job, err := a.store.AcquireProvisionerJob(ctx, database.AcquireProvisionerJobParams{
 				StartedAt: sql.NullTime{
 					Time:  dbtime.Now(),
 					Valid: true,
@@ -298,7 +295,7 @@ func (a *Acquirer) subscribe() {
 		bkoff := backoff.WithContext(eb, a.ctx)
 		var cancel context.CancelFunc
 		err := backoff.Retry(func() error {
-			cancelFn, err := a.ps.SubscribeWithErr(EventJobPosted, a.jobPosted)
+			cancelFn, err := a.ps.SubscribeWithErr(provisionerjobs.EventJobPosted, a.jobPosted)
 			if err != nil {
 				a.logger.Warn(a.ctx, "failed to subscribe to job postings", slog.Error(err))
 				return err
@@ -335,7 +332,7 @@ func (a *Acquirer) jobPosted(ctx context.Context, message []byte, err error) {
 		a.logger.Warn(a.ctx, "unhandled pubsub error", slog.Error(err))
 		return
 	}
-	posting := JobPosting{}
+	posting := provisionerjobs.JobPosting{}
 	err = json.Unmarshal(message, &posting)
 	if err != nil {
 		a.logger.Error(a.ctx, "unable to parse job posting",
@@ -457,7 +454,7 @@ type domain struct {
 	acquirees map[chan<- struct{}]*acquiree
 }

-func (d domain) contains(p JobPosting) bool {
+func (d domain) contains(p provisionerjobs.JobPosting) bool {
 	if !slices.Contains(d.pt, p.ProvisionerType) {
 		return false
 	}
@@ -485,8 +482,3 @@ func (d domain) poll(dur time.Duration) {
 		}
 	}
 }
-
-type JobPosting struct {
-	ProvisionerType database.ProvisionerType `json:"type"`
-	Tags            map[string]string        `json:"tags"`
-}
@@ -18,6 +18,7 @@ import (
 	"cdr.dev/slog/sloggers/slogtest"
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbfake"
+	"github.com/coder/coder/v2/coderd/database/provisionerjobs"
 	"github.com/coder/coder/v2/coderd/database/pubsub"
 	"github.com/coder/coder/v2/coderd/provisionerdserver"
 	"github.com/coder/coder/v2/testutil"
@@ -316,12 +317,12 @@ func TestAcquirer_UnblockOnCancel(t *testing.T) {

 func postJob(t *testing.T, ps pubsub.Pubsub, pt database.ProvisionerType, tags provisionerdserver.Tags) {
 	t.Helper()
-	msg, err := json.Marshal(provisionerdserver.JobPosting{
+	msg, err := json.Marshal(provisionerjobs.JobPosting{
 		ProvisionerType: pt,
 		Tags:            tags,
 	})
 	require.NoError(t, err)
-	err = ps.Publish(provisionerdserver.EventJobPosted, msg)
+	err = ps.Publish(provisionerjobs.EventJobPosted, msg)
 	require.NoError(t, err)
 }

@@ -11,7 +11,6 @@ import (
 	"reflect"
 	"strconv"
 	"strings"
-	"sync"
 	"sync/atomic"
 	"time"

@@ -44,16 +43,18 @@ import (
 	sdkproto "github.com/coder/coder/v2/provisionersdk/proto"
 )

-var (
-	lastAcquire      time.Time
-	lastAcquireMutex sync.RWMutex
-)
+// DefaultAcquireJobLongPollDur is the time the (deprecated) AcquireJob rpc waits to try to obtain a job before
+// canceling and returning an empty job.
+const DefaultAcquireJobLongPollDur = time.Second * 5

 type Options struct {
 	OIDCConfig     httpmw.OAuth2Config
 	GitAuthConfigs []*gitauth.Config
 	// TimeNowFn is only used in tests
 	TimeNowFn func() time.Time
+
+	// AcquireJobLongPollDur is used in tests
+	AcquireJobLongPollDur time.Duration
 }

 type server struct {
@@ -62,9 +63,10 @@ type server struct {
 	Logger                      slog.Logger
 	Provisioners                []database.ProvisionerType
 	GitAuthConfigs              []*gitauth.Config
-	Tags                        json.RawMessage
+	Tags                        Tags
 	Database                    database.Store
 	Pubsub                      pubsub.Pubsub
+	Acquirer                    *Acquirer
 	Telemetry                   telemetry.Reporter
 	Tracer                      trace.Tracer
 	QuotaCommitter              *atomic.Pointer[proto.QuotaCommitter]
@@ -73,10 +75,11 @@ type server struct {
 	UserQuietHoursScheduleStore *atomic.Pointer[schedule.UserQuietHoursScheduleStore]
 	DeploymentValues            *codersdk.DeploymentValues

-	AcquireJobDebounce time.Duration
-	OIDCConfig         httpmw.OAuth2Config
+	OIDCConfig httpmw.OAuth2Config

 	TimeNowFn func() time.Time
+
+	acquireJobLongPollDur time.Duration
 }

 // We use the null byte (0x00) in generating a canonical map key for tags, so
@@ -108,9 +111,10 @@ func NewServer(
 	id uuid.UUID,
 	logger slog.Logger,
 	provisioners []database.ProvisionerType,
-	tags json.RawMessage,
+	tags Tags,
 	db database.Store,
 	ps pubsub.Pubsub,
+	acquirer *Acquirer,
 	tel telemetry.Reporter,
 	tracer trace.Tracer,
 	quotaCommitter *atomic.Pointer[proto.QuotaCommitter],
@@ -118,7 +122,6 @@ func NewServer(
 	templateScheduleStore *atomic.Pointer[schedule.TemplateScheduleStore],
 	userQuietHoursScheduleStore *atomic.Pointer[schedule.UserQuietHoursScheduleStore],
 	deploymentValues *codersdk.DeploymentValues,
-	acquireJobDebounce time.Duration,
 	options Options,
 ) (proto.DRPCProvisionerDaemonServer, error) {
 	// Panic early if pointers are nil
@@ -137,6 +140,18 @@ func NewServer(
 	if deploymentValues == nil {
 		return nil, xerrors.New("deploymentValues is nil")
 	}
+	if acquirer == nil {
+		return nil, xerrors.New("acquirer is nil")
+	}
+	if tags == nil {
+		return nil, xerrors.Errorf("tags is nil")
+	}
+	if err := tags.Valid(); err != nil {
+		return nil, xerrors.Errorf("invalid tags: %w", err)
+	}
+	if options.AcquireJobLongPollDur == 0 {
+		options.AcquireJobLongPollDur = DefaultAcquireJobLongPollDur
+	}
 	return &server{
 		AccessURL:                   accessURL,
 		ID:                          id,
@@ -146,6 +161,7 @@ func NewServer(
 		Tags:                        tags,
 		Database:                    db,
 		Pubsub:                      ps,
+		Acquirer:                    acquirer,
 		Telemetry:                   tel,
 		Tracer:                      tracer,
 		QuotaCommitter:              quotaCommitter,
@@ -153,9 +169,9 @@ func NewServer(
 		TemplateScheduleStore:       templateScheduleStore,
 		UserQuietHoursScheduleStore: userQuietHoursScheduleStore,
 		DeploymentValues:            deploymentValues,
-		AcquireJobDebounce:          acquireJobDebounce,
 		OIDCConfig:                  options.OIDCConfig,
 		TimeNowFn:                   options.TimeNowFn,
+		acquireJobLongPollDur:       options.AcquireJobLongPollDur,
 	}, nil
 }

@@ -169,50 +185,119 @@ func (s *server) timeNow() time.Time {
 }

 // AcquireJob queries the database to lock a job.
+//
+// Deprecated: This method is only available for back-level provisioner daemons.
 func (s *server) AcquireJob(ctx context.Context, _ *proto.Empty) (*proto.AcquiredJob, error) {
 	//nolint:gocritic // Provisionerd has specific authz rules.
 	ctx = dbauthz.AsProvisionerd(ctx)
-	// This prevents loads of provisioner daemons from consistently
-	// querying the database when no jobs are available.
-	//
-	// The debounce only occurs when no job is returned, so if loads of
-	// jobs are added at once, they will start after at most this duration.
-	lastAcquireMutex.RLock()
-	if !lastAcquire.IsZero() && time.Since(lastAcquire) < s.AcquireJobDebounce {
-		s.Logger.Debug(ctx, "debounce acquire job", slog.F("debounce", s.AcquireJobDebounce), slog.F("last_acquire", lastAcquire))
-		lastAcquireMutex.RUnlock()
-		return &proto.AcquiredJob{}, nil
-	}
-	lastAcquireMutex.RUnlock()
-	// This marks the job as locked in the database.
-	job, err := s.Database.AcquireProvisionerJob(ctx, database.AcquireProvisionerJobParams{
-		StartedAt: sql.NullTime{
-			Time:  dbtime.Now(),
-			Valid: true,
-		},
-		WorkerID: uuid.NullUUID{
-			UUID:  s.ID,
-			Valid: true,
-		},
-		Types: s.Provisioners,
-		Tags:  s.Tags,
-	})
-	if errors.Is(err, sql.ErrNoRows) {
-		// The provisioner daemon assumes no jobs are available if
-		// an empty struct is returned.
-		lastAcquireMutex.Lock()
-		lastAcquire = dbtime.Now()
-		lastAcquireMutex.Unlock()
+	// Since AcquireJob blocks until a job is available, we set a long (5s by default) timeout.  This allows back-level
+	// provisioner daemons to gracefully shut down within a few seconds, but keeps them from rapidly polling the
+	// database.
+	acqCtx, acqCancel := context.WithTimeout(ctx, s.acquireJobLongPollDur)
+	defer acqCancel()
+	job, err := s.Acquirer.AcquireJob(acqCtx, s.ID, s.Provisioners, s.Tags)
+	if xerrors.Is(err, context.DeadlineExceeded) {
+		s.Logger.Debug(ctx, "successful cancel")
 		return &proto.AcquiredJob{}, nil
 	}
 	if err != nil {
 		return nil, xerrors.Errorf("acquire job: %w", err)
 	}
 	s.Logger.Debug(ctx, "locked job from database", slog.F("job_id", job.ID))
+	return s.acquireProtoJob(ctx, job)
+}

+type jobAndErr struct {
+	job database.ProvisionerJob
+	err error
+}
+
+// AcquireJobWithCancel queries the database to lock a job.
+func (s *server) AcquireJobWithCancel(stream proto.DRPCProvisionerDaemon_AcquireJobWithCancelStream) (retErr error) {
+	//nolint:gocritic // Provisionerd has specific authz rules.
+	streamCtx := dbauthz.AsProvisionerd(stream.Context())
+	defer func() {
+		closeErr := stream.Close()
+		s.Logger.Debug(streamCtx, "closed stream", slog.Error(closeErr))
+		if retErr == nil {
+			retErr = closeErr
+		}
+	}()
+	acqCtx, acqCancel := context.WithCancel(streamCtx)
+	defer acqCancel()
+	recvCh := make(chan error, 1)
+	go func() {
+		_, err := stream.Recv() // cancel is the only message
+		recvCh <- err
+	}()
+	jec := make(chan jobAndErr, 1)
+	go func() {
+		job, err := s.Acquirer.AcquireJob(acqCtx, s.ID, s.Provisioners, s.Tags)
+		jec <- jobAndErr{job: job, err: err}
+	}()
+	var recvErr error
+	var je jobAndErr
+	select {
+	case recvErr = <-recvCh:
+		acqCancel()
+		je = <-jec
+	case je = <-jec:
+	}
+	if xerrors.Is(je.err, context.Canceled) {
+		s.Logger.Debug(streamCtx, "successful cancel")
+		err := stream.Send(&proto.AcquiredJob{})
+		if err != nil {
+			// often this is just because the other side hangs up and doesn't wait for the cancel, so log at INFO
+			s.Logger.Info(streamCtx, "failed to send empty job", slog.Error(err))
+			return err
+		}
+		return nil
+	}
+	if je.err != nil {
+		return xerrors.Errorf("acquire job: %w", je.err)
+	}
+	logger := s.Logger.With(slog.F("job_id", je.job.ID))
+	logger.Debug(streamCtx, "locked job from database")
+
+	if recvErr != nil {
+		logger.Error(streamCtx, "recv error and failed to cancel acquire job", slog.Error(recvErr))
+		// Well, this is awkward.  We hit an error receiving from the stream, but didn't cancel before we locked a job
+		// in the database.  We need to mark this job as failed so the end user can retry if they want to.
+		err := s.Database.UpdateProvisionerJobWithCompleteByID(
+			context.Background(),
+			database.UpdateProvisionerJobWithCompleteByIDParams{
+				ID: je.job.ID,
+				CompletedAt: sql.NullTime{
+					Time:  dbtime.Now(),
+					Valid: true,
+				},
+				Error: sql.NullString{
+					String: "connection to provisioner daemon broken",
+					Valid:  true,
+				},
+			})
+		if err != nil {
+			logger.Error(streamCtx, "error updating failed job", slog.Error(err))
+		}
+		return recvErr
+	}
+
+	pj, err := s.acquireProtoJob(streamCtx, je.job)
+	if err != nil {
+		return err
+	}
+	err = stream.Send(pj)
+	if err != nil {
+		s.Logger.Error(streamCtx, "failed to send job", slog.Error(err))
+		return err
+	}
+	return nil
+}
+
+func (s *server) acquireProtoJob(ctx context.Context, job database.ProvisionerJob) (*proto.AcquiredJob, error) {
 	// Marks the acquired job as failed with the error message provided.
 	failJob := func(errorMessage string) error {
-		err = s.Database.UpdateProvisionerJobWithCompleteByID(ctx, database.UpdateProvisionerJobWithCompleteByIDParams{
+		err := s.Database.UpdateProvisionerJobWithCompleteByID(ctx, database.UpdateProvisionerJobWithCompleteByIDParams{
 			ID: job.ID,
 			CompletedAt: sql.NullTime{
 				Time:  dbtime.Now(),
@@ -4,12 +4,19 @@ import (
 	"context"
 	"database/sql"
 	"encoding/json"
+	"io"
 	"net/url"
 	"strings"
+	"sync"
 	"sync/atomic"
 	"testing"
 	"time"

+	"golang.org/x/xerrors"
+	"storj.io/drpc"
+
+	"cdr.dev/slog"
+
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -57,400 +64,411 @@ func testUserQuietHoursScheduleStore() *atomic.Pointer[schedule.UserQuietHoursSc
 	return ptr
 }

+func TestAcquireJob_LongPoll(t *testing.T) {
+	t.Parallel()
+	srv, _, _ := setup(t, false, &overrides{acquireJobLongPollDuration: time.Microsecond})
+	job, err := srv.AcquireJob(context.Background(), nil)
+	require.NoError(t, err)
+	require.Equal(t, &proto.AcquiredJob{}, job)
+}
+
+func TestAcquireJobWithCancel_Cancel(t *testing.T) {
+	t.Parallel()
+	srv, _, _ := setup(t, false, nil)
+	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
+	defer cancel()
+	fs := newFakeStream(ctx)
+	errCh := make(chan error)
+	go func() {
+		errCh <- srv.AcquireJobWithCancel(fs)
+	}()
+	fs.cancel()
+	select {
+	case <-ctx.Done():
+		t.Fatal("timed out waiting for AcquireJobWithCancel")
+	case err := <-errCh:
+		require.NoError(t, err)
+	}
+	job, err := fs.waitForJob()
+	require.NoError(t, err)
+	require.NotNil(t, job)
+	require.Equal(t, "", job.JobId)
+}
+
 func TestAcquireJob(t *testing.T) {
 	t.Parallel()
-	t.Run("Debounce", func(t *testing.T) {
-		t.Parallel()
-		db := dbfake.New()
-		ps := pubsub.NewInMemory()
-		srv, err := provisionerdserver.NewServer(
-			&url.URL{},
-			uuid.New(),
-			slogtest.Make(t, nil),
-			[]database.ProvisionerType{database.ProvisionerTypeEcho},
-			nil,
-			db,
-			ps,
-			telemetry.NewNoop(),
-			trace.NewNoopTracerProvider().Tracer("noop"),
-			&atomic.Pointer[proto.QuotaCommitter]{},
-			mockAuditor(),
-			testTemplateScheduleStore(),
-			testUserQuietHoursScheduleStore(),
-			&codersdk.DeploymentValues{},
-			time.Hour,
-			provisionerdserver.Options{},
-		)
-		require.NoError(t, err)
-		job, err := srv.AcquireJob(context.Background(), nil)
-		require.NoError(t, err)
-		require.Equal(t, &proto.AcquiredJob{}, job)
-		_, err = db.InsertProvisionerJob(context.Background(), database.InsertProvisionerJobParams{
-			ID:            uuid.New(),
-			InitiatorID:   uuid.New(),
-			Provisioner:   database.ProvisionerTypeEcho,
-			StorageMethod: database.ProvisionerStorageMethodFile,
-			Type:          database.ProvisionerJobTypeTemplateVersionDryRun,
-		})
-		require.NoError(t, err)
-		job, err = srv.AcquireJob(context.Background(), nil)
-		require.NoError(t, err)
-		require.Equal(t, &proto.AcquiredJob{}, job)
-	})
-	t.Run("NoJobs", func(t *testing.T) {
-		t.Parallel()
-		srv, _, _ := setup(t, false, nil)
-		job, err := srv.AcquireJob(context.Background(), nil)
-		require.NoError(t, err)
-		require.Equal(t, &proto.AcquiredJob{}, job)
-	})
-	t.Run("InitiatorNotFound", func(t *testing.T) {
-		t.Parallel()
-		srv, db, _ := setup(t, false, nil)
-		_, err := db.InsertProvisionerJob(context.Background(), database.InsertProvisionerJobParams{
-			ID:            uuid.New(),
-			InitiatorID:   uuid.New(),
-			Provisioner:   database.ProvisionerTypeEcho,
-			StorageMethod: database.ProvisionerStorageMethodFile,
-			Type:          database.ProvisionerJobTypeTemplateVersionDryRun,
-		})
-		require.NoError(t, err)
-		_, err = srv.AcquireJob(context.Background(), nil)
-		require.ErrorContains(t, err, "sql: no rows in result set")
-	})
-	t.Run("WorkspaceBuildJob", func(t *testing.T) {
-		t.Parallel()
-		// Set the max session token lifetime so we can assert we
-		// create an API key with an expiration within the bounds of the
-		// deployment config.
-		dv := &codersdk.DeploymentValues{MaxTokenLifetime: clibase.Duration(time.Hour)}
-		gitAuthProvider := "github"
-		srv, db, ps := setup(t, false, &overrides{
-			deploymentValues: dv,
-			gitAuthConfigs: []*gitauth.Config{{
-				ID:           gitAuthProvider,
-				OAuth2Config: &testutil.OAuth2Config{},
-			}},
-		})
-		ctx := context.Background()

-		user := dbgen.User(t, db, database.User{})
-		link := dbgen.UserLink(t, db, database.UserLink{
-			LoginType:        database.LoginTypeOIDC,
-			UserID:           user.ID,
-			OAuthExpiry:      dbtime.Now().Add(time.Hour),
-			OAuthAccessToken: "access-token",
-		})
-		dbgen.GitAuthLink(t, db, database.GitAuthLink{
-			ProviderID: gitAuthProvider,
-			UserID:     user.ID,
-		})
-		template := dbgen.Template(t, db, database.Template{
-			Name:        "template",
-			Provisioner: database.ProvisionerTypeEcho,
-		})
-		file := dbgen.File(t, db, database.File{CreatedBy: user.ID})
-		versionFile := dbgen.File(t, db, database.File{CreatedBy: user.ID})
-		version := dbgen.TemplateVersion(t, db, database.TemplateVersion{
-			TemplateID: uuid.NullUUID{
-				UUID:  template.ID,
-				Valid: true,
-			},
-			JobID: uuid.New(),
-		})
-		err := db.UpdateTemplateVersionGitAuthProvidersByJobID(ctx, database.UpdateTemplateVersionGitAuthProvidersByJobIDParams{
-			JobID:            version.JobID,
-			GitAuthProviders: []string{gitAuthProvider},
-			UpdatedAt:        dbtime.Now(),
-		})
-		require.NoError(t, err)
-		// Import version job
-		_ = dbgen.ProvisionerJob(t, db, database.ProvisionerJob{
-			ID:            version.JobID,
-			InitiatorID:   user.ID,
-			FileID:        versionFile.ID,
-			Provisioner:   database.ProvisionerTypeEcho,
-			StorageMethod: database.ProvisionerStorageMethodFile,
-			Type:          database.ProvisionerJobTypeTemplateVersionImport,
-			Input: must(json.Marshal(provisionerdserver.TemplateVersionImportJob{
-				TemplateVersionID: version.ID,
-				UserVariableValues: []codersdk.VariableValue{
-					{Name: "second", Value: "bah"},
-				},
-			})),
-		})
-		_ = dbgen.TemplateVersionVariable(t, db, database.TemplateVersionVariable{
-			TemplateVersionID: version.ID,
-			Name:              "first",
-			Value:             "first_value",
-			DefaultValue:      "default_value",
-			Sensitive:         true,
-		})
-		_ = dbgen.TemplateVersionVariable(t, db, database.TemplateVersionVariable{
-			TemplateVersionID: version.ID,
-			Name:              "second",
-			Value:             "second_value",
-			DefaultValue:      "default_value",
-			Required:          true,
-			Sensitive:         false,
-		})
-		workspace := dbgen.Workspace(t, db, database.Workspace{
-			TemplateID: template.ID,
-			OwnerID:    user.ID,
-		})
-		build := dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
-			WorkspaceID:       workspace.ID,
-			BuildNumber:       1,
-			JobID:             uuid.New(),
-			TemplateVersionID: version.ID,
-			Transition:        database.WorkspaceTransitionStart,
-			Reason:            database.BuildReasonInitiator,
-		})
-		_ = dbgen.ProvisionerJob(t, db, database.ProvisionerJob{
-			ID:            build.ID,
-			InitiatorID:   user.ID,
-			Provisioner:   database.ProvisionerTypeEcho,
-			StorageMethod: database.ProvisionerStorageMethodFile,
-			FileID:        file.ID,
-			Type:          database.ProvisionerJobTypeWorkspaceBuild,
-			Input: must(json.Marshal(provisionerdserver.WorkspaceProvisionJob{
-				WorkspaceBuildID: build.ID,
-			})),
-		})
-
-		startPublished := make(chan struct{})
-		var closed bool
-		closeStartSubscribe, err := ps.Subscribe(codersdk.WorkspaceNotifyChannel(workspace.ID), func(_ context.Context, _ []byte) {
-			if !closed {
-				close(startPublished)
-				closed = true
+	// These test acquiring a single job without canceling, and tests both AcquireJob (deprecated) and
+	// AcquireJobWithCancel as the way to get the job.
+	cases := []struct {
+		name    string
+		acquire func(context.Context, proto.DRPCProvisionerDaemonServer) (*proto.AcquiredJob, error)
+	}{
+		{name: "Deprecated", acquire: func(ctx context.Context, srv proto.DRPCProvisionerDaemonServer) (*proto.AcquiredJob, error) {
+			return srv.AcquireJob(ctx, nil)
+		}},
+		{name: "WithCancel", acquire: func(ctx context.Context, srv proto.DRPCProvisionerDaemonServer) (*proto.AcquiredJob, error) {
+			fs := newFakeStream(ctx)
+			err := srv.AcquireJobWithCancel(fs)
+			if err != nil {
+				return nil, err
 			}
+			return fs.waitForJob()
+		}},
+	}
+	for _, tc := range cases {
+		tc := tc
+		t.Run(tc.name+"_InitiatorNotFound", func(t *testing.T) {
+			t.Parallel()
+			srv, db, _ := setup(t, false, nil)
+			ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
+			defer cancel()
+			_, err := db.InsertProvisionerJob(context.Background(), database.InsertProvisionerJobParams{
+				ID:            uuid.New(),
+				InitiatorID:   uuid.New(),
+				Provisioner:   database.ProvisionerTypeEcho,
+				StorageMethod: database.ProvisionerStorageMethodFile,
+				Type:          database.ProvisionerJobTypeTemplateVersionDryRun,
+			})
+			require.NoError(t, err)
+			_, err = tc.acquire(ctx, srv)
+			require.ErrorContains(t, err, "sql: no rows in result set")
 		})
-		require.NoError(t, err)
-		defer closeStartSubscribe()
+		t.Run(tc.name+"_WorkspaceBuildJob", func(t *testing.T) {
+			t.Parallel()
+			// Set the max session token lifetime so we can assert we
+			// create an API key with an expiration within the bounds of the
+			// deployment config.
+			dv := &codersdk.DeploymentValues{MaxTokenLifetime: clibase.Duration(time.Hour)}
+			gitAuthProvider := "github"
+			srv, db, ps := setup(t, false, &overrides{
+				deploymentValues: dv,
+				gitAuthConfigs: []*gitauth.Config{{
+					ID:           gitAuthProvider,
+					OAuth2Config: &testutil.OAuth2Config{},
+				}},
+			})
+			ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
+			defer cancel()

-		var job *proto.AcquiredJob
+			user := dbgen.User(t, db, database.User{})
+			link := dbgen.UserLink(t, db, database.UserLink{
+				LoginType:        database.LoginTypeOIDC,
+				UserID:           user.ID,
+				OAuthExpiry:      dbtime.Now().Add(time.Hour),
+				OAuthAccessToken: "access-token",
+			})
+			dbgen.GitAuthLink(t, db, database.GitAuthLink{
+				ProviderID: gitAuthProvider,
+				UserID:     user.ID,
+			})
+			template := dbgen.Template(t, db, database.Template{
+				Name:        "template",
+				Provisioner: database.ProvisionerTypeEcho,
+			})
+			file := dbgen.File(t, db, database.File{CreatedBy: user.ID})
+			versionFile := dbgen.File(t, db, database.File{CreatedBy: user.ID})
+			version := dbgen.TemplateVersion(t, db, database.TemplateVersion{
+				TemplateID: uuid.NullUUID{
+					UUID:  template.ID,
+					Valid: true,
+				},
+				JobID: uuid.New(),
+			})
+			err := db.UpdateTemplateVersionGitAuthProvidersByJobID(ctx, database.UpdateTemplateVersionGitAuthProvidersByJobIDParams{
+				JobID:            version.JobID,
+				GitAuthProviders: []string{gitAuthProvider},
+				UpdatedAt:        dbtime.Now(),
+			})
+			require.NoError(t, err)
+			// Import version job
+			_ = dbgen.ProvisionerJob(t, db, ps, database.ProvisionerJob{
+				ID:            version.JobID,
+				InitiatorID:   user.ID,
+				FileID:        versionFile.ID,
+				Provisioner:   database.ProvisionerTypeEcho,
+				StorageMethod: database.ProvisionerStorageMethodFile,
+				Type:          database.ProvisionerJobTypeTemplateVersionImport,
+				Input: must(json.Marshal(provisionerdserver.TemplateVersionImportJob{
+					TemplateVersionID: version.ID,
+					UserVariableValues: []codersdk.VariableValue{
+						{Name: "second", Value: "bah"},
+					},
+				})),
+			})
+			_ = dbgen.TemplateVersionVariable(t, db, database.TemplateVersionVariable{
+				TemplateVersionID: version.ID,
+				Name:              "first",
+				Value:             "first_value",
+				DefaultValue:      "default_value",
+				Sensitive:         true,
+			})
+			_ = dbgen.TemplateVersionVariable(t, db, database.TemplateVersionVariable{
+				TemplateVersionID: version.ID,
+				Name:              "second",
+				Value:             "second_value",
+				DefaultValue:      "default_value",
+				Required:          true,
+				Sensitive:         false,
+			})
+			workspace := dbgen.Workspace(t, db, database.Workspace{
+				TemplateID: template.ID,
+				OwnerID:    user.ID,
+			})
+			build := dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
+				WorkspaceID:       workspace.ID,
+				BuildNumber:       1,
+				JobID:             uuid.New(),
+				TemplateVersionID: version.ID,
+				Transition:        database.WorkspaceTransitionStart,
+				Reason:            database.BuildReasonInitiator,
+			})
+			_ = dbgen.ProvisionerJob(t, db, ps, database.ProvisionerJob{
+				ID:            build.ID,
+				InitiatorID:   user.ID,
+				Provisioner:   database.ProvisionerTypeEcho,
+				StorageMethod: database.ProvisionerStorageMethodFile,
+				FileID:        file.ID,
+				Type:          database.ProvisionerJobTypeWorkspaceBuild,
+				Input: must(json.Marshal(provisionerdserver.WorkspaceProvisionJob{
+					WorkspaceBuildID: build.ID,
+				})),
+			})
+
+			startPublished := make(chan struct{})
+			var closed bool
+			closeStartSubscribe, err := ps.Subscribe(codersdk.WorkspaceNotifyChannel(workspace.ID), func(_ context.Context, _ []byte) {
+				if !closed {
+					close(startPublished)
+					closed = true
+				}
+			})
+			require.NoError(t, err)
+			defer closeStartSubscribe()
+
+			var job *proto.AcquiredJob
+
+			for {
+				// Grab jobs until we find the workspace build job. There is also
+				// an import version job that we need to ignore.
+				job, err = tc.acquire(ctx, srv)
+				require.NoError(t, err)
+				if _, ok := job.Type.(*proto.AcquiredJob_WorkspaceBuild_); ok {
+					break
+				}
+			}
+
+			<-startPublished
+
+			got, err := json.Marshal(job.Type)
+			require.NoError(t, err)
+
+			// Validate that a session token is generated during the job.
+			sessionToken := job.Type.(*proto.AcquiredJob_WorkspaceBuild_).WorkspaceBuild.Metadata.WorkspaceOwnerSessionToken
+			require.NotEmpty(t, sessionToken)
+			toks := strings.Split(sessionToken, "-")
+			require.Len(t, toks, 2, "invalid api key")
+			key, err := db.GetAPIKeyByID(ctx, toks[0])
+			require.NoError(t, err)
+			require.Equal(t, int64(dv.MaxTokenLifetime.Value().Seconds()), key.LifetimeSeconds)
+			require.WithinDuration(t, time.Now().Add(dv.MaxTokenLifetime.Value()), key.ExpiresAt, time.Minute)
+
+			want, err := json.Marshal(&proto.AcquiredJob_WorkspaceBuild_{
+				WorkspaceBuild: &proto.AcquiredJob_WorkspaceBuild{
+					WorkspaceBuildId: build.ID.String(),
+					WorkspaceName:    workspace.Name,
+					VariableValues: []*sdkproto.VariableValue{
+						{
+							Name:      "first",
+							Value:     "first_value",
+							Sensitive: true,
+						},
+						{
+							Name:  "second",
+							Value: "second_value",
+						},
+					},
+					GitAuthProviders: []*sdkproto.GitAuthProvider{{
+						Id:          gitAuthProvider,
+						AccessToken: "access_token",
+					}},
+					Metadata: &sdkproto.Metadata{
+						CoderUrl:                      (&url.URL{}).String(),
+						WorkspaceTransition:           sdkproto.WorkspaceTransition_START,
+						WorkspaceName:                 workspace.Name,
+						WorkspaceOwner:                user.Username,
+						WorkspaceOwnerEmail:           user.Email,
+						WorkspaceOwnerOidcAccessToken: link.OAuthAccessToken,
+						WorkspaceId:                   workspace.ID.String(),
+						WorkspaceOwnerId:              user.ID.String(),
+						TemplateId:                    template.ID.String(),
+						TemplateName:                  template.Name,
+						TemplateVersion:               version.Name,
+						WorkspaceOwnerSessionToken:    sessionToken,
+					},
+				},
+			})
+			require.NoError(t, err)
+
+			require.JSONEq(t, string(want), string(got))
+
+			// Assert that we delete the session token whenever
+			// a stop is issued.
+			stopbuild := dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
+				WorkspaceID:       workspace.ID,
+				BuildNumber:       2,
+				JobID:             uuid.New(),
+				TemplateVersionID: version.ID,
+				Transition:        database.WorkspaceTransitionStop,
+				Reason:            database.BuildReasonInitiator,
+			})
+			_ = dbgen.ProvisionerJob(t, db, ps, database.ProvisionerJob{
+				ID:            stopbuild.ID,
+				InitiatorID:   user.ID,
+				Provisioner:   database.ProvisionerTypeEcho,
+				StorageMethod: database.ProvisionerStorageMethodFile,
+				FileID:        file.ID,
+				Type:          database.ProvisionerJobTypeWorkspaceBuild,
+				Input: must(json.Marshal(provisionerdserver.WorkspaceProvisionJob{
+					WorkspaceBuildID: stopbuild.ID,
+				})),
+			})
+
+			stopPublished := make(chan struct{})
+			closeStopSubscribe, err := ps.Subscribe(codersdk.WorkspaceNotifyChannel(workspace.ID), func(_ context.Context, _ []byte) {
+				close(stopPublished)
+			})
+			require.NoError(t, err)
+			defer closeStopSubscribe()

-		for {
 			// Grab jobs until we find the workspace build job. There is also
 			// an import version job that we need to ignore.
-			job, err = srv.AcquireJob(ctx, nil)
+			job, err = tc.acquire(ctx, srv)
 			require.NoError(t, err)
-			if _, ok := job.Type.(*proto.AcquiredJob_WorkspaceBuild_); ok {
-				break
-			}
-		}
+			_, ok := job.Type.(*proto.AcquiredJob_WorkspaceBuild_)
+			require.True(t, ok, "acquired job not a workspace build?")

-		<-startPublished
+			<-stopPublished

-		got, err := json.Marshal(job.Type)
-		require.NoError(t, err)
+			// Validate that a session token is deleted during a stop job.
+			sessionToken = job.Type.(*proto.AcquiredJob_WorkspaceBuild_).WorkspaceBuild.Metadata.WorkspaceOwnerSessionToken
+			require.Empty(t, sessionToken)
+			_, err = db.GetAPIKeyByID(ctx, key.ID)
+			require.ErrorIs(t, err, sql.ErrNoRows)
+		})

-		// Validate that a session token is generated during the job.
-		sessionToken := job.Type.(*proto.AcquiredJob_WorkspaceBuild_).WorkspaceBuild.Metadata.WorkspaceOwnerSessionToken
-		require.NotEmpty(t, sessionToken)
-		toks := strings.Split(sessionToken, "-")
-		require.Len(t, toks, 2, "invalid api key")
-		key, err := db.GetAPIKeyByID(ctx, toks[0])
-		require.NoError(t, err)
-		require.Equal(t, int64(dv.MaxTokenLifetime.Value().Seconds()), key.LifetimeSeconds)
-		require.WithinDuration(t, time.Now().Add(dv.MaxTokenLifetime.Value()), key.ExpiresAt, time.Minute)
+		t.Run(tc.name+"_TemplateVersionDryRun", func(t *testing.T) {
+			t.Parallel()
+			srv, db, ps := setup(t, false, nil)
+			ctx := context.Background()

-		want, err := json.Marshal(&proto.AcquiredJob_WorkspaceBuild_{
-			WorkspaceBuild: &proto.AcquiredJob_WorkspaceBuild{
-				WorkspaceBuildId: build.ID.String(),
-				WorkspaceName:    workspace.Name,
-				VariableValues: []*sdkproto.VariableValue{
-					{
-						Name:      "first",
-						Value:     "first_value",
-						Sensitive: true,
-					},
-					{
-						Name:  "second",
-						Value: "second_value",
+			user := dbgen.User(t, db, database.User{})
+			version := dbgen.TemplateVersion(t, db, database.TemplateVersion{})
+			file := dbgen.File(t, db, database.File{CreatedBy: user.ID})
+			_ = dbgen.ProvisionerJob(t, db, ps, database.ProvisionerJob{
+				InitiatorID:   user.ID,
+				Provisioner:   database.ProvisionerTypeEcho,
+				StorageMethod: database.ProvisionerStorageMethodFile,
+				FileID:        file.ID,
+				Type:          database.ProvisionerJobTypeTemplateVersionDryRun,
+				Input: must(json.Marshal(provisionerdserver.TemplateVersionDryRunJob{
+					TemplateVersionID: version.ID,
+					WorkspaceName:     "testing",
+				})),
+			})
+
+			job, err := tc.acquire(ctx, srv)
+			require.NoError(t, err)
+
+			got, err := json.Marshal(job.Type)
+			require.NoError(t, err)
+
+			want, err := json.Marshal(&proto.AcquiredJob_TemplateDryRun_{
+				TemplateDryRun: &proto.AcquiredJob_TemplateDryRun{
+					Metadata: &sdkproto.Metadata{
+						CoderUrl:      (&url.URL{}).String(),
+						WorkspaceName: "testing",
 					},
 				},
-				GitAuthProviders: []*sdkproto.GitAuthProvider{{
-					Id:          gitAuthProvider,
-					AccessToken: "access_token",
-				}},
-				Metadata: &sdkproto.Metadata{
-					CoderUrl:                      (&url.URL{}).String(),
-					WorkspaceTransition:           sdkproto.WorkspaceTransition_START,
-					WorkspaceName:                 workspace.Name,
-					WorkspaceOwner:                user.Username,
-					WorkspaceOwnerEmail:           user.Email,
-					WorkspaceOwnerOidcAccessToken: link.OAuthAccessToken,
-					WorkspaceId:                   workspace.ID.String(),
-					WorkspaceOwnerId:              user.ID.String(),
-					TemplateId:                    template.ID.String(),
-					TemplateName:                  template.Name,
-					TemplateVersion:               version.Name,
-					WorkspaceOwnerSessionToken:    sessionToken,
+			})
+			require.NoError(t, err)
+			require.JSONEq(t, string(want), string(got))
+		})
+		t.Run(tc.name+"_TemplateVersionImport", func(t *testing.T) {
+			t.Parallel()
+			srv, db, ps := setup(t, false, nil)
+			ctx := context.Background()
+
+			user := dbgen.User(t, db, database.User{})
+			file := dbgen.File(t, db, database.File{CreatedBy: user.ID})
+			_ = dbgen.ProvisionerJob(t, db, ps, database.ProvisionerJob{
+				FileID:        file.ID,
+				InitiatorID:   user.ID,
+				Provisioner:   database.ProvisionerTypeEcho,
+				StorageMethod: database.ProvisionerStorageMethodFile,
+				Type:          database.ProvisionerJobTypeTemplateVersionImport,
+			})
+
+			job, err := tc.acquire(ctx, srv)
+			require.NoError(t, err)
+
+			got, err := json.Marshal(job.Type)
+			require.NoError(t, err)
+
+			want, err := json.Marshal(&proto.AcquiredJob_TemplateImport_{
+				TemplateImport: &proto.AcquiredJob_TemplateImport{
+					Metadata: &sdkproto.Metadata{
+						CoderUrl: (&url.URL{}).String(),
+					},
 				},
-			},
+			})
+			require.NoError(t, err)
+			require.JSONEq(t, string(want), string(got))
 		})
-		require.NoError(t, err)
+		t.Run(tc.name+"_TemplateVersionImportWithUserVariable", func(t *testing.T) {
+			t.Parallel()
+			srv, db, ps := setup(t, false, nil)

-		require.JSONEq(t, string(want), string(got))
+			user := dbgen.User(t, db, database.User{})
+			version := dbgen.TemplateVersion(t, db, database.TemplateVersion{})
+			file := dbgen.File(t, db, database.File{CreatedBy: user.ID})
+			_ = dbgen.ProvisionerJob(t, db, ps, database.ProvisionerJob{
+				FileID:        file.ID,
+				InitiatorID:   user.ID,
+				Provisioner:   database.ProvisionerTypeEcho,
+				StorageMethod: database.ProvisionerStorageMethodFile,
+				Type:          database.ProvisionerJobTypeTemplateVersionImport,
+				Input: must(json.Marshal(provisionerdserver.TemplateVersionImportJob{
+					TemplateVersionID: version.ID,
+					UserVariableValues: []codersdk.VariableValue{
+						{Name: "first", Value: "first_value"},
+					},
+				})),
+			})

-		// Assert that we delete the session token whenever
-		// a stop is issued.
-		stopbuild := dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
-			WorkspaceID:       workspace.ID,
-			BuildNumber:       2,
-			JobID:             uuid.New(),
-			TemplateVersionID: version.ID,
-			Transition:        database.WorkspaceTransitionStop,
-			Reason:            database.BuildReasonInitiator,
-		})
-		_ = dbgen.ProvisionerJob(t, db, database.ProvisionerJob{
-			ID:            stopbuild.ID,
-			InitiatorID:   user.ID,
-			Provisioner:   database.ProvisionerTypeEcho,
-			StorageMethod: database.ProvisionerStorageMethodFile,
-			FileID:        file.ID,
-			Type:          database.ProvisionerJobTypeWorkspaceBuild,
-			Input: must(json.Marshal(provisionerdserver.WorkspaceProvisionJob{
-				WorkspaceBuildID: stopbuild.ID,
-			})),
-		})
+			ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
+			defer cancel()

-		stopPublished := make(chan struct{})
-		closeStopSubscribe, err := ps.Subscribe(codersdk.WorkspaceNotifyChannel(workspace.ID), func(_ context.Context, _ []byte) {
-			close(stopPublished)
-		})
-		require.NoError(t, err)
-		defer closeStopSubscribe()
+			job, err := tc.acquire(ctx, srv)
+			require.NoError(t, err)

-		// Grab jobs until we find the workspace build job. There is also
-		// an import version job that we need to ignore.
-		job, err = srv.AcquireJob(ctx, nil)
-		require.NoError(t, err)
-		_, ok := job.Type.(*proto.AcquiredJob_WorkspaceBuild_)
-		require.True(t, ok, "acquired job not a workspace build?")
+			got, err := json.Marshal(job.Type)
+			require.NoError(t, err)

-		<-stopPublished
-
-		// Validate that a session token is deleted during a stop job.
-		sessionToken = job.Type.(*proto.AcquiredJob_WorkspaceBuild_).WorkspaceBuild.Metadata.WorkspaceOwnerSessionToken
-		require.Empty(t, sessionToken)
-		_, err = db.GetAPIKeyByID(ctx, key.ID)
-		require.ErrorIs(t, err, sql.ErrNoRows)
-	})
-
-	t.Run("TemplateVersionDryRun", func(t *testing.T) {
-		t.Parallel()
-		srv, db, _ := setup(t, false, nil)
-		ctx := context.Background()
-
-		user := dbgen.User(t, db, database.User{})
-		version := dbgen.TemplateVersion(t, db, database.TemplateVersion{})
-		file := dbgen.File(t, db, database.File{CreatedBy: user.ID})
-		_ = dbgen.ProvisionerJob(t, db, database.ProvisionerJob{
-			InitiatorID:   user.ID,
-			Provisioner:   database.ProvisionerTypeEcho,
-			StorageMethod: database.ProvisionerStorageMethodFile,
-			FileID:        file.ID,
-			Type:          database.ProvisionerJobTypeTemplateVersionDryRun,
-			Input: must(json.Marshal(provisionerdserver.TemplateVersionDryRunJob{
-				TemplateVersionID: version.ID,
-				WorkspaceName:     "testing",
-			})),
-		})
-
-		job, err := srv.AcquireJob(ctx, nil)
-		require.NoError(t, err)
-
-		got, err := json.Marshal(job.Type)
-		require.NoError(t, err)
-
-		want, err := json.Marshal(&proto.AcquiredJob_TemplateDryRun_{
-			TemplateDryRun: &proto.AcquiredJob_TemplateDryRun{
-				Metadata: &sdkproto.Metadata{
-					CoderUrl:      (&url.URL{}).String(),
-					WorkspaceName: "testing",
+			want, err := json.Marshal(&proto.AcquiredJob_TemplateImport_{
+				TemplateImport: &proto.AcquiredJob_TemplateImport{
+					UserVariableValues: []*sdkproto.VariableValue{
+						{Name: "first", Sensitive: true, Value: "first_value"},
+					},
+					Metadata: &sdkproto.Metadata{
+						CoderUrl: (&url.URL{}).String(),
+					},
 				},
-			},
+			})
+			require.NoError(t, err)
+			require.JSONEq(t, string(want), string(got))
 		})
-		require.NoError(t, err)
-		require.JSONEq(t, string(want), string(got))
-	})
-	t.Run("TemplateVersionImport", func(t *testing.T) {
-		t.Parallel()
-		srv, db, _ := setup(t, false, nil)
-		ctx := context.Background()
-
-		user := dbgen.User(t, db, database.User{})
-		file := dbgen.File(t, db, database.File{CreatedBy: user.ID})
-		_ = dbgen.ProvisionerJob(t, db, database.ProvisionerJob{
-			FileID:        file.ID,
-			InitiatorID:   user.ID,
-			Provisioner:   database.ProvisionerTypeEcho,
-			StorageMethod: database.ProvisionerStorageMethodFile,
-			Type:          database.ProvisionerJobTypeTemplateVersionImport,
-		})
-
-		job, err := srv.AcquireJob(ctx, nil)
-		require.NoError(t, err)
-
-		got, err := json.Marshal(job.Type)
-		require.NoError(t, err)
-
-		want, err := json.Marshal(&proto.AcquiredJob_TemplateImport_{
-			TemplateImport: &proto.AcquiredJob_TemplateImport{
-				Metadata: &sdkproto.Metadata{
-					CoderUrl: (&url.URL{}).String(),
-				},
-			},
-		})
-		require.NoError(t, err)
-		require.JSONEq(t, string(want), string(got))
-	})
-	t.Run("TemplateVersionImportWithUserVariable", func(t *testing.T) {
-		t.Parallel()
-		srv, db, _ := setup(t, false, nil)
-
-		user := dbgen.User(t, db, database.User{})
-		version := dbgen.TemplateVersion(t, db, database.TemplateVersion{})
-		file := dbgen.File(t, db, database.File{CreatedBy: user.ID})
-		_ = dbgen.ProvisionerJob(t, db, database.ProvisionerJob{
-			FileID:        file.ID,
-			InitiatorID:   user.ID,
-			Provisioner:   database.ProvisionerTypeEcho,
-			StorageMethod: database.ProvisionerStorageMethodFile,
-			Type:          database.ProvisionerJobTypeTemplateVersionImport,
-			Input: must(json.Marshal(provisionerdserver.TemplateVersionImportJob{
-				TemplateVersionID: version.ID,
-				UserVariableValues: []codersdk.VariableValue{
-					{Name: "first", Value: "first_value"},
-				},
-			})),
-		})
-
-		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
-		defer cancel()
-
-		job, err := srv.AcquireJob(ctx, nil)
-		require.NoError(t, err)
-
-		got, err := json.Marshal(job.Type)
-		require.NoError(t, err)
-
-		want, err := json.Marshal(&proto.AcquiredJob_TemplateImport_{
-			TemplateImport: &proto.AcquiredJob_TemplateImport{
-				UserVariableValues: []*sdkproto.VariableValue{
-					{Name: "first", Sensitive: true, Value: "first_value"},
-				},
-				Metadata: &sdkproto.Metadata{
-					CoderUrl: (&url.URL{}).String(),
-				},
-			},
-		})
-		require.NoError(t, err)
-		require.JSONEq(t, string(want), string(got))
-	})
+	}
 }

 func TestUpdateJob(t *testing.T) {
@@ -1142,7 +1160,7 @@ func TestCompleteJob(t *testing.T) {
 					Transition:        c.transition,
 					Reason:            database.BuildReasonInitiator,
 				})
-				job := dbgen.ProvisionerJob(t, db, database.ProvisionerJob{
+				job := dbgen.ProvisionerJob(t, db, ps, database.ProvisionerJob{
 					FileID: file.ID,
 					Type:   database.ProvisionerJobTypeWorkspaceBuild,
 					Input: must(json.Marshal(provisionerdserver.WorkspaceProvisionJob{
@@ -1390,7 +1408,7 @@ func TestCompleteJob(t *testing.T) {
 					Transition:        c.transition,
 					Reason:            database.BuildReasonInitiator,
 				})
-				job := dbgen.ProvisionerJob(t, db, database.ProvisionerJob{
+				job := dbgen.ProvisionerJob(t, db, ps, database.ProvisionerJob{
 					FileID: file.ID,
 					Type:   database.ProvisionerJobTypeWorkspaceBuild,
 					Input: must(json.Marshal(provisionerdserver.WorkspaceProvisionJob{
@@ -1662,10 +1680,14 @@ type overrides struct {
 	templateScheduleStore       *atomic.Pointer[schedule.TemplateScheduleStore]
 	userQuietHoursScheduleStore *atomic.Pointer[schedule.UserQuietHoursScheduleStore]
 	timeNowFn                   func() time.Time
+	acquireJobLongPollDuration  time.Duration
 }

 func setup(t *testing.T, ignoreLogErrors bool, ov *overrides) (proto.DRPCProvisionerDaemonServer, database.Store, pubsub.Pubsub) {
 	t.Helper()
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+	logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug)
 	db := dbfake.New()
 	ps := pubsub.NewInMemory()
 	deploymentValues := &codersdk.DeploymentValues{}
@@ -1674,6 +1696,7 @@ func setup(t *testing.T, ignoreLogErrors bool, ov *overrides) (proto.DRPCProvisi
 	tss := testTemplateScheduleStore()
 	uqhss := testUserQuietHoursScheduleStore()
 	var timeNowFn func() time.Time
+	pollDur := time.Duration(0)
 	if ov != nil {
 		if ov.deploymentValues != nil {
 			deploymentValues = ov.deploymentValues
@@ -1705,6 +1728,7 @@ func setup(t *testing.T, ignoreLogErrors bool, ov *overrides) (proto.DRPCProvisi
 		if ov.timeNowFn != nil {
 			timeNowFn = ov.timeNowFn
 		}
+		pollDur = ov.acquireJobLongPollDuration
 	}

 	srv, err := provisionerdserver.NewServer(
@@ -1712,9 +1736,10 @@ func setup(t *testing.T, ignoreLogErrors bool, ov *overrides) (proto.DRPCProvisi
 		srvID,
 		slogtest.Make(t, &slogtest.Options{IgnoreErrors: ignoreLogErrors}),
 		[]database.ProvisionerType{database.ProvisionerTypeEcho},
-		nil,
+		provisionerdserver.Tags{},
 		db,
 		ps,
+		provisionerdserver.NewAcquirer(ctx, logger.Named("acquirer"), db, ps),
 		telemetry.NewNoop(),
 		trace.NewNoopTracerProvider().Tracer("noop"),
 		&atomic.Pointer[proto.QuotaCommitter]{},
@@ -1722,12 +1747,11 @@ func setup(t *testing.T, ignoreLogErrors bool, ov *overrides) (proto.DRPCProvisi
 		tss,
 		uqhss,
 		deploymentValues,
-		// Negative values cause the debounce to never kick in.
-		-time.Minute,
 		provisionerdserver.Options{
-			GitAuthConfigs: gitAuthConfigs,
-			TimeNowFn:      timeNowFn,
-			OIDCConfig:     &oauth2.Config{},
+			GitAuthConfigs:        gitAuthConfigs,
+			TimeNowFn:             timeNowFn,
+			OIDCConfig:            &oauth2.Config{},
+			AcquireJobLongPollDur: pollDur,
 		},
 	)
 	require.NoError(t, err)
@@ -1740,3 +1764,95 @@ func must[T any](value T, err error) T {
 	}
 	return value
 }
+
+var (
+	errUnimplemented = xerrors.New("unimplemented")
+	errClosed        = xerrors.New("closed")
+)
+
+type fakeStream struct {
+	ctx        context.Context
+	c          *sync.Cond
+	closed     bool
+	canceled   bool
+	sendCalled bool
+	job        *proto.AcquiredJob
+}
+
+func newFakeStream(ctx context.Context) *fakeStream {
+	return &fakeStream{
+		ctx: ctx,
+		c:   sync.NewCond(&sync.Mutex{}),
+	}
+}
+
+func (s *fakeStream) Send(j *proto.AcquiredJob) error {
+	s.c.L.Lock()
+	defer s.c.L.Unlock()
+	s.sendCalled = true
+	s.job = j
+	s.c.Broadcast()
+	return nil
+}
+
+func (s *fakeStream) Recv() (*proto.CancelAcquire, error) {
+	s.c.L.Lock()
+	defer s.c.L.Unlock()
+	for !(s.canceled || s.closed) {
+		s.c.Wait()
+	}
+	if s.canceled {
+		return &proto.CancelAcquire{}, nil
+	}
+	return nil, io.EOF
+}
+
+// Context returns the context associated with the stream. It is canceled
+// when the Stream is closed and no more messages will ever be sent or
+// received on it.
+func (s *fakeStream) Context() context.Context {
+	return s.ctx
+}
+
+// MsgSend sends the Message to the remote.
+func (*fakeStream) MsgSend(drpc.Message, drpc.Encoding) error {
+	return errUnimplemented
+}
+
+// MsgRecv receives a Message from the remote.
+func (*fakeStream) MsgRecv(drpc.Message, drpc.Encoding) error {
+	return errUnimplemented
+}
+
+// CloseSend signals to the remote that we will no longer send any messages.
+func (*fakeStream) CloseSend() error {
+	return errUnimplemented
+}
+
+// Close closes the stream.
+func (s *fakeStream) Close() error {
+	s.c.L.Lock()
+	defer s.c.L.Unlock()
+	s.closed = true
+	s.c.Broadcast()
+	return nil
+}
+
+func (s *fakeStream) waitForJob() (*proto.AcquiredJob, error) {
+	s.c.L.Lock()
+	defer s.c.L.Unlock()
+	for !(s.sendCalled || s.closed) {
+		s.c.Wait()
+	}
+	if s.sendCalled {
+		return s.job, nil
+	}
+	return nil, errClosed
+}
+
+func (s *fakeStream) cancel() {
+	s.c.L.Lock()
+	defer s.c.L.Unlock()
+	s.canceled = true
+	s.c.Broadcast()
+}