feat: reinitialize agents when a prebuilt workspace is claimed #17475
introduce unit testable abstraction layers
- Loading branch information
commit a22b4143e1e224accd982fdba2a7f758a43b37c4
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,102 @@ | ||
| package prebuilds | ||
|
|
||
| import ( | ||
| "context" | ||
| "fmt" | ||
| "net/http" | ||
|
|
||
| "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" | ||
| ) | ||
|
|
||
| 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 { | ||
| 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) { | ||
| select { | ||
| case <-ctx.Done(): | ||
| return | ||
| case <-inner.Done(): | ||
| return | ||
| 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 { | ||
| return func() {}, nil, xerrors.Errorf("failed to subscribe to prebuild claimed channel: %w", err) | ||
| } | ||
| defer cancelSub() | ||
| return func() { cancelSub() }, reinitEvents, 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) | ||
|
Contributor
There was a problem hiding this comment. We can fix this later, but I think this might be inappropriate at this layer. |
||
| 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)) | ||
|
SasSwart marked this conversation as resolved.
Outdated
|
||
| } | ||
| } | ||
| } | ||
| } | ||
|
|
||
| type MockClaimCoordinator interface{} | ||
|
|
||
| type ClaimListener interface{} | ||
| type PostgresClaimListener struct{} | ||
|
|
||
| type AgentReinitializer interface{} | ||
| type SSEAgentReinitializer struct{} | ||
|
|
||
| type ClaimCoordinator interface { | ||
| ClaimListener | ||
| AgentReinitializer | ||
| } | ||
|
SasSwart marked this conversation as resolved.
Outdated
|
||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.