Graceful Shutdown [isolated commits]

lib/kue.js

@@ -74,6 +74,7 @@ exports.createQueue = function(){
 
 function Queue() {
   this.client = redis.createClient();
+  this.workers = [];
 }
 
 /**
@@ -179,23 +180,57 @@ Queue.prototype.setting = function(name, fn){
  */
 
 Queue.prototype.process = function(type, n, fn){
-  var self = this;
-
   if ('function' == typeof n) fn = n, n = 1;
 
   while (n--) {
-    (function(worker){
-      worker.on('error', function(err){
-        self.emit('error', err);
-      });
-
-      worker.on('job complete', function(job){
-        self.client.incrby('q:stats:work-time', job.duration);
-      });
-    })(new Worker(this, type).start(fn));
+    var worker = new Worker(this, type).start(fn);
+
+    worker.on('error', function(err){
+      this.emit('error', err);
+    });
+
+    worker.on('job complete', function(job){
+      this.client.incrby('q:stats:work-time', job.duration);
+    });
+
+    // Save worker so we can access it later
+    this.workers.push(worker);
   }
 };
 
+/**
+ * Graceful shutdown
+ *
+ * @param {Function} fn callback
+ * @return {Queue} for chaining
+ * @api public
+ */
+
+Queue.prototype.shutdown = function(fn, timeout) {
+  var origFn = fn || function(){}
+    , n = this.workers.length;
+    , self = this
+
+  // Wrap `fn` to only call after all workers finished
+  fn = function(err) {
+    if (err) return origFn(err);
+    if (! --n) {
+        self.workers = [];
+        origFn.apply(null, arguments);
+    }
+  }
+
+  // Shut down workers 1 by 1
+  this.workers.forEach(function(worker) {
+    worker.shutdown(function(err) {
+      if (err) return fn(err);
+      fn();
+    }, timeout);
+  });
+
+  return this;
+};
+
 /**
  * Get the job types present and callback `fn(err, types)`.
  *

lib/queue/worker.js

@@ -40,6 +40,8 @@ function Worker(queue, type) {
   this.type = type;
   this.client = Worker.client || (Worker.client = redis.createClient());
   this.interval = 1000;
+  this.running = true;
+  this.job = null;
 }
 
 /**
@@ -59,6 +61,7 @@ Worker.prototype.__proto__ = EventEmitter.prototype;
 
 Worker.prototype.start = function(fn){
   var self = this;
+  if (!this.running) return;
   self.getJob(function(err, job){
     if (err) self.error(err, job);
     if (!job || err) return process.nextTick(function(){ self.start(fn); });
@@ -123,7 +126,9 @@ Worker.prototype.process = function(job, fn){
   var self = this
     , start = new Date;
   job.active();
+  this.job = job;
   fn(job, function(err){
+    self.job = null;
     if (err) return self.failed(job, err, fn);
     job.complete();
     job.set('duration', job.duration = new Date - start);
@@ -177,3 +182,33 @@ Worker.prototype.getJob = function(fn){
     });
   });
 };
+
+/**
+ * Gracefully shut down the worker
+ *
+ * @param {Function} fn
+ * @api private
+ */
+
+Worker.prototype.shutdown = function(fn, timeout) {
+  var self = this;
+
+  if (!this.running) return fn();
+  this.running = false;
+
+  // Close redis connection (if zpopping, break it)
+  this.client.end();
+
+  // As soon as we're free, signal that we're done
+  if (!this.job) return fn();
+  this.once('job complete', fn);
+
+  if (timeout) {
+    setTimeout(function() {
+      self.job.failed();
+      self.job = null;
+      fn();
+    }, timeout);
+  }
+};
+