Repository: brightcove/hot-shots
Pull request #129 (2019-10-08 07:12:39)

Abstract transports
.. was looking into adding another transport (raw stream), however the way different protocols were managed was quite hectic. Doing abstraction as a first step toward that.


Patch:
diff --git a/lib/constants.js b/lib/constants.js
new file mode 100644
index 0000000..1cc21af
--- /dev/null
+++ b/lib/constants.js
@@ -0,0 +1,6 @@
+exports.PROTOCOL = {
+  TCP: 'tcp',
+  UDS: 'uds',
+  UDP: 'udp',
+  STREAM: 'stream'
+};
diff --git a/lib/statsd.js b/lib/statsd.js
index ab644ed..09cd4ed 100644
--- a/lib/statsd.js
+++ b/lib/statsd.js
@@ -1,18 +1,11 @@
-const dgram = require('dgram'),
-  util = require('util'),
+const util = require('util'),
   dns = require('dns'),
-  net = require('net'),
   helpers = require('./helpers'),
   applyStatsFns = require('./statsFunctions');
 
-// installed below, only if needed
-let unixDgram;
+const { PROTOCOL } = require('./constants');
+const createTransport = require('./transport');
 
-const TCP = 'tcp';
-const UDS = 'uds';
-// UDP constant not needed/used, defaults to udp
-
-const UDS_PATH_DEFAULT = '/var/run/datadog/dsd.socket';
 const UDS_DEFAULT_GRACEFUL_RESTART_LIMIT = 1000;
 
 /**
@@ -44,53 +37,6 @@ const Client = function (host, port, prefix, suffix, globalize, cacheDns, mock,
     };
   }
 
-  const createSocket = function createSocket(instance, args) {
-    let socket;
-    let errMessage;
-
-    if (args.protocol === TCP) {
-      try {
-        socket = net.connect(args.port, args.host);
-        socket.setKeepAlive(true);
-      } catch (err) {
-        errMessage = `Could not establish connection to ${args.host}:${args.port}: ${err}`;
-      }
-    } else if (args.protocol === UDS) {
-      try {
-        // this will not always be available, as noted in the error message below
-        unixDgram = require('unix-dgram'); // eslint-disable-line global-require
-      }
-      catch (err) {
-        errMessage = 'The library unix_dgram, needed for the uds protocol to work, is not installed. You need to pick another protocol to use hot-shots.  See the hot-shots README for additional details.';
-      }
-
-      if (unixDgram) {
-        const udsPath = args.path ? args.path : UDS_PATH_DEFAULT;
-        try {
-          socket = unixDgram.createSocket('unix_dgram');
-          socket.connect(udsPath);
-        } catch (err) {
-          errMessage = `Could not establish UDS connection to ${udsPath}: ${err}`;
-        }
-      }
-    } else {
-      try {
-        socket = dgram.createSocket('udp4');
-      } catch (err) {
-        errMessage = `Could not create socket: ${err}`;
-      }
-    }
-
-    if (errMessage) {
-      if (instance.errorHandler) {
-        instance.errorHandler(new Error(errMessage));
-      } else {
-        console.error(errMessage);
-      }
-    }
-
-    return socket;
-  };
 
   // hidden global_tags option for backwards compatibility
   options.globalTags = options.globalTags || options.global_tags;
@@ -100,7 +46,7 @@ const Client = function (host, port, prefix, suffix, globalize, cacheDns, mock,
   this.port = options.port || parseInt(process.env.DD_DOGSTATSD_PORT, 10) || 8125;
   this.prefix = options.prefix || '';
   this.suffix = options.suffix || '';
-  this.socket = options.isChild ? options.socket : createSocket(this, {
+  this.socket = options.isChild ? options.socket : createTransport(this, {
     host: this.host,
     path: options.path,
     port: this.port,
@@ -160,7 +106,7 @@ const Client = function (host, port, prefix, suffix, globalize, cacheDns, mock,
   // will gracefully (attempt) to re-open the socket with a small delay
   // options.udsGracefulRestartRateLimit is the minimum time (ms) between creating sockets
   // does not support options.isChild (how to re-create a socket you didn't create?)
-  if (!options.isChild && options.protocol === 'uds' && options.udsGracefulErrorHandling) {
+  if (!options.isChild && options.protocol === PROTOCOL.UDS && options.udsGracefulErrorHandling) {
     const socketCreateLimit = options.udsGracefulRestartRateLimit || UDS_DEFAULT_GRACEFUL_RESTART_LIMIT; // only recreate once per second
     const lastSocketCreateTime = Date.now();
     this.socket.on('error', (err) => {
@@ -174,7 +120,7 @@ const Client = function (host, port, prefix, suffix, globalize, cacheDns, mock,
               this.socket.removeListener('error', this.errorHandler);
             }
             this.socket.close();
-            this.socket = createSocket(this, {
+            this.socket = createTransport(this, {
               host: this.host,
               path: options.path,
               port: this.port,
@@ -402,8 +348,8 @@ Client.prototype.sendMessage = function (message, callback) {
     return;
   }
 
-  if (this.protocol === TCP && message.lastIndexOf('\n') !== message.length - 1) {
-    message += '\n';
+  if (!this.socket) {
+    return callback(new Error('No protocol set. Check previous errors for details.'));
   }
 
   const handleCallback = (err) => {
@@ -418,23 +364,16 @@ Client.prototype.sendMessage = function (message, callback) {
       if (this.errorHandler) {
         this.errorHandler(errFormatted);
       } else {
-        console.log(String(errFormatted));
+        console.error(String(errFormatted));
         // emit error ourselves on the socket for backwards compatibility
         this.socket.emit('error', errFormatted);
       }
     }
   };
 
-  const buf = Buffer.from(message);
   try {
     this.messagesInFlight++;
-    if (this.protocol === TCP) {
-      this.socket.write(buf, 'ascii', handleCallback);
-    } else if (this.protocol === UDS) {
-      this.socket.send(buf, handleCallback);
-    } else {
-      this.socket.send(buf, 0, buf.length, this.port, this.host, handleCallback);
-    }
+    this.socket.send(Buffer.from(message), handleCallback);
   } catch (err) {
     handleCallback(err);
   }
@@ -527,19 +466,7 @@ Client.prototype._close = function (callback) {
   }
 
   try {
-    if (this.protocol === TCP) {
-      this.socket.destroy();
-    } else if (this.protocol === UDS) {
-      this.socket.close();
-      // this close is synchronous, and the socket will not emit a
-      // close event, so we can callback right away
-      if (callback) {
-        callback();
-      }
-    }
-    else {
-      this.socket.close();
-    }
+    this.socket.close();
   } catch (err) {
     handleErr(err);
   }
diff --git a/lib/transport.js b/lib/transport.js
new file mode 100644
index 0000000..739af6d
--- /dev/null
+++ b/lib/transport.js
@@ -0,0 +1,99 @@
+const dgram = require('dgram');
+const net = require('net');
+const { PROTOCOL } = require('./constants');
+
+// Imported below, only if needed
+let unixDgram;
+
+const UDS_PATH_DEFAULT = '/var/run/datadog/dsd.socket';
+
+// interface Transport {
+//   emit(name: string, payload: any):void;
+//   on(name: string, listener: Function):void;
+//   removeListener(name: string, listener: Function):void;
+//   send(buf: Buffer, callback: Function):void;
+//   close():void;
+// }
+const createTcpTransport = args => {
+  const socket = net.connect(args.port, args.host);
+  socket.setKeepAlive(true);
+  return {
+    emit: socket.emit.bind(socket),
+    on: socket.on.bind(socket),
+    removeListener: socket.removeListener.bind(socket),
+    send: (buf, callback) => {
+      let msg = buf.toString();
+      if (msg.length > 0 && msg[msg.length - 1] !== '\n') {
+        msg += '\n';
+      }
+      socket.write(msg, 'ascii', callback);
+    },
+    close: () => socket.destroy()
+  };
+};
+
+const createUdpTransport = args => {
+  const socket = dgram.createSocket('udp4');
+  return {
+    emit: socket.emit.bind(socket),
+    on: socket.on.bind(socket),
+    removeListener: socket.removeListener.bind(socket),
+    send: (buf, callback) => socket.send(buf, 0, buf.length, args.port, args.host, callback),
+    close: socket.close.bind(socket)
+  };
+};
+
+const createUdsTransport = args => {
+  try {
+    // This will not always be available, as noted in the error message below
+    unixDgram = require('unix-dgram'); // eslint-disable-line global-require
+  } catch (err) {
+    throw new Error(
+      'The library `unix_dgram`, needed for the uds protocol to work, is not installed. ' +
+        'You need to pick another protocol to use hot-shots. ' +
+        'See the hot-shots README for additional details.'
+    );
+  }
+  const udsPath = args.path ? args.path : UDS_PATH_DEFAULT;
+  const socket = unixDgram.createSocket('unix_dgram');
+  socket.connect(udsPath);
+
+  return {
+    emit: socket.emit.bind(socket),
+    on: socket.on.bind(socket),
+    removeListener: socket.removeListener.bind(socket),
+    send: socket.send.bind(socket),
+    close: () => {
+      socket.close();
+      // close is synchronous, and the socket will not emit a
+      // close event, so we can callback right away
+      socket.emit('close');
+    }
+  };
+};
+
+module.exports = (instance, args) => {
+  let transport = null;
+  const protocol = args.protocol || PROTOCOL.UDP;
+
+  try {
+    if (protocol === PROTOCOL.TCP) {
+      transport = createTcpTransport(args);
+    } else if (protocol === PROTOCOL.UDS) {
+      transport = createUdsTransport(args);
+    } else if (protocol === PROTOCOL.UDP) {
+      transport = createUdpTransport(args);
+    } else {
+      throw new Error(`Unsupported protocol '${protocol}'`);
+    }
+    transport.type = protocol;
+  } catch (e) {
+    if (instance.errorHandler) {
+      instance.errorHandler(e);
+    } else {
+      console.error(e);
+    }
+  }
+
+  return transport;
+};


Test patch:
diff --git a/test/errorHandling.js b/test/errorHandling.js
index caa109d..cf3e965 100644
--- a/test/errorHandling.js
+++ b/test/errorHandling.js
@@ -127,18 +127,18 @@ describe('#errorHandling', () => {
       it('should errback for an unresolvable host', done => {
         // this does not work for tcp, which throws an error during setup
         // that needs errorHandler or a socket.on('error') handler
-        if (serverType === 'tcp') {
+        if (serverType !== 'udp') {
           return done();
         }
 
-       statsd = createHotShotsClient({
+        statsd = createHotShotsClient({
           host: '...',
           protocol: serverType
         }, clientType);
 
         statsd.send('test title', [], error => {
           assert.ok(error);
-          assert.equal(error.code, serverType === 'uds' ? '-2' : 'ENOTFOUND');
+          assert.equal(error.code, 'ENOTFOUND');
           // skip closing, because the unresolvable host hangs
           statsd = null;
           done();
@@ -146,6 +146,10 @@ describe('#errorHandling', () => {
       });
 
       it('should use errorHandler for an unresolvable host', done => {
+        // Test is not applicable for uds/stream transports
+        if (!['tcp', 'udp'].includes(serverType)) {
+          return done();
+        }
         statsd = createHotShotsClient({
           host: '...',
           protocol: serverType,
@@ -163,6 +167,10 @@ describe('#errorHandling', () => {
       });
 
       it('should throw error on socket for an unresolvable host', done => {
+        // Test is not applicable for uds/stream transports
+        if (!['tcp', 'udp'].includes(serverType)) {
+          return done();
+        }
         statsd = createHotShotsClient({
           host: '...',
           protocol: serverType
diff --git a/test/helpers/helpers.js b/test/helpers/helpers.js
index 21ba156..f932944 100644
--- a/test/helpers/helpers.js
+++ b/test/helpers/helpers.js
@@ -5,7 +5,7 @@ const fs = require('fs');
 const os = require('os');
 const process = require('process');
 const StatsD = require('../../lib/statsd.js');
-
+const { TCP, UDP, UDS } = require('../../lib/constants').PROTOCOL;
 let unixDgram;
 try {
   // this will not always be available
@@ -18,9 +18,6 @@ catch (err) {
 const CLIENT = 'client';
 const CHILD_CLIENT = 'child client';
 const CHILD_CHILD_CLIENT = 'child child client';
-const TCP = 'tcp';
-const UDP = 'udp';
-const UDS = 'uds';
 const TCP_BROKEN = 'tcp_broken';
 const UDS_BROKEN = 'uds_broken';
 
@@ -120,7 +117,7 @@ function createServer(serverType, onListening) {
     server.bind(0, '127.0.0.1');
   }
   else if (serverType === UDS) {
-    // we always have to manually unlink the test socket
+    // We always have to manually unlink the test socket
     if (fs.existsSync(UDS_TEST_PATH)) { // eslint-disable-line no-sync
       fs.unlinkSync(UDS_TEST_PATH); // eslint-disable-line no-sync
     }
@@ -133,7 +130,7 @@ function createServer(serverType, onListening) {
       onListening('127.0.0.1');
     });
     server.on('error', (err) => {
-      console.log('uds connection failed', err);
+      console.error('Error: uds connection failed', err);
       onListening('127.0.0.1');
     });
     server.bind(UDS_TEST_PATH);
@@ -203,6 +200,8 @@ function createServer(serverType, onListening) {
  * @param {*} clientType
  */
 function createHotShotsClient(args, clientType) {
+  // FIXME: This is inconsistent with the rest.
+  // Fn `createServer` should be returning the path instead.
   if (args.protocol === UDS) {
     args.path = UDS_TEST_PATH;
   }
diff --git a/test/init.js b/test/init.js
index 6e0b093..378519e 100644
--- a/test/init.js
+++ b/test/init.js
@@ -1,7 +1,5 @@
 const assert = require('assert');
-const dgram = require('dgram');
 const dns = require('dns');
-const net = require('net');
 
 const StatsD = require('../lib/statsd');
 const helpers = require('./helpers/helpers.js');
@@ -234,7 +232,7 @@ describe('#init', () => {
 
   it('should create a socket variable that is an instance of dgram.Socket', () => {
     statsd = createHotShotsClient({}, clientType);
-    assert.ok(statsd.socket instanceof dgram.Socket);
+    assert.equal(statsd.socket.type, 'udp');
     skipClose = true;
   });
 
@@ -245,7 +243,7 @@ describe('#init', () => {
         port: address.port,
         protocol: 'tcp'
       }, clientType);
-      assert.ok(statsd.socket instanceof net.Socket);
+      assert.equal(statsd.socket.type, 'tcp');
       done();
     });
   });