From 7bfeaa1b207d8c5e90918104adf3978f5d5c603b Mon Sep 17 00:00:00 2001
From: Jaco Greeff <jacogr@gmail.com>
Date: Thu, 25 Oct 2018 14:45:51 +0200
Subject: [PATCH] Auto re-subscribe upon re-connect (#340)

* Auto re-subscribe upon re-connect

* Update FIXME

* newHead -> subscribeNewHead in example
---
 .../promise/02_listen_to_blocks/index.js      |  2 +-
 packages/rpc-provider/src/ws/index.ts         | 45 ++++++++++++++++---
 .../rpc-provider/src/ws/unsubscribe.spec.js   |  4 +-
 3 files changed, 41 insertions(+), 10 deletions(-)

diff --git a/docs/examples/promise/02_listen_to_blocks/index.js b/docs/examples/promise/02_listen_to_blocks/index.js
index a3ecd30c78..271384f916 100755
--- a/docs/examples/promise/02_listen_to_blocks/index.js
+++ b/docs/examples/promise/02_listen_to_blocks/index.js
@@ -10,7 +10,7 @@ async function main () {
   // Subscribe to the new headers on-chain. The callback is fired when new headers
   // are found, the call itself returns a promise with a subscription that can be
   // used to unsubscribe from the newHead subscription
-  const subscriptionId = await api.rpc.chain.newHead((header) => {
+  const subscriptionId = await api.rpc.chain.subscribeNewHead((header) => {
     console.log(`best #${header.blockNumber}`);
   });
 
diff --git a/packages/rpc-provider/src/ws/index.ts b/packages/rpc-provider/src/ws/index.ts
index 75d3c648c6..73307a83a0 100644
--- a/packages/rpc-provider/src/ws/index.ts
+++ b/packages/rpc-provider/src/ws/index.ts
@@ -59,13 +59,13 @@ const l = logger('api-ws');
  */
 export default class WsProvider implements WSProviderInterface {
   private _eventemitter: EventEmitter;
+  private _isConnected: boolean = false;
   private autoConnect: boolean;
   private coder: Coder;
   private endpoint: string;
   private handlers: {
-    [index: number]: WsState$Awaiting
+    [index: string]: WsState$Awaiting
   };
-  private _isConnected: boolean;
   private queued: {
     [index: string]: string
   };
@@ -85,7 +85,6 @@ export default class WsProvider implements WSProviderInterface {
     this.autoConnect = autoConnect;
     this.coder = new Coder();
     this.endpoint = endpoint;
-    this._isConnected = false;
     this.handlers = {};
     this.queued = {};
     this.subscriptions = {};
@@ -133,6 +132,9 @@ export default class WsProvider implements WSProviderInterface {
 
   /**
    * @summary Send JSON data using WebSockets to configured HTTP Endpoint or queue.
+   * @param method The RPC methods to execute
+   * @param params Encoded paramaters as appliucable for the method
+   * @param subscription Subscription details (internally used)
    */
   async send (method: string, params: Array<any>, subscription?: SubscriptionHandler): Promise<any> {
     return new Promise((resolve, reject): void => {
@@ -202,7 +204,15 @@ export default class WsProvider implements WSProviderInterface {
   async unsubscribe (type: string, method: string, id: number): Promise<boolean> {
     const subscription = `${type}::${id}`;
 
-    assert(!isUndefined(this.subscriptions[subscription]), `Unable to find active subscription=${subscription}`);
+    // FIXME This now could happen with re-subscriptions. The issue is that with a re-sub
+    // the assigned id now does not match what the API user originally received. It has
+    // a slight complication in solving - since we cannot rely on the send id, but rather
+    // need to find the actual subscription id to map it
+    if (isUndefined(this.subscriptions[subscription])) {
+      l.warn(`Unable to find active subscription=${subscription}`);
+
+      return false;
+    }
 
     delete this.subscriptions[subscription];
 
@@ -301,9 +311,32 @@ export default class WsProvider implements WSProviderInterface {
     this._isConnected = true;
     this.emit('connected');
 
+    this.sendQueue();
+    this.resubscribe();
+
+    return true;
+  }
+
+  private resubscribe (): void {
+    const subscriptions = this.subscriptions;
+
+    this.subscriptions = {};
+
+    Object.keys(subscriptions).forEach(async (id) => {
+      const { callback, method, params, type } = subscriptions[id];
+
+      try {
+        await this.subscribe(type, method, params, callback);
+      } catch (error) {
+        l.error(error);
+      }
+    });
+  }
+
+  private sendQueue (): void {
     Object.keys(this.queued).forEach((id) => {
       try {
-        // @ts-ignore checked above
+        // @ts-ignore we have done the websocket check in onSocketOpen, if an issue, will catch it
         this.websocket.send(
           this.queued[id]
         );
@@ -313,7 +346,5 @@ export default class WsProvider implements WSProviderInterface {
         l.error(error);
       }
     });
-
-    return true;
   }
 }
diff --git a/packages/rpc-provider/src/ws/unsubscribe.spec.js b/packages/rpc-provider/src/ws/unsubscribe.spec.js
index a086e733d5..5932f92818 100644
--- a/packages/rpc-provider/src/ws/unsubscribe.spec.js
+++ b/packages/rpc-provider/src/ws/unsubscribe.spec.js
@@ -78,8 +78,8 @@ describe('subscribe', () => {
       .then((id) => {
         return ws.unsubscribe('test', 'subscribe_test', 111);
       })
-      .catch((error) => {
-        expect(error.message).toMatch(/find active subscription/);
+      .then((result) => {
+        expect(result).toBe(false);
       });
   });
 });