WIP: claim triggering manifest push to agent

Signed-off-by: Danny Kopping <danny@coder.com>

coderd/agentapi/api.go

@@ -97,6 +97,7 @@ func New(opts Options) *API {
 		DerpMapFn:                opts.DerpMapFn,
 		WorkspaceID:              opts.WorkspaceID,
 		Log:                      opts.Log.Named("manifests"),
+		Pubsub:                   opts.Pubsub,
 	}
 
 	api.AnnouncementBannerAPI = &AnnouncementBannerAPI{

coderd/agentapi/manifest.go

@@ -4,6 +4,8 @@ import (
 	"cdr.dev/slog"
 	"context"
 	"database/sql"
+	"fmt"
+	"github.com/coder/coder/v2/coderd/database/pubsub"
 	"net/url"
 	"strings"
 	"time"
@@ -35,6 +37,7 @@ type ManifestAPI struct {
 	AgentFn   func(context.Context) (database.WorkspaceAgent, error)
 	Database  database.Store
 	DerpMapFn func() *tailcfg.DERPMap
+	Pubsub    pubsub.Pubsub
 	Log       slog.Logger
 }
 
@@ -47,21 +50,52 @@ func (a *ManifestAPI) StreamManifests(in *agentproto.GetManifestRequest, stream
 		}
 	}()
 
+	updates := make(chan struct{}, 1)
+
+	unsub, err := a.Pubsub.Subscribe(ManifestUpdateChannel(a.WorkspaceID), func(ctx context.Context, _ []byte) {
+		a.Log.Info(ctx, "received 'prebuild claimed' event for workspace, pushing down new manifest", slog.F("workspace_id", a.WorkspaceID.String()))
+		select {
+		case <-streamCtx.Done():
+			return
+		case <-ctx.Done():
+			return
+		case updates <- struct{}{}:
+		}
+	})
+	if err != nil {
+		return xerrors.Errorf("subscribe to 'prebuild claimed' event: %w", err)
+	}
+	defer unsub()
+
 	for {
-		resp, err := a.GetManifest(streamCtx, in)
+		manifest, err := a.GetManifest(streamCtx, in)
 		if err != nil {
 			return xerrors.Errorf("receive manifest: %w", err)
 		}
 
-		err = stream.Send(resp)
+		a.Log.Debug(streamCtx, "pushing manifest to workspace", slog.F("workspace_id", a.WorkspaceID))
+
+		// Send first retrieved manifest.
+		err = stream.Send(manifest)
 		if err != nil {
 			return xerrors.Errorf("send manifest: %w", err)
 		}
 
-		time.Sleep(time.Second * 5)
+		// ...then wait until triggered by prebuild claim completion.
+		// At this stage, a prebuild will have been claimed by a user and the agent will need to be reconfigured.
+		select {
+		case <-updates:
+			a.Log.Info(streamCtx, "received manifest update request", slog.F("workspace_id", a.WorkspaceID))
+		case <-streamCtx.Done():
+			return xerrors.Errorf("stream close: %w", streamCtx.Err())
+		}
 	}
 }
 
+func ManifestUpdateChannel(id uuid.UUID) string {
+	return fmt.Sprintf("prebuild_claimed_%s", id)
+}
+
 func (a *ManifestAPI) GetManifest(ctx context.Context, _ *agentproto.GetManifestRequest) (*agentproto.Manifest, error) {
 	workspaceAgent, err := a.AgentFn(ctx)
 	if err != nil {

coderd/provisionerdserver/provisionerdserver.go

@@ -6,6 +6,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"github.com/coder/coder/v2/coderd/agentapi"
 	"net/http"
 	"net/url"
 	"reflect"
@@ -1707,6 +1708,17 @@ func (s *server) CompleteJob(ctx context.Context, completed *proto.CompletedJob)
 		if err != nil {
 			return nil, xerrors.Errorf("update workspace: %w", err)
 		}
+
+		// If this job was initiated by the prebuilds user and the job is not a prebuild, then it MUST be the claim run.
+		// TODO: maybe add some specific metadata to indicate this rather than imputing it.
+		if input.IsPrebuildClaimByUser != uuid.Nil {
+			s.Logger.Info(ctx, "workspace prebuild successfully claimed by user",
+				slog.F("user", input.IsPrebuildClaimByUser.String()),
+				slog.F("workspace_id", workspace.ID))
+			if err := s.Pubsub.Publish(agentapi.ManifestUpdateChannel(workspace.ID), nil); err != nil {
+				s.Logger.Error(ctx, "failed to publish message to workspace agent to pull new manifest", slog.Error(err))
+			}
+		}
 	case *proto.CompletedJob_TemplateDryRun_:
 		for _, resource := range jobType.TemplateDryRun.Resources {
 			s.Logger.Info(ctx, "inserting template dry-run job resource",
@@ -2356,9 +2368,10 @@ type TemplateVersionImportJob struct {
 
 // WorkspaceProvisionJob is the payload for the "workspace_provision" job type.
 type WorkspaceProvisionJob struct {
-	WorkspaceBuildID uuid.UUID `json:"workspace_build_id"`
-	DryRun           bool      `json:"dry_run"`
-	IsPrebuild       bool      `json:"is_prebuild,omitempty"`
+	WorkspaceBuildID      uuid.UUID `json:"workspace_build_id"`
+	DryRun                bool      `json:"dry_run"`
+	IsPrebuild            bool      `json:"is_prebuild,omitempty"`
+	IsPrebuildClaimByUser uuid.UUID `json:"is_prebuild_claim_by,omitempty"`
 	// RunningWorkspaceAgentID is *only* used for prebuilds. We pass it down when we want to rebuild a prebuilt workspace
 	// but not generate a new agent token. The provisionerdserver will retrieve this token and push it down to
 	// the provisioner (and ultimately to the `coder_agent` resource in the Terraform provider) where it will be

coderd/workspaces.go

@@ -685,6 +685,7 @@ func createWorkspace(
 
 		if claimedWorkspace != nil {
 			workspaceID = claimedWorkspace.ID
+			initiatorID = prebuilds.PrebuildOwnerUUID
 		} else {
 			// Workspaces are created without any versions.
 			minimumWorkspace, err := db.InsertWorkspace(ctx, database.InsertWorkspaceParams{
@@ -727,6 +728,10 @@ func createWorkspace(
 			builder = builder.VersionID(req.TemplateVersionID)
 		}
 
+		if claimedWorkspace != nil {
+			builder = builder.MarkPrebuildClaimBy(owner.ID)
+		}
+
 		workspaceBuild, provisionerJob, provisionerDaemons, err = builder.Build(
 			ctx,
 			db,

coderd/wsbuilder/wsbuilder.go

@@ -72,8 +72,10 @@ type Builder struct {
 	lastBuildJob                 *database.ProvisionerJob
 	parameterNames               *[]string
 	parameterValues              *[]string
-	prebuild                     bool
-	runningWorkspaceAgentID      uuid.UUID
+
+	prebuild                bool
+	prebuildClaimBy         uuid.UUID
+	runningWorkspaceAgentID uuid.UUID
 
 	verifyNoLegacyParametersOnce bool
 }
@@ -176,6 +178,12 @@ func (b Builder) MarkPrebuild() Builder {
 	return b
 }
 
+func (b Builder) MarkPrebuildClaimBy(userID uuid.UUID) Builder {
+	// nolint: revive
+	b.prebuildClaimBy = userID
+	return b
+}
+
 // RunningWorkspaceAgentID is only used for prebuilds; see the associated field in `provisionerdserver.WorkspaceProvisionJob`.
 func (b Builder) RunningWorkspaceAgentID(id uuid.UUID) Builder {
 	// nolint: revive
@@ -311,6 +319,7 @@ func (b *Builder) buildTx(authFunc func(action policy.Action, object rbac.Object
 		WorkspaceBuildID:        workspaceBuildID,
 		LogLevel:                b.logLevel,
 		IsPrebuild:              b.prebuild,
+		IsPrebuildClaimByUser:   b.prebuildClaimBy,
 		RunningWorkspaceAgentID: b.runningWorkspaceAgentID,
 	})
 	if err != nil {