From 1d15b5580f7b5b7bea63fa8efab43c07f8004a81 Mon Sep 17 00:00:00 2001
From: Jeff McAffer <jeff.mcaffer@microsoft.com>
Date: Fri, 6 Jan 2017 23:34:03 -0800
Subject: [PATCH] handle benched tokens

---
 lib/githubFetcher.js | 82 +++++++++++++++++++++++++++-----------------
 1 file changed, 51 insertions(+), 31 deletions(-)

diff --git a/lib/githubFetcher.js b/lib/githubFetcher.js
index 964eb01..37c7ec1 100644
--- a/lib/githubFetcher.js
+++ b/lib/githubFetcher.js
@@ -36,7 +36,14 @@ class GitHubFetcher {
     const self = this;
     const etagPromise = checkEtag ? this.store.etag(request.type, request.url) : Q(null);
     return etagPromise.then(etag => {
-      return this._getFromGitHub(request, etag).then(githubResponse => {
+      let options = self._addTokenOption(request, { headers: {} });
+      if (typeof options === 'number') {
+        // if we get back a number, all tokens have been benched so we have to requeue and wait.
+        return self._requeueBenched(request, options);
+      }
+      options = self._addEtagOption(options, etag);
+      options = self._addTypeOptions(request, options);
+      return this._getFromGitHub(request, options).then(githubResponse => {
         const status = githubResponse.statusCode;
         if (status !== 200 && status !== 304) {
           if (status === 409 || status === 204) {
@@ -46,11 +53,7 @@ class GitHubFetcher {
           // and wait a bit before processing more requests
           const remaining = parseInt(githubResponse.headers['x-ratelimit-remaining'], 10) || 0;
           if (status === 403 && remaining === 0) {
-            const delay = self.options.forbiddenDelay || 120000;
-            request.exhaustToken(Date.now() + delay);
-            request.delay(delay);
-            request.addMeta({ forbiddenDelay: delay });
-            return request.markRequeue('Throttled', `GitHub throttled ${request.url}`);
+            return self._requeueThrottled(request);
           }
           throw new Error(`Code ${status} for ${request.url}`);
         }
@@ -74,11 +77,7 @@ class GitHubFetcher {
     });
   }
 
-  _getFromGitHub(request, etag) {
-    const options = this._getOptions(request);
-    if (etag) {
-      options.headers['If-None-Match'] = etag;
-    }
+  _getFromGitHub(request, options) {
     const start = Date.now();
     const deferred = Q.defer();
     this.fetchQueue.push({ url: request.url, options: options }, (error, response) => {
@@ -100,13 +99,29 @@ class GitHubFetcher {
     try {
       this._incrementMetric('fetch');
       this.requestor.get(spec.url, spec.options).then(
-        response => callback(null, response),
-        error => callback(error));
+        response =>
+          callback(null, response),
+        error =>
+          callback(error));
     } catch (e) {
       callback(e);
     }
   }
 
+  _requeueBenched(request, benchTime) {
+    request.delayUntil(benchTime);
+    request.addMeta({ benchDelay: benchTime - Date.now() });
+    return request.markRequeue('Benched', `Wait for token while getting ${request.url}`);
+  }
+
+  _requeueThrottled(request) {
+    const delay = this.options.forbiddenDelay || 120000;
+    request.exhaustToken(Date.now() + delay);
+    request.delay(delay);
+    request.addMeta({ forbiddenDelay: delay });
+    return request.markRequeue('Throttled', `GitHub throttled ${request.url}`);
+  }
+
   _fetchFromStorage(request) {
     const start = Date.now();
     return this.store.get(request.type, request.url).then(
@@ -138,7 +153,7 @@ class GitHubFetcher {
     // This code is not designed to handle the 403 scenarios.  That is handled by the retry logic.
     const remaining = parseInt(response.headers['x-ratelimit-remaining'], 10) || 0;
     request.addMeta({ remaining: remaining });
-    const tokenLowerBound = this.options ? (this.options.tokenLowerBound || 50) : 50;
+    const tokenLowerBound = this.options.tokenLowerBound || 50;
     if (remaining < tokenLowerBound) {
       const reset = parseInt(response.headers['x-ratelimit-reset'], 10) || 0;
       const delay = Math.max(0, reset - Date.now());
@@ -170,23 +185,15 @@ class GitHubFetcher {
     return request;
   }
 
-  _getOptions(request) {
-    const result = { headers: {} };
-    const token = this._getToken(request);
-    if (token) {
-      result.headers.authorization = `token ${token}`;
-      request.exhaustToken = until => this.tokenFactory.exhaust(token, until);
-    } else {
-      throw new Error(`No API tokens available for ${request.toString()}`);
+
+  _addEtagOption(options, etag) {
+    if (etag) {
+      options.headers['If-None-Match'] = etag;
     }
-    const headers = this._getHeaders(request);
-    if (headers) {
-      Object.assign(result.headers, headers);
-    }
-    return result;
+    return options;
   }
 
-  _getToken(request) {
+  _addTokenOption(request, options) {
     const traits = this._getTypeDetails(request.type).tokenTraits || [];
     const additionalTraits = [];
     if (request.context.repoType) {
@@ -196,12 +203,25 @@ class GitHubFetcher {
       // if this is a retry, elevate the token to avoid any permissions issues
       additionalTraits.push('private', 'admin');
     }
-    return this.tokenFactory.getToken(traits.concat(additionalTraits));
+    const token = this.tokenFactory.getToken(traits.concat(additionalTraits));
+    if (!token) {
+      throw new Error(`No API tokens available for ${request.toString()}`);
+    }
+    if (typeof token === 'number') {
+      return token;
+    }
+    options.headers.authorization = `token ${token}`;
+    request.exhaustToken = until => this.tokenFactory.exhaust(token, until);
+    return options;
   }
 
-  _getHeaders(request) {
+  _addTypeOptions(request, options) {
     const typeDetails = this._getTypeDetails(request.type);
-    return typeDetails.headers;
+    const headers = typeDetails.headers;
+    if (headers) {
+      Object.assign(options.headers, headers);
+    }
+    return options;
   }
 
   _incrementMetric(operation) {