Add configurable heartbeat to WebSocketOrderbookChannel

6 files changed, 42 insertions, 2 deletions

diff --git a/packages/connect/CHANGELOG.md b/packages/connect/CHANGELOG.md
index 18c808e5e..b187267e5 100644
--- a/packages/connect/CHANGELOG.md
+++ b/packages/connect/CHANGELOG.md
@@ -3,6 +3,7 @@
 ## v0.6.0 - _TBD, 2018_
 
     * Add pagination options to HttpClient methods (#393)
+    * Add heartbeat configuration to WebSocketOrderbookChannel constructor (#393)
 
 ## v0.5.7 - _February 9, 2018_
 
diff --git a/packages/connect/src/index.ts b/packages/connect/src/index.ts
index 344a32e28..bb42384f9 100644
--- a/packages/connect/src/index.ts
+++ b/packages/connect/src/index.ts
@@ -17,4 +17,5 @@ export {
     TokenPairsItem,
     TokenPairsRequestOpts,
     TokenTradeInfo,
+    WebSocketOrderbookChannelConfig,
 } from './types';
diff --git a/packages/connect/src/schemas/schemas.ts b/packages/connect/src/schemas/schemas.ts
index 0b8b798a9..b9a8472fb 100644
--- a/packages/connect/src/schemas/schemas.ts
+++ b/packages/connect/src/schemas/schemas.ts
@@ -3,6 +3,7 @@ import { orderBookRequestSchema } from './orderbook_request_schema';
 import { ordersRequestOptsSchema } from './orders_request_opts_schema';
 import { pagedRequestOptsSchema } from './paged_request_opts_schema';
 import { tokenPairsRequestOptsSchema } from './token_pairs_request_opts_schema';
+import { webSocketOrderbookChannelConfigSchema } from './websocket_orderbook_channel_config_schema';
 
 export const schemas = {
     feesRequestSchema,
@@ -10,4 +11,5 @@ export const schemas = {
     ordersRequestOptsSchema,
     pagedRequestOptsSchema,
     tokenPairsRequestOptsSchema,
+    webSocketOrderbookChannelConfigSchema,
 };
diff --git a/packages/connect/src/schemas/websocket_orderbook_channel_config_schema.ts b/packages/connect/src/schemas/websocket_orderbook_channel_config_schema.ts
new file mode 100644
index 000000000..2e9987e63
--- /dev/null
+++ b/packages/connect/src/schemas/websocket_orderbook_channel_config_schema.ts
@@ -0,0 +1,7 @@
+export const webSocketOrderbookChannelConfigSchema = {
+    id: '/WebSocketOrderbookChannelConfig',
+    type: 'object',
+    properties: {
+        heartbeatIntervalMs: { type: 'number' },
+    },
+};
diff --git a/packages/connect/src/types.ts b/packages/connect/src/types.ts
index 970eff498..5f837b0b3 100644
--- a/packages/connect/src/types.ts
+++ b/packages/connect/src/types.ts
@@ -44,6 +44,13 @@ export interface OrderbookChannel {
 }
 
 /*
+ * heartbeatInterval: Interval in milliseconds that the orderbook channel should ping the underlying websocket. Default: 15000
+ */
+export interface WebSocketOrderbookChannelConfig {
+    heartbeatIntervalMs?: number;
+}
+
+/*
  * baseTokenAddress: The address of token designated as the baseToken in the currency pair calculation of price
  * quoteTokenAddress: The address of token designated as the quoteToken in the currency pair calculation of price
  * snapshot: If true, a snapshot of the orderbook will be sent before the updates to the orderbook
diff --git a/packages/connect/src/ws_orderbook_channel.ts b/packages/connect/src/ws_orderbook_channel.ts
index 822a022f4..564e8686b 100644
--- a/packages/connect/src/ws_orderbook_channel.ts
+++ b/packages/connect/src/ws_orderbook_channel.ts
@@ -3,6 +3,7 @@ import { schemas } from '@0xproject/json-schemas';
 import * as _ from 'lodash';
 import * as WebSocket from 'websocket';
 
+import { schemas as clientSchemas } from './schemas/schemas';
 import {
     OrderbookChannel,
     OrderbookChannelHandler,
@@ -10,9 +11,12 @@ import {
     OrderbookChannelSubscriptionOpts,
     WebsocketClientEventType,
     WebsocketConnectionEventType,
+    WebSocketOrderbookChannelConfig,
 } from './types';
 import { orderbookChannelMessageParser } from './utils/orderbook_channel_message_parser';
 
+const DEFAULT_HEARTBEAT_INTERVAL_MS = 15000;
+
 /**
  * This class includes all the functionality related to interacting with a websocket endpoint
  * that implements the standard relayer API v0
@@ -21,15 +25,25 @@ export class WebSocketOrderbookChannel implements OrderbookChannel {
     private _apiEndpointUrl: string;
     private _client: WebSocket.client;
     private _connectionIfExists?: WebSocket.connection;
+    private _heartbeatTimerIfExists?: NodeJS.Timer;
     private _subscriptionCounter = 0;
+    private _heartbeatIntervalMs: number;
     /**
      * Instantiates a new WebSocketOrderbookChannel instance
      * @param   url                 The relayer API base WS url you would like to interact with
+     * @param   url                 The configuration object. Look up the type for the description.
      * @return  An instance of WebSocketOrderbookChannel
      */
-    constructor(url: string) {
+    constructor(url: string, config?: WebSocketOrderbookChannelConfig) {
         assert.isUri('url', url);
+        if (!_.isUndefined(config)) {
+            assert.doesConformToSchema('config', config, clientSchemas.webSocketOrderbookChannelConfigSchema);
+        }
         this._apiEndpointUrl = url;
+        this._heartbeatIntervalMs =
+            _.isUndefined(config) || _.isUndefined(config.heartbeatIntervalMs)
+                ? DEFAULT_HEARTBEAT_INTERVAL_MS
+                : config.heartbeatIntervalMs;
         this._client = new WebSocket.client();
     }
     /**
@@ -63,7 +77,7 @@ export class WebSocketOrderbookChannel implements OrderbookChannel {
                 connection.on(WebsocketConnectionEventType.Error, wsError => {
                     handler.onError(this, subscriptionOpts, wsError);
                 });
-                connection.on(WebsocketConnectionEventType.Close, () => {
+                connection.on(WebsocketConnectionEventType.Close, (code: number, desc: string) => {
                     handler.onClose(this, subscriptionOpts);
                 });
                 connection.on(WebsocketConnectionEventType.Message, message => {
@@ -80,6 +94,9 @@ export class WebSocketOrderbookChannel implements OrderbookChannel {
         if (!_.isUndefined(this._connectionIfExists)) {
             this._connectionIfExists.close();
         }
+        if (!_.isUndefined(this._heartbeatTimerIfExists)) {
+            clearInterval(this._heartbeatTimerIfExists);
+        }
     }
     private _getConnection(callback: (error?: Error, connection?: WebSocket.connection) => void) {
         if (!_.isUndefined(this._connectionIfExists) && this._connectionIfExists.connected) {
@@ -87,6 +104,11 @@ export class WebSocketOrderbookChannel implements OrderbookChannel {
         } else {
             this._client.on(WebsocketClientEventType.Connect, connection => {
                 this._connectionIfExists = connection;
+                if (this._heartbeatIntervalMs !== 0) {
+                    this._heartbeatTimerIfExists = setInterval(() => {
+                        connection.ping('');
+                    }, this._heartbeatIntervalMs);
+                }
                 callback(undefined, this._connectionIfExists);
             });
             this._client.on(WebsocketClientEventType.ConnectFailed, error => {