From 322be961934bf7be7d61464c226bdbe90171878b Mon Sep 17 00:00:00 2001
From: Tracy <31976538+tracy-codes@users.noreply.github.com>
Date: Fri, 21 Apr 2023 11:08:45 -0700
Subject: [PATCH] Add rate limit with Bottleneck

---
 package.json         |  3 ++-
 src/helpers/index.ts | 49 ++++++++++++++++++++++++++------------------
 yarn.lock            | 13 ++++++++----
 3 files changed, 40 insertions(+), 25 deletions(-)

diff --git a/package.json b/package.json
index fe15776..03379bd 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
     "name": "@shadow-drive/cli",
-    "version": "0.4.2",
+    "version": "0.4.3",
     "main": "shdw-drive.js",
     "bin": {
         "shdw-drive": "./dist/shdw-drive.js"
@@ -18,6 +18,7 @@
         "@solana/spl-token": "^0.2.0",
         "@solana/web3.js": "^1.41.0",
         "bigint-conversion": "^2.2.1",
+        "bottleneck": "^2.19.5",
         "bs58": "^5.0.0",
         "cli-progress": "^3.11.2",
         "commander": "^9.2.0",
diff --git a/src/helpers/index.ts b/src/helpers/index.ts
index 8a8ac1a..77f278c 100644
--- a/src/helpers/index.ts
+++ b/src/helpers/index.ts
@@ -16,6 +16,7 @@ import {
 import { Program } from "@project-serum/anchor";
 import { SHDW_DRIVE_ENDPOINT, programId } from "../constants";
 import fetch from "node-fetch";
+import Bottleneck from "bottleneck";
 
 export function loadWalletKey(keypair: string): Keypair {
     if (!keypair || keypair == "") {
@@ -117,6 +118,10 @@ export async function getFormattedStorageAccounts(
     key: anchor.web3.PublicKey,
     totalAccounts: number
 ): Promise<[Array<any>, Array<anchor.web3.PublicKey>]> {
+    const limiter = new Bottleneck({
+        minTime: 50,
+        maxConcurrent: 10,
+    });
     let accountsToFetch: anchor.web3.PublicKey[] = [];
 
     for (let i = 0; i <= totalAccounts; i++) {
@@ -136,28 +141,32 @@ export async function getFormattedStorageAccounts(
 
     await Promise.all(
         accountsToFetch.map(async (account) => {
-            const storageAccountDetails = await fetch(
-                `${SHDW_DRIVE_ENDPOINT}/storage-account-info`,
-                {
-                    method: "POST",
-                    headers: {
-                        "Content-Type": "application/json",
-                    },
-                    body: JSON.stringify({
-                        storage_account: account.toString(),
-                    }),
+            try {
+                const storageAccountDetails = await limiter.schedule(() =>
+                    fetch(`${SHDW_DRIVE_ENDPOINT}/storage-account-info`, {
+                        method: "POST",
+                        headers: {
+                            "Content-Type": "application/json",
+                        },
+                        body: JSON.stringify({
+                            storage_account: account.toString(),
+                        }),
+                    })
+                );
+                // TODO proper handling of this fetch request
+                const storageAccountDetailsJson =
+                    await storageAccountDetails.json();
+                if (
+                    storageAccountDetailsJson.identifier !== null &&
+                    typeof storageAccountDetailsJson.identifier !== "undefined"
+                ) {
+                    accounts.push(storageAccountDetailsJson);
                 }
-            );
-            // TODO proper handling of this fetch request
-            const storageAccountDetailsJson =
-                await storageAccountDetails.json();
-            if (
-                storageAccountDetailsJson.identifier !== null &&
-                typeof storageAccountDetailsJson.identifier !== "undefined"
-            ) {
-                accounts.push(storageAccountDetailsJson);
+                return storageAccountDetailsJson;
+            } catch (e) {
+                log.error(e);
+                return null;
             }
-            return storageAccountDetailsJson;
         })
     );
     let alist1 = accounts.map((account: any, idx: number) => {
diff --git a/yarn.lock b/yarn.lock
index 5175292..f9c6388 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -114,10 +114,10 @@
     bn.js "^5.1.2"
     buffer-layout "^1.2.0"
 
-"@shadow-drive/sdk@^3.0.10":
-  version "3.0.10"
-  resolved "https://registry.yarnpkg.com/@shadow-drive/sdk/-/sdk-3.0.10.tgz#affb6e379de3be84dcd0465480e3841310549568"
-  integrity sha512-OYyh3DTn83Y2wN8Cx/fqOVkoZyz/CVj52nTJ7Gdi+QRV+WxzTj259sz6K7iJMRJ5s0OlKD6nSz1mnc19gwCcJQ==
+"@shadow-drive/sdk@^3.1.1":
+  version "3.3.0"
+  resolved "https://registry.yarnpkg.com/@shadow-drive/sdk/-/sdk-3.3.0.tgz#c506c6ca7d22fb6a3dbb334df0d8314a9ce392e3"
+  integrity sha512-lgL2kG0xFnWJ19YyNVI0Rd7ke1Kdtlxj8NDUNoNxyoqPa7RqAXb4wVrrgzc17iULqNOoSAAbR0o3YTRYnBDltA==
   dependencies:
     "@project-serum/anchor" "^0.24.2"
     "@solana/spl-token" "^0.2.0"
@@ -416,6 +416,11 @@ borsh@^0.7.0:
     bs58 "^4.0.0"
     text-encoding-utf-8 "^1.0.2"
 
+bottleneck@^2.19.5:
+  version "2.19.5"
+  resolved "https://registry.yarnpkg.com/bottleneck/-/bottleneck-2.19.5.tgz#5df0b90f59fd47656ebe63c78a98419205cadd91"
+  integrity sha512-VHiNCbI1lKdl44tGrhNfU3lup0Tj/ZBMJB5/2ZbNXRCPuRCO7ed2mgcK4r17y+KB2EfuYuRaVlwNbAeaWGSpbw==
+
 braces@^3.0.2:
   version "3.0.2"
   resolved "https://registry.npmjs.org/braces/-/braces-3.0.2.tgz"