Nothing the control queue carries is lost while the store restarts: an enrolment claims its token and spends it last, and is asked to try again; build results, upgrades and catch-ups are handed back, bounded (novox/hq issue 083)

internal/inventory/nodes.go

@@ -212,6 +212,60 @@ var ErrTokenRefused = errors.New("that token cannot be used")
 // The update is the check: one statement that both finds a live token and marks it used, so two
 // simultaneous redemptions of one secret cannot both succeed. Reading first and writing second
 // would leave exactly that gap.
+// ClaimLease is how long a claimed token is held for the one presenter that claimed it. Long
+// enough for an enrolment to be tried again through a store restart; short enough that a host
+// which gave up and was started over, with keys of its own, is not kept waiting long.
+const ClaimLease = 2 * time.Minute
+
+// ErrTokenInUse is a token another presenter holds a claim on right now. Not a refusal: the claim
+// lapses, and asking again after it is the answer.
+var ErrTokenInUse = errors.New("the token is being used by another enrolment")
+
+// Claim takes a token for one presenter — `by`, which names the key presenting it — for the length
+// of a lease, and says which node it enrols. The same presenter may claim it again, as may anyone
+// once the lease has lapsed; nothing is spent until Spend (novox/hq 04-ISSUES/083).
+func (i *Inventory) Claim(ctx context.Context, secret, by string) (Node, error) {
+	var id string
+	err := i.store.Pool().QueryRow(ctx,
+		`update enrolment_token set claimed_by = $2, claimed_until = now() + $3::interval
+		  where secret = $1 and redeemed is null and expires > now()
+		    and (claimed_by is null or claimed_by = $2 or claimed_until < now())
+		  returning node`, hashSecret(secret), by, ClaimLease.String()).Scan(&id)
+	if errors.Is(err, pgx.ErrNoRows) {
+		// Unusable, or held by someone else — told apart, because the second passes.
+		var held bool
+		probe := i.store.Pool().QueryRow(ctx,
+			`select true from enrolment_token
+			  where secret = $1 and redeemed is null and expires > now()`, hashSecret(secret)).Scan(&held)
+		if probe == nil && held {
+			return Node{}, ErrTokenInUse
+		}
+		return Node{}, ErrTokenRefused
+	}
+	if err != nil {
+		return Node{}, err
+	}
+	var n Node
+	err = i.store.Pool().QueryRow(ctx,
+		`select id, name, created from node where id = $1`, id).Scan(&n.ID, &n.Name, &n.Created)
+	return n, err
+}
+
+// Spend makes a claimed token used, only for the presenter holding the claim. The last write of an
+// enrolment, so a token is spent exactly when the node it enrolled is complete.
+func (i *Inventory) Spend(ctx context.Context, secret, by string) error {
+	tag, err := i.store.Pool().Exec(ctx,
+		`update enrolment_token set redeemed = now()
+		  where secret = $1 and redeemed is null and claimed_by = $2`, hashSecret(secret), by)
+	if err != nil {
+		return err
+	}
+	if tag.RowsAffected() == 0 {
+		return ErrTokenRefused
+	}
+	return nil
+}
+
 func (i *Inventory) Redeem(ctx context.Context, secret string) (Node, error) {
 	var id string
 	err := i.store.Pool().QueryRow(ctx,