Apply expected consumption for all limits in one write

api/src/main/repositories/mongoose/ContractRepository.ts

@@ -185,6 +185,46 @@ class ContractRepository extends RepositoryBase {
     return contract ? toPlainObject<LeanContract>(contract.toJSON()) : null;
   }
 
+  /**
+   * Add to several usage levels of one contract in a single atomic update.
+   *
+   * `$inc` is evaluated by the database against the stored document rather than
+   * against a copy the process read earlier, so concurrent increments compose
+   * instead of overwriting one another. Read-modify-write cannot do this from
+   * the application: two callers who read the same value both write the same
+   * total, and one consumption disappears.
+   *
+   * Every path is required to exist by the filter, so a limit that is not part
+   * of the contract matches no document and is reported to the caller instead
+   * of being created by the update - `$inc` would otherwise happily add the
+   * field. That check and the increment are one operation, so a limit cannot be
+   * validated and then vanish before the write.
+   *
+   * @param increments usage level path (`service.limit`) to amount to add.
+   * @returns the contract as it is after the increment, or null when the filter
+   *          matched nothing.
+   */
+  async incrementUsageLevels(
+    userId: string,
+    increments: Record<string, number>
+  ): Promise<LeanContract | null> {
+    const filter: Record<string, unknown> = { 'userContact.userId': userId };
+    const inc: Record<string, number> = {};
+
+    for (const [path, amount] of Object.entries(increments)) {
+      filter[`usageLevels.${path}.consumed`] = { $exists: true };
+      inc[`usageLevels.${path}.consumed`] = amount;
+    }
+
+    const contract = await ContractMongoose.findOneAndUpdate(
+      filter,
+      { $inc: inc },
+      { new: true }
+    );
+
+    return contract ? toPlainObject<LeanContract>(contract.toJSON()) : null;
+  }
+
   async changeServiceName(oldServiceName: string, newServiceName: string, organizationId: string): Promise<number> {
     const oldServiceKey = oldServiceName.toLowerCase();
     const newServiceKey = newServiceName.toLowerCase();

api/src/main/services/ContractService.ts

@@ -459,38 +459,95 @@ class ContractService {
     usageLimitId: string,
     expectedConsumption: number
   ): Promise<void> {
-    let contract = await this.cacheService.get(`contracts.${userId}`);
+    await this._applyExpectedConsumptions(userId, { [usageLimitId]: expectedConsumption });
+  }
 
-    if (!contract) {
-      contract = await this.contractRepository.findByUserId(userId);
+  /**
+   * Apply several expected consumptions to a contract in one atomic update.
+   *
+   * Read-modify-write loses consumptions, and it does so in two ways. Applying
+   * the limits of one evaluation in turn made each application read the whole
+   * contract, change a single usage level in its own copy and write the whole
+   * contract back, so an evaluation touching two limits recorded one of them.
+   * Batching the limits into a single read and write fixes that, but not the
+   * case of two requests arriving together: both read the same consumed value,
+   * both write the same total, and one consumption is gone.
+   *
+   * The database is the only place that can settle this, so the increment is
+   * handed to it as `$inc` and evaluated against the stored document. Whatever
+   * order concurrent calls arrive in, every one of them is added.
+   */
+  async _applyExpectedConsumptions(
+    userId: string,
+    expectedConsumptions: Record<string, number>
+  ): Promise<void> {
+    const usageLimitIds = Object.keys(expectedConsumptions);
+    if (usageLimitIds.length === 0) {
+      return;
     }
 
-    if (!contract) {
-      throw new Error(`Contract with userId ${userId} not found`);
+    const targets = usageLimitIds.map(usageLimitId => {
+      const serviceName: string = usageLimitId.split('-')[0];
+      const usageLimit: string = usageLimitId.split('-')[1];
+
+      return { serviceName, usageLimit, amount: expectedConsumptions[usageLimitId] };
+    });
+
+    const increments: Record<string, number> = {};
+    for (const { serviceName, usageLimit, amount } of targets) {
+      increments[`${serviceName}.${usageLimit}`] = amount;
+    }
+
+    const updatedContract = await this.contractRepository.incrementUsageLevels(userId, increments);
+
+    if (!updatedContract) {
+      // The update requires the contract and every named usage level to exist,
+      // so it matched nothing. Which of the two is missing only matters for the
+      // message, and is worth a read to get right.
+      await this._explainMissingUsageLevels(userId, targets);
     }
 
-    const serviceName: string = usageLimitId.split('-')[0];
-    const usageLimit: string = usageLimitId.split('-')[1];
+    const appliedAt = new Date().getTime();
 
-    if (contract.usageLevels[serviceName][usageLimit]) {
+    for (const { serviceName, usageLimit, amount } of targets) {
+      // What the level held before this call, for `_revertExpectedConsumption`.
+      // Derived from the result rather than from a prior read, so it is this
+      // caller's own contribution that gets taken back even if others landed in
+      // between.
       await this.cacheService.set(
-        `${new Date().getTime()}.usageLevels.${userId}.${serviceName}.${usageLimit}`,
-        contract.usageLevels[serviceName][usageLimit].consumed,
+        `${appliedAt}.usageLevels.${userId}.${serviceName}.${usageLimit}`,
+        updatedContract!.usageLevels[serviceName][usageLimit].consumed - amount,
         120
       ); // 120 secs = 2 mins
+    }
 
-      contract.usageLevels[serviceName][usageLimit].consumed += expectedConsumption;
+    await this.cacheService.set(`contracts.${userId}`, updatedContract, 3600, true); // Cache for 1 hour
+  }
 
-      const updatedContract = await this.contractRepository.update(userId, contract);
+  /**
+   * Say which part of the contract was missing, having established that one was.
+   */
+  private async _explainMissingUsageLevels(
+    userId: string,
+    targets: { serviceName: string; usageLimit: string }[]
+  ): Promise<never> {
+    const contract = await this.contractRepository.findByUserId(userId);
 
-      if (!updatedContract) {
-        throw new Error(`Failed to update contract for userId ${userId}`);
-      }
+    if (!contract) {
+      throw new Error(`Contract with userId ${userId} not found`);
+    }
 
-      await this.cacheService.set(`contracts.${userId}`, updatedContract, 3600, true); // Cache for 1 hour
-    } else {
-      throw new Error(`Usage level ${usageLimit} not found in contract for userId ${userId}`);
+    const missing = targets.find(
+      ({ serviceName, usageLimit }) => !contract.usageLevels[serviceName]?.[usageLimit]
+    );
+
+    if (missing) {
+      throw new Error(
+        `Usage level ${missing.usageLimit} not found in contract for userId ${userId}`
+      );
     }
+
+    throw new Error(`Failed to update contract for userId ${userId}`);
   }
 
   async _revertExpectedConsumption(

api/src/main/utils/feature-evaluation/featureEvaluation.ts

@@ -61,20 +61,21 @@ async function evaluateFeature(
         }
       }
 
-      // Then apply all consumptions after validation has passed
+      // Then apply all consumptions after validation has passed.
+      //
+      // In one call rather than one per limit: each application reads the whole
+      // contract, increments one usage level and writes the whole contract
+      // back, so running them concurrently made every one of them start from
+      // the same state and only the last write survive.
       if (options.userId) {
         const contractService: ContractService = container.resolve('contractService');
 
-        const limits = Object.keys(featureEvaluation.used);
-        await Promise.all(
-          limits.map(limit =>
-            contractService._applyExpectedConsumption(
-              options.userId!,
-              limit,
-              expectedConsumption[limit]
-            )
-          )
-        );
+        const consumptions: Record<string, number> = {};
+        for (const limit of Object.keys(featureEvaluation.used)) {
+          consumptions[limit] = expectedConsumption[limit];
+        }
+
+        await contractService._applyExpectedConsumptions(options.userId, consumptions);
       }
     }
   }