Merge pull request #1040 from amvanbaren/load-extensions-rate-limiter

cli: rate limit requests in load-test-extensions
eclipse · Nov 14, 2024 · 7b22dd3 · 7b22dd3
2 parents 8004ec8 + 760d45e
commit 7b22dd3
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 16 deletions.
diff --git a/cli/package.json b/cli/package.json
@@ -53,6 +53,7 @@
         "@typescript-eslint/eslint-plugin": "^5.44.0",
         "@typescript-eslint/parser": "^5.44.0",
         "eslint": "^8.28.0",
+        "limiter": "^2.1.0",
         "rimraf": "^3.0.2",
         "typescript": "^4.3.2"
     },

diff --git a/cli/scripts/load-test-extensions.js b/cli/scripts/load-test-extensions.js
@@ -13,6 +13,7 @@ const { Registry, DEFAULT_URL } = require('../lib/registry');
 const { makeDirs } = require('../lib/util');
 const path = require('path');
 const fs = require('fs');
+const { RateLimiter } = require('limiter');
 
 let searchSize = 100;
 if (process.argv.length >= 3) {
@@ -26,40 +27,46 @@ if (process.argv.length >= 4) {
 
 async function loadTestExtensions() {
     const publicReg = new Registry();
-    const localReg = new Registry({registryUrl: process.env.OVSX_REGISTRY_URL ?? 'http://localhost:8080'});
+    const publicLimiter = new RateLimiter({ tokensPerInterval: 15, interval: 'second' });
+    const localReg = new Registry({ registryUrl: process.env.OVSX_REGISTRY_URL ?? 'http://localhost:8080' });
+    const localLimiter = new RateLimiter({ tokensPerInterval: 15, interval: 'second' });
     /** @type {{ extensions: import('../lib/registry').Extension[] } & import('../lib/registry').Response} */
+    await publicLimiter.removeTokens(1);
     const search = await publicReg.getJson(new URL(`${DEFAULT_URL}/api/-/search?size=${searchSize}`));
     if (search.error) {
         console.error(search.error);
         process.exit(1);
     }
     console.log(`Found ${search.extensions.length} extensions in ${DEFAULT_URL}`);
     for (const ext of search.extensions) {
+        await publicLimiter.removeTokens(1);
         const meta = await publicReg.getMetadata(ext.namespace, ext.name);
         if (meta.error) {
             console.error(`\u274c  ${meta.error}`);
             continue;
         }
-        const fileName = await download(publicReg, meta);
+        const fileName = await download(publicReg, meta, publicLimiter);
         try {
-          const nsResult = await localReg.createNamespace(meta.namespace, accessToken);
-          console.log(nsResult.success);
+            await localLimiter.removeTokens(1);
+            const nsResult = await localReg.createNamespace(meta.namespace, accessToken);
+            console.log(nsResult.success);
         } catch (error) {
-          if (!error.message.startsWith('Namespace already exists')) {
-            console.error(error);
-            process.exit(1);
-          }
+            if (!error.message.startsWith('Namespace already exists')) {
+                console.error(error);
+                process.exit(1);
+            }
         }
 
         try {
-          const published = await localReg.publish(fileName, accessToken);
-          if (published.namespace && published.name) {
-            console.log(`\u2713  Published ${published.namespace}.${published.name}@${published.version}`);
-          }
+            await localLimiter.removeTokens(1);
+            const published = await localReg.publish(fileName, accessToken);
+            if (published.namespace && published.name) {
+                console.log(`\u2713  Published ${published.namespace}.${published.name}@${published.version}`);
+            }
         } catch (error) {
-          if (!error.message.endsWith('is already published.')) {
-              console.error(`\u274c  ${error}`);
-          }
+            if (!error.message.endsWith('is already published.')) {
+                console.error(`\u274c  ${error}`);
+            }
         }
     }
 }
@@ -68,7 +75,7 @@ async function loadTestExtensions() {
  * @param {Registry} registry 
  * @param {import('../lib/registry').Extension} extension 
  */
-async function download(registry, extension) {
+async function download(registry, extension, rateLimiter) {
     const downloadUrl = extension.files.download;
     const fileNameIndex = downloadUrl.lastIndexOf('/');
     const fileName = downloadUrl.substring(fileNameIndex + 1);
@@ -78,6 +85,7 @@ async function download(registry, extension) {
     }
     await makeDirs(path.dirname(filePath));
     console.log(`Downloading ${extension.namespace}.${extension.name}@${extension.version} to ${filePath}`);
+    await rateLimiter.removeTokens(1);
     await registry.download(filePath, new URL(downloadUrl));
     return filePath;
 }

diff --git a/cli/yarn.lock b/cli/yarn.lock
@@ -2068,6 +2068,13 @@ __metadata:
   languageName: node
   linkType: hard
 
+"just-performance@npm:4.3.0":
+  version: 4.3.0
+  resolution: "just-performance@npm:4.3.0"
+  checksum: 10/5e5f0664f357d291aac9ccca6357b997254600f2d2dd429e9b63e058035c44b6b1c73dbae51f58af7918c0701db4b20d4f365ed61e77d87c92a0183efb77a72a
+  languageName: node
+  linkType: hard
+
 "jwa@npm:^1.4.1":
   version: 1.4.1
   resolution: "jwa@npm:1.4.1"
@@ -2147,6 +2154,15 @@ __metadata:
   languageName: node
   linkType: hard
 
+"limiter@npm:^2.1.0":
+  version: 2.1.0
+  resolution: "limiter@npm:2.1.0"
+  dependencies:
+    just-performance: "npm:4.3.0"
+  checksum: 10/d6292c2a576fe7c4664652a64d286cb3701dd18e1878e01207840d809c62ce6db46ec6ee5e2ab53d548cb56152d0aeeb8555e32fe9177fcb4b04983ef04f805e
+  languageName: node
+  linkType: hard
+
 "linkify-it@npm:^5.0.0":
   version: 5.0.0
   resolution: "linkify-it@npm:5.0.0"
@@ -2642,6 +2658,7 @@ __metadata:
     follow-redirects: "npm:^1.14.6"
     is-ci: "npm:^2.0.0"
     leven: "npm:^3.1.0"
+    limiter: "npm:^2.1.0"
     rimraf: "npm:^3.0.2"
     semver: "npm:^7.6.0"
     tmp: "npm:^0.2.3"