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

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