feat: slot change callback (provides real-time fork information)

web3.js/src/connection.js

@@ -318,6 +318,23 @@ const ProgramAccountNotificationResult = struct({
   result: ProgramAccountInfoResult,
 });
 
+/**
+ * @private
+ */
+const SlotInfo = struct({
+  parent: 'number',
+  slot: 'number',
+  root: 'number',
+});
+
+/***
+ * Expected JSON RPC response for the "slotNotification" message
+ */
+const SlotNotificationResult = struct({
+  subscription: 'number',
+  result: SlotInfo,
+});
+
 /**
  * Expected JSON RPC response for the "getProgramAccounts" message
  */
@@ -569,6 +586,19 @@ type ProgramAccountSubscriptionInfo = {
   subscriptionId: null | number, // null when there's no current server subscription id
 };
 
+/**
+ * @private
+ */
+type SlotSubscriptionInfo = {
+  callback: SlotChangeCallback,
+  subscriptionId: null | number, // null when there's no current server subscription id
+};
+
+/**
+ * Callback function for slot change notifications
+ */
+export type SlotChangeCallback = (slotInfo: SlotInfo) => void;
+
 /**
  * Signature status: Success
  *
@@ -618,6 +648,10 @@ export class Connection {
     [number]: ProgramAccountSubscriptionInfo,
   } = {};
   _programAccountChangeSubscriptionCounter: number = 0;
+  _slotSubscriptions: {
+    [number]: SlotSubscriptionInfo,
+  } = {};
+  _slotSubscriptionCounter: number = 0;
 
   /**
    * Establish a JSON RPC connection
@@ -657,6 +691,10 @@ export class Connection {
       'programNotification',
       this._wsOnProgramAccountNotification.bind(this),
     );
+    this._rpcWebSocket.on(
+      'slotNotification',
+      this._wsOnSlotNotification.bind(this),
+    );
   }
 
   /**
@@ -1256,7 +1294,12 @@ export class Connection {
     const programKeys = Object.keys(
       this._programAccountChangeSubscriptions,
     ).map(Number);
-    if (accountKeys.length === 0 && programKeys.length === 0) {
+    const slotKeys = Object.keys(this._slotSubscriptions).map(Number);
+    if (
+      accountKeys.length === 0 &&
+      programKeys.length === 0 &&
+      slotKeys.length === 0
+    ) {
       this._rpcWebSocket.close();
       return;
     }
@@ -1268,6 +1311,9 @@ export class Connection {
       for (let id of programKeys) {
         this._programAccountChangeSubscriptions[id].subscriptionId = null;
       }
+      for (let id of slotKeys) {
+        this._slotSubscriptions[id].subscriptionId = null;
+      }
       this._rpcWebSocket.connect();
       return;
     }
@@ -1307,6 +1353,18 @@ export class Connection {
         }
       }
     }
+    for (let id of slotKeys) {
+      const {subscriptionId} = this._slotSubscriptions[id];
+      if (subscriptionId === null) {
+        try {
+          this._slotSubscriptions[
+            id
+          ].subscriptionId = await this._rpcWebSocket.call('slotSubscribe', []);
+        } catch (err) {
+          console.log(`slotSubscribe error: ${err.message}`);
+        }
+      }
+    }
   }
 
   /**
@@ -1455,6 +1513,69 @@ export class Connection {
     }
   }
 
+  /**
+   * @private
+   */
+  _wsOnSlotNotification(notification: Object) {
+    const res = SlotNotificationResult(notification);
+    if (res.error) {
+      throw new Error(res.error.message);
+    }
+    assert(typeof res.result !== 'undefined');
+    const {parent, slot, root} = res.result;
+
+    const keys = Object.keys(this._slotSubscriptions).map(Number);
+    for (let id of keys) {
+      const sub = this._slotSubscriptions[id];
+      if (sub.subscriptionId === res.subscription) {
+        sub.callback({
+          parent,
+          slot,
+          root,
+        });
+        return true;
+      }
+    }
+  }
+
+  /**
+   * Register a callback to be invoked upon slot changes
+   *
+   * @param callback Function to invoke whenever the slot changes
+   * @return subscription id
+   */
+  onSlotChange(callback: SlotChangeCallback): number {
+    const id = ++this._slotSubscriptionCounter;
+    this._slotSubscriptions[id] = {
+      callback,
+      subscriptionId: id,
+    };
+    this._updateSubscriptions();
+    return id;
+  }
+
+  /**
+   * Deregister a slot notification callback
+   *
+   * @param id subscription id to deregister
+   */
+  async removeSlotChangeListener(id: number): Promise<void> {
+    if (this._slotSubscriptions[id]) {
+      const {subscriptionId} = this._slotSubscriptions[id];
+      delete this._slotSubscriptions[id];
+      if (subscriptionId !== null) {
+        try {
+          await this._rpcWebSocket.call('slotUnsubscribe', [subscriptionId]);
+        } catch (err) {
+          console.log('slotUnsubscribe error:', err.message);
+        }
+      }
+      this._updateSubscriptions();
+    } else {
+      throw new Error(`Unknown slot change id: ${id}`);
+    }
+  }
+
   _argsWithCommitment(args: Array<any>, override: ?Commitment): Array<any> {
     const commitment = override || this._commitment;
     if (commitment) {