May 14, 2025 · Apr 21, 2025 · Apr 21, 2025 · Apr 23, 2025 · Apr 24, 2025 · Apr 24, 2025
diff --git a/coderd/prebuilds/claim.go b/coderd/prebuilds/claim.go

 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}
 }
 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}
 }
 }

 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)

 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
 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
 "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"
 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

 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

 "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"
 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():
 }
 }
 }

 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
 }
Original file line number	Diff line number	Diff line change
Expand Up		@@ -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}
		}
Expand All		@@ -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}
		}
Expand All		@@ -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)
SasSwart marked this conversation as resolved. Show resolvedHide resolved
		cancelSub, err := p.ps.Subscribe(agentsdk.PrebuildClaimedChannel(workspaceID), func(inner context.Context, id []byte) {
Copy link Contributor spikecurtisMay 2, 2025 Choose a reason for hiding this comment The reason will be displayed to describe this comment to others.Learn more. As we talked about on a call --- the pubsub is not considered reliable transport, and we can miss events. This can lead to a situation where the agent misses the reinit signal and never reinitializes, even though it has been claimed. The deep problem here is that we are using the PubSub to send material information (who the new owner is), rather than just a kick that there is new information available (the workspace has a new owner). In the latter case, when there is an error, we can recover by re-querying the database to find the owner, and then decide whether we need to signal the agent with new information. This requires the handler to keep track of the last owner it sent, but that's a trivial amount of memory to keep. I don't necessarily think this needs to be fixed in this PR, since the plan is to move to a new "stream of manifest" architecture, but as we implement that, we need to keep error handling in mind on both sides: coderd recovers from a pubsub error by querying the database (or closing the connection if the database query fails), and then deciding whether there is something new to send. The agent recovers from a dropped connection by redialing, and then checking the new manifest against it's existing one to see if it needs to take any action. Copy link ContributorAuthor SasSwartMay 12, 2025 Choose a reason for hiding this comment The reason will be displayed to describe this comment to others.Learn more. Thanks for identifying this. Let's defer it beyond this PR to the next release if there are no objections. I'd like to get to this as part of the manifest streaming work.
		claimantID, err := uuid.ParseBytes(id)
Expand DownExpand Up		@@ -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
		}
Original file line number	Diff line number	Diff line change
Expand Up		@@ -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:
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -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"
Expand DownExpand Up		@@ -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,
SasSwart marked this conversation as resolved. Show resolvedHide resolved
		})
		if err != nil {
		s.Logger.Error(ctx, "failed to publish workspace claim event", slog.Error(err))
		}
		}
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -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)
SasSwart marked this conversation as resolved. Show resolvedHide resolved
		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.
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -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"
Expand DownExpand Up		@@ -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(100time.Millisecond, 10time.Second); retrier.Wait(ctx); {
		logger.Debug(ctx, "waiting for agent reinitialization instructions")
		reinitEvent, err := client.WaitForReinit(ctx)
SasSwart marked this conversation as resolved. Show resolvedHide resolved
		if err != nil {
		logger.Error(ctx, "failed to wait for agent reinitialization instructions", slog.Error(err))
		continue
		}
SasSwart marked this conversation as resolved. Show resolvedHide resolved
		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")
SasSwart marked this conversation as resolved. Show resolvedHide resolved
		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():
Expand DownExpand Up		@@ -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(100time.Millisecond, 10time.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
		}