Skip to content

Commit 2d2f1a3

Browse files
deansheatherdeansheather
committed
deregister
1 parent dcf072e commit 2d2f1a3

File tree

4 files changed

+137
-14
lines changed

4 files changed

+137
-14
lines changed

‎enterprise/coderd/coderd.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,7 @@ func New(ctx context.Context, options *Options) (*API, error) {
119119
)
120120
r.Post("/issue-signed-app-token", api.workspaceProxyIssueSignedAppToken)
121121
r.Post("/register", api.workspaceProxyRegister)
122+
r.Post("/deregister", api.workspaceProxyDeregister)
122123
})
123124
r.Route("/{workspaceproxy}", func(r chi.Router) {
124125
r.Use(

‎enterprise/coderd/workspaceproxy.go

Lines changed: 84 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@ import (
3434
// This is useful when a proxy is created or deleted. Errors will be logged.
3535
func (api *API) forceWorkspaceProxyHealthUpdate(ctx context.Context) {
3636
if err := api.ProxyHealth.ForceUpdate(ctx); err != nil {
37-
api.Logger.Error(ctx, "force proxy health update", slog.Error(err))
37+
api.Logger.Warn(ctx, "force proxy health update", slog.Error(err))
3838
}
3939
}
4040

@@ -316,17 +316,17 @@ func (api *API) workspaceProxyIssueSignedAppToken(rw http.ResponseWriter, r *htt
316316
// in the database and returns a signed token that can be used to authenticate
317317
// tokens.
318318
//
319-
// This is called periodically by the proxy in the background (once per minute
320-
// per replica) to ensure that the proxy is still registered and the
321-
// corresponding replica table entry is refreshed.
319+
// This is called periodically by the proxy in the background (every 30s per
320+
// replica) to ensure that the proxy is still registered and the corresponding
321+
// replica table entry is refreshed.
322322
//
323323
// @Summary Register workspace proxy
324324
// @ID register-workspace-proxy
325325
// @Security CoderSessionToken
326326
// @Accept json
327327
// @Produce json
328328
// @Tags Enterprise
329-
// @Param request body wsproxysdk.RegisterWorkspaceProxyRequest true "Issue signed app token request"
329+
// @Param request body wsproxysdk.RegisterWorkspaceProxyRequest true "Register workspace proxy request"
330330
// @Success 201 {object} wsproxysdk.RegisterWorkspaceProxyResponse
331331
// @Router /workspaceproxies/me/register [post]
332332
// @x-apidocgen {"skip": true}
@@ -367,6 +367,13 @@ func (api *API) workspaceProxyRegister(rw http.ResponseWriter, r *http.Request)
367367
}
368368
}
369369

370+
if req.ReplicaID == uuid.Nil {
371+
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
372+
Message: "Replica ID is invalid.",
373+
})
374+
return
375+
}
376+
370377
// TODO: get region ID
371378
var regionID int32 = 1234
372379

@@ -472,6 +479,78 @@ func (api *API) workspaceProxyRegister(rw http.ResponseWriter, r *http.Request)
472479
go api.forceWorkspaceProxyHealthUpdate(api.ctx)
473480
}
474481

482+
// @Summary Deregister workspace proxy
483+
// @ID deregister-workspace-proxy
484+
// @Security CoderSessionToken
485+
// @Accept json
486+
// @Tags Enterprise
487+
// @Param request body wsproxysdk.DeregisterWorkspaceProxyRequest true "Deregister workspace proxy request"
488+
// @Success 204
489+
// @Router /workspaceproxies/me/deregister [post]
490+
// @x-apidocgen {"skip": true}
491+
func (api *API) workspaceProxyDeregister(rw http.ResponseWriter, r *http.Request) {
492+
ctx := r.Context()
493+
494+
var req wsproxysdk.DeregisterWorkspaceProxyRequest
495+
if !httpapi.Read(ctx, rw, r, &req) {
496+
return
497+
}
498+
499+
err := api.Database.InTx(func(db database.Store) error {
500+
now := time.Now()
501+
replica, err := db.GetReplicaByID(ctx, req.ReplicaID)
502+
if err != nil {
503+
return xerrors.Errorf("get replica: %w", err)
504+
}
505+
506+
if replica.StoppedAt.Valid && !replica.StartedAt.IsZero() {
507+
// TODO: sadly this results in 500 when it should be 400
508+
return xerrors.Errorf("replica %s is already marked stopped", replica.ID)
509+
}
510+
511+
replica, err = db.UpdateReplica(ctx, database.UpdateReplicaParams{
512+
ID: replica.ID,
513+
UpdatedAt: now,
514+
StartedAt: replica.StartedAt,
515+
StoppedAt: sql.NullTime{
516+
Valid: true,
517+
Time: now,
518+
},
519+
RelayAddress: replica.RelayAddress,
520+
RegionID: replica.RegionID,
521+
Hostname: replica.Hostname,
522+
Version: replica.Version,
523+
Error: replica.Error,
524+
DatabaseLatency: replica.DatabaseLatency,
525+
Primary: replica.Primary,
526+
})
527+
if err != nil {
528+
return xerrors.Errorf("update replica: %w", err)
529+
}
530+
531+
return nil
532+
}, nil)
533+
if httpapi.Is404Error(err) {
534+
httpapi.ResourceNotFound(rw)
535+
return
536+
}
537+
if err != nil {
538+
httpapi.InternalServerError(rw, err)
539+
return
540+
}
541+
542+
// Publish a replicasync event with a nil ID so every replica (yes, even the
543+
// current replica) will refresh its replicas list.
544+
err = api.Pubsub.Publish(replicasync.PubsubEvent, []byte(uuid.Nil.String()))
545+
if err != nil {
546+
httpapi.InternalServerError(rw, err)
547+
return
548+
}
549+
550+
rw.WriteHeader(http.StatusNoContent)
551+
go api.forceWorkspaceProxyHealthUpdate(api.ctx)
552+
}
553+
475554
// reconnectingPTYSignedToken issues a signed app token for use when connecting
476555
// to the reconnecting PTY websocket on an external workspace proxy. This is set
477556
// by the client as a query parameter when connecting.

‎enterprise/wsproxy/wsproxy.go

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -324,7 +324,7 @@ func (s *Server) Close() error {
324324
s.cancel()
325325

326326
var err error
327-
registerDoneWaitTicker := time.NewTicker(3 * time.Second)
327+
registerDoneWaitTicker := time.NewTicker(11 * time.Second) // the attempt timeout is 10s
328328
select {
329329
case <-registerDoneWaitTicker.C:
330330
err = multierror.Append(err, xerrors.New("timed out waiting for registerDone"))
@@ -335,6 +335,7 @@ func (s *Server) Close() error {
335335
if appServerErr != nil {
336336
err = multierror.Append(err, appServerErr)
337337
}
338+
s.SDKClient.SDKClient.HTTPClient.CloseIdleConnections()
338339
return err
339340
}
340341

‎enterprise/wsproxy/wsproxysdk/wsproxysdk.go

Lines changed: 50 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -202,6 +202,29 @@ func (c *Client) RegisterWorkspaceProxy(ctx context.Context, req RegisterWorkspa
202202
return resp, json.NewDecoder(res.Body).Decode(&resp)
203203
}
204204

205+
type DeregisterWorkspaceProxyRequest struct {
206+
// ReplicaID is a unique identifier for the replica of the proxy that is
207+
// deregistering. It should be generated by the client on startup and
208+
// should've already been passed to the register endpoint.
209+
ReplicaID uuid.UUID `json:"replica_id"`
210+
}
211+
212+
func (c *Client) DeregisterWorkspaceProxy(ctx context.Context, req DeregisterWorkspaceProxyRequest) error {
213+
res, err := c.Request(ctx, http.MethodPost,
214+
"/api/v2/workspaceproxies/me/deregister",
215+
req,
216+
)
217+
if err != nil {
218+
return xerrors.Errorf("make request: %w", err)
219+
}
220+
defer res.Body.Close()
221+
222+
if res.StatusCode != http.StatusNoContent {
223+
return codersdk.ReadBodyAsError(res)
224+
}
225+
return nil
226+
}
227+
205228
type RegisterWorkspaceProxyLoopOpts struct {
206229
Logger slog.Logger
207230
Request RegisterWorkspaceProxyRequest
@@ -240,7 +263,9 @@ type RegisterWorkspaceProxyLoopOpts struct {
240263
// stop immediately and the context error will be returned to the FailureFn.
241264
//
242265
// The returned channel will be closed when the loop stops and can be used to
243-
// ensure the loop is dead before continuing.
266+
// ensure the loop is dead before continuing. When a fatal error is encountered,
267+
// the proxy will be deregistered (with the same ReplicaID and AttemptTimeout)
268+
// before calling the FailureFn.
244269
func (c *Client) RegisterWorkspaceProxyLoop(ctx context.Context, opts RegisterWorkspaceProxyLoopOpts) (RegisterWorkspaceProxyResponse, <-chan struct{}, error) {
245270
if opts.Interval == 0 {
246271
opts.Interval = 30 * time.Second
@@ -259,8 +284,25 @@ func (c *Client) RegisterWorkspaceProxyLoop(ctx context.Context, opts RegisterWo
259284
return nil
260285
}
261286
}
262-
if opts.FailureFn == nil {
263-
opts.FailureFn = func(_ error) {}
287+
288+
failureFn := func(err error) {
289+
// We have to use background context here because the original context
290+
// may be canceled.
291+
deregisterCtx, cancel := context.WithTimeout(context.Background(), opts.AttemptTimeout)
292+
defer cancel()
293+
deregisterErr := c.DeregisterWorkspaceProxy(deregisterCtx, DeregisterWorkspaceProxyRequest{
294+
ReplicaID: opts.Request.ReplicaID,
295+
})
296+
if deregisterErr != nil {
297+
opts.Logger.Error(ctx,
298+
"failed to deregister workspace proxy with Coder primary (it will be automatically deregistered shortly)",
299+
slog.F("err", deregisterErr),
300+
)
301+
}
302+
303+
if opts.FailureFn != nil {
304+
opts.FailureFn(err)
305+
}
264306
}
265307

266308
originalRes, err := c.RegisterWorkspaceProxy(ctx, opts.Request)
@@ -279,7 +321,7 @@ func (c *Client) RegisterWorkspaceProxyLoop(ctx context.Context, opts RegisterWo
279321
for {
280322
select {
281323
case <-ctx.Done():
282-
opts.FailureFn(ctx.Err())
324+
failureFn(ctx.Err())
283325
return
284326
case <-ticker.C:
285327
}
@@ -305,24 +347,24 @@ func (c *Client) RegisterWorkspaceProxyLoop(ctx context.Context, opts RegisterWo
305347
)
306348

307349
if failedAttempts > opts.MaxFailureCount {
308-
opts.FailureFn(xerrors.Errorf("exceeded re-registration failure count of %d: last error: %w", opts.MaxFailureCount, err))
350+
failureFn(xerrors.Errorf("exceeded re-registration failure count of %d: last error: %w", opts.MaxFailureCount, err))
309351
return
310352
}
311353
}
312354
failedAttempts = 0
313355

314356
if res.AppSecurityKey != originalRes.AppSecurityKey {
315-
opts.FailureFn(xerrors.New("app security key has changed, proxy must be restarted"))
357+
failureFn(xerrors.New("app security key has changed, proxy must be restarted"))
316358
return
317359
}
318360
if res.DERPMeshKey != originalRes.DERPMeshKey {
319-
opts.FailureFn(xerrors.New("DERP mesh key has changed, proxy must be restarted"))
361+
failureFn(xerrors.New("DERP mesh key has changed, proxy must be restarted"))
320362
return
321363
}
322364

323365
err = opts.CallbackFn(ctx, res)
324366
if err != nil {
325-
opts.FailureFn(xerrors.Errorf("callback fn returned error: %w", err))
367+
failureFn(xerrors.Errorf("callback fn returned error: %w", err))
326368
return
327369
}
328370

0 commit comments

Comments
 (0)