From 072602b8e410b81a8c467aa575e316a55aa40124 Mon Sep 17 00:00:00 2001
From: Dave Gramlich <callmehiphop@gmail.com>
Date: Thu, 28 Apr 2016 13:13:26 -0400
Subject: [PATCH] core: implement retry strategy for grpc requests

---
 lib/common/grpc-service.js  | 42 +++++++++++++++++-------
 test/common/grpc-service.js | 64 +++++++++++++++++++++++++++++++++++++
 2 files changed, 94 insertions(+), 12 deletions(-)

diff --git a/lib/common/grpc-service.js b/lib/common/grpc-service.js
index 48e48b7b4e0..43bf715646b 100644
--- a/lib/common/grpc-service.js
+++ b/lib/common/grpc-service.js
@@ -25,6 +25,7 @@ var grpc = require('grpc');
 var is = require('is');
 var nodeutil = require('util');
 var path = require('path');
+var retryRequest = require('retry-request');
 
 /**
  * @type {module:common/service}
@@ -109,7 +110,8 @@ var HTTP_ERROR_CODE_MAP = {
 
   14: {
     code: 503,
-    message: 'Service Unavailable'
+    message: 'Service Unavailable',
+    shouldRetry: true
   },
 
   15: {
@@ -123,6 +125,12 @@ var HTTP_ERROR_CODE_MAP = {
   }
 };
 
+/**
+ * @const {number} - Allowed number of a request retries
+ * @private
+ */
+var MAX_RETRIES = 2;
+
 /**
  * Service is a base class, meant to be inherited from by a "service," like
  * BigQuery or Storage.
@@ -241,19 +249,29 @@ GrpcService.prototype.request = function(protoOpts, reqOpts, callback) {
     grpcOpts.deadline = new Date(Date.now() + protoOpts.timeout);
   }
 
-  service[protoOpts.method](reqOpts, function(err, resp) {
-    if (err) {
-      if (HTTP_ERROR_CODE_MAP[err.code]) {
-        var httpError = HTTP_ERROR_CODE_MAP[err.code];
-        err.code = httpError.code;
-      }
+  var attempts = 0;
 
-      callback(err);
-      return;
-    }
+  (function makeRequest() {
+    service[protoOpts.method](reqOpts, function(err, resp) {
+      if (err) {
+        if (HTTP_ERROR_CODE_MAP[err.code]) {
+          var httpError = HTTP_ERROR_CODE_MAP[err.code];
+
+          if (httpError.shouldRetry && attempts++ < MAX_RETRIES) {
+            setTimeout(makeRequest, retryRequest.getNextRetryDelay(attempts));
+            return;
+          }
+
+          err.code = httpError.code;
+        }
+
+        callback(err);
+        return;
+      }
 
-    callback(null, resp);
-  }, null, grpcOpts);
+      callback(null, resp);
+    }, null, grpcOpts);
+  }());
 };
 
 /**
diff --git a/test/common/grpc-service.js b/test/common/grpc-service.js
index bb7676692f3..edc07e4e5e3 100644
--- a/test/common/grpc-service.js
+++ b/test/common/grpc-service.js
@@ -550,6 +550,70 @@ describe('GrpcService', function() {
         });
       });
     });
+
+    describe('retrying requests', function() {
+      var UNAVAILABLE = 14;
+      var _setTimeout;
+
+      before(function() {
+        _setTimeout = global.setTimeout;
+        global.setTimeout = function(func) {
+          func();
+        };
+      });
+
+      after(function() {
+        global.setTimeout = _setTimeout;
+      });
+
+      it('should retry if the service is unavailable', function(done) {
+        var callCount = 0;
+
+        grpcService.protos.Service = {
+          service: function() {
+            return {
+              method: function(reqOpts, callback) {
+                var err = null;
+
+                if (++callCount < 2) {
+                  err = { code: UNAVAILABLE };
+                }
+
+                callback(err);
+              }
+            };
+          }
+        };
+
+        grpcService.request(PROTO_OPTS, REQ_OPTS, function(err) {
+          assert.ifError(err);
+          assert.strictEqual(callCount, 2);
+          done();
+        });
+      });
+
+      it('should retry a maximum of 2 times before failing', function(done) {
+        var callCount = 0;
+
+        grpcService.protos.Service = {
+          service: function() {
+            return {
+              method: function(reqOpts, callback) {
+                callCount += 1;
+                callback({ code: UNAVAILABLE });
+              }
+            };
+          }
+        };
+
+        grpcService.request(PROTO_OPTS, REQ_OPTS, function(err) {
+          assert.strictEqual(err.code, 503);
+          // 1 for the original request + 2 retries
+          assert.strictEqual(callCount, 3);
+          done();
+        });
+      });
+    });
   });
 
   describe('convertValue_', function() {