Optimize balance fetching for gated actions (#6145)

* optimize fetching for refresh route * fix group test * add type
hicommonwealth · Jan 4, 2024 · 451890f · 451890f
1 parent d424dc3
commit 451890f
Show file tree

Hide file tree

Showing 3 changed files with 94 additions and 70 deletions.
diff --git a/packages/commonwealth/server/util/requirementsModule/refreshMembershipsForAddress.ts b/packages/commonwealth/server/util/requirementsModule/refreshMembershipsForAddress.ts
@@ -1,10 +1,13 @@
 import moment from 'moment';
-import { Sequelize } from 'sequelize';
+import { FindOptions, Op, Sequelize } from 'sequelize';
 import { MEMBERSHIP_REFRESH_TTL_SECONDS } from '../../config';
 import { DB } from '../../models';
 import { AddressAttributes } from '../../models/address';
 import { GroupAttributes } from '../../models/group';
-import { MembershipInstance } from '../../models/membership';
+import {
+  MembershipAttributes,
+  MembershipInstance,
+} from '../../models/membership';
 import { TokenBalanceCache } from '../tokenBalanceCache/tokenBalanceCache';
 import { OptionsWithBalances } from '../tokenBalanceCache/types';
 import { makeGetBalancesOptions } from './makeGetBalancesOptions';
@@ -25,7 +28,66 @@ export async function refreshMembershipsForAddress(
   groups: GroupAttributes[],
   cacheRefresh: boolean,
 ): Promise<MembershipInstance[]> {
-  const getBalancesOptions = makeGetBalancesOptions(groups, [address]);
+  const findAllQuery: FindOptions<MembershipAttributes> = {
+    where: {
+      group_id: {
+        [Op.in]: groups.map((g) => g.id),
+      },
+      address_id: address.id,
+    },
+    include: [
+      {
+        model: models.Group,
+        as: 'group',
+      },
+    ],
+  };
+  const existingMemberships = await models.Membership.findAll(findAllQuery);
+
+  const membershipsToCreate: MembershipAttributes[] = [];
+  const membershipsToUpdate: MembershipAttributes[] = [];
+  const freshMemberships: MembershipAttributes[] = [];
+
+  for (const group of groups) {
+    const membership = existingMemberships.find(
+      (m) => m.group_id === group.id && m.address_id === address.id,
+    );
+
+    // membership does not exist
+    if (!membership) {
+      membershipsToCreate.push({
+        group_id: group.id,
+        address_id: address.id,
+        last_checked: null,
+      });
+      continue;
+    }
+
+    // membership exists
+
+    if (!cacheRefresh) {
+      const expiresAt = moment(membership.last_checked).add(
+        MEMBERSHIP_REFRESH_TTL_SECONDS,
+        'seconds',
+      );
+      if (moment().isBefore(expiresAt)) {
+        // membership is fresh
+        freshMemberships.push(membership);
+        continue;
+      }
+    }
+
+    // membership is stale
+    membershipsToUpdate.push(membership);
+  }
+
+  // only fetch balances for groups with stale membership
+  const groupsToFetchBalance = groups.filter(
+    (g) => !freshMemberships.find((m) => m.group_id === g.id),
+  );
+  const getBalancesOptions = makeGetBalancesOptions(groupsToFetchBalance, [
+    address,
+  ]);
   const balances = await Promise.all(
     getBalancesOptions.map(async (options) => {
       return {
@@ -38,79 +100,46 @@ export async function refreshMembershipsForAddress(
     }),
   );
 
-  // update membership for each group
-  const updatedMemberships = await Promise.all(
-    groups.map(async (group) => {
-      const membership = await models.Membership.findOne({
-        where: {
-          group_id: group.id,
-          address_id: address.id,
-        },
-        include: [
-          {
-            model: models.Group,
-            as: 'group',
-          },
-        ],
-      });
-
-      if (!cacheRefresh && membership) {
-        // membership exists
-        const expiresAt = moment(membership.last_checked).add(
-          MEMBERSHIP_REFRESH_TTL_SECONDS,
-          'seconds',
-        );
-        if (moment().isBefore(expiresAt)) {
-          // membership is fresh, don't recompute
-          return membership;
-        }
-        // membership is stale, recompute
-        return recomputeMembership(
-          models,
-          membership,
-          group,
-          address,
-          balances,
-        );
-      }
-
-      // membership does not exist, create it and recompute
-      return recomputeMembership(models, membership, group, address, balances);
-    }),
+  const toBulkCreate = [...membershipsToUpdate, ...membershipsToCreate].map(
+    (m) =>
+      computeMembership(
+        groups.find((g) => g.id === m.group_id),
+        address,
+        balances,
+      ),
   );
 
-  return updatedMemberships;
+  await models.Membership.bulkCreate(toBulkCreate, {
+    updateOnDuplicate: ['reject_reason', 'last_checked'],
+  });
+
+  // must query again to get newly created values after bulkCreate
+  return models.Membership.findAll(findAllQuery);
 }
 
 /**
- * recomputeMembership checks the membership against the requirements,
+ * computeMembership checks the membership against the requirements,
  * updates (or creates) the membership and returns it
  * @param membership The membership to recompute
  * @param group The group of the membership
  * @param address The user address
- * @returns MembershipInstance
+ * @returns MembershipAttributes
  */
-async function recomputeMembership(
-  models: DB,
-  membership: MembershipInstance | null,
+function computeMembership(
   group: GroupAttributes,
   address: AddressAttributes,
   balances: OptionsWithBalances[],
-): Promise<MembershipInstance> {
+): MembershipAttributes {
   const { requirements } = group;
-  const { isValid, messages } = await validateGroupMembership(
+  const { isValid, messages } = validateGroupMembership(
     address.address,
     requirements,
     balances,
   );
-  const computedMembership = {
+  return {
     group_id: group.id,
     address_id: address.id,
     reject_reason: isValid ? null : messages,
     last_checked: Sequelize.literal('CURRENT_TIMESTAMP') as any,
   };
-  if (!membership) {
-    return models.Membership.create(computedMembership);
-  }
-  return membership.update(computedMembership);
 }
diff --git a/packages/commonwealth/server/util/requirementsModule/validateGroupMembership.ts b/packages/commonwealth/server/util/requirementsModule/validateGroupMembership.ts
@@ -21,32 +21,28 @@ export type ValidateGroupMembershipResponse = {
  * @param balances address balances
  * @returns ValidateGroupMembershipResponse validity and messages on requirements that failed
  */
-export default async function validateGroupMembership(
+export default function validateGroupMembership(
   userAddress: string,
   requirements: Requirement[],
   balances: OptionsWithBalances[],
   numRequiredRequirements: number = 0,
-): Promise<ValidateGroupMembershipResponse> {
+): ValidateGroupMembershipResponse {
   const response: ValidateGroupMembershipResponse = {
     isValid: true,
     messages: [],
   };
   let allowListOverride = false;
   let numRequirementsMet = 0;
 
-  const checks = requirements.map(async (requirement) => {
+  requirements.forEach((requirement) => {
     let checkResult: { result: boolean; message: string };
     switch (requirement.rule) {
       case 'threshold': {
-        checkResult = await _thresholdCheck(
-          userAddress,
-          requirement.data,
-          balances,
-        );
+        checkResult = _thresholdCheck(userAddress, requirement.data, balances);
         break;
       }
       case 'allow': {
-        checkResult = await _allowlistCheck(
+        checkResult = _allowlistCheck(
           userAddress,
           requirement.data as AllowlistData,
         );
@@ -74,8 +70,6 @@ export default async function validateGroupMembership(
     }
   });
 
-  await Promise.all(checks);
-
   if (allowListOverride) {
     // allow if address is whitelisted
     return { isValid: true };
@@ -92,11 +86,11 @@ export default async function validateGroupMembership(
   return response;
 }
 
-async function _thresholdCheck(
+function _thresholdCheck(
   userAddress: string,
   thresholdData: ThresholdData,
   balances: OptionsWithBalances[],
-): Promise<{ result: boolean; message: string }> {
+): { result: boolean; message: string } {
   try {
     let balanceSourceType: BalanceSourceType;
     let contractAddress: string;
@@ -181,10 +175,10 @@ async function _thresholdCheck(
   }
 }
 
-async function _allowlistCheck(
+function _allowlistCheck(
   userAddress: string,
   allowlistData: AllowlistData,
-): Promise<{ result: boolean; message: string }> {
+): { result: boolean; message: string } {
   try {
     const result = allowlistData.allow.includes(userAddress);
     return {

diff --git a/packages/commonwealth/test/unit/server_controllers/server_groups_controller.spec.ts b/packages/commonwealth/test/unit/server_controllers/server_groups_controller.spec.ts
@@ -99,6 +99,7 @@ const createMockedGroupsController = () => {
       },
       count: async () => memberships.length,
       destroy: async () => {},
+      bulkCreate: async () => {},
     },
     CommunityRole: {
       findAll: async () => [