add tests for agent reinitialization

coder · SasSwart · May 14, 2025 · Apr 21, 2025 · Apr 21, 2025 · Apr 23, 2025
commit 580420197f6a736e784dcd3c2f14360c9a716fc6
diff --git a/coderd/prebuilds/claim.go b/coderd/prebuilds/claim.go
@@ -2,23 +2,16 @@ package prebuilds
 
 import (
 	"context"
-	"net/http"
 	"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/coderd/httpapi"
-	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/codersdk/agentsdk"
 )
 
-type WorkspaceClaimPublisher interface {
-	PublishWorkspaceClaim(agentsdk.ReinitializationEvent)
-}
-
 func NewPubsubWorkspaceClaimPublisher(ps pubsub.Pubsub) *PubsubWorkspaceClaimPublisher {
 	return &PubsubWorkspaceClaimPublisher{ps: ps}
 }
@@ -35,10 +28,6 @@ func (p PubsubWorkspaceClaimPublisher) PublishWorkspaceClaim(claim agentsdk.Rein
 	return nil
 }
 
-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}
 }
@@ -49,6 +38,12 @@ type PubsubWorkspaceClaimListener struct {
 }
 
 func (p PubsubWorkspaceClaimListener) ListenForWorkspaceClaims(ctx context.Context, workspaceID uuid.UUID) (func(), <-chan agentsdk.ReinitializationEvent, error) {
+	select {
+	case <-ctx.Done():
+		return func() {}, nil, ctx.Err()
+	default:
+	}
+
 	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)
@@ -91,52 +86,3 @@ func (p PubsubWorkspaceClaimListener) ListenForWorkspaceClaims(ctx context.Conte
 
 	return cancel, workspaceClaims, nil
 }
-
-func StreamAgentReinitEvents(ctx context.Context, logger slog.Logger, rw http.ResponseWriter, r *http.Request, reinitEvents <-chan agentsdk.ReinitializationEvent) {
-	sseSendEvent, sseSenderClosed, err := httpapi.ServerSentEventSender(rw, r)
-	if err != nil {
-		httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
-			Message: "Internal error setting up server-sent events.",
-			Detail:  err.Error(),
-		})
-		return
-	}
-	// Prevent handler from returning until the sender is closed.
-	defer func() {
-		<-sseSenderClosed
-	}()
-
-	// An initial ping signals to the requester that the server is now ready
-	// and the client can begin servicing a channel with data.
-	_ = sseSendEvent(codersdk.ServerSentEvent{
-		Type: codersdk.ServerSentEventTypePing,
-	})
-
-	for {
-		select {
-		case <-ctx.Done():
-			return
-		case reinitEvent := <-reinitEvents:
-			err = sseSendEvent(codersdk.ServerSentEvent{
-				Type: codersdk.ServerSentEventTypeData,
-				Data: reinitEvent,
-			})
-			if err != nil {
-				logger.Warn(ctx, "failed to send SSE response to trigger reinit", slog.Error(err))
-			}
-		}
-	}
-}
-
-type MockClaimCoordinator interface{}
-
-type ClaimListener interface{}
-type PostgresClaimListener struct{}
-
-type AgentReinitializer interface{}
-type SSEAgentReinitializer struct{}
-
-type ClaimCoordinator interface {
-	ClaimListener
-	AgentReinitializer
-}
diff --git a/coderd/prebuilds/claim_test.go b/coderd/prebuilds/claim_test.go
@@ -79,12 +79,12 @@ func TestPubsubWorkspaceClaimListener(t *testing.T) {
 		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()
 
+		cancel()
 		// Channel should be closed immediately due to context cancellation
 		select {
 		case _, ok := <-claims:

diff --git a/coderd/provisionerdserver/provisionerdserver.go b/coderd/provisionerdserver/provisionerdserver.go
@@ -44,6 +44,7 @@ import (
 	"github.com/coder/coder/v2/coderd/tracing"
 	"github.com/coder/coder/v2/coderd/wspubsub"
 	"github.com/coder/coder/v2/codersdk"
+	"github.com/coder/coder/v2/codersdk/agentsdk"
 	"github.com/coder/coder/v2/codersdk/drpc"
 	"github.com/coder/coder/v2/provisioner"
 	"github.com/coder/coder/v2/provisionerd/proto"
@@ -1749,7 +1750,12 @@ func (s *server) CompleteJob(ctx context.Context, completed *proto.CompletedJob)
 				slog.F("user", input.PrebuildClaimedByUser.String()),
 				slog.F("workspace_id", workspace.ID))
 
-			if err := prebuilds.PublishWorkspaceClaim(ctx, s.Pubsub, workspace.ID, input.PrebuildClaimedByUser); err != nil {
+			err = prebuilds.NewPubsubWorkspaceClaimPublisher(s.Pubsub).PublishWorkspaceClaim(agentsdk.ReinitializationEvent{
+				UserID:      input.PrebuildClaimedByUser,
+				WorkspaceID: workspace.ID,
+				Reason:      agentsdk.ReinitializeReasonPrebuildClaimed,
+			})
+			if err != nil {
 				s.Logger.Error(ctx, "failed to publish workspace claim event", slog.Error(err))
 			}
 		}

diff --git a/coderd/workspaceagents.go b/coderd/workspaceagents.go
@@ -1181,14 +1181,24 @@ func (api *API) workspaceAgentReinit(rw http.ResponseWriter, r *http.Request) {
 
 	log.Info(ctx, "agent waiting for reinit instruction")
 
-	cancel, reinitEvents, err := prebuilds.ListenForWorkspaceClaims(ctx, log, api.Pubsub, workspace.ID)
+	cancel, reinitEvents, err := prebuilds.NewPubsubWorkspaceClaimListener(api.Pubsub, log).ListenForWorkspaceClaims(ctx, workspace.ID)
 	if err != nil {
 		log.Error(ctx, "failed to subscribe to prebuild claimed channel", slog.Error(err))
 		httpapi.InternalServerError(rw, xerrors.New("failed to subscribe to prebuild claimed channel"))
 		return
 	}
 	defer cancel()
-	prebuilds.StreamAgentReinitEvents(ctx, log, rw, r, reinitEvents)
+
+	transmitter := agentsdk.NewSSEAgentReinitTransmitter(log, rw, r)
+
+	err = transmitter.Transmit(ctx, reinitEvents)
+	if err != nil {
+		log.Error(ctx, "failed to stream agent reinit events", slog.Error(err))
+		httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
+			Message: "Internal error streaming agent reinitialization events.",
+			Detail:  err.Error(),
+		})
+	}
 }
 
 // convertProvisionedApps converts applications that are in the middle of provisioning process.

diff --git a/codersdk/agentsdk/agentsdk.go b/codersdk/agentsdk/agentsdk.go
@@ -24,6 +24,7 @@ import (
 
 	"github.com/coder/coder/v2/agent/proto"
 	"github.com/coder/coder/v2/apiversion"
+	"github.com/coder/coder/v2/coderd/httpapi"
 	"github.com/coder/coder/v2/codersdk"
 	drpcsdk "github.com/coder/coder/v2/codersdk/drpc"
 	tailnetproto "github.com/coder/coder/v2/tailnet/proto"
@@ -730,8 +731,86 @@ func (c *Client) WaitForReinit(ctx context.Context) (*ReinitializationEvent, err
 		return nil, codersdk.ReadBodyAsError(res)
 	}
 
-	nextEvent := codersdk.ServerSentEventReader(ctx, res.Body)
+	reinitEvent, err := NewSSEAgentReinitReceiver(res.Body).Receive(ctx)
+	if err != nil {
+		return nil, xerrors.Errorf("listening for reinitialization events: %w", err)
+	}
+	return reinitEvent, nil
+}
+
+func WaitForReinitLoop(ctx context.Context, logger slog.Logger, client *Client) <-chan ReinitializationEvent {
+	reinitEvents := make(chan ReinitializationEvent)
+
+	go func() {
+		for retrier := retry.New(100*time.Millisecond, 10*time.Second); retrier.Wait(ctx); {
+			logger.Debug(ctx, "waiting for agent reinitialization instructions")
+			reinitEvent, err := client.WaitForReinit(ctx)
+			if err != nil {
+				logger.Error(ctx, "failed to wait for agent reinitialization instructions", slog.Error(err))
+				continue
+			}
+			select {
+			case <-ctx.Done():
+				close(reinitEvents)
+				return
+			case reinitEvents <- *reinitEvent:
+			}
+		}
+	}()
+
+	return reinitEvents
+}
+
+func NewSSEAgentReinitTransmitter(logger slog.Logger, rw http.ResponseWriter, r *http.Request) *SSEAgentReinitTransmitter {
+	return &SSEAgentReinitTransmitter{logger: logger, rw: rw, r: r}
+}
+
+type SSEAgentReinitTransmitter struct {
+	rw     http.ResponseWriter
+	r      *http.Request
+	logger slog.Logger
+}
+
+func (s *SSEAgentReinitTransmitter) Transmit(ctx context.Context, reinitEvents <-chan ReinitializationEvent) error {
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	default:
+	}
+
+	sseSendEvent, sseSenderClosed, err := httpapi.ServerSentEventSender(s.rw, s.r)
+	if err != nil {
+		return xerrors.Errorf("failed to create sse transmitter: %w", err)
+	}
 
+	for {
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case <-sseSenderClosed:
+			return xerrors.New("sse connection closed")
+		case reinitEvent := <-reinitEvents:
+			err := sseSendEvent(codersdk.ServerSentEvent{
+				Type: codersdk.ServerSentEventTypeData,
+				Data: reinitEvent,
+			})
+			if err != nil {
+				s.logger.Warn(ctx, "failed to send SSE response to trigger reinit", slog.Error(err))
+			}
+		}
+	}
+}
+
+func NewSSEAgentReinitReceiver(r io.ReadCloser) *SSEAgentReinitReceiver {
+	return &SSEAgentReinitReceiver{r: r}
+}
+
+type SSEAgentReinitReceiver struct {
+	r io.ReadCloser
+}
+
+func (s *SSEAgentReinitReceiver) Receive(ctx context.Context) (*ReinitializationEvent, error) {
+	nextEvent := codersdk.ServerSentEventReader(ctx, s.r)
 	for {
 		select {
 		case <-ctx.Done():
@@ -763,26 +842,3 @@ func (c *Client) WaitForReinit(ctx context.Context) (*ReinitializationEvent, err
 		}
 	}
 }
-
-func WaitForReinitLoop(ctx context.Context, logger slog.Logger, client *Client) <-chan ReinitializationEvent {
-	reinitEvents := make(chan ReinitializationEvent)
-
-	go func() {
-		for retrier := retry.New(100*time.Millisecond, 10*time.Second); retrier.Wait(ctx); {
-			logger.Debug(ctx, "waiting for agent reinitialization instructions")
-			reinitEvent, err := client.WaitForReinit(ctx)
-			if err != nil {
-				logger.Error(ctx, "failed to wait for agent reinitialization instructions", slog.Error(err))
-				continue
-			}
-			select {
-			case <-ctx.Done():
-				close(reinitEvents)
-				return
-			case reinitEvents <- *reinitEvent:
-			}
-		}
-	}()
-
-	return reinitEvents
-}