From f171e54ac9c86a34d6df0062639136a5148cb296 Mon Sep 17 00:00:00 2001
From: Yanlong Wang <yanlong.wang@naiver.org>
Date: Tue, 14 May 2024 17:25:59 +0800
Subject: [PATCH] fix: log charge amount

---
 .../functions/src/cloud-functions/crawler.ts  | 20 ++++++++++++++++---
 1 file changed, 17 insertions(+), 3 deletions(-)

diff --git a/backend/functions/src/cloud-functions/crawler.ts b/backend/functions/src/cloud-functions/crawler.ts
index 611dbeb..bba8229 100644
--- a/backend/functions/src/cloud-functions/crawler.ts
+++ b/backend/functions/src/cloud-functions/crawler.ts
@@ -364,14 +364,13 @@ ${authMixin}`,
                 { contentType: 'text/plain', envelope: null }
             );
         }
-
         if (uid) {
             const user = await auth.assertUser();
             if (!(user.wallet.total_balance > 0)) {
                 throw new InsufficientBalanceError(`Account balance not enough to run this query, please recharge.`);
             }
 
-            await this.rateLimitControl.simpleRPCUidBasedLimit(rpcReflect, uid, ['CRAWL'],
+            const apiRoll = await this.rateLimitControl.simpleRPCUidBasedLimit(rpcReflect, uid, ['CRAWL'],
                 [
                     // 1000 requests per minute
                     new Date(Date.now() - 60 * 1000), 1000
@@ -383,15 +382,26 @@ ${authMixin}`,
                     auth.reportUsage(chargeAmount, 'reader-crawl').catch((err) => {
                         this.logger.warn(`Unable to report usage for ${uid}`, { err: marshalErrorLike(err) });
                     });
+                    apiRoll._ref?.set({
+                        chargeAmount,
+                    }, { merge: true }).catch((err) => this.logger.warn(`Failed to log charge amount in apiRoll`, { err }));
                 }
             });
         } else if (ctx.req.ip) {
-            await this.rateLimitControl.simpleRpcIPBasedLimit(rpcReflect, ctx.req.ip, ['CRAWL'],
+            const apiRoll = await this.rateLimitControl.simpleRpcIPBasedLimit(rpcReflect, ctx.req.ip, ['CRAWL'],
                 [
                     // 100 requests per minute
                     new Date(Date.now() - 60 * 1000), 100
                 ]
             );
+
+            rpcReflect.finally(() => {
+                if (chargeAmount) {
+                    apiRoll._ref?.set({
+                        chargeAmount,
+                    }, { merge: true }).catch((err) => this.logger.warn(`Failed to log charge amount in apiRoll`, { err }));
+                }
+            });
         }
 
         let urlToCrawl;
@@ -653,6 +663,10 @@ ${authMixin}`,
     }
 
     getChargeAmount(formatted: { [k: string]: any; }) {
+        if (!formatted) {
+            return undefined;
+        }
+
         const textContent = formatted?.content || formatted?.text || formatted?.html;
 
         if (typeof textContent === 'string') {