feat: add force refresh of license entitlements (#9155)

* feat: add force refresh of license entitlements
* send "going away" mesasge on licenses pubsub on close
* Add manual refresh to licenses page

coderd/apidoc/docs.go

@@ -1024,6 +1024,31 @@ const docTemplate = `{
                 }
             }
         },
+        "/licenses/refresh-entitlements": {
+            "post": {
+                "security": [
+                    {
+                        "CoderSessionToken": []
+                    }
+                ],
+                "produces": [
+                    "application/json"
+                ],
+                "tags": [
+                    "Organizations"
+                ],
+                "summary": "Update license entitlements",
+                "operationId": "update-license-entitlements",
+                "responses": {
+                    "201": {
+                        "description": "Created",
+                        "schema": {
+                            "$ref": "#/definitions/codersdk.Response"
+                        }
+                    }
+                }
+            }
+        },
         "/licenses/{id}": {
             "delete": {
                 "security": [
@@ -8068,6 +8093,10 @@ const docTemplate = `{
                 "has_license": {
                     "type": "boolean"
                 },
+                "refreshed_at": {
+                    "type": "string",
+                    "format": "date-time"
+                },
                 "require_telemetry": {
                     "type": "boolean"
                 },

coderd/apidoc/swagger.json

@@ -880,6 +880,27 @@
         }
       }
     },
+    "/licenses/refresh-entitlements": {
+      "post": {
+        "security": [
+          {
+            "CoderSessionToken": []
+          }
+        ],
+        "produces": ["application/json"],
+        "tags": ["Organizations"],
+        "summary": "Update license entitlements",
+        "operationId": "update-license-entitlements",
+        "responses": {
+          "201": {
+            "description": "Created",
+            "schema": {
+              "$ref": "#/definitions/codersdk.Response"
+            }
+          }
+        }
+      }
+    },
     "/licenses/{id}": {
       "delete": {
         "security": [
@@ -7223,6 +7244,10 @@
         "has_license": {
           "type": "boolean"
         },
+        "refreshed_at": {
+          "type": "string",
+          "format": "date-time"
+        },
         "require_telemetry": {
           "type": "boolean"
         },

codersdk/deployment.go

@@ -103,6 +103,7 @@ type Entitlements struct {
 	HasLicense       bool                    `json:"has_license"`
 	Trial            bool                    `json:"trial"`
 	RequireTelemetry bool                    `json:"require_telemetry"`
+	RefreshedAt      time.Time               `json:"refreshed_at" format:"date-time"`
 }
 
 func (c *Client) Entitlements(ctx context.Context) (Entitlements, error) {

docs/api/enterprise.md

@@ -134,6 +134,7 @@ curl -X GET http://coder-server:8080/api/v2/entitlements \
     }
   },
   "has_license": true,
+  "refreshed_at": "2019-08-24T14:15:22Z",
   "require_telemetry": true,
   "trial": true,
   "warnings": ["string"]

docs/api/organizations.md

@@ -49,6 +49,44 @@ curl -X POST http://coder-server:8080/api/v2/licenses \
 
 To perform this operation, you must be authenticated. [Learn more](authentication.md).
 
+## Update license entitlements
+
+### Code samples
+
+```shell
+# Example request using curl
+curl -X POST http://coder-server:8080/api/v2/licenses/refresh-entitlements \
+  -H 'Accept: application/json' \
+  -H 'Coder-Session-Token: API_KEY'
+```
+
+`POST /licenses/refresh-entitlements`
+
+### Example responses
+
+> 201 Response
+
+```json
+{
+  "detail": "string",
+  "message": "string",
+  "validations": [
+    {
+      "detail": "string",
+      "field": "string"
+    }
+  ]
+}
+```
+
+### Responses
+
+| Status | Meaning                                                      | Description | Schema                                           |
+| ------ | ------------------------------------------------------------ | ----------- | ------------------------------------------------ |
+| 201    | [Created](https://tools.ietf.org/html/rfc7231#section-6.3.2) | Created     | [codersdk.Response](schemas.md#codersdkresponse) |
+
+To perform this operation, you must be authenticated. [Learn more](authentication.md).
+
 ## Create organization
 
 ### Code samples

docs/api/schemas.md

@@ -2674,6 +2674,7 @@ AuthorizationObject can represent a "set" of objects, such as: all workspaces in
     }
   },
   "has_license": true,
+  "refreshed_at": "2019-08-24T14:15:22Z",
   "require_telemetry": true,
   "trial": true,
   "warnings": ["string"]
@@ -2688,6 +2689,7 @@ AuthorizationObject can represent a "set" of objects, such as: all workspaces in
 | `features`          | object                               | false    |              |             |
 | » `[any property]`  | [codersdk.Feature](#codersdkfeature) | false    |              |             |
 | `has_license`       | boolean                              | false    |              |             |
+| `refreshed_at`      | string                               | false    |              |             |
 | `require_telemetry` | boolean                              | false    |              |             |
 | `trial`             | boolean                              | false    |              |             |
 | `warnings`          | array of string                      | false    |              |             |

enterprise/coderd/coderd.go

@@ -130,6 +130,7 @@ func New(ctx context.Context, options *Options) (_ *API, err error) {
 		})
 		r.Route("/licenses", func(r chi.Router) {
 			r.Use(apiKeyMiddleware)
+			r.Post("/refresh-entitlements", api.postRefreshEntitlements)
 			r.Post("/", api.postLicense)
 			r.Get("/", api.licenses)
 			r.Delete("/{id}", api.deleteLicense)
@@ -403,10 +404,13 @@ type API struct {
 }
 
 func (api *API) Close() error {
-	api.cancel()
+	// Replica manager should be closed first. This is because the replica
+	// manager updates the replica's table in the database when it closes.
+	// This tells other Coderds that it is now offline.
 	if api.replicaManager != nil {
 		_ = api.replicaManager.Close()
 	}
+	api.cancel()
 	if api.derpMesh != nil {
 		_ = api.derpMesh.Close()
 	}
@@ -802,6 +806,17 @@ func (api *API) runEntitlementsLoop(ctx context.Context) {
 	updates := make(chan struct{}, 1)
 	subscribed := false
 
+	defer func() {
+		// If this function ends, it means the context was cancelled and this
+		// coderd is shutting down. In this case, post a pubsub message to
+		// tell other coderd's to resync their entitlements. This is required to
+		// make sure things like replica counts are updated in the UI.
+		// Ignore the error, as this is just a best effort. If it fails,
+		// the system will eventually recover as replicas timeout
+		// if their heartbeats stop. The best effort just tries to update the
+		// UI faster if it succeeds.
+		_ = api.Pubsub.Publish(PubsubEventLicenses, []byte("going away"))
+	}()
 	for {
 		select {
 		case <-ctx.Done():

enterprise/coderd/license/license.go

@@ -225,6 +225,7 @@ func Entitlements(
 			entitlements.Features[featureName] = feature
 		}
 	}
+	entitlements.RefreshedAt = now
 
 	return entitlements, nil
 }

enterprise/coderd/licenses.go

@@ -8,6 +8,7 @@ import (
 	_ "embed"
 	"encoding/base64"
 	"encoding/json"
+	"fmt"
 	"net/http"
 	"strconv"
 	"strings"
@@ -150,6 +151,75 @@ func (api *API) postLicense(rw http.ResponseWriter, r *http.Request) {
 	httpapi.Write(ctx, rw, http.StatusCreated, convertLicense(dl, rawClaims))
 }
 
+// postRefreshEntitlements forces an `updateEntitlements` call and publishes
+// a message to the PubsubEventLicenses topic to force other replicas
+// to update their entitlements.
+// Updates happen automatically on a timer, however that time is every 10 minutes,
+// and we want to be able to force an update immediately in some cases.
+//
+// @Summary Update license entitlements
+// @ID update-license-entitlements
+// @Security CoderSessionToken
+// @Produce json
+// @Tags Organizations
+// @Success 201 {object} codersdk.Response
+// @Router /licenses/refresh-entitlements [post]
+func (api *API) postRefreshEntitlements(rw http.ResponseWriter, r *http.Request) {
+	ctx := r.Context()
+
+	// If the user cannot create a new license, then they cannot refresh entitlements.
+	// Refreshing entitlements is a way to force a refresh of the license, so it is
+	// equivalent to creating a new license.
+	if !api.AGPL.Authorize(r, rbac.ActionCreate, rbac.ResourceLicense) {
+		httpapi.Forbidden(rw)
+		return
+	}
+
+	// Prevent abuse by limiting how often we allow a forced refresh.
+	now := time.Now()
+	if diff := now.Sub(api.entitlements.RefreshedAt); diff < time.Minute {
+		wait := time.Minute - diff
+		rw.Header().Set("Retry-After", strconv.Itoa(int(wait.Seconds())))
+		httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
+			Message: fmt.Sprintf("Entitlements already recently refreshed, please wait %d seconds to force a new refresh", int(wait.Seconds())),
+			Detail:  fmt.Sprintf("Last refresh at %s", now.UTC().String()),
+		})
+		return
+	}
+
+	err := api.replicaManager.UpdateNow(ctx)
+	if err != nil {
+		httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
+			Message: "Failed to sync replicas",
+			Detail:  err.Error(),
+		})
+		return
+	}
+
+	err = api.updateEntitlements(ctx)
+	if err != nil {
+		httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
+			Message: "Failed to update entitlements",
+			Detail:  err.Error(),
+		})
+		return
+	}
+
+	err = api.Pubsub.Publish(PubsubEventLicenses, []byte("refresh"))
+	if err != nil {
+		api.Logger.Error(context.Background(), "failed to publish forced entitlement update", slog.Error(err))
+		httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
+			Message: "Failed to publish forced entitlement update. Other replicas might not be updated.",
+			Detail:  err.Error(),
+		})
+		return
+	}
+
+	httpapi.Write(ctx, rw, http.StatusOK, codersdk.Response{
+		Message: "Entitlements updated",
+	})
+}
+
 // @Summary Get licenses
 // @ID get-licenses
 // @Security CoderSessionToken

site/src/api/api.ts

@@ -808,6 +808,10 @@ export const putWorkspaceExtension = async (
   })
 }
 
+export const refreshEntitlements = async (): Promise<void> => {
+  await axios.post("/api/v2/licenses/refresh-entitlements")
+}
+
 export const getEntitlements = async (): Promise<TypesGen.Entitlements> => {
   try {
     const response = await axios.get("/api/v2/entitlements")
@@ -821,6 +825,7 @@ export const getEntitlements = async (): Promise<TypesGen.Entitlements> => {
         require_telemetry: false,
         trial: false,
         warnings: [],
+        refreshed_at: "",
       }
     }
     throw ex

site/src/api/typesGenerated.ts

@@ -416,6 +416,7 @@ export interface Entitlements {
   readonly has_license: boolean
   readonly trial: boolean
   readonly require_telemetry: boolean
+  readonly refreshed_at: string
 }
 
 // From codersdk/deployment.go

site/src/pages/DeploySettingsPage/LicensesSettingsPage/LicensesSettingsPage.tsx

@@ -9,14 +9,20 @@ import useToggle from "react-use/lib/useToggle"
 import { pageTitle } from "utils/page"
 import { entitlementsMachine } from "xServices/entitlements/entitlementsXService"
 import LicensesSettingsPageView from "./LicensesSettingsPageView"
+import { getErrorMessage } from "api/errors"
 
 const LicensesSettingsPage: FC = () => {
   const queryClient = useQueryClient()
-  const [entitlementsState] = useMachine(entitlementsMachine)
-  const { entitlements } = entitlementsState.context
+  const [entitlementsState, sendEvent] = useMachine(entitlementsMachine)
+  const { entitlements, getEntitlementsError } = entitlementsState.context
   const [searchParams, setSearchParams] = useSearchParams()
   const success = searchParams.get("success")
   const [confettiOn, toggleConfettiOn] = useToggle(false)
+  if (getEntitlementsError) {
+    displayError(
+      getErrorMessage(getEntitlementsError, "Failed to fetch entitlements"),
+    )
+  }
 
   const { mutate: removeLicenseApi, isLoading: isRemovingLicense } =
     useMutation(removeLicense, {
@@ -58,6 +64,10 @@ const LicensesSettingsPage: FC = () => {
         licenses={licenses}
         isRemovingLicense={isRemovingLicense}
         removeLicense={(licenseId: number) => removeLicenseApi(licenseId)}
+        refreshEntitlements={() => {
+          const x = sendEvent("REFRESH")
+          return !x.context.getEntitlementsError
+        }}
       />
     </>
   )

site/src/pages/DeploySettingsPage/LicensesSettingsPage/LicensesSettingsPageView.tsx

@@ -2,6 +2,7 @@ import Button from "@mui/material/Button"
 import { makeStyles, useTheme } from "@mui/styles"
 import Skeleton from "@mui/material/Skeleton"
 import AddIcon from "@mui/icons-material/AddOutlined"
+import RefreshIcon from "@mui/icons-material/Refresh"
 import { GetLicensesResponse } from "api/api"
 import { Header } from "components/DeploySettingsLayout/Header"
 import { LicenseCard } from "components/LicenseCard/LicenseCard"
@@ -11,6 +12,8 @@ import Confetti from "react-confetti"
 import { Link } from "react-router-dom"
 import useWindowSize from "react-use/lib/useWindowSize"
 import MuiLink from "@mui/material/Link"
+import { displaySuccess } from "components/GlobalSnackbar/utils"
+import Tooltip from "@mui/material/Tooltip"
 
 type Props = {
   showConfetti: boolean
@@ -20,6 +23,7 @@ type Props = {
   licenses?: GetLicensesResponse[]
   isRemovingLicense: boolean
   removeLicense: (licenseId: number) => void
+  refreshEntitlements?: () => boolean
 }
 
 const LicensesSettingsPageView: FC<Props> = ({
@@ -30,6 +34,7 @@ const LicensesSettingsPageView: FC<Props> = ({
   licenses,
   isRemovingLicense,
   removeLicense,
+  refreshEntitlements,
 }) => {
   const styles = useStyles()
   const { width, height } = useWindowSize()
@@ -55,13 +60,29 @@ const LicensesSettingsPageView: FC<Props> = ({
           description="Manage licenses to unlock Enterprise features."
         />
 
-        <Button
-          component={Link}
-          to="/deployment/licenses/add"
-          startIcon={<AddIcon />}
-        >
-          Add a license
-        </Button>
+        <Stack direction="row" spacing={2}>
+          <Button
+            component={Link}
+            to="/deployment/licenses/add"
+            startIcon={<AddIcon />}
+          >
+            Add a license
+          </Button>
+          <Tooltip title="Refresh license entitlements. This is done automatically every 10 minutes.">
+            <Button
+              onClick={() => {
+                if (refreshEntitlements) {
+                  if (refreshEntitlements()) {
+                    displaySuccess("Successfully refreshed licenses")
+                  }
+                }
+              }}
+              startIcon={<RefreshIcon />}
+            >
+              Refresh
+            </Button>
+          </Tooltip>
+        </Stack>
       </Stack>
 
       {isLoading && <Skeleton variant="rectangular" height={200} />}

site/src/testHelpers/entities.ts

@@ -1450,6 +1450,7 @@ export const MockEntitlements: TypesGen.Entitlements = {
   features: withDefaultFeatures({}),
   require_telemetry: false,
   trial: false,
+  refreshed_at: "2022-05-20T16:45:57.122Z",
 }
 
 export const MockEntitlementsWithWarnings: TypesGen.Entitlements = {
@@ -1458,6 +1459,7 @@ export const MockEntitlementsWithWarnings: TypesGen.Entitlements = {
   has_license: true,
   trial: false,
   require_telemetry: false,
+  refreshed_at: "2022-05-20T16:45:57.122Z",
   features: withDefaultFeatures({
     user_limit: {
       enabled: true,
@@ -1482,6 +1484,7 @@ export const MockEntitlementsWithAuditLog: TypesGen.Entitlements = {
   has_license: true,
   require_telemetry: false,
   trial: false,
+  refreshed_at: "2022-05-20T16:45:57.122Z",
   features: withDefaultFeatures({
     audit_log: {
       enabled: true,
@@ -1496,6 +1499,7 @@ export const MockEntitlementsWithScheduling: TypesGen.Entitlements = {
   has_license: true,
   require_telemetry: false,
   trial: false,
+  refreshed_at: "2022-05-20T16:45:57.122Z",
   features: withDefaultFeatures({
     advanced_template_scheduling: {
       enabled: true,

site/src/xServices/deploymentStats/deploymentStatsMachine.ts

@@ -4,6 +4,7 @@ import { assign, createMachine } from "xstate"
 
 export const deploymentStatsMachine = createMachine(
   {
+    /** @xstate-layout N4IgpgJg5mDOIC5QTABwDYHsCeBbMAdgC4DKRAhkbALLkDGAFgJYFgB0sFVAxBJq2xYA3TAGt2KDDnzEylGvWYDO8hMMx1KTfgG0ADAF19BxKFSZYTItoKmQAD0QBWAMwuANCGyIATAHYATgBfIM9JLDxCUi4FRhZ2FR4wACdkzGS2DEoAM3TcNnDpKLkqWjjlGLUCEU1rXUNjO3NLOtskB0QAvU9vBAAWVxCwtAiZaPkypXYmCHQwbgAlAFEAGQB5AEEAEUb25qsbO0cEAEY9AA4exDOANhDQkAJMFHh2wsjZGMn4posD-iOiD6Piup3ObCcQxA7zGJViUw4MV+LUO7WOfT03S81xOLihMOKX0U8UEszAyP+bVAxxc5z6oJ8AT6EPuQSAA */
     id: "deploymentStatsMachine",
     predictableActionArguments: true,
 

site/src/xServices/entitlements/entitlementsXService.ts

@@ -9,6 +9,7 @@ export type EntitlementsContext = {
 
 export const entitlementsMachine = createMachine(
   {
+    /** @xstate-layout N4IgpgJg5mDOIC5RgHYBcCWaA2YC2qasAsgIYDGAFhimAHQBOYAZk7JQMQQD2tdNAN24Brek1ZxKAUXRZcBdLADaABgC6iUAAdusLBl6aQAD0QAWAGwAOOgGYAjACYArABoQAT0SOA7AE46C0czP2cVHzNIlRV7HwBfOPdCOXxCEgpqPnE2TjAGBm4GOi1sUjRmQrxGFhyZTBxUxVUNJBAdPUxDVtMESxsHF3cvBFtbAKcQsIiomPjE8FkGhSIyKhp6GDRMFCg6lOXYLl56QRENsDQ9pbTmo3b9LtAe2wtnOmcBi1Hos0c-Hx8Q0QVnsdBCfj8Vh8UKsYVscySi3kaVWmXOWxouyRjSIHDyBSKJTKFQYVU2V2RTXUd10DxQRmer3en2+Kl+-0BnkQsTM7yCZimkTM0ViCUR9UpKwy634EFwHAASlIAGJKgDKAAlbq17p16d1uZCgb1HBZ3n54fZIrCgrZHGKFhKcek1nx8YVFSr1VrqTraXqjMN7KE6CoLT4rWYbY4HJyev5QabnBYzA4LX5XmYEvMUNwIHAjMlropUesaR0DPqnogALRR401s3RaKORwqWyw2Ehe3zIuSl1o6oSdjlukMxDwlR0HwWAFsiz-PwqZz-Y0r3m2AXhaGhPxmHzOB1952lvibbZYp0HUcBg0m42wuwp5wfeHP9s98X7FHSvgYOVgDelbjggKjGiENgrpa1rJjGn6Ot+Ja-vQsAAK7kOQcDwH6FaPCYiBQXQSYpmmYyZsa9gqI4oYDM49gWMGPimrOR7Ygcp70O6DBAXhPSEcRqbBmRzhmBRIZhtBUawbG2ZxEAA */
     id: "entitlementsMachine",
     predictableActionArguments: true,
     tsTypes: {} as import("./entitlementsXService.typegen").Typegen0,
@@ -22,13 +23,27 @@ export const entitlementsMachine = createMachine(
     },
     initial: "gettingEntitlements",
     states: {
+      refresh: {
+        invoke: {
+          id: "refreshEntitlements",
+          src: "refreshEntitlements",
+          onDone: {
+            target: "gettingEntitlements",
+          },
+          onError: {
+            target: "error",
+            actions: ["assignGetEntitlementsError"],
+          },
+        },
+        entry: "clearGetEntitlementsError",
+      },
       gettingEntitlements: {
         entry: "clearGetEntitlementsError",
         invoke: {
           id: "getEntitlements",
           src: "getEntitlements",
           onDone: {
-            target: "success",
+            target: "idle",
             actions: ["assignEntitlements"],
           },
           onError: {
@@ -37,11 +52,18 @@ export const entitlementsMachine = createMachine(
           },
         },
       },
+      idle: {
+        on: {
+          REFRESH: "refresh",
+        },
+      },
       success: {
         type: "final",
       },
       error: {
-        type: "final",
+        on: {
+          REFRESH: "refresh",
+        },
       },
     },
   },
@@ -51,13 +73,18 @@ export const entitlementsMachine = createMachine(
         entitlements: (_, event) => event.data,
       }),
       assignGetEntitlementsError: assign({
-        getEntitlementsError: (_, event) => event.data,
+        getEntitlementsError: (_, event) => {
+          return event.data
+        },
       }),
       clearGetEntitlementsError: assign({
         getEntitlementsError: (_) => undefined,
       }),
     },
     services: {
+      refreshEntitlements: async () => {
+        return API.refreshEntitlements()
+      },
       getEntitlements: async () => {
         // Entitlements is injected by the Coder server into the HTML document.
         const entitlements = document.querySelector(