Merge pull request #624 from mendableai/nsc/check-credits-optimization

Optimize check credits func w/ caching

.github/workflows/ci.yml

@@ -28,6 +28,7 @@ env:
   HYPERDX_API_KEY: ${{ secrets.HYPERDX_API_KEY }}
   HDX_NODE_BETA_MODE: 1
   FIRE_ENGINE_BETA_URL: ${{ secrets.FIRE_ENGINE_BETA_URL }}
+  USE_DB_AUTHENTICATION: ${{ secrets.USE_DB_AUTHENTICATION }}
 
 
 jobs:

apps/api/src/services/billing/credit_billing.ts

@@ -5,7 +5,7 @@ import { supabase_service } from "../supabase";
 import { Logger } from "../../lib/logger";
 import { getValue, setValue } from "../redis";
 import { redlock } from "../redlock";
-
+import * as Sentry from "@sentry/node";
 
 const FREE_CREDITS = 500;
 
@@ -176,9 +176,24 @@ export async function supaCheckTeamCredits(team_id: string, credits: number) {
     return { success: true, message: "Preview team, no credits used", remainingCredits: Infinity };
   }
 
-  // Retrieve the team's active subscription and check for available coupons concurrently
+
-  const [{ data: subscription, error: subscriptionError }, { data: coupons }] =
+  let cacheKeySubscription = `subscription_${team_id}`;
-    await Promise.all([
+  let cacheKeyCoupons = `coupons_${team_id}`;
+
+  // Try to get data from cache first
+  const [cachedSubscription, cachedCoupons] = await Promise.all([
+    getValue(cacheKeySubscription),
+    getValue(cacheKeyCoupons)
+  ]);
+
+  let subscription, subscriptionError, coupons;
+
+  if (cachedSubscription && cachedCoupons) {
+    subscription = JSON.parse(cachedSubscription);
+    coupons = JSON.parse(cachedCoupons);
+  } else {
+    // If not in cache, retrieve from database
+    const [subscriptionResult, couponsResult] = await Promise.all([
       supabase_service
         .from("subscriptions")
         .select("id, price_id, current_period_start, current_period_end")
@@ -192,6 +207,18 @@ export async function supaCheckTeamCredits(team_id: string, credits: number) {
         .eq("status", "active"),
     ]);
 
+    subscription = subscriptionResult.data;
+    subscriptionError = subscriptionResult.error;
+    coupons = couponsResult.data;
+
+    // Cache the results for a minute, sub can be null and that's fine
+    await setValue(cacheKeySubscription, JSON.stringify(subscription), 60); // Cache for 1 minute, even if null
+    
+    if (coupons) {
+      await setValue(cacheKeyCoupons, JSON.stringify(coupons), 60); // Cache for 1 minute
+    }
+  }
+
   let couponCredits = 0;
   if (coupons && coupons.length > 0) {
     couponCredits = coupons.reduce(
@@ -212,41 +239,54 @@ export async function supaCheckTeamCredits(team_id: string, credits: number) {
 
     let creditUsages;
     let creditUsageError;
-    let retries = 0;
+    let totalCreditsUsed = 0;
-    const maxRetries = 3;
+    const cacheKeyCreditUsage = `credit_usage_${team_id}`;
-    const retryInterval = 2000; // 2 seconds
 
-    while (retries < maxRetries) {
+    // Try to get credit usage from cache
-      const result = await supabase_service
+    const cachedCreditUsage = await getValue(cacheKeyCreditUsage);
-        .from("credit_usage")
-        .select("credits_used")
-        .is("subscription_id", null)
-        .eq("team_id", team_id);
 
-      creditUsages = result.data;
+    if (cachedCreditUsage) {
-      creditUsageError = result.error;
+      totalCreditsUsed = parseInt(cachedCreditUsage);
+    } else {
+      let retries = 0;
+      const maxRetries = 3;
+      const retryInterval = 2000; // 2 seconds
 
-      if (!creditUsageError) {
+      while (retries < maxRetries) {
-        break;
+        const result = await supabase_service
+          .from("credit_usage")
+          .select("credits_used")
+          .is("subscription_id", null)
+          .eq("team_id", team_id);
+
+        creditUsages = result.data;
+        creditUsageError = result.error;
+
+        if (!creditUsageError) {
+          break;
+        }
+
+        retries++;
+        if (retries < maxRetries) {
+          await new Promise(resolve => setTimeout(resolve, retryInterval));
+        }
       }
 
-      retries++;
+      if (creditUsageError) {
-      if (retries < maxRetries) {
+        Logger.error(`Credit usage error after ${maxRetries} attempts: ${creditUsageError}`);
-        await new Promise(resolve => setTimeout(resolve, retryInterval));
+        throw new Error(
+          `Failed to retrieve credit usage for team_id: ${team_id}`
+        );
       }
-    }
 
-    if (creditUsageError) {
+      totalCreditsUsed = creditUsages.reduce(
-      Logger.error(`Credit usage error after ${maxRetries} attempts: ${creditUsageError}`);
+        (acc, usage) => acc + usage.credits_used,
-      throw new Error(
+        0
-        `Failed to retrieve credit usage for team_id: ${team_id}`
       );
-    }
 
-    const totalCreditsUsed = creditUsages.reduce(
+      // Cache the result for 30 seconds
-      (acc, usage) => acc + usage.credits_used,
+      await setValue(cacheKeyCreditUsage, totalCreditsUsed.toString(), 30);
-      0
+    }
-    );
 
     Logger.info(`totalCreditsUsed: ${totalCreditsUsed}`);
 
@@ -312,7 +352,7 @@ export async function supaCheckTeamCredits(team_id: string, credits: number) {
 
         if (creditUsages && creditUsages.length > 0) {
           totalCreditsUsed = creditUsages[0].total_credits_used;
-          await setValue(cacheKey, totalCreditsUsed.toString(), 1800); // Cache for 30 minutes
+          await setValue(cacheKey, totalCreditsUsed.toString(), 500); // Cache for 8 minutes
           // Logger.info(`Cache set for credit usage: ${totalCreditsUsed}`);
         }
       }
@@ -325,17 +365,38 @@ export async function supaCheckTeamCredits(team_id: string, credits: number) {
 
   // Adjust total credits used by subtracting coupon value
   const adjustedCreditsUsed = Math.max(0, totalCreditsUsed - couponCredits);
-  // Get the price details
-  const { data: price, error: priceError } = await supabase_service
-    .from("prices")
-    .select("credits")
-    .eq("id", subscription.price_id)
-    .single();
 
-  if (priceError) {
+  // Get the price details from cache or database
-    throw new Error(
+  const priceCacheKey = `price_${subscription.price_id}`;
-      `Failed to retrieve price for price_id: ${subscription.price_id}`
+  let price;
-    );
+
+  try {
+    const cachedPrice = await getValue(priceCacheKey);
+    if (cachedPrice) {
+      price = JSON.parse(cachedPrice);
+    } else {
+      const { data, error: priceError } = await supabase_service
+        .from("prices")
+        .select("credits")
+        .eq("id", subscription.price_id)
+        .single();
+
+      if (priceError) {
+        throw new Error(
+          `Failed to retrieve price for price_id: ${subscription.price_id}`
+        );
+      }
+
+      price = data;
+      // There are only 21 records, so this is super fine
+      // Cache the price for a long time (e.g., 1 day)
+      await setValue(priceCacheKey, JSON.stringify(price), 86400);
+    }
+  } catch (error) {
+    Logger.error(`Error retrieving or caching price: ${error}`);
+    Sentry.captureException(error);
+    // If errors, just assume it's a big number so user don't get an error
+    price = { credits: 1000000 };
   }
 
   const creditLimit = price.credits;

apps/api/src/services/queue-jobs.ts

@@ -67,6 +67,6 @@ export function waitForJob(jobId: string, timeout: number) {
           reject((await getScrapeQueue().getJob(jobId)).failedReason);
         }
       }
-    }, 1000);
+    }, 500);
   })
 }