coder · Emyrk · Aug 22, 2023 · Aug 17, 2023 · Aug 17, 2023 · Aug 17, 2023
diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go
diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json
diff --git a/codersdk/deployment.go b/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) {

diff --git a/docs/api/enterprise.md b/docs/api/enterprise.md
diff --git a/docs/api/organizations.md b/docs/api/organizations.md
diff --git a/docs/api/schemas.md b/docs/api/schemas.md
diff --git a/enterprise/coderd/coderd.go b/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()
 	}
@@ -799,6 +803,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():

diff --git a/enterprise/coderd/license/license.go b/enterprise/coderd/license/license.go
@@ -225,6 +225,7 @@ func Entitlements(
 			entitlements.Features[featureName] = feature
 		}
 	}
+	entitlements.RefreshedAt = now
 
 	return entitlements, nil
 }

diff --git a/enterprise/coderd/licenses.go b/enterprise/coderd/licenses.go
@@ -8,6 +8,7 @@ import (
 	_ "embed"
 	"encoding/base64"
 	"encoding/json"
+	"fmt"
 	"net/http"
 	"strconv"
 	"strings"
@@ -150,6 +151,71 @@ 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))
+		// don't fail the HTTP request, since we did write it successfully to the database
+	}
+
+	httpapi.Write(ctx, rw, http.StatusOK, codersdk.Response{
+		Message: "Entitlements updated",
+	})
+}
+
 // @Summary Get licenses
 // @ID get-licenses
 // @Security CoderSessionToken

diff --git a/site/src/api/api.ts b/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

diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts
diff --git a/site/src/pages/DeploySettingsPage/LicensesSettingsPage/LicensesSettingsPage.tsx b/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
+        }}
       />
     </>
   )
-Original file line number
+Diff line change
@@ Expand Up / @@ -225,6 +225,7 @@ func Entitlements( @@
     			entitlements.Features[featureName] = feature
     		}
     	}
+    	entitlements.RefreshedAt = now
     	return entitlements, nil
     }
@@ Expand Down @@