In-memory reaper. TtlSet needs tests.

lib/units/reaper/index.js

@@ -4,18 +4,43 @@ var zmq = require('zmq')
 var logger = require('../../util/logger')
 var wire = require('../../wire')
 var wireutil = require('../../wire/util')
+var wirerouter = require('../../wire/router')
 var dbapi = require('../../db/api')
 var lifecycle = require('../../util/lifecycle')
 var srv = require('../../util/srv')
 
+var TtlSet = require('./util/ttlset')
+
 module.exports = function(options) {
   var log = logger.createLogger('reaper')
-    , timer
+  var ttlset = new TtlSet(options.heartbeatTimeout)
 
   if (options.name) {
     logger.setGlobalIdentifier(options.name)
   }
 
+  // Input
+  var sub = zmq.socket('sub')
+  Promise.map(options.endpoints.sub, function(endpoint) {
+    return srv.resolve(endpoint).then(function(records) {
+      return srv.attempt(records, function(record) {
+        log.info('Receiving input from "%s"', record.url)
+        sub.connect(record.url)
+        return Promise.resolve(true)
+      })
+    })
+  })
+  .catch(function(err) {
+    log.fatal('Unable to connect to sub endpoint', err)
+    lifecycle.fatal()
+  })
+
+  // Establish always-on channels
+  ;[wireutil.global].forEach(function(channel) {
+    log.info('Subscribing to permanent channel "%s"', channel)
+    sub.subscribe(channel)
+  })
+
   // Output
   var push = zmq.socket('push')
   Promise.map(options.endpoints.push, function(endpoint) {
@@ -32,39 +57,55 @@ module.exports = function(options) {
     lifecycle.fatal()
   })
 
-  function reap() {
-    dbapi.getDeadDevices(options.heartbeatTimeout)
-      .then(function(cursor) {
-        return Promise.promisify(cursor.toArray, cursor)()
-          .then(function(list) {
-            list.forEach(function(device) {
-              log.info('Reaping device "%s"', device.serial)
-              push.send([
-                wireutil.global
-              , wireutil.envelope(new wire.DeviceAbsentMessage(
-                  device.serial
-                ))
-              ])
-            })
-          })
-      })
-      .catch(function(err) {
-        log.error('Failed to load device list: ', err.message, err.stack)
-        lifecycle.fatal()
-      })
-  }
+  sub.on('message', wirerouter()
+    .on(wire.DevicePresentMessage, function(channel, message) {
+      ttlset.bump(message.serial, Date.now())
+    })
+    .on(wire.DeviceHeartbeatMessage, function(channel, message) {
+      ttlset.bump(message.serial, Date.now())
+    })
+    .on(wire.DeviceAbsentMessage, function(channel, message) {
+      ttlset.drop(message.serial)
+    })
+    .handler())
+
+  ttlset.on('timeout', function(serial) {
+    log.info('Reaping device "%s" due to heartbeat timeout', serial)
+    push.send([
+      wireutil.global
+    , wireutil.envelope(new wire.DeviceAbsentMessage(
+        serial
+      ))
+    ])
+  })
+
+  dbapi.loadPresentDevices()
+    .then(function(cursor) {
+      return Promise.promisify(cursor.toArray, cursor)()
+        .then(function(list) {
+          var now = Date.now()
+          list.forEach(function(device) {
+            ttlset.bump(device.serial, now)
+          })
+        })
+    })
+    .catch(function(err) {
+      log.error('Failed to load device list: ', err.stack)
+      lifecycle.fatal()
+    })
 
-  timer = setInterval(reap, options.reapInterval)
   log.info('Reaping devices with no heartbeat')
 
   lifecycle.observe(function() {
-    clearTimeout(timer)
+    [push, sub].forEach(function(sock) {
+      try {
+        sock.close()
+      }
+      catch (err) {
+        // No-op
+      }
+    })
 
-    try {
-      push.close()
-    }
-    catch (err) {
-      // No-op
-    }
+    ttlset.stop()
   })
 }

lib/units/reaper/util/ttlset.js

@@ -0,0 +1,108 @@
+var util = require('util')
+
+var EventEmitter = require('eventemitter3').EventEmitter
+
+function TtlItem(value) {
+  this.next = null
+  this.prev = null
+  this.time = null
+  this.value = value
+}
+
+function TtlSet(ttl) {
+  this.head = null
+  this.tail = null
+  this.mapping = Object.create(null)
+  this.ttl = ttl
+  this.timer = null
+}
+
+util.inherits(TtlSet, EventEmitter)
+
+TtlSet.prototype.bump = function(value, time) {
+  var item = this._remove(this.mapping[value]) || this._create(value)
+
+  item.time = time || Date.now()
+  item.prev = this.tail
+
+  this.tail = item
+
+  if (!this.head) {
+    this.head = item
+    this._scheduleCheck()
+  }
+}
+
+TtlSet.prototype.drop = function(value) {
+  this._drop(this.mapping[value])
+}
+
+TtlSet.prototype.stop = function() {
+  clearTimeout(this.timer)
+}
+
+TtlSet.prototype._scheduleCheck = function() {
+  clearTimeout(this.timer)
+  if (this.head) {
+    var delay = Math.max(0, this.ttl - (Date.now() - this.head.time))
+    this.timer = setTimeout(this._check.bind(this), delay)
+  }
+}
+
+TtlSet.prototype._check = function() {
+  var now = Date.now()
+
+  var item
+  while ((item = this.head)) {
+    if (now - item.time > this.ttl) {
+      this._drop(item)
+      this.emit('timeout', item.value)
+    }
+    else {
+      break
+    }
+  }
+
+  this._scheduleCheck()
+}
+
+TtlSet.prototype._create = function(value) {
+  var item = new TtlItem(value)
+  this.mapping[value] = item
+  return item
+}
+
+TtlSet.prototype._drop = function(item) {
+  if (item) {
+    this._remove(item)
+    delete this.mapping[item.value]
+  }
+}
+
+TtlSet.prototype._remove = function(item) {
+  if (!item) {
+    return null
+  }
+
+  if (item.prev) {
+    item.prev.next = item.next
+  }
+
+  if (item.next) {
+    item.next.prev = item.prev
+  }
+
+  if (item === this.head) {
+    this.head = item.next
+  }
+
+  if (item === this.tail) {
+    this.tail = item.prev
+  }
+
+  item.next = item.prev = null
+
+  return item
+}
+
+module.exports = TtlSet