pubsub: allow aborting an active HTTP request + auto-abort when subscription closed

Author: stephenplusplus

lib/common/util.js

@@ -323,6 +323,7 @@ function makeAuthenticatedRequestFactory(config) {
   function makeAuthenticatedRequest(reqOpts, options) {
     var stream;
     var reqConfig = extend({}, config);
+    var activeRequest_;
 
     if (!options) {
       stream = duplexify();
@@ -345,7 +346,8 @@ function makeAuthenticatedRequestFactory(config) {
       if (options && options.onAuthenticated) {
         options.onAuthenticated(null, authenticatedReqOpts);
       } else {
-        util.makeRequest(authenticatedReqOpts, reqConfig, options);
+        activeRequest_ =
+          util.makeRequest(authenticatedReqOpts, reqConfig, options);
       }
     }
 
@@ -360,6 +362,15 @@ function makeAuthenticatedRequestFactory(config) {
     if (stream) {
       return stream;
     }
+
+    return {
+      abort: function() {
+        if (activeRequest_) {
+          activeRequest_.abort();
+          activeRequest_ = null;
+        }
+      }
+    };
   }
 
   makeAuthenticatedRequest.getCredentials =
@@ -426,7 +437,7 @@ function makeRequest(reqOpts, config, callback) {
 
     dup.abort = requestStream.abort;
   } else {
-    retryRequest(reqOpts, options, function(err, response, body) {
+    return retryRequest(reqOpts, options, function(err, response, body) {
       util.handleResp(err, response, body, callback);
     });
   }

lib/pubsub/subscription.js

@@ -509,7 +509,7 @@ Subscription.prototype.pull = function(options, callback) {
     options.maxResults = MAX_EVENTS_LIMIT;
   }
 
-  this.request({
+  this.activeRequest_ = this.request({
     timeout: this.timeout,
     method: 'POST',
     uri: ':pull',
@@ -518,6 +518,8 @@ Subscription.prototype.pull = function(options, callback) {
       maxMessages: options.maxResults
     }
   }, function(err, response) {
+    self.activeRequest_ = null;
+
     if (err) {
       if (err.code === 'ETIMEDOUT' && !err.connect) {
         // Simulate a server timeout where no messages were received.
@@ -616,6 +618,10 @@ Subscription.prototype.listenForEvents_ = function() {
   this.on('removeListener', function(event) {
     if (event === 'message' && --self.messageListeners === 0) {
       self.closed = true;
+
+      if (self.activeRequest_) {
+        self.activeRequest_.abort();
+      }
     }
   });
 };

package.json

@@ -105,7 +105,7 @@
     "protobufjs": "^5.0.1",
     "pumpify": "^1.3.3",
     "request": "^2.53.0",
-    "retry-request": "^1.2.1",
+    "retry-request": "^1.2.3",
     "split-array-stream": "^1.0.0",
     "stream-events": "^1.0.1",
     "string-format-obj": "^1.0.0",

test/common/util.js

@@ -789,6 +789,35 @@ describe('common/util', function() {
           makeAuthenticatedRequest(reqOpts, done);
         });
 
+        it('should return abort() from the active request', function(done) {
+          var retryRequest = {
+            abort: done
+          };
+
+          utilOverrides.makeRequest = function() {
+            return retryRequest;
+          };
+
+          var makeAuthenticatedRequest = util.makeAuthenticatedRequestFactory();
+          makeAuthenticatedRequest(reqOpts, assert.ifError).abort();
+        });
+
+        it('should only abort() once', function(done) {
+          var retryRequest = {
+            abort: done // Will throw if called more than once.
+          };
+
+          utilOverrides.makeRequest = function() {
+            return retryRequest;
+          };
+
+          var makeAuthenticatedRequest = util.makeAuthenticatedRequestFactory();
+          var request = makeAuthenticatedRequest(reqOpts, assert.ifError);
+
+          request.abort(); // done()
+          request.abort(); // done()
+        });
+
         it('should provide stream to makeRequest', function(done) {
           var stream;
 
@@ -916,6 +945,17 @@ describe('common/util', function() {
         retryRequestOverride = testCustomRetryRequestConfig(done);
         util.makeRequest(reqOpts, customRetryRequestConfig);
       });
+
+      it('should return the instance of retryRequest', function() {
+        var requestInstance = {};
+
+        retryRequestOverride = function() {
+          return requestInstance;
+        };
+
+        var request = util.makeRequest(reqOpts, assert.ifError);
+        assert.strictEqual(request, requestInstance);
+      });
     });
 
     describe('stream mode', function() {

test/pubsub/subscription.js

@@ -497,6 +497,33 @@ describe('Subscription', function() {
       subscription.pull(assert.ifError);
     });
 
+    it('should store the active request', function() {
+      var requestInstance = {};
+
+      subscription.request = function() {
+        return requestInstance;
+      };
+
+      subscription.pull(assert.ifError);
+      assert.strictEqual(subscription.activeRequest_, requestInstance);
+    });
+
+    it('should clear the active request', function(done) {
+      var requestInstance = {};
+
+      subscription.request = function(reqOpts, callback) {
+        setImmediate(function() {
+          callback(null, {});
+          assert.strictEqual(subscription.activeRequest_, null);
+          done();
+        });
+
+        return requestInstance;
+      };
+
+      subscription.pull(assert.ifError);
+    });
+
     it('should pass error to callback', function(done) {
       var error = new Error('Error.');
       subscription.request = function(reqOpts, callback) {
@@ -825,6 +852,17 @@ describe('Subscription', function() {
       // 0 listeners: sub should be closed.
       assert.strictEqual(subscription.closed, true);
     });
+
+    it('should abort the HTTP request when listeners removed', function(done) {
+      subscription.startPulling_ = util.noop;
+
+      subscription.activeRequest_ = {
+        abort: done
+      };
+
+      subscription.on('message', util.noop);
+      subscription.removeAllListeners();
+    });
   });
 
   describe('startPulling_', function() {