Merge pull request #67 from ChainSafe/tuyen/message-id-function

Provide a way to override message id function

Author: wemeetagain

README.md

@@ -67,6 +67,8 @@ Options is an optional object with the following key-value pairs:
 
 * **`fallbackToFloodsub`**: boolean identifying whether the node should fallback to the floodsub protocol, if another connecting peer does not support gossipsub (defaults to **true**).
 * **`emitSelf`**: boolean identifying whether the node should emit to self on publish, in the event of the topic being subscribed (defaults to **false**).
+* **`msgIdFn`**: a function with signature `(message) => string` defining the message id given a message, used internally to deduplicate gossip (defaults to `(message) => message.from + message.seqno.toString('hex')`)
+* **`messageCache`**: optional, a customized `MessageCache` instance, see the implementation for the interface.
 
 For the remaining API, see https://github.com/libp2p/js-libp2p-pubsub
 

src/index.js

@@ -22,6 +22,8 @@ class GossipSub extends BasicPubsub {
    * @param {bool} [options.emitSelf] if publish should emit to self, if subscribed, defaults to false
    * @param {bool} [options.gossipIncoming] if incoming messages on a subscribed topic should be automatically gossiped, defaults to true
    * @param {bool} [options.fallbackToFloodsub] if dial should fallback to floodsub, defaults to true
+   * @param {function} [options.msgIdFn] override the default message id function
+   * @param {Object} [options.messageCache] override the default MessageCache
    * @constructor
    */
   constructor (peerInfo, registrar, options = {}) {
@@ -72,11 +74,16 @@ class GossipSub extends BasicPubsub {
      */
     this.control = new Map()
 
+    /**
+     * Use the overriden mesgIdFn or the default one.
+     */
+    this._msgIdFn = options.msgIdFn || this.defaultMsgIdFn
+
     /**
      * A message cache that contains the messages for last few hearbeat ticks
      *
      */
-    this.messageCache = new MessageCache(constants.GossipSubHistoryGossip, constants.GossipSubHistoryLength)
+    this.messageCache = options.messageCache || new MessageCache(constants.GossipSubHistoryGossip, constants.GossipSubHistoryLength, this._msgIdFn)
 
     /**
      * A heartbeat timer that maintains the mesh
@@ -374,6 +381,17 @@ class GossipSub extends BasicPubsub {
     })
   }
 
+  /**
+   * Override the default implementation in BasicPubSub.
+   * If we don't provide msgIdFn in constructor option, it's the same.
+   * @override
+   * @param {rpc.RPC.Message} msg the message object
+   * @returns {string} message id as string
+   */
+  getMsgId (msg) {
+    return this._msgIdFn(msg)
+  }
+
   _publish (messages) {
     messages.forEach((msgObj) => {
       this.messageCache.put(msgObj)

src/messageCache.js

@@ -1,7 +1,5 @@
 'use strict'
 
-const { utils } = require('libp2p-pubsub')
-
 class CacheEntry {
   /**
    * @param {String} msgID
@@ -19,10 +17,11 @@ class MessageCache {
   /**
    * @param {Number} gossip
    * @param {Number} history
+   * @param {msgIdFn} msgIdFn a function that returns message id from a message
    *
    * @constructor
    */
-  constructor (gossip, history) {
+  constructor (gossip, history, msgIdFn) {
     /**
      * @type {Map<string, RPC.Message>}
      */
@@ -40,6 +39,11 @@ class MessageCache {
      * @type {Number}
      */
     this.gossip = gossip
+
+    /**
+     * @type {Function}
+     */
+    this.msgIdFn = msgIdFn
   }
 
   /**
@@ -49,11 +53,20 @@ class MessageCache {
    * @returns {void}
    */
   put (msg) {
-    const msgID = utils.msgId(msg.from, msg.seqno)
+    const msgID = this.getMsgId(msg)
     this.msgs.set(msgID, msg)
     this.history[0].push(new CacheEntry(msgID, msg.topicIDs))
   }
 
+  /**
+   * Get message id of message.
+   * @param {rpc.RPC.Message} msg
+   * @returns {string}
+   */
+  getMsgId (msg) {
+    return this.msgIdFn(msg)
+  }
+
   /**
    * Retrieves a message from the cache by its ID, if it is still present
    *

src/pubsub.js

@@ -71,6 +71,13 @@ class BasicPubSub extends Pubsub {
     this._options = _options
 
     this._onRpc = this._onRpc.bind(this)
+
+    /**
+     * The default msgID implementation
+     * @param {rpc.RPC.Message} msg the message object
+     * @returns {string} message id as string
+     */
+    this.defaultMsgIdFn = (msg) => utils.msgId(msg.from, msg.seqno)
   }
 
   /**
@@ -166,14 +173,14 @@ class BasicPubSub extends Pubsub {
     if (msgs.length) {
       msgs.forEach(async message => {
         const msg = utils.normalizeInRpcMessage(message)
-        const seqno = utils.msgId(msg.from, msg.seqno)
+        const msgID = this.getMsgId(msg)
 
         // Ignore if we've already seen the message
-        if (this.seenCache.has(seqno)) {
+        if (this.seenCache.has(msgID)) {
           return
         }
 
-        this.seenCache.put(seqno)
+        this.seenCache.put(msgID)
 
         // Ensure the message is valid before processing it
         let isValid
@@ -399,6 +406,15 @@ class BasicPubSub extends Pubsub {
     return Array.from(this.subscriptions)
   }
 
+  /**
+   * Child class can override this.
+   * @param {rpc.RPC.Message} msg the message object
+   * @returns {string} message id as string
+   */
+  getMsgId (msg) {
+    return this.defaultMsgIdFn(msg)
+  }
+
   _emitMessages (topics, messages) {
     topics.forEach((topic) => {
       if (!this.subscriptions.has(topic)) {

test/messageCache.spec.js

@@ -18,7 +18,7 @@ const getMsgID = (msg) => {
 }
 
 describe('Testing Message Cache Operations', () => {
-  const messageCache = new MessageCache(3, 5)
+  const messageCache = new MessageCache(3, 5, getMsgID)
   const testMessages = []
 
   before(() => {