feat: reinitialize agents when a prebuilt workspace is claimed (#17475)

This pull request allows coder workspace agents to be reinitialized when
a prebuilt workspace is claimed by a user. This facilitates the transfer
of ownership between the anonymous prebuilds system user and the new
owner of the workspace.

Only a single agent per prebuilt workspace is supported for now, but
plumbing has already been done to facilitate the seamless transition to
multi-agent support.

---------

Signed-off-by: Danny Kopping <dannykopping@gmail.com>
Co-authored-by: Danny Kopping <dannykopping@gmail.com>

coderd/prebuilds/claim.go

@@ -0,0 +1,82 @@
+package prebuilds
+
+import (
+	"context"
+	"sync"
+
+	"github.com/google/uuid"
+	"golang.org/x/xerrors"
+
+	"cdr.dev/slog"
+	"github.com/coder/coder/v2/coderd/database/pubsub"
+	"github.com/coder/coder/v2/codersdk/agentsdk"
+)
+
+func NewPubsubWorkspaceClaimPublisher(ps pubsub.Pubsub) *PubsubWorkspaceClaimPublisher {
+	return &PubsubWorkspaceClaimPublisher{ps: ps}
+}
+
+type PubsubWorkspaceClaimPublisher struct {
+	ps pubsub.Pubsub
+}
+
+func (p PubsubWorkspaceClaimPublisher) PublishWorkspaceClaim(claim agentsdk.ReinitializationEvent) error {
+	channel := agentsdk.PrebuildClaimedChannel(claim.WorkspaceID)
+	if err := p.ps.Publish(channel, []byte(claim.Reason)); err != nil {
+		return xerrors.Errorf("failed to trigger prebuilt workspace agent reinitialization: %w", err)
+	}
+	return nil
+}
+
+func NewPubsubWorkspaceClaimListener(ps pubsub.Pubsub, logger slog.Logger) *PubsubWorkspaceClaimListener {
+	return &PubsubWorkspaceClaimListener{ps: ps, logger: logger}
+}
+
+type PubsubWorkspaceClaimListener struct {
+	logger slog.Logger
+	ps     pubsub.Pubsub
+}
+
+// ListenForWorkspaceClaims subscribes to a pubsub channel and sends any received events on the chan that it returns.
+// pubsub.Pubsub does not communicate when its last callback has been called after it has been closed. As such the chan
+// returned by this method is never closed. Call the returned cancel() function to close the subscription when it is no longer needed.
+// cancel() will be called if ctx expires or is canceled.
+func (p PubsubWorkspaceClaimListener) ListenForWorkspaceClaims(ctx context.Context, workspaceID uuid.UUID, reinitEvents chan<- agentsdk.ReinitializationEvent) (func(), error) {
+	select {
+	case <-ctx.Done():
+		return func() {}, ctx.Err()
+	default:
+	}
+
+	cancelSub, err := p.ps.Subscribe(agentsdk.PrebuildClaimedChannel(workspaceID), func(inner context.Context, reason []byte) {
+		claim := agentsdk.ReinitializationEvent{
+			WorkspaceID: workspaceID,
+			Reason:      agentsdk.ReinitializationReason(reason),
+		}
+
+		select {
+		case <-ctx.Done():
+			return
+		case <-inner.Done():
+			return
+		case reinitEvents <- claim:
+		}
+	})
+	if err != nil {
+		return func() {}, xerrors.Errorf("failed to subscribe to prebuild claimed channel: %w", err)
+	}
+
+	var once sync.Once
+	cancel := func() {
+		once.Do(func() {
+			cancelSub()
+		})
+	}
+
+	go func() {
+		<-ctx.Done()
+		cancel()
+	}()
+
+	return cancel, nil
+}