chore: Implement workspace proxy going away (graceful shutdown) (#7459)

* chore: Implement workspace proxy going away

When a workspace proxy shuts down, the health status of that
proxy should immediately be updated. This is purely a courtesy
and technically not required

coderd/apidoc/docs.go

@@ -5093,6 +5093,34 @@ const docTemplate = `{
                 }
             }
         },
+        "/workspaceproxies/me/goingaway": {
+            "post": {
+                "security": [
+                    {
+                        "CoderSessionToken": []
+                    }
+                ],
+                "produces": [
+                    "application/json"
+                ],
+                "tags": [
+                    "Enterprise"
+                ],
+                "summary": "Workspace proxy going away",
+                "operationId": "workspace-proxy-going-away",
+                "responses": {
+                    "201": {
+                        "description": "Created",
+                        "schema": {
+                            "$ref": "#/definitions/codersdk.Response"
+                        }
+                    }
+                },
+                "x-apidocgen": {
+                    "skip": true
+                }
+            }
+        },
         "/workspaceproxies/me/issue-signed-app-token": {
             "post": {
                 "security": [
@@ -8419,13 +8447,13 @@ const docTemplate = `{
         "codersdk.ProxyHealthStatus": {
             "type": "string",
             "enum": [
-                "reachable",
+                "ok",
                 "unreachable",
                 "unhealthy",
                 "unregistered"
             ],
             "x-enum-varnames": [
-                "ProxyReachable",
+                "ProxyHealthy",
                 "ProxyUnreachable",
                 "ProxyUnhealthy",
                 "ProxyUnregistered"

coderd/apidoc/swagger.json

@@ -4481,6 +4481,30 @@
         }
       }
     },
+    "/workspaceproxies/me/goingaway": {
+      "post": {
+        "security": [
+          {
+            "CoderSessionToken": []
+          }
+        ],
+        "produces": ["application/json"],
+        "tags": ["Enterprise"],
+        "summary": "Workspace proxy going away",
+        "operationId": "workspace-proxy-going-away",
+        "responses": {
+          "201": {
+            "description": "Created",
+            "schema": {
+              "$ref": "#/definitions/codersdk.Response"
+            }
+          }
+        },
+        "x-apidocgen": {
+          "skip": true
+        }
+      }
+    },
     "/workspaceproxies/me/issue-signed-app-token": {
       "post": {
         "security": [
@@ -7528,9 +7552,9 @@
     },
     "codersdk.ProxyHealthStatus": {
       "type": "string",
-      "enum": ["reachable", "unreachable", "unhealthy", "unregistered"],
+      "enum": ["ok", "unreachable", "unhealthy", "unregistered"],
       "x-enum-varnames": [
-        "ProxyReachable",
+        "ProxyHealthy",
         "ProxyUnreachable",
         "ProxyUnhealthy",
         "ProxyUnregistered"

codersdk/workspaceproxy.go

@@ -15,9 +15,9 @@ import (
 type ProxyHealthStatus string
 
 const (
-	// ProxyReachable means the proxy access url is reachable and returns a healthy
+	// ProxyHealthy means the proxy access url is reachable and returns a healthy
 	// status code.
-	ProxyReachable ProxyHealthStatus = "reachable"
+	ProxyHealthy ProxyHealthStatus = "ok"
 	// ProxyUnreachable means the proxy access url is not responding.
 	ProxyUnreachable ProxyHealthStatus = "unreachable"
 	// ProxyUnhealthy means the proxy access url is responding, but there is some

docs/api/enterprise.md

@@ -1192,7 +1192,7 @@ curl -X GET http://coder-server:8080/api/v2/workspaceproxies \
         "errors": ["string"],
         "warnings": ["string"]
       },
-      "status": "reachable"
+      "status": "ok"
     },
     "updated_at": "2019-08-24T14:15:22Z",
     "url": "string",
@@ -1234,7 +1234,7 @@ Status Code **200**
 
 | Property | Value          |
 | -------- | -------------- |
-| `status` | `reachable`    |
+| `status` | `ok`           |
 | `status` | `unreachable`  |
 | `status` | `unhealthy`    |
 | `status` | `unregistered` |
@@ -1289,7 +1289,7 @@ curl -X POST http://coder-server:8080/api/v2/workspaceproxies \
       "errors": ["string"],
       "warnings": ["string"]
     },
-    "status": "reachable"
+    "status": "ok"
   },
   "updated_at": "2019-08-24T14:15:22Z",
   "url": "string",
@@ -1342,7 +1342,7 @@ curl -X GET http://coder-server:8080/api/v2/workspaceproxies/{workspaceproxy} \
       "errors": ["string"],
       "warnings": ["string"]
     },
-    "status": "reachable"
+    "status": "ok"
   },
   "updated_at": "2019-08-24T14:15:22Z",
   "url": "string",
@@ -1453,7 +1453,7 @@ curl -X PATCH http://coder-server:8080/api/v2/workspaceproxies/{workspaceproxy}
       "errors": ["string"],
       "warnings": ["string"]
     },
-    "status": "reachable"
+    "status": "ok"
   },
   "updated_at": "2019-08-24T14:15:22Z",
   "url": "string",

docs/api/schemas.md

@@ -3476,7 +3476,7 @@ Parameter represents a set value for the scope.
 ## codersdk.ProxyHealthStatus
 
 ```json
-"reachable"
+"ok"
 ```
 
 ### Properties
@@ -3485,7 +3485,7 @@ Parameter represents a set value for the scope.
 
 | Value          |
 | -------------- |
-| `reachable`    |
+| `ok`           |
 | `unreachable`  |
 | `unhealthy`    |
 | `unregistered` |
@@ -5361,7 +5361,7 @@ Parameter represents a set value for the scope.
       "errors": ["string"],
       "warnings": ["string"]
     },
-    "status": "reachable"
+    "status": "ok"
   },
   "updated_at": "2019-08-24T14:15:22Z",
   "url": "string",
@@ -5393,7 +5393,7 @@ Parameter represents a set value for the scope.
     "errors": ["string"],
     "warnings": ["string"]
   },
-  "status": "reachable"
+  "status": "ok"
 }
 ```
 

enterprise/cli/proxyserver.go

@@ -249,6 +249,7 @@ func (*RootCmd) proxyServer() *clibase.Cmd {
 			if err != nil {
 				return xerrors.Errorf("create workspace proxy: %w", err)
 			}
+			closers.Add(func() { _ = proxy.Close() })
 
 			shutdownConnsCtx, shutdownConns := context.WithCancel(ctx)
 			defer shutdownConns()

enterprise/coderd/coderd.go

@@ -113,6 +113,7 @@ func New(ctx context.Context, options *Options) (*API, error) {
 				)
 				r.Post("/issue-signed-app-token", api.workspaceProxyIssueSignedAppToken)
 				r.Post("/register", api.workspaceProxyRegister)
+				r.Post("/goingaway", api.workspaceProxyGoingAway)
 			})
 			r.Route("/{workspaceproxy}", func(r chi.Router) {
 				r.Use(
@@ -239,7 +240,7 @@ func New(ctx context.Context, options *Options) (*API, error) {
 	if api.AGPL.Experiments.Enabled(codersdk.ExperimentMoons) {
 		// Proxy health is a moon feature.
 		api.ProxyHealth, err = proxyhealth.New(&proxyhealth.Options{
-			Interval:   time.Minute * 1,
+			Interval:   options.ProxyHealthInterval,
 			DB:         api.Database,
 			Logger:     options.Logger.Named("proxyhealth"),
 			Client:     api.HTTPClient,

enterprise/coderd/coderdenttest/coderdenttest.go

@@ -48,6 +48,7 @@ type Options struct {
 	EntitlementsUpdateInterval time.Duration
 	SCIMAPIKey                 []byte
 	UserWorkspaceQuota         int
+	ProxyHealthInterval        time.Duration
 }
 
 // New constructs a codersdk client connected to an in-memory Enterprise API instance.
@@ -74,6 +75,7 @@ func NewWithAPI(t *testing.T, options *Options) (*codersdk.Client, io.Closer, *c
 		Options:                    oop,
 		EntitlementsUpdateInterval: options.EntitlementsUpdateInterval,
 		Keys:                       Keys,
+		ProxyHealthInterval:        options.ProxyHealthInterval,
 	})
 	assert.NoError(t, err)
 	setHandler(coderAPI.AGPL.RootHandler)

enterprise/coderd/workspaceproxy.go

@@ -68,11 +68,7 @@ func (api *API) regions(rw http.ResponseWriter, r *http.Request) {
 				continue
 			}
 
-			health, ok := proxyHealth[proxy.ID]
-			if !ok {
-				health.Status = proxyhealth.Unknown
-			}
-
+			health := proxyHealth[proxy.ID]
 			regions = append(regions, codersdk.Region{
 				ID:               proxy.ID,
 				Name:             proxy.Name,
@@ -423,7 +419,7 @@ func (api *API) workspaceProxyRegister(rw http.ResponseWriter, r *http.Request)
 		//	Log:     api.Logger,
 		//	Request: r,
 		//	Action:  database.AuditActionWrite,
-		//})
+		// })
 	)
 	// aReq.Old = proxy
 	// defer commitAudit()
@@ -473,6 +469,33 @@ func (api *API) workspaceProxyRegister(rw http.ResponseWriter, r *http.Request)
 	go api.forceWorkspaceProxyHealthUpdate(api.ctx)
 }
 
+// workspaceProxyGoingAway is used to tell coderd that the workspace proxy is
+// shutting down and going away. The main purpose of this function is for the
+// health status of the workspace proxy to be more quickly updated when we know
+// that the proxy is going to be unhealthy. This does not delete the workspace
+// or cause any other side effects.
+// If the workspace proxy comes back online, even without a register, it will
+// be found healthy again by the normal checks.
+// @Summary Workspace proxy going away
+// @ID workspace-proxy-going-away
+// @Security CoderSessionToken
+// @Produce json
+// @Tags Enterprise
+// @Success 201 {object} codersdk.Response
+// @Router /workspaceproxies/me/goingaway [post]
+// @x-apidocgen {"skip": true}
+func (api *API) workspaceProxyGoingAway(rw http.ResponseWriter, r *http.Request) {
+	ctx := r.Context()
+
+	// Force a health update to happen immediately. The proxy should
+	// not return a successful response if it is going away.
+	go api.forceWorkspaceProxyHealthUpdate(api.ctx)
+
+	httpapi.Write(ctx, rw, http.StatusOK, codersdk.Response{
+		Message: "OK",
+	})
+}
+
 // reconnectingPTYSignedToken issues a signed app token for use when connecting
 // to the reconnecting PTY websocket on an external workspace proxy. This is set
 // by the client as a query parameter when connecting.
@@ -588,6 +611,9 @@ func convertProxies(p []database.WorkspaceProxy, statuses map[uuid.UUID]proxyhea
 }
 
 func convertProxy(p database.WorkspaceProxy, status proxyhealth.ProxyStatus) codersdk.WorkspaceProxy {
+	if status.Status == "" {
+		status.Status = proxyhealth.Unknown
+	}
 	return codersdk.WorkspaceProxy{
 		ID:               p.ID,
 		Name:             p.Name,

enterprise/coderd/workspaceproxy_test.go

@@ -7,6 +7,7 @@ import (
 	"net/http/httputil"
 	"net/url"
 	"testing"
+	"time"
 
 	"github.com/google/uuid"
 	"github.com/moby/moby/pkg/namesgenerator"
@@ -172,6 +173,69 @@ func TestRegions(t *testing.T) {
 		require.Error(t, err)
 		require.Empty(t, regions)
 	})
+
+	t.Run("GoingAway", func(t *testing.T) {
+		t.Parallel()
+
+		dv := coderdtest.DeploymentValues(t)
+		dv.Experiments = []string{
+			string(codersdk.ExperimentMoons),
+			"*",
+		}
+
+		db, pubsub := dbtestutil.NewDB(t)
+
+		ctx := testutil.Context(t, testutil.WaitLong)
+
+		client, closer, api := coderdenttest.NewWithAPI(t, &coderdenttest.Options{
+			Options: &coderdtest.Options{
+				AppHostname:      appHostname,
+				Database:         db,
+				Pubsub:           pubsub,
+				DeploymentValues: dv,
+			},
+			// The interval is set to 1 hour so the proxy health
+			// check will never happen manually. All checks will be
+			// forced updates.
+			ProxyHealthInterval: time.Hour,
+		})
+		t.Cleanup(func() {
+			_ = closer.Close()
+		})
+		_ = coderdtest.CreateFirstUser(t, client)
+		_ = coderdenttest.AddLicense(t, client, coderdenttest.LicenseOptions{
+			Features: license.Features{
+				codersdk.FeatureWorkspaceProxy: 1,
+			},
+		})
+
+		const proxyName = "testproxy"
+		proxy := coderdenttest.NewWorkspaceProxy(t, api, client, &coderdenttest.ProxyOptions{
+			Name: proxyName,
+		})
+		_ = proxy
+
+		require.Eventuallyf(t, func() bool {
+			proxy, err := client.WorkspaceProxyByName(ctx, proxyName)
+			if err != nil {
+				// We are testing the going away, not the initial healthy.
+				// Just force an update to change this to healthy.
+				_ = api.ProxyHealth.ForceUpdate(ctx)
+				return false
+			}
+			return proxy.Status.Status == codersdk.ProxyHealthy
+		}, testutil.WaitShort, testutil.IntervalFast, "proxy never became healthy")
+
+		_ = proxy.Close()
+		// The proxy should tell the primary on close that is is no longer healthy.
+		require.Eventuallyf(t, func() bool {
+			proxy, err := client.WorkspaceProxyByName(ctx, proxyName)
+			if err != nil {
+				return false
+			}
+			return proxy.Status.Status == codersdk.ProxyUnhealthy
+		}, testutil.WaitShort, testutil.IntervalFast, "proxy never became unhealthy after close")
+	})
 }
 
 func TestWorkspaceProxyCRUD(t *testing.T) {

enterprise/wsproxy/wsproxy.go

@@ -265,6 +265,12 @@ func New(ctx context.Context, opts *Options) (*Server, error) {
 
 func (s *Server) Close() error {
 	s.cancel()
+
+	// A timeout to prevent the SDK from blocking the server shutdown.
+	tmp, cancel := context.WithTimeout(context.Background(), 3*time.Second)
+	defer cancel()
+	_ = s.SDKClient.WorkspaceProxyGoingAway(tmp)
+
 	return s.AppServer.Close()
 }
 
@@ -294,6 +300,16 @@ func (s *Server) healthReport(rw http.ResponseWriter, r *http.Request) {
 	ctx := r.Context()
 	var report codersdk.ProxyHealthReport
 
+	// This is to catch edge cases where the server is shutting down, but might
+	// still serve a web request that returns "healthy". This is mainly just for
+	// unit tests, as shutting down the test webserver is tied to the lifecycle
+	// of the test. In practice, the webserver is tied to the lifecycle of the
+	// app, so the webserver AND the proxy will be shut down at the same time.
+	if s.ctx.Err() != nil {
+		httpapi.Write(r.Context(), rw, http.StatusInternalServerError, "workspace proxy in middle of shutting down")
+		return
+	}
+
 	// Hit the build info to do basic version checking.
 	primaryBuild, err := s.SDKClient.SDKClient.BuildInfo(ctx)
 	if err != nil {

enterprise/wsproxy/wsproxysdk/wsproxysdk.go

@@ -170,3 +170,19 @@ func (c *Client) RegisterWorkspaceProxy(ctx context.Context, req RegisterWorkspa
 	var resp RegisterWorkspaceProxyResponse
 	return resp, json.NewDecoder(res.Body).Decode(&resp)
 }
+
+func (c *Client) WorkspaceProxyGoingAway(ctx context.Context) error {
+	res, err := c.Request(ctx, http.MethodPost,
+		"/api/v2/workspaceproxies/me/goingaway",
+		nil,
+	)
+	if err != nil {
+		return xerrors.Errorf("make request: %w", err)
+	}
+	defer res.Body.Close()
+
+	if res.StatusCode != http.StatusOK {
+		return codersdk.ReadBodyAsError(res)
+	}
+	return nil
+}

site/src/api/typesGenerated.ts

@@ -1495,12 +1495,12 @@ export const ProvisionerTypes: ProvisionerType[] = ["echo", "terraform"]
 
 // From codersdk/workspaceproxy.go
 export type ProxyHealthStatus =
-  | "reachable"
+  | "ok"
   | "unhealthy"
   | "unreachable"
   | "unregistered"
 export const ProxyHealthStatuses: ProxyHealthStatus[] = [
-  "reachable",
+  "ok",
   "unhealthy",
   "unreachable",
   "unregistered",