Merge pull request #47 from JupiterOne/fix-rate-limit-throttling

Fix rate limit throttling

CHANGELOG.md

@@ -8,6 +8,15 @@ and this project adheres to
 
 ## [Unreleased]
 
+## [2.0.5] - 2021-10-07
+
+### Fixed
+
+- Fixed retry logic after encountering 429 rate limit errors. Previously, the
+  `x-ratelimit-retryafter` header was not properly respected because the header
+  returned an epoch time in seconds, and we compared this to the current epoch
+  time in milliseconds.
+
 ## [2.0.4] - 2021-08-30
 
 ### Removed

package.json

@@ -1,6 +1,6 @@
 {
   "name": "@jupiterone/graph-crowdstrike",
-  "version": "2.0.4",
+  "version": "2.0.5",
   "description": "A template for JupiterOne graph converters.",
   "main": "src/index.js",
   "types": "src/index.d.ts",
@@ -28,7 +28,6 @@
   "dependencies": {
     "@jupiterone/integration-sdk-core": "^6.16.1",
     "@jupiterone/integration-sdk-runtime": "^6.16.1",
-    "await-timeout": "^1.0.1",
     "node-fetch": "^2.6.0"
   },
   "devDependencies": {

src/crowdstrike/FalconAPIClient.test.ts

@@ -127,9 +127,9 @@ describe('executeAPIRequest', () => {
       name: 'executeAPIRequest429',
     });
 
-    const requestTimes: number[] = [];
+    const requestTimesInMs: number[] = [];
     recording.server.any().on('request', (_req, _event) => {
-      requestTimes.push(Date.now());
+      requestTimesInMs.push(Date.now());
     });
 
     recording.server
@@ -141,7 +141,7 @@ describe('executeAPIRequest', () => {
 
     await createClient().authenticate();
 
-    expect(requestTimes.length).toBe(2);
+    expect(requestTimesInMs.length).toBe(2);
   });
 
   test('waits until retryafter on 429 response', async () => {
@@ -150,20 +150,20 @@ describe('executeAPIRequest', () => {
       name: 'executeAPIRequest429',
     });
 
-    const requestTimes: number[] = [];
+    const requestTimesInMs: number[] = [];
     recording.server.any().on('request', (_req, _event) => {
-      requestTimes.push(Date.now());
+      requestTimesInMs.push(Date.now());
     });
 
-    const retryAfter = Date.now() + 1000;
+    const retryAfterTimeInSeconds = Date.now() / 1000 + 1; // server responds with epoch time in seconds, Date.now() returns epoch time in ms
     recording.server
       .any()
       .times(1)
       .intercept((_req, res) => {
         res
           .status(429)
           .setHeaders({
-            'x-ratelimit-retryafter': String(retryAfter),
+            'x-ratelimit-retryafter': String(retryAfterTimeInSeconds),
           })
           .json({
             errors: [
@@ -177,8 +177,8 @@ describe('executeAPIRequest', () => {
 
     await createClient().authenticate();
 
-    expect(requestTimes.length).toBe(2);
-    expect(requestTimes[1]).toBeGreaterThan(retryAfter);
+    expect(requestTimesInMs.length).toBe(2);
+    expect(requestTimesInMs[1]).toBeGreaterThan(retryAfterTimeInSeconds * 1000);
   });
 
   test('retries 429 response limited times', async () => {
@@ -187,16 +187,17 @@ describe('executeAPIRequest', () => {
       name: 'executeAPIRequest429limit',
     });
 
-    const requestTimes: number[] = [];
+    const requestTimesInMs: number[] = [];
     recording.server.any().on('request', (_req, _event) => {
-      requestTimes.push(Date.now());
+      requestTimesInMs.push(Date.now());
     });
 
+    const retryAfterTimeInSeconds = Date.now() / 1000 - 10; // server responds with epoch time in seconds, Date.now() returns epoch time in ms
     recording.server.any().intercept((_req, res) => {
       res
         .status(429)
         .setHeaders({
-          'x-ratelimit-retryafter': String(Date.now() - 10),
+          'x-ratelimit-retryafter': String(retryAfterTimeInSeconds),
         })
         .json({
           errors: [
@@ -218,7 +219,7 @@ describe('executeAPIRequest', () => {
 
     await expect(client.authenticate()).rejects.toThrowError(/2/);
 
-    expect(requestTimes.length).toBe(2);
+    expect(requestTimesInMs.length).toBe(2);
   });
 
   test('throttles at specified reserveLimit', async () => {

src/crowdstrike/FalconAPIClient.ts

@@ -1,4 +1,3 @@
-import Timeout from 'await-timeout';
 import fetch, { RequestInfo, RequestInit, Response } from 'node-fetch';
 import { URLSearchParams } from 'url';
 
@@ -18,6 +17,10 @@ import {
 } from './types';
 import { IntegrationLogger } from '@jupiterone/integration-sdk-core';
 
+async function sleep(ms) {
+  return new Promise((resolve) => setTimeout(resolve, ms));
+}
+
 type APIRequest = {
   url: string;
   exec: () => Promise<Response>;
@@ -256,20 +259,24 @@ export class FalconAPIClient {
   private async executeAPIRequestWithRateLimitRetries<T>(
     request: APIRequest,
   ): Promise<APIResponse> {
-    const config = request.rateLimitConfig;
-
     let attempts = 0;
     let rateLimitState = request.rateLimitState;
 
     do {
-      const tryAfterCooldown =
-        rateLimitState.limitRemaining <= config.reserveLimit
-          ? Date.now() + config.cooldownPeriod
-          : 0;
-
-      const tryAfter = Math.max(rateLimitState.retryAfter, tryAfterCooldown);
+      if (
+        rateLimitState.limitRemaining <= request.rateLimitConfig.reserveLimit
+      ) {
+        this.logger.info(
+          {
+            rateLimitState,
+            rateLimitConfig: request.rateLimitConfig,
+          },
+          'Rate limit remaining is less than reserve limit. Waiting for cooldown period.',
+        );
+        await sleep(request.rateLimitConfig.cooldownPeriod);
+      }
 
-      const response = await tryAPIRequest(request.exec, tryAfter);
+      const response = await request.exec();
 
       rateLimitState = {
         limitRemaining:
@@ -292,6 +299,21 @@ export class FalconAPIClient {
         };
       }
 
+      if (response.status === 429) {
+        const epochTimeNow = Date.now() / 1000;
+        const timeToSleepInSeconds = rateLimitState.retryAfter - epochTimeNow;
+        this.logger.info(
+          {
+            epochTimeNow,
+            timeToSleepInSeconds,
+            rateLimitState,
+            rateLimitConfig: request.rateLimitConfig,
+          },
+          'Encountered 429 response. Waiting to retry request.',
+        );
+        await sleep(timeToSleepInSeconds * 1000);
+      }
+
       attempts += 1;
       this.logger.warn(
         {
@@ -308,17 +330,6 @@ export class FalconAPIClient {
   }
 }
 
-async function tryAPIRequest(
-  request: () => Promise<Response>,
-  tryAfter: number,
-): Promise<Response> {
-  const now = Date.now();
-  if (tryAfter > now) {
-    await Timeout.set(tryAfter - now);
-  }
-  return request();
-}
-
 function isValidToken(token: OAuth2Token): boolean {
   return !!(token && token.expiresAt > Date.now());
 }

yarn.lock

@@ -1223,11 +1223,6 @@ atob@^2.1.2:
   resolved "https://registry.yarnpkg.com/atob/-/atob-2.1.2.tgz#6d9517eb9e030d2436666651e86bd9f6f13533c9"
   integrity sha512-Wm6ukoaOGJi/73p/cl2GvLjTI5JM1k/O14isD73YML8StrH/7/lRFgmg8nICZgD3bZZvjwCGxtMOD3wWNAu8cg==
 
-await-timeout@^1.0.1:
-  version "1.1.1"
-  resolved "https://registry.yarnpkg.com/await-timeout/-/await-timeout-1.1.1.tgz#d42062ee6bc4eb271fe4d4f851eb658dae7e3906"
-  integrity sha512-gsDXAS6XVc4Jt+7S92MPX6Noq69bdeXUPEaXd8dk3+yVr629LTDLxNt4j1ycBbrU+AStK2PhKIyNIM+xzWMVOQ==
-
 aws-sdk@^2.184.0:
   version "2.976.0"
   resolved "https://registry.yarnpkg.com/aws-sdk/-/aws-sdk-2.976.0.tgz#1e9d0d359698876eaa952c7c6223eac4c354ae2e"