test workspace claim pubsub

coder · SasSwart · May 14, 2025 · Apr 21, 2025 · Apr 21, 2025 · Apr 23, 2025
commit 9bbd2c75842ea47ed72637b6f2ab492be1568cf8
diff --git a/cli/agent.go b/cli/agent.go
@@ -331,6 +331,7 @@ func (r *RootCmd) workspaceAgent() *serpent.Command {
 			}
 
 			reinitEvents := agentsdk.WaitForReinitLoop(ctx, logger, client)
+
 			var (
 				lastErr  error
 				mustExit bool
@@ -379,7 +380,7 @@ func (r *RootCmd) workspaceAgent() *serpent.Command {
 					mustExit = true
 				case event := <-reinitEvents:
 					logger.Warn(ctx, "agent received instruction to reinitialize",
-						slog.F("message", event.Message), slog.F("reason", event.Reason))
+						slog.F("user_id", event.UserID), slog.F("workspace_id", event.WorkspaceID), slog.F("reason", event.Reason))
 				}
 
 				lastErr = agnt.Close()

diff --git a/coderd/prebuilds/claim.go b/coderd/prebuilds/claim.go
@@ -2,8 +2,8 @@ package prebuilds
 
 import (
 	"context"
-	"fmt"
 	"net/http"
+	"sync"
 
 	"github.com/google/uuid"
 	"golang.org/x/xerrors"
@@ -15,41 +15,81 @@ import (
 	"github.com/coder/coder/v2/codersdk/agentsdk"
 )
 
-func PublishWorkspaceClaim(ctx context.Context, ps pubsub.Pubsub, workspaceID, userID uuid.UUID) error {
-	channel := agentsdk.PrebuildClaimedChannel(workspaceID)
-	if err := ps.Publish(channel, []byte(userID.String())); err != nil {
+type WorkspaceClaimPublisher interface {
+	PublishWorkspaceClaim(agentsdk.ReinitializationEvent)
+}
+
+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.UserID.String())); err != nil {
 		return xerrors.Errorf("failed to trigger prebuilt workspace agent reinitialization: %w", err)
 	}
 	return nil
 }
 
-func ListenForWorkspaceClaims(ctx context.Context, logger slog.Logger, ps pubsub.Pubsub, workspaceID uuid.UUID) (func(), <-chan agentsdk.ReinitializationEvent, error) {
-	reinitEvents := make(chan agentsdk.ReinitializationEvent, 1)
-	cancelSub, err := ps.Subscribe(agentsdk.PrebuildClaimedChannel(workspaceID), func(inner context.Context, id []byte) {
+type WorkspaceClaimListener interface {
+	ListenForWorkspaceClaims(ctx context.Context, workspaceID uuid.UUID) (func(), <-chan agentsdk.ReinitializationEvent, error)
+}
+
+func NewPubsubWorkspaceClaimListener(ps pubsub.Pubsub, logger slog.Logger) *PubsubWorkspaceClaimListener {
+	return &PubsubWorkspaceClaimListener{ps: ps, logger: logger}
+}
+
+type PubsubWorkspaceClaimListener struct {
+	logger slog.Logger
+	ps     pubsub.Pubsub
+}
+
+func (p PubsubWorkspaceClaimListener) ListenForWorkspaceClaims(ctx context.Context, workspaceID uuid.UUID) (func(), <-chan agentsdk.ReinitializationEvent, error) {
+	workspaceClaims := make(chan agentsdk.ReinitializationEvent, 1)
+	cancelSub, err := p.ps.Subscribe(agentsdk.PrebuildClaimedChannel(workspaceID), func(inner context.Context, id []byte) {
+		claimantID, err := uuid.ParseBytes(id)
+		if err != nil {
+			p.logger.Error(ctx, "invalid prebuild claimed channel payload", slog.F("input", string(id)))
+			return
+		}
+		claim := agentsdk.ReinitializationEvent{
+			UserID:      claimantID,
+			WorkspaceID: workspaceID,
+			Reason:      agentsdk.ReinitializeReasonPrebuildClaimed,
+		}
 		select {
 		case <-ctx.Done():
 			return
 		case <-inner.Done():
 			return
+		case workspaceClaims <- claim:
 		default:
 		}
-
-		claimantID, err := uuid.ParseBytes(id)
-		if err != nil {
-			logger.Error(ctx, "invalid prebuild claimed channel payload", slog.F("input", string(id)))
-			return
-		}
-		// TODO: turn this into a <- uuid.UUID
-		reinitEvents <- agentsdk.ReinitializationEvent{
-			Message: fmt.Sprintf("prebuild claimed by user: %s", claimantID),
-			Reason:  agentsdk.ReinitializeReasonPrebuildClaimed,
-		}
 	})
+
 	if err != nil {
+		close(workspaceClaims)
 		return func() {}, nil, xerrors.Errorf("failed to subscribe to prebuild claimed channel: %w", err)
 	}
-	defer cancelSub()
-	return func() { cancelSub() }, reinitEvents, nil
+
+	var once sync.Once
+	cancel := func() {
+		once.Do(func() {
+			cancelSub()
+			close(workspaceClaims)
+		})
+	}
+
+	go func() {
+		<-ctx.Done()
+		cancel()
+	}()
+
+	return cancel, workspaceClaims, nil
 }
 
 func StreamAgentReinitEvents(ctx context.Context, logger slog.Logger, rw http.ResponseWriter, r *http.Request, reinitEvents <-chan agentsdk.ReinitializationEvent) {

diff --git a/coderd/prebuilds/claim_test.go b/coderd/prebuilds/claim_test.go
@@ -0,0 +1,191 @@
+package prebuilds_test
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"golang.org/x/xerrors"
+
+	"cdr.dev/slog/sloggers/slogtest"
+	"github.com/coder/coder/v2/coderd/database/pubsub"
+	"github.com/coder/coder/v2/coderd/prebuilds"
+	"github.com/coder/coder/v2/codersdk/agentsdk"
+	"github.com/coder/coder/v2/testutil"
+)
+
+func TestPubsubWorkspaceClaimPublisher(t *testing.T) {
+	t.Parallel()
+	t.Run("publish claim", func(t *testing.T) {
+		t.Parallel()
+
+		ps := pubsub.NewInMemory()
+		publisher := prebuilds.NewPubsubWorkspaceClaimPublisher(ps)
+
+		workspaceID := uuid.New()
+		userID := uuid.New()
+
+		userIDCh := make(chan uuid.UUID, 1)
+		channel := agentsdk.PrebuildClaimedChannel(workspaceID)
+		cancel, err := ps.Subscribe(channel, func(ctx context.Context, message []byte) {
+			userIDCh <- uuid.MustParse(string(message))
+		})
+		require.NoError(t, err)
+		defer cancel()
+
+		claim := agentsdk.ReinitializationEvent{
+			UserID:      userID,
+			WorkspaceID: workspaceID,
+			Reason:      agentsdk.ReinitializeReasonPrebuildClaimed,
+		}
+		err = publisher.PublishWorkspaceClaim(claim)
+		require.NoError(t, err)
+
+		// Verify the message was published
+		select {
+		case gotUserID := <-userIDCh:
+			require.Equal(t, userID, gotUserID)
+		case <-time.After(testutil.WaitShort):
+			t.Fatal("timeout waiting for claim")
+		}
+	})
+
+	t.Run("fail to publish claim", func(t *testing.T) {
+		t.Parallel()
+
+		ps := &brokenPubsub{}
+
+		publisher := prebuilds.NewPubsubWorkspaceClaimPublisher(ps)
+		claim := agentsdk.ReinitializationEvent{
+			UserID:      uuid.New(),
+			WorkspaceID: uuid.New(),
+			Reason:      agentsdk.ReinitializeReasonPrebuildClaimed,
+		}
+
+		err := publisher.PublishWorkspaceClaim(claim)
+		require.Error(t, err)
+	})
+}
+
+func TestPubsubWorkspaceClaimListener(t *testing.T) {
+	t.Parallel()
+	t.Run("stops listening if context canceled", func(t *testing.T) {
+		t.Parallel()
+
+		ps := pubsub.NewInMemory()
+		listener := prebuilds.NewPubsubWorkspaceClaimListener(ps, slogtest.Make(t, nil))
+
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+
+		cancelFunc, claims, err := listener.ListenForWorkspaceClaims(ctx, uuid.New())
+		require.NoError(t, err)
+		defer cancelFunc()
+
+		// Channel should be closed immediately due to context cancellation
+		select {
+		case _, ok := <-claims:
+			assert.False(t, ok)
+		case <-time.After(testutil.WaitShort):
+			t.Fatal("timeout waiting for closed channel")
+		}
+	})
+
+	t.Run("stops listening if cancel func is called", func(t *testing.T) {
+		t.Parallel()
+
+		ps := pubsub.NewInMemory()
+		listener := prebuilds.NewPubsubWorkspaceClaimListener(ps, slogtest.Make(t, nil))
+
+		cancelFunc, claims, err := listener.ListenForWorkspaceClaims(context.Background(), uuid.New())
+		require.NoError(t, err)
+
+		cancelFunc()
+		select {
+		case _, ok := <-claims:
+			assert.False(t, ok)
+		case <-time.After(testutil.WaitShort):
+			t.Fatal("timeout waiting for closed channel")
+		}
+	})
+
+	t.Run("finds claim events for its workspace", func(t *testing.T) {
+		t.Parallel()
+
+		ps := pubsub.NewInMemory()
+		listener := prebuilds.NewPubsubWorkspaceClaimListener(ps, slogtest.Make(t, nil))
+
+		workspaceID := uuid.New()
+		userID := uuid.New()
+		cancelFunc, claims, err := listener.ListenForWorkspaceClaims(context.Background(), workspaceID)
+		require.NoError(t, err)
+		defer cancelFunc()
+
+		// Publish a claim
+		channel := agentsdk.PrebuildClaimedChannel(workspaceID)
+		err = ps.Publish(channel, []byte(userID.String()))
+		require.NoError(t, err)
+
+		// Verify we receive the claim
+		select {
+		case claim := <-claims:
+			assert.Equal(t, userID, claim.UserID)
+			assert.Equal(t, workspaceID, claim.WorkspaceID)
+			assert.Equal(t, agentsdk.ReinitializeReasonPrebuildClaimed, claim.Reason)
+		case <-time.After(time.Second):
+			t.Fatal("timeout waiting for claim")
+		}
+	})
+
+	t.Run("ignores claim events for other workspaces", func(t *testing.T) {
+		t.Parallel()
+
+		ps := pubsub.NewInMemory()
+		listener := prebuilds.NewPubsubWorkspaceClaimListener(ps, slogtest.Make(t, nil))
+
+		workspaceID := uuid.New()
+		otherWorkspaceID := uuid.New()
+		cancelFunc, claims, err := listener.ListenForWorkspaceClaims(context.Background(), workspaceID)
+		require.NoError(t, err)
+		defer cancelFunc()
+
+		// Publish a claim for a different workspace
+		channel := agentsdk.PrebuildClaimedChannel(otherWorkspaceID)
+		err = ps.Publish(channel, []byte(uuid.New().String()))
+		require.NoError(t, err)
+
+		// Verify we don't receive the claim
+		select {
+		case <-claims:
+			t.Fatal("received claim for wrong workspace")
+		case <-time.After(100 * time.Millisecond):
+			// Expected - no claim received
+		}
+	})
+
+	t.Run("communicates the error if it can't subscribe", func(t *testing.T) {
+		t.Parallel()
+
+		ps := &brokenPubsub{}
+		listener := prebuilds.NewPubsubWorkspaceClaimListener(ps, slogtest.Make(t, nil))
+
+		_, _, err := listener.ListenForWorkspaceClaims(context.Background(), uuid.New())
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), "failed to subscribe to prebuild claimed channel")
+	})
+}
+
+type brokenPubsub struct {
+	pubsub.Pubsub
+}
+
+func (brokenPubsub) Subscribe(_ string, _ pubsub.Listener) (func(), error) {
+	return nil, xerrors.New("broken")
+}
+
+func (brokenPubsub) Publish(_ string, _ []byte) error {
+	return xerrors.New("broken")
+}
diff --git a/codersdk/agentsdk/agentsdk.go b/codersdk/agentsdk/agentsdk.go
@@ -696,8 +696,9 @@ const (
 )
 
 type ReinitializationEvent struct {
-	Message string                 `json:"message"`
-	Reason  ReinitializationReason `json:"reason"`
+	WorkspaceID uuid.UUID
+	UserID      uuid.UUID
+	Reason      ReinitializationReason `json:"reason"`
 }
 
 func PrebuildClaimedChannel(id uuid.UUID) string {
@@ -707,7 +708,6 @@ func PrebuildClaimedChannel(id uuid.UUID) string {
 // WaitForReinit polls a SSE endpoint, and receives an event back under the following conditions:
 // - ping: ignored, keepalive
 // - prebuild claimed: a prebuilt workspace is claimed, so the agent must reinitialize.
-// NOTE: the caller is responsible for closing the events chan.
 func (c *Client) WaitForReinit(ctx context.Context) (*ReinitializationEvent, error) {
 	// TODO: allow configuring httpclient
 	c.SDK.HTTPClient.Timeout = time.Hour * 24
@@ -733,9 +733,6 @@ func (c *Client) WaitForReinit(ctx context.Context) (*ReinitializationEvent, err
 	nextEvent := codersdk.ServerSentEventReader(ctx, res.Body)
 
 	for {
-		// TODO (Sasswart): I don't like that we do this select at the start and at the end.
-		// nextEvent should return an error if the context is canceled, but that feels like a larger refactor.
-		// if it did, we'd only have the select at the end of the loop.
 		select {
 		case <-ctx.Done():
 			return nil, ctx.Err()